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

[wdykas]

What does this PR do?

Keep filtered policy support information until the actor loss can exclude invalid tokens consistently.

Previously, current/previous logprob preparation replaced out-of-support -inf values with zero but discarded the corresponding validity mask. The actor then treated those zeros as log probability one. In GSPO, one such token can alter the response-wide ratio and clip an otherwise valid token's gradient to zero.

Preserve support information through ordinary and packed logprob preparation and Megatron/Automodel previous-logprob computation. Previous logprobs keep -inf only on valid tokens; prompt and padding positions are still zeroed (Automodel padding now uses torch.where, avoiding -inf * 0 = NaN). At the actor loss, combine current, previous and generation validity with the existing token/sample masks before sanitizing values and calculating ratios, diagnostics and reductions. Fully excluded responses contribute no actor loss. Keep separate masks and generation logprobs for unfiltered reference KL so actor support exclusion does not silently drop valid KL terms.

Issues

No linked GitHub issue.

Consumers of prev_logprobs outside the actor loss must not see the preserved -inf: KL-in-reward (Reinforce++/GAE) and the OPD advantage give zero contribution at tokens outside the policy support. The unused mask_out_neg_inf_logprobs helper is replaced by mask_filtered_logprobs_outside_tokens.

Sequence-level logprob-error filtering excludes unsupported negative-infinity tokens from its diagnostic mask when top-k/top-p filtering is active. This applies to rollout-side GRPO/PPO filtering (including the single controller) and the single-forward in-loss check. An unsupported token therefore does not reject the whole response or discard otherwise valid reference-KL terms. NaN and positive infinity remain visible; without top-k/top-p, negative infinity still fails a finite sequence threshold. The generic compute_seq_logprob_errors helper keeps its strict semantics. policy_support_excluded_tokens reports the raw count removed from the actor mask in both full and minimal metrics.

Usage

No new configuration. Applies to token-level and sequence-level importance ratios, including GSPO. Filtered-out positions no longer influence actor ratios or gradients.

Sequence averages use the surviving actor tokens. Existing global rollout-batch normalization factors remain unchanged; this does not introduce per-microbatch survivor renormalization. The existing explicit in-loss survivor-normalization mode receives updated survivor counts.

Validation

• 125 focused tests passed covering filtered-support loss/gradients, reference KL, fully excluded responses, OPD/GAE advantages, pinned full/minimal loss compatibility, and single-forward sequence filtering. Includes the CUDA/NCCL packed-loss regression and Megatron previous-logprob preparation.
• The final expanded edge-case suite passed all 62 tests, including metric/gradient equivalence and filtered/nonfinite behavior through GRPO, GRPO-sync and PPO helpers.
• All 35 single-controller advantage-stage tests passed; 23 existing sequence-filtering/GRPO trainer regressions and both synchronous/asynchronous PPO trainer masking tests passed.
• Regression coverage includes the reported unsupported-token example, filtered versus unfiltered negative infinity, NaN and positive infinity, and token/sequence loss and metric equivalence to explicit exclusion.
• Ruff lint/format and diff whitespace checks passed.
• The Automodel-marked regression and three compiled CUDA cases were not run in this validation; full distributed convergence/recipe comparisons remain outstanding.

Before your PR is "Ready for review"

• Read contributor guidelines and signed off commits.
• Added regression tests.
• Ran targeted tests and lint checks.
• Full unit and functional suites run locally.
• Documented normalization and KL semantics above.

Additional Information

This does not make training and inference logits identical or align their top-p implementations; it handles tokens outside policy support correctly when a mismatch occurs. CPU tests patch only a hardcoded target .cuda() placement; filtering, loss and autograd math run normally. The packed regression runs on GPU without that placement patch.

CI:L0 only; no CI:L1 label is added.

[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.

[wdykas]

/ok to test 9c10c40

[wdykas]

/ok to test f40a22d

[wdykas]

/claude review

[wdykas]

Claude Code Review (reposted) — automated review from workflow run 36876625666 at f40a22d3. The bot finished but never posted its result, so the output is copied here verbatim.

---

I've completed a thorough review of PR #4367, reading the full diff, the surrounding loss-function and advantage-estimator source, all downstream prev_logprobs/generation_logprobs consumers, and the new tests, with the testing skill context loaded.

Review summary

LGTM. This is a clean, carefully-scoped fix with strong test coverage. I verified the following and found no critical issues:

• Helper rename is complete — no stale references to mask_out_neg_inf_logprobs; mask_filtered_logprobs_outside_tokens is imported/used consistently in both automodel/train.py and megatron/train.py.
• All -inf-leaking consumers are guarded. The three KL-in-reward / OPD sites in advantage_estimator.py (GRPO compute_advantage L310, GAE _build_token_level_rewards L448, OPD L694) each wrap with torch.where(torch.isfinite(...), ..., 0.0) before the value feeds rewards/advantages. The OPD metrics are computed from the already-sanitized distill_advantages, so last_metrics stays finite (matches the test assertion).
• The actor-loss masking order is sound. policy_valid is built from all three distributions, then token_mask/sample_mask/mask are reduced and curr/prev/generation are zeroed after. Fully-excluded sequences yield token_mask-all-zero → masked_mean → 0/(0+1e-8) with no NaN (the invalid_source="all" test confirms zero gradient).
• The KL-path decoupling is the right fix. kl_generation_logprobs is zeroed only by the original kl_mask (retaining real values on support-excluded-but-valid tokens), while generation_logprobs used by the actor is zeroed by the reduced mask — so switching the on-policy KL importance weight to kl_generation_logprobs (L707) correctly preserves KL terms the actor dropped. kl_token_mask/kl_sample_mask/kl_mask likewise retain the pre-exclusion token set in the KL reductions.
• post_attention_mask is torch.bool, so the new torch.where(post_attention_mask, ...) in automodel/train.py (replacing the -inf * 0 = NaN multiply) is valid.
• Seq-logprob-error interaction is as documented: -inf prev now yields an inf error that still fails any finite threshold (rejection preserved), with only the logged max/mean possibly going inf — acknowledged in the PR body.

Test coverage is comprehensive: token- and sequence-level, each invalid source, fully-excluded responses, unfiltered-KL preservation, a real GPU packed path asserting the exact [0.5, 0, -0.5] gradient, and the automodel prompt/padding-vs-valid distinction. The Megatron train.py edit isn't directly unit-tested, but it only calls the shared (tested) helper and would require a heavier env — consistent with repo norms.

Nothing to flag.

[terrykong]

@mehraakash could you help review?

[wdykas]

@mehraakash any update?

[mehraakash]

LGTM actually!
Ran the tests on GPUs, changes are okay.

[mehraakash]

LGTM

[wdykas]

@mehraakash thanks for the quick response! I will resolve conflicts for merge

[wdykas]

/ok to test 984e315

[wdykas]

/ok to test 3faf55a

[wdykas]

/ok to test ca01e85

[terrykong]

Thanks for this. Building the actor mask in one place in the loss, with separate KL masks, is a clean design, and torch.where instead of multiplying -inf by a mask is the right tool.

First: the PR has merge conflicts with main. Please rebase and re-run CI. (All L0 lanes are green at ca01e85; functional tests did not run.)

Please rewrite the description against main today. #4154 (on main since 2026-10-06) already removes tokens outside the top-k/top-p set from token_mask: the curr side in prepare_loss_input (loss_input.py#L269-L280, linked at the base) and the prev side in the worker, written back by the driver (grpo.py#L3604-L3607). So the actor ratio bug in "Previously …" is already fixed on main. Compared with main, this PR:

• keeps reference KL on those tokens — a real improvement, since KL uses unfiltered logprobs;
• counts them in global_valid_toks, so the actor loss is slightly scaled down (data.py#L1199-L1204). "Existing global rollout-batch normalization factors remain unchanged" is no longer accurate;
• makes the sequence-error check reject whole responses (see the sequence-error comment on loss_input.py#L270);
• drops fully excluded responses.

It also says "all seven changed files"; there are 13.

Most important: the sequence-error regression on loss_input.py#L270. With that fixed and the now-unused token_mask return path removed (comment on megatron/train.py#L869-L873), this PR becomes a net simplification over #4154.

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4367-unsupported-token.html

Generated by Claude Code

[wdykas]

/ok to test 60fb026

[wdykas]

/ok to test 6b6f332

[wdykas]

/ok to test d9168c7

[wdykas]

/ok to test 72c972f

[wdykas]

/ok to test 372c041

[terrykong]

lgtm. just needs tests to be fixed

[wdykas]

/ok to test 926728e

[wdykas]

/ok to test 4334212