PaddlePaddle
diff --git a/‎python/paddle/distributed/auto_parallel/static/reshard.py
Copy file name to clipboard
+307-185Lines changed: 307 additions & 185 deletions b/‎python/paddle/distributed/auto_parallel/static/reshard.py
Copy file name to clipboard
+307-185Lines changed: 307 additions & 185 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/utils.py
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎python/paddle/distributed/auto_parallel/static/utils.py
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
diff --git a/‎python/paddle/distributed/passes/auto_parallel_gradient_merge.py
Copy file name to clipboard
+38-12Lines changed: 38 additions & 12 deletions b/‎python/paddle/distributed/passes/auto_parallel_gradient_merge.py
Copy file name to clipboard
+38-12Lines changed: 38 additions & 12 deletions
diff --git a/‎python/paddle/distributed/passes/auto_parallel_recompute.py
Copy file name to clipboard
+12-2Lines changed: 12 additions & 2 deletions b/‎python/paddle/distributed/passes/auto_parallel_recompute.py
Copy file name to clipboard
+12-2Lines changed: 12 additions & 2 deletions
diff --git a/‎python/paddle/distributed/passes/auto_parallel_sequence_parallel_optimization.py
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎python/paddle/distributed/passes/auto_parallel_sequence_parallel_optimization.py
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎python/paddle/distributed/passes/auto_parallel_sharding.py
Copy file name to clipboard
+19-2Lines changed: 19 additions & 2 deletions b/‎python/paddle/distributed/passes/auto_parallel_sharding.py
Copy file name to clipboard
+19-2Lines changed: 19 additions & 2 deletions
@@ -2210,10 +2210,9 @@ def insert_dependencies_for_vars(
     for post_var in post_vars:
         assert block.has_var(post_var.name)
 
+    post_dist_attr = dist_context.get_tensor_dist_attr_for_program(post_vars[0])
     if process_mesh is None:
-        process_mesh = dist_context.get_tensor_dist_attr_for_program(
-            post_vars[0]
-        ).process_mesh
+        process_mesh = post_dist_attr.process_mesh
     assert process_mesh is not None
 
     use_nop = True
@@ -2245,6 +2244,7 @@ def insert_dependencies_for_vars(
         depend_op_dist_attr.impl_type = "default"
         depend_op_dist_attr.process_mesh = process_mesh
         depend_op_dist_attr.is_recompute = is_recompute
+        depend_op_dist_attr.chunk_id = post_dist_attr.chunk_id
         for input_varname in depend_op.desc.input_arg_names():
             var = block.var(input_varname)
             mapping = dist_context.get_tensor_dist_attr_for_program(
 
@@ -52,10 +52,6 @@ def _remove_and_get_optimizer_op(main_program, dist_context):
             new_op_desc.copy_from(op.desc)
             removed_op_idx.append(idx)
 
-            # del op from dist_context
-            # if dist_context:
-            #     dist_context.del_dist_op_for_program(op)
-
     for idx in removed_op_idx[::-1]:
         main_block._remove_op(idx, sync=False)
     main_block._sync_with_cpp()
@@ -203,6 +199,7 @@ def _append_gradient_merge_backward_op(
                     gradient_merge_var,
                     ref_dims_mapping,
                     ref_process_mesh,
+                    chunk_id=ref_dist_attr.chunk_id,
                 )
 
                 # Add persistable gradient variables in startup_program
@@ -246,6 +243,7 @@ def _append_gradient_merge_backward_op(
                     ref_process_mesh,
                     ref_dims_mapping,
                     dist_context,
+                    chunk_id=ref_dist_attr.chunk_id,
                 )
 
                 del grad_to_params_grads[out_name]
@@ -288,37 +286,55 @@ def true_apply_gradient():
                     },
                 )
                 scale_op._set_attr(OP_ROLE_KEY, OpRole.Optimize)
+                ref_dist_attr = dist_context.get_tensor_dist_attr_for_program(
+                    new_grad
+                )
+                naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
+                    scale_op,
+                    ref_dist_attr.process_mesh,
+                    ref_dist_attr.dims_mapping,
+                    dist_context,
+                    chunk_id=ref_dist_attr.chunk_id,
+                )
 
         # append optimizer ops
         for opt_op_idx in range(optimize_ops_block.desc.op_size()):
             op_desc = optimize_ops_block.desc.op(opt_op_idx)
             new_op_desc = cur_block.desc.append_op()
             new_op_desc.copy_from(op_desc)
