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

[tianyi-zhang-02]

What does this PR do?

Removes positions whose current-policy logprob was -inf from actor-side reductions without narrowing the reference-policy KL.

mask_out_neg_inf_logprobs now returns the keep mask. Regular and packed callers publish it as curr_logprobs_keep_mask, and ClippedPGLossFn derives a separate actor mask instead of changing token_mask.

The diagnostic means remain exact across uneven microbatches and DP ranks: the loss emits additive numerator fragments plus the retained actor-token count, and GRPO, synchronous GRPO, PPO, and SingleController divide once after global aggregation.

Quantity	Mask	Normalization
actor loss	retained current-policy tokens	existing global valid-token scale
ratio / clamped-ratio / entropy diagnostics	retained current-policy tokens	globally summed retained-token count
reference-policy KL and mismatch diagnostics	original token mask	unchanged

Prev-policy keep-mask propagation would require widening LogprobOutputSpec and remains out of scope.

Validation

Current head: 3ebc1e860cb9f95e5eda38be18a69324cfcf5df7, including main at fd7112c374af23b3bcc392637fb737495f8a3f31.

On 2026-09-21, all four changed unit-test files passed on Runpod Secure Cloud: 161 passed, 0 failed/skipped. Environment: 1× H100 80 GB, nvcr.io/nvidia/nemo-rl:v0.7.0, driver 580.126.09, CUDA 13.0, PyTorch 2.11.0+cu130. The current-main test fixture required pytest==9.0.3 and pinned nemo-lens commit b85578f, installed only in the temporary Pod. No multi-GPU functional run was attempted in this single-GPU session.

The earlier CPU-only validation on head cbfd57d454fdb8e6f963b4d241bece9c2b6f2df7 also passed five exact-value targeted tests (including a SingleController caller with four uneven microbatch/rank fragments) and 65 related CPU tests; 44 GPU-only tests were skipped then.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tianyi-zhang-02]

One consequence of this change that the description did not spell out, and that a reviewer should decide on rather than discover.

Folding the narrowed mask into token_mask does not only affect the actor term. The reference-policy KL reduces with the same mask, so it narrows too — and it is computed from curr_logprobs_unfiltered, which is finite at those positions. The terms being dropped are real ones.

Happy to change it if you would rather the KL stayed on the full mask; flagging it because it is the kind of thing that is easy to miss in a diff this size.

[tianyi-zhang-02]

Followed up on the open question in my earlier comment rather than leaving it to review.

Folding the keep mask into token_mask narrowed every reduction in the loss, and two of them should not narrow.

The reference-policy KL reduces with token_mask, but it is computed from curr_logprobs_unfiltered, which is finite at exactly these positions. So the terms being dropped are real. Worse, they are dropped selectively: a token gets filtered because the training policy gives it near-zero mass, while the reference checkpoint may not. That is where the two distributions disagree most — the part of the KL you least want to lose.

token_mult_prob_error, gen_kl_error, policy_kl_error and the JS divergence read only prev and generation logprobs. The substitution never touches them, so narrowing there just hides mismatch the metrics exist to surface.

Latest push publishes the mask as curr_logprobs_keep_mask and applies it only to things derived from curr_logprobs: actor loss (both reduction modes), the GSPO sequence-level ratio, the entropy approximation, the VAPO NLL term, and the probs_ratio metrics. The substituted 0.0 fabricates a ratio of exp(-prev_logprob) there, which would show up as a spurious spike in probs_ratio_max.

global_valid_toks still comes from the full token mask, so filtered positions contribute zero rather than reweighting the rest — same as before.

The test asserts both halves. It goes red under either mutation: ignoring the keep mask, or folding it back into token_mask :)

[tianyi-zhang-02]

One more thing the description does not say, and should.

mask_out_neg_inf_logprobs has five call sites, not two. Besides the two in loss/utils.py for curr_logprobs, it is called on prev_logprobs from automodel/train.py, megatron/train.py and dtensor_policy_worker.py. Those three take the same substituted 0.0, and they discard the keep mask. Deliberately, but it is a real gap rather than a non-issue.

It cannot be closed here because the narrowed mask has no way to travel with prev_logprobs. LogprobOutputSpec carries logprobs only, so widening it is a separate change with its own blast radius.

