fix(loss): exclude invalid policy support from actor ratios

docs/guides/grpo.md

@@ -651,7 +651,7 @@ When top-p or top-k filtering is enabled, the following conventions apply:
 - **`curr_logprobs` and `prev_logprobs`** are computed *with* filtering applied, for compatibility with the actor loss.
 - **`reference_policy_logprobs`** is computed *without* filtering (see the `use_reference_model` in the policy worker).
 - **KL divergence** uses `curr_logprobs_unfiltered`(`curr_logprobs` *without* filtering) so that it is consistent with the reference policy logprobs.
-- **`-inf` positions** caused by the training-side and inference-side disagreeing on the top-k/top-p filtered set are dropped from `token_mask`. They contribute to neither the actor loss nor the KL penalty, as the `token_mask` is applied to `curr_logprobs_unfiltered`.
+- **Invalid policy support** caused by disagreement on the top-k/top-p filtered set is excluded from actor ratios and actor-loss reduction. The shared `token_mask` stays unchanged so the unfiltered reference KL still uses valid response tokens.
 
 Under tensor parallelism (TP), enabling top-p or top-k adds communication overhead. The vocabulary is sharded across GPUs (vocab-parallel), while top-p and top-k require full-vocabulary probabilities. A naive all-gather of logits would require large additional memory. The implementation therefore switches to a batch–sequence-parallel layout via all-to-all communication, applies filtering over the full vocabulary, then switches back, avoiding materialization of the full vocabulary on any single rank.
 

nemo_rl/algorithms/advantage_estimator.py

@@ -313,6 +313,9 @@ def compute_advantage(
                 logprobs_reference,
                 kl_type=self.kl_type,
             )
+            # Tokens outside the filtered policy support are excluded from the
+            # actor loss; keep their -inf logprobs from poisoning the batch.
+            kl = torch.where(torch.isfinite(logprobs_policy), kl, 0.0)
             adv = adv - self.kl_coef * kl
 
         # global normalization across the batch
@@ -448,6 +451,9 @@ def _build_token_level_rewards(
             and logprobs_reference is not None
         ):
             kl = calculate_kl(logprobs_policy, logprobs_reference, self.kl_type)
+            # Tokens outside the filtered policy support are excluded from the
+            # actor loss; keep their -inf logprobs from poisoning the rewards.
+            kl = torch.where(torch.isfinite(logprobs_policy), kl, 0.0)
             token_level_rewards = token_level_rewards - self.kl_coef * kl
 
         # Place terminal reward at the last response token (last mask=1
@@ -720,6 +726,13 @@ def compute_advantage(
             )
             distill_advantages = (proximal_teacher_logprobs - prev_logprobs).detach()
 
+        # Exclude unsupported student tokens from the baseline and metrics as
+        # well as the actor signal, including when using a proximal teacher.
+        student_support = torch.isfinite(prev_logprobs)
+        mask = mask * student_support.to(mask.dtype)
+        teacher_student_gap = torch.where(student_support, teacher_student_gap, 0.0)
+        distill_advantages = torch.where(student_support, distill_advantages, 0.0)
+
         if self.subtract_global_baseline:
             valid_advantages = torch.masked_select(distill_advantages, mask.bool())
             if valid_advantages.numel() > 0:

nemo_rl/algorithms/loss/loss_functions.py

@@ -654,6 +654,30 @@ def __call__(
                     mask.bool(), data["reference_policy_logprobs"][:, 1:], 0.0
                 )
 
