adapted for sequence parallel (#163)

2025-09-25 03:31:56 +00:00 · 2022-01-20 13:44:51 +08:00
parent a2e649da39
commit e2089c5c15
17 changed files with 432 additions and 119 deletions
--- a/colossalai/engine/schedule/_pipeline_schedule.py
+++ b/colossalai/engine/schedule/_pipeline_schedule.py
@@ -222,7 +222,6 @@ class PipelineSchedule(BaseSchedule):

        assert forward_only or return_loss, \
            'The argument \'return_loss\' has to be True when \'forward_only\' is False, but got False.'
-
        self.load_batch(data_iter)
        num_warmup_microbatches = \
            (gpc.get_world_size(ParallelMode.PIPELINE) -