【Dist-Dy2Stat】fix bug of process_mesh in concat (#65820)

liym27 · web-flow · commit f254aa19e3ac · 2024-07-09T11:18:14.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/reshard.py b/python/paddle/distributed/auto_parallel/static/reshard.py
@@ -2188,12 +2188,12 @@ def is_grad(name):
                             self.dist_context,
                             out_var,
                             [-1] * len(out_var.shape),
-                            src_tensor_attr.process_mesh,
+                            dst_input_attr[0],  # process_mesh
                             chunk_id=src_tensor_attr.chunk_id,
                         )
                     naive_set_dist_op_attr_for_program_by_mesh(
                         op,
-                        src_tensor_attr.process_mesh,
+                        dst_input_attr[0],  # process_mesh
                         self.dist_context,
                         chunk_id=src_tensor_attr.chunk_id,
                     )