+        # Reference KL uses the unfiltered policy and retains its own token set.
+        # Actor support mismatches must not discard otherwise valid KL terms.
+        kl_token_mask = token_mask
+        kl_sample_mask = sample_mask
+        kl_mask = mask
+        kl_generation_logprobs = torch.where(kl_mask.bool(), generation_logprobs, 0.0)
+
+        # Keep support information until all three policy distributions are
+        # available. Replacing -inf with zero earlier turns an impossible token
+        # into log-probability one in token/sequence importance ratios.
+        policy_valid = (
+            torch.isfinite(curr_logprobs)
+            & torch.isfinite(prev_logprobs)
+            & torch.isfinite(generation_logprobs)
+        )
+        token_mask = token_mask * policy_valid.to(token_mask.dtype)
+        sample_mask = sample_mask * token_mask.bool().any(dim=-1).to(sample_mask.dtype)
+        mask = token_mask * sample_mask.unsqueeze(-1)
+        curr_logprobs = torch.where(mask.bool(), curr_logprobs, 0.0)
+        prev_logprobs = torch.where(mask.bool(), prev_logprobs, 0.0)
+        generation_logprobs = torch.where(mask.bool(), generation_logprobs, 0.0)
+        if self.seq_logprob_error_in_loss:
+            seq_error_metrics["seq_logprob_error_valid_tokens"] = mask.sum().item()
+            seq_error_metrics["seq_logprob_error_valid_seqs"] = sample_mask.sum().item()
         # Calculate KL regularization.
         if self.reference_policy_kl_penalty != 0:
             # When top-k/top-p filtering is enabled, we need special handling for KL:
