feat: top-p/top-k train sampling with native sampling replay

[mikasenghaas]

Summary

Builds on vLLM 0.28.0 from #3430. Supersedes #3235 — same feature and same user API, but the capture side is vLLM's native sampling-mask support (vllm#49577, released in 0.28.0) instead of prime-rl's engine/IPC monkey patches.

Adds top-p and top-k sampling support for train rollouts (both were hardcoded off). Truncated sampling renormalizes the rollout distribution over the surviving "kept set" of tokens; rollout logprobs already reflect that (logprobs_mode = "processed_logprobs"), but the trainer normalizes over the full vocabulary — so every importance ratio is biased and runs with truncated sampling collapse. This PR records the kept set at sampling time and renormalizes trainer logprobs over the same set: DeepSeek V3.2's "Keep Sampling Mask" (arXiv:2512.02556 §3.1).

Usage

[orchestrator.train.sampling]
top_p = 0.95
top_k = 20   # optional — defaulted to 512 when truncation is on (bounds the kept sets)

That's the whole config — there are no replay flags. Truncated train sampling (top_p < 1 and/or top_k) implies sampling replay end to end:

• Every truncating policy-sourced sampling config gets a top-k bound (top_k respected if set, else defaulted to 512; values above 512 are rejected — the trainer pads each micro batch's masks to the largest kept set, so the bound caps trainer memory). Truncation knobs must be the typed fields — smuggling them via extra_body is rejected, as is temperature = 0. Frozen-source envs are exempt.
• inference.enable_return_sampling_mask is auto-set and persisted into per-node configs; it maps to vLLM's --return-sampling-mask and forces the V2 model runner. Hand-setting is only for standalone-launched servers.
• The trainer is data-driven: it replays masks whenever a batch carries them. The orchestrator enforces that truncating envs actually produce masks (fails fast if the server isn't capturing).
• opd/opsd are rejected at config time (reference logprobs are full-vocab prefill scores).

How it works

Inference — no prime-rl patches. vLLM's --return-sampling-mask records torch.isfinite(processed_logits) after top-k/top-p/min-p filtering and returns it on /inference/v1/generate choices as sampling_mask: list[list[int]] (one list of surviving vocab ids per completion token). The stock tokens endpoint already emits the field, so PrimeRlServingTokens is untouched.

Native-capture constraints, surfaced as config validation / documented behavior:

• Requires the V2 model runner — the launcher sets VLLM_USE_V2_MODEL_RUNNER=1 under the flag, and the flag is rejected together with enable_return_routed_experts (router replay is V1-only).
• Capture is engine-wide: while it is on, vLLM rejects any request with temperature <= 0 or without an effective top_k > 0. The rl entrypoint warns when eval sources share the engine — eval sampling must set top_k (the model's generation config often supplies one) and a non-zero temperature.

Transport (renderers → verifiers → orchestrator → trainer):

• renderers#144 — surface sampling_mask on the generate result (drops the base64 kept_tokens splice from the feat: top-p/top-k train sampling with sampling replay #3235 lineage; no server emits it anymore).
• verifiers#2460 — KeptTokens.from_sampling_mask converts to flat int32 ids/counts arrays; graph attribution validates alignment; Branch.kept_tokens unchanged.
• prime-rl: KeptTokens {ids, counts} (int32 bytes, CSR-style) on TrainingSample/MicroBatch, appended last to keep the positional wire layout stable; packed/truncated/padded alongside the other per-token streams; tensorized as [1, seq, max_kept] with -1 padding.

Trainer (unchanged from #3235):

• Masked positions compute logprob = logits[label]/T - logsumexp(logits[kept]/T) in both the chunked fused LM head (backward restricted to kept ids) and the vanilla path. Positions without a mask use full-vocab logprobs.
• Singleton kept sets give logprob 0 and exactly zero gradient (the entropy-preserving property).
• Entropy stays full-vocab (it's a collapse diagnostic).
• Gemma-family softcapped lm_heads fail loudly on their head assert.

Verification

• uv run ruff check / ruff format --check; uv run pytest tests/unit/test_configs.py tests/unit/train tests/unit/orchestrator tests/unit/inference minus tests/unit/train/models: 300 passed, 4 skipped, 1 failed (test_qwen3_vl_e2e, pre-existing fix: token_id-formatted logprob tokens in the qwen3-vl fake engine #3161; the tests/unit/train/models backward tests crash identically on the base branch on this box).

End-to-end on reverse-text (Qwen3-0.6B-Reverse-Text-SFT, 20 steps, 1 trainer + 1 inference GPU, filesystem rollout transport for payload inspection), all runs from this branch:

• Baseline (no truncation) — regression check: reward 0.19 → 0.79, 0% rollout errors, 100% trainable, mismatch KL 0.0007–0.0132. The engine starts without return_sampling_mask; all 162 orch→trainer micro batches carry kept_tokens = None (wire layout unchanged).
• top_p = 0.95 (auto top_k = 512, with warning): reward 0.23 → 0.85, 0% rollout errors, mismatch KL 0.0006–0.0090 — inside the baseline's band, i.e. replay is exact. Engine logs return_sampling_mask: True + V2 model runner. All 162 micro batches carry masks; per loss-masked token: mean kept 11.3, median 2, p99 169, max 388 ≤ 512, 48.4% singletons, 100.00% mask coverage; sampled token contained in its own kept set at 58,570/58,570 checked positions.
• top_p = 0.95, top_k = 20: reward 0.21 → 0.87, 0% rollout errors, mismatch KL 0.0006–0.0126. All 162 micro batches carry masks; mean kept 3.7, median 2, p99 17, max 19 ≤ 20, 100.00% coverage; membership 60,854/60,854.

W&B: reverse-text/reverse-text-native-replay-{baseline,topp095,topp095-topk20}. Payload sizes were checked by decoding every rollouts/step_*/rank_*.bin: counts always aligns with input_ids and sum(counts) == len(ids) in every micro batch; the mask streams add 40% (top_k 512) / 22% (top_k 20) to reverse-text's tiny batches — the share shrinks with real context lengths since prompt tokens carry no masks.

🤖 Generated with Claude Code