| 1295 | ////////////////// Backward(grad) api impls ////////////////////// |
| 1296 | |
| 1297 | void embedding_grad_impl(const Tensor& x, |
| 1298 | const Tensor& weight, |
| 1299 | const Tensor& out_grad, |
| 1300 | int64_t padding_idx, |
| 1301 | bool sparse, |
| 1302 | Tensor* weight_grad) { |
| 1303 | DataType kernel_data_type = ParseDataType(weight); |
| 1304 | auto kernel_key_set = ParseKernelKeyByInputArgs(weight); |
| 1305 | auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey(); |
| 1306 | VLOG(6) << "embedding_grad API kernel key: [" << kernel_key.backend() << ", " |
| 1307 | << kernel_key.layout() << ", " << kernel_data_type << "]"; |
| 1308 | |
| 1309 | if (phi::DenseTensor::classof(weight.impl().get()) || |
| 1310 | phi::distributed::DistTensor::classof(weight.impl().get())) { |
| 1311 | std::string kernel_name = |
| 1312 | sparse ? "embedding_sparse_grad" : "embedding_grad"; |
| 1313 | auto kernel_result = |
| 1314 | phi::KernelFactory::Instance().SelectKernelOrThrowError( |
| 1315 | kernel_name, |
| 1316 | {kernel_key.backend(), kernel_key.layout(), kernel_data_type}); |
| 1317 | const auto& kernel = kernel_result.kernel; |
| 1318 | VLOG(6) << kernel_name << " API kernel: " << kernel; |
| 1319 | |
| 1320 | auto* dev_ctx = GetDeviceContextByBackend( |
| 1321 | kernel_result.has_fallback_cpu ? Backend::CPU : kernel_key.backend()); |
| 1322 | |
| 1323 | #ifdef PADDLE_WITH_DISTRIBUTE |
| 1324 | bool run_auto_parallel = AllInputsAreDistTensor(x, weight, out_grad); |
| 1325 | // Auto Parallel condition |
| 1326 | if (run_auto_parallel) { |
| 1327 | bool rank_is_in_current_mesh = true; |
| 1328 | auto mesh = |
| 1329 | std::static_pointer_cast<phi::distributed::DistTensor>(x.impl()) |
| 1330 | ->dist_attr() |
| 1331 | .process_mesh(); |
| 1332 | rank_is_in_current_mesh = phi::distributed::IsCurRankInMesh(mesh); |
| 1333 | |
| 1334 | // 1. InferSpmd (Infer DistAttr of Inputs&Outputs) |
| 1335 | auto meta_dist_input_x = MakeDistMetaTensor(*x.impl()); |
| 1336 | auto meta_dist_input_weight = MakeDistMetaTensor(*weight.impl()); |
| 1337 | auto meta_dist_input_out_grad = MakeDistMetaTensor(*out_grad.impl()); |
| 1338 | auto spmd_info = |
| 1339 | phi::distributed::EmbeddingGradInferSpmd(meta_dist_input_x, |
| 1340 | meta_dist_input_weight, |
| 1341 | meta_dist_input_out_grad, |
| 1342 | padding_idx, |
| 1343 | sparse); |
| 1344 | |
| 1345 | // 2. Create Temporary Output & Prepare Dist and Dense Output |
| 1346 | std::shared_ptr<phi::distributed::DistTensor> shared_dist_out = |
| 1347 | CreateKernelDistOutput( |
| 1348 | weight_grad, !rank_is_in_current_mesh, spmd_info.second[0]); |
| 1349 | phi::distributed::DistTensor* dist_out = shared_dist_out.get(); |
| 1350 | phi::DenseTensor* dense_out = dist_out->unsafe_mutable_value(); |
| 1351 | if (dense_out && !rank_is_in_current_mesh && !dist_out->defined()) { |
| 1352 | *dense_out = phi::DenseTensor( |
| 1353 | std::make_shared<phi::Allocation>( |
| 1354 | nullptr, 0, phi::distributed::GetDefaultPlace()), |
nothing calls this directly
no test coverage detected