fix the limitation of fthenb schedule (#60134) (#60315)

PaddlePaddle · Dec 26, 2023 · 3d2b2de · 3d2b2de
1 parent a726569
commit 3d2b2de
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 14 deletions.
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -1566,13 +1566,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
@@ -162,17 +162,16 @@ def distributed_model(model):
             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
                 )