fix(grpo): drop top-k/top-p -inf from token_mask

nemo_rl/algorithms/grpo.py

@@ -3602,9 +3602,12 @@ def _grpo_train_impl(
                     )
 
                     if not skip_prev_logprobs:
-                        train_data["prev_logprobs"] = policy.get_logprobs(
-                            logprob_data, timer=timer
-                        )["logprobs"]
+                        prev_lp_result = policy.get_logprobs(logprob_data, timer=timer)
+                        train_data["prev_logprobs"] = prev_lp_result["logprobs"]
+                        # When top-k/top-p filtering is enabled, the worker returns a mask that has
+                        # -inf positions zeroed. Propagate it so the loss reduction skips them.
+                        if "token_mask" in prev_lp_result:
+                            train_data["token_mask"] = prev_lp_result["token_mask"]
                     else:
                         print(
                             "▶ Skipping prev_logprobs (force_on_policy_ratio=True)...",
@@ -5405,9 +5408,11 @@ def _flush_collector_telemetry() -> None:
                     ),
                 ):
                     if not skip_prev_logprobs:
-                        train_data["prev_logprobs"] = policy.get_logprobs(
-                            train_data, timer=timer
-                        )["logprobs"]
+                        prev_lp_result = policy.get_logprobs(train_data, timer=timer)
+                        train_data["prev_logprobs"] = prev_lp_result["logprobs"]
+                        # Propagate the top-k/top-p neginf token mask.
+                        if "token_mask" in prev_lp_result:
+                            train_data["token_mask"] = prev_lp_result["token_mask"]
                     else:
                         train_data["prev_logprobs"] = torch.zeros_like(
                             train_data["generation_logprobs"]

nemo_rl/algorithms/loss/loss_input.py

@@ -272,7 +272,13 @@ def prepare_loss_input(
             # 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 = mask_out_neg_inf_logprobs(logprobs, mask[:, 1:], "curr_logprobs")
+            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 reference policy KL penalty
             if (
@@ -544,7 +550,11 @@ def prepare_packed_loss_input(
     # use filtered curr_logprobs for actor loss, but keep unfiltered values for KL.
     if need_top_k_or_top_p_filtering(sampling_params):
         mask = data["token_mask"] * data["sample_mask"].unsqueeze(-1)
-        logprobs = mask_out_neg_inf_logprobs(logprobs, mask[:, 1:], "curr_logprobs")
+        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
 
         if (
             hasattr(loss_fn, "reference_policy_kl_penalty")

nemo_rl/algorithms/ppo.py

@@ -1659,12 +1659,16 @@ def ppo_train(
                         {
                             "input_ids": train_data["input_ids"],
                             "input_lengths": train_data["input_lengths"],
+                            "token_mask": train_data["token_mask"],
+                            "sample_mask": train_data["sample_mask"],
                             **extra_multimodal_data,
                         }
                     )
-                    train_data["prev_logprobs"] = policy.get_logprobs(
-                        logprob_data, timer=timer
-                    )["logprobs"]
+                    prev_lp_result = policy.get_logprobs(logprob_data, timer=timer)
+                    train_data["prev_logprobs"] = prev_lp_result["logprobs"]
+                    # Propagate the top-k/top-p neginf mask so the loss skips these positions.
+                    if "token_mask" in prev_lp_result:
+                        train_data["token_mask"] = prev_lp_result["token_mask"]
 
                     if not master_config.ppo.skip_reference_policy_logprobs_calculation:
                         train_data["reference_policy_logprobs"] = (
@@ -2686,12 +2690,16 @@ def _raise_if_collector_stopped(waiting_for: str) -> None:
                         {
                             "input_ids": train_data["input_ids"],
                             "input_lengths": train_data["input_lengths"],
+                            "token_mask": train_data["token_mask"],
+                            "sample_mask": train_data["sample_mask"],
                             **extra_multimodal_data,
                         }
                     )
-                    train_data["prev_logprobs"] = policy.get_logprobs(
-                        logprob_data, timer=timer
-                    )["logprobs"]
+                    prev_lp_result = policy.get_logprobs(logprob_data, timer=timer)
+                    train_data["prev_logprobs"] = prev_lp_result["logprobs"]
+                    # Propagate the top-k/top-p neginf mask so the loss skips these positions.
+                    if "token_mask" in prev_lp_result:
+                        train_data["token_mask"] = prev_lp_result["token_mask"]
                     if not master_config.ppo.skip_reference_policy_logprobs_calculation:
                         train_data["reference_policy_logprobs"] = (
                             policy.get_reference_policy_logprobs(

nemo_rl/algorithms/utils.py

@@ -300,7 +300,7 @@ def masked_mean(
 
 def mask_out_neg_inf_logprobs(
     logprobs: torch.Tensor, mask: torch.Tensor, logprobs_name: str
-) -> torch.Tensor:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Mask out negative infinity log probabilities.
 
     Handling sampling mask mismatch:
@@ -315,7 +315,7 @@ def mask_out_neg_inf_logprobs(
         logprobs_name: Name of the logprobs tensor. Used for printing warning messages.
 
     Returns:
-        Masked log probabilities.
+        Tuple of (masked log probabilities, finite-position indicator).
     """
     is_neginf = torch.isinf(logprobs)
     neginf_count = (is_neginf & mask.bool()).sum().item()
@@ -325,10 +325,11 @@ def mask_out_neg_inf_logprobs(
             "(policy top-k/top-p mismatch). Masking out these positions."
         )
 
-    mask = mask * (~is_neginf).float()
-    logprobs = torch.where(mask.bool(), logprobs, 0.0)
+    finite_mask = (~is_neginf).float()
+    effective_mask = mask * finite_mask
+    logprobs = torch.where(effective_mask.bool(), logprobs, 0.0)
 
-    return logprobs
+    return logprobs, finite_mask
 
 
 def masked_var(

nemo_rl/data_plane/worker_mixin.py

@@ -922,7 +922,8 @@ def get_logprobs_presharded(
         """Per-rank logprob entrypoint. Fetch → packing prep → run → write back.
 
         Returns ``None`` — the per-token tensor is committed to TQ via
-        :meth:`_write_back_result_field` under ``prev_logprobs``.
+        :meth:`_write_back_result_field` under ``prev_logprobs``;
+        when the worker narrows ``token_mask``, that column is rewritten too.
         Callers fetch it through :meth:`TQPolicy.read_from_dataplane` —
         skipping the Ray plasma roundtrip on the (B, S) tensor.
         ``del result`` drops the local reference before returning so the
@@ -940,6 +941,14 @@ def get_logprobs_presharded(
             result_key="logprobs",
             tq_field="prev_logprobs",
         )
+        # Top-k/top-p filtering narrows token_mask at -inf positions.
+        if "token_mask" in result:
+            self._write_back_result_field(
+                meta,
+                result,
+                result_key="token_mask",
+                tq_field="token_mask",
+            )
         del result
 
     @accepts_trace_context

nemo_rl/models/automodel/train.py

@@ -386,7 +386,11 @@ def forward_with_post_processing_fn(
             sequence_dim=sequence_dim,
         )
         if isinstance(post_processing_fn, LogprobsPostProcessor):
-            metrics = {"logprobs": result}
+            logprobs_result, updated_token_mask = result
+            result = logprobs_result
+            metrics = {"logprobs": logprobs_result}
+            if updated_token_mask is not None:
+                metrics["token_mask"] = updated_token_mask
         else:
             vals, idx = result
             metrics = {"topk_logits": vals, "topk_indices": idx}
@@ -705,7 +709,7 @@ def __call__(
         *,
         cp_sharder: Optional[ContextParallelSharder],
         sequence_dim: int = 1,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """Compute token log probabilities from logits.
 
         Args:
@@ -719,7 +723,9 @@ def __call__(
             sequence_dim: Sequence dimension
 
         Returns:
-            Token log probabilities tensor [batch_size, seq_length]
+            (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).
         """
         input_lengths = data_dict["input_lengths"]
 
@@ -788,13 +794,15 @@ def __call__(
             token_logprobs = token_logprobs * post_attention_mask
 
         # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
+        updated_token_mask: Optional[torch.Tensor] = 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 = mask_out_neg_inf_logprobs(
+            token_logprobs, finite_mask = mask_out_neg_inf_logprobs(
                 token_logprobs, mask, "prev_logprobs"
             )
+            updated_token_mask = data_dict["token_mask"] * finite_mask
 
-        return token_logprobs
+        return token_logprobs, updated_token_mask
 
     def _compute_local_logprobs(
         self,

nemo_rl/models/megatron/train.py

@@ -799,7 +799,8 @@ 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})
+            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
         """
         unpacked_input_ids = data_dict["input_ids"]
 
@@ -845,17 +846,20 @@ def processor_fn_inner(output_tensor):
             )
 
             # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
+            result_dict: dict[str, torch.Tensor] = {}
             if need_top_k_or_top_p_filtering(self.sampling_params):
                 mask = data_dict["token_mask"] * data_dict["sample_mask"].unsqueeze(-1)
-                token_logprobs = mask_out_neg_inf_logprobs(
+                token_logprobs, finite_mask = mask_out_neg_inf_logprobs(
                     token_logprobs, mask, "prev_logprobs"
                 )
+                result_dict["token_mask"] = (data_dict["token_mask"] * finite_mask)[
+                    :, :original_seq_length
+                ]
 
             token_logprobs = token_logprobs[:, :original_seq_length]
+            result_dict["logprobs"] = token_logprobs
 
-            return torch.tensor(0.0, device=token_logprobs.device), {
-                "logprobs": token_logprobs
-            }
+            return torch.tensor(0.0, device=token_logprobs.device), result_dict
 
         return processor_fn_inner
 

nemo_rl/models/policy/interfaces.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
-from typing import Any, Optional, TypedDict
+from typing import Any, NotRequired, Optional, TypedDict
 
 import ray
 import torch
@@ -24,9 +24,19 @@
 
 
 class LogprobOutputSpec(TypedDict):
-    """logprobs: Tensor of log probabilities."""
+    """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).
+    """
 
     logprobs: torch.Tensor
+    token_mask: NotRequired[torch.Tensor]
 
 
 class ReferenceLogprobOutputSpec(TypedDict):

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -28,7 +28,10 @@
 from torch import nn
 from torch.distributed.tensor import DTensor
 
-from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
+from nemo_rl.algorithms.logits_sampling_utils import (
+    TrainingSamplingParams,
+    need_top_k_or_top_p_filtering,
+)
 from nemo_rl.algorithms.loss.interfaces import LossFunction
 from nemo_rl.algorithms.metric_utils import LEARNING_RATE_KEY
 from nemo_rl.data_plane.worker_mixin import TQWorkerMixin
@@ -86,6 +89,7 @@
 from nemo_rl.utils.grad_norm import warn_if_inf_grad_norm
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
 from nemo_rl.utils.packed_tensor import packed_broadcast_producer
+from nemo_rl.utils.tensor_ops import pad_and_concat
 from nemo_rl.utils.timer import Timer
 
 
@@ -609,6 +613,7 @@ def get_logprobs(
         sequence_dim, seq_dim_size = check_sequence_dim(data)
 
         all_log_probs = []
+        all_token_masks: list[torch.Tensor] = []
         self.model.eval()
 
         # Create logprobs post-processor
@@ -657,19 +662,21 @@ def get_logprobs(
                     continue
 
                 all_log_probs.append(token_logprobs)
+                if "token_mask" in _metrics:
+                    all_token_masks.append(_metrics["token_mask"])
 
         # Concatenate all batches
         return_data = BatchedDataDict[LogprobOutputSpec]()
 
-        all_log_probs_padded = []
-        for lp in all_log_probs:
-            padding_needed = seq_dim_size - lp.shape[1]
-            if padding_needed > 0:
-                lp = torch.nn.functional.pad(
-                    lp, (0, padding_needed), mode="constant", value=0.0
-                )
-            all_log_probs_padded.append(lp)
-        return_data["logprobs"] = torch.cat(all_log_probs_padded, dim=0).cpu()
+        return_data["logprobs"] = pad_and_concat(
+            all_log_probs, target_len=seq_dim_size
+        ).cpu()
+        # Taken from config so every DP rank emits the same keys.
+        if need_top_k_or_top_p_filtering(self.sampling_params):
+            # Pad token_mask with 0 so padded positions are excluded from the loss.
+            return_data["token_mask"] = pad_and_concat(
+                all_token_masks, target_len=seq_dim_size
+            ).cpu()
 
         self.timer.stop("get_logprobs")
         return return_data
@@ -820,31 +827,11 @@ def get_topk_logits(
 
         ret = BatchedDataDict[Any]()
         # Pad each micro-batch result on sequence dim to common length (S), similar to get_logprobs
-        all_topk_vals_padded = []
-        all_topk_idx_padded = []
-        target_seq_len = seq_dim_size
-        for vals, idx in zip(out_topk_vals, out_topk_idx):
-            pad_needed = target_seq_len - vals.shape[1]
-            if pad_needed > 0:
-                # pad along sequence dimension (second dim): (last_dim_pad_left, last_dim_pad_right, seq_pad_left, seq_pad_right, batch_pad_left, batch_pad_right)
-                vals = torch.nn.functional.pad(
-                    vals, (0, 0, 0, pad_needed, 0, 0), mode="constant", value=0.0
-                )
-                idx = torch.nn.functional.pad(
-                    idx, (0, 0, 0, pad_needed, 0, 0), mode="constant", value=0
-                )
-            all_topk_vals_padded.append(vals)
-            all_topk_idx_padded.append(idx)
-
-        ret["topk_logits"] = (
-            torch.cat(all_topk_vals_padded, dim=0)
-            if len(all_topk_vals_padded) > 1
-            else all_topk_vals_padded[0]
+        ret["topk_logits"] = pad_and_concat(
+            out_topk_vals, target_len=seq_dim_size
         ).cpu()
-        ret["topk_indices"] = (
-            torch.cat(all_topk_idx_padded, dim=0)
-            if len(all_topk_idx_padded) > 1
-            else all_topk_idx_padded[0]
+        ret["topk_indices"] = pad_and_concat(
+            out_topk_idx, target_len=seq_dim_size
         ).cpu()
         return ret