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

[tdene]

What does this PR do ?

Fix issue with IS and topp/topk masking; see offline discussion.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[tdene]

/ok to test 5f2b68f

[tdene]

/ok to test bc23541

[yuki-97]

Thanks for pushing both implementations up rather than picking one — having d4d8a09 and bc23541 side by side made the tradeoff much easier to judge. The megatron side derives has_token_mask from config rather than from the data so every PP rank agrees before the broadcast (megatron_policy_worker.py:2202-2203).

Let's go with this implementation. Carrying the mask explicitly means a downstream consumer only has to read token_mask to be correct, whereas d4d8a09 also requires every consumer of prev_logprobs to know to test isfinite — an unwritten rule that is easy to miss.

Scope: the new token_mask entry on LogprobOutputSpec and its propagation through the three policy workers, the GRPO/PPO drivers, and prepare_loss_input / prepare_packed_loss_input.

Checked:

• Every implementation of get_logprobs and every caller of it, for who propagates the new key and who drops it.
• Shape and layout of the new mask across sequence packing, CP, the dynamic-batching reorder, the PP broadcast, and the DP from_batches merge.
• Whether the in-place data["token_mask"] rewrite reaches the loss fn on all four wrapper paths, and whether it leaks across microbatches or PPO epochs.
• The reference-policy path, where use_reference_model disables top-k/top-p, so no analogous key is needed there.
• Existing coverage — tests/functional/grpo_topp_topk.sh and grpo_megatron_generation_topp_topk.sh do exercise the headline config on dtensor v2 and megatron.

Remaining comments: two callers that do not propagate the new key — the first of them decides whether the "just read token_mask" guarantee actually holds on the SingleController path — one on the reference-KL reduction losing the same positions, and three nits (a pre-existing test gap, a duplicated _pad_and_concat, and a presence check keyed off data rather than config).

[tdene]

Team review of PR #4154 at ab86ede (agents: rl-expert, bug-finder, test-agent, comment-reviewer, design-reviewer, devil-advocate; every inline finding was independently re-verified by the leader, and all CPU-runnable checks were executed against this head).

Overall: the token_mask design is implemented consistently across the three producers, the four driver sites and the TQ mixin; the five addressed threads from the previous review (T1, T3, T4, T5, T6) are complete at this head, and the rebase port is byte-identical to the reviewed hunks apart from the pinned-memory return adaptation. One correctness issue (inline, two one-line casts): the TQ write-back changes the token_mask column's dtype, which TransferQueue rejects silently; it regresses the two shipped topp0.9-topk50-tq_mooncake nightlies.

Merge gates (not code asks):

• No CICD NeMo RL run exists on this head; the last one was on bc23541 (2026-09-17, Lfast, pre-rebase, 51 main commits ago). The commit that addressed the review touches 12 files, so /ok to test ab86edeab3e9c8b597f04de822ca1ee9e38e3d3e is needed before anything is CI-proven. The GPU-gated test_clipped_pg_loss_kl_penalty[1-...] parametrization and tests/unit/utils/test_tensor_ops.py have never executed in CI; the gate is the literal PASSED lines in the L0_Unit_Tests_Algorithms / L0_Unit_Tests_Other logs. (A CPU copy of the GPU test passes here with the hand-computed values, and a mutant with the narrowing removed fails, so a pass is expected.)
• CI:Lfast is adequate for the merge round: it runs both headline functionals (grpo_topp_topk.sh, grpo_megatron_generation_topp_topk.sh) plus the PPO/distillation/topk paths touched by pad_and_concat, and CI:L1 would add only one more filtering test (grpo_multiturn.sh, dtensor, top_k=8000; the megatron multiturn twin is disabled). What no tier covers pre-merge is the data plane with filtering on — the two shipped nightlies that do (nightly.txt:189-190) are exactly the configuration the inline dtype finding breaks.
• No tier runs megatron PP>1 with top-k/top-p (every top-k functional and both TQ nightlies are PP=1), so the PP broadcast of the new key is untested anywhere; low risk, since broadcast_tensor ships shape and dtype metadata.
• The packing/CP layout cell of the previous review's "Checked" list was re-derived at this head except where #4106 (megatron/data.py) landed after that review; the fast-tagged top-k/top-p functionals on this head are the current-stack evidence for it.

Context — no action:

• Order of consumers is identical on the Ray drivers, grpo_sync and SingleController (logprobs → seq-logprob-error masking → advantages → train), so every path sees the narrowed mask; the only divergence is metrics-only (response_advantages selection uses the rollout mask in grpo.py:3815-3819 and the narrowed one in grpo_sync.py:1168).
• On PPO the write-back precedes GAE, so if the final response token is a -inf position the terminal reward lands one token earlier — consistent with the design.
• test_clipped_pg_loss_kl_penalty computes global_valid_toks after prepare_loss_input (=1) while workers all-reduce it before the forward from the un-narrowed mask (=3), so 0.0368 is valid loss arithmetic rather than the value a worker would log.
• T2 (reference-KL at -inf positions) stays an open discussion with yuki-97; the inline docs comment on loss_input.py:278 is the minimal outcome if the current behaviour stands.

Linter: pre-commit run --all-files at this head: every hook passes (ruff, ruff-format, taplo, pyrefly, minimize-check).

Devil's advocate: 6 confirmed (the dtype finding upgraded to critical) | 0 disputed | 1 downgraded and dropped (a dtensor-worker predicate hoist with no concrete divergence scenario). PR necessity: needed — a non-trivial functional change that removes the exp(-gen) importance-weight blow-up at -inf positions.

Generated by Claude Code

[tdene]

/ok to test 5f1b1af

[tdene]

/ok to test 10da9fa

[yuki-97]

Scope: re-review at 10da9fa — the fixes since the previous pass (ab86ede, 5f1b1af) and the new -topp0.9-topk50-is nightly.

Checked:

• Previous threads — each fix re-read at head; the resolved ones are replied to with the commit that carried them.
• Data-plane write-back of token_mask — packed and ordered like prev_logprobs, and grpo_sync, train_presharded and the SingleController advantage stage all read it only after the logprob stage. The dtype cast keeps whatever dtype the column was written with (int64 from the GRPO rollout, float32 from the SingleController reassembler).
• pad_and_concat — equivalent to every copy it replaced (2D, 3D top-k, teacher payload, both value workers).
• New nightly — use_importance_sampling_correction: true on top of the Megatron top-p/top-k recipe, and the nightly budget cap moves to exactly the total with the new entry.

Remaining comments are on documenting, in the GRPO guide and next to the unfiltered KL recompute, that -inf positions are also dropped from the reference-KL term — keeping the current behaviour.

nit, not blocking merging — the data-plane stage→column note in schema.py:44-46 does not yet mention the token_mask rewrite.

[tdene]

/ok to test 12485ed

[tdene]

/ok to test 2cb0fe0

[tdene]

/ok to test ba53070