fix the moe_loss_coeff bug

2023-08-15 18:47:19 +08:00 · 2023-08-15 18:47:19 +08:00 · 92a31732f9
parent 8c7d868f01
commit 92a31732f9
1 changed files with 1 additions and 0 deletions
--- a/internlm/core/scheduler/pipeline_scheduler.py
+++ b/internlm/core/scheduler/pipeline_scheduler.py
@ -532,6 +532,7 @@ class PipelineScheduler(BaseScheduler):
                return_tensors,
                return_output_label=return_output_label,
                accum_loss=accum_loss,
+                moe_loss_coeff=moe_loss_coeff,
            )

            if gpc.is_last_rank(ParallelMode.PIPELINE):