+            op_dist_attr = dist_context.get_op_dist_attr_for_program_with_id(
+                new_op_desc.original_id()
+            )
 
             # update input/output
             for input_name in new_op_desc.input_arg_names():
                 if input_name in grad_to_gradient_merge:
+                    in_dims_mapping = op_dist_attr.get_input_dims_mapping(
+                        input_name
+                    )
                     new_op_desc._rename_input(
                         input_name, grad_to_gradient_merge[input_name]
                     )
+                    op_dist_attr.set_input_dims_mapping(
+                        grad_to_gradient_merge[input_name], in_dims_mapping
+                    )
 
             for output_name in new_op_desc.output_arg_names():
                 if output_name in grad_to_gradient_merge:
+                    out_dims_mapping = op_dist_attr.get_output_dims_mapping(
+                        output_name
+                    )
                     new_op_desc._rename_output(
                         output_name, grad_to_gradient_merge[output_name]
                     )
+                    op_dist_attr.set_output_dims_mapping(
+                        grad_to_gradient_merge[output_name], out_dims_mapping
+                    )
 
             # remove op_role_var
             if new_op_desc.has_attr(OP_ROLE_VAR_KEY):
                 new_op_desc.remove_attr(OP_ROLE_VAR_KEY)
 
-            # op's update Grad
-            if core.grad_var_suffix() in new_op_desc.input_arg_names():
-                grad_value = new_op_desc.input("Grad")[0]
-                # TODO FIXME(xym) support fp16
-                grad_merge_value = grad_value + '@MERGE'
-                new_op_desc.set_input("Grad", [grad_merge_value])
-
         main_program.global_block()._sync_with_cpp()
         cur_block._sync_with_cpp()
 
@@ -332,7 +348,7 @@ def true_apply_gradient():
         # clear gradient_merge_vars
         # NOTE(zhaoyingli): Must use 'set_value' op in pir to assign 0-value for persistable var.
         for _, new_grad in new_params_to_grads:
-            cur_block.append_op(
+            set_value_op = cur_block.append_op(
                 type="set_value",
                 inputs={"Input": [new_grad]},
                 outputs={"Out": [new_grad]},
@@ -347,6 +363,16 @@ def true_apply_gradient():
                     OP_ROLE_KEY: OpRole.Optimize,
                 },
             )
+            ref_dist_attr = dist_context.get_tensor_dist_attr_for_program(
+                new_grad
+            )
+            naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
+                set_value_op,
+                ref_dist_attr.process_mesh,
+                ref_dist_attr.dims_mapping,
+                dist_context,
+                chunk_id=ref_dist_attr.chunk_id,
+            )
 
     paddle.static.nn.cond(cond_var, true_fn=true_apply_gradient, false_fn=None)
     cond_op = main_program.global_block().ops[-1]
 
@@ -168,7 +168,11 @@ def modify_forward_desc_for_recompute(self, dist_context):
             ref_dims_mapping = [-1]
             ref_process_mesh = cur_op_dist_attr.process_mesh
             seed_var_dist_attr = set_var_dist_attr(
-                dist_context, seed_var, ref_dims_mapping, ref_process_mesh
+                dist_context,
+                seed_var,
+                ref_dims_mapping,
+                ref_process_mesh,
+                chunk_id=cur_op_dist_attr.chunk_id,
             )
 
             seed = (
@@ -187,7 +191,11 @@ def modify_forward_desc_for_recompute(self, dist_context):
             seed_op._set_attr('op_namescope', cur_op.attr('op_namescope'))
             # set new seed op's dist_attr
             naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
-                seed_op, ref_process_mesh, ref_dims_mapping, dist_context
+                seed_op,
+                ref_process_mesh,
+                ref_dims_mapping,
+                dist_context,
+                chunk_id=cur_op_dist_attr.chunk_id,
             )
 
             # modify dropout op's desc
@@ -494,6 +502,7 @@ def _apply_single_impl(self, main_program, startup_program, context):
                             rc_var,
                             ref_dims_mapping,
                             ref_process_mesh,
+                            chunk_id=cur_op_dist_attr.chunk_id,
                         )
             # get recomputed segment's descs
             segment_descs = _add_needed_descs_to_block(
@@ -611,6 +620,7 @@ def set_op_dist_attr(self, op, old_dist_attr, var_name_dict):
         new_dist_attr.impl_idx = old_dist_attr.impl_idx
         new_dist_attr.impl_type = old_dist_attr.impl_type
         new_dist_attr.process_mesh = old_dist_attr.process_mesh
+        new_dist_attr.chunk_id = old_dist_attr.chunk_id
         for input in old_dist_attr.inputs_dist_attrs.keys():
             if input in var_name_dict.keys():
                 in_dist_attr = old_dist_attr.inputs_dist_attrs[input]
 
@@ -145,7 +145,10 @@ def is_valid_split_op(idx, block):
             )
             ref_process_mesh = allreduce_input_dist_attr.process_mesh
             naive_set_dist_op_attr_for_program_by_mesh(
-                new_op, ref_process_mesh, self.dist_context
+                new_op,
+                ref_process_mesh,
+                self.dist_context,
+                chunk_id=allreduce_input_dist_attr.chunk_id,
             )
 
         # remove vars
 
@@ -497,6 +497,7 @@ def _insert_optimizer_broadcasts(self, main_block, startup_block):
                     param_dist_attr.process_mesh,
                     param_dist_attr.dims_mapping,
                     self._dist_context,
+                    chunk_id=param_dist_attr.chunk_id,
                 )
         main_block._sync_with_cpp()
 
@@ -610,13 +611,23 @@ def _shard_parameter(self, main_block, startup_block):
                                 input_var
                             )
                         )
-                        out_var_dist_attr = set_var_dist_attr(
+                        set_var_dist_attr(
                             self._dist_context,
                             new_var,
                             ref_dist_attr.dims_mapping,
                             ref_dist_attr.process_mesh,
+                            chunk_id=ref_dist_attr.chunk_id,
+                        )
+                        op_dist_attr = (
+                            self._dist_context.get_op_dist_attr_for_program(op)
+                        )
+                        input_dist_attr = op_dist_attr.get_input_dist_attr(
+                            input_name
                         )
                         op._rename_input(input_name, broadcast_varname)
+                        op_dist_attr.set_input_dist_attr(
+                            broadcast_varname, input_dist_attr
+                        )
 
                     _insert_init_and_broadcast_op(
                         main_block,
@@ -1408,6 +1419,7 @@ def _insert_init_and_broadcast_op(
         broadcast_var_dist_attr.process_mesh,
         broadcast_var_dist_attr.dims_mapping,
         dist_context,
+        chunk_id=broadcast_var_dist_attr.chunk_id,
     )
     if local_rank != root_rank:
         new_op = block._insert_op_without_sync(
@@ -1425,6 +1437,7 @@ def _insert_init_and_broadcast_op(
             broadcast_var_dist_attr.process_mesh,
             broadcast_var_dist_attr.dims_mapping,
             dist_context,
+            chunk_id=broadcast_var_dist_attr.chunk_id,
         )
 
 
@@ -1458,7 +1471,11 @@ def _insert_reduce_op(
         block.var(reduce_var)
     )
     naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
-        new_op, dist_attr.process_mesh, dist_attr.dims_mapping, dist_context
+        new_op,
+        dist_attr.process_mesh,
+        dist_attr.dims_mapping,
+        dist_context,
+        chunk_id=dist_attr.chunk_id,
     )
     new_op._set_attr('op_namescope', '/' + ParallelMode.DataParallel)
     return new_op
Original file line number	Diff line number	Diff line change
`@@ -145,7 +145,10 @@ def is_valid_split_op(idx, block):`
`145`	`145`	`)`
`146`	`146`	`ref_process_mesh = allreduce_input_dist_attr.process_mesh`
`147`	`147`	`naive_set_dist_op_attr_for_program_by_mesh(`
`148`		`- new_op, ref_process_mesh, self.dist_context`
	`148`	`+ new_op,`
	`149`	`+ ref_process_mesh,`
	`150`	`+ self.dist_context,`
	`151`	`+ chunk_id=allreduce_input_dist_attr.chunk_id,`
`149`	`152`	`)`
`150`	`153`
`151`	`154`	`# remove vars`