[AutoParallel] Fit vpp for gradient_merge pass (#60560)

* add dist attr * add op namescope
PaddlePaddle · Jan 9, 2024 · 5e35868 · 5e35868
1 parent 10a11d8
commit 5e35868
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -233,6 +233,7 @@ def _append_gradient_merge_backward_op(
                     attrs={
                         'axis': -1,
                         OP_ROLE_KEY: OpRole.Backward,
+                        "op_namescope": "/auto_parallel/gradient_merge",
                     },
                 )
 
@@ -376,8 +377,16 @@ def true_apply_gradient():
             )
 
     paddle.static.nn.cond(cond_var, true_fn=true_apply_gradient, false_fn=None)
+    cond_dist_attr = dist_context.get_tensor_dist_attr_for_program(cond_var)
     cond_op = main_program.global_block().ops[-1]
     cond_op._set_attr(OP_ROLE_KEY, OpRole.Optimize)
+    naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
+        cond_op,
+        process_mesh=cond_dist_attr.process_mesh,
+        ref_mapping=cond_dist_attr.dims_mapping,
+        ctx=dist_context,
+        chunk_id=cond_dist_attr.chunk_id,
+    )
 
 
 def parse_program(