From 8407c203a3bd51e3e8a68017a3b7981de445e704 Mon Sep 17 00:00:00 2001 From: zhanglei Date: Tue, 22 Aug 2023 10:53:21 +0800 Subject: [PATCH] refactor code --- internlm/core/scheduler/pipeline_scheduler.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/internlm/core/scheduler/pipeline_scheduler.py b/internlm/core/scheduler/pipeline_scheduler.py index d3a1bd6..055bb92 100644 --- a/internlm/core/scheduler/pipeline_scheduler.py +++ b/internlm/core/scheduler/pipeline_scheduler.py @@ -240,8 +240,16 @@ class PipelineScheduler(BaseScheduler): """ return step_id - def _forward_step(self, engine, input_obj, return_tensors, return_output_label=True, - accum_loss=None, accum_moe_loss=None, moe_loss_coeff=1.0): + def _forward_step( + self, + engine, + input_obj, + return_tensors, + return_output_label=True, + accum_loss=None, + accum_moe_loss=None, + moe_loss_coeff=1.0, + ): """ Forward step for passed-in model. If it is the first stage, the input tensor is obtained from data_iterator, otherwise the passed-in input_obj is used. @@ -620,8 +628,9 @@ class PipelineScheduler(BaseScheduler): return output, label, accum_loss, accum_moe_loss - def forward_backward_step(self, engine, data_iter, forward_only=False, return_loss=True, - return_output_label=True, moe_loss_coeff=1.0): + def forward_backward_step( + self, engine, data_iter, forward_only=False, return_loss=True, return_output_label=True, moe_loss_coeff=1.0 + ): """Runs non-interleaved 1F1B schedule, with communication between pipeline stages. Returns a tuple with losses if the last stage, an empty tuple otherwise. @@ -1286,8 +1295,9 @@ class InterleavedPipelineScheduler(PipelineScheduler): # 3. Cooldown self._run_cooldown_loop(engine, num_microsteps, num_1f1b_micropairs=num_1f1b_micropairs) - def forward_backward_step(self, engine, data_iter, forward_only=False, return_loss=True, - return_output_label=True, moe_loss_coeff=1.0): + def forward_backward_step( + self, engine, data_iter, forward_only=False, return_loss=True, return_output_label=True, moe_loss_coeff=1.0 + ): """Run interleaved 1F1B schedule (model split into model chunks), with communication between pipeline stages as needed.