feat: Add chunked linear ce loss function from hidden states

nemo_rl/algorithms/loss_functions.py

@@ -1332,3 +1332,144 @@ def __call__(
         }
 
         return kl_loss, metrics
+
+
+class NLLLinearCEFusionLoss(LossFunction):
+    """Negative Log Likelihood Loss function with linear CE fusion."""
+
+    loss_type = LossType.TOKEN_LEVEL
+
+    def __call__(
+        self,
+        token_logprobs: Tensor,
+        data: BatchedDataDict[Any],
+        global_valid_seqs: Tensor | None,
+        global_valid_toks: Tensor,
+        vocab_parallel_rank: Optional[int] = None,
+        vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+        context_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+        dpo_loss: bool = False,
+        dpo_average_log_probs: bool = False,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        # Get the next token hidden states for each position
+        token_mask = data["token_mask"][:, 1:]
+        sample_mask = data["sample_mask"]
+        mask = token_mask * sample_mask.unsqueeze(-1)
+        seq_index = data.get("seq_index", None)
+        token_logprobs = token_logprobs.to(torch.float32)
+        # Gather the logprobs for the actual next tokens
+        # only support megatron lm tensor parallel for now
+        if vocab_parallel_group is not None:
+            assert vocab_parallel_rank is not None, (
+                "vocab_parallel_rank must be provided when vocab_parallel_group is provided"
+            )
+            token_logprobs = token_logprobs[:, : data["input_ids"].shape[1] - 1]
+        elif isinstance(token_logprobs, torch.distributed.tensor.DTensor):
+            raise NotImplementedError(
+                "Distributed tensor not supported for linear CE fusion loss"
+            )
+        else:
+            raise NotImplementedError(
+                "Non-distributed tensor not supported for linear CE fusion loss"
+            )
+
+        if dpo_loss:
+            ## shape: [batch_size]
+            num_unmasked_tokens = torch.sum(mask, -1)
+            ## multiply by sample_mask to zero out invalid samples
+            loss = -torch.sum(token_logprobs * mask, dim=-1)
+            if dpo_average_log_probs:
+                loss = loss / num_unmasked_tokens.clamp(min=1)
+        else:
+            ## single scalar loss
+            ## scale by the total number of tokens in the batch
+            loss = -masked_mean(
+                token_logprobs,
+                mask,
+                global_normalization_factor=global_valid_toks,
+            )
+
+        return loss, {
+            "loss": loss.item() if loss.ndim == 0 else loss,
+            "num_unmasked_tokens": mask.sum().item(),
+            "num_valid_samples": sample_mask.sum().item(),
+        }
+
+
+class SequencePackingNLLLinearCEFusionLossWrapper:
+    def __init__(
+        self,
+        loss_fn: NLLLinearCEFusionLoss,
+        cu_seqlens_q: Tensor,
+        cu_seqlens_q_padded: Optional[Tensor] = None,
+    ):
+        self.loss_fn = loss_fn
+        self.cu_seqlens_q = cu_seqlens_q
+        self.cu_seqlens_q_padded = cu_seqlens_q_padded
+
+    def __call__(
+        self,
+        token_logprobs: 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
+        unpadded_seq_lengths = self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1]
+        if self.cu_seqlens_q_padded is not None:
+            padded_cu_seqlens = self.cu_seqlens_q_padded
+            padded_seq_lengths = (
+                self.cu_seqlens_q_padded[1:] - self.cu_seqlens_q_padded[:-1]
+            )
+        else:
+            padded_cu_seqlens = unpadded_cu_seqlens
+            padded_seq_lengths = unpadded_seq_lengths
+        seq_starts = padded_cu_seqlens[:-1]
+        seq_ends = padded_cu_seqlens[1:]
+
+        loss_accum = 0
+        metrics_accum = {}
+        for seq_idx in range(len(seq_starts)):
+            seq_start = seq_starts[seq_idx].item()
+            seq_end = seq_ends[seq_idx].item()
+
+            # get sequence and unpad all 'data' tensors. The data dict is a BatchedDataDict of unpacked tensors
+            seq_data = data.slice(seq_idx, seq_idx + 1)
+            unpadded_seq_data = {}
+            for k, v in seq_data.items():
+                if isinstance(v, torch.Tensor) and v.ndim > 1 and v.shape[1] > 1:
+                    unpadded_seq_data[k] = v[:, : unpadded_seq_lengths[seq_idx]]
+                else:
+                    unpadded_seq_data[k] = v
+
+            # get next_token_logits
+            cp_size = (
+                1
+                if context_parallel_group is None
+                else torch.distributed.get_world_size(context_parallel_group)
+            )
+            logit_slice_idxs = slice(
+                seq_start // cp_size,
+                (seq_start + padded_seq_lengths[seq_idx]) // cp_size,
+            )
+            token_logprobs_slice = token_logprobs[:, logit_slice_idxs]
+            loss, metrics = self.loss_fn(
+                token_logprobs_slice,
+                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():
+                if k not in metrics_accum:
+                    metrics_accum[k] = 0
+                metrics_accum[k] += v
+
+        return loss_accum, metrics_accum

nemo_rl/algorithms/sft.py

@@ -22,7 +22,7 @@
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
 from nemo_rl.algorithms.loss_functions import (
-    NLLLoss,
+    NLLLoss, NLLLinearCEFusionLoss
 )
 from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed
 from nemo_rl.data import DataConfig
@@ -210,7 +210,7 @@ def setup(
     # print the node IP and GPU ID of the policy workers for debugging
     policy.print_node_ip_and_gpu_id()
 
-    loss_fn = NLLLoss()
+    loss_fn = NLLLinearCEFusionLoss() if policy_config["megatron_cfg"].get("enabled", False) and policy_config["megatron_cfg"].get("use_linear_ce_fusion_loss", False) else NLLLoss()
     print("  ✓ Model initialized")
 
     print("\n" + "=" * 60)