@@ -671,7 +695,7 @@ def __call__(
             if self.use_on_policy_kl_approximation:
                 # See: docs/guides/grpo.md#on-policy-kl-approximation
                 kl_importance_weights = torch.exp(
-                    curr_logprobs_unfiltered - generation_logprobs
+                    curr_logprobs_unfiltered - kl_generation_logprobs
                 )
             else:
                 kl_importance_weights = torch.exp(
@@ -694,12 +718,12 @@ def __call__(
             # Reduce KL loss
             if self.loss_type == LossType.TOKEN_LEVEL:
                 kl = masked_mean(
-                    kl, mask, global_normalization_factor=global_valid_toks
+                    kl, kl_mask, global_normalization_factor=global_valid_toks
                 )
             else:
                 kl = masked_mean(
-                    masked_mean(kl, token_mask, dim=-1),
-                    sample_mask,
+                    masked_mean(kl, kl_token_mask, dim=-1),
+                    kl_sample_mask,
                     global_normalization_factor=global_valid_seqs,
                 )
         else:

nemo_rl/algorithms/loss/loss_input.py

@@ -28,7 +28,6 @@
     map_teacher_logits_to_draft_vocab,
     reconstruct_opd_full_teacher_logits,
 )
-from nemo_rl.algorithms.utils import mask_out_neg_inf_logprobs
 from nemo_rl.algorithms.x_token.loss_utils import (
     prepare_xtoken_cross_tokenizer_loss_input,
 )
@@ -268,19 +267,8 @@ def prepare_loss_input(
 
         # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
         if need_top_k_or_top_p_filtering(sampling_params):
-            # mask out negative infinity logprobs
-            # prev_logprobs is already masked out in the previous step
-            mask = data["token_mask"] * data["sample_mask"].unsqueeze(-1)
-            logprobs, curr_finite_mask = mask_out_neg_inf_logprobs(
-                logprobs, mask[:, 1:], "curr_logprobs"
-            )
-            # Propagate the neg-inf mask so the loss reduction skips these positions.
-            # Without this change, the IS weight exp(prev-gen) becomes exp(-gen).
-            data["token_mask"] = data["token_mask"].clone()
-            data["token_mask"][:, 1:] = data["token_mask"][:, 1:] * curr_finite_mask
-
-            # compute unfiltered logprobs for the reference-KL penalty;
-            # -inf positions were dropped from token_mask above, so the KL skips them.
+            # Preserve support information until the actor loss builds its mask.
+            # compute unfiltered logprobs for reference policy KL penalty
             if (
                 hasattr(loss_fn, "reference_policy_kl_penalty")
                 and loss_fn.reference_policy_kl_penalty != 0
@@ -549,15 +537,8 @@ def prepare_packed_loss_input(
 
     # Match prepare_loss_input behavior for top-k/top-p filtered training:
     # use filtered curr_logprobs for actor loss, but keep unfiltered values for KL.
-    # Note that `-inf` positions are dropped from `token_mask`, so terms like KL skip them.
     if need_top_k_or_top_p_filtering(sampling_params):
-        mask = data["token_mask"] * data["sample_mask"].unsqueeze(-1)
-        logprobs, curr_finite_mask = mask_out_neg_inf_logprobs(
-            logprobs, mask[:, 1:], "curr_logprobs"
-        )
-        data["token_mask"] = data["token_mask"].clone()
-        data["token_mask"][:, 1:] = data["token_mask"][:, 1:] * curr_finite_mask
-
+        # Preserve support information until the actor loss builds its mask.
         if (
             hasattr(loss_fn, "reference_policy_kl_penalty")
             and loss_fn.reference_policy_kl_penalty != 0

nemo_rl/algorithms/utils.py

@@ -298,6 +298,26 @@ def masked_mean(
     return torch.sum(values * mask, dim=dim) / (normalization_factor + 1e-8)
 
 
+def mask_filtered_logprobs_outside_tokens(
+    logprobs: torch.Tensor, mask: torch.Tensor
+) -> torch.Tensor:
+    """Zero top-k/top-p filtered log probabilities outside ``mask``.
+
+    Prompt and padding tokens are often outside the filtered support and would
+    otherwise carry -inf (or NaN after multiplying by a mask) to every consumer.
+    Valid tokens keep -inf so the actor loss can exclude policy support
+    mismatches instead of treating them as log-probability zero.
+
+    Args:
+        logprobs: Filtered log probabilities.
+        mask: Valid-token mask (token_mask * sample_mask).
+
+    Returns:
+        Log probabilities with zeros outside ``mask``.
+    """
+    return torch.where(mask.bool(), logprobs, 0.0)
+
+
 def mask_out_neg_inf_logprobs(
     logprobs: torch.Tensor, mask: torch.Tensor, logprobs_name: str
 ) -> tuple[torch.Tensor, torch.Tensor]:

nemo_rl/data_plane/schema.py

@@ -60,7 +60,7 @@
 # Tensor fields in the train partition. Rollout writes the input
 # subset on first put; later stages add prev_logprobs /
 # reference_policy_logprobs (workers) and advantages (driver);
-# under top-k/top-p the prev-logprob stage also rewrites token_mask.
+# the prev-logprob stage may also return a backend-provided token_mask.
 DP_TRAIN_FIELDS = (
     "input_ids",
     "input_lengths",

nemo_rl/models/automodel/train.py

@@ -49,7 +49,7 @@
 )
 from nemo_rl.algorithms.loss import SequencePackingLossWrapper, prepare_loss_input
 from nemo_rl.algorithms.loss.interfaces import LossFunction, LossInputType
-from nemo_rl.algorithms.utils import mask_out_neg_inf_logprobs
+from nemo_rl.algorithms.utils import mask_filtered_logprobs_outside_tokens
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import (
     distributed_vocab_topk,
@@ -728,8 +728,8 @@ def __call__(
 
         Returns:
             (token log probabilities tensor [batch_size, seq_length],
-             updated token_mask [batch_size, seq_length] with -inf positions zeroed,
-             or None when top-k/top-p filtering is disabled).
+             optional token_mask). Filtering preserves the input mask and leaves
+             policy support information in logprobs for the actor loss.
         """
         input_lengths = data_dict["input_lengths"]
 
@@ -798,20 +798,18 @@ def __call__(
             for i, length in enumerate(input_lengths):
                 # For right-padded sequence, set 1s at the beginning of the sequence
                 post_attention_mask[i, :length] = 1
-            token_logprobs = token_logprobs * post_attention_mask
+            # torch.where: multiplying a filtered -inf by the mask would give NaN.
+            token_logprobs = torch.where(post_attention_mask, token_logprobs, 0.0)
 
-        # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
-        updated_token_mask: Optional[torch.Tensor] = None
+        # Preserve -inf only on valid tokens outside the policy support. The
+        # actor loss excludes these positions before sanitizing logprobs.
+        output_token_mask = None
         if need_top_k_or_top_p_filtering(self.sampling_params):
             mask = data_dict["token_mask"] * data_dict["sample_mask"].unsqueeze(-1)
-            token_logprobs, finite_mask = mask_out_neg_inf_logprobs(
-                token_logprobs, mask, "prev_logprobs"
-            )
-            updated_token_mask = (data_dict["token_mask"] * finite_mask).to(
-                data_dict["token_mask"].dtype
-            )
+            token_logprobs = mask_filtered_logprobs_outside_tokens(token_logprobs, mask)
+            output_token_mask = data_dict["token_mask"]
 
-        return token_logprobs, updated_token_mask
+        return token_logprobs, output_token_mask
 
     def _compute_local_logprobs(
         self,

nemo_rl/models/megatron/train.py

@@ -52,7 +52,7 @@
 from nemo_rl.algorithms.loss.draft import DEFAULT_DRAFT_TOKEN_CHUNK_SIZE
 from nemo_rl.algorithms.loss.interfaces import LossFunction
 from nemo_rl.algorithms.loss.utils import _pack_input_ids
-from nemo_rl.algorithms.utils import mask_out_neg_inf_logprobs
+from nemo_rl.algorithms.utils import mask_filtered_logprobs_outside_tokens
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import (
     allgather_cp_sharded_tensor,
@@ -807,8 +807,7 @@ def __call__(
             original_seq_length: Sequence width before dense padding was applied
 
         Returns:
-            Callable: Function that takes output tensor and returns (dummy_loss, {"logprobs": token_logprobs}),
-                plus a "token_mask" entry narrowed at -inf positions when top-k/top-p filtering is on
+            Callable: Function that takes output tensor and returns (dummy_loss, {"logprobs": token_logprobs})
         """
         unpacked_input_ids = data_dict["input_ids"]
         cu_seqlens_padded_cpu = None
@@ -858,19 +857,20 @@ def processor_fn_inner(output_tensor):
                 [torch.zeros_like(token_logprobs[:, :1]), token_logprobs], dim=1
             )
 
-            # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
-            result_dict: dict[str, torch.Tensor] = {}
+            # Preserve -inf only on valid tokens outside the policy support. The
+            # actor loss excludes these positions before sanitizing logprobs.
             if need_top_k_or_top_p_filtering(self.sampling_params):
                 mask = data_dict["token_mask"] * data_dict["sample_mask"].unsqueeze(-1)
-                token_logprobs, finite_mask = mask_out_neg_inf_logprobs(
-                    token_logprobs, mask, "prev_logprobs"
+                token_logprobs = mask_filtered_logprobs_outside_tokens(
+                    token_logprobs, mask
                 )
-                result_dict["token_mask"] = (data_dict["token_mask"] * finite_mask).to(
-                    data_dict["token_mask"].dtype
-                )[:, :original_seq_length]
 
             token_logprobs = token_logprobs[:, :original_seq_length]
-            result_dict["logprobs"] = token_logprobs
+            result_dict = {"logprobs": token_logprobs}
+            if need_top_k_or_top_p_filtering(self.sampling_params):
+                result_dict["token_mask"] = data_dict["token_mask"][
+                    :, :original_seq_length
+                ]
 
             return torch.tensor(0.0, device=token_logprobs.device), result_dict
 

nemo_rl/models/policy/interfaces.py

@@ -26,13 +26,10 @@
 class LogprobOutputSpec(TypedDict):
     """logprobs: Tensor of log probabilities.
 
-    token_mask: Optional. Present only when top-k/top-p filtering is enabled. The input
-    data["token_mask"] AND-ed with the finite-position mask returned by
-    mask_out_neg_inf_logprobs. Callers MUST write this back where the loss reads
-    token_mask from: train_data["token_mask"] (GRPO/PPO drivers) or the TQ token_mask
-    column (TQWorkerMixin.get_logprobs_presharded), so the loss skips positions where
-    the training-policy logits gave -inf for the vLLM-sampled token (top-k/top-p
-    mismatch).
+    token_mask: Optional backend-provided validity mask. Callers write this back
+    to training data or the transfer queue. Filtered policy support is preserved
+    as -inf in logprobs and excluded by the actor loss without narrowing the
+    shared mask used by unfiltered reference KL.
     """
 
     logprobs: torch.Tensor