[Auto-parallel] fix custom op in dynamic mode (#72663)

waliwali777 · web-flow · commit d8dfc9d66357 · 2025-05-26T14:15:48.000+08:00
diff --git a/paddle/fluid/eager/accumulation/accumulation_node.cc b/paddle/fluid/eager/accumulation/accumulation_node.cc
@@ -35,7 +35,7 @@ static void CopyOrAddTensor(paddle::Tensor* tensor,
     VLOG(3) << "Move Tensor ptr: " << t.impl();
     *tensor = t;
   } else {
-    if (!tensor->defined() || !tensor->has_allocation()) {
+    if (!tensor->defined() || !tensor->initialized()) {
       // Simply copy tensor->impl
       VLOG(3) << "Move Tensor ptr: " << t.impl();
       *tensor = t;
diff --git a/paddle/fluid/eager/custom_operator/custom_operator_utils.cc b/paddle/fluid/eager/custom_operator/custom_operator_utils.cc
@@ -803,6 +803,26 @@ void run_custom_op_impl(const paddle::OpMetaInfo& op_info,
   phi::distributed::ProcessMesh current_process_mesh = std::get<2>(result);
   auto& spmd_info = std::get<3>(result);
   if (!rank_is_in_current_mesh) {
+    std::vector<Tensor>* output_all = ctx.AllMutableOutput();
+    for (size_t i = 0; i < output_all->size(); ++i) {
+      auto& tensor = output_all->at(i);
+      phi::distributed::TensorDistAttr dist_attr;
+      if (!spmd_info.second.empty()) {
+        dist_attr = PADDLE_GET_CONST(phi::distributed::TensorDistAttr,
+                                     spmd_info.second[i]);
+      } else {
+        std::vector<int64_t> shape = common::vectorize(output_dims[i]);
+        dist_attr.set_default_dims_mapping(shape);
+        dist_attr.set_process_mesh(current_process_mesh);
+      }
+      auto dense_tensor =
+          std::dynamic_pointer_cast<phi::DenseTensor>(tensor.impl());
+      if (dense_tensor != nullptr) {
+        auto dist_t = std::make_shared<phi::distributed::DistTensor>(
+            dense_tensor, output_dims[i], dist_attr);
+        tensor.set_impl(dist_t);
+      }
+    }
     return;
   }
 #endif
