[RL] P2P RDT weight sync - #43375
Open
hao-aaron wants to merge 83 commits into
Open
Conversation
Contributor
|
Documentation preview: https://vllm--43375.org.readthedocs.build/en/43375/ |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces the Ray Direct Transport (RDT) weight transfer backend, which allows vLLM inference workers to pull weights directly from a trainer Ray actor using NIXL. The implementation includes the new RDTWeightTransferEngine, its registration in the factory, and updates to the Ray executor to support tensor transport. Additionally, two example scripts demonstrate standard and elastic RL scaling using this backend. Review feedback suggested several enhancements for robustness and performance, including enforcing the Ray executor backend requirement, safely handling private Ray attributes, improving error messages for missing producer methods, and implementing a fetch-ahead mechanism to overlap weight transfers with loading.
5 tasks
Signed-off-by: hao-aaron <ahao@anyscale.com>
Signed-off-by: hao-aaron <ahao@anyscale.com>
Signed-off-by: hao-aaron <ahao@anyscale.com>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: hao-aaron <ahao@anyscale.com>
Signed-off-by: hao-aaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
SumanthRH
reviewed
Aug 21, 2026
Contributor
There was a problem hiding this comment.
The number of buffers here is not documented properly:
It's 2 buffers per gather for the trainer ranks and then 2 buffers per consumer in the serve producer by default. Controlled by DEFAULT_GATHER_LOOKAHEAD and num_rdt_buffers
I would want a clear description of this in the docs, probably in limitations. We should also mention future plans to migrate to remote GPU -> CPU memory transfer since that was just added: ray-project/ray#64815
Contributor
There was a problem hiding this comment.
Let's add this after you add your optimization of using only one common buffer for all consumers
Contributor
|
/ci run |
|
❌ @SumanthRH, Only reviewers with write access can use CI commands before CI is delegated to the PR author. |
Signed-off-by: haoaaron <ahao@anyscale.com>
SumanthRH
reviewed
Aug 21, 2026
Comment on lines
+509
to
+523
| # ``ray.get_actor`` resolves by NAME, in a process that never | ||
| # imported the producer's class, and ``enable_tensor_transport`` | ||
| # lives in that class's metadata: the creating process passes | ||
| # ``meta.enable_tensor_transport``, and Ray infers it from the | ||
| # class's ``@ray.method(tensor_transport=...)`` decorators. A | ||
| # name-resolved handle therefore always reports False and the | ||
| # dispatch guard rejects the pull, even though the trainer set the | ||
| # option. Inherent to resolving by name across processes, not a | ||
| # version bug; it fails on Ray 2.56.0. | ||
| # | ||
| # Forcing it skips no validation. Ray's next guard, | ||
| # ``actor_has_tensor_transport``, asks the LIVE actor whether it can | ||
| # build a NIXL agent and still runs, and ``_spawn_server`` hardcodes | ||
| # the actor option, so a misconfigured producer cannot reach here. | ||
| actor._ray_enable_tensor_transport = True |
Contributor
There was a problem hiding this comment.
In summary:
This code gets the actor handle for the trainer by resolving by actor name;
actor = ray.get_actor(
chosen_name,
namespace=init_info.trainer_actor_namespace,
)However, the _ray_enable_tensor_transport attribute is not set on the actor. Seems like a limitation of how we resolve actors with RDT right now @Sparks0219
SumanthRH
reviewed
Aug 21, 2026
Comment on lines
+91
to
+120
| # Cross-deployment serve-slot sharing. | ||
| # | ||
| # Consumers whose ids differ by a multiple of ``workers_per_replica`` are the | ||
| # same worker of different inference deployments: identical parallel config, | ||
| # identical baked plan, identical chunk sequence, byte-identical pack layout. A | ||
| # serve ring per consumer therefore costs one full ring per deployment on the | ||
| # producer's GPU, and repeats the pack once per deployment for identical bytes. | ||
| # | ||
| # NIXL reads are one-sided and non-destructive, so R readers can read ONE | ||
| # registered slot concurrently. What a producer cannot observe is when a reader | ||
| # has FINISHED, so the release edge comes from the consumer's ISSUE order, which | ||
| # it sends as ``seq``: the pipeline drains pull i before issuing i+K, so slot | ||
| # ``seq % K`` was last packed for ``seq - K``, whose read is over. Under sharing | ||
| # the same holds, because generation ``seq`` is only packed once every live | ||
| # sharer has arrived at it, so each has drained ``seq - K``. | ||
| # | ||
| # Deriving the slot from a per-call counter on THIS side instead (execution | ||
| # order) is wrong: Ray may start a consumer's K concurrent produce calls in any | ||
| # order, so the call that executes K-before another can be a pull that is still | ||
| # being read, and its slot gets repacked underneath the reader. Silent, showing | ||
| # up only as a logprob drift. | ||
| # | ||
| # Hence one rendezvous per generation, keyed by ``seq``: the group's live sharers | ||
| # meet there, the LAST to arrive packs, and all of them return that one blob. | ||
| # | ||
| # Sharers that do NOT pull the same chunks would rendezvous on sequences whose | ||
| # bytes differ, so the plan is compared at init instead | ||
| # (``reserve_serve_buffer``'s ``plan_digest``), where it is one error naming both | ||
| # consumers before any byte moves. A sharer that dies mid-sync stalls its group, | ||
| # exactly as a dead consumer already stalls the per-group free barrier; the next |
Contributor
There was a problem hiding this comment.
Can we also add a note here on the limitation of cross deployment slot sharing here? We are introducing some synchronization here across replicas to make this word. Ideally in P2P we don't have any such synchronization, but this is a limitation given the current model of GPU -> GPU memory transfers with RDT with limited buffer sizes on GPU.
Let's add a TODO to explore GPU -> CPU memory transfer that could eliminate this
Upstream vllm-project/vllm main at 7a2fdba (325 commits since b652ded). Two conflicts, both in files the RDT work also touches: - .buildkite/test_areas/distributed.yaml: upstream relabelled every step to ":nvidia: (<gpu>) <name>". Kept the Sharded RDT step and renamed it to match. - docs/training/weight_transfer/README.md: auto-merged; upstream's new Rust-frontend paragraph listed the transports, so sharded-RDT was added to that list. Nothing the RDT code imports moved: of its upstream dependencies only config/parallel.py and v1/executor/multiproc_executor.py changed in the range, and world_size_across_dp / FutureWrapper / MultiprocExecutor / WorkerProc are all still there. Tests: PYTHONPATH=<fork> pytest tests/distributed/test_sharded_rdt_{producer,plan,trainer}.py --noconftest -> 202 passed, 10 skipped (the documented baseline). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #85072 for commit |
SumanthRH
reviewed
Aug 21, 2026
Comment on lines
+554
to
+558
| rings = self._serve_rings.setdefault(sg, [None] * self._nring) | ||
| for i in range(self._nring): | ||
| slot = rings[i] | ||
| if slot is None or slot.numel() < alloc: | ||
| rings[i] = self._new_serve_buffer(alloc) |
Contributor
There was a problem hiding this comment.
Wait, so we still have one serve ring per consumer per trainer rank?
Your current fix ensures that we don't scale with number of replicas, but it will still scale with the number of workers per replica?
Signed-off-by: haoaaron <ahao@anyscale.com>
Contributor
|
There are some CI failures @hao-aaron . Can you take a look? Was the test written for L4s? |
Contributor
|
Only the unit/ integration tests have run. The E2E examples are yet to run in CI. |
Contributor
|
Overall the changes here LGTM as a first step for a RDT engine pending CI fixes |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #85120 for commit |
Member
|
/ci run |
|
✅ Triggered Buildkite CI #85128 for commit |
…nto rdt-weight-sync
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #85157 for commit |
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
4 participants
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.
original design doc: https://docs.google.com/document/d/103ppQh3VxOR6Njzw0PQn6vGIaDFOAxTpiS9_Na68zN4/edit?tab=t.0
Summary
A fourth weight-transfer backend,
sharded_rdt. Every existing backend pusheswhole parameters to every inference worker. This one lets each worker pull
only the slice it actually consumes under tensor and expert parallelism, over
NIXL / Ray Direct Transport. A large MoE moves roughly
total_bytes / num_workersper worker instead of
total_bytes.The mechanism, in one paragraph: at init we run
model.load_weightsonce againstzero-storage placeholder tensors and record what each of vLLM's own weight
loaders does —
narrow,t,view, … — as a replayable op chain perdestination slice. Every later sync is pure replay: the worker asks a trainer rank
for exactly those slices, and they land directly in the
layerwise reload buffers.
No full HF tensor is ever materialized on a worker.
Measured on Qwen3-235B-A22B, 8 trainer GPUs → 8 inference GPUs (H100): the full
472 GB moves in a 3.3 s warm sync under DP8+EP8 at 43 GiB/s, vs 6.1 s at
26–28 GiB/s for the same model under TP8. GLM-4.5-Air holds the same wire rate;
Kimi K2 (1T, FP8) is validated from a raw sharded checkpoint. Full table below.
docs/training/weight_transfer/sharded_rdt.mdvllm/distributed/weight_transfer/base.py.../sharded_rdt_lazy.py.../sharded_rdt_common.py.../sharded_rdt_engine.py.../sharded_rdt_trainer.pytests/distributed/test_sharded_rdt_*.pyexamples/rl/*rdt*.pyHow it works
1. Slices are tracked through vLLM's own weight loaders
A weight loader normally receives a full HF tensor and slices out this worker's
part. We hand it a
LazyRDTTensorinstead: a_make_wrapper_subclasstensor thatanswers
.shape/.dtype/.size()but owns no storage. Every allowlistedview or shape op returns a new lazy with that op appended to a recorded chain,
and
copy_is the sink that terminates it.The chain is the wire format:
The trainer replays it as
getattr(tensor, op)(*args, **kwargs)against its liveparameter and sends the result.
Anything a loader does that needs real data — arithmetic,
.to(),.item(),.data, bool-mask indexing — falls off the allowlist, reaches__torch_dispatch__, and raises at init. That is the intended behavior:failing during setup beats silently transferring the wrong bytes.
SUPPORTED_OPSis a single table both sides derive from, so the recorder and the replayer cannot
drift apart.
2. Discovery happens once
Running the loaders is expensive, so it happens exactly once, at
init_transfer_engine, as a dry run with every parameter on meta. Nothingtransfers; we only record, per leaf module, which slice feeds which
as_strideddestination region. Every later sync is replay with noload_weights, no lazy dispatch, no discovery.A module only contributes its recording if it fully loaded (copied numel ≥
get_layer_size). A partially-recorded module would leave unwritten regions forfinalizeto initialize, so baking it would scatter garbage.There is no fallback path. A name whose
copy_fired but which produced norecording fails the plan build at init, naming the offending weights. Earlier
revisions carried a per-slice fallback load; it was deleted (see
Removed before review).
3. Received slices land directly in the layerwise reload buffers
The engine drives layerwise reload itself, in
start_weight_update/finish_weight_update. Because each destination is already recorded as anas_stridedregion, an arriving slice is copied straight into the layer beingreloaded. Each layer is quantized and copied into its persistent kernel storage as
soon as its last slice lands, on a dedicated stream, so quant overlaps the next
layer's transfer.
4. Gathers and pulls are pipelined
The trainer usually cannot serve its parameters as they sit: FSDP shards them, and
even an EP-split trainer has to assemble a whole expert. So each sync still runs
gather collectives — but a layer at a time.
A gather group is one decoder layer (the name list is cut on
model.layers.<N>.boundaries, leaving the embeddings and the final norm /lm_headas one group each). The trainer gathers a layer, publishes it —immediately pullable — and moves on while the consumers pull the previous one.
When every consumer signals it is done with a layer, the trainer drops it and
gains a credit to gather another.
gather_lookaheadbounds that: at mostgather_lookahead + 1layers are residenton the trainer. The default of 1 keeps the next layer gathered and serveable
while the current one is being pulled, which hides the handoff at the 2-layer
memory floor.
5. A trainer rank need not hold the whole model
Each rank declares what it holds via
WeightSource.held_names(); the fleetall-gathers those declarations at
trainer_initand transposes them into acompact table (the distinct owner sets, plus a per-name index into them). Consumers
route each pull to a rank that actually holds the name.
One declaration expresses every layout — pipeline stages, expert parallelism,
both at once, or a shape that fits neither. Consumers block-and-rotate across the
ranks that hold a name so no single trainer NIC becomes the bottleneck.
Public API
Two new entries in existing registries and one new init-info dataclass. Nothing
else is added to the public surface.
Inference side — a plain backend selector; everything else arrives from the
trainer at the init handshake:
Trainer side — identical in shape to the NCCL and IPC engines:
Requirements
distributed_executor_backend="ray"— workers must be Ray actorsnixlinstalled in the environment shared by trainer and workersSUPPORTED_OPSrecorded plan, and a silent mismatch would load weights into the wrong expert
slots
Shared code touched
Four files outside the new backend. All additive; no existing backend changes
behavior.
vllm/distributed/weight_transfer/base.py(+154) — three additions to theABCs, each with a working default so existing sources and engines are unaffected:
layerwise_groups(names)— the pre / per-decoder-layer / post partition. Itlives in
base.pyrather than in the backend because it defines what a groupindex means for any
WeightSource.WeightSource.held_names()→Collection[str] | None, defaulting toNone("this rank holds everything"), plus
groups()anditer_groups()derived fromit.
iter_groups()also exists for cost: driving materialization per groupinstead of per tensor turns ~37k generator resumes into ~95 on a per-expert MoE.
WeightTransferEngine.defers_processing(defaultFalse) anddrain_pending()(default no-op). This engine pipelines its GPUpost-processing onto background threads, so
update_weightscannot synchronizethe device; the flag tells a caller that has taken over the update tail that it
must drain first.
vllm/v1/executor/ray_executor_v2.py(+9) — Ray requires the calling actorto opt into tensor transport, so worker actors get
enable_tensor_transport=Truewhen and only when the backend issharded_rdt.vllm/config/weight_transfer.py(+1/−1) andfactory.py(+12) — thebackend literal and the two lazy factory registrations, matching the existing
pattern exactly.
tools/pre_commit/check_torch_cuda.py(+1) — allowlistssharded_rdt_engine.py, which needs realtorch.cudastreams and events for thereceive-slot handshake (same exemption
ipc_engine.pyalready has).Tests
142 tests across three new files, all CPU/meta — no GPU, no Ray, no NIXL. The
planning core is pure (baked plan + partition → static plan), which is what makes
that possible.
Three of these are load-bearing rather than routine:
TestPackedLayoutguards the one invariant no runtime check can catch. Bothsides compute the packed byte layout independently; if they disagree the bytes
still arrive exactly as sent and only the carving differs, so weights are
silently wrong. The test transcribes the producer's rule independently and
asserts the consumer's matches, on a mixed-dtype group whose sizes do not land
on a 16 B boundary.
TestSignalCompletenesspins that every gather group is signalled exactlyonce per sync, under every chunking shape. A missed signal wedges the trainer's
gather loop; a doubled one frees a layer while a pull is still reading it.
TestWeightSourceGroupContractpins the new ABC defaults against a sourcethat holds only part of the model.
Commands run
python -m pytest tests/distributed/test_sharded_rdt_plan.py \ tests/distributed/test_sharded_rdt_trainer.py \ tests/distributed/test_sharded_rdt_producer.py \ tests/distributed/test_weight_transfer.py -q # 223 passed, 26 skipped in 73s pre-commit run --all-files # all hooks pass (ruff-check, ruff-format, typos, markdownlint-cli2, # mypy-3.10, check-spdx-header, check_torch_cuda)Of the 26 skips, 7 belong to this PR — the trainer-side tests that need a real
device (CUDA-IPC export in the gather loop, serve-ring packing). They pass on a
GPU host. The other 19 are pre-existing GPU skips in
test_weight_transfer.pyforthe NCCL, IPC and sparse backends.
Worth noting what does not skip: all 77 consumer-plan tests and all 41
producer-protocol tests run on CPU, because the planning core is pure and the
serve actor is exercised through a local (non-Ray) instance.
End-to-end validation
Both examples were run on 2 nodes × 8×H100:
WeightSource, no HF materializationModel evaluation: this is a weight transport, not a modeling change — the
bytes that land are the trainer's parameters, and correctness is "the inference
engine produces the trainer's model." That is verified end to end in the examples
(coherent generations after sync vs. garbage from dummy weights before it) and at
the byte level by
TestPackedLayout. No accuracy-affecting kernel or model codeis touched, so
tests/evals/was not run.Performance
8×H100 trainer → 8×H100 inference, measured at the engine level:
The two Qwen3-235B columns move the same 472 GB — the difference is entirely in
how it is cut. Under DP+EP each consumer holds full experts for exactly one
trainer EP coordinate, so it pulls ~2 slices per layer (one expert plus one
replicated) instead of a TP slice of every coordinate. Four times fewer, larger
pulls is what takes the wire rate from 26–28 GiB/s to 43: the transfer stops being
per-slice latency-bound and starts being bandwidth-bound, which is the whole point
of the packed-chunk design.
GLM-4.5-Air is the same shape at smaller scale and lands at the same wire rate,
which is the useful signal — 42 GiB/s is the fabric, not a model-specific result.
The bill on 80 GB cards is real and worth stating: weights are 67.8 GiB/rank under
DP+EP (vs ~59 under TP8, since attention and embeddings are unsliced), the DP+EP
fused-MoE workspace scales with
max_num_batched_tokens × dp_size, and the 2.5GiB of receive arenas need
enforce_eagerto fit. Treat DP+EP at this scale as aneager-mode configuration for now.
Known limitations
urgent follow ups:
less urgent:
gpu_memory_utilization. Like NCCL and NIXLinternals they are not counted, so a fraction that leaves no headroom OOMs at
the first sync even though the engine came up healthy. Documented with a warning
box;
arena_presize_gbexists to size them deterministically.first-class: a currently-loading hook instead of monkeypatched loader stamps, a
dry-run mode instead of bypassing
online_process_loader, and anabort_layerwise_reloadinstead of a hand-rolled restore. Marked in_bake's docstring; a good follow-up.shared dicts are correct only because the free barrier orders access. It has not
bitten, but a change to free timing or actor concurrency could expose it.
_all_gather_ownedis a full all-gather where a gather-to-rank-0 would do,except that every rank retains the table so a restarted consumer can rejoin via
get_worker_init_payload()without a collective. ~12 KB/rank at 100k names, sonot urgent.
Documentation
docs/training/weight_transfer/sharded_rdt.md(new) — the backend page,matching the structure of
nccl.mdandipc.md.docs/training/weight_transfer/base.md(+114) — theWeightSourcesectionis reframed as the adapter for whatever shape your trainer's weights are in,
and now documents
held_names()and gather groups. The example is aMegatron-Bridge source, since that is the shape framework authors actually start
from.
README.md— the backend table.AI assistance
This change was developed with AI assistance (Claude). The design decisions,
the measured results, the multi-node validation runs, and every line of the diff
were reviewed by the submitting human, who can defend the change end to end.