From e506c0ab93a707ef4fa9c2c384da3bee24464e20 Mon Sep 17 00:00:00 2001 From: andsonder Date: Thu, 4 Jan 2024 04:55:48 +0000 Subject: [PATCH 1/2] add dist attr --- .../distributed/passes/auto_parallel_gradient_merge.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py index c793639c5ba01..7e4f8894d90f4 100644 --- a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py +++ b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py @@ -375,8 +375,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( From eb3d08ee045afadc1df3947f38aada4043d597b8 Mon Sep 17 00:00:00 2001 From: andsonder Date: Mon, 8 Jan 2024 06:01:17 +0000 Subject: [PATCH 2/2] add op namescope --- python/paddle/distributed/passes/auto_parallel_gradient_merge.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py index 7e4f8894d90f4..719f733305cac 100644 --- a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py +++ b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py @@ -232,6 +232,7 @@ def _append_gradient_merge_backward_op( attrs={ 'axis': -1, OP_ROLE_KEY: OpRole.Backward, + "op_namescope": "/auto_parallel/gradient_merge", }, )