fix seq packing

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

nemo_rl/algorithms/loss_functions.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from typing import Any, NotRequired, Optional, TypedDict, TypeVar
+from typing import Any, Callable, NotRequired, Optional, TypedDict, TypeVar
 
 import torch
 import torch.distributed
@@ -802,22 +802,27 @@ class SequencePackingLossWrapper:
     def __init__(
         self,
         loss_fn: LossFunction,
+        prepare_fn: Callable[Any, Any],
         cu_seqlens_q: Tensor,
         cu_seqlens_q_padded: Optional[Tensor] = None,
+        vocab_parallel_rank: Optional[int] = None,
+        vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+        context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     ):
         self.loss_fn = loss_fn
+        self.prepare_fn = prepare_fn
         self.cu_seqlens_q = cu_seqlens_q
         self.cu_seqlens_q_padded = cu_seqlens_q_padded
+        self.vocab_parallel_rank = vocab_parallel_rank
+        self.vocab_parallel_group = vocab_parallel_group
+        self.context_parallel_group = context_parallel_group
 
     def __call__(
         self,
         next_token_logits: Tensor,
         data: BatchedDataDict[Any],
         global_valid_seqs: Tensor | None,
         global_valid_toks: Tensor | None,
-        vocab_parallel_rank: Optional[int] = None,
-        vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
-        context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
     ) -> tuple[Tensor, dict[str, Any]]:
         """Wraps a loss function to handle sequence packing by doing one sequence at a time to avoid excessive padding."""
         unpadded_cu_seqlens = self.cu_seqlens_q
@@ -851,8 +856,8 @@ def __call__(
             # get next_token_logits
             cp_size = (
                 1
-                if context_parallel_group is None
-                else torch.distributed.get_world_size(context_parallel_group)
+                if self.context_parallel_group is None
+                else torch.distributed.get_world_size(self.context_parallel_group)
             )
             logit_start = seq_start // cp_size
             logit_end = (seq_start + padded_seq_lengths[seq_idx]) // cp_size
@@ -861,14 +866,14 @@ def __call__(
                 1, logit_start, logit_length
             )
 
+            # prepare data for loss function
+            loss_fn_args = self.prepare_fn(next_token_logits_slice, unpadded_seq_data)
+
             loss, metrics = self.loss_fn(
-                next_token_logits_slice,
+                *loss_fn_args,
                 unpadded_seq_data,
                 global_valid_seqs,
                 global_valid_toks,
-                vocab_parallel_rank=vocab_parallel_rank,
-                vocab_parallel_group=vocab_parallel_group,
-                context_parallel_group=context_parallel_group,
             )
             loss_accum += loss
             for k, v in metrics.items():

nemo_rl/models/automodel/train.py

@@ -505,6 +505,12 @@ def __call__(
         Returns:
             Tuple of (loss, metrics)
         """
+        from nemo_rl.algorithms.loss_functions import (
+            ClippedPGLossFn,
+            DPOLossFn,
+            NLLLoss,
+        )
+
         # Handle CP redistribution
         if self.cp_size > 1:
             _, mb = prepare_data_for_cp(
@@ -514,30 +520,45 @@ def __call__(
                 logits, self.device_mesh, self.cp_mesh, sequence_dim
             )
 
-        # Compute logprobs from logits
-        logprobs = get_logprobs_from_logits(
-            input_ids=mb["input_ids"],
-            next_token_logits=logits,
-            seq_index=mb.get("seq_index", None),
-        )
-        del logits
+        # Prepare data for loss function
+        def prepare_for_loss_fn(
+            logits: torch.Tensor, mb: BatchedDataDict[Any]
+        ) -> tuple[Any]:
+            if isinstance(self.loss_fn, (ClippedPGLossFn, NLLLoss, DPOLossFn)):
+                logprobs = get_logprobs_from_logits(
+                    input_ids=mb["input_ids"],
+                    next_token_logits=logits,
+                    seq_index=mb.get("seq_index", None),
+                )
+
+                loss_fn_args = (logprobs,)
+
+            # TODO: PreferenceLoss, DistillationLossFn
+
+            return loss_fn_args
 
         # Wrap loss function for sequence packing if needed
         if self.enable_seq_packing:
             loss_fn_ = SequencePackingLossWrapper(
                 loss_fn=self.loss_fn,
+                prepare_fn=prepare_for_loss_fn,
                 cu_seqlens_q=processed_inputs.flash_attn_kwargs.cu_seqlens_q,
                 cu_seqlens_q_padded=processed_inputs.flash_attn_kwargs.cu_seqlens_q,
             )
+            loss, loss_metrics = loss_fn_(
+                logits,
+                mb,
+                global_valid_seqs,
+                global_valid_toks,
+            )
         else:
-            loss_fn_ = self.loss_fn
-
-        loss, loss_metrics = loss_fn_(
-            logprobs,
-            mb,
-            global_valid_seqs,
-            global_valid_toks,
-        )
+            loss_fn_args = prepare_for_loss_fn(logits, mb)
+            loss, loss_metrics = self.loss_fn(
+                *loss_fn_args,
+                mb,
+                global_valid_seqs,
+                global_valid_toks,
+            )
 
         return loss, loss_metrics