[AutoParallel] Fit pp and vpp for local layer (#70831)

AndSonder · winter-wang · web-flow · commit ad9acb0770a4 · 2025-01-20T11:18:02.000+08:00
* add local layer api

* update

* add test

* update LocalLayer

* recover dy_local_view_compute test

* recover cmake list

* update notes

* update local layer

* update notes

* fix typing check

* fix typing check

* fix typing check

* remove spec

* fix

* fix type checking

* fix type checking

* fix example code style

* support local view for pp and vpp

* remove print

* fix

* fix

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc b/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
@@ -165,8 +165,14 @@ void ProcessDistBlock(pir::Block* block) {
 
     int64_t chunk_id = -1;
     if (op_item->HasAttribute(kAttrOpDistAttr)) {
-      chunk_id = op_item->attribute<OperationDistAttribute>(kAttrOpDistAttr)
-                     .chunk_id();
+      if (op_item->HasAttribute("chunk_id")) {
+        chunk_id = op_item->attribute("chunk_id")
+                       .dyn_cast<pir::Int32Attribute>()
+                       .data();
+      } else {
+        chunk_id = op_item->attribute<OperationDistAttribute>(kAttrOpDistAttr)
+                       .chunk_id();
+      }
       op_item->erase_attribute(kAttrOpDistAttr);
     }
     op_item->set_attribute("chunk_id", pir::Int32Attribute::get(ctx, chunk_id));
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -1047,7 +1047,9 @@ def infer_dist_skip_op_chunk_id(op):
 
     # TODO(Ruibiao): Reorganize these unclear codes about chunk_id
     def get_op_chunk_id(op):
-        if op.dist_attr is None:
+        if op.has_attr("chunk_id"):
+            op_chunk_id = op.chunk_id
+        elif op.dist_attr is None:
             op_chunk_id = -1
             if op.name() in dist_skip_op_list:
                 op_chunk_id = infer_dist_skip_op_chunk_id(op)
@@ -1077,7 +1079,10 @@ def get_op_chunk_id(op):
         fwd_op_chunk_id = get_op_chunk_id(op)
 
         for bwd_op in state.op_to_opgrad[op]:
-            if bwd_op.dist_attr is None:
+            if op.has_attr("chunk_id"):
+                bwd_op.set_int_attr("chunk_id", op.chunk_id)
+                continue
+            elif bwd_op.dist_attr is None:
                 continue
 
             if bwd_op.name() in ["pd_op.add_", "pd_op.add_n_"]:
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -425,7 +425,13 @@ def remove_other_rank_op_pass(dist_program):
         # pruning op and value not belong to cur rank
         def prune_op(block):
             cur_rank = paddle.distributed.get_rank()
-            for op in block.ops[::-1]:
+            reverse_block_ops = block.ops[::-1]
+            skip_idx = 0
+            for idx, op in enumerate(reverse_block_ops):
+                if idx < skip_idx:
+                    continue
+                skip_idx += 1
+
                 if op.name() == "dist_op.moe_sub_mesh_tensors":
                     replace_moe_sub_mesh_tensors(op)
                     continue
@@ -459,6 +465,22 @@ def prune_op(block):
                     # update pylayer op's inputs
                     op.as_pylayer_op().update_input()
                     continue
+                elif op.name() == "dist_op.dtensor_from_local":
+                    dtensor_to_local_idx = idx
+                    while (
+                        reverse_block_ops[dtensor_to_local_idx].name()
+                        != "dist_op.dtensor_to_local"
+                    ):
+                        dtensor_to_local_idx += 1
+                    if (
+                        op.dist_attr
+                        and cur_rank
+                        not in op.dist_attr.process_mesh.process_ids
+                    ):
+                        for i in range(idx, dtensor_to_local_idx + 1):
+                            reverse_block_ops[i].erase()
+                    skip_idx = dtensor_to_local_idx + 1
+                    continue
                 elif op.name() in partition_skip_op_list:
                     can_delete = True
                     for val in op.results():
@@ -821,7 +843,7 @@ def _get_seg_struct_names(ops, seg_method):
     seg_op_mesh = collections.OrderedDict()
 
     for i in range(fwd_start_op_index, fwd_end_op_index + 1):
-        if ops[i].name() in dist_skip_op_list:
+        if ops[i].dist_attr is None:
             continue
 
         struct_name = _extract_seg_method(ops[i], seg_method)
@@ -847,7 +869,7 @@ def _analyze_use_custom_mesh(ops, seg_method, pp_degree):
     seg_pp_stages = [-1]
 
     for op in ops:
-        if _extract_seg_method(op, seg_method) and "pd_op" in op.name():
+        if _extract_seg_method(op, seg_method) and op.dist_attr:
             op_mesh = op.dist_attr.process_mesh
             pp_stage = get_pp_stage_by_process_mesh(op_mesh, pp_degree)
             if pp_stage is None:
@@ -866,7 +888,13 @@ def _analyze_use_custom_mesh(ops, seg_method, pp_degree):
     return non_use_custom_mesh
 
 
-def _set_process_mesh_and_chunk_id(op, chunk_process_mesh, chunk_id, set_mesh):
+def _set_process_mesh_and_chunk_id(
+    op,
+    chunk_process_mesh,
+    chunk_id,
+    set_input_mesh=False,
+    set_output_mesh=False,
+):
     def set_var_origin_op_process_mesh(var_origin_op):
         var_origin_op_input_attr = var_origin_op.dist_attr.operands()
         var_origin_op_output_attr = var_origin_op.dist_attr.results()
@@ -1013,6 +1041,7 @@ def set_process_mesh(vars, attrs, process_mesh):
 
     op_input_vars = op.operands_source()
     op_output_vars = op.results()
+
     # NOTE(zhangwl):dist_skip_op donnot have op_mesh
     op_mesh = None
     if op.name() in dist_skip_op_list:
@@ -1030,10 +1059,13 @@ def set_process_mesh(vars, attrs, process_mesh):
     op_mesh = op_dist_attr.process_mesh
     op_input_attrs = op_dist_attr.operands()
     op_output_attrs = op_dist_attr.results()
+
     # if op in seq_chunk , vpp need set var and op chunk_process_mesh and chunk_id
-    if set_mesh:
+    if set_input_mesh:
         set_process_mesh(op_input_vars, op_input_attrs, chunk_process_mesh)
+    if set_output_mesh:
         set_process_mesh(op_output_vars, op_output_attrs, chunk_process_mesh)
+    if set_input_mesh or set_output_mesh:
         op_mesh = chunk_process_mesh
 
     op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
@@ -1069,6 +1101,7 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
         dist_program.global_block().ops, seg_method
     )
     ops = dist_program.global_block().ops
+
     # Step2: analysis whether the pp_stage is non-decreasing among segments
     # 1. if non_use_custom_mesh is True, the ops' process_mesh will be changed by vpp strategy
     # 2. if non_use_custom_mesh is False, the ops's process_mesh will not be changed.
@@ -1079,9 +1112,9 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
     seg_chunk_ids = [i // pp_degree for i in range(num_chunks)]
     seg_parts = [0]
     last_struct_name = None
-    stage_ids = (
-        []
-    )  # stage_ids[i] represents the stage number assigned to the i-th layer.
+    # stage_ids[i] represents the stage number assigned to the i-th layer.
+    stage_ids = []
+
     for idx, op in enumerate(ops):
         if len(seg_parts) == len(seg_struct_names):
             break
@@ -1096,12 +1129,14 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
         if struct_name == seg_struct_names[len(seg_parts)]:
             seg_parts.append(idx)
     seg_parts.append(len(ops))
+
     pp_stage_layer_nums = [0] * pp_degree
     for i in stage_ids:
         pp_stage_layer_nums[i] = pp_stage_layer_nums[i] + 1
     assert all(
         value >= vpp_degree for value in pp_stage_layer_nums
     ), "The number of layers on each pp_stage must not be less than the vpp_degree in the pp_stage to ensure that each chunk contains at least one layer."
+
     seg_layer_num = [0] * num_chunks
     for pp_stage in range(
         0, pp_degree
@@ -1112,6 +1147,7 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
             virtual_chunk_id = i % vpp_degree
             real_chunk_id = (virtual_chunk_id) * pp_degree + pp_stage
             seg_layer_num[real_chunk_id] = seg_layer_num[real_chunk_id] + 1
+
     # Step4: Set the process_mesh of each op
     seg_id = 0
     reshard_ops = []
@@ -1139,12 +1175,42 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
             f"start op: [{ops[start_idx].name()}], end op: [{ops[end_idx - 1].name()}]"
         )
 
+        skip_idx = 0
         for idx in range(start_idx, end_idx):
+            if idx < skip_idx:
+                continue
+
+            is_seg_op = _extract_seg_method(ops[idx], seg_method) is not None
+            set_mesh = non_use_custom_mesh & is_seg_op
+
             if ops[idx].name() == "dist_op.reshard":
                 reshard_ops.append(ops[idx])
                 continue
+            elif ops[idx].name() == "dist_op.dtensor_to_local":
+                _set_process_mesh_and_chunk_id(
+                    ops[idx],
+                    process_mesh,
+                    chunk_id,
+                    set_input_mesh=set_mesh,
+                )
+                dtensor_from_local_idx = idx + 1
+                while (
+                    ops[dtensor_from_local_idx].name()
+                    != "dist_op.dtensor_from_local"
+                ):
+                    dtensor_from_local_idx += 1
+                for local_op_idx in range(idx + 1, dtensor_from_local_idx):
+                    ops[local_op_idx].set_int_attr("chunk_id", chunk_id)
+
+                _set_process_mesh_and_chunk_id(
+                    ops[dtensor_from_local_idx],
+                    process_mesh,
+                    chunk_id,
+                    set_output_mesh=set_mesh,
+                )
+                skip_idx = dtensor_from_local_idx + 1
+                continue
 
-            is_seg_op = _extract_seg_method(ops[idx], seg_method) is not None
             for sub_block in ops[idx].blocks():
                 # TODO(luchang): support condition block
                 pass
@@ -1153,8 +1219,10 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
                 ops[idx],
                 process_mesh,
                 chunk_id,
-                non_use_custom_mesh & is_seg_op,
+                set_input_mesh=set_mesh,
+                set_output_mesh=set_mesh,
             )
+            skip_idx = idx + 1
 
     # Step5: set right process_mesh for reshard op
     for op in reshard_ops:
@@ -1243,13 +1311,23 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
 
 def check_chunk_id(dist_program):
     all_ops = dist_program.global_block().ops
-
     for idx, op in enumerate(all_ops):
         if op.op_role in [int(OpRole.Forward), int(OpRole.Backward)]:
             if op.name() in dist_skip_op_list:
                 continue
 
-            if op.dist_attr.chunk_id == -1:
+            if op.has_attr("chunk_id"):
+                # op between dtensor_from_local and dtensor_to_local will
+                # be assigned a chunk_id attribute.
+                continue
+            elif op.name() in [
+                "dist_op.dtensor_from_local",
+                "dist_op.dtensor_to_local",
+            ]:
+                # dtensor_from_local and dtensor_to_local ops will be removed after
+                # converting the program to a dense program.
+                continue
+            elif op.dist_attr.chunk_id == -1:
                 if op.name() in ["pd_op.data", "builtin.parameter"]:
                     op.dist_attr = copy_op_attr_with_new_member(
                         op.dist_attr, new_chunk_id=0
diff --git a/python/paddle/distributed/auto_parallel/static/utils.py b/python/paddle/distributed/auto_parallel/static/utils.py
@@ -2374,8 +2374,8 @@ def get_sub_process_mesh_by_program(dist_program):
     all_ops = dist_program.global_block().ops
     process_meshes = []
 
-    for op in all_ops:
-        if "pd_op" in op.name():
+    for idx, op in enumerate(all_ops):
+        if "pd_op" in op.name() and op.dist_attr:
             process_mesh = op.dist_attr.process_mesh
             if process_mesh not in process_meshes:
                 process_meshes.append(process_mesh)
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -817,7 +817,7 @@ def get_chunk_id(op_idx):
                 return op.dist_attr.chunk_id
         else:
             if op.has_attr("chunk_id"):
-                return op.attrs()["chunk_id"]
+                return op.chunk_id
             else:
                 return -1
 
@@ -836,11 +836,8 @@ def get_chunk_id(op_idx):
         for used_op in all_used_ops:
             if used_op.dist_attr and used_op.dist_attr.chunk_id != -1:
                 return used_op.dist_attr.chunk_id != -1
-            elif (
-                used_op.has_attr("chunk_id")
-                and used_op.attrs()["chunk_id"] != -1
-            ):
-                return used_op.attrs()["chunk_id"]
+            elif used_op.has_attr("chunk_id") and used_op.chunk_id != -1:
+                return used_op.chunk_id
 
     return -1
 
@@ -1966,7 +1963,7 @@ def _pir_split_matmul_grad_to_matmul(block, matmul_grad_id):
     # When the rank of input matrix is 3, MatmulGradKernel use reshape to fold the first two dimensions of x and out_grad (see FoldInitDims in matmul_grad_kernel_impl.h), and then calls blas.Matmul to calculate y_grad.
     # If we directly append matmul op to calculate y_grad without FoldInitDims, blas.BatchedGEMM is actually called in MatmulKernel, which has a larger cost than using blas.Matmul after dimension folding.
     # Therefore, we imitate MatmulGradKernel here by inserting reshape op before matmul.
-    chunk_id = matmul_grad_op.attrs()["chunk_id"]
+    chunk_id = matmul_grad_op.chunk_id
 
     paddle.pir.set_insertion_point_after(matmul_grad_op)
     new_x = paddle._C_ops.reshape(x, new_x_dims)
diff --git a/test/auto_parallel/pir/CMakeLists.txt b/test/auto_parallel/pir/CMakeLists.txt
@@ -84,6 +84,9 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 300)
   py_test_modules(test_pir_moe_utils_api MODULES test_pir_moe_utils_api ENVS
                   FLAGS_enable_pir_api=1)
+  py_test_modules(
+    test_pipeline_scheduler_vpp_local_view MODULES
+    test_pipeline_scheduler_vpp_local_view ENVS FLAGS_enable_pir_api=1)
 endif()
 py_test_modules(test_pir_1f1b_plan MODULES test_pir_1f1b_plan ENVS
                 FLAGS_enable_pir_api=1)
diff --git a/test/auto_parallel/pir/test_pipeline_scheduler_vpp_local_view.py b/test/auto_parallel/pir/test_pipeline_scheduler_vpp_local_view.py
@@ -0,0 +1,36 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestVPPLocalView(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=300,
+        )
+        self._default_envs = {"dtype": "float32", "seed": "2024"}
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_pp(self):
+        self.run_test_case(
+            "vpp_pass_unittest_local_view_pir.py",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/auto_parallel/pir/vpp_pass_unittest_local_view_pir.py b/test/auto_parallel/pir/vpp_pass_unittest_local_view_pir.py
