Make returning moe_losses optional in SequentialWrapper (#1336)

* - Adds a flag to check whether to return moe losses in SequentialWrapper * pre-commit --------- Co-authored-by: Nawras Alnaasan <[email protected]> Co-authored-by: Quentin Anthony <[email protected]>
EleutherAI · Feb 8, 2025 · d7cc7aa · d7cc7aa
1 parent f7a5a6f
commit d7cc7aa
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/megatron/model/utils.py b/megatron/model/utils.py
@@ -141,7 +141,12 @@ 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 +217,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):
         """