[AutoParallel] Fit vpp for amp (#68545)

AndSonder · web-flow · commit a521f82a4b6d · 2024-10-09T19:00:56.000+08:00
* add enable send recv

* fix for vpp when open fused pass

* fix

* fix

* Update ir_backward.py

* format code

* fix vpp when using amp master grad

* recover third_party

* fix grad merge

* recover some changes
diff --git a/python/paddle/amp/grad_scaler.py b/python/paddle/amp/grad_scaler.py
@@ -238,7 +238,19 @@ def scale(self, var: Tensor) -> Tensor:
                 var = var.astype('float32')
             if not self._use_dynamic_loss_scaling:
                 return var
-            return var * self._scale
+            scale_out = paddle._C_ops.multiply(var, self._scale)
+            multiply_op = scale_out.get_defining_op()
+            src_var_op = var.get_defining_op()
+            if multiply_op.dist_attr and src_var_op.dist_attr:
+                multiply_op.dist_attr = (
+                    paddle.base.libpaddle.pir.create_op_dist_attribute(
+                        multiply_op.dist_attr.process_mesh,
+                        multiply_op.dist_attr.operands(),
+                        multiply_op.dist_attr.results(),
+                        src_var_op.dist_attr.chunk_id,
+                    )
+                )
+            return scale_out
 
         # NOTE(lizhiyu): We hack here to avoid changing the `dist_attr` of `self._scale` of 'no-calculation-rank'
         if not self._enable or not var._is_initialized():
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -732,7 +732,9 @@ def _parallel_pir(self, mode):
             self._strategy.pipeline.enable
             and self._strategy.pipeline.schedule_mode == "VPP"
         ):
-            complete_chunk_id(dist_program, self._strategy.pipeline)
+            complete_chunk_id(
+                dist_program, startup_program, self._strategy.pipeline
+            )
 
         # Step 1.2: pir backward
         if mode == "train" and self._loss and self._optimizer:
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -24,6 +24,7 @@
 from paddle.base.log_helper import get_logger
 from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.distributed.passes.pass_base import PassContext, new_pass
+from paddle.distributed.passes.pass_utils import infer_chunk_id
 
 from .mix_to_dist_pass import dist_skip_op_list
 from .process_group import get_process_group
@@ -811,7 +812,7 @@ def set_process_mesh(vars, attrs):
     )
 
 
-def complete_chunk_id(dist_program, pipeline_strategy):
+def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
     if not pipeline_strategy.enable:
         return
 
@@ -975,6 +976,14 @@ def complete_chunk_id(dist_program, pipeline_strategy):
     # Step6: add reshard op between pipeline chunks
     apply_partition_pass(dist_program)
 
+    for op in startup_program.global_block().ops:
+        if op.name() == "builtin.set_parameter":
+            param_name = op.str_attr("parameter_name")
+            startup_param = op.operand_source(0)
+            param = dist_program.get_parameter_value_by_name(param_name)
+            if param.dist_attr():
+                startup_param.update_dist_attr(param.dist_attr())
+
 
 def check_chunk_id(dist_program):
     all_ops = dist_program.global_block().ops
@@ -993,22 +1002,23 @@ def check_chunk_id(dist_program):
                     all_used_ops = op.result(0).all_used_ops()
                     for used_op in all_used_ops:
                         if used_op.dist_attr.chunk_id != -1:
-                            op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
-                                op.dist_attr.process_mesh,
-                                op.dist_attr.operands(),
-                                op.dist_attr.results(),
-                                used_op.dist_attr.chunk_id,
+                            op.dist_attr = copy_op_attr_with_new_member(
+                                op.dist_attr,
+                                new_chunk_id=used_op.dist_attr.chunk_id,
                             )
                             break
-                    if op.dist_attr.chunk_id == -1:
-                        raise ValueError(
-                            f"The chunk_id of op[{op.name()}] is not set. Please check the chunk_id setting."
-                        )
+
                 else:
-                    raise ValueError(
-                        f"The chunk_id of op[{op.name()}] is not set. Please check the chunk_id setting."
+                    op_chunk_id = infer_chunk_id(idx, all_ops)
+                    op.dist_attr = copy_op_attr_with_new_member(
+                        op.dist_attr, new_chunk_id=op_chunk_id
                     )
 
+            if op.dist_attr.chunk_id == -1:
+                raise ValueError(
+                    f"The chunk_id of op[{op.name()}] is not set. Please check the chunk_id setting."
+                )
+
 
 def check_order(op_list, order):
     pointer = 0
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/sub_to_global_mesh_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/sub_to_global_mesh_func.py
@@ -63,14 +63,17 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
 
         if cur_rank == root_rank:
             # the root rank will broadcast the src_value to other ranks
+            chunk_id = -1
+            if src_value.get_defining_op().dist_attr:
+                chunk_id = src_value.get_defining_op().dist_attr.chunk_id
             tmp_value = paddle._C_ops.share_data_(src_value)
             value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
                 src_value.type(), src_value.dist_attr()
             )
             tmp_value.set_type(value_type)
             op = tmp_value.get_defining_op()
             op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
-                src_mesh, [src_dist_attr], [src_dist_attr]
+                src_mesh, [src_dist_attr], [src_dist_attr], chunk_id
             )
         else:
             # create the buffer on other ranks for receving the data
