fix(loss): exclude -inf logprobs from the actor reduction, not the KL

nemo_rl/algorithms/loss/loss_functions.py

@@ -174,6 +174,9 @@ class ClippedPGLossDataDict(TypedDict):
     reference_policy_logprobs: torch.Tensor
     token_mask: torch.Tensor
     sample_mask: torch.Tensor
+    # Set by ``prepare_loss_input`` when top-k/top-p filtering leaves -inf
+    # logprobs; shape [B, S-1], already in the next-token frame.
+    curr_logprobs_keep_mask: NotRequired[torch.Tensor]
     __extra__: Any
 
 
@@ -396,6 +399,17 @@ def __call__(
 
         mask = token_mask * sample_mask.unsqueeze(-1)
 
+        # Positions whose ``curr_logprobs`` were filtered out carry a
+        # substituted value, so every quantity derived from ``curr_logprobs``
+        # reduces over ``actor_mask`` instead. Everything else keeps the full
+        # mask: the reference-policy KL reads ``curr_logprobs_unfiltered`` and
+        # the mismatch diagnostics read only prev/generation logprobs, so
+        # narrowing them would drop real terms -- and drop them selectively,
+        # since a token is filtered precisely where the policies disagree.
+        keep_mask = data.get("curr_logprobs_keep_mask")
+        actor_token_mask = token_mask if keep_mask is None else token_mask * keep_mask
+        actor_mask = actor_token_mask * sample_mask.unsqueeze(-1)
+
         # For truly on-policy training, use curr_logprobs as prev_logprobs
         # This avoids computing prev_logprobs upstream
         if self.force_on_policy_ratio:
@@ -528,7 +542,7 @@ def __call__(
             if self.sequence_level_importance_ratios:
                 seq_log_ratio_mean = masked_mean(
                     log_ratios,
-                    token_mask,
+                    actor_token_mask,
                     dim=-1,
                 ).unsqueeze(-1)
                 seq_ratio = seq_log_ratio_mean.exp()
@@ -678,14 +692,14 @@ def __call__(
         if self.loss_type == LossType.TOKEN_LEVEL:
             actor_loss = masked_mean(
                 importance_weights_to_use * clip_loss,
-                mask,
+                actor_mask,
                 global_normalization_factor=global_valid_toks,
             )
         else:
             actor_loss = masked_mean(
                 masked_mean(
                     importance_weights_to_use * clip_loss,
-                    token_mask,
+                    actor_token_mask,
                     dim=-1,
                 ),
                 sample_mask,
@@ -712,7 +726,7 @@ def __call__(
         with torch.no_grad():
             seq_entropy_approx = -masked_mean(
                 torch.exp(curr_logprobs - generation_logprobs) * curr_logprobs,
-                mask,
+                actor_mask,
                 global_normalization_factor=global_valid_toks,
             )
 
@@ -723,7 +737,7 @@ def __call__(
         nll_loss = torch.tensor(0.0, device=mask.device)
         if self.positive_example_nll_weight > 0 and "rewards" in data:
             correct_sample_mask = (data["rewards"] > 0).float()  # [batch]
-            correct_mask = mask * correct_sample_mask.unsqueeze(-1)
+            correct_mask = actor_mask * correct_sample_mask.unsqueeze(-1)
             correct_valid_toks = correct_mask.sum()
             if correct_valid_toks > 0:
                 nll_loss = masked_mean(
@@ -736,18 +750,18 @@ def __call__(
         with torch.no_grad():
             probs_ratio = masked_mean(
                 ratios.detach(),
-                mask,
+                actor_mask,
                 global_normalization_factor=global_valid_toks,
             ).item()
             probs_ratio_clamped = masked_mean(
                 ratios_clamped.detach(),
-                mask,
+                actor_mask,
                 global_normalization_factor=global_valid_toks,
             ).item()
 
             # Calculate min/max values for ratios (only for valid tokens)
-            masked_ratios = ratios.detach()[mask.bool()]
-            masked_ratios_clamped = ratios_clamped.detach()[mask.bool()]
+            masked_ratios = ratios.detach()[actor_mask.bool()]
+            masked_ratios_clamped = ratios_clamped.detach()[actor_mask.bool()]
 
             # Handle edge case where there might be no valid tokens
             if masked_ratios.numel() > 0:

nemo_rl/algorithms/loss/utils.py

@@ -98,7 +98,14 @@ 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, keep = mask_out_neg_inf_logprobs(
+                logprobs, mask[:, 1:], "curr_logprobs"
+            )
+            # Publish the narrowing instead of folding it into ``token_mask``:
+            # the actor term must drop these positions, but the
+            # reference-policy KL reduces with ``token_mask`` and is computed
+            # from ``curr_logprobs_unfiltered``, which is finite here.
+            data["curr_logprobs_keep_mask"] = keep
 
             # compute unfiltered logprobs for reference policy KL penalty
             if (
@@ -359,7 +366,14 @@ 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, keep = mask_out_neg_inf_logprobs(
+            logprobs, mask[:, 1:], "curr_logprobs"
+        )
+        # Publish the narrowing instead of folding it into ``token_mask``:
+        # the actor term must drop these positions, but the
+        # reference-policy KL reduces with ``token_mask`` and is computed
+        # from ``curr_logprobs_unfiltered``, which is finite here.
+        data["curr_logprobs_keep_mask"] = keep
 
         if (
             hasattr(loss_fn, "reference_policy_kl_penalty")

nemo_rl/algorithms/utils.py

@@ -226,7 +226,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:
@@ -235,13 +235,22 @@ def mask_out_neg_inf_logprobs(
     token X may fall outside the training policy's top-k/p set -> curr_logprobs[X] = -inf, prev_logprobs[X] = -inf
     Detect positions with -inf in any logprobs (generation_logprobs is always finite for valid tokens)
 
+    The substituted ``0.0`` is *not* a neutral value: it means log p = 0, i.e.
+    p = 1, while ``generation_logprobs`` at the same position is finite and
+    typically around -5. Anything that compares the two -- the importance ratio,
+    ``token_mult_prob_error``, the sequence-level weights -- reads a large
+    fabricated difference unless the position is also dropped from the reduction
+    mask. Callers must therefore fold ``keep`` into their own mask; returning it
+    is what makes the warning above true.
+
     Args:
         logprobs: Log probabilities.
         mask: Mask.
         logprobs_name: Name of the logprobs tensor. Used for printing warning messages.
 
     Returns:
-        Masked log probabilities.
+        ``(logprobs, keep)``, where ``keep`` is 0 at positions whose logprob was
+        -inf and 1 elsewhere, with the same shape as ``mask``.
     """
     is_neginf = torch.isinf(logprobs)
     neginf_count = (is_neginf & mask.bool()).sum().item()
@@ -251,10 +260,11 @@ def mask_out_neg_inf_logprobs(
             "(policy top-k/top-p mismatch). Masking out these positions."
         )
 
-    mask = mask * (~is_neginf).float()
+    keep = (~is_neginf).to(mask.dtype)
+    mask = mask * keep
     logprobs = torch.where(mask.bool(), logprobs, 0.0)
 
-    return logprobs
+    return logprobs, keep
 
 
 def masked_var(

nemo_rl/models/automodel/train.py

@@ -791,7 +791,12 @@ def __call__(
         # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
         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(
+            # The narrowed mask cannot travel with prev_logprobs today:
+            # LogprobOutputSpec carries logprobs only. The curr side publishes
+            # its narrowing as curr_logprobs_keep_mask, which covers the actor
+            # reduction; a prev-side -inf at a position the curr side kept is
+            # left for a follow-up that widens the spec.
+            token_logprobs, _ = mask_out_neg_inf_logprobs(
                 token_logprobs, mask, "prev_logprobs"
             )
 

nemo_rl/models/megatron/train.py

@@ -673,7 +673,12 @@ def processor_fn_inner(output_tensor):
             # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
             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(
+                # The narrowed mask cannot travel with prev_logprobs today:
+                # LogprobOutputSpec carries logprobs only. The curr side publishes its
+                # narrowing as curr_logprobs_keep_mask, which covers the actor
+                # reduction; a prev-side -inf at a position the curr side kept is left
+                # for a follow-up that widens the spec.
+                token_logprobs, _ = mask_out_neg_inf_logprobs(
                     token_logprobs, mask, "prev_logprobs"
                 )
 

nemo_rl/models/policy/workers/dtensor_policy_worker.py

@@ -1334,7 +1334,12 @@ def get_logprobs(
         # handle top-k/top-p filtering for logprobs, only used for ClippedPGLossFn now
         if need_top_k_or_top_p_filtering(self.sampling_params):
             mask = data["token_mask"] * data["sample_mask"].unsqueeze(-1)
-            token_logprobs = mask_out_neg_inf_logprobs(
+            # The narrowed mask cannot travel with prev_logprobs today:
+            # LogprobOutputSpec carries logprobs only. The curr side publishes its
+            # narrowing as curr_logprobs_keep_mask, which covers the actor reduction; a
+            # prev-side -inf at a position the curr side kept is left for a follow-up
+            # that widens the spec.
+            token_logprobs, _ = mask_out_neg_inf_logprobs(
                 token_logprobs, mask, "prev_logprobs"
             )
 

tests/unit/algorithms/test_loss_functions.py

@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
+import math
 
 import pytest
 import torch
 
+import nemo_rl.algorithms.loss.utils as loss_utils
+from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
 from nemo_rl.algorithms.loss import (
     ClippedPGLossConfig,
     ClippedPGLossFn,
@@ -43,6 +46,91 @@
 )
 
 
+def test_prepare_loss_input_excludes_filtered_neg_inf_logprobs(monkeypatch):
+    """The caller must persist the narrowed mask used by the actor loss."""
+    data = BatchedDataDict(
+        {
+            "input_ids": torch.tensor([[0, 1, 2, 3]]),
+            "token_mask": torch.ones(1, 4),
+            "sample_mask": torch.ones(1),
+        }
+    )
+    filtered_logprobs = torch.tensor([[-0.5, float("-inf"), -1.5]])
+    monkeypatch.setattr(
+        loss_utils,
+        "get_next_token_logprobs_from_logits",
+        lambda **_: filtered_logprobs,
+    )
+
+    loss_input, updated_data = prepare_loss_input(
+        torch.empty(1, 4, 4),
+        data,
+        ClippedPGLossFn(ClippedPGLossConfig(reference_policy_kl_penalty=0.0)),
+        sampling_params=TrainingSamplingParams(top_k=1),
+    )
+
+    assert loss_input["next_token_logprobs"].tolist() == [[-0.5, 0.0, -1.5]]
+    assert updated_data["curr_logprobs_keep_mask"].tolist() == [[1.0, 0.0, 1.0]]
+    # token_mask must not be narrowed: the reference-policy KL reduces with it
+    # and is computed from unfiltered logprobs, which are finite here.
+    assert updated_data["token_mask"].tolist() == [[1.0, 1.0, 1.0, 1.0]]
+
+
+def test_filtered_positions_leave_the_actor_term_but_not_the_kl():
+    """Dropping a filtered position must not weaken the reference-policy KL.
+
+    The substituted 0.0 only corrupts the actor term. ``curr_logprobs_unfiltered``
+    is finite at the same position, so the KL there is real -- and it is exactly
+    where the policies disagree most, since the token was filtered for having
+    near-zero probability under the training policy. Narrowing the KL reduction
+    would systematically remove its largest contributions.
+    """
+    batch, seq = 1, 4
+
+    def make_data(keep_mask):
+        data = BatchedDataDict(
+            {
+                "token_mask": torch.ones(batch, seq),
+                "sample_mask": torch.ones(batch),
+                "advantages": torch.full((batch, seq), 0.5),
+                "prev_logprobs": torch.full((batch, seq), -1.0),
+                "generation_logprobs": torch.full((batch, seq), -1.0),
+                "reference_policy_logprobs": torch.full((batch, seq), -2.0),
+                # large divergence at the position the training policy filtered
+                "curr_logprobs_unfiltered": torch.tensor([[-0.5, -0.6, -3.0]]),
+            }
+        )
+        if keep_mask is not None:
+            data["curr_logprobs_keep_mask"] = keep_mask
+        return data
+
+    curr_logprobs = torch.tensor([[-0.5, -0.6, 0.0]])  # third is the substitute
+    loss_fn = ClippedPGLossFn(ClippedPGLossConfig(reference_policy_kl_penalty=0.5))
+    global_valid_toks = torch.tensor(3.0)
+    global_valid_seqs = torch.tensor(1.0)
+
+    keep = torch.tensor([[1.0, 1.0, 0.0]])
+    loss_all, metrics_all = loss_fn(
+        curr_logprobs, make_data(None), global_valid_seqs, global_valid_toks
+    )
+    loss_kept, metrics_kept = loss_fn(
+        curr_logprobs, make_data(keep), global_valid_seqs, global_valid_toks
+    )
+
+    assert metrics_kept["kl_penalty"] == pytest.approx(metrics_all["kl_penalty"]), (
+        "the keep mask must not reach the KL reduction"
+    )
+    assert metrics_kept["token_mult_prob_error"] == pytest.approx(
+        metrics_all["token_mult_prob_error"]
+    ), "mismatch diagnostics read prev/generation logprobs and stay un-narrowed"
+
+    # The actor side must move: the substituted logprob fabricates a ratio of
+    # exp(0.0 - (-1.0)) = e at the filtered position.
+    assert metrics_all["probs_ratio_max"] == pytest.approx(math.e)
+    assert metrics_kept["probs_ratio_max"] < metrics_all["probs_ratio_max"]
+    assert loss_kept.item() != pytest.approx(loss_all.item())
+
+
 def setup_dpo_loss_test_data(vocab_size=16, batch_size=1):
     seq_len = 4
     data = {

tests/unit/algorithms/test_sequence_packing_fusion.py

@@ -26,6 +26,7 @@
 import pytest
 import torch
 
+import nemo_rl.algorithms.loss.utils as loss_utils
 from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
 from nemo_rl.algorithms.loss import (
     ClippedPGLossConfig,
@@ -38,6 +39,40 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 
 
+def test_prepare_packed_loss_input_excludes_filtered_neg_inf_logprobs(monkeypatch):
+    """The fused caller must publish the narrowed mask used by the actor loss."""
+    data = BatchedDataDict(
+        {
+            "input_ids": torch.tensor([[0, 1, 2, 3]]),
+            "token_mask": torch.ones(1, 4),
+            "sample_mask": torch.ones(1),
+        }
+    )
+    filtered_logprobs = torch.tensor([[-0.5, float("-inf"), -1.5]])
+    monkeypatch.setattr(
+        loss_utils,
+        "from_parallel_logits_to_logprobs_packed_sequences",
+        lambda *_args, **_kwargs: filtered_logprobs,
+    )
+
+    loss_input, updated_data = prepare_packed_loss_input(
+        torch.empty(1, 4, 4),
+        data,
+        ClippedPGLossFn(ClippedPGLossConfig(reference_policy_kl_penalty=0.0)),
+        cu_seqlens_q=torch.tensor([0, 4], dtype=torch.int32),
+        cu_seqlens_q_padded=torch.tensor([0, 4], dtype=torch.int32),
+        vocab_parallel_rank=0,
+        vocab_parallel_group=object(),
+        sampling_params=TrainingSamplingParams(top_k=1),
+    )
+
+    assert loss_input["next_token_logprobs"].tolist() == [[-0.5, 0.0, -1.5]]
+    assert updated_data["curr_logprobs_keep_mask"].tolist() == [[1.0, 0.0, 1.0]]
+    # ``token_mask`` stays intact: it also reduces the reference-policy KL,
+    # which reads the unfiltered logprobs and is finite at this position.
+    assert updated_data["token_mask"].tolist() == [[1.0, 1.0, 1.0, 1.0]]
+
+
 def _setup_2d_process_groups(rank, world_size, cp_size, tp_size):
     """Create 2D (cp, tp) process groups.