feat: support top-k / top-p sampling with trainer-side replay

[mikasenghaas]

Summary

• Add top_k and top_p to TrainSamplingConfig so they can be enabled for training rollouts (previously only EvalSamplingConfig exposed them).
• Plumb the sampling args from the orchestrator to the trainer (temperature / top_k / top_p on TrainingSample and MicroBatch).
• Replay the same truncation on the trainer side when computing logprobs, so the importance ratio against vLLM's processed_logprobs stays unbiased.

Closes #2600.

Why this needs trainer-side replay

prime-rl already configures vLLM with logprobs_mode="processed_logprobs" (inference.py:412), so inference_logprobs returned by vLLM are computed from the truncated, renormalized distribution after top-k / top-p. If we enable top-k / top-p sampling but the trainer keeps computing logprobs over the full vocab, the importance ratio (trainer / inference) is biased.

Implementation notes

• apply_top_k_top_p (in trainer/rl/loss.py) mirrors vLLM's apply_top_k_top_p_pytorch reference. Verified bit-exact against the reference (same -inf mask positions, zero value diff) across top-k only, top-p only, and combined.
• The helper accepts optional labels. We restore the label token's original logit after masking so FP-precision boundary cases (the sampled token falling just outside the trainer's top-k under bf16 forward) don't push its logprob to -inf and blow up the masked loss. Without this guard, even step 0 with identical inference / trainer weights hits Loss: inf because of bf16 ranking jitter at the boundary.
• Sampling args are stored as scalars on MicroBatch (one value per micro batch, not per-token). The packer enforces that samples sharing a micro batch also share (training_mode, temperature, top_k, top_p) — non-restrictive in practice because a training run shares its sampling config.
• The fused LM head (FusedOutputLinear) rejects non-trivial truncation: the chunked kernel streams over vocab chunks and can't find a global top-k threshold without materializing the full logits. The vanilla LM head (the RL default) is unaffected.
• Entropy is computed on the pre-truncation distribution so the metric stays comparable across runs that do / don't truncate.

Multi-run compatibility

MultiPacker.pack() groups samples by run_idx and calls prepare_batch separately for each run (packer.py:320-330). The scalar-sampling-args constraint then only needs to hold within a single run, which it does naturally since one run = one RL config = one sampling config. Verified end-to-end:

• Per-run isolation: each run's scalars flow through to its own micro batches; the trainer broadcasts each batch's scalar inside its loop, so truncation matches what inference applied for that run.
• LoRA routing preserved: idxs=[run_idx] * len(run_samples) sets lora_num_tokens[run_idx] = N correctly.
• Mixed envs within a run: if two envs in the same run use different sampling configs, the packer splits cleanly into separate micro batches. Packing efficiency drops slightly in this case, correctness is preserved.
• Dummy batches added for worker-count divisibility deepcopy from a real batch, so they inherit the same (temperature, top_k, top_p) and the trainer applies truncation normally — the label-safety guard keeps logprobs finite and loss_mask=False zeros the loss contribution.

Future work: defer-to-vLLM-defaults

TrainSamplingConfig currently requires explicit values for top_k / top_p (defaults -1 / 1.0, both disabled). Supporting None — "use whatever vLLM applies, which comes from the model's generation_config.json" — would be nice for ergonomics, but vLLM doesn't echo the resolved sampling params back in its response, so the orchestrator wouldn't know what to record on the TrainingSample, and the trainer wouldn't know what to mirror. The clean implementation is to load the model's GenerationConfig orchestrator-side at startup, substitute None → default value before sending, then both vLLM and the trainer use the same explicit values. Left out for now.

Verification

Ran reverse-text RL with Qwen3-0.6B (default generation config: top_k=20, top_p=0.95) on 2 GPUs. Mismatch KL is small and stable, on the same order as a no-truncation baseline — confirming the trainer's truncation matches what vLLM did:

step	with top_k=20 + top_p=0.95	baseline (no truncation)
0	0.0007	0.0010
1	0.0009	0.0012
2	0.0018	0.0014
3	0.0071	0.0054
4	0.0064	0.0045
5	0.0019	0.0021

For comparison, running with inference-side top-k / top-p but trainer-side truncation disabled gives ~2-3x larger Mismatch KL at step 0 (0.0025 vs 0.0007) even with mild truncation (top_k=20 keeps ~99% of mass on Qwen). With more aggressive truncation the gap would be much wider.

All 33 unit tests in tests/unit/orchestrator/test_batch.py, test_trajectories.py, test_teacher_logprobs.py, and tests/unit/train/rl/test_packer.py pass.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes sampling parameter plumbing and logprob computation in the RL trainer; mistakes could bias importance ratios or break training when using fused LM heads.

Overview
Adds train-time top_k/top_p sampling controls (and enforces temperature > 0) and plumbs these sampling args from orchestrator rollouts through TrainingSample/MicroBatch into the trainer.

Updates batching/packing to treat sampling args as scalar per microbatch (no more per-token temperature lists), splitting microbatches when (temperature, top_k, top_p, training_mode) differ.

Implements trainer-side replay of vLLM-style top-k/top-p truncation via apply_top_k_top_p before computing logprobs, and explicitly rejects truncation when using the fused LM head; tests are updated accordingly.

^{Reviewed by Cursor Bugbot for commit a12741d. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 0548ac5. Configure here.}

[cursor]

Top-p masking applied to logits instead of raw probabilities

Medium Severity

The top-p implementation computes sorted_probs = sorted_logits.softmax(dim=-1) then uses cumulative probabilities to create a mask, but then applies that mask to sorted_logits (the pre-softmax values) using masked_fill. After masking some logits to -inf, when selective_log_softmax later recomputes log_softmax, the renormalized probabilities will differ from what vLLM computed because the softmax denominator changes. In contrast, vLLM's apply_top_k_top_p_pytorch determines which tokens to remove via the same cumulative probability logic but the key difference is that vLLM computes processed_logprobs on the already-truncated distribution in a single pass, whereas here the softmax used to determine the mask boundary and the final log_softmax in selective_log_softmax are two separate computations over the same truncated set—which is actually equivalent. On closer inspection the math is consistent since the same tokens end up masked in both the boundary-finding softmax and the final log_softmax. However, there is a real discrepancy: after top-k masking sets some logits to -inf, the top-p branch recomputes softmax on these partially-masked logits, finding the cumulative boundary on the renormalized (post-top-k) distribution. If vLLM applies top-p on the original pre-top-k distribution rather than post-top-k, the mask boundaries would differ. The PR claims bit-exact matching against vLLM's reference but vLLM applies both filters to the same original logits distribution simultaneously rather than sequentially.

 

^{Reviewed by Cursor Bugbot for commit 0548ac5. Configure here.}

[mikasenghaas]

Bit-exact verified against apply_top_k_top_p_pytorch with both top_k and top_p set (0 mask disagreements, 0 value diff). vLLM applies top-p the same way — sort, mask top-k to -inf, then softmax(post-top-k-sorted_logits) + cumsum + mask top-p tail. The order is equivalent.