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

[mikasenghaas]

Summary

• Add typed top-p and top-k settings for policy rollouts.
• Capture vLLM 0.28 sampling masks and replay the truncated distribution in the trainer.
• Bound sampling masks with a default top_k = 512 when truncation is enabled.
• Use Model Runner V2 for router replay on standard deployments.
• Support router replay and sampling replay together on Model Runner V2.
• Keep NIXL P/D router replay on V1 and reject the unsupported combined P/D mode.
• Use SamplingMask and sampling_mask across Renderers, Verifiers, and prime-rl.
• Pin merged Renderers and require its published 0.1.12.dev2 build.
• Pin Verifiers to the latest main after its sampling-mask merge.
• Document standalone capture, the NIXL replay matrix, and sampling-mask layout examples.
• Shard sampling masks with labels for multimodal context-parallel training.
• Reject mixed enabled and disabled top-k settings across live-policy train sources.

This builds on vLLM 0.28.0 from #3430. It supersedes #3235. The capture side uses vLLM's native sampling-mask support from vllm#49577.

Usage

[orchestrator.train.sampling]
top_p = 0.95
top_k = 20

There are no replay flags for sampling replay. Truncated policy sampling enables mask capture and trainer replay automatically.

Set trainer.enable_router_replay = true to combine router replay with sampling replay. Standard deployments use Model Runner V2 for both captures. Disaggregated NIXL deployments keep router replay on V1 and cannot combine both replay modes.

Behavior

• Policy sampling with top_p < 1 or top_k gets a bounded sampling mask.
• Truncation without an explicit top_k defaults to 512.
• Different truncating sources can use different top-p values; each gets a positive top-k.
• Mixed top-k capture modes across live-policy sources are rejected.
• Values above 512 are rejected to bound trainer memory.
• The trainer renormalizes each sampled token over the same mask as inference.
• Frozen-source environments do not require masks.
• opd and opsd reject truncated policy sampling because their reference scores use the full vocabulary.
• vLLM rejects sampling-mask requests with temperature <= 0 or without an effective top_k > 0.

Runs

GLM-4.5-Air on scaleswe on 4 nodes

Run	Runner	Router replay	Top-p	Effective top-k	Nodes
glm-air-v1-baseline	V1	Off	1.0	Off	6
glm-air-v2-baseline	V2	Off	1.0	Off	6
glm-air-v2-router-replay	V2	On	1.0	Off	6
glm-air-v2-top-p-0.95	V2	Off	0.95	512	6
glm-air-v2-router-replay-top-p-0.95	V2	On	0.95	512	6

Sampling replay leads to

• no reduction in step 0 KL mismatch
• more stable KL mismatch in conjunction with router replay
• more stable entropy (not rising)
• slower step time (follow-up perf investigation TBD)

Verification

• uv run pytest -q tests/unit/test_configs.py -k 'policy_sources': 2 passed.

• uv run pytest -q tests/unit/test_configs.py -k 'not test_load_configs': 69 passed and 67 deselected.

• Pre-commit checks passed for the changed config and test files.

• uv run pytest tests/unit/train/rl/test_fused_lm_head.py tests/unit/train/rl/test_loss.py -q: 19 passed and 1 skipped after the context-parallel fix.

• uv run pytest tests/unit/test_configs.py::test_combined_replay_uses_v2_runner -q: passed after the documentation review.

• uv sync --all-extras: passed with merged Renderers and Verifiers dependencies.

• uv run pytest tests/unit/orchestrator tests/unit/inference tests/unit/train/rl/test_loss.py tests/unit/train/rl/test_fused_lm_head.py --ignore=tests/unit/orchestrator/test_qwen3_vl_e2e.py -q: 108 passed and 1 skipped.

• The excluded Qwen3-VL test has an existing fake-response fixture mismatch.

• uv run rl @ examples/basic/reverse-text/rl.toml --max-steps 5 --orchestrator.train.sampling.top-p 0.95: completed five trainer and orchestrator steps on the final dependency chain.

• The end-to-end run used Model Runner V2, effective top_k = 512, and sampling-mask capture. Final mismatch KL was 0.0035 with no rollout errors.

• Renderers: uv run pytest tests/test_client.py -q: 21 passed and 5 skipped. Ruff and format checks pass.

• Latest Verifiers main: uv run pytest deps/verifiers/tests/v1/test_graph.py deps/verifiers/tests/v1/test_trace.py -q: 21 passed. Ruff, format, type, and pre-commit checks pass.

• Reverse-text baseline, top-p 0.95, and top-p 0.95 plus top-k 20 runs completed with 100% sampling-mask coverage.

W&B: reverse-text/reverse-text-native-replay-{baseline,topp095,topp095-topk20}.

🤖 Generated with Claude Code

---

Note

High Risk
Changes core RL logprob/importance-ratio math and vLLM runner selection; misconfiguration or missing masks can bias training or fail at runtime, though validation and runtime checks aim to catch incompatible modes.

Overview
Adds sampling replay so policy rollouts with top_p < 1 or top_k stay aligned with trainer importance ratios: vLLM 0.28 returns per-token sampling masks (enable_return_sampling_mask / --return-sampling-mask), and the trainer renormalizes logprobs over the same mask instead of the full vocabulary.

Config & orchestration: TrainSamplingConfig gains typed top_p and top_k (truncation via extra_body is rejected). The rl entrypoint auto-enables mask capture when policy sampling truncates; unbounded truncation defaults top_k = 512, values above 512 are rejected, temperature = 0 and opd/opsd are blocked. Policy train sources must agree on top-k capture mode (engine-wide). Inference: new enable_return_sampling_mask; vLLM env setup prefers V2 for sampling capture and for router replay on standard deployments, while disaggregated NIXL P/D keeps routed-expert capture on V1 and rejects router + sampling replay together.

Data path: SamplingMask on TrainingSample / MicroBatch, encoding in trajectories, packing/padding in the trainer batch builder, and a hard error in TrainSink if truncated rollouts lack masks. Trainer: masks flow through CP sharding; fused lm_head and vanilla-path selective_log_softmax_with_sampling_mask compute mask-renormalized logprobs (Gemma softcap heads explicitly unsupported). Docs cover sampling replay; renderers is bumped for mask support in the rollout stack.

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

[samsja]

From Sami — written by Codex:

Sampling replay currently breaks for multimodal batches with context parallelism. In the deferred VLM path, labels are sharded outside the model, but sampling_mask is left at the full sequence length because it is only sharded when not defer_vlm_cp_to_model. The model then internally shards the hidden states, so the LM head receives local hidden states/labels with a global sampling-mask tensor. The fused path will fail while reshaping the mask, and the vanilla path will see incompatible sequence dimensions.

The sampling mask should be sharded alongside labels even when VLM input sharding is deferred.

[mikasenghaas]

Fixed the multimodal context-parallel sampling-mask mismatch in bd24f08. Sampling masks now shard with labels even when VLM input sharding is deferred. A pre-forward assertion also verifies mask-label alignment. Focused LM-head and loss tests pass: 19 passed, 1 skipped.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

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

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