fix the limitation of fthenb schedule (#60134)

ForFishes · web-flow · commit 86031cdb0a3a · 2023-12-19T14:10:55.000+08:00
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -1288,13 +1288,9 @@ def forward_backward_pipeline(
         self._forward_only = forward_only
 
         assert (
-            self.accumulate_steps >= self.num_stages
-        ), "accumulate_steps({}) should be larger than num_stages({}) for pipeline with interleave".format(
-            self.accumulate_steps, self.num_stages
-        )
-        assert (
-            self.accumulate_steps < 2 * self.num_stages
-        ), "accumulate_steps({}) should be smaller than 2 * num_stages({}) for pipeline with interleave".format(
+            self.accumulate_steps == self.num_stages
+            or self.accumulate_steps % self.num_stages != 0
+        ), "accumulate_steps({}) and num_stages({}) should be a multiple or accumulate_steps % num_stages == 0".format(
             self.accumulate_steps, self.num_stages
         )
 
diff --git a/python/paddle/distributed/fleet/model.py b/python/paddle/distributed/fleet/model.py
@@ -168,17 +168,16 @@ def forward(self, x):
             accumulate_steps = strategy.pipeline_configs['accumulate_steps']
             pp_degree = fleet_env._hcg.get_pipe_parallel_world_size()
             if (
-                accumulate_steps >= pp_degree
-                and accumulate_steps < pp_degree * 2
+                accumulate_steps > pp_degree
+                and accumulate_steps % pp_degree == 0
             ):
-                # NOTE(shenliang03): Hacky for unbalanced pipeline parallel with interleave
-                # Currently, we only support pp_degree <= accumulate_steps < 2 * pp_degree
-                model = PipelineParallelWithInterleaveFthenB(
+                # interleave pipeline
+                model = PipelineParallelWithInterleave(
                     model, fleet_env._hcg, strategy=strategy
                 )
             else:
-                # interleave pipeline
-                model = PipelineParallelWithInterleave(
+                # NOTE(shenliang03): Hacky for unbalanced pipeline parallel with interleave
+                model = PipelineParallelWithInterleaveFthenB(
                     model, fleet_env._hcg, strategy=strategy
                 )
 