What that leaves: if the previous policy had a -inf where the current policy kept a token, prev_logprobs there is 0.0 while curr_logprobs is finite. The ratio at that token becomes exp(curr) rather than exp(curr - prev). Nothing here narrows it away, since the keep mask this publishes is the curr-side one.

Two reasons I left it. Both sides are filtered under the same sampling_params, so the positions coincide except where the policy moved between rollout and update. And closing it properly means widening the spec, which I would rather do in its own PR than smuggle into a mask fix. Happy to do that follow-up, or fold it in here if you would rather have one change.

The three call sites carry a comment saying the same thing, so it is not silent in the code either :)

[yuki-97]

@jinglinglingling to review

[jinglinglingling]

LGTM. No concerns from my side.

[tianyi-zhang-02]

Merged current main in. Should be clean merge w/ no conflicts.

[tianyi-zhang-02]

Ran this on Linux, since the body notes the two caller tests could not execute on the macOS host. All six new tests pass, and test_loss_functions.py + test_sequence_packing_fusion.py + test_utils.py come to 106 passed at this head. Python 3.13.13, torch 2.11.0+cu130, RTX 4090.

One correction for the merge description: the body says the keep mask is "folded into token_mask in both the regular and packed loss-input paths", but bc9b49e moved away from that. The code publishes curr_logprobs_keep_mask and leaves token_mask untouched, which is what keeps the reference-policy KL reducing over real terms. That distinction is the substance of the change, so it would be a shame for the older wording to be what lands in the log.

Three things I measured while checking it.

1. force_on_policy_ratio aliases prev to curr at :415, above the diagnostics, so the substituted 0.0 does reach them and the premise at :406 does not hold in that config. With π_gen = exp(-5) at the filtered position, token_mult_prob_error reads 50.14 whether or not keep is passed. Under sequence_level_importance_ratios the sequence weight reads exp(5) = 148.41, and that one scales the gradient rather than a log line. Fixed in #3778, stacked on this.

2. global_valid_toks still counts the positions the numerator now drops. It is all-reduced from the un-narrowed token_mask at automodel/data.py:414, upstream of the filtering, so the actor_mask reductions this PR introduces divide by a denominator that includes them. On a 3-token microbatch with one position dropped, probs_ratio goes 1.0 → 0.667. At realistic filtered fractions it is a sub-percent scale on the actor term, so it is not urgent, but it does arrive with this PR and I could not tell from the diff whether it was intended.

3. The deferred prev-side case is not diagnostics-only. The worker comments are right that it needs LogprobOutputSpec widened, but two of the reductions it reaches move the gradient: seq_lp_diff at :584, whose per-sequence scalar multiplies every token, and the seq-mask-tis gate at :659, which can zero a whole sequence. Quantified below on the real ClippedPGLossFn, B=4 S=512, median of 8 seeds, force_on_policy_ratio off throughout:

At 1% of valid tokens substituted, gen_kl_error reads 783× its honest value and the GSPO sequence weight 2.3e4×. Happy to take the spec widening as a separate PR if you want it.

cc @jinglinglingling

[tianyi-zhang-02]

Follow-up on final SHA cbfd57d454fdb8e6f963b4d241bece9c2b6f2df7 after correcting the retained-token denominator:

• ratio, clamped-ratio, and entropy now aggregate raw numerators plus the narrowed actor-token count across microbatches/ranks, then divide once;
• actor-loss scaling is unchanged, and the reference-policy KL still uses the full mask;
• the caller-level regression uses four uneven SingleController fragments and checks the exact global result;
• on the Runpod image (Python 3.13.14, torch 2.11.0+cu130, CPU-only), 5 targeted and 65 related tests pass; 44 GPU-only cases skip; Ruff check/format pass.

cc @jinglinglingling — this is the final behavior after the earlier reviewed head moved.

[tianyi-zhang-02]

Closing to avoid competing with the active top-k/top-p masking direction in #4154 (and the related masked-value safety work in #4120). This PR intentionally made a different KL-mask choice, so the right next step is maintainer clarification rather than keeping two policy-loss semantics open in parallel.