- Adds a flag to check whether to return moe losses in SequentialWrapper

EleutherAI · Feb 7, 2025 · b1e3a7d · b1e3a7d
1 parent f7a5a6f
commit b1e3a7d
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/megatron/model/utils.py b/megatron/model/utils.py
@@ -141,7 +141,7 @@ def train_mode(self):
         recursive_setattr(self.sequential, "training", True)
 
     def forward(
-        self, forward_input, curriculum_seqlen=None, labels=None, neox_args=None
+        self, forward_input, curriculum_seqlen=None, labels=None, neox_args=None, return_moe_losses=False
     ):
 
         if self.batch_fn:
@@ -212,7 +212,10 @@ def exec_func(*inputs):
                     )
                 else:
                     x = exec_range_func(start_idx, end_idx)(*x)
-        return x, moe_losses
+        if return_moe_losses:
+            return x, moe_losses
+        else:
+            return x
 
     def clear_cache(self):
         """