Skip to content

feat: write router-replay ids and sampling masks by handle on multi-node runs - #3820

Closed
faresobeid wants to merge 4 commits into
feat/payload-segmentsfrom
feat/payload-by-handle
Closed

faresobeid wants to merge 4 commits into
feat/payload-segmentsfrom
feat/payload-by-handle

Conversation

@faresobeid

@faresobeid faresobeid commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Reworked onto vLLM 0.31 (stacked on #3819 → #3888)

Branch feat/payload-by-handle now holds one commit a4a98c5f2 on top of #3819's 1bcca4115. This PR is closed, so GitHub still shows the old head.

  • 0.31 port: the inline path keys on engine_client.vllm_config.aux_output_config.enable_return_routed_experts (0.31 moved the flag off ModelConfig). The writer is unchanged: 0.31's engine output still carries routed_experts (concatenated AuxOutput chunks from routed_experts_prompt_start) and sampling_mask.token_ids.
  • Discarded rollouts no longer keep their files. TrainSink.process_group deletes 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 by constant_trainer_batch_size. _drop_stale deletes 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.
  • Still inline under disaggregated P/D (needs router#60 and a router pin bump) and with score centering (rejected when payload_root is set explicitly).
  • Checks: pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/inference: 307 passed; ruff clean.
  • GPU (2026-10-06, vLLM 0.31.0, Qwen3-30B-A3B TP1, BI, 16 requests of 500-8000 tokens, 256 tokens, conc 8, same seeds with and without 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.
  • Before merging, move the deps/renderers (a79a10f) and deps/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) and deps/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.

  • Switch (one config field): orchestrator.payload_root. rl sets it to <run dir>/payloads (absolute) when router replay or sampling-mask capture is on, the deployment is multi_node, and inference is not disaggregated. Inference needs no config: the directory travels with each request.
  • Request: the dispatcher gives each live-sourced train rollout payload_dir = <payload_root>/v<k>, where k is the rollout's own dispatch policy version (self.policy.version when it calls env.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._sampling puts it in extra_body.extra_args; verifiers' wire_args() flattens extra_body into the generate sampling_params, and renderers forwards those verbatim, so verifiers and renderers need no request-side change. Checked ad hoc: vLLM's GenerateRequest parses it into sampling_params.extra_args["payload_dir"], merged with any user extra_args.
  • Writer (serving_tokens.py): a request with extra_args.payload_dir gets _PayloadCapture, which takes routed_experts / sampling_mask off the engine output (upstream encodes nothing inline), writes both into one file <payload_dir>/<uuid>.bin in a thread, closes it before responding (NFS close-to-open), and returns payload segments on the choice. Routing row i is position routed_experts_prompt_start + i, as uint8/uint16 (compact_routed_experts, split out of serialize_routed_experts). Mask row i is completion token i at prompt_len + i, int32 [tokens, widest row] padded with -1. Requests without payload_dir (eval, frozen-source rollouts, standalone clients) keep today's inline path unchanged.
  • Orchestrator (trajectories.py): trace_to_samples takes branch.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, like Branch.routed_experts returning None.
  • GC (Orchestrator.on_policy_update): after inference applies v{step}, delete v<k> for k < min_fresh_version(step + 1, max_off_policy_steps). Why that bound is exact:
    • The trainer reads batch s and then publishes v{s} (dataloader.get_batch() comes before weight_sender.broadcast(model, step=progress.step) in train.py), and the orchestrator sender and trainer receiver advance one batch per step from the same index. So when v{step} is applied, every batch up to step has been read.
    • Batch b ≥ step + 1 ships only rollouts whose group started at min_fresh_version(b) or later (TrainSink._drop_stale, the hard guarantee), and a rollout's k is at least its group's start.
    • Fresh run: the root is wiped. Resumed run: the root is kept, because a trainer that outlived the orchestrator may still read old batches; the startup sync calls the same hook with v{resume_step}. A directory that a still-running stale rollout recreates after GC is caught by the next collection.
  • Docs: docs/inference.md § Payloads by handle.

Not supported with disaggregated P/D yet: rl never enables it there, because the PD router merges only inline routing (PrimeIntellect-ai/router#60). Setting orchestrator.payload_root by hand on P/D loses prefill routing, which trace_to_samples reports 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 awaiting post_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 check and ruff format --check pass.
  • Ad hoc: GC with max_off_policy_steps=2, v0..v11 present and step=10 keeps v8..v11. _sample_payload repeats the last routing row for the final token (reader output equals the inline construction) and drops routing with a gap.
  • Config: examples/advanced/qwen3-30b-a3b/swe.toml (multi-node) + --trainer.enable-router-replay sets orchestrator.payload_root to <run dir>/payloads. Without router replay it stays None. Single-node reverse-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).

Landing order

  1. feat(client): pass through by-handle payload segments renderers#171 and feat(v1): by-handle payload segments on traces verifiers#2765
  2. feat: by-handle payload segments: reader, packing, CP-window reads, GC #3819 (reader/packing, no behavior change)
  3. feat: write router-replay ids and sampling masks by handle on multi-node runs #3820 (writer, switch and GC; re-pin deps/renderers / deps/verifiers to the merged commits first)
  4. feat(pd): stitch by-handle routed_experts payload segments router#60 (independent), then let rl turn by-handle on for disaggregated P/D

Follow-ups


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. The rl entrypoint auto-sets it to <run dir>/payloads when router replay or sampling-mask capture is on, deployment is multi_node, and inference is not disaggregated P/D; hand-setting it with disaggregated P/D or trainer.loss.score_centering is rejected.

Train rollouts get extra_args.payload_dir = <root>/v<policy version> through the dispatcher and env sampling config. vLLM’s PrimeRlServingTokens writes routed_experts and sampling_mask into per-choice .bin files and returns payload segments (field, file, offset, pos, shape); requests without payload_dir keep the existing inline path. compact_routed_experts is shared between inline serialization and file writes.

The orchestrator wipes payload_root on a fresh run, clips branch payloads in trace_to_samples (including final-token routing row repeat), accepts masks via payload segments in TrainSink, and deletes v<k> directories older than min_fresh_version after 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.

@cursor cursor Bot left a comment •

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

@faresobeid
faresobeid force-pushed the feat/payload-by-handle branch from b430fc3 to 40bb9e0 Compare October 3, 2026 23:28

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Fix All in Cursor

❌ 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.

@faresobeid
faresobeid force-pushed the feat/payload-by-handle branch from f7b0763 to 19ec1ae Compare October 4, 2026 06:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants