Repository navigation
feat: write router-replay ids and sampling masks by handle on multi-node runs - #3820
Closed
faresobeid wants to merge 4 commits into
Closed
faresobeid wants to merge 4 commits into
faresobeid wants to merge 4 commits into
Conversation
faresobeid
force-pushed
the
feat/payload-by-handle
branch
from
October 3, 2026 22:06
d540660 to
f3fd1bd
Compare
This was referenced Oct 3, 2026
faresobeid
force-pushed
the
feat/payload-by-handle
branch
from
October 3, 2026 22:11
f3fd1bd to
b430fc3
Compare
faresobeid
marked this pull request as ready for review
October 3, 2026 23:06
faresobeid
force-pushed
the
feat/payload-by-handle
branch
from
October 3, 2026 23:28
b430fc3 to
40bb9e0
Compare
There was a problem hiding this comment.
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 40bb9e0. Configure here.
This was referenced Oct 4, 2026
faresobeid
force-pushed
the
feat/payload-by-handle
branch
from
October 4, 2026 06:47
f7b0763 to
19ec1ae
Compare
… router merges by-handle routing
This was referenced Oct 4, 2026
This was referenced Oct 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

Reworked onto vLLM 0.31 (stacked on #3819 → #3888)
Branch
feat/payload-by-handlenow holds one commita4a98c5f2on top of #3819's1bcca4115. This PR is closed, so GitHub still shows the old head.engine_client.vllm_config.aux_output_config.enable_return_routed_experts(0.31 moved the flag offModelConfig). The writer is unchanged: 0.31's engine output still carriesrouted_experts(concatenated AuxOutput chunks fromrouted_experts_prompt_start) andsampling_mask.token_ids.TrainSink.process_groupdeletes the payload files that no queued sample points into, right after the group is processed: stale-voided groups, groups with no trainable survivors or rejected by the curriculum, untrainable branches, and samples pruned byconstant_trainer_batch_size._drop_staledeletes the files of queued traces it sweeps. The per-version GC (v<k>deleted once no unread batch can reference it) still handles shipped batches and turns written by in-flight rollouts that were cancelled.payload_rootis set explicitly).pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/inference: 307 passed; ruff clean.payload_dir): 16/16 file rows equal to the inline routing (pos == start) and masks (pos == prompt_len). Response size 33.9 MB → 0.57 MB; mean latency by-handle 5.12/4.43 s vs inline 6.43/4.60 s. Trainer tensors from those rows: 140/140 bitwise equal to inline.deps/renderers(a79a10f) anddeps/verifiers(410667d6) pins to merged commits. Both are one commit on top of chore: bump vllm to 0.31.0 #3888's pins.Stacked on #3819 (landing order below); rebased on #3819 @ 99fa8c8. Before merging, move the
deps/renderers(a79a10f) anddeps/verifiers(410667d6) pins to the merged commits.What
This PR turns on passing router-replay ids and sampling masks by handle on multi-node runs, and owns the payload directory's lifecycle. Inline stays everywhere else.
orchestrator.payload_root.rlsets it to<run dir>/payloads(absolute) when router replay or sampling-mask capture is on, the deployment ismulti_node, and inference is not disaggregated. Inference needs no config: the directory travels with each request.payload_dir = <payload_root>/v<k>, wherekis the rollout's own dispatch policy version (self.policy.versionwhen it callsenv.run). That is never older than its group's start, which the staleness sweep keys on, so the GC bound below holds even for members dispatched many versions after their group started.Env._samplingputs it inextra_body.extra_args; verifiers'wire_args()flattensextra_bodyinto the generatesampling_params, and renderers forwards those verbatim, so verifiers and renderers need no request-side change. Checked ad hoc: vLLM'sGenerateRequestparses it intosampling_params.extra_args["payload_dir"], merged with any userextra_args.serving_tokens.py): a request withextra_args.payload_dirgets_PayloadCapture, which takesrouted_experts/sampling_maskoff the engine output (upstream encodes nothing inline), writes both into one file<payload_dir>/<uuid>.binin a thread, closes it before responding (NFS close-to-open), and returnspayloadsegments on the choice. Routing row i is positionrouted_experts_prompt_start + i, as uint8/uint16 (compact_routed_experts, split out ofserialize_routed_experts). Mask row i is completion token i atprompt_len + i, int32[tokens, widest row]padded with -1. Requests withoutpayload_dir(eval, frozen-source rollouts, standalone clients) keep today's inline path unchanged.trajectories.py):trace_to_samplestakesbranch.payload(verifiers#2765), clips it to the sample and appends a 1-row segment for the final token pointing at the previous row's bytes, which is what the inline path's placeholder copies, so tensors match inline. A branch with any other routing gap drops its routing and logs a warning, likeBranch.routed_expertsreturning None.Orchestrator.on_policy_update): after inference appliesv{step}, deletev<k>fork < min_fresh_version(step + 1, max_off_policy_steps). Why that bound is exact:sand then publishesv{s}(dataloader.get_batch()comes beforeweight_sender.broadcast(model, step=progress.step)intrain.py), and the orchestrator sender and trainer receiver advance one batch per step from the same index. So whenv{step}is applied, every batch up tostephas been read.b ≥ step + 1ships only rollouts whose group started atmin_fresh_version(b)or later (TrainSink._drop_stale, the hard guarantee), and a rollout'skis at least its group's start.v{resume_step}. A directory that a still-running stale rollout recreates after GC is caught by the next collection.docs/inference.md§ Payloads by handle.Not supported with disaggregated P/D yet:
rlnever enables it there, because the PD router merges only inline routing (PrimeIntellect-ai/router#60). Settingorchestrator.payload_rootby hand on P/D loses prefill routing, whichtrace_to_samplesreports with a warning and drops.Validation
tests/unit/inference/test_serving_tokens.py: one new test checks that by-handle post-process writes routing + mask into one file, returns the right segments, and clears the inline fields. The existing inline test is unchanged apart from awaitingpost_process.pytest tests/unit/test_configs.py tests/unit/orchestrator tests/unit/inference tests/unit/train/rl -m "not gpu"(with verifiers#2765 on the path): 318 passed.ruff checkandruff format --checkpass.max_off_policy_steps=2,v0..v11present andstep=10keepsv8..v11._sample_payloadrepeats the last routing row for the final token (reader output equals the inline construction) and drops routing with a gap.examples/advanced/qwen3-30b-a3b/swe.toml(multi-node) +--trainer.enable-router-replaysetsorchestrator.payload_rootto<run dir>/payloads. Without router replay it stays None. Single-nodereverse-text+ router replay stays inline.GPU validation (2026-10-04)
Setup: 8x H200, vLLM 0.30.0, one integration branch (#3832 ⊃ #3831 ⊃ #3830, + #3794, #3793, #3820 ⊃ #3819) with router#61 + router#60. Exact comparisons use
VLLM_BATCH_INVARIANT=1(BI) with Qwen/Qwen3-30B-A3B (TP1, 48 layers, top-8) and the same token ids and seeds (temperature 1, top_p 0.95, top_k 50).rlby handle (single node,orchestrator.payload_rootset, CP=2, mini-glm-moe, reverse-text, 6 steps): first attempt crashed at step 1 (TrainSinkchecked only the inlinesampling_mask). Fixed in 6dcf1de; then 6 steps, 0% errors, mismatch_kl mean 6.7-7.3e-5 (same as inline), 435 payload files / 67 MB in v0, v2..v5.trainer.loss.score_centeringon,rldoes not auto-setorchestrator.payload_root(payloads stay inline), and it rejects a hand-setpayload_roottogether with score centering. Both are dropped once feat: by-handle payload fields table (TRR weights, score-centering logprobs) #3836 and feat: exact score centering over the sampling mask (arXiv:2609.20807) #3794's writer hook land.Landing order
deps/renderers/deps/verifiersto the merged commits first)rlturn by-handle on for disaggregated P/DFollow-ups
payloadsegments with the inline merge's rule, then letrlsetpayload_rootfor disaggregated P/D too.[ids | weight bits]by handle (or a separaterouter_weightsfield, 5 B/slot instead of 8).compact_routed_expertsmust not narrow it.sampling_mask_logprobsfield, split in the writer.Note
Medium Risk
Changes the multi-node rollout→trainer data path and shared-filesystem lifecycle; misconfigured GC or NFS visibility could drop replay data, though eval/single-node inline behavior is unchanged and incompatible modes are config-blocked.
Overview
Multi-node RL can now pass router-replay ids and sampling masks by file handle instead of huge inline HTTP bodies, via new
orchestrator.payload_root. Therlentrypoint auto-sets it to<run dir>/payloadswhen router replay or sampling-mask capture is on, deployment ismulti_node, and inference is not disaggregated P/D; hand-setting it with disaggregated P/D ortrainer.loss.score_centeringis rejected.Train rollouts get
extra_args.payload_dir = <root>/v<policy version>through the dispatcher and env sampling config. vLLM’sPrimeRlServingTokenswritesrouted_expertsandsampling_maskinto per-choice.binfiles and returnspayloadsegments (field, file, offset, pos, shape); requests withoutpayload_dirkeep the existing inline path.compact_routed_expertsis shared between inline serialization and file writes.The orchestrator wipes
payload_rooton a fresh run, clips branch payloads intrace_to_samples(including final-token routing row repeat), accepts masks via payload segments inTrainSink, and deletesv<k>directories older thanmin_fresh_versionafter each policy update. Docs add a “Payloads by handle” section.Reviewed by Cursor Bugbot for commit a9d4c17. Bugbot is set up for automated code reviews on this repo. Configure here.