What happens to a sampled token that the training policy filtered out — on main, and with this PR

Explainer for NVIDIA-NeMo/RL #4367, which builds on #4154 (already on main). Permalinks pin the PR head ca01e85; "main" links pin the PR's base 8661b47.

Before
On main, #4154 already removes such a token from token_mask, so the actor loss skips it — but so do the reference KL and the sequence-error check.
This PR
Keeps token_mask whole and leaves -inf in the logprobs; the loss builds its own actor mask. Reference KL now keeps the token.
Status
KL part: fixed by this PR. New problem: open — not fixed in this PR. When seq_logprob_error_threshold is set, one such token now throws away the whole response. Raised in the review.
Background. With top-k or top-p sampling, vLLM samples a token from its filtered set. The trainer recomputes the same filtering, but the two can disagree near the cut-off. Then the trainer gives the sampled token log-probability -inf: in its view the token could not be sampled. NeMo-RL compares three logprobs per token: curr, prev and generation. Dotted words have a short explanation: hover or tap them.

Before this PR, the worker replaced that -inf with 0 and removed the token from token_mask (main: loss_input.py, main: grpo.py write-back). token_mask is shared by everything after it, so the token also left the reference KL. This PR keeps -inf on valid tokens (megatron/train.py) and lets the loss drop it from the actor terms only (loss_functions.py). But one more reader of prev_logprobs runs before the loss and was not changed: the sequence-error check.

One response, four tokens, token 3 filtered out

The response: 4 tokens. vLLM sampled token 3, but the trainer's top-k/top-p set does not contain it. generation prev -0.3-0.4 -0.9-0.9 -2.0-inf -0.5-0.5 token 1token 2token 3token 4 token counts here token skipped token -inf, or whole response dropped main (with #4154) this PR Actor loss loss_functions.py:667 token 3 skipped same — token 3 skipped Reference KL loss_functions.py:657 token 3's KL term lost kept (KL uses unfiltered logprobs) — the fix Tokens counted megatron/data.py:1199 global_valid_toks = 3 global_valid_toks = 4 actor loss now × 3/4 (small) Sequence-error check (limit 2.0) utils.py:264 error = 1.035 response kept error = inf whole response dropped from actor loss AND KL error = mean over counted tokens of exp(|generation − prev|). Main: (e^0.1 + 1 + 1) / 3 = 1.035 ≤ 2.0, so the response stays. PR: token 3 still counts, |−2.0 − (−inf)| = inf, so the error is inf and fails any limit. The logged max/mean error also read inf, even with no limit set. The same check runs in the GRPO, async GRPO, sync-TQ, PPO and SingleController drivers, and inside the loss.

🔴 review comment — loss_input.py:270

What would fix it — proposed, not in the PR

Where the fix lands: RL only. the callers of compute_seq_logprob_errors (the drivers and prepare_loss_input). No vLLM change: vLLM's own logprob is right; the trainer's filter is the side that disagrees.

Nothing below is implemented or run. The check must skip a -inf token only when top-k/top-p is on. With filtering off, a -inf is a real bug, and existing tests (test_single_forward_grpo.py) require that such a response is rejected. Skipping it unconditionally breaks one of them (NaN gradients).

compute_seq_logprob_errors itself stays as it is. The caller narrows the mask first, because only the caller knows whether filtering is on. The driver side, as one example:

@@ nemo_rl/algorithms/grpo.py L2822-L2855 · compute_and_apply_seq_logprob_error_masking
     train_data: BatchedDataDict,
     rewards: torch.Tensor,
     seq_logprob_error_threshold: Optional[float],
+    filtering_on: bool = False,
 ) -> dict:
     ...
     prev_logprobs = train_data["prev_logprobs"][:, 1:]
     generation_logprobs = train_data["generation_logprobs"][:, 1:]
+    if filtering_on:
+        # Under top-k/top-p, -inf means "outside the trainer's filtered set":
+        # skip that token (as main did), not the whole response. NaN still fails.
+        token_mask = token_mask * (~torch.isneginf(prev_logprobs)).to(token_mask.dtype)
     seq_mult_prob_error, valid_seq_mask = compute_seq_logprob_errors(

Each driver passes filtering_on=need_top_k_or_top_p_filtering(...) with its sampling settings — the same test the workers use to decide whether to filter (loss_input.py:269). For the in-loss check (loss_functions.py:625), prepare_loss_input builds the same "inside the top-k/top-p set" mask in that branch and the loss ANDs it into token_mask before the check. With either, the example above gives error 1.035 again and the response stays. The loss's own actor mask should use the same gated mask, with -inf only, so a NaN still shows up.

🔴 review comment — loss_input.py:270  🟠 review comment — loss_functions.py:667

So what. The PR's goal — keep the KL on token 3 — is right, and the actor loss treats token 3 the same as main. But for anyone who sets seq_logprob_error_threshold together with top-k/top-p, one disagreeing token now throws away the whole response, and long responses are the most likely to contain one. No shipped recipe sets both, so CI will not catch it. Fix the check as above, and add a test with filtering on.