[RL] Avoid full-vocab gather for TP policy statistics - #4685

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

Thanks @tianyu-l -- I kept batch-invariant mode on the existing full-gather path and added a dedicated CPU regression test that pins the patched vLLM call to compute_logprobs(logits, labels) without global_vocab_size. The external-contributor workflows are currently awaiting maintainer approval.

i've updated #4582, this should not affect your changes

@claude I'm confused by the input sharding contract of the target function. See my latest comments. Please explain.

Claude finished @tianyu-l's task in 4m 20s —— View job

Input sharding contract of compute_logprobs

I traced the full trainer and generator paths. Short version: the contract is not static "always vocab-parallel". It is selected at runtime by two things -- whether a TP mesh is active on the SPMD stack, and whether global_vocab_size is passed. The patched vLLM call satisfies the unsharded branch, so there's no violation. Details below.

The three regimes compute_logprobs dispatches on

The gate is tp_size = spmd_mesh_size("tp") (loss.py:397), which reads the thread-local mesh stack via current_spmd_mesh() (spmd_types.py:148-156). If no SPMD mesh context is active, it returns 1 even when the process is running TP.

So the answer to "is this function assuming vocab parallel or not?" is: only when a TP mesh is on the stack. The tp_size > 1 branches assume local shards; the tp_size == 1 fallthrough assumes full logits. This was already true before the PR -- the PR only adds the first row (a no-gather branch guarded by global_vocab_size is not None). When global_vocab_size is None, behavior is byte-identical to the old code, so the legacy contract is preserved, not violated.

Why the patched vLLM path is the tp_size == 1 (full-logits) branch

Two independent reasons, either one sufficient:

1. No active SPMD mesh around the sampler. The wrapper only enters a mesh context inside compute_logits/forward via with self.spmd_context(): (vllm_wrapper.py:491), and the mesh stack is pushed only inside that context manager (spmd_types.py:193). vLLM's v2 sampler calls the patched compute_token_logprobs (batch_invariance.py:59) downstream of compute_logits, outside any spmd_context. So current_spmd_mesh() is None -> spmd_mesh_size("tp") == 1 -> the plain branch. The patched call also never passes global_vocab_size, so even the row-1 guard is off.

2. The logits are already gathered to full vocab. You're right that decoder_sharding.py:358 keeps the LM head vocab-parallel even in inference. But compute_logits gathers S(-1) -> R before returning to vLLM (vllm_wrapper.py:496-505), exactly to satisfy the "vllm expects unsharded logits" comment you cited. So by the time the patched kernel runs, the tensor is full [N, V] -- matching what vLLM's native compute_token_logprobs assumes (the logprob.py#L89 behavior you pointed at).

Both facts agree: the patched replacement receives full, unsharded logits, and compute_logprobs correctly routes it through the non-vocab-parallel branch. Passing global_vocab_size there would be the actual bug -- it would select the sharded algorithm on already-gathered full logits and produce wrong reductions -- which is precisely why the batch-invariant path leaves it unset.

On "it used to assume so, but this code violates it"

The pre-PR compute_logprobs was not "always vocab parallel" either -- under tp_size > 1 it gathered S(-1) -> I and then ran plain CE, and under tp_size == 1 it ran plain CE directly. The PR keeps both of those unchanged and inserts an opt-in shard-local branch ahead of them. The input each caller must supply is unchanged for every pre-existing caller.

Suggested doc fix

The one thing worth improving is the docstring: it currently describes the algorithm but not the input contract. Given the confusion, it would help to state the shape contract explicitly per regime, e.g.:

logits is the local vocab shard [T, V_local] when a TP mesh is active (global_vocab_size then selects gather-free vs. legacy-gather); it is the full [T, V] (or [N, V]) tensor when no TP mesh is active, as in the vLLM generator path where compute_logits has already all-gathered the vocab shards.

Add input-contract docs to compute_logprobs →

• Read compute_logprobs and cross_entropy_loss in loss.py
• Read vllm_wrapper.py compute_logits gather
• Read batch_invariance.py patching path
• Read spmd_types.py to confirm the runtime mesh gate
• Explain the contract and answer @tianyu-l

Summary: the contract is runtime-selected, not static. TP mesh active -> local shards; no TP mesh (generator/vLLM patch) -> full logits. The vLLM path hits the full-logits branch both because the sampler runs outside any spmd_context and because compute_logits already all-gathers to Replicate, so nothing is violated. Only a docstring clarification is recommended.
• branch feature/vocab-parallel-policy-stats

@claude check my latest comment. In particular

• Is my comment correct that under batch invariant patching, trainer is using _LossParallelCrossEntropy and _VocabParallelEntropy, whereas vllm generator is using F.cross_entropy and logsumexp + softmax, respectively?
• If so, why could we still achieve batch invariance -- is it because reduction="none" and reduceOp=dist.ReduceOp.MAX.name itself makes numerics sharing invaraint?

Claude finished @tianyu-l's task in 3m 11s —— View job

Batch invariance under the patched logprob path

Short answer: the premise is inverted. Under the shipped batch-invariant configs, the trainer does not use _LossParallelCrossEntropy / _VocabParallelEntropy. Both sides run F.cross_entropy and logsumexp+softmax on full logits, and that identity -- not a numerical property of the MAX all-reduce -- is what preserves batch invariance.

Q1: Does the trainer use the vocab-parallel path in batch-invariant mode?

No. The batch-invariant recipes deliberately leave global_vocab_size unset:

• rl_grpo_qwen3_0_6b_varlen_batch_invariant -> loss_fn=GRPOLoss.Config(), default global_vocab_size=None (alphabet_sort/config_registry.py:1008)
• rl_grpo_qwen3_0_6b_flex_batch_invariant -> dataclasses.replace(loss_config.loss_fn, global_vocab_size=None) (:255)
• rl_grpo_qwen3_5_9b_varlen_batch_invariant -> same replace(..., global_vocab_size=None) (:1130)

So in compute_logprobs, the guard tp_size > 1 and global_vocab_size is not None (loss.py:398) is false on the trainer (TP mesh is active, but global_vocab_size is None). It falls through to the legacy branch (loss.py:416):

The generator hits tp_size == 1 because the patched compute_token_logprobs runs downstream of compute_logits, outside any spmd_context, and compute_logits has already all-gathered the vocab shards to full [N, V]. So both sides execute the same PyTorch ops on the same full-vocab logits. _LossParallelCrossEntropy / _VocabParallelEntropy are on neither side.

Q2: So why is batch invariance preserved -- is it the MAX all-reduce / reduction="none"?

No -- it's preserved because the vocab-parallel path is simply not taken. It is not the case that the sharded reductions coincidentally match the single-tensor ops. They would not be bitwise identical:

• The MAX all-reduce (amax per shard -> ReduceOp.MAX) is exact, since max is associative/idempotent -- the global max is bitwise identical regardless of sharding.
• But the softmax denominator is a SUM: the vocab-parallel path computes sum(exp(shifted)) over each shard and then ReduceOp.SUM across ranks (loss.py:152-155), whereas the generator's F.cross_entropy / logsumexp sums over the full contiguous vocab in one reduction. Floating-point addition is not associative, so partial-sums-then-all-reduce differs from one contiguous sum in the low bits. Same for the entropy weighted-sum (_VocabParallelEntropy, loss.py:261-267).

That divergence is exactly why the batch-invariant recipes opt out: they keep the trainer on the full-gather + F.cross_entropy sequence so its summation grouping matches the generator's bit-for-bit (bit_wise/logprob_diff == 0).

Consequence for your review suggestion

This is the tension with your earlier ask to "when tp_size > 1, assert global_vocab_size is not None and maybe remove this all-gather path." Doing that would break batch invariance: the batch-invariant trainer runs TP=2 with global_vocab_size=None and depends on the legacy gather branch surviving. Your later comment ("we should still fallback to all-gather path under batch invariant mode") is the correct resolution, and this PR already achieves it -- via config (global_vocab_size unset) rather than via a runtime is_in_batch_invariant_mode() check.

One thing worth deciding explicitly: right now the fallback is keyed on global_vocab_size is None, which is a config-level convention, not an enforced invariant. If you'd prefer a hard guarantee, cross_entropy_loss/compute_logprobs could additionally check is_in_batch_invariant_mode() (distributed/utils.py:259) and force the gather path even if a config mistakenly set global_vocab_size. That would make the "batch-invariant => full-gather" contract robust to config drift rather than relying on every recipe author remembering to leave it unset.

• Read compute_logprobs / cross_entropy_loss / _VocabParallelEntropy
• Read batch_invariance.py patch path
• Read DAPO/GRPO losses + global_vocab_size flow
• Confirm batch-invariant configs leave global_vocab_size unset
• Answer both questions
  • branch feature/vocab-parallel-policy-stats

Addressed the latest review in b750d3f. GRPOLoss.Config now inherits DAPOLoss.Config instead of repeating global_vocab_size. compute_logprobs now names and documents the vocab-parallel selection explicitly, checks is_in_batch_invariant_mode(), and keeps the full-vocabulary gather path with a warning even if a recipe mistakenly sets global_vocab_size. I did not assert that every TP call must set it: batch-invariant TP depends on the existing unset-size gather path. The vLLM patch docstrings now state that its input is full logits and the size is intentionally omitted. A CPU regression covers the misconfigured batch-invariant case. Locally, the two runnable focused tests pass; ten optional-dependency cases skip. Formatting, AST, lint, docstring and spelling hooks pass. Project Pyrefly remains limited by missing local dependencies and an invalid configured search path; the targeted check reports only missing vLLM and the pre-existing torch.stack union diagnostic. The new-head Lint, CPU Unit Test, and GPU Unit Tests require maintainer workflow approval; only three Meta checks have passed on this head so far.

Addressed the latest review in b0ac8d10. compute_logprobs now branches first on TP and runtime batch-invariant mode: batch-invariant TP always gathers full-vocabulary logits, while ordinary TP requires global_vocab_size and uses the vocab-parallel implementation. I removed the warning and stopped using global_vocab_size=None as a mode selector; all batch-invariant recipes now retain their model vocabulary size. The vLLM docstring also states why its full-logit call goes directly through the replicated path.

Added regression coverage for both runtime branches and expanded recipe checks. Focused result: 3 passed, 13 optional-dependency cases skipped. Formatting, AST, flake8, docstring, spelling, and link hooks pass. The broader test_loss.py still cannot collect with the local stable PyTorch because it lacks nightly DataParallelMeshDims; targeted Pyrefly remains limited by the repository’s invalid local search path and missing optional vLLM/renderers dependencies.

Also addressed the older test-placement threads in 143cff18: both policy-statistics tests now live under torchtitan/experiments/rl/tests/, and the RL integration workflow points to the new config-test path. Before the path-only move, the locally runnable focused cases passed; collection from the new package path requires the optional vLLM dependency, which the RL CI job installs.