[RL][Performance] Avoid full-vocabulary TP gather for selected-token logprobs

[chengcuiping]

Problem

Megatron-RL currently computes selected-token log probabilities by requesting gathered vocabulary logits from the model:

runtime_gather_output=True

get_logprobs() then applies selective_log_softmax to the full [batch, sequence, vocabulary] tensor on every tensor-parallel rank.

For large vocabularies and long RL sequences, this materializes and communicates substantially more data than needed: GRPO actor/reference consumers only require the log probability of one selected token per sequence position.

Relevant code: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/rl/rl_utils.py

Tagging @NVIDIA/mcore-oncall for design guidance.

Proposed solution

Add an opt-in Megatron-RL path that:

1. keeps the LM-head output vocabulary-sharded;
2. uses the existing GPT output_processor contract;
3. computes selected-token log probabilities as the negative token-level cross entropy using the existing native megatron.core.tensor_parallel.vocab_parallel_cross_entropy;
4. returns only [batch, sequence - 1] selected log probabilities;
5. preserves the existing gathered-logits implementation as the default and compatibility fallback.

For label smoothing 0 and identical targets:

log_softmax(logits)[target] == -cross_entropy(logits, target)

The native vocab-parallel cross-entropy implementation already owns target-shard masking, numerically stable normalization, tensor-parallel reductions, and differentiable backward.

The proposed configuration is:

--rl-use-vocab-parallel-selected-logprobs

It would initially be default-off.

Prototype correctness

The prototype has been tested with Qwen3-4B, BF16, TP=1/2/4.

Results include:

• selected-logprob maximum absolute error: 1.9073486328125e-06;
• FP32 operator logprob maximum absolute error: 2.38418579e-07;
• actor gradient minimum cosine similarity: 0.9998191;
• masked positions have exactly zero gradient;
• reference/no-grad and actor/backward paths pass;
• packed-sequence execution passes;
• 20 consecutive distributed iterations complete without divergence or hangs;
• a controlled native GRPO workflow passes 10 optimized updates.

The native workflow exercises:

• compute_logprobs_batch;
• train_rl.forward_step;
• calculate_grpo_loss;
• reference/old-policy logprobs;
• actor forward and backward;
• gradient synchronization;
• optimizer updates.

Prototype performance

Controlled Qwen3-4B TP4 native GRPO update, sequence length 2048, response length 1024, global batch size 32:

Configuration	Microbatch	Microbatches/update	Median update time
Gathered baseline	8	4	8.8432 s
Vocab-parallel selected logprobs	8	4	8.5022 s
Vocab-parallel selected logprobs	16	2	8.2641 s

Observed results:

• common-shape A8 to D8 uplift: 3.86%;
• capacity-realized A8 to D16 uplift: 6.55%;
• candidate-grid maximum stable tokens/microbatch: 65,536 → 131,072;
• logical removal of the full-vocabulary logits collective on the optimized path.

The same-shape global peak allocated memory did not improve in the native GRPO update (approximately 1.95% higher in the optimized measurement), because transformer activations dominate the global peak. The intended claim is therefore increased feasible microbatch capacity and removal of full-vocabulary materialization, not general end-to-end memory reduction.

These measurements use a controlled frozen rollout and are not presented as a complete Agentic-RL lifecycle benchmark.

Compatibility and fallbacks

The existing gathered path remains available. Modes not yet proven compatible will explicitly fall back rather than silently changing behavior, including as applicable:

• CUDA Graph capture/replay;
• batch-invariant execution;
• unvalidated CP/PP configurations;
• non-standard full-logit consumers;
• non-zero label smoothing.

The implementation does not add a fused kernel and does not depend on Liger Kernel.

Alternatives considered

1. Continue gathering full vocabulary logits and optimize only selective_log_softmax. This does not remove the dominant [B, S, V] materialization and TP gather.
2. Add a new fused LM-head/cross-entropy kernel. This is a larger and overlapping kernel-level change and is unnecessary for the initial RL consumer optimization.
3. Gather logits only on one rank. This reduces replicated post-processing but retains the full-vocabulary communication and leader memory requirement.

Design questions

Before submitting the implementation PR, I would appreciate guidance on:

1. Is using the existing GPT output_processor plus native vocab_parallel_cross_entropy the preferred integration boundary for Megatron-RL?
2. Should this initially remain behind a default-off RL flag, or should it be selected automatically when the configuration is supported?
3. Is an explicit fallback to the existing gathered path acceptable for CUDA Graph and other unvalidated compatibility modes?

I have a working implementation, distributed correctness tests, and controlled native GRPO measurements ready for a draft PR after design alignment.