支持vpp去尾操作 (#70164)

zty-king · web-flow · commit 16cd479e037b · 2024-12-13T14:26:38.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -1035,9 +1035,10 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
     )
     ops = dist_program.global_block().ops
 
-    assert (
-        len(seg_struct_names) % num_chunks == 0
-    ), f"The number of layers[{seg_method}] ({len(seg_struct_names)}) should be divided by part number ({num_chunks})."
+    assert (len(seg_struct_names) % num_chunks == 0) or (
+        (len(seg_struct_names) + 1) % num_chunks == 0
+        and (len(seg_struct_names) + 1) // num_chunks != 1
+    ), f"The number of layers[{seg_method}] ({len(seg_struct_names)}) should be divisible by part number ({num_chunks}),or ({len(seg_struct_names)} + 1) should be divisible by {num_chunks} and not equal to {num_chunks}."
 
     # 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
@@ -1047,7 +1048,10 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
     # Step3: Get op index boundary, pp_stage, chunk_id, struct_names of each segment
     seg_pp_stages = [i % pp_degree for i in range(num_chunks)]
     seg_chunk_ids = [i // pp_degree for i in range(num_chunks)]
-    seg_layer_num = len(seg_struct_names) // num_chunks
+    seg_layer_num = [0] * num_chunks
+    for j in range(0, len(seg_struct_names)):
+        i = j % num_chunks
+        seg_layer_num[i] = seg_layer_num[i] + 1
     seg_parts = [0]
 
     for idx, op in enumerate(ops):
@@ -1061,16 +1065,21 @@ def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
     # Step4: Set the process_mesh of each op
     seg_id = 0
     reshard_ops = []
+    previous_seg_parts_end_idx = 0
     for seg_id in range(num_chunks):
-        start_idx = seg_parts[seg_id * seg_layer_num]
-        end_idx = seg_parts[seg_id * seg_layer_num + seg_layer_num]
+        start_idx = seg_parts[previous_seg_parts_end_idx]
+        end_idx = seg_parts[previous_seg_parts_end_idx + seg_layer_num[seg_id]]
         pp_stage = seg_pp_stages[seg_id]
         chunk_id = seg_chunk_ids[seg_id]
         struct_name = ",".join(
             seg_struct_names[
-                seg_id * seg_layer_num : seg_id * seg_layer_num + seg_layer_num
+                previous_seg_parts_end_idx : previous_seg_parts_end_idx
+                + seg_layer_num[seg_id]
             ]
         )
+        previous_seg_parts_end_idx = (
+            previous_seg_parts_end_idx + seg_layer_num[seg_id]
+        )
         process_mesh = sub_process_meshes[pp_stage]
 
         _logger.info(
diff --git a/test/auto_parallel/pir/vpp_pass_unittest_pir.py b/test/auto_parallel/pir/vpp_pass_unittest_pir.py
@@ -102,6 +102,7 @@ def __init__(
         dropout_ratio=0.1,
         initializer_range=0.02,
         manual=True,
+        hidden_layer=4,
     ):
         super().__init__()
 
@@ -112,7 +113,9 @@ def __init__(
         if manual:
             self.layer_to_mesh = [PP_MESH_0, PP_MESH_1, PP_MESH_0, PP_MESH_1]
         else:
-            self.layer_to_mesh = [PP_MESH_0, PP_MESH_0, PP_MESH_1, PP_MESH_1]
+            self.layer_to_mesh = [PP_MESH_0] * (
+                hidden_layer - hidden_layer // 2
+            ) + [PP_MESH_1] * (hidden_layer // 2)
 
         self.layers = nn.LayerList(
             [
@@ -123,7 +126,7 @@ def __init__(
                     weight_attr,
                     mesh=self.layer_to_mesh[i],
                 )
-                for i in range(4)
+                for i in range(hidden_layer)
             ]
         )
 
@@ -217,11 +220,12 @@ def run_pipeline(
         manual=True,
         enable_send_recv_overlap=False,
         batch_size=BATCH_SIZE,
+        hidden_layer=4,
     ):
         self.init()
 
         strategy = apply_pass(schedule_mode, acc_step, enable_send_recv_overlap)
-        model = MLPLayer(manual=manual)
+        model = MLPLayer(manual=manual, hidden_layer=hidden_layer)
         opt = paddle.optimizer.AdamW(
             learning_rate=0.00001, parameters=model.parameters()
         )
@@ -267,6 +271,11 @@ def test_pp_pass(self):
         )
         self.check_result(Non_uniform_loss_vpp, Non_uniform_loss_vpp_manual)
         self.check_result(loss_fthenb, Non_uniform_loss_vpp)
+        # Tail-removed-vpp
+        Tail_removed_loss_vpp = self.run_pipeline(
+            schedule_mode="VPP", acc_step=4, manual=False, hidden_layer=7
+        )
+        self.check_result(Tail_removed_loss_vpp, loss_vpp)
 
     def check_result(self, loss1, loss2):
         return np.array_equal(loss1, loss2)
