add do_not_average_loss arg

Signed-off-by: ashors1 <[email protected]>

Author: ashors1

nemo_rl/models/megatron/train.py

@@ -180,6 +180,7 @@ def megatron_forward_backward(
     defer_fp32_logits: Optional[bool] = None,
     global_valid_seqs: Optional[torch.Tensor] = None,
     global_valid_toks: Optional[torch.Tensor] = None,
+    do_not_average_loss: bool = False,
 ) -> Any:
     """
     Execute forward and backward passes using Megatron's utilities.
@@ -222,6 +223,7 @@ def megatron_forward_backward(
         micro_batch_size=mbs,
         decoder_seq_length=seq_length,
         forward_only=forward_only,
+        do_not_average_loss=do_not_average_loss,
     )
 
 class LossPostProcessor:

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -963,10 +963,10 @@ def train(
                         mbs=micro_batch_size,
                         post_processing_fn=loss_fn_wrapped,
                         forward_only=eval_mode,
-                        #do_not_average_loss=True, ## TODO!
                         defer_fp32_logits=self.defer_fp32_logits,
                         global_valid_seqs=global_valid_seqs,
                         global_valid_toks=global_valid_toks,
+                        do_not_average_loss=True,
                     )
 
                 # Empty unused memory.