[RL] P2P RDT weight sync (#43375)
Aaron Hao committed
Aug 23, 2026 at 09:42 UTC
6cd9713463f98442a60f21f9388d16884144f62f
24 files changed
+8409
-29
.buildkite/test_areas/distributed.yaml
+25
-1
@@ -17,6 +17,21 @@ steps:
17
- pytest -v -s distributed/test_shm_buffer.py
18
- pytest -v -s distributed/test_shm_storage.py
19
20
+- label: ":nvidia: (L4) Sharded RDT Weight Transfer"
21
+ key: sharded-rdt-weight-transfer
22
+ timeout_in_minutes: 20
23
+ working_dir: "/vllm-workspace/tests"
24
+ num_devices: 1
25
+ source_file_dependencies:
26
+ - vllm/distributed/weight_transfer/
27
+ - tests/distributed/test_sharded_rdt_plan.py
28
+ - tests/distributed/test_sharded_rdt_producer.py
29
+ - tests/distributed/test_sharded_rdt_trainer.py
30
+ commands:
31
+ - pytest -v -s distributed/test_sharded_rdt_plan.py
32
+ - pytest -v -s distributed/test_sharded_rdt_producer.py
33
+ - pytest -v -s distributed/test_sharded_rdt_trainer.py
34
+
35
- label: ":nvidia: (L4) Distributed DP Basic"
36
key: distributed-dp-tests-2-gpus
37
timeout_in_minutes: 35
@@ -107,7 +122,7 @@ steps:
122
123
- label: ":nvidia: (L4) Distributed Torchrun + Examples"
124
key: distributed-torchrun-examples-4-gpus
110
- timeout_in_minutes: 30
125
+ timeout_in_minutes: 45
126
working_dir: "/vllm-workspace"
127
num_devices: 4
128
source_file_dependencies:
@@ -137,6 +152,15 @@ steps:
152
# rlhf examples (each launches its own `vllm serve` and tears it down)
153
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
154
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
155
+ # sharded RDT: 2 FSDP2 trainer ranks -> 2 vLLM DP ranks with EP. Its data
156
+ # plane is NIXL, which ships only in the kv-connector set, and the script
157
+ # also drops the nixl-cu wheel variant that does not match this image. Kept
158
+ # last so the commands above run in an unmodified environment. RDT_MODEL swaps
159
+ # the script's reference MoE for a tiny one: no weights are downloaded either
160
+ # way, but every trainer rank builds the model before sharding it, and the
161
+ # default does not fit an L4.
162
+ - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
163
+ - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
164
165
- label: ":nvidia: (L4) Distributed DP Extended"
166
key: distributed-dp-tests-4-gpus
docs/training/weight_transfer/README.md
+3
-2
@@ -34,6 +34,7 @@ engine drives on your behalf:
34
| [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
35
| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
36
| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
37
+| [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) |
38
39
## Quickstart
40
@@ -48,7 +49,7 @@ from vllm.config import WeightTransferConfig
49
50
llm = LLM(
51
model="my-model",
51
- weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl"
52
+ weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl", "sharded_rdt"
53
)
54
```
55
@@ -158,7 +159,7 @@ When running vLLM as an HTTP server, the following endpoints are available for w
159
!!! note
160
The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
161
161
-The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, or sparse-NCCL transport.
162
+The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, sparse-NCCL, or sharded-RDT transport.
163
164
## Extending the System
165
docs/training/weight_transfer/base.md
+136
-15
@@ -20,8 +20,14 @@ vice versa, so the registries stay independent.
20
21
### WeightSource
22
23
-A `WeightSource` is a **re-iterable** source of the trainer's weights with two
24
-channels:
23
+**This is the adapter for whatever shape your trainer's weights are in.**
24
+`WeightSource` defines how you extract your weights for your specific framework.
25
+
26
+What an engine gets is always the same: HF-format parameter names, and tensors
27
+already materialized to their full (unsharded) shape. Whatever gathering,
28
+re-fusing, dequantizing or renaming that takes belongs inside the source.
29
+
30
+A `WeightSource` is **re-iterable** and has two required channels:
31
32
- **`metadata() -> list[ParamMeta]`** — the name, wire dtype, and full shape of
33
every parameter, *without transferring anything*. Cheap when shapes are known
@@ -51,9 +57,7 @@ class ParamMeta:
57
Materializing is typically a collective, so **every trainer rank must iterate the
58
same source in the same order, in lockstep**, or ranks deadlock. `metadata()` can
59
itself be a collective for custom producers, so it too runs on every rank — only
54
-the sender ships the result. Under pipeline parallelism a rank may not own a
55
-parameter at all; iterating still drives the collective, and the yielded tensor
56
-is only meaningful on the sender.
60
+the sender ships the result.
61
62
`iter(source)` must yield a *fresh* pass each round.
63
@@ -70,35 +74,133 @@ from vllm.distributed.weight_transfer import ModuleSource
74
source = ModuleSource(model)
75
```
76
73
-#### Custom Sources
77
+#### Custom sources
78
75
-Subclass `WeightSource` when the weights you want to send require additional processing to convert to a HF compatible format.
79
+Subclass `WeightSource` when the weights need work to reach HF format — a
80
+framework-specific export, a re-fusing step, a dtype cast.
81
82
```python
83
from vllm.distributed.weight_transfer import ParamMeta, WeightSource
79
-from vllm.distributed.weight_transfer.base import materialize_full_tensor
84
85
82
-class MyExportSource(WeightSource):
83
- def __init__(self, model):
84
- self._model = model
86
+class MegatronBridgeSource(WeightSource):
87
+ """Megatron model -> HF names, via a bridge that gathers TP/PP/EP internally
88
+ and returns full tensors on every rank."""
89
+
90
+ def __init__(self, bridge, module, dtype):
91
+ self._bridge, self._module, self._dtype = bridge, module, dtype
92
self._meta: list[ParamMeta] | None = None
93
94
+ def _export(self):
95
+ return self._bridge.export_hf_weights(self._module)
96
+
97
def metadata(self) -> list[ParamMeta]:
98
# Cache: for producers that must materialize to learn shapes, this is
99
# the expensive channel. Runs on every rank (it may be a collective).
100
if self._meta is None:
101
self._meta = [
92
- ParamMeta(name, t.dtype, tuple(t.shape)) for name, t in self._export()
102
+ ParamMeta(name, self._dtype, tuple(t.shape))
103
+ for name, t in self._export()
104
]
105
return self._meta
106
107
def __iter__(self):
108
# Must yield exactly what metadata() declared, in the same order.
109
for name, tensor in self._export():
99
- yield name, materialize_full_tensor(tensor)
110
+ yield name, tensor.to(self._dtype).detach().contiguous()
111
+```
112
+
113
+#### `held_names()`: partial ownership
114
+
115
+By default every rank is assumed to be able to produce every parameter — which is
116
+what the source above does, since its bridge gathers across all parallelism before
117
+yielding. That is simple and always correct, but it means the gather cost is paid
118
+in full on every rank.
119
+
120
+Optionally, you can override `held_names()` when the ranks are split so each holds only part of the
121
+model. It returns the
122
+parameter names this rank holds (or `None`, the default, for all of them):
123
+
124
+```python
125
+def held_names(self):
126
+ # This pipeline stage's layers, and within them only this EP rank's experts.
127
+ return self._my_stage_names - self._foreign_expert_names
128
```
129
130
+This covers various trainer layouts — pipeline stages (a rank holds some layers),
131
+expert parallelism (a rank holds some experts), both at once, or a shape that fits
132
+neither. Backends that can route per parameter (see
133
+[sharded RDT](sharded_rdt.md)) then pull each name from a rank that actually
134
+holds it.
135
+
136
+Three requirements come with overriding it:
137
+
138
+- **`metadata()` must still describe the whole model on every rank.** Only the
139
+ sender's metadata reaches the inference side, so a rank that reported just its
140
+ own share would leave the rest silently un-transferred. Sharded RDT cross-checks
141
+ this across ranks at init.
142
+- **Every name must be held by at least one rank**, or it can never be served.
143
+ The engine raises at init naming the first orphan.
144
+- **Iteration yields `None` for a name this rank does not hold.** The name still
145
+ appears, in metadata order, so the order check stays aligned across ranks — only
146
+ the data is absent. Claiming a name and then yielding `None` for it is an error
147
+ the engine reports by name.
148
+
149
+!!! warning "Partial ownership only works with sharded RDT"
150
+ Only a backend that routes per parameter honours `held_names()`.
151
+ Broadcast backends ignore it and send every name from every rank, so
152
+ declaring partial ownership there changes nothing.
153
+
154
+#### Gather groups
155
+
156
+Some backends transfer a layer at a time rather than a model at a time, so they
157
+partition `metadata()` into **gather groups**. `layerwise_groups` keys each name
158
+on the **outermost index segment** it contains, so **a group is one decoder
159
+layer**, with runs of un-indexed names (the embeddings, the final norm,
160
+`lm_head`) forming groups of their own where they appear:
161
+
162
+```text
163
+group 0 model.embed_tokens.weight
164
+group 1 model.layers.0.* <- one decoder layer
165
+group 2 model.layers.1.*
166
+...
167
+group N+1 model.norm.weight, lm_head.weight
168
+```
169
+
170
+Keying on the index rather than a literal prefix means no per-architecture table:
171
+`model.layers.0.`, `model.language_model.layers.0.` (recent Qwen text
172
+checkpoints), `transformer.h.0.` (GPT-2, Falcon), `backbone.layers.0.` (Mamba)
173
+and a vision tower's `visual.blocks.0.` all partition the same way. The index
174
+taken is the *outermost* one, which keeps a MoE layer whole — a per-expert name
175
+like `model.layers.3.mlp.experts.7.w1` keys on the layer, not the expert.
176
+
177
+Group index *g* means the same layer on every rank and every consumer, because
178
+every side derives it from one rank's `metadata()`. That agreement is what lets a
179
+backend bound its buffers to one layer and free a layer once everyone is done
180
+with it.
181
+
182
+!!! warning "A leaf module's sources must all fall in one group"
183
+ The sharded-RDT engine frees a group as soon as its last chunk lands, so a
184
+ module split across groups would park a pull until the stall watchdog fires.
185
+ The default partition guarantees this; a `groups()` override must preserve
186
+ it.
187
+
188
+Two hooks follow from this, both with working defaults:
189
+
190
+- **`groups()`** — this rank's groups, in metadata order. The default is
191
+ `layerwise_groups(metadata())` restricted to the groups holding at least one
192
+ held name. A group with nothing held here is skipped entirely.
193
+- **`iter_groups()`** — the same stream, batched one group at a time. The default
194
+ drives `__iter__` and batches its output, checking that names arrive in metadata
195
+ order. Override it when your framework can produce a whole group in one step:
196
+ materializing is usually a collective, and driving it per group rather than per
197
+ tensor turns ~37k generator resumes into ~95 on a per-expert MoE model.
198
+
199
+Because `metadata()` order defines the partition, **all names sharing a layer
200
+index must be contiguous in it**. A source whose natural export order interleaves
201
+layers — bucketing all the MoE experts together, say — has to reorder before
202
+returning.
203
+
204
### VLLMWeightSyncClient
205
206
**This is the adapter for however your RL stack reaches vLLM.** Many RL frameworks wrap
@@ -372,7 +474,7 @@ Subclasses must implement five methods:
474
The base class provides:
475
476
1. `__init__`, taking `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`), and `model` (`nn.Module`).
375
-2. `update_weights(update_info_dict)`, a thin wrapper for `receive_weights`: it parses the dict into the typed dataclass, calls `receive_weights`, and synchronizes the device.
477
+2. `update_weights(update_info_dict)`, a thin wrapper for `receive_weights`: it parses the dict into the typed dataclass, calls `receive_weights`, and synchronizes the device — unless the engine sets `defers_processing`, below.
478
3. `parse_init_info` / `parse_update_info`, which convert API-level dicts into the typed dataclasses and raise `ValueError` on a bad payload.
479
4. `set_weight_update_target` / `reset_weight_update_target`, used to retarget an update at the speculative draft model.
480
@@ -383,6 +485,25 @@ The base class provides:
485
update info carries only per-round metadata. This is what makes a
486
trainer/worker mismatch unrepresentable.
487
488
+!!! note "`defers_processing`: when a returned update means queued, not applied"
489
+ An engine that pipelines its GPU post-processing onto background threads
490
+ cannot let `update_weights` synchronize the device — that would block on
491
+ those threads and serialize the pipeline. Such an engine sets the class
492
+ attribute `defers_processing = True`, omits the per-update sync, and
493
+ guarantees completion in `finish_weight_update` instead.
494
+
495
+ Callers that go through `finish_weight_update` need do nothing; the engine
496
+ drains there. A caller that drives the tail itself — running its own
497
+ `finalize_layerwise_reload`, say — must check the flag and call
498
+ `drain_pending()` first, because with it set a returned `update_weights`
499
+ means *queued*, not *applied*. `drain_pending()` is idempotent, and a no-op
500
+ on an engine that processes synchronously, so it is always safe to call.
501
+
502
+ [Sharded RDT](sharded_rdt.md) is the built-in engine that sets it: it
503
+ scatters and quantizes on background threads with their own CUDA streams, so
504
+ its `drain_pending()` joins both queues and syncs both streams before
505
+ `finalize_layerwise_reload` runs.
506
+
507
### Request Classes
508
509
The API-level request classes provide backend-agnostic serialization using plain dictionaries.
@@ -497,7 +618,7 @@ Once registered, users select your backend via `WeightTransferConfig(backend="my
618
619
### WeightTransferEngineFactory
620
500
-The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
621
+The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, `sparse_nccl` and `sharded_rdt`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
622
623
```python
624
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
docs/training/weight_transfer/sharded_rdt.md
new
+126
@@ -0,0 +1,126 @@
1
+# Sharded RDT Engine
2
+
3
+The sharded RDT weight transfer engine moves weights point-to-point over [NIXL](https://github.com/ai-dynamo/nixl) using Ray Direct Transport (RDT), Ray's zero-copy tensor transport between actors. It is **pull-based**: the inference workers initiate every transfer and each one asks only for the *slice* it consumes under tensor and expert parallelism. A large MoE model therefore moves roughly `total_bytes / num_workers` per worker instead of the `total_bytes` a broadcast costs.
4
+
5
+## When to Use Sharded RDT
6
+
7
+- Very large models where broadcasting whole parameters is the bottleneck — typically an MoE served with expert parallelism, where each worker owns a small fraction of the experts
8
+- Trainer and inference on **separate GPUs**, over a fabric NIXL supports (InfiniBand, RoCE, EFA)
9
+- Trainers that are themselves sharded, including **pipeline-parallel** ones where a rank holds only part of the model
10
+
11
+Requirements:
12
+
13
+- `distributed_executor_backend="ray"` — the workers must be Ray actors
14
+- Ray `>= 2.56.0` on both the trainer and the workers
15
+- `nixl` installed in the environment shared by trainer and workers
16
+- Weight loaders that stay inside the supported op set (below)
17
+- EPLB (`enable_eplb=true`) is rejected: it rearranges experts at runtime, which invalidates the recorded plan
18
+
19
+## How It Works
20
+
21
+### Slices are tracked through vLLM's own weight loaders
22
+
23
+A weight loader normally receives a full HF-format tensor and slices out the part this worker needs. Instead of sending it one, the engine hands it a `FakeRDTTensor`: a zero-storage tensor that answers `.shape` / `.dtype` / `.size()` but holds no data. Every view or slice op the loader calls returns a new fake with that op appended to a recorded chain, and `copy_` is the sink that ends it.
24
+
25
+That chain is the wire format. `("model.layers.0.w", (("narrow", (0, 512, 512), ()), ("t", (), ())))` tells the trainer: take this tensor, `narrow` it, transpose it, send the result. The trainer replays it with `getattr(tensor, op)(*args, **kwargs)`.
26
+
27
+Discovery is expensive, so it happens **once**, at `init_transfer_engine`, as a dry run over `model.load_weights` with every parameter on meta. Nothing is transferred; the engine just records, per leaf module, which slice feeds which destination region. Every later sync is pure replay.
28
+
29
+Anything a loader does that needs real data — arithmetic, `.to()`, `.float()`, `.item()`, `.data`, bool-mask indexing — falls outside the allowlist and raises at init. That is deliberate: failing loudly during setup beats silently transferring the wrong bytes. `SUPPORTED_OPS` in `sharded_rdt_common.py` is the single table both sides derive from, so the recorder and the replayer cannot drift.
30
+
31
+### Received slices land directly in the layerwise reload buffers
32
+
33
+The engine drives [layerwise reload](../layerwise.md) itself, in `start_weight_update` / `finish_weight_update`. Because the dry run already recorded each destination as an `as_strided` region of its parameter, an arriving slice is copied straight into the layer being reloaded — no full HF tensor is ever materialized on the worker, and no second pass over `load_weights` runs. Each layer is quantized and copied into its persistent kernel storage as soon as its last slice lands.
34
+
35
+### Gathers and pulls are pipelined, which is what `gather_lookahead` bounds
36
+
37
+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, not a model at a time.
38
+
39
+**A gather group is one decoder layer.** The parameter list is keyed on the outermost index segment of each name, which leaves runs of un-indexed names — the embeddings before the first layer, the final norm and `lm_head` after the last — as groups of their own:
40
+
41
+```text
42
+group 0 model.embed_tokens.weight
43
+group 1 model.layers.0.* <- one decoder layer
44
+group 2 model.layers.1.*
45
+...
46
+group N+1 model.norm.weight, lm_head.weight
47
+```
48
+
49
+The layer is the unit of everything that follows: the trainer gathers a layer, publishes it (immediately pullable), and moves on to the next while the consumers pull the one it just published. Once every consumer has signalled that it is done with a layer, the trainer drops it and gains a credit to gather another. `gather_lookahead` is how far ahead of the consumers that loop may run, so at most `gather_lookahead + 1` layers are resident on the trainer at a time. The default of 1 keeps the next layer gathered and pullable while the current one is being pulled — enough to hide the handoff without doubling trainer memory. Raise it only if one layer's gather is slower than its pulls.
50
+
51
+Because a layer is also the unit the consumers free and the unit the receive buffers are sized against, it is what keeps memory bounded on both sides: without it the whole model would be one transfer, and both sides would have to hold their full share of it at once. Keying on the index rather than a fixed `model.layers.` prefix is what makes that hold across naming conventions — a VLM's `model.language_model.layers.`, GPT-2's `transformer.h.`, a vision tower's `visual.blocks.`. Sources can control the partition; see [gather groups](base.md#gather-groups).
52
+
53
+### Ownership
54
+
55
+A trainer rank need not hold the whole model. Each one declares what it holds through [`WeightSource.held_names()`](base.md#held_names-partial-ownership), the fleet all-gathers those declarations at `trainer_init`, and the consumers route each pull to a rank that actually holds the name. Pipeline stages, expert parallelism, and combinations of the two are all the same declaration. Consumers spread their pulls across the ranks that hold a name, so no single trainer NIC becomes the bottleneck.
56
+
57
+## Inference Side
58
+
59
+```python
60
+from vllm import LLM
61
+from vllm.config import WeightTransferConfig
62
+
63
+llm = LLM(
64
+ model="my-model",
65
+ weight_transfer_config=WeightTransferConfig(backend="sharded_rdt"),
66
+ distributed_executor_backend="ray",
67
+)
68
+```
69
+
70
+```bash
71
+vllm serve my-model \
72
+ --distributed-executor-backend ray \
73
+ --weight-transfer-config '{"backend": "sharded_rdt"}'
74
+```
75
+
76
+Everything else — which producers exist, how the model splits into layer groups, the ownership table — arrives from the trainer at the init handshake.
77
+
78
+!!! warning "Size the receive buffers before choosing `gpu_memory_utilization`"
79
+ Each worker holds `num_rdt_buffers` receive buffers, each large enough for the biggest single slice batch it pulls. Like NCCL and NIXL internals, they do **not** count against `gpu_memory_utilization`, so a fraction that leaves no headroom OOMs at the first sync even though the engine came up healthy. The buffer size is driven by the largest atomic slice — for an untied vocab matrix on a worker that holds it unsliced, that is the whole embedding.
80
+
81
+## Trainer Side
82
+
83
+```python
84
+from vllm.distributed.weight_transfer import (
85
+ ModuleSource,
86
+ HTTPVLLMWeightSyncClient,
87
+ WeightTransferTrainerFactory,
88
+)
89
+from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
90
+ ShardedRDTTrainerInitInfo,
91
+)
92
+
93
+engine = WeightTransferTrainerFactory.trainer_init(
94
+ init_info=ShardedRDTTrainerInitInfo(
95
+ rank=rank, # rank 0 is the sender
96
+ num_consumers=8, # inference workers, fleet-wide
97
+ trainer_actor_namespace="my_namespace", # must be visible to the workers
98
+ ),
99
+ client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
100
+ source=ModuleSource(model),
101
+)
102
+
103
+engine.send_weights() # once per sync, on every trainer rank
104
+```
105
+
106
+`trainer_init` and `send_weights` run on **every** trainer rank: each one owns a serve actor and takes part in the gathers, while only rank 0 drives the inference-side handshake. Any [`VLLMWeightSyncClient`](base.md#vllmweightsyncclient) works.
107
+
108
+To adapt a trainer that is not a plain `nn.Module` — a Megatron export, a raw sharded checkpoint — subclass [`WeightSource`](base.md#weightsource).
109
+
110
+### `ShardedRDTTrainerInitInfo`
111
+
112
+| Field | Default | Description |
113
+| ----- | ------- | ----------- |
114
+| `rank` | — | Keyword-only. This trainer rank; **0 is the sender** |
115
+| `num_consumers` | — | Inference workers across the whole fleet (TP × DP) |
116
+| `trainer_actor_namespace` | `None` | Ray namespace for the serve actors; the workers resolve them by name here |
117
+| `num_rdt_buffers` | 2 | Ring depth on both sides |
118
+| `buffer_presize_gb` | 0.0 | Pre-size each buffer slot, in GiB. Set it to cover the largest atomic slice |
119
+| `gather_lookahead` | 1 | Gathered-but-unfreed layers the gather loop runs ahead by |
120
+| `stall_timeout_s` | 300.0 | Seconds without progress before the sync fails. A liveness backstop for a consumer that dies mid-sync, not a latency target |
121
+
122
+## Examples
123
+
124
+- [Small MoE on 4 GPUs](../../../examples/rl/rlhf_sharded_rdt_small_ep.py) — 2 FSDP2 trainer ranks → 2 vLLM DP ranks with expert parallelism, one node. It pairs a trainer fleet with a separate inference fleet, the only arrangement this backend supports, and asserts that the sync moved the weights and that a second sync leaves generation unchanged — so it runs unattended in CI
125
+
126
+It keeps the trainer deliberately small — just enough FSDP2 to make the weights real — so the file stays about the weight sync rather than about the trainer. For a full RL trainer, SkyRL integrates this backend with Megatron (PP-local gathering and expert-stack fusion for MoE) alongside FSDP: [NovaSky-AI/SkyRL#1753](https://github.com/NovaSky-AI/SkyRL/pull/1753).
examples/rl/rdt_vllm_serve.py
new
+152
@@ -0,0 +1,152 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Driver-side helpers to run the inference fleet as a standard ``vllm serve``
4
+process for the sharded-RDT examples.
5
+
6
+The trainer reaches the server's weight-sync control plane over the RLHF HTTP
7
+routes (``HTTPVLLMWeightSyncClient``); generation uses ``/v1/completions``. RDT
8
+still needs the workers to be RayExecutorV2 tensor-transport actors sharing the
9
+trainer's Ray cluster, so the server is launched with the Ray v2 executor and
10
+inherits this process's env (``address=auto``). These helpers run only on the
11
+driver — they are not shipped to or imported by the Ray actors.
12
+"""
13
+
14
+import json
15
+import os
16
+import shutil
17
+import subprocess
18
+import sys
19
+import time
20
+
21
+import ray
22
+import requests
23
+
24
+
25
+def launch_vllm_serve(
26
+ model: str,
27
+ *,
28
+ tensor_parallel_size: int = 1,
29
+ data_parallel_size: int = 1,
30
+ enable_expert_parallel: bool = False,
31
+ port: int = 8000,
32
+ gpu_memory_utilization: float = 0.7,
33
+ extra_args: list[str] | None = None,
34
+) -> subprocess.Popen:
35
+ """Start ``vllm serve`` for the RDT inference fleet.
36
+
37
+ ``VLLM_SERVER_DEV_MODE`` exposes the RLHF weight-sync routes; the Ray v2
38
+ executor makes the workers tensor-transport actors (RDT's data plane). The
39
+ child inherits our env, so it joins the same Ray cluster (RAY_ADDRESS below)
40
+ and picks up the editable vLLM install via the venv interpreter.
41
+ """
42
+ vllm_bin = os.path.join(os.path.dirname(sys.executable), "vllm")
43
+ if not os.path.exists(vllm_bin):
44
+ vllm_bin = shutil.which("vllm") or "vllm"
45
+ cmd = [
46
+ vllm_bin,
47
+ "serve",
48
+ model,
49
+ "--port",
50
+ str(port),
51
+ "--enforce-eager",
52
+ "--load-format",
53
+ "dummy",
54
+ "--gpu-memory-utilization",
55
+ str(gpu_memory_utilization),
56
+ "--distributed-executor-backend",
57
+ "ray",
58
+ "--weight-transfer-config",
59
+ json.dumps({"backend": "sharded_rdt"}),
60
+ ]
61
+ if tensor_parallel_size > 1:
62
+ cmd += ["--tensor-parallel-size", str(tensor_parallel_size)]
63
+ # Only engage the DP-ray backend when actually data-parallel: DP=1 with
64
+ # data_parallel_backend=ray forces vLLM's DP-placement path, which fails.
65
+ if data_parallel_size > 1:
66
+ cmd += [
67
+ "--data-parallel-size",
68
+ str(data_parallel_size),
69
+ "--data-parallel-backend",
70
+ "ray",
71
+ ]
72
+ if enable_expert_parallel:
73
+ cmd.append("--enable-expert-parallel")
74
+ if extra_args:
75
+ cmd += extra_args
76
+ env = dict(
77
+ os.environ,
78
+ VLLM_SERVER_DEV_MODE="1",
79
+ VLLM_USE_RAY_V2_EXECUTOR_BACKEND="1",
80
+ )
81
+ # Pin the child to THIS Ray cluster. vLLM's ray DP backend reaches Ray
82
+ # through a bare ``ray.init()``, which STARTS a cluster rather than
83
+ # attaching, so wherever we started Ray ourselves (a bare CI runner) the
84
+ # workers would come up in a second cluster and never see the producer
85
+ # actors -- the failure is an actor lookup that cannot succeed. A bare
86
+ # ``ray.init()`` does honour RAY_ADDRESS, so hand it over explicitly.
87
+ if ray.is_initialized():
88
+ env["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
89
+ # Let the server's Ray workers use the ambient interpreter/install rather
90
+ # than a workspace snapshot working_dir. On an Anyscale dev workspace the
91
+ # RAY_RUNTIME_ENV_HOOK injects a git snapshot as the workers' working_dir,
92
+ # which shadows an editable vLLM install (its compiled .so is gitignored, so
93
+ # the snapshot lacks it). No-op off Anyscale (the var is unset there).
94
+ env.pop("RAY_RUNTIME_ENV_HOOK", None)
95
+ print(f"[serve] launching: {' '.join(cmd)}", flush=True)
96
+ return subprocess.Popen(cmd, env=env)
97
+
98
+
99
+def wait_for_server(
100
+ endpoint: str, proc: subprocess.Popen, timeout: float = 1800
101
+) -> None:
102
+ """Block until ``/health`` returns 200, failing fast if the server exits."""
103
+ deadline = time.time() + timeout
104
+ while time.time() < deadline:
105
+ if proc.poll() is not None:
106
+ raise RuntimeError(f"vllm serve exited early (code {proc.returncode})")
107
+ try:
108
+ if requests.get(f"{endpoint}/health", timeout=5).status_code == 200:
109
+ print("[serve] server is healthy", flush=True)
110
+ return
111
+ except requests.RequestException:
112
+ pass
113
+ time.sleep(3)
114
+ raise RuntimeError("vllm serve did not become healthy in time")
115
+
116
+
117
+def http_generate(
118
+ endpoint: str, model: str, prompts: list[str], max_tokens: int = 16
119
+) -> list[str]:
120
+ """Greedy ``/v1/completions`` for each prompt; returns the generated texts."""
121
+ outs = []
122
+ for p in prompts:
123
+ r = requests.post(
124
+ f"{endpoint}/v1/completions",
125
+ json={
126
+ "model": model,
127
+ "prompt": p,
128
+ "max_tokens": max_tokens,
129
+ "temperature": 0,
130
+ },
131
+ timeout=120,
132
+ )
133
+ r.raise_for_status()
134
+ outs.append(r.json()["choices"][0]["text"])
135
+ return outs
136
+
137
+
138
+def pause_generation(endpoint: str, mode: str = "abort") -> None:
139
+ requests.post(f"{endpoint}/pause", params={"mode": mode}, timeout=60)
140
+
141
+
142
+def resume_generation(endpoint: str) -> None:
143
+ requests.post(f"{endpoint}/resume", timeout=60)
144
+
145
+
146
+def shutdown_server(proc: subprocess.Popen) -> None:
147
+ print("[serve] shutting down vllm serve", flush=True)
148
+ proc.terminate()
149
+ try:
150
+ proc.wait(timeout=30)
151
+ except subprocess.TimeoutExpired:
152
+ proc.kill()
examples/rl/rdt_weight_source.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The ``WeightSource`` for ``rlhf_sharded_rdt_small_ep.py``.
4
+
5
+A trainer's producer job is to present its weights under the names the inference
6
+side expects, which for vLLM means HF CHECKPOINT names. Transformers does not
7
+store MoE experts that way -- it fuses each layer's into ``[E, ...]`` tensors --
8
+so that is the one conversion this file does. Kept beside the example rather than
9
+inside it so the example stays about the sync itself, alongside
10
+``rdt_vllm_serve.py``.
11
+"""
12
+
13
+import regex as re
14
+import torch
15
+
16
+from vllm.distributed.weight_transfer import ParamMeta, WeightSource
17
+from vllm.distributed.weight_transfer.base import materialize_full_tensor
18
+
19
+# Transformers fuses a layer's experts into ``[E, ...]`` tensors; the checkpoint
20
+# stores them per expert. This matches the fused parameter names.
21
+_FUSED_EXPERT_RE = re.compile(r"^(.*\.experts)\.(gate_up_proj|down_proj)$")
22
+
23
+
24
+class CheckpointNameSource(WeightSource):
25
+ """`WeightSource` over an FSDP2 model that publishes HF CHECKPOINT names.
26
+
27
+ Transformers keeps each layer's experts fused as ``[E, ...]`` tensors, but
28
+ vLLM's MoE loaders are written against the per-expert checkpoint entries --
29
+ only ``qwen{2,3}_moe`` also accept the fused names, so publishing them
30
+ straight from ``named_parameters()`` limits the example to one family and
31
+ fails the consumer's bake on ``experts.gate_up_proj`` everywhere else.
32
+ Splitting them back is also what a real trainer does, since converting an
33
+ internal layout to checkpoint names is the normal producer job.
34
+
35
+ Every rank publishes the whole model: fused params are all-gathered and then
36
+ sliced per expert. Ownership stays uniform (no ``held_names``) -- serving
37
+ only the experts a rank holds is a real trainer's concern, not an example's.
38
+ """
39
+
40
+ def __init__(self, module: torch.nn.Module) -> None:
41
+ self._module = module
42
+
43
+ @staticmethod
44
+ def _expand(name: str, param: torch.Tensor) -> list[tuple[str, tuple]]:
45
+ """The checkpoint entries ``name`` contributes, as (name, shape).
46
+
47
+ ``gate_up_proj`` is ``[E, 2I, H]`` with gate in rows ``:I`` and up in
48
+ ``I:``; ``down_proj`` is ``[E, H, I]``. Anything else is already a
49
+ checkpoint name and passes through.
50
+ """
51
+ m = _FUSED_EXPERT_RE.match(name)
52
+ if m is None:
53
+ return [(name, tuple(param.shape))]
54
+ prefix, kind = m.group(1), m.group(2)
55
+ experts, rows, cols = param.shape
56
+ if kind == "down_proj":
57
+ return [
58
+ (f"{prefix}.{e}.down_proj.weight", (rows, cols)) for e in range(experts)
59
+ ]
60
+ half = rows // 2
61
+ return [
62
+ (f"{prefix}.{e}.{proj}_proj.weight", (half, cols))
63
+ for e in range(experts)
64
+ for proj in ("gate", "up")
65
+ ]
66
+
67
+ def metadata(self) -> list[ParamMeta]:
68
+ return [
69
+ ParamMeta(entry_name, param.dtype, shape)
70
+ for name, param in self._module.named_parameters()
71
+ for entry_name, shape in self._expand(name, param)
72
+ ]
73
+
74
+ def __iter__(self):
75
+ for name, param in self._module.named_parameters():
76
+ m = _FUSED_EXPERT_RE.match(name)
77
+ if m is None:
78
+ yield name, materialize_full_tensor(param)
79
+ continue
80
+ # One gather per FUSED param, not per expert: the views below are
81
+ # into that one tensor, and every expert of a layer rides in the
82
+ # same gather group, so it stays resident exactly as long as needed.
83
+ prefix, kind = m.group(1), m.group(2)
84
+ full = materialize_full_tensor(param)
85
+ if kind == "down_proj":
86
+ for e in range(full.shape[0]):
87
+ yield f"{prefix}.{e}.down_proj.weight", full[e]
88
+ else:
89
+ half = full.shape[1] // 2
90
+ for e in range(full.shape[0]):
91
+ yield f"{prefix}.{e}.gate_proj.weight", full[e, :half, :]
92
+ yield f"{prefix}.{e}.up_proj.weight", full[e, half:, :]
examples/rl/rlhf_sharded_rdt_small_ep.py
new
+270
@@ -0,0 +1,270 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Sharded-RDT weight sync on FOUR GPUs, one node: 2 FSDP2 trainer ranks -> 2
4
+vLLM DP ranks with expert parallelism.
5
+
6
+The reference example for the backend, sized to fit a CI runner:
7
+``CheckpointNameSource`` over an FSDP2 model, one serve actor per rank, the
8
+HTTP control plane, and assertions so it can run unattended.
9
+
10
+Why a MoE model rather than a small dense one: the point of sharded RDT is that
11
+each inference worker pulls only the slices it consumes, so the example is only
12
+meaningful when no rank holds the whole model. That needs experts on both sides.
13
+``Qwen/Qwen1.5-MoE-A2.7B`` is the reference shape, a real MoE of the same family
14
+as the 30B model the big example uses. Its WEIGHTS are never read: the trainer
15
+builds the config and the server runs ``--load-format dummy``, so the sync is
16
+what is under test and nothing is downloaded but the config and tokenizer. A
17
+smaller MoE via ``RDT_MODEL`` changes nothing about what is exercised, which is
18
+what CI runs.
19
+
20
+The trainer publishes CHECKPOINT names, not the names its own modules carry:
21
+Transformers fuses each layer's experts into ``[E, ...]`` tensors
22
+(``mlp.experts.gate_up_proj``/``down_proj``) while the checkpoint stores them per
23
+expert, and vLLM's MoE loaders read the per-expert form. ``CheckpointNameSource``
24
+splits them back, which is the same conversion a real trainer does when it maps
25
+its internal layout onto checkpoint names. Publishing the fused names instead
26
+works only for the Qwen families, whose loaders happen to accept them.
27
+
28
+FSDP2 shards those fused tensors on dim 0 — the expert dimension — so after
29
+sharding no rank stores every expert, and the inference side gets real EP from
30
+``--enable-expert-parallel`` over DP2. Two simplifications an example can afford
31
+and a real trainer cannot: each rank BUILDS the whole model before
32
+``fully_shard`` splits it (fine for a tiny config, so nothing here streams a
33
+checkpoint), and iteration all-gathers every parameter, so ownership stays
34
+uniform. Serving only the experts a rank holds (``held_names``) is where a real
35
+trainer earns its keep and is deliberately out of scope.
36
+
37
+Model naming is load-bearing: names are mapped onto checkpoint keys one-to-one,
38
+so a model whose HF module names have drifted from its checkpoint (as
39
+GraniteMoe's router has) will not work here.
40
+
41
+What it checks (a smoke test, not a numerical equivalence check):
42
+ - generation CHANGES after the first sync — the weights really moved, so a
43
+ silently-skipped sync fails instead of printing plausible text;
44
+ - generation is UNCHANGED by a second sync — replay is stable.
45
+
46
+Run:
47
+ python examples/rl/rlhf_sharded_rdt_small_ep.py
48
+
49
+Needs 4 GPUs on one node, each big enough for the whole model, since every rank
50
+builds it before ``fully_shard`` splits it: the default is ~29 GiB in bf16, so
51
+40 GiB cards and up. Set ``RDT_MODEL`` to a tiny MoE on smaller GPUs. Joins an
52
+existing Ray cluster if there is one and otherwise starts its own.
53
+"""
54
+
55
+import os
56
+import sys
57
+import time
58
+
59
+import ray
60
+import torch
61
+import torch.distributed as dist
62
+from ray.util.placement_group import placement_group, placement_group_table
63
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
64
+from rdt_vllm_serve import (
65
+ http_generate,
66
+ launch_vllm_serve,
67
+ pause_generation,
68
+ resume_generation,
69
+ shutdown_server,
70
+ wait_for_server,
71
+)
72
+from rdt_weight_source import CheckpointNameSource
73
+from torch.distributed.fsdp import fully_shard
74
+from transformers import AutoConfig, AutoModelForCausalLM
75
+
76
+from vllm.distributed.weight_transfer import (
77
+ HTTPVLLMWeightSyncClient,
78
+ WeightTransferTrainerFactory,
79
+)
80
+from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
81
+ ShardedRDTTrainerInitInfo,
82
+)
83
+from vllm.utils.network_utils import get_open_port
84
+
85
+# Weights are built, never loaded, so this costs a config fetch -- but every rank
86
+# materializes the model before sharding, so the default wants a 40 GiB card.
87
+MODEL_NAME = os.environ.get("RDT_MODEL", "Qwen/Qwen1.5-MoE-A2.7B")
88
+RAY_NAMESPACE = "sharded_rdt_small_ep_example"
89
+VLLM_PORT = int(os.environ.get("RDT_VLLM_PORT", "8100"))
90
+VLLM_ENDPOINT = f"http://127.0.0.1:{VLLM_PORT}"
91
+
92
+FSDP_WORLD_SIZE = 2
93
+INFERENCE_TP_SIZE = 1
94
+INFERENCE_DP_SIZE = 2
95
+NUM_INFERENCE_CONSUMERS = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE
96
+
97
+PROMPTS = ["The capital of France is", "The future of AI is"]
98
+
99
+
100
+@ray.remote(num_gpus=1)
101
+class FSDPTrainWorker:
102
+ """One FSDP2 rank per GPU. Identical to the 8-GPU example's worker: the
103
+ sharded-RDT engine owns the NIXL serve surface, so this stays a plain Ray
104
+ actor with no producer mixin."""
105
+
106
+ def __init__(self, model_name, rank, world_size, master_addr, master_port):
107
+ self.rank = rank
108
+ self.engine = None
109
+ os.environ["MASTER_ADDR"] = master_addr
110
+ os.environ["MASTER_PORT"] = str(master_port)
111
+ dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
112
+ torch.accelerator.set_device_index(0)
113
+
114
+ # Same seed on every rank, so the shards compose into ONE model rather
115
+ # than a blend of independent inits -- the second sync's stability check
116
+ # is only meaningful against coherent weights.
117
+ torch.manual_seed(0)
118
+ config = AutoConfig.from_pretrained(model_name)
119
+ with torch.device("cuda"):
120
+ model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16)
121
+ for layer in model.model.layers:
122
+ fully_shard(layer)
123
+ fully_shard(model)
124
+ self.model = model
125
+
126
+ def get_rank(self):
127
+ return self.rank
128
+
129
+ def setup_engine(self, vllm_endpoint: str):
130
+ self.engine = WeightTransferTrainerFactory.trainer_init(
131
+ ShardedRDTTrainerInitInfo(
132
+ rank=self.rank,
133
+ num_consumers=NUM_INFERENCE_CONSUMERS,
134
+ trainer_actor_namespace=RAY_NAMESPACE,
135
+ num_rdt_buffers=int(os.environ.get("NUM_RDT_BUFFERS", "2")),
136
+ buffer_presize_gb=float(os.environ.get("RDT_BUFFER_PRESIZE_GB", "0")),
137
+ ),
138
+ client=HTTPVLLMWeightSyncClient(vllm_endpoint),
139
+ source=CheckpointNameSource(self.model),
140
+ )
141
+
142
+ def sync_weights(self):
143
+ self.engine.send_weights()
144
+
145
+
146
+def main():
147
+ # Ship a minimal working_dir (this example dir) so Ray actors do NOT inherit
148
+ # a workspace snapshot that shadows the editable vLLM install.
149
+ runtime_env: dict[str, object] = {
150
+ "py_executable": sys.executable,
151
+ "working_dir": os.path.dirname(os.path.abspath(__file__)),
152
+ }
153
+ forwarded = {
154
+ k: os.environ[k]
155
+ for k in (
156
+ "NCCL_CUMEM_ENABLE",
157
+ "VLLM_NCCL_SO_PATH",
158
+ "LD_PRELOAD",
159
+ "LD_LIBRARY_PATH",
160
+ )
161
+ if k in os.environ
162
+ }
163
+ if forwarded:
164
+ runtime_env["env_vars"] = forwarded
165
+ if not ray.is_initialized():
166
+ try:
167
+ ray.init(address="auto", runtime_env=runtime_env, namespace=RAY_NAMESPACE)
168
+ except (ConnectionError, RuntimeError):
169
+ # No cluster to join (a bare CI runner): start one on this node.
170
+ # `launch_vllm_serve` hands the child this cluster's address, which
171
+ # it needs -- vLLM's ray DP backend would otherwise start its own.
172
+ ray.init(runtime_env=runtime_env, namespace=RAY_NAMESPACE)
173
+
174
+ print(
175
+ f"[init] model {MODEL_NAME}, {FSDP_WORLD_SIZE} trainer + "
176
+ f"{INFERENCE_DP_SIZE} inference GPUs"
177
+ )
178
+
179
+ # Reserve the trainer's GPUs BEFORE launching the server, so the server's DP
180
+ # workers take the node's remaining two. STRICT_PACK keeps the trainer's
181
+ # all-gathers on one node's NVLink -- here that is the only node anyway.
182
+ trainer_pg = placement_group(
183
+ [{"GPU": 1, "CPU": 1}] * FSDP_WORLD_SIZE, strategy="STRICT_PACK"
184
+ )
185
+ ray.get(trainer_pg.ready())
186
+ pg_node_id = next(
187
+ iter(placement_group_table(trainer_pg)["bundles_to_node_id"].values())
188
+ )
189
+ master_addr = next(
190
+ n["NodeManagerAddress"] for n in ray.nodes() if n["NodeID"] == pg_node_id
191
+ )
192
+
193
+ @ray.remote(
194
+ num_cpus=0,
195
+ scheduling_strategy=PlacementGroupSchedulingStrategy(
196
+ placement_group=trainer_pg
197
+ ),
198
+ )
199
+ def _free_port():
200
+ return get_open_port()
201
+
202
+ master_port = ray.get(_free_port.remote())
203
+
204
+ workers = [
205
+ FSDPTrainWorker.options(
206
+ scheduling_strategy=PlacementGroupSchedulingStrategy(
207
+ placement_group=trainer_pg, placement_group_bundle_index=rank
208
+ ),
209
+ ).remote(MODEL_NAME, rank, FSDP_WORLD_SIZE, master_addr, master_port)
210
+ for rank in range(FSDP_WORLD_SIZE)
211
+ ]
212
+ ray.get([w.get_rank.remote() for w in workers])
213
+ print(f"[init] {FSDP_WORLD_SIZE} FSDP trainer workers ready (seeded weights).")
214
+
215
+ server = launch_vllm_serve(
216
+ MODEL_NAME,
217
+ tensor_parallel_size=INFERENCE_TP_SIZE,
218
+ data_parallel_size=INFERENCE_DP_SIZE,
219
+ enable_expert_parallel=True,
220
+ port=VLLM_PORT,
221
+ extra_args=["--max-model-len", "2048", "--max-num-seqs", "8"],
222
+ )
223
+ try:
224
+ wait_for_server(VLLM_ENDPOINT, server)
225
+
226
+ before = http_generate(VLLM_ENDPOINT, MODEL_NAME, PROMPTS)
227
+ print("[generate] BEFORE sync (vLLM's own dummy init):")
228
+ for p, t in zip(PROMPTS, before):
229
+ print(f" {p!r} -> {t!r}")
230
+
231
+ print("[sync] building trainer engines (bake on the sender)...")
232
+ t0 = time.perf_counter()
233
+ ray.get([w.setup_engine.remote(VLLM_ENDPOINT) for w in workers])
234
+ print(f"[sync] engine setup (incl. bake) took {time.perf_counter() - t0:.2f}s")
235
+
236
+ pause_generation(VLLM_ENDPOINT)
237
+ t0 = time.perf_counter()
238
+ ray.get([w.sync_weights.remote() for w in workers])
239
+ print(f"[sync] iter 0 took {time.perf_counter() - t0:.3f}s")
240
+ resume_generation(VLLM_ENDPOINT)
241
+
242
+ after = http_generate(VLLM_ENDPOINT, MODEL_NAME, PROMPTS)
243
+ print("[generate] AFTER sync (the trainer's weights):")
244
+ for p, t in zip(PROMPTS, after):
245
+ print(f" {p!r} -> {t!r}")
246
+
247
+ # Sync again; replay must be stable, so generation must not move.
248
+ pause_generation(VLLM_ENDPOINT)
249
+ t0 = time.perf_counter()
250
+ ray.get([w.sync_weights.remote() for w in workers])
251
+ print(f"[sync] iter 1 took {time.perf_counter() - t0:.3f}s")
252
+ resume_generation(VLLM_ENDPOINT)
253
+ again = http_generate(VLLM_ENDPOINT, MODEL_NAME, PROMPTS)
254
+
255
+ # Greedy decoding, so both comparisons are exact.
256
+ assert after != before, (
257
+ f"generation did not change across the weight sync, so the sync "
258
+ f"moved nothing: {before!r}"
259
+ )
260
+ assert again == after, (
261
+ f"a second sync changed generation, so replay is not stable: "
262
+ f"{after!r} then {again!r}"
263
+ )
264
+ print("[ok] weights moved on sync 0 and replay was stable on sync 1")
265
+ finally:
266
+ shutdown_server(server)
267
+
268
+
269
+if __name__ == "__main__":
270
+ main()
requirements/test/cpu.txt
+1
-1
@@ -895,7 +895,7 @@ pyzmq==27.1.0
895
# via -r requirements/test/../common.txt
896
rapidfuzz==3.12.1
897
# via jiwer
898
-ray==2.48.0
898
+ray==2.56.1
899
# via -r requirements/test/cuda.in
900
redis==5.2.0
901
# via tensorizer
requirements/test/cuda.in
+1
-1
@@ -22,7 +22,7 @@ vector_quantize_pytorch # required for minicpmo_26 test
22
vocos # required for minicpmo_26 test
23
peft>=0.19.1 # required for phi-4-mm test
24
pqdm
25
-ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
25
+ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56)
26
sentence-transformers>=5.2.0 # required for embedding tests
27
soundfile # required for audio tests
28
jiwer # required for audio tests
requirements/test/cuda.txt
+1
-1
@@ -988,7 +988,7 @@ pyzmq==27.1.0
988
# -r requirements/test/../common.txt
989
rapidfuzz==3.12.1
990
# via jiwer
991
-ray==2.48.0
991
+ray==2.56.1
992
# via -r requirements/test/cuda.in
993
redis==5.2.0
994
# via tensorizer
requirements/test/nightly-torch.txt
+1
-1
@@ -16,7 +16,7 @@ librosa # required for audio tests
16
vocos # required for minicpmo_26 test
17
peft
18
pqdm
19
-ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
19
+ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56)
20
sentence-transformers>=5.2.0 # required for embedding tests
21
soundfile # required for audio tests
22
jiwer # required for audio tests
tests/distributed/test_sharded_rdt_plan.py
new
+1753
@@ -0,0 +1,1753 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Consumer-side unit tests for the sharded-RDT backend.
4
+
5
+The consumer engine's planning core is pure — it turns the baked plan plus the
6
+driver's gather-group partition into a static `_CallPlan` with no pulls, no
7
+model and no Ray — so most of it is exercised here on CPU/meta tensors only.
8
+These tests pin the values the engine produces today (chunk boundaries, packed
9
+byte offsets, recorded op chains) so a refactor that changes them fails loudly
10
+instead of silently shipping different bytes.
11
+
12
+`TestBakeOnARealModel` at the end is the exception: it needs a GPU, because the
13
+bake only exists to be driven through a real `model.load_weights`, and the parts
14
+that matter there — that the monkeypatched loaders are put back, and that the
15
+recorded op chains reproduce what the loaders did — cannot be faked.
16
+
17
+`tests/distributed/test_sharded_rdt_trainer.py` covers the trainer (producer)
18
+side; `test_sharded_rdt_producer.py` covers its serve-actor protocol.
19
+"""
20
+
21
+from types import SimpleNamespace
22
+
23
+import pytest
24
+import torch
25
+
26
+from vllm.distributed.weight_transfer.base import layerwise_groups
27
+from vllm.distributed.weight_transfer.sharded_rdt_common import (
28
+ ALLOWED_OPS,
29
+ SUPPORTED_OPS,
30
+ RdtRouter,
31
+ assign_producer_indices,
32
+ buffer_alloc_bytes,
33
+)
34
+from vllm.distributed.weight_transfer.sharded_rdt_engine import (
35
+ ShardedRDTWeightTransferEngine,
36
+ ShardedRDTWeightTransferInitInfo,
37
+ _dtype_from_name,
38
+)
39
+from vllm.distributed.weight_transfer.sharded_rdt_fake import (
40
+ BakeSink,
41
+ FakeRDTTensor,
42
+ _Scatter,
43
+ _UnsupportedFakeOp,
44
+)
45
+
46
+META = torch.device("meta")
47
+
48
+
49
+# ---------------------------------------------------------------------------
50
+# Harness
51
+# ---------------------------------------------------------------------------
52
+
53
+
54
+class _FakeLayer:
55
+ """Stand-in for a leaf module. The planner only stores it and keys on
56
+ ``id()``; nothing in the plan path touches its attributes."""
57
+
58
+ def __init__(self, tag: str) -> None:
59
+ self.tag = tag
60
+
61
+ def __repr__(self) -> str:
62
+ return f"_FakeLayer({self.tag!r})"
63
+
64
+
65
+def _module(layer, copies):
66
+ """What the bake records for one leaf module: a scatter list whose entries
67
+ all carry that module as their ``layer``."""
68
+ for c in copies:
69
+ c.layer = layer
70
+ return copies
71
+
72
+
73
+def _copy(
74
+ name,
75
+ layer_param="weight",
76
+ *,
77
+ offset=0,
78
+ shape=(4,),
79
+ ops=(),
80
+ layer=None,
81
+ dtype=torch.bfloat16,
82
+):
83
+ """A `_Scatter` as the bake would record it: source key, owning layer, and
84
+ the meta destination region."""
85
+ stride = (
86
+ (1,)
87
+ if len(shape) == 1
88
+ else tuple(
89
+ int(torch.empty(shape, device=META).stride()[i]) for i in range(len(shape))
90
+ )
91
+ )
92
+ return _Scatter(
93
+ layer=layer,
94
+ param_name=layer_param,
95
+ src=(name, tuple(ops)),
96
+ offset=offset,
97
+ shape=tuple(shape),
98
+ stride=stride,
99
+ dtype=dtype,
100
+ )
101
+
102
+
103
+def _planner(baked, *, name_meta, live=None, held_by=None, group_lens=None):
104
+ """A planning-only engine.
105
+
106
+ `__init__` needs a config/model/device the planner never reads, so build the
107
+ instance without it and set exactly the state `_build_call_plan` consumes:
108
+ the baked plan, the name metadata, the live-name set, and a bound router.
109
+
110
+ ``held_by``: name -> the ranks holding it. Names left out are held by every
111
+ producer (the replicated class 0). Default: one producer holding everything,
112
+ the degenerate router. Stand-in ranks 0-3 let plan tests assert chunk
113
+ structure; owner VALUES are pinned by TestCallPlanRouting.
114
+ """
115
+ names = list(name_meta)
116
+ if held_by:
117
+ n_prod = max(r for owners in held_by.values() for r in owners) + 1
118
+ owner_sets = [list(range(n_prod))]
119
+ seen = {tuple(range(n_prod)): 0}
120
+ classes = []
121
+ for n in names:
122
+ key = tuple(sorted(held_by.get(n, range(n_prod))))
123
+ if key not in seen:
124
+ seen[key] = len(owner_sets)
125
+ owner_sets.append(list(key))
126
+ classes.append(seen[key])
127
+ else:
128
+ n_prod, owner_sets, classes = 1, [[0]], [0] * len(names)
129
+ eng = object.__new__(ShardedRDTWeightTransferEngine)
130
+ eng._name_to_plan = dict(baked)
131
+ eng._name_meta = dict(name_meta)
132
+ eng._live_names = set(live or name_meta)
133
+ router = RdtRouter(
134
+ n_prod, 1, owner_sets, classes, names, group_lens or [1] * len(names)
135
+ )
136
+ router.bind([object()] * n_prod, [object()] * n_prod, 0)
137
+ eng._router = router
138
+ return eng
139
+
140
+
141
+def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4):
142
+ """`(baked, name_meta, names, group_lens)` for a pre / N-layer / post model
143
+ where every gather group is one fully-baked leaf module."""
144
+ names = ["embed.weight"]
145
+ names += [f"model.layers.{i}.w" for i in range(n_layers)]
146
+ names += ["norm.weight"]
147
+ baked = {}
148
+ for n in names:
149
+ layer = _FakeLayer(n)
150
+ baked[n] = _module(layer, [_copy(n, shape=(numel,))])
151
+ name_meta = {n: (dtype, [numel]) for n in names}
152
+ group_lens = [1] * len(names)
153
+ return baked, name_meta, names, group_lens
154
+
155
+
156
+# ---------------------------------------------------------------------------
157
+# FakeRDTTensor: op-chain recording
158
+# ---------------------------------------------------------------------------
159
+
160
+
161
+class TestFakeOpChains:
162
+ """Every allowlisted op must append itself to the chain and hand back a
163
+ child whose shape/dtype PyTorch itself computed. The chain is the wire
164
+ format the producer replays, so its exact contents are load-bearing."""
165
+
166
+ def _fake(self, shape=(4, 6), dtype=torch.bfloat16, sink=None):
167
+ return FakeRDTTensor(
168
+ name="w", shape=torch.Size(shape), dtype=dtype, device=META, sink=sink
169
+ )
170
+
171
+ def test_bare_fake_has_metadata_but_no_chain(self):
172
+ t = self._fake()
173
+ assert t.shape == (4, 6)
174
+ assert t.dtype is torch.bfloat16
175
+ assert t._key() == ("w", ())
176
+
177
+ @pytest.mark.parametrize(
178
+ "call,expect_op,expect_shape",
179
+ [
180
+ (lambda t: t.narrow(0, 1, 2), ("narrow", (0, 1, 2), ()), (2, 6)),
181
+ (lambda t: t.view(6, 4), ("view", (6, 4), ()), (6, 4)),
182
+ (lambda t: t.reshape(24), ("reshape", (24,), ()), (24,)),
183
+ (lambda t: t[1], ("__getitem__", (1,), ()), (6,)),
184
+ (lambda t: t.unsqueeze(0), ("unsqueeze", (0,), ()), (1, 4, 6)),
185
+ (lambda t: t.transpose(0, 1), ("transpose", (0, 1), ()), (6, 4)),
186
+ (lambda t: t.t(), ("t", (), ()), (6, 4)),
187
+ (lambda t: t.permute(1, 0), ("permute", (1, 0), ()), (6, 4)),
188
+ (lambda t: t.flatten(), ("flatten", (), ()), (24,)),
189
+ (lambda t: t.contiguous(), ("contiguous", (), ()), (4, 6)),
190
+ ],
191
+ )
192
+ def test_single_return_op_appends_one_spec(self, call, expect_op, expect_shape):
193
+ child = call(self._fake())
194
+ assert isinstance(child, FakeRDTTensor)
195
+ assert child._key() == ("w", (expect_op,))
196
+ assert tuple(child.shape) == expect_shape
197
+
198
+ def test_squeeze_records_its_argument(self):
199
+ child = self._fake(shape=(1, 4)).squeeze(0)
200
+ assert child._key() == ("w", (("squeeze", (0,), ()),))
201
+ assert tuple(child.shape) == (4,)
202
+
203
+ def test_chains_compose_in_call_order(self):
204
+ child = self._fake().t().narrow(0, 0, 3).flatten()
205
+ assert child._key() == (
206
+ "w",
207
+ (
208
+ ("t", (), ()),
209
+ ("narrow", (0, 0, 3), ()),
210
+ ("flatten", (), ()),
211
+ ),
212
+ )
213
+ assert tuple(child.shape) == (12,)
214
+
215
+ def test_kwargs_are_frozen_sorted_for_hashability(self):
216
+ child = self._fake().narrow(dim=0, start=1, length=2)
217
+ (op, args, kwargs) = child._key()[1][0]
218
+ assert (op, args) == ("narrow", ())
219
+ assert kwargs == (("dim", 0), ("length", 2), ("start", 1))
220
+ hash(child._key()) # the chain is used as a dict key on both sides
221
+
222
+ @pytest.mark.parametrize("op,n", [("chunk", 2), ("unbind", 4)])
223
+ def test_multi_return_op_emits_one_child_per_output(self, op, n):
224
+ """chunk/unbind hand back a tuple; each child carries the base op plus a
225
+ trailing __getitem__(i) so the producer can index the replayed result."""
226
+ parts = getattr(self._fake(), op)(*((n,) if op == "chunk" else ()), 0)
227
+ assert isinstance(parts, tuple)
228
+ assert len(parts) == n
229
+ for i, part in enumerate(parts):
230
+ base, index = part._key()[1]
231
+ assert base[0] == op
232
+ assert index == ("__getitem__", (i,), ())
233
+
234
+ def test_op_chain_is_hashable_as_a_fetch_key(self):
235
+ keys = {self._fake().t()._key(), self._fake().t()._key()}
236
+ assert len(keys) == 1, "equal chains must collapse — they dedup pull keys"
237
+
238
+
239
+class TestFakeUnsupportedOps:
240
+ """Anything that needs real data must fail loudly at bake time rather than
241
+ silently transferring the wrong bytes."""
242
+
243
+ def _fake(self):
244
+ return FakeRDTTensor(
245
+ name="w",
246
+ shape=torch.Size((4,)),
247
+ dtype=torch.bfloat16,
248
+ device=META,
249
+ sink=None,
250
+ )
251
+
252
+ @pytest.mark.parametrize(
253
+ "call",
254
+ [
255
+ lambda t: t.to(torch.float32),
256
+ lambda t: t.float(),
257
+ lambda t: t + 1,
258
+ lambda t: t * 2,
259
+ lambda t: t.sum(),
260
+ ],
261
+ ids=["to", "float", "add", "mul", "sum"],
262
+ )
263
+ def test_data_dependent_ops_raise(self, call):
264
+ with pytest.raises(_UnsupportedFakeOp):
265
+ call(self._fake())
266
+
267
+ def test_error_names_the_weight_and_the_chain(self):
268
+ with pytest.raises(_UnsupportedFakeOp) as exc:
269
+ self._fake().narrow(0, 0, 2).float()
270
+ msg = str(exc.value)
271
+ assert "'w'" in msg
272
+ assert "narrow" in msg
273
+
274
+ def test_unsupported_is_a_notimplementederror(self):
275
+ """Callers distinguish "this backend can't handle the loader" from bugs."""
276
+ assert issubclass(_UnsupportedFakeOp, NotImplementedError)
277
+
278
+
279
+class TestBakeRecording:
280
+ """During the dry run the fake's ``copy_`` is the data sink: it records the
281
+ source chain plus the meta destination's strided region and moves nothing."""
282
+
283
+ def _recorder_and_fake(self, shape=(4, 6)):
284
+ rec = BakeSink()
285
+ fake = FakeRDTTensor(
286
+ name="w",
287
+ shape=torch.Size(shape),
288
+ dtype=torch.bfloat16,
289
+ device=META,
290
+ sink=rec,
291
+ )
292
+ return rec, fake
293
+
294
+ def test_copy_records_the_destination_region(self):
295
+ rec, fake = self._recorder_and_fake()
296
+ layer = _FakeLayer("q_proj")
297
+ param = torch.empty((8, 6), dtype=torch.bfloat16, device=META)
298
+ dest = param.narrow(0, 4, 4) # the second half of a fused param
299
+ rec.current = (layer, "weight")
300
+ dest.copy_(fake)
301
+
302
+ (recorded,) = rec.copies_by_layer[layer]
303
+ assert recorded.src == ("w", ())
304
+ assert recorded.param_name == "weight"
305
+ assert recorded.offset == dest.storage_offset() == 24
306
+ assert recorded.shape == (4, 6)
307
+ assert recorded.stride == (6, 1)
308
+
309
+ def test_copy_marks_the_source_name_live(self):
310
+ rec, fake = self._recorder_and_fake()
311
+ rec.current = (_FakeLayer("l"), "weight")
312
+ torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake)
313
+ assert rec.copied_names == {"w"}
314
+
315
+ def test_unattributed_copy_is_live_but_unrecorded(self):
316
+ """A copy_ with no loader stamp cannot be attributed to a param, so its
317
+ module must fall back to the plain load — but the name still moved data."""
318
+ rec, fake = self._recorder_and_fake()
319
+ rec.current = None
320
+ torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake)
321
+ assert rec.copied_names == {"w"}
322
+ assert dict(rec.copies_by_layer) == {}
323
+
324
+ def test_copies_are_grouped_by_module_in_call_order(self):
325
+ rec, _ = self._recorder_and_fake()
326
+ layer = _FakeLayer("gate_up")
327
+ for i, name in enumerate(("gate", "up")):
328
+ fake = FakeRDTTensor(
329
+ name=name,
330
+ shape=torch.Size((4,)),
331
+ dtype=torch.bfloat16,
332
+ device=META,
333
+ sink=rec,
334
+ )
335
+ param = torch.empty((8,), dtype=torch.bfloat16, device=META)
336
+ rec.current = (layer, "weight")
337
+ param.narrow(0, 4 * i, 4).copy_(fake)
338
+ assert [c.src[0] for c in rec.copies_by_layer[layer]] == ["gate", "up"]
339
+ assert [c.offset for c in rec.copies_by_layer[layer]] == [0, 4]
340
+
341
+ def test_recording_a_sliced_source(self):
342
+ """The chain on the source and the region on the dest are independent."""
343
+ rec, fake = self._recorder_and_fake(shape=(8, 6))
344
+ layer = _FakeLayer("k_proj")
345
+ rec.current = (layer, "weight")
346
+ param = torch.empty((4, 6), dtype=torch.bfloat16, device=META)
347
+ param.copy_(fake.narrow(0, 2, 4))
348
+ (recorded,) = rec.copies_by_layer[layer]
349
+ assert recorded.src == ("w", (("narrow", (0, 2, 4), ()),))
350
+ assert recorded.shape == (4, 6)
351
+
352
+ def test_the_recorded_dtype_is_the_chains_output_not_the_sources(self):
353
+ """``view(dtype)`` is allowlisted, so a chain can reinterpret dtype. The
354
+ producer packs what the replay yields, so the record must carry the
355
+ POST-chain dtype: taking it from the source name's metadata instead sizes
356
+ the slice with the wrong itemsize and shifts every later slice in the
357
+ chunk, carving the packed blob differently on the two sides."""
358
+ rec, fake = self._recorder_and_fake(shape=(4,))
359
+ layer = _FakeLayer("reinterpreted")
360
+ rec.current = (layer, "weight")
361
+ viewed = fake.view(torch.float32) # 4 x bf16 -> 2 x f32
362
+ param = torch.empty((2,), dtype=torch.float32, device=META)
363
+ param.copy_(viewed)
364
+
365
+ (recorded,) = rec.copies_by_layer[layer]
366
+ assert recorded.dtype is torch.float32, "recorded the source dtype"
367
+ assert recorded.shape == (2,)
368
+ # What the source name alone would have said, and what the producer
369
+ # would NOT have sent.
370
+ assert 2 * torch.float32.itemsize == 8 != 2 * torch.bfloat16.itemsize
371
+
372
+ def test_a_broadcasting_copy_is_refused_at_bake(self):
373
+ """The packed layout sizes each slice from the DESTINATION shape while
374
+ the producer packs what the chain yields, and nothing downstream can
375
+ catch a mismatch — the consumer's buffer view is exactly prod(dest.shape)
376
+ elements, so it reshapes cleanly over bytes laid out at other offsets.
377
+ So the bake refuses rather than recording a slice it cannot carve."""
378
+ rec, fake = self._recorder_and_fake(shape=(1, 6))
379
+ rec.current = (_FakeLayer("broadcast"), "weight")
380
+ param = torch.empty((4, 6), dtype=torch.bfloat16, device=META)
381
+ with pytest.raises(_UnsupportedFakeOp, match="broadcasting"):
382
+ param.copy_(fake)
383
+
384
+
385
+# ---------------------------------------------------------------------------
386
+# Consumer identity
387
+# ---------------------------------------------------------------------------
388
+
389
+
390
+class TestConsumerIdentity:
391
+ """Every worker in the fleet needs a DISTINCT id in 0..C-1: it selects the
392
+ worker's producer block and keys the producer's per-consumer serve ring, so
393
+ a collision silently serves two workers out of one ring."""
394
+
395
+ def _engine(self, *, dp_index, rank, world_size):
396
+ eng = object.__new__(ShardedRDTWeightTransferEngine)
397
+ eng.parallel_config = SimpleNamespace(
398
+ data_parallel_index=dp_index, rank=rank, world_size=world_size
399
+ )
400
+ return eng
401
+
402
+ def _ids(self, *, num_consumers, workers, replica_rank=0, num_replicas=1):
403
+ """Consumer ids for a DP-only engine whose workers are ``dp_index``es."""
404
+ info = ShardedRDTWeightTransferInitInfo(
405
+ num_consumers=num_consumers,
406
+ replica_rank=replica_rank,
407
+ num_replicas=num_replicas,
408
+ )
409
+ out = []
410
+ for dp_index in workers:
411
+ eng = self._engine(dp_index=dp_index, rank=0, world_size=1)
412
+ eng._num_consumers_override = num_consumers
413
+ out.append(eng._resolve_consumer_id(info))
414
+ return out
415
+
416
+ def test_one_engine_indexes_by_dp_and_tp(self):
417
+ """dense-via-TP and MoE-via-DP+EP both flatten to 0..C-1."""
418
+ eng = self._engine(dp_index=0, rank=3, world_size=8)
419
+ eng._num_consumers_override = 8
420
+ info = ShardedRDTWeightTransferInitInfo(num_consumers=8)
421
+ assert eng._resolve_consumer_id(info) == 3
422
+
423
+ eng = self._engine(dp_index=5, rank=0, world_size=1)
424
+ eng._num_consumers_override = 8
425
+ assert eng._resolve_consumer_id(info) == 5
426
+
427
+ def test_independent_engines_offset_into_distinct_ranges(self):
428
+ """A fleet of separate engines restarts _global_worker_index at 0 in
429
+ each, so without the replica offset every engine would claim 0..w-1.
430
+ The driver sets replica_rank on the payload; these are the ids it buys."""
431
+ ids = [
432
+ self._ids(num_consumers=8, workers=range(4), replica_rank=r, num_replicas=2)
433
+ for r in (0, 1)
434
+ ]
435
+ assert ids == [[0, 1, 2, 3], [4, 5, 6, 7]]
436
+ assert len(set(ids[0] + ids[1])) == 8, "every worker distinct fleet-wide"
437
+
438
+ def test_the_default_is_a_single_engine_with_no_offset(self):
439
+ assert self._ids(num_consumers=4, workers=range(4)) == [0, 1, 2, 3]
440
+
441
+
442
+# ---------------------------------------------------------------------------
443
+# RdtRouter: per-name ownership -> one producer per pull
444
+# ---------------------------------------------------------------------------
445
+
446
+
447
+def _router(
448
+ num_producers,
449
+ num_consumers,
450
+ owner_sets=None,
451
+ name_owner_class=None,
452
+ names=None,
453
+ group_lens=None,
454
+):
455
+ """A router over ``n`` single-name groups unless told otherwise, so a test
456
+ can state expectations in group indices."""
457
+ if names is None:
458
+ n = len(name_owner_class or []) or (len(group_lens or []) or 1)
459
+ names = [f"g{i}" for i in range(n)]
460
+ if group_lens is None:
461
+ group_lens = [1] * len(names)
462
+ return RdtRouter(
463
+ num_producers, num_consumers, owner_sets, name_owner_class, names, group_lens
464
+ )
465
+
466
+
467
+class TestRdtRouter:
468
+ """Who serves each name.
469
+
470
+ A wrong answer is not a wrong number but a hang or a loud misroute: a
471
+ consumer pulling from a producer that never gathered the name trips its
472
+ served-names guard. Freeing does NOT route this way — every consumer signals
473
+ every owner of a group and the producers count to the live total — so these
474
+ pin PULL routing, plus ``group_owners`` for the free fan-out.
475
+ """
476
+
477
+ def test_identity_when_fleets_match(self):
478
+ r = _router(8, 8, group_lens=[1] * 6)
479
+ assert all(r.producer_for(3, f"g{g}") == 3 for g in range(6))
480
+ assert all(r.producer_for(c, "g0") == c for c in range(8))
481
+
482
+ def test_every_producer_nic_carries_traffic(self):
483
+ """16 producers / 8 consumers: the block rule spreads each consumer's
484
+ pulls over its block, alternating by group, so no producer sits idle.
485
+ Plain `consumer_id % len(owners)` would idle half of them."""
486
+ r = _router(16, 8, group_lens=[1] * 95)
487
+ names = [f"g{i}" for i in range(95)]
488
+ for c in range(8):
489
+ block = assign_producer_indices(16, 8, c)
490
+ assert {r.producer_for(c, n) for n in names} == set(block)
491
+ assert [r.producer_for(0, f"g{g}") for g in range(4)] == [0, 1, 0, 1]
492
+ assert {r.producer_for(c, n) for c in range(8) for n in names} == set(range(16))
493
+
494
+ def test_fan_in_shares_one_producer(self):
495
+ r = _router(2, 8, group_lens=[1] * 5)
496
+ names = [f"g{i}" for i in range(5)]
497
+ assert [{r.producer_for(c, n) for n in names} for c in range(8)] == [
498
+ {c // 4} for c in range(8)
499
+ ]
500
+
501
+ def test_pipeline_stages_route_to_the_owning_stage(self):
502
+ """2 stages x 8 ranks: groups 0-2 on stage 0, 3-5 on stage 1. Two owner
503
+ sets, one per stage; each consumer reaches both."""
504
+ r = _router(
505
+ 16,
506
+ 8,
507
+ owner_sets=[list(range(8)), list(range(8, 16))],
508
+ name_owner_class=[0, 0, 0, 1, 1, 1],
509
+ )
510
+ for c in range(8):
511
+ picks = [r.producer_for(c, f"g{g}") for g in range(6)]
512
+ assert picks == [c] * 3 + [c + 8] * 3
513
+ assert r.group_owners(0) == list(range(8))
514
+ assert r.group_owners(3) == list(range(8, 16))
515
+ r.validate()
516
+
517
+ def test_expert_names_route_to_the_holding_ranks(self):
518
+ """A name held by only some ranks routes to those; a replicated name
519
+ keeps the full set."""
520
+ r = RdtRouter(
521
+ 8,
522
+ 4,
523
+ [[0, 1, 2, 3], [2, 3], [4, 5, 6, 7], [6, 7]],
524
+ [0, 1, 2, 3],
525
+ ["repl0", "expert0", "repl1", "expert1"],
526
+ [2, 2],
527
+ )
528
+ r.validate()
529
+ assert r.owners("expert0") == [2, 3]
530
+ assert r.owners("expert1") == [6, 7]
531
+ assert r.owners("repl1") == [4, 5, 6, 7]
532
+ for c in range(4):
533
+ assert r.producer_for(c, "expert0") in (2, 3)
534
+ assert r.producer_for(c, "repl1") in (4, 5, 6, 7)
535
+
536
+ def test_consumers_spread_over_an_owner_set(self):
537
+ """Several ranks hold a name (its TP peers); the block rule must spread
538
+ consumers across them, not funnel through one."""
539
+ r = _router(
540
+ 4, 4, owner_sets=[[0, 1]], name_owner_class=[0] * 6, group_lens=[1] * 6
541
+ )
542
+ served = {r.producer_for(c, f"g{g}") for c in range(4) for g in range(6)}
543
+ assert served == {0, 1}
544
+
545
+ def test_a_group_frees_at_every_owner_of_any_of_its_names(self):
546
+ """The free barrier is per group, so its fan-out is the UNION over the
547
+ group's names: a rank holding none of them is not signalled."""
548
+ r = RdtRouter(4, 2, [[0, 1], [2, 3]], [0, 1, 0], ["a", "b", "c"], [2, 1])
549
+ assert r.group_owners(0) == [0, 1, 2, 3]
550
+ assert r.group_owners(1) == [0, 1]
551
+
552
+ def test_an_unowned_name_raises(self):
553
+ r = RdtRouter(2, 2, [[]], [0], ["w"], [1])
554
+ with pytest.raises(ValueError, match="has no owner"):
555
+ r.producer_for(0, "w")
556
+
557
+ def test_validate_rejects_an_empty_owner_set(self):
558
+ with pytest.raises(ValueError, match="empty"):
559
+ RdtRouter(
560
+ 4, 2, [[0, 1], [], [2, 3]], [0, 1, 2], ["a", "b", "c"], [3]
561
+ ).validate()
562
+
563
+ def test_validate_rejects_an_out_of_range_owner(self):
564
+ with pytest.raises(ValueError, match="out of range"):
565
+ RdtRouter(2, 2, [[0, 5]], [0], ["a"], [1]).validate()
566
+
567
+ def test_validate_rejects_an_unresolvable_class(self):
568
+ with pytest.raises(ValueError, match="owner class"):
569
+ RdtRouter(4, 2, [[0, 1]], [0, 7], ["a", "b"], [2]).validate()
570
+
571
+
572
+class TestReplicaOverlay:
573
+ """[RDT-SHARE-SLOTS] Several inference deployments carve the same blocks.
574
+
575
+ The producer can serve R deployments out of ONE slot only if the R copies of
576
+ a worker meet on the same producer asking for the same bytes. Identical
577
+ plans give the second half for free; this is the first half. Carving over
578
+ the whole fleet instead spreads the multi-owner (replicated) names of
579
+ different deployments onto different producers, so a producer ends up
580
+ holding a ring per distinct worker index and nothing can merge.
581
+ """
582
+
583
+ def _fleet(
584
+ self,
585
+ num_producers,
586
+ workers,
587
+ replicas,
588
+ owner_sets=None,
589
+ name_owner_class=None,
590
+ n_names=6,
591
+ ):
592
+ names = [f"g{i}" for i in range(n_names)]
593
+ return (
594
+ RdtRouter(
595
+ num_producers,
596
+ workers * replicas,
597
+ owner_sets,
598
+ name_owner_class,
599
+ names,
600
+ [1] * n_names,
601
+ workers_per_replica=workers,
602
+ ),
603
+ names,
604
+ )
605
+
606
+ def test_every_replica_of_a_worker_resolves_the_same_producer(self):
607
+ r, names = self._fleet(16, 8, 4)
608
+ for worker in range(8):
609
+ for name in names:
610
+ picks = {r.producer_for(worker + rep * 8, name) for rep in range(4)}
611
+ assert len(picks) == 1, (
612
+ f"worker {worker} split across producers {picks}"
613
+ )
614
+
615
+ def test_adding_a_deployment_does_not_move_the_first_one(self):
616
+ """The strongest form: an existing consumer's routes are byte-identical
617
+ before and after another deployment is provisioned."""
618
+ one, names = self._fleet(16, 8, 1)
619
+ two, _ = self._fleet(16, 8, 2)
620
+ for consumer_id in range(8):
621
+ assert [one.producer_for(consumer_id, n) for n in names] == [
622
+ two.producer_for(consumer_id, n) for n in names
623
+ ]
624
+
625
+ def test_one_deployment_is_the_historical_carve(self):
626
+ """Default (no width) and width == the whole fleet must agree with the
627
+ plain block rule, so a single-deployment fleet is untouched."""
628
+ plain = _router(16, 8, group_lens=[1] * 6)
629
+ explicit, names = self._fleet(16, 8, 1)
630
+ for consumer_id in range(8):
631
+ block = assign_producer_indices(16, 8, consumer_id)
632
+ for name in names:
633
+ assert plain.producer_for(consumer_id, name) == explicit.producer_for(
634
+ consumer_id, name
635
+ )
636
+ assert explicit.producer_for(consumer_id, name) in block
637
+
638
+ def test_the_overlay_holds_for_partially_owned_names(self):
639
+ """Pipeline stages and experts too: two owner sets, one narrow. The
640
+ narrow one already forced agreement (one owner, no block to carve); the
641
+ wide one is what the carve decides."""
642
+ r, names = self._fleet(
643
+ 8,
644
+ 4,
645
+ 2,
646
+ owner_sets=[[0, 1, 2, 3], [5]],
647
+ name_owner_class=[0, 1, 0, 1, 0, 1],
648
+ )
649
+ for worker in range(4):
650
+ for name in names:
651
+ assert r.producer_for(worker, name) == r.producer_for(worker + 4, name)
652
+ assert {r.producer_for(c, "g1") for c in range(8)} == {5}
653
+
654
+ def test_a_width_that_does_not_divide_the_fleet_raises(self):
655
+ """A uniform fleet is asserted upstream in ``get_world_size``; if it ever
656
+ is not, the overlay would map two workers of one deployment onto the same
657
+ block index, so refuse rather than serve the wrong bytes."""
658
+ with pytest.raises(ValueError, match="does not divide"):
659
+ RdtRouter(4, 6, None, None, ["a"], [1], workers_per_replica=4)
660
+
661
+ def test_pulls_still_carry_the_global_consumer_id(self):
662
+ """The carve uses the index WITHIN a deployment, but the wire must keep
663
+ the fleet-global id: it is what lets the producer tell the sharers of a
664
+ slot apart, and count their arrivals separately."""
665
+
666
+ class _Method:
667
+ def __init__(self):
668
+ self.calls = []
669
+
670
+ def remote(self, keys, consumer_id, seq):
671
+ self.calls.append((keys, consumer_id, seq))
672
+
673
+ r, _ = self._fleet(4, 2, 2)
674
+ methods = [_Method() for _ in range(4)]
675
+ r.bind([object()] * 4, methods, consumer_id=3) # replica 1, worker 1
676
+ r.pull(2, ["k"], 7)
677
+
678
+ assert methods[2].calls == [(["k"], 3, 7)]
679
+
680
+
681
+# ---------------------------------------------------------------------------
682
+# Packed layout: the cross-process invariant
683
+# ---------------------------------------------------------------------------
684
+
685
+
686
+def _producer_pack_offsets(slices):
687
+ """The producer's rule, transcribed from ``rdt_produce_weights_batched``
688
+ (sharded_rdt_trainer.py): 16B-aligned offsets in specs order.
689
+
690
+ Kept as an independent implementation on purpose: the consumer computing the
691
+ same offsets is the invariant that makes the packed blob readable, and this
692
+ is the only thing that catches the two rules drifting apart. A divergence is
693
+ silent at runtime — the bytes still arrive, they are just carved up wrong.
694
+ """
695
+ pack_cur = 0
696
+ offsets = []
697
+ for numel, dtype in slices:
698
+ off = (pack_cur + 15) & ~15
699
+ pack_cur = off + numel * dtype.itemsize
700
+ offsets.append(off)
701
+ return offsets, pack_cur
702
+
703
+
704
+class TestPackedLayout:
705
+ def test_consumer_layout_matches_the_producer_rule(self):
706
+ """Mixed dtypes, sizes that do not land on 16B — i.e. the Kimi-style
707
+ group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment."""
708
+ specs = [
709
+ ("w.fp8", "float8_e4m3fn", [17]),
710
+ ("w.scale", "float32", [3]),
711
+ ("w.norm", "bfloat16", [5]),
712
+ ("w.big", "bfloat16", [1024]),
713
+ ]
714
+ baked, name_meta = {}, {}
715
+ layer = _FakeLayer("mixed")
716
+ copies = []
717
+ for name, dtype_name, shape in specs:
718
+ name_meta[name] = (dtype_name, shape)
719
+ copies.append(
720
+ _copy(name, shape=tuple(shape), dtype=_dtype_from_name(dtype_name))
721
+ )
722
+ mod_plan = _module(layer, copies)
723
+ for name, _d, _s in specs:
724
+ baked[name] = mod_plan
725
+
726
+ eng = _planner(baked, name_meta=name_meta)
727
+ plan = eng._build_call_plan([n for n, _, _ in specs], [len(specs)])
728
+ (chunk,) = plan.chunks
729
+
730
+ want_offsets, want_bytes = _producer_pack_offsets(
731
+ [(s[0], _dtype_from_name(d)) for _n, d, s in specs]
732
+ )
733
+ assert [off for off, _dt, _n, _shape in chunk.pack_layout] == want_offsets
734
+ assert chunk.pack_bytes == want_bytes
735
+
736
+ def test_every_offset_is_16b_aligned(self):
737
+ specs = [(f"w{i}", "bfloat16", [i * 3 + 1]) for i in range(6)]
738
+ layer = _FakeLayer("l")
739
+ copies = [_copy(n, shape=tuple(s)) for n, _d, s in specs]
740
+ mod_plan = _module(layer, copies)
741
+ eng = _planner(
742
+ {n: mod_plan for n, _d, _s in specs},
743
+ name_meta={n: (d, s) for n, d, s in specs},
744
+ )
745
+ (chunk,) = eng._build_call_plan([n for n, _, _ in specs], [len(specs)]).chunks
746
+ assert all(off % 16 == 0 for off, _dt, _n, _s in chunk.pack_layout)
747
+
748
+ def test_layout_carries_dtype_numel_and_shape_for_the_view_carve(self):
749
+ baked, name_meta, names, group_lens = _one_module_per_layer(1, numel=6)
750
+ eng = _planner(baked, name_meta=name_meta)
751
+ chunk = eng._build_call_plan(names, group_lens).chunks[0]
752
+ (off, dt, numel, shape) = chunk.pack_layout[0]
753
+ assert (off, dt, numel, shape) == (0, torch.bfloat16, 6, (6,))
754
+
755
+ def test_duplicate_source_keys_are_packed_once(self):
756
+ """Two destinations fed by the SAME (name, chain) pull one slice."""
757
+ layer = _FakeLayer("l")
758
+ shared = ("w", ())
759
+ copies = [
760
+ _copy("w", "a", shape=(4,)),
761
+ _copy("w", "b", shape=(4,)),
762
+ ]
763
+ assert copies[0].src == shared and copies[1].src == shared
764
+ eng = _planner(
765
+ {"w": _module(layer, copies)},
766
+ name_meta={"w": ("bfloat16", [4])},
767
+ )
768
+ (chunk,) = eng._build_call_plan(["w"], [1]).chunks
769
+ assert len(chunk.scatters) == 2
770
+ assert chunk.keys == [shared]
771
+ assert len(chunk.pack_layout) == 1
772
+
773
+
774
+# ---------------------------------------------------------------------------
775
+# Chunk splitting
776
+# ---------------------------------------------------------------------------
777
+
778
+
779
+class TestChunkModuleScatters:
780
+ """The chunk cut: one chunk per distinct owner class present in the copies,
781
+ ascending by class index. Derived purely from the bake plus the ownership
782
+ table, so vLLM's expert placement needs no cases."""
783
+
784
+ def test_unstamped_copies_form_one_chunk(self):
785
+ layer = _FakeLayer("l")
786
+ copies = [_copy(f"w{i}", shape=(4,)) for i in range(5)]
787
+ eng = _planner({}, name_meta={f"w{i}": ("bfloat16", [4]) for i in range(5)})
788
+ chunks = eng._chunk_module_scatters([_module(layer, copies)])
789
+ assert [ci for ci, _s in chunks] == [0]
790
+ assert len(chunks[0][1]) == 5
791
+
792
+ def test_one_chunk_per_owner_class_ascending(self):
793
+ """Classes are numbered by first appearance in metadata order, so the
794
+ replicated class (0) leads and the rest follow — the ordering the old
795
+ "-1 first, then ascending stamp" rule produced, now falling out of the
796
+ table instead of a special case."""
797
+ layer = _FakeLayer("moe")
798
+ names = ["norm", "e5", "e0", "e2"]
799
+ copies = [_copy(n, shape=(4,)) for n in names]
800
+ eng = _planner(
801
+ {},
802
+ name_meta={n: ("bfloat16", [4]) for n in names},
803
+ held_by={"e0": [0], "e2": [1], "e5": [2]},
804
+ )
805
+ chunks = eng._chunk_module_scatters([_module(layer, copies)])
806
+ assert [ci for ci, _s in chunks] == [0, 1, 2, 3]
807
+ assert [[s.src[0] for s in sc] for _ci, sc in chunks] == [
808
+ ["norm"],
809
+ ["e5"],
810
+ ["e0"],
811
+ ["e2"],
812
+ ]
813
+
814
+ def test_copy_order_within_a_chunk_is_bake_order(self):
815
+ """Scatter order within a chunk must follow the bake, not the stamp
816
+ sort: the packed layout and the producer's replay agree on keys order."""
817
+ layer = _FakeLayer("moe")
818
+ names = ["e3", "e1", "e2"] # one stamp, deliberately unsorted names
819
+ copies = [_copy(n, shape=(4,)) for n in names]
820
+ eng = _planner(
821
+ {},
822
+ name_meta={n: ("bfloat16", [4]) for n in names},
823
+ held_by={n: [0] for n in names},
824
+ )
825
+ ((_ci, scatters),) = eng._chunk_module_scatters([_module(layer, copies)])
826
+ assert [s.src[0] for s in scatters] == names
827
+
828
+ def test_stamps_interleave_across_modules_without_reordering_within(self):
829
+ """Two modules' copies bucket by stamp; within each bucket the order is
830
+ module-then-bake order."""
831
+ a, b = _FakeLayer("a"), _FakeLayer("b")
832
+ mods = [
833
+ _module(a, [_copy("a0", shape=(4,)), _copy("a1", shape=(4,))]),
834
+ _module(b, [_copy("b0", shape=(4,)), _copy("b1", shape=(4,))]),
835
+ ]
836
+ eng = _planner(
837
+ {},
838
+ name_meta={n: ("bfloat16", [4]) for n in ("a0", "a1", "b0", "b1")},
839
+ held_by={"a0": [0], "a1": [1], "b0": [0], "b1": [1]},
840
+ )
841
+ chunks = eng._chunk_module_scatters(mods)
842
+ assert [(ci, [s.src[0] for s in sc]) for ci, sc in chunks] == [
843
+ (1, ["a0", "b0"]),
844
+ (2, ["a1", "b1"]),
845
+ ]
846
+
847
+ def test_scatters_carry_their_own_dtype(self):
848
+ """dtype rides the record from the bake, where it is the fake's dtype
849
+ AFTER its op chain — not a plan-time lookup of the source name, which
850
+ would be wrong for any chain that reinterprets dtype."""
851
+ layer = _FakeLayer("l")
852
+ eng = _planner({}, name_meta={"w": ("float32", [10])})
853
+ ((_er, scatters),) = eng._chunk_module_scatters(
854
+ [_module(layer, [_copy("w", shape=(10,), dtype=torch.float32)])]
855
+ )
856
+ (sc,) = scatters
857
+ assert sc.dtype is torch.float32
858
+ assert sc.layer is layer
859
+
860
+
861
+# ---------------------------------------------------------------------------
862
+# _build_call_plan
863
+# ---------------------------------------------------------------------------
864
+
865
+
866
+class TestBuildCallPlan:
867
+ def test_one_chunk_per_group_when_unstamped(self):
868
+ baked, name_meta, names, group_lens = _one_module_per_layer(3)
869
+ eng = _planner(baked, name_meta=name_meta)
870
+ plan = eng._build_call_plan(names, group_lens)
871
+ assert len(plan.chunks) == len(group_lens) == 5
872
+ assert plan.pre_free == []
873
+
874
+ def test_stamps_multiply_chunks_within_a_group(self):
875
+ layer = _FakeLayer("l")
876
+ names = [f"w{i}" for i in range(4)]
877
+ copies = [_copy(n, shape=(4,)) for n in names]
878
+ mod_plan = _module(layer, copies)
879
+ eng = _planner(
880
+ {n: mod_plan for n in names},
881
+ name_meta={n: ("bfloat16", [4]) for n in names},
882
+ held_by={"w2": [0], "w3": [1]},
883
+ )
884
+ plan = eng._build_call_plan(names, [4])
885
+ assert len(plan.chunks) == 3 # replicated, w2's class, w3's class
886
+
887
+ def test_materialize_fires_on_a_modules_first_chunk_only(self):
888
+ """Empty HF params are allocated once per module, by construction —
889
+ including a FusedMoE-like module whose copies span owner-class chunks."""
890
+ layer = _FakeLayer("spanning")
891
+ names = [f"w{i}" for i in range(4)]
892
+ copies = [_copy(n, shape=(4,)) for n in names]
893
+ mod_plan = _module(layer, copies)
894
+ eng = _planner(
895
+ {n: mod_plan for n in names},
896
+ name_meta={n: ("bfloat16", [4]) for n in names},
897
+ held_by={"w0": [0], "w1": [0], "w2": [1], "w3": [1]},
898
+ )
899
+ plan = eng._build_call_plan(names, [4])
900
+ assert [c.materialize for c in plan.chunks] == [[layer], []]
901
+
902
+ def test_quant_defers_to_a_modules_last_chunk(self):
903
+ layer = _FakeLayer("spanning")
904
+ names = [f"w{i}" for i in range(4)]
905
+ copies = [_copy(n, shape=(4,)) for n in names]
906
+ mod_plan = _module(layer, copies)
907
+ eng = _planner(
908
+ {n: mod_plan for n in names},
909
+ name_meta={n: ("bfloat16", [4]) for n in names},
910
+ held_by={"w0": [0], "w1": [0], "w2": [1], "w3": [1]},
911
+ )
912
+ plan = eng._build_call_plan(names, [4])
913
+ assert [c.quant for c in plan.chunks] == [[], [layer]]
914
+
915
+ def test_free_signal_fires_on_each_groups_last_chunk(self):
916
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
917
+ eng = _planner(baked, name_meta=name_meta)
918
+ plan = eng._build_call_plan(names, group_lens)
919
+ assert [c.free for c in plan.chunks] == [[0], [1], [2], [3]]
920
+
921
+ def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self):
922
+ """The load-bearing case: with the group cut across owner-class chunks the
923
+ signal must hang off the LAST one. Signaling earlier lets the producers
924
+ drop the gather buffers while a later chunk's RDMA is still reading."""
925
+ layer = _FakeLayer("l")
926
+ names = [f"w{i}" for i in range(4)]
927
+ copies = [_copy(n, shape=(4,)) for n in names]
928
+ mod_plan = _module(layer, copies)
929
+ eng = _planner(
930
+ {n: mod_plan for n in names},
931
+ name_meta={n: ("bfloat16", [4]) for n in names},
932
+ held_by={"w2": [0], "w3": [1]},
933
+ )
934
+ plan = eng._build_call_plan(names, [4])
935
+ assert len(plan.chunks) == 3
936
+ assert [c.free for c in plan.chunks] == [[], [], [0]]
937
+
938
+ def test_free_signal_timing_across_two_split_groups(self):
939
+ """Two groups x 2 owner-class chunks each: signals land on each group's
940
+ last chunk only."""
941
+ names_a = [f"model.layers.0.w{i}" for i in range(4)]
942
+ names_b = [f"model.layers.1.w{i}" for i in range(4)]
943
+ baked, name_meta, stamps = {}, {}, {}
944
+ for names in (names_a, names_b):
945
+ layer = _FakeLayer(names[0])
946
+ mod_plan = _module(layer, [_copy(n, shape=(4,)) for n in names])
947
+ for i, n in enumerate(names):
948
+ baked[n] = mod_plan
949
+ name_meta[n] = ("bfloat16", [4])
950
+ stamps[n] = [i // 2] # two owner classes per group
951
+ eng = _planner(baked, name_meta=name_meta, held_by=stamps)
952
+ plan = eng._build_call_plan(names_a + names_b, [4, 4])
953
+ assert [c.free for c in plan.chunks] == [[], [0], [], [1]]
954
+
955
+ def test_a_group_with_nothing_local_is_signaled_at_sync_start(self):
956
+ """The owners still published it, so this consumer must signal it —
957
+ immediately, before the pipeline (signal-before-publish is tolerated),
958
+ never hung off another group's chunk."""
959
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
960
+ del baked["model.layers.1.w"] # nothing baked, and not live => no pull
961
+ eng = _planner(
962
+ baked, name_meta=name_meta, live=set(names) - {"model.layers.1.w"}
963
+ )
964
+ plan = eng._build_call_plan(names, group_lens)
965
+ assert len(plan.chunks) == 3
966
+ assert plan.pre_free == [2]
967
+ assert [c.free for c in plan.chunks] == [[0], [1], [3]]
968
+
969
+ def test_leading_groups_with_nothing_local_go_to_pre_free(self):
970
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
971
+ del baked["embed.weight"]
972
+ eng = _planner(baked, name_meta=name_meta, live=set(names) - {"embed.weight"})
973
+ plan = eng._build_call_plan(names, group_lens)
974
+ assert plan.pre_free == [0]
975
+ assert len(plan.chunks) == 3
976
+
977
+ def test_modules_are_deduped_within_a_group(self):
978
+ """Several names of one fused module map to the same scatter list; the
979
+ plan must not chunk it twice."""
980
+ layer = _FakeLayer("fused")
981
+ names = ["qkv.q", "qkv.k", "qkv.v"]
982
+ copies = [_copy(n, shape=(4,)) for n in names]
983
+ mod_plan = _module(layer, copies)
984
+ eng = _planner(
985
+ {n: mod_plan for n in names},
986
+ name_meta={n: ("bfloat16", [4]) for n in names},
987
+ )
988
+ plan = eng._build_call_plan(names, [3])
989
+ (chunk,) = plan.chunks
990
+ assert chunk.materialize == [layer]
991
+ assert len(chunk.scatters) == 3
992
+
993
+ def test_plan_is_a_pure_function_of_its_inputs(self):
994
+ """Called twice with the same arguments it must produce the same plan —
995
+ that is what lets the engine build it once and reuse it every sync."""
996
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
997
+ eng = _planner(baked, name_meta=name_meta)
998
+ a = eng._build_call_plan(names, group_lens)
999
+ b = eng._build_call_plan(names, group_lens)
1000
+ assert [c.pack_bytes for c in a.chunks] == [c.pack_bytes for c in b.chunks]
1001
+ assert [c.keys for c in a.chunks] == [c.keys for c in b.chunks]
1002
+ assert [c.free for c in a.chunks] == [c.free for c in b.chunks]
1003
+ assert a.pre_free == b.pre_free
1004
+
1005
+ def test_single_producer_routes_every_chunk_to_local_owner_zero(self):
1006
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
1007
+ eng = _planner(baked, name_meta=name_meta)
1008
+ plan = eng._build_call_plan(names, group_lens)
1009
+ assert [c.owner for c in plan.chunks] == [0] * len(plan.chunks)
1010
+
1011
+
1012
+class TestUnbakedNamesGuard:
1013
+ """A live name with no baked plan has no way to load — there is no fallback,
1014
+ and its pull would target groups the pipeline already freed — so the plan
1015
+ build fails at init, naming the names."""
1016
+
1017
+ def test_a_live_unbaked_name_fails_the_plan_build(self):
1018
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
1019
+ del baked["norm.weight"] # unbaked but live
1020
+ eng = _planner(baked, name_meta=name_meta, live=set(names))
1021
+ with pytest.raises(RuntimeError, match="norm.weight"):
1022
+ eng._build_call_plan(names, group_lens)
1023
+
1024
+ def test_never_copied_names_are_dropped_entirely(self):
1025
+ """Experts owned by another EP rank no-op in their loader: not an error,
1026
+ and not a chunk — their group is pre-freed if nothing else fills it."""
1027
+ baked, name_meta, names, group_lens = _one_module_per_layer(2)
1028
+ del baked["norm.weight"]
1029
+ eng = _planner(baked, name_meta=name_meta, live=set(names) - {"norm.weight"})
1030
+ plan = eng._build_call_plan(names, group_lens)
1031
+ assert len(names) - 1 == len(plan.chunks)
1032
+ assert plan.pre_free == [len(group_lens) - 1]
1033
+
1034
+
1035
+class TestCallPlanRouting:
1036
+ """Under partial ownership each chunk must be pulled from a producer that
1037
+ actually holds every name in it. Consumers bind EVERY producer (signals fan
1038
+ out to all owners), so local index == trainer rank."""
1039
+
1040
+ def _routed_planner(
1041
+ self, held_by_group, num_consumers, consumer_id, *, held_by=None
1042
+ ):
1043
+ """``held_by_group[g]`` = the ranks holding group g's names, so a test
1044
+ can still speak in groups; ``held_by`` overrides individual names."""
1045
+ baked, name_meta, names, group_lens = _one_module_per_layer(
1046
+ len(held_by_group) - 2
1047
+ )
1048
+ per_name = {n: list(held_by_group[gi]) for gi, n in enumerate(names)}
1049
+ per_name.update(held_by or {})
1050
+ n_prod = max(max(o) for o in per_name.values()) + 1
1051
+ owner_sets: list[list[int]] = []
1052
+ seen: dict[tuple, int] = {}
1053
+ classes = []
1054
+ for n in names:
1055
+ key = tuple(sorted(per_name[n]))
1056
+ if key not in seen:
1057
+ seen[key] = len(owner_sets)
1058
+ owner_sets.append(list(key))
1059
+ classes.append(seen[key])
1060
+ eng = _planner(baked, name_meta=name_meta)
1061
+ router = RdtRouter(
1062
+ n_prod, num_consumers, owner_sets, classes, names, group_lens
1063
+ )
1064
+ router.validate()
1065
+ router.bind([object()] * n_prod, [object()] * n_prod, consumer_id)
1066
+ eng._router = router
1067
+ return eng, names, group_lens
1068
+
1069
+ def test_two_stage_ownership_splits_chunks_between_owners(self):
1070
+ """Two PP stages, 4 groups: stage 0 holds the first half, stage 1 the
1071
+ second. Every chunk goes to a holder of its names."""
1072
+ eng, names, group_lens = self._routed_planner([[0], [0], [1], [1]], 1, 0)
1073
+ plan = eng._build_call_plan(names, group_lens)
1074
+ assert [c.owner for c in plan.chunks] == [0, 0, 1, 1]
1075
+
1076
+ def test_a_multi_owner_group_still_resolves_to_one_producer(self):
1077
+ """Every chunk is served by exactly ONE producer: splitting a pull only
1078
+ multiplies produce calls, since the consumer's NIC bounds it."""
1079
+ eng, names, group_lens = self._routed_planner([[0, 1]] * 4, 1, 0)
1080
+ plan = eng._build_call_plan(names, group_lens)
1081
+ assert all(0 <= c.owner < 2 for c in plan.chunks)
1082
+
1083
+ def test_an_expert_name_routes_to_the_rank_holding_it(self):
1084
+ """PP ∩ EP: stage 0 = ranks 0-1, stage 1 = ranks 2-3, and one expert
1085
+ name inside a stage-1 group is held by rank 3 alone. The name's own owner
1086
+ set is what routes it."""
1087
+ eng, names, group_lens = self._routed_planner(
1088
+ [[0, 1], [0, 1], [2, 3], [2, 3]],
1089
+ 1,
1090
+ 0,
1091
+ held_by={"model.layers.1.w": [3]},
1092
+ )
1093
+ plan = eng._build_call_plan(names, group_lens)
1094
+ assert plan.chunks[2].owner == 3
1095
+ assert plan.chunks[0].owner in (0, 1)
1096
+ assert plan.chunks[3].owner in (2, 3)
1097
+
1098
+ def test_an_unheld_name_is_rejected_at_init(self):
1099
+ """A name no rank holds can never be served. The trainer raises when it
1100
+ transposes the holdings; the router's validate() is the consumer-side
1101
+ backstop, and either way it fails at init rather than at first pull."""
1102
+ with pytest.raises(ValueError, match="empty"):
1103
+ RdtRouter(2, 1, [[]], [0], ["w"], [1]).validate()
1104
+
1105
+
1106
+class TestSignalCompleteness:
1107
+ """Every gather group is signaled exactly once — after its last chunk when
1108
+ the worker pulls from it, at sync start otherwise — for ANY expert
1109
+ placement. A group signaled twice over-credits the barrier; a group never
1110
+ signaled parks a producer credit and hangs end_sync."""
1111
+
1112
+ def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2):
1113
+ """pre / n_layers MoE layers / post. Each layer: one norm (its own
1114
+ module, stamp -1) + n_experts expert names stamped ``e // n_local``.
1115
+ The worker's bake covers only ``worker_experts`` (its placement);
1116
+ foreign expert names never copied => dropped entirely, as in the real
1117
+ bake."""
1118
+ n_local = n_experts // ep_size
1119
+ names = ["embed.weight"]
1120
+ group_lens = [1]
1121
+ baked = {
1122
+ "embed.weight": _module(
1123
+ _FakeLayer("embed"), [_copy("embed.weight", shape=(4,))]
1124
+ )
1125
+ }
1126
+ name_meta = {"embed.weight": ("bfloat16", [4])}
1127
+ stamps: dict = {}
1128
+ live = {"embed.weight"}
1129
+ for li in range(n_layers):
1130
+ norm = f"model.layers.{li}.norm"
1131
+ enames = [f"model.layers.{li}.experts.{e}.w" for e in range(n_experts)]
1132
+ names += [norm] + enames
1133
+ group_lens.append(1 + n_experts)
1134
+ name_meta[norm] = ("bfloat16", [4])
1135
+ live.add(norm)
1136
+ baked[norm] = _module(_FakeLayer(f"norm{li}"), [_copy(norm, shape=(4,))])
1137
+ fused = _module(
1138
+ _FakeLayer(f"moe{li}"),
1139
+ [
1140
+ _copy(f"model.layers.{li}.experts.{e}.w", shape=(4,))
1141
+ for e in worker_experts
1142
+ ],
1143
+ )
1144
+ for e in range(n_experts):
1145
+ n = f"model.layers.{li}.experts.{e}.w"
1146
+ name_meta[n] = ("bfloat16", [4])
1147
+ stamps[n] = [e // n_local] # held by that EP rank alone
1148
+ if e in worker_experts:
1149
+ baked[n] = fused
1150
+ live.add(n)
1151
+ names.append("lm_head.weight")
1152
+ group_lens.append(1)
1153
+ baked["lm_head.weight"] = _module(
1154
+ _FakeLayer("head"), [_copy("lm_head.weight", shape=(4,))]
1155
+ )
1156
+ name_meta["lm_head.weight"] = ("bfloat16", [4])
1157
+ live.add("lm_head.weight")
1158
+ eng = _planner(baked, name_meta=name_meta, live=live, held_by=stamps)
1159
+ return eng, names, group_lens
1160
+
1161
+ @staticmethod
1162
+ def _signals(plan):
1163
+ return sorted(list(plan.pre_free) + [gi for c in plan.chunks for gi in c.free])
1164
+
1165
+ def test_linear_placement_signals_every_group_exactly_once(self):
1166
+ eng, names, group_lens = self._moe_planner(worker_experts=[0, 1, 2, 3])
1167
+ plan = eng._build_call_plan(names, group_lens)
1168
+ assert self._signals(plan) == list(range(len(group_lens)))
1169
+
1170
+ def test_round_robin_placement_signals_every_group_exactly_once(self):
1171
+ eng, names, group_lens = self._moe_planner(worker_experts=[0, 2, 4, 6])
1172
+ plan = eng._build_call_plan(names, group_lens)
1173
+ assert self._signals(plan) == list(range(len(group_lens)))
1174
+
1175
+ def test_placement_changes_only_the_chunk_count_never_the_signals(self):
1176
+ """linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6
1177
+ hits all 4. More chunks per group, identical signal set."""
1178
+ plans = {}
1179
+ for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])):
1180
+ eng, names, group_lens = self._moe_planner(worker_experts=experts)
1181
+ plans[label] = eng._build_call_plan(names, group_lens)
1182
+ # per MoE group: the replicated chunk + one per owner class present
1183
+ assert (
1184
+ len(plans["linear"].chunks) == 2 + 2 * 3
1185
+ ) # embed+head + 2 layers x (1 + 2)
1186
+ assert len(plans["round_robin"].chunks) == 2 + 2 * 5 # 2 layers x (1 + 4)
1187
+ assert self._signals(plans["linear"]) == self._signals(plans["round_robin"])
1188
+
1189
+ def test_the_fused_module_materializes_first_and_quants_last_across_chunks(self):
1190
+ """The FusedMoE shape: one module's copies span every owner-class chunk of
1191
+ its group; materialize on its first, quant on its last."""
1192
+ eng, names, group_lens = self._moe_planner(
1193
+ worker_experts=[0, 2, 4, 6], n_layers=1
1194
+ )
1195
+ plan = eng._build_call_plan(names, group_lens)
1196
+ # group 1 = [norm chunk, coord 0..3 chunks] at chunk indices 1..5
1197
+ fused_chunks = [
1198
+ c
1199
+ for c in plan.chunks[1:6]
1200
+ if any("experts" in sc.src[0] for sc in c.scatters)
1201
+ ]
1202
+ assert len(fused_chunks) == 4
1203
+ (fused_layer,) = fused_chunks[0].materialize
1204
+ assert fused_layer.tag == "moe0"
1205
+ assert fused_chunks[0].quant == []
1206
+ assert fused_chunks[-1].quant == [fused_layer]
1207
+
1208
+
1209
+# ---------------------------------------------------------------------------
1210
+# Shared helpers
1211
+# ---------------------------------------------------------------------------
1212
+
1213
+
1214
+class TestBufferAllocBytes:
1215
+ def test_rounds_up_to_a_coarse_256mb_boundary(self):
1216
+ assert buffer_alloc_bytes(1) == 256 << 20
1217
+ assert buffer_alloc_bytes(256 << 20) == 256 << 20
1218
+ assert buffer_alloc_bytes((256 << 20) + 1) == 512 << 20
1219
+
1220
+ def test_presize_is_a_floor_not_a_cap(self):
1221
+ assert buffer_alloc_bytes(1, presize=3 << 30) == 3 << 30
1222
+ big = 4 << 30
1223
+ assert buffer_alloc_bytes(big, presize=1 << 30) == big
1224
+
1225
+ def test_never_returns_less_than_requested(self):
1226
+ for nbytes in (1, 1 << 20, (700 << 20) + 3):
1227
+ assert buffer_alloc_bytes(nbytes) >= nbytes
1228
+
1229
+
1230
+class TestLayerwiseGroups:
1231
+ def test_partition_is_one_group_per_layer_between_unindexed_blocks(self):
1232
+ names = [
1233
+ "embed.weight",
1234
+ "model.layers.0.a",
1235
+ "model.layers.0.b",
1236
+ "model.layers.1.a",
1237
+ "norm.weight",
1238
+ ]
1239
+ assert layerwise_groups(names) == [
1240
+ ["embed.weight"],
1241
+ ["model.layers.0.a", "model.layers.0.b"],
1242
+ ["model.layers.1.a"],
1243
+ ["norm.weight"],
1244
+ ]
1245
+
1246
+ def test_flattening_a_partition_returns_the_input_order(self):
1247
+ """Un-indexed names group by POSITION, not by name class — so one that
1248
+ looks like a pre-block name still lands after the layers it follows.
1249
+ Group index therefore means the same thing on every rank."""
1250
+ names = [
1251
+ "embed.weight",
1252
+ "model.layers.0.a",
1253
+ "embed.weight.tied", # after the layers => its own trailing block
1254
+ ]
1255
+ groups = layerwise_groups(names)
1256
+ assert [n for g in groups for n in g] == names
1257
+
1258
+ def test_layers_are_ordered_by_index_not_appearance(self):
1259
+ names = ["model.layers.10.a", "model.layers.2.a"]
1260
+ assert layerwise_groups(names) == [
1261
+ ["model.layers.2.a"],
1262
+ ["model.layers.10.a"],
1263
+ ]
1264
+
1265
+ def test_stacks_keep_prefix_appearance_order_and_sort_within(self):
1266
+ """Two stacks stay in the order they appear — the vision tower before the
1267
+ text stack — while each sorts by index internally."""
1268
+ names = [
1269
+ "visual.blocks.1.w",
1270
+ "visual.blocks.0.w",
1271
+ "model.language_model.layers.1.w",
1272
+ "model.language_model.layers.0.w",
1273
+ ]
1274
+ assert layerwise_groups(names) == [
1275
+ ["visual.blocks.0.w"],
1276
+ ["visual.blocks.1.w"],
1277
+ ["model.language_model.layers.0.w"],
1278
+ ["model.language_model.layers.1.w"],
1279
+ ]
1280
+
1281
+ def test_the_post_block_lands_last_however_early_it_arrives(self):
1282
+ """Megatron-Bridge streams the last pipeline stage's output block before
1283
+ its layers. Sweeping un-indexed names after the first layer into a
1284
+ trailing group is what lets the gather loop walk groups ascending."""
1285
+ names = ["model.layers.0.w", "model.norm.weight", "model.layers.1.w"]
1286
+ assert layerwise_groups(names) == [
1287
+ ["model.layers.0.w"],
1288
+ ["model.layers.1.w"],
1289
+ ["model.norm.weight"],
1290
+ ]
1291
+
1292
+ def test_names_sharing_a_layer_coalesce_when_the_source_interleaves(self):
1293
+ """A raw-checkpoint source may yield in shard-packing order. Coalescing on
1294
+ the key keeps that from shattering one layer into many groups."""
1295
+ names = [
1296
+ "model.layers.0.a",
1297
+ "model.layers.1.a",
1298
+ "model.layers.0.b",
1299
+ ]
1300
+ assert layerwise_groups(names) == [
1301
+ ["model.layers.0.a", "model.layers.0.b"],
1302
+ ["model.layers.1.a"],
1303
+ ]
1304
+
1305
+ def test_no_empty_groups(self):
1306
+ for names in ([], ["only.pre"], ["model.layers.0.a"]):
1307
+ assert all(g for g in layerwise_groups(names))
1308
+
1309
+ def test_a_model_with_no_layers_is_one_pre_group(self):
1310
+ assert layerwise_groups(["a", "b"]) == [["a", "b"]]
1311
+
1312
+ @pytest.mark.parametrize(
1313
+ "prefix",
1314
+ [
1315
+ "model.layers.", # Llama, Qwen3, GLM-4.5, DeepSeek, Mixtral
1316
+ "model.language_model.layers.", # Qwen3.5/3.6, Qwen2.5-VL, Gemma3
1317
+ "language_model.layers.", # Kimi-K2.5, Nemotron-VL
1318
+ "thinker.model.layers.", # Qwen2.5/3-Omni
1319
+ "transformer.h.", # GPT-2, GPT-J, Bloom, Falcon
1320
+ "transformer.blocks.", # MPT, DBRX
1321
+ "transformer.encoder.layers.", # ChatGLM
1322
+ "backbone.layers.", # Mamba, Mamba2
1323
+ "gpt_neox.layers.", # GPT-NeoX
1324
+ "model.decoder.layers.", # OPT
1325
+ "model.h.", # Exaone
1326
+ "model.layers.layers.", # Plamo3
1327
+ "bert.encoder.layer.", # BERT, RoBERTa
1328
+ "visual.blocks.", # Qwen-VL vision tower
1329
+ "vision_tower.vision_model.encoder.layers.", # SigLIP/CLIP towers
1330
+ "audio_tower.layers.", # Whisper-style audio towers
1331
+ ],
1332
+ )
1333
+ def test_every_supported_naming_convention_partitions_per_layer(self, prefix):
1334
+ """One decoder layer per group under any prefix, with no per-architecture
1335
+ table. A literal `model.layers.` match silently yields one whole-model
1336
+ group for most of these."""
1337
+ names = [f"{prefix}{i}.self_attn.q_proj.weight" for i in range(3)]
1338
+ assert layerwise_groups(names) == [[n] for n in names]
1339
+
1340
+ def test_a_moe_layer_stays_whole(self):
1341
+ """The OUTERMOST index wins, so per-expert names group by layer. Splitting
1342
+ on the expert would cut FusedMoE — one leaf module — across groups."""
1343
+ names = [
1344
+ "model.layers.0.mlp.experts.0.w1.weight",
1345
+ "model.layers.0.mlp.experts.1.w1.weight",
1346
+ "model.layers.1.mlp.experts.0.w1.weight",
1347
+ ]
1348
+ assert layerwise_groups(names) == [
1349
+ [
1350
+ "model.layers.0.mlp.experts.0.w1.weight",
1351
+ "model.layers.0.mlp.experts.1.w1.weight",
1352
+ ],
1353
+ ["model.layers.1.mlp.experts.0.w1.weight"],
1354
+ ]
1355
+
1356
+ def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self):
1357
+ """Qwen-VL layout. Today's literal `model.layers.` match puts every one of
1358
+ these in ONE group — the whole model — because the text stack ships as
1359
+ `model.language_model.layers.`."""
1360
+ names = [
1361
+ "visual.patch_embed.proj.weight",
1362
+ "visual.blocks.0.attn.qkv.weight",
1363
+ "visual.blocks.1.attn.qkv.weight",
1364
+ "visual.merger.mlp.weight",
1365
+ "model.language_model.layers.0.self_attn.q_proj.weight",
1366
+ "lm_head.weight",
1367
+ ]
1368
+ assert layerwise_groups(names) == [
1369
+ ["visual.patch_embed.proj.weight"],
1370
+ ["visual.blocks.0.attn.qkv.weight"],
1371
+ ["visual.blocks.1.attn.qkv.weight"],
1372
+ ["model.language_model.layers.0.self_attn.q_proj.weight"],
1373
+ ["visual.merger.mlp.weight", "lm_head.weight"],
1374
+ ]
1375
+
1376
+ def test_two_stacks_sharing_an_index_do_not_merge(self):
1377
+ """The key carries the prefix, so vision block 0 and text layer 0 are
1378
+ different groups even though both are index 0."""
1379
+ names = [
1380
+ "visual.blocks.0.attn.qkv.weight",
1381
+ "model.language_model.layers.0.self_attn.q_proj.weight",
1382
+ ]
1383
+ assert layerwise_groups(names) == [[n] for n in names]
1384
+
1385
+
1386
+class TestDtypeFromName:
1387
+ def test_resolves_wire_names(self):
1388
+ assert _dtype_from_name("bfloat16") is torch.bfloat16
1389
+ assert _dtype_from_name("float8_e4m3fn") is torch.float8_e4m3fn
1390
+
1391
+ def test_rejects_a_non_dtype_attribute(self):
1392
+ with pytest.raises(ValueError):
1393
+ _dtype_from_name("nn") # torch.nn exists but is not a dtype
1394
+
1395
+ def test_rejects_an_unknown_name(self):
1396
+ with pytest.raises(ValueError):
1397
+ _dtype_from_name("float9")
1398
+
1399
+
1400
+class TestOpAllowlistAgreement:
1401
+ """The consumer records op chains from ``SUPPORTED_OPS``; the producer
1402
+ replays a chain only if every op is in ``ALLOWED_OPS``. The two must describe
1403
+ the same contract, so they are derived from one table.
1404
+
1405
+ A divergence is silent until first pull: an op the consumer can emit but the
1406
+ producer rejects bakes fine at init, then fails a whole sync later with
1407
+ "disallowed op".
1408
+ """
1409
+
1410
+ def test_the_two_allowlists_describe_the_same_contract(self):
1411
+ assert set(SUPPORTED_OPS.values()) == set(ALLOWED_OPS)
1412
+
1413
+ def test_every_bake_recordable_op_is_replayable(self):
1414
+ """Anything the bake can record must be serveable, or the failure lands a
1415
+ whole sync later than the mistake."""
1416
+ assert not set(SUPPORTED_OPS.values()) - set(ALLOWED_OPS)
1417
+
1418
+ def test_the_producer_allows_nothing_the_consumer_cannot_emit(self):
1419
+ """The allowlist is also a guard against a misbehaving or spoofed
1420
+ consumer invoking arbitrary methods on trainer tensors, so it must not be
1421
+ wider than the ops the bake can actually produce. ``to`` in particular
1422
+ would let a replay change dtype or device."""
1423
+ assert not set(ALLOWED_OPS) - set(SUPPORTED_OPS.values())
1424
+ assert "to" not in ALLOWED_OPS
1425
+
1426
+ def test_transpose_via_t_is_serveable(self):
1427
+ """The specific regression: a chain recorded from ``.t()`` must pass the
1428
+ producer's guard."""
1429
+ fake = FakeRDTTensor(
1430
+ name="w",
1431
+ shape=torch.Size((4, 6)),
1432
+ dtype=torch.bfloat16,
1433
+ device=META,
1434
+ sink=None,
1435
+ )
1436
+ (op, _args, _kwargs) = fake.t()._key()[1][0]
1437
+ assert op in ALLOWED_OPS
1438
+
1439
+
1440
+# ---------------------------------------------------------------------------
1441
+# Bake against a real model (GPU)
1442
+# ---------------------------------------------------------------------------
1443
+
1444
+TINY_QWEN = "Qwen/Qwen3-0.6B"
1445
+
1446
+# One small model per attention/FFN shape the bake has to survive. The point is
1447
+# the checkpoint->module mapping, which differs structurally between them: GQA
1448
+# fuses Q/K/V of unequal head counts into one param, MLA replaces them with the
1449
+# compressed kv_a/kv_b pair, and MoE turns per-expert checkpoint entries into
1450
+# stacked expert params.
1451
+#
1452
+# Keep them small: `load_format="dummy"` skips the download but not the
1453
+# allocation, so every parameter is still built on the device at full size.
1454
+BAKE_MODELS = [
1455
+ pytest.param(TINY_QWEN, id="gqa-dense"),
1456
+ pytest.param("TitanML/tiny-mixtral", id="moe"),
1457
+ # MLA (kv_lora_rank set, so `ModelConfig.use_mla`) plus routed experts. Only
1458
+ # the STRUCTURE is under test -- `_bake_probe` synthesizes the checkpoint
1459
+ # from the safetensors headers and never reads a weight value -- so a random
1460
+ # 0.06B model does the job, as long as it has the MLA head shape the prefill
1461
+ # backends accept: (qk_nope=128, qk_rope=64, v=128).
1462
+ #
1463
+ # Not an fp8 checkpoint: `process_weights_after_loading` swaps the vLLM
1464
+ # parameter subclasses out on those, and the plain second `load_weights`
1465
+ # this test compares against then cannot run.
1466
+ pytest.param("hmellor/tiny-random-DeepseekV2ForCausalLM", id="mla-moe"),
1467
+]
1468
+
1469
+
1470
+class TestRequiresTheRayExecutor:
1471
+ """The engine's data plane is Ray's, so a non-Ray executor cannot work. The
1472
+ check is at construction because the alternative is an opaque failure during
1473
+ the first handshake, long after the misconfiguration."""
1474
+
1475
+ def _construct(self, backend):
1476
+ cfg = SimpleNamespace(
1477
+ parallel_config=SimpleNamespace(distributed_executor_backend=backend)
1478
+ )
1479
+ eng = object.__new__(ShardedRDTWeightTransferEngine)
1480
+ # Only the guard is under test, so stand in for the base __init__.
1481
+ with pytest.MonkeyPatch.context() as m:
1482
+ m.setattr(
1483
+ ShardedRDTWeightTransferEngine.__bases__[0],
1484
+ "__init__",
1485
+ lambda self, *a, **k: None,
1486
+ )
1487
+ ShardedRDTWeightTransferEngine.__init__(eng, None, cfg, META, None)
1488
+ return eng
1489
+
1490
+ @pytest.mark.parametrize("backend", ["uni", "mp", "external_launcher"])
1491
+ def test_a_non_ray_executor_is_refused(self, backend):
1492
+ with pytest.raises(ValueError, match="requires distributed_executor"):
1493
+ self._construct(backend)
1494
+
1495
+ def test_ray_is_accepted(self):
1496
+ self._construct("ray")
1497
+
1498
+ def test_a_custom_executor_class_is_left_alone(self):
1499
+ """A `type[Executor]` override is deliberate and unjudgeable here, so it
1500
+ must not be rejected for merely not being the string "ray"."""
1501
+
1502
+ class _CustomExecutor:
1503
+ pass
1504
+
1505
+ self._construct(_CustomExecutor)
1506
+
1507
+
1508
+def _bake_probe(self, model_name): # noqa: PLR0915 (runs in the vLLM worker)
1509
+ """Drive `_bake` against the worker's real model and check the three
1510
+ properties. Returns a small summary; every assertion fires here, in the
1511
+ worker, so a failure surfaces as the RPC's exception.
1512
+
1513
+ Self-contained on purpose: this is pickled to the worker, so it imports what
1514
+ it needs rather than closing over this module's globals.
1515
+ """
1516
+ import torch
1517
+ from huggingface_hub import get_safetensors_metadata
1518
+
1519
+ from vllm.distributed.weight_transfer.sharded_rdt_common import ALLOWED_OPS
1520
+ from vllm.distributed.weight_transfer.sharded_rdt_engine import (
1521
+ ShardedRDTWeightTransferEngine,
1522
+ ShardedRDTWeightTransferInitInfo,
1523
+ )
1524
+ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
1525
+
1526
+ model = self.model_runner.model
1527
+ device = next(model.parameters()).device
1528
+
1529
+ # ---- the checkpoint, exactly as a trainer would describe and serve it ----
1530
+ # Only the safetensors HEADERS are fetched (a few KB); the values are
1531
+ # synthesized. Nothing here depends on real weights -- both the reference
1532
+ # load and the replay read this same dict, so the comparison is unaffected,
1533
+ # and it keeps a multi-GB download out of CI.
1534
+ st_dtype = {
1535
+ "F64": torch.float64,
1536
+ "F32": torch.float32,
1537
+ "F16": torch.float16,
1538
+ "BF16": torch.bfloat16,
1539
+ "F8_E4M3": torch.float8_e4m3fn,
1540
+ "F8_E5M2": torch.float8_e5m2,
1541
+ "I64": torch.int64,
1542
+ "I32": torch.int32,
1543
+ "I8": torch.int8,
1544
+ "U8": torch.uint8,
1545
+ "BOOL": torch.bool,
1546
+ }
1547
+ meta = get_safetensors_metadata(model_name)
1548
+ specs = {}
1549
+ for file_meta in meta.files_metadata.values():
1550
+ for name, t in file_meta.tensors.items():
1551
+ specs[name] = (st_dtype[t.dtype], tuple(t.shape))
1552
+
1553
+ gen = torch.Generator(device="cpu").manual_seed(0)
1554
+ ckpt: dict[str, torch.Tensor] = {}
1555
+ for name, (dt, shape) in specs.items():
1556
+ if dt.is_floating_point:
1557
+ # Small magnitudes so any narrowing cast stays finite.
1558
+ v = torch.randn(shape, generator=gen, dtype=torch.float32) * 0.1
1559
+ ckpt[name] = v.to(dtype=dt)
1560
+ else:
1561
+ ckpt[name] = torch.zeros(shape, dtype=dt)
1562
+ names = sorted(ckpt)
1563
+ dtype_names = [str(ckpt[n].dtype).split(".")[-1] for n in names]
1564
+ shapes = [list(ckpt[n].shape) for n in names]
1565
+
1566
+ def _feed():
1567
+ return iter([(n, ckpt[n]) for n in names])
1568
+
1569
+ # ---- reference: what the model's OWN loader produces from that checkpoint --
1570
+ # Poison first, so "which params did load_weights actually write?" is
1571
+ # answerable. Some params have no checkpoint source at all -- derived ones,
1572
+ # or ones process_weights_after_loading builds -- and those keep their NaNs
1573
+ # and are legitimately outside the plan. That set is what the plan must
1574
+ # cover, not every parameter of the model.
1575
+ def _poison():
1576
+ with torch.no_grad():
1577
+ for prm in model.parameters():
1578
+ if prm.is_floating_point():
1579
+ prm.fill_(float("nan"))
1580
+ else:
1581
+ prm.zero_()
1582
+
1583
+ _poison()
1584
+ model.load_weights(_feed())
1585
+ loaded = {
1586
+ n
1587
+ for n, prm in model.named_parameters()
1588
+ if not (prm.is_floating_point() and torch.isnan(prm).any())
1589
+ }
1590
+ # On the host: the weights themselves already fill the GPU.
1591
+ reference = {
1592
+ n: prm.detach().to("cpu", copy=True)
1593
+ for n, prm in model.named_parameters()
1594
+ if n in loaded
1595
+ }
1596
+
1597
+ # ---- pre-bake state, for the restoration check --------------------------
1598
+ param_name_of = {}
1599
+ for mod_name, mod in model.named_modules():
1600
+ for pname in get_layer_tensors(mod):
1601
+ param_name_of[(id(mod), pname)] = f"{mod_name}.{pname}".lstrip(".")
1602
+ loaders_before = {}
1603
+ for mod in model.modules():
1604
+ for pname, tensor in get_layer_tensors(mod).items():
1605
+ loaders_before[(id(mod), pname)] = getattr(tensor, "weight_loader", None)
1606
+
1607
+ # ---- bake ---------------------------------------------------------------
1608
+ eng = object.__new__(ShardedRDTWeightTransferEngine)
1609
+ eng.model = model
1610
+ eng.device = device
1611
+ eng._name_to_plan = {}
1612
+ eng._name_meta = {}
1613
+ eng._live_names = set()
1614
+ eng._bake(
1615
+ ShardedRDTWeightTransferInitInfo(
1616
+ names=names, dtype_names=dtype_names, shapes=shapes
1617
+ )
1618
+ )
1619
+
1620
+ # (1) the bake ran and produced a plan
1621
+ assert eng._name_to_plan, "bake produced no plan"
1622
+ assert eng._live_names, "bake recorded no live names"
1623
+ n_modules = len({id(v) for v in eng._name_to_plan.values()})
1624
+
1625
+ # (2) every intermediate change to model state was undone
1626
+ for mod in model.modules():
1627
+ for pname, tensor in get_layer_tensors(mod).items():
1628
+ key = (id(mod), pname)
1629
+ loader = getattr(tensor, "weight_loader", None)
1630
+ assert not hasattr(loader, "_rdt_stamp_inner"), (
1631
+ f"recording stamp left on {param_name_of.get(key, key)}"
1632
+ )
1633
+ assert loader is loaders_before[key], (
1634
+ f"weight_loader not restored on {param_name_of.get(key, key)}"
1635
+ )
1636
+ assert not tensor.is_meta, (
1637
+ f"{param_name_of.get(key, key)} left on meta after the dry run"
1638
+ )
1639
+ for pname, p in model.named_parameters():
1640
+ assert torch.equal(p.detach().cpu(), reference[pname]), f"bake mutated {pname}"
1641
+
1642
+ # (3) replaying the recorded op chains reproduces that same load
1643
+ scatters = {id(v): v for v in eng._name_to_plan.values()}
1644
+
1645
+ # No deferred-attention layer may be a plan target. The engine's quant
1646
+ # thread calls `info.reset()` on every module it processes, which makes
1647
+ # finalize_layerwise_reload SKIP that module -- and finalize is the only
1648
+ # thing that calls the LAYER's process_weights_after_loading, which is where
1649
+ # MLA derives W_UK_T/W_UV from kv_b_proj. An attention layer in the plan
1650
+ # would therefore leave those derived weights stale after a sync, silently.
1651
+ from vllm.model_executor.layers.attention import is_deferred_attention_layer
1652
+
1653
+ attn_in_plan = sorted(
1654
+ {
1655
+ type(c.layer).__name__
1656
+ for plan in scatters.values()
1657
+ for c in plan
1658
+ if is_deferred_attention_layer(c.layer)
1659
+ }
1660
+ )
1661
+ assert not attn_in_plan, (
1662
+ f"plan targets deferred-attention layer(s) {attn_in_plan}; the quant "
1663
+ f"thread would reset them and finalize would skip their "
1664
+ f"process_weights_after_loading, staling any derived weights"
1665
+ )
1666
+
1667
+ touched = set()
1668
+ _poison()
1669
+ with torch.no_grad():
1670
+ for plan in scatters.values():
1671
+ for c in plan:
1672
+ src_name, chain = c.src
1673
+ t = ckpt[src_name]
1674
+ for op, args, kw in chain:
1675
+ assert op in ALLOWED_OPS, f"disallowed op {op!r}"
1676
+ t = getattr(t, op)(*args, **dict(kw))
1677
+ dest = getattr(c.layer, c.param_name)
1678
+ dest.as_strided(c.shape, c.stride, c.offset).copy_(t)
1679
+ touched.add(param_name_of[(id(c.layer), c.param_name)])
1680
+
1681
+ assert loaded, "the reference load wrote nothing"
1682
+ # The plan owns everything with a checkpoint source. Anything else the model
1683
+ # ends up holding is DERIVED -- under MLA, `process_weights_after_loading`
1684
+ # splits kv_b_proj into W_UK_T/W_UV via `replace_parameter` -- and belongs to
1685
+ # finalize_layerwise_reload, which `finish_weight_update` runs after the
1686
+ # replay. So the uncovered set must contain only params owned by a module
1687
+ # that post-processes; a plain loadable weight going missing still fails.
1688
+ post_processed = {
1689
+ id(mod)
1690
+ for mod in model.modules()
1691
+ if hasattr(mod, "process_weights_after_loading")
1692
+ }
1693
+ modules_by_name = dict(model.named_modules())
1694
+ uncovered = sorted(loaded - touched)
1695
+ derived: list[str] = []
1696
+ unexplained: list[str] = []
1697
+ for n in uncovered:
1698
+ mod_name = n.rsplit(".", 1)[0]
1699
+ mod = modules_by_name.get(mod_name)
1700
+ is_derived = mod is not None and id(mod) in post_processed
1701
+ (derived if is_derived else unexplained).append(n)
1702
+ assert not unexplained, (
1703
+ f"the plan misses {len(unexplained)} param(s) load_weights wrote that "
1704
+ f"nothing derives: {unexplained[:5]}"
1705
+ )
1706
+ mismatched = [
1707
+ n
1708
+ for n, prm in model.named_parameters()
1709
+ if n in loaded and not torch.equal(prm.detach().cpu(), reference[n])
1710
+ ]
1711
+ assert not mismatched, f"replay differs from load_weights for: {mismatched[:5]}"
1712
+
1713
+ return {
1714
+ "names": len(names),
1715
+ "baked": len(eng._name_to_plan),
1716
+ "live": len(eng._live_names),
1717
+ "modules": n_modules,
1718
+ "loaded_params": len(loaded),
1719
+ "derived_not_in_plan": len(derived),
1720
+ "replayed_params": len(touched),
1721
+ "total_params": sum(1 for _ in model.named_parameters()),
1722
+ }
1723
+
1724
+
1725
+@pytest.mark.skipif(
1726
+ not torch.cuda.is_available(), reason="the bake drives a real model on GPU"
1727
+)
1728
+class TestBakeOnARealModel:
1729
+ """`_bake` against a real vLLM model, which the CPU tests above cannot reach:
1730
+ they drive the planner with hand-built scatters, so nothing there exercises
1731
+ the dry run through a real `load_weights` — the part that monkeypatches the
1732
+ model and has to put it back.
1733
+ """
1734
+
1735
+ @pytest.mark.parametrize("model", BAKE_MODELS)
1736
+ def test_bake_records_a_replayable_plan_and_restores_the_model(
1737
+ self, vllm_runner, monkeypatch, model
1738
+ ):
1739
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
1740
+ with vllm_runner(
1741
+ model, load_format="dummy", enforce_eager=True, max_model_len=1024
1742
+ ) as llm:
1743
+ (summary,) = llm.llm.collective_rpc(_bake_probe, args=(model,))
1744
+ print(f"bake summary [{model}]: {summary}")
1745
+ # Non-vacuity: the checks above are only worth something if the plan
1746
+ # covers the model. Fusion means far fewer leaf modules than names --
1747
+ # QKV and gate/up for GQA, the stacked experts for MoE.
1748
+ assert summary["baked"] >= 0.95 * summary["names"], summary
1749
+ assert summary["live"] >= 0.95 * summary["names"], summary
1750
+ assert 0 < summary["modules"] < summary["baked"], summary
1751
+ # The probe already asserts the plan covers every param load_weights
1752
+ # wrote; this just keeps that from being vacuously few.
1753
+ assert summary["loaded_params"] > 0.5 * summary["total_params"], summary
tests/distributed/test_sharded_rdt_producer.py
new
+1214
@@ -0,0 +1,1214 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Protocol tests for the sharded-RDT producer server (`_RDTProducerServer`).
4
+
5
+The publish -> serve -> free_group -> release lifecycle is the subsystem's
6
+concurrency core: a group whose barrier never completes stalls `end_sync`
7
+forever, and a group freed too early drops CUDA-IPC storage while a consumer is
8
+still reading it. Freeing is a PER-GROUP BARRIER: every live consumer signals
9
+`free_group(gi)` at every owner of the group exactly once per sync, and the
10
+group is released when the count reaches the live total handed to `begin_sync`
11
+— one uniform integer, no routed per-producer targets.
12
+
13
+These tests drive the real server. Everything here is CPU-clean: the IPC rebuild
14
+is stubbed (`rebuild_cuda_tensor`) and the serve path is only entered as far as
15
+the routing guard, which fires before any CUDA work.
16
+
17
+`rdt_produce_weights_batched`'s pack/copy half needs a real GPU buffer; it stays
18
+covered by the GPU tests in `test_sharded_rdt_trainer.py`.
19
+"""
20
+
21
+import contextlib
22
+import gc
23
+import threading
24
+import time
25
+from concurrent.futures import ThreadPoolExecutor
26
+from concurrent.futures import TimeoutError as FuturesTimeout
27
+
28
+import pytest
29
+import torch
30
+
31
+import vllm.distributed.weight_transfer.sharded_rdt_trainer as trainer_mod
32
+from vllm.distributed.weight_transfer.sharded_rdt_common import (
33
+ buffer_alloc_bytes,
34
+)
35
+from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
36
+ DEFAULT_GATHER_LOOKAHEAD,
37
+ ShardedRDTTrainerInitInfo,
38
+ ShardedRDTTrainerWeightTransferEngine,
39
+ _RDTProducerServer,
40
+)
41
+
42
+# Publish/free are keyed by GROUP INDEX; the names ride the publish entries.
43
+GI_A, GI_B = 0, 1
44
+GROUP_A = ("model.layers.0.w", "model.layers.0.b")
45
+GROUP_B = ("model.layers.1.w",)
46
+
47
+
48
+@pytest.fixture
49
+def server_factory(monkeypatch):
50
+ """Build real `_RDTProducerServer` instances without a GPU.
51
+
52
+ Three process-global side effects in `__init__` have to be neutralized:
53
+ accelerator discovery (no driver here), the Ray NIXL monkeypatch (fail-soft
54
+ already, but pointless), and `gc.freeze()` — which would otherwise leak a
55
+ frozen GC generation into the pytest process for the rest of the session.
56
+ """
57
+ monkeypatch.setattr(
58
+ torch.accelerator, "current_device_index", lambda *a, **kw: 0, raising=False
59
+ )
60
+ monkeypatch.setattr(gc, "freeze", lambda: None)
61
+
62
+ # The IPC rebuild: hand back a plain CPU uint8 storage of the requested
63
+ # size. publish_group then does the real .view(dtype) + as_strided on it, so
64
+ # the view reconstruction under test is genuine.
65
+ def _fake_rebuild(*args):
66
+ return torch.arange(args[0], dtype=torch.uint8) if args else torch.empty(0)
67
+
68
+ monkeypatch.setattr(trainer_mod, "rebuild_cuda_tensor", _fake_rebuild)
69
+
70
+ built = []
71
+
72
+ def _make(**kwargs):
73
+ kwargs.setdefault("num_rdt_buffers", 2)
74
+ kwargs.setdefault("buffer_presize_gb", 0.0)
75
+ kwargs.setdefault("gather_lookahead", 2)
76
+ server = _RDTProducerServer(**kwargs)
77
+ built.append(server)
78
+ return server
79
+
80
+ yield _make
81
+ for server in built:
82
+ server.shutdown()
83
+
84
+
85
+def _entries(names, *, nbytes=64, dtype="bfloat16"):
86
+ """The `(storages, views)` payload the engine's gather loop ships: one
87
+ storage export shared by every name, plus a per-name as_strided view spec.
88
+
89
+ `storages` values stand in for `reduce_tensor`'s args. Only the arity
90
+ matters here: `publish_group` overwrites index 6 with the server's own
91
+ device index, so the tuple must be at least 7 long. The stubbed rebuild
92
+ reads the byte count from index 0.
93
+ """
94
+ sid = 1
95
+ storages = {sid: (nbytes, None, None, None, None, None, -1, None)}
96
+ itemsize = getattr(torch, dtype).itemsize
97
+ per = (nbytes // itemsize) // max(1, len(names))
98
+ views = {name: (sid, dtype, [per], [1], i * per) for i, name in enumerate(names)}
99
+ return storages, views
100
+
101
+
102
+def _publish(server, gi, names, **kw):
103
+ return server.publish_group(gi, _entries(names, **kw))
104
+
105
+
106
+class TestPublishAndRebuild:
107
+ def test_publish_makes_every_name_serveable(self, server_factory):
108
+ server = server_factory()
109
+ server.begin_sync(1)
110
+ _publish(server, GI_A, GROUP_A)
111
+ assert set(server._cache) == set(GROUP_A)
112
+
113
+ def test_rebuilt_views_have_the_requested_dtype_and_geometry(self, server_factory):
114
+ server = server_factory()
115
+ server.begin_sync(1)
116
+ _publish(server, GI_A, GROUP_A, nbytes=64, dtype="bfloat16")
117
+ for name in GROUP_A:
118
+ t = server._cache[name]
119
+ assert t.dtype is torch.bfloat16
120
+ assert tuple(t.shape) == (16,) # 64 bytes / 2 names / 2 bytes
121
+
122
+ def test_views_of_one_storage_do_not_overlap(self, server_factory):
123
+ """One IPC export per storage, one as_strided view per name — the whole
124
+ point of the storage/view split. Overlapping views would serve the same
125
+ bytes for two different weights."""
126
+ server = server_factory()
127
+ server.begin_sync(1)
128
+ _publish(server, GI_A, GROUP_A, nbytes=64)
129
+ a, b = (server._cache[n] for n in GROUP_A)
130
+ assert a.storage_offset() != b.storage_offset()
131
+ assert a.data_ptr() != b.data_ptr()
132
+
133
+ def test_publish_records_the_groups_names_for_release(self, server_factory):
134
+ server = server_factory()
135
+ server.begin_sync(3)
136
+ _publish(server, GI_A, GROUP_A)
137
+ assert server._group_names[GI_A] == list(GROUP_A)
138
+
139
+ def test_publish_records_the_group_as_unfreed(self, server_factory):
140
+ server = server_factory()
141
+ server.begin_sync(1)
142
+ _publish(server, GI_A, GROUP_A)
143
+ assert server._inflight_groups == [GI_A]
144
+
145
+ def test_the_default_lookahead_is_one(self):
146
+ """1 is the sweet spot: group N+1 is gathered AND pullable while N is
147
+ being pulled, with resident memory at its floor of 2 groups — the bound
148
+ the larger-model runs size against."""
149
+ assert DEFAULT_GATHER_LOOKAHEAD == 1
150
+
151
+
152
+class TestFreeBarrier:
153
+ def test_single_consumer_signal_releases_the_group(self, server_factory):
154
+ server = server_factory()
155
+ server.begin_sync(1)
156
+ _publish(server, GI_A, GROUP_A)
157
+ server.free_group(GI_A)
158
+ assert server._inflight_groups == []
159
+ assert set(server._cache) == set()
160
+
161
+ def test_group_is_held_until_the_last_live_consumer_signals(self, server_factory):
162
+ """Every live consumer signals every owner; releasing on the first
163
+ signal would drop storage other consumers are still reading."""
164
+ server = server_factory()
165
+ server.begin_sync(3)
166
+ _publish(server, GI_A, GROUP_A)
167
+ held_after_each_signal = []
168
+ for _ in range(3):
169
+ server.free_group(GI_A)
170
+ held_after_each_signal.append(GI_A in server._inflight_groups)
171
+ assert held_after_each_signal == [True, True, False]
172
+
173
+ def test_intermediate_signals_do_not_release(self, server_factory):
174
+ server = server_factory()
175
+ server.begin_sync(3)
176
+ _publish(server, GI_A, GROUP_A)
177
+ server.free_group(GI_A)
178
+ assert server._inflight_groups == [GI_A]
179
+ assert set(server._cache) == set(GROUP_A)
180
+ server.free_group(GI_A)
181
+ assert server._inflight_groups == [GI_A]
182
+ server.free_group(GI_A)
183
+ assert server._inflight_groups == []
184
+
185
+ def test_a_signal_that_beats_its_publish_is_completed_at_publish(
186
+ self, server_factory
187
+ ):
188
+ """A consumer with nothing to pull from a group signals it at sync
189
+ start, which can precede the publish. The publish must then release the
190
+ group rather than wait for a signal that will never come again."""
191
+ server = server_factory()
192
+ server.begin_sync(1)
193
+ server.free_group(GI_A)
194
+ _publish(server, GI_A, GROUP_A)
195
+ assert server._inflight_groups == []
196
+ assert server._freed_pending == [GI_A]
197
+
198
+ def test_early_signals_from_every_consumer_complete_at_publish(
199
+ self, server_factory
200
+ ):
201
+ server = server_factory()
202
+ server.begin_sync(3)
203
+ for _ in range(3):
204
+ server.free_group(GI_A)
205
+ _publish(server, GI_A, GROUP_A)
206
+ assert server._freed_pending == [GI_A]
207
+
208
+ def test_freed_groups_are_reported_back_to_the_engine(self, server_factory):
209
+ """The engine holds the CUDA-IPC export refs and may only drop them once
210
+ the server says the group is gone — and it hears that ONLY through
211
+ ``wait_freed`` / ``end_sync``, never a publish return (a freed notice
212
+ riding an unharvested async publish while the engine blocks in
213
+ ``wait_freed`` would wedge the loop)."""
214
+ server = server_factory()
215
+ server.begin_sync(1)
216
+ _publish(server, GI_A, GROUP_A)
217
+ server.free_group(GI_A)
218
+ assert server.wait_freed() == [GI_A]
219
+ assert server._freed_pending == []
220
+
221
+ def test_a_signal_for_an_unpublished_group_releases_nothing(self, server_factory):
222
+ server = server_factory()
223
+ server.begin_sync(1)
224
+ _publish(server, GI_A, GROUP_A)
225
+ server.free_group(99)
226
+ assert server._inflight_groups == [GI_A]
227
+ assert set(server._cache) == set(GROUP_A)
228
+
229
+ def test_release_drops_cache_and_count_entries(self, server_factory):
230
+ server = server_factory()
231
+ server.begin_sync(1)
232
+ _publish(server, GI_A, GROUP_A)
233
+ _publish(server, GI_B, GROUP_B)
234
+ server.free_group(GI_A)
235
+ assert set(server._cache) == set(GROUP_B)
236
+ assert server._free_counts.get(GI_A) is None
237
+ assert server._group_names.get(GI_A) is None
238
+
239
+
240
+class _CountingSource:
241
+ """`iter_groups` yielding one CPU-tensor group per call, built lazily so a
242
+ group's tensors exist only once the engine's credit gate lets it gather."""
243
+
244
+ def __init__(self, groups):
245
+ self._groups = groups
246
+
247
+ def iter_groups(self):
248
+ for names in self._groups:
249
+ yield list(names), [torch.zeros(16, dtype=torch.float32) for _ in names]
250
+
251
+
252
+class _ResidencyDict(dict):
253
+ """The engine's `_inflight`, instrumented: `max_len` is the high-water mark
254
+ of groups whose CUDA-IPC export refs were alive at once — the trainer's
255
+ resident-memory bound in groups."""
256
+
257
+ def __init__(self):
258
+ super().__init__()
259
+ self.max_len = 0
260
+
261
+ def __setitem__(self, key, value):
262
+ super().__setitem__(key, value)
263
+ self.max_len = max(self.max_len, len(self))
264
+
265
+
266
+def _loop_engine(server, n_groups, *, lookahead):
267
+ """A real engine wired straight to a real (non-Ray) server, with only the
268
+ state `_run_gather_loop` reads. The `_rpc` seam dispatches inline;
269
+ `_publish_async` sees no `.remote` and runs inline too."""
270
+ groups = [(f"model.layers.{gi}.w",) for gi in range(n_groups)]
271
+ e = ShardedRDTTrainerWeightTransferEngine.__new__(
272
+ ShardedRDTTrainerWeightTransferEngine
273
+ )
274
+ e._init_info = ShardedRDTTrainerInitInfo(
275
+ rank=0, num_consumers=1, gather_lookahead=lookahead
276
+ )
277
+ e.source = _CountingSource(groups)
278
+ e._server = server
279
+ e._rpc = lambda method, *a: getattr(server, method)(*a)
280
+ e._groups = [list(g) for g in groups]
281
+ e._owned_idx = list(range(n_groups))
282
+ e._held_names = None
283
+ e._inflight = _ResidencyDict()
284
+ return e
285
+
286
+
287
+class TestGatherCredit:
288
+ """The memory bound gates GATHERS, not publishes: a gathered group is
289
+ serveable immediately, and the ENGINE's loop stops gathering while more than
290
+ `gather_lookahead` groups are unfreed. So at most `lookahead + 1` groups are
291
+ resident — two at the default — while group N+1 is pulled the instant N's
292
+ pulls finish."""
293
+
294
+ @pytest.fixture
295
+ def gather_engine(self, monkeypatch):
296
+ """Neutralize the two CUDA touches in the export path so the REAL
297
+ gather loop runs on CPU: `.cuda()` becomes identity, and the storage
298
+ export ships the `(nbytes, ..., device, ...)` tuple the fixture's
299
+ stubbed `rebuild_cuda_tensor` expects."""
300
+ monkeypatch.setattr(torch.Tensor, "cuda", lambda self: self)
301
+ monkeypatch.setattr(
302
+ trainer_mod,
303
+ "reduce_tensor",
304
+ lambda base: (None, (base.numel(), None, None, None, None, None, -1, None)),
305
+ )
306
+ return _loop_engine
307
+
308
+ def test_publish_never_blocks(self, server_factory):
309
+ """Publishing must not park on a credit: the whole point is that a
310
+ gathered group becomes serveable immediately. Three publishes with no
311
+ frees, inline on this thread — a block here hangs the test."""
312
+ server = server_factory(gather_lookahead=1)
313
+ server.begin_sync(1)
314
+ for gi in range(3):
315
+ _publish(server, gi, (f"model.layers.{gi}.w",))
316
+ assert server._inflight_groups == [0, 1, 2]
317
+
318
+ def test_lookahead_one_holds_at_most_two_groups_resident(
319
+ self, server_factory, gather_engine
320
+ ):
321
+ """THE memory invariant: at lookahead=1 the trainer never holds more
322
+ than 2 groups of gathered weights, no matter how the consumer paces —
323
+ and the pipeline still overlaps, because the consumer here refuses to
324
+ free group N until N+1 is already published (pullable). A gate an
325
+ off-by-one too tight deadlocks this test; too loose fails max_len."""
326
+ server = server_factory(gather_lookahead=1)
327
+ n_groups = 8
328
+ engine = gather_engine(server, n_groups, lookahead=1)
329
+
330
+ failures: list[str] = []
331
+
332
+ def _consumer():
333
+ for gi in range(n_groups):
334
+ # N+1 must be pullable before N frees
335
+ target = min(gi + 1, n_groups - 1)
336
+ deadline = time.monotonic() + 10
337
+ while time.monotonic() < deadline:
338
+ with server._cache_cond:
339
+ if target in server._group_names:
340
+ break
341
+ time.sleep(0.002)
342
+ else:
343
+ failures.append(
344
+ f"group {target} never published while {gi} was unfreed"
345
+ )
346
+ return
347
+ server.free_group(gi)
348
+
349
+ consumer = threading.Thread(target=_consumer, daemon=True)
350
+ consumer.start()
351
+ engine._run_gather_loop(update_future=None, live_count=1)
352
+ consumer.join(timeout=10)
353
+
354
+ assert not consumer.is_alive() and not failures, failures
355
+ assert engine._inflight.max_len <= 2, (
356
+ f"{engine._inflight.max_len} groups were resident at lookahead=1"
357
+ )
358
+ assert engine._inflight == {} and server._cache == {}, (
359
+ "everything must be freed by end_sync"
360
+ )
361
+
362
+ def test_lookahead_zero_holds_exactly_one_group(
363
+ self, server_factory, gather_engine
364
+ ):
365
+ """0 is the serialized baseline: the loop waits for group N to be freed
366
+ before gathering N+1, so exactly ONE group is ever resident and nothing
367
+ overlaps. A floor of 1 on the bound would silently pipeline instead."""
368
+ server = server_factory(gather_lookahead=0)
369
+ engine = gather_engine(server, 4, lookahead=0)
370
+
371
+ def _run():
372
+ # unwound via set_gather_error below
373
+ with contextlib.suppress(Exception):
374
+ engine._run_gather_loop(update_future=None, live_count=1)
375
+
376
+ loop = threading.Thread(target=_run, daemon=True)
377
+ loop.start()
378
+ deadline = time.monotonic() + 5
379
+ while time.monotonic() < deadline and not server._inflight_groups:
380
+ time.sleep(0.005)
381
+ time.sleep(0.1) # a would-be second gather gets every chance to run
382
+ try:
383
+ assert engine._inflight.max_len == 1
384
+ assert server._inflight_groups == [0]
385
+ finally:
386
+ server.set_gather_error("test over") # unblock the parked wait_freed
387
+ loop.join(timeout=5)
388
+
389
+ def test_the_residency_bound_scales_with_the_lookahead(
390
+ self, server_factory, gather_engine
391
+ ):
392
+ """lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing
393
+ consumer sees exactly 3 groups gathered before the loop parks."""
394
+ server = server_factory(gather_lookahead=2)
395
+ engine = gather_engine(server, 8, lookahead=2)
396
+
397
+ def _run():
398
+ # unwound via set_gather_error below
399
+ with contextlib.suppress(Exception):
400
+ engine._run_gather_loop(update_future=None, live_count=1)
401
+
402
+ loop = threading.Thread(target=_run, daemon=True)
403
+ loop.start()
404
+ deadline = time.monotonic() + 5
405
+ while time.monotonic() < deadline and len(server._inflight_groups) < 3:
406
+ time.sleep(0.005)
407
+ time.sleep(0.1) # would-be 4th gather gets every chance to (wrongly) run
408
+ try:
409
+ assert engine._inflight.max_len == 3
410
+ assert server._inflight_groups == [0, 1, 2]
411
+ finally:
412
+ server.set_gather_error("test over") # unblock the parked wait_freed
413
+ loop.join(timeout=10)
414
+ assert not loop.is_alive()
415
+
416
+ def test_wait_freed_returns_the_freed_backlog(self, server_factory):
417
+ server = server_factory()
418
+ server.begin_sync(1)
419
+ _publish(server, GI_A, GROUP_A)
420
+ _publish(server, GI_B, GROUP_B)
421
+ server.free_group(GI_A)
422
+ assert server.wait_freed() == [GI_A]
423
+ server.free_group(GI_B)
424
+ assert server.wait_freed() == [GI_B]
425
+
426
+ def test_a_gather_error_releases_a_blocked_wait_freed(self, server_factory):
427
+ """Otherwise a trainer-side failure on another rank deadlocks this one
428
+ inside its credit gate. It must RAISE, not return empty — an empty
429
+ return would spin the engine straight back into the wait."""
430
+ server = server_factory()
431
+ server.begin_sync(1)
432
+ _publish(server, GI_A, GROUP_A)
433
+
434
+ errors: list = []
435
+ done = threading.Event()
436
+
437
+ def _wait():
438
+ try:
439
+ server.wait_freed()
440
+ except BaseException as e: # noqa: BLE001
441
+ errors.append(e)
442
+ done.set()
443
+
444
+ threading.Thread(target=_wait, daemon=True).start()
445
+ assert not done.wait(timeout=0.3), (
446
+ "wait_freed must block while nothing is freed"
447
+ )
448
+
449
+ server.set_gather_error("rank 3 export failed")
450
+ assert done.wait(timeout=5), "set_gather_error must unblock wait_freed"
451
+ assert errors and "rank 3 export failed" in str(errors[0])
452
+
453
+
454
+class TestEndSync:
455
+ def test_end_sync_returns_immediately_when_everything_is_freed(
456
+ self, server_factory
457
+ ):
458
+ server = server_factory()
459
+ server.begin_sync(1)
460
+ _publish(server, GI_A, GROUP_A)
461
+ server.free_group(GI_A)
462
+ assert server.end_sync() == [GI_A]
463
+
464
+ def test_end_sync_waits_for_an_outstanding_group(self, server_factory):
465
+ server = server_factory()
466
+ server.begin_sync(1)
467
+ _publish(server, GI_A, GROUP_A)
468
+
469
+ returned = []
470
+ done = threading.Event()
471
+
472
+ def _end():
473
+ returned.append(server.end_sync())
474
+ done.set()
475
+
476
+ t = threading.Thread(target=_end, daemon=True)
477
+ t.start()
478
+ assert not done.wait(timeout=0.3), "end_sync must wait for the barrier"
479
+
480
+ server.free_group(GI_A)
481
+ assert done.wait(timeout=5)
482
+ t.join(timeout=5)
483
+ assert returned == [[GI_A]]
484
+
485
+ def test_a_gather_error_releases_a_blocked_end_sync(self, server_factory):
486
+ server = server_factory()
487
+ server.begin_sync(1)
488
+ _publish(server, GI_A, GROUP_A)
489
+
490
+ done = threading.Event()
491
+
492
+ def _end_then_signal():
493
+ server.end_sync()
494
+ done.set()
495
+
496
+ threading.Thread(target=_end_then_signal, daemon=True).start()
497
+ assert not done.wait(timeout=0.3)
498
+ server.set_gather_error("boom")
499
+ assert done.wait(timeout=5)
500
+
501
+
502
+class TestBeginSync:
503
+ def test_begin_sync_sets_the_barrier_target(self, server_factory):
504
+ server = server_factory()
505
+ server.begin_sync(5)
506
+ assert server._live_count == 5
507
+
508
+ def test_the_target_is_floored_at_one(self, server_factory):
509
+ """A bare/zero call must not make every group free instantly (or divide
510
+ the barrier by zero)."""
511
+ server = server_factory()
512
+ server.begin_sync(0)
513
+ assert server._live_count == 1
514
+
515
+ def test_a_degraded_sync_lowers_the_target(self, server_factory):
516
+ """FT degraded sync: the live count is the WHOLE degraded-sync
517
+ mechanism on this side — the group releases after the live consumers
518
+ alone."""
519
+ server = server_factory()
520
+ server.begin_sync(2) # 4 provisioned, 2 alive
521
+ _publish(server, GI_A, GROUP_A)
522
+ server.free_group(GI_A)
523
+ assert server._inflight_groups == [GI_A]
524
+ server.free_group(GI_A)
525
+ assert server._inflight_groups == []
526
+
527
+ def test_begin_sync_clears_per_sync_state(self, server_factory):
528
+ server = server_factory()
529
+ server.begin_sync(2)
530
+ _publish(server, GI_A, GROUP_A)
531
+ server.free_group(GI_A)
532
+
533
+ server.begin_sync(2)
534
+ assert server._inflight_groups == []
535
+ assert server._freed_pending == []
536
+ assert server._free_counts == {}
537
+ assert server._group_names == {}
538
+
539
+ def test_a_straggler_signal_cannot_credit_the_next_sync(self, server_factory):
540
+ """begin_sync resets the counts, so a signal that leaked past the sync
541
+ boundary would over-credit — the reason the consumer drains its fired
542
+ signals before finishing. Pin the reset side of that contract."""
543
+ server = server_factory()
544
+ server.begin_sync(2)
545
+ _publish(server, GI_A, GROUP_A)
546
+ server.free_group(GI_A) # one of two consumers; sync then aborts/ends
547
+ server.begin_sync(2)
548
+ _publish(server, GI_A, GROUP_A)
549
+ server.free_group(GI_A)
550
+ assert server._inflight_groups == [GI_A], (
551
+ "one signal after the reset must not release"
552
+ )
553
+
554
+ def test_begin_sync_clears_a_previous_gather_error(self, server_factory):
555
+ server = server_factory()
556
+ server.set_gather_error("last sync failed")
557
+ server.begin_sync(1)
558
+ assert server._gather_error is None
559
+
560
+ def test_the_packed_destination_cache_survives_the_sync_boundary(
561
+ self, server_factory
562
+ ):
563
+ """The packed layout repeats every sync — that is what makes caching the
564
+ destination views worth ~5ms per 384-spec group."""
565
+ server = server_factory()
566
+ server._pack_dsts[("sentinel",)] = (0, [])
567
+ server.begin_sync(1)
568
+ assert ("sentinel",) in server._pack_dsts
569
+
570
+
571
+class TestServedNamesGuard:
572
+ """A producer serves only the names it holds: its stage's groups, and
573
+ within them its own EP coordinate's experts plus the replicated names.
574
+ Without this guard a misrouted pull would block forever in the cache wait
575
+ for a name this rank never publishes."""
576
+
577
+ def test_a_pull_for_an_unserved_name_fails_loudly(self, server_factory):
578
+ server = server_factory(served_names=list(GROUP_A))
579
+ server.begin_sync(1)
580
+ with pytest.raises(RuntimeError, match="wrong producer"):
581
+ server.rdt_produce_weights_batched([(GROUP_B[0], ())], consumer_id=0)
582
+
583
+ def test_the_error_names_the_unserved_weights(self, server_factory):
584
+ server = server_factory(served_names=list(GROUP_A))
585
+ server.begin_sync(1)
586
+ with pytest.raises(RuntimeError, match=GROUP_B[0]):
587
+ server.rdt_produce_weights_batched([(GROUP_B[0], ())], consumer_id=0)
588
+
589
+ def test_served_names_none_accepts_anything(self, server_factory):
590
+ """Gather-to-all: every producer holds every group, so there is nothing
591
+ to guard. The call proceeds to the cache wait (not exercised here)."""
592
+ server = server_factory(served_names=None)
593
+ assert server._served_names is None
594
+
595
+
596
+class TestStallWatchdog:
597
+ """The liveness bound on the producer's three blocking waits.
598
+
599
+ Engine-death detection is generation-driven, and no generation is in flight
600
+ during a weight sync, so a consumer that dies INSIDE the sync window has no
601
+ detector. It never sends its `free_group` signals, the group's barrier never
602
+ completes, and `wait_freed` / `end_sync` / the serve cache wait block
603
+ forever — which stops this rank iterating its WeightSource, which is a
604
+ collective, which wedges every other trainer rank in NCCL with no exception
605
+ anywhere.
606
+
607
+ The watchdog converts that into one real error. These use a tiny timeout;
608
+ the production default is 300s.
609
+ """
610
+
611
+ def test_a_blocked_wait_freed_fails_instead_of_hanging(self, server_factory):
612
+ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
613
+ server.begin_sync(1)
614
+ # the engine's credit gate now waits; nobody will signal
615
+ _publish(server, GI_A, GROUP_A)
616
+
617
+ errors: list = []
618
+ done = threading.Event()
619
+
620
+ def _wait():
621
+ try:
622
+ server.wait_freed()
623
+ except BaseException as e: # noqa: BLE001
624
+ errors.append(e)
625
+ done.set()
626
+
627
+ threading.Thread(target=_wait, daemon=True).start()
628
+
629
+ assert done.wait(timeout=10), (
630
+ "wait_freed must give up rather than block forever"
631
+ )
632
+ assert isinstance(server._gather_error, RuntimeError)
633
+ assert "RDT stall" in str(server._gather_error)
634
+ assert errors, "the engine must get an exception, not an empty credit list"
635
+
636
+ def test_end_sync_fails_when_a_consumer_never_signals(self, server_factory):
637
+ """The exact mid-sync death signature: the group is published, one of
638
+ the two live consumers that owed a signal is gone."""
639
+ server = server_factory(stall_timeout_s=0.3)
640
+ server.begin_sync(2)
641
+ _publish(server, GI_A, GROUP_A)
642
+ server.free_group(GI_A) # one of the two consumers signals; the other died
643
+
644
+ server.end_sync()
645
+ assert server._gather_error is not None
646
+ assert "RDT stall" in str(server._gather_error)
647
+
648
+ def test_the_error_reaches_every_blocked_waiter(self, server_factory):
649
+ """The watchdog fires on the existing `set_gather_error` channel, so one
650
+ stall unwinds the whole rank through one path rather than each waiter
651
+ timing out separately — here both a consumer parked in the serve cache
652
+ wait and the engine parked in its credit gate."""
653
+ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
654
+ server.begin_sync(1)
655
+ _publish(server, GI_A, GROUP_A)
656
+
657
+ errors: list = []
658
+
659
+ def _pull():
660
+ try:
661
+ server.rdt_produce_weights_batched([(GROUP_B[0], ())], consumer_id=0)
662
+ except BaseException as e: # noqa: BLE001
663
+ errors.append(("pull", e))
664
+
665
+ def _wait():
666
+ try:
667
+ server.wait_freed()
668
+ except BaseException as e: # noqa: BLE001
669
+ errors.append(("credit", e))
670
+
671
+ puller = threading.Thread(target=_pull, daemon=True)
672
+ waiter = threading.Thread(target=_wait, daemon=True)
673
+ puller.start()
674
+ waiter.start()
675
+ puller.join(timeout=10)
676
+ waiter.join(timeout=10)
677
+ assert not puller.is_alive() and not waiter.is_alive()
678
+ assert {who for who, _ in errors} == {"pull", "credit"}
679
+ assert all("gather errored" in str(e) for _, e in errors)
680
+
681
+ def test_progress_anywhere_keeps_a_slow_sync_alive(self, server_factory):
682
+ """A consumer that is slow but signaling must never trip it: the stamp
683
+ is global to the producer, so a steady trickle of signals holds the
684
+ whole rank open. Nine sequential publishes with a 0.3s timeout each take
685
+ longer in total than the timeout, and none may fire."""
686
+ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
687
+ server.begin_sync(1)
688
+ for gi in range(9):
689
+ _publish(server, gi, (f"model.layers.{gi}.w",))
690
+ time.sleep(0.1)
691
+ server.free_group(gi)
692
+ assert server._gather_error is None
693
+ server.end_sync()
694
+ assert server._gather_error is None
695
+
696
+ def test_begin_sync_resets_the_progress_stamp(self, server_factory):
697
+ """Syncs are minutes apart. Without the reset the first publish of sync N+1
698
+ would measure its stall from somewhere inside sync N and fire immediately."""
699
+ server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
700
+ server.begin_sync(1)
701
+ _publish(server, GI_A, GROUP_A)
702
+ server.free_group(GI_A)
703
+ server.end_sync()
704
+
705
+ time.sleep(0.5) # the idle gap between syncs, > the timeout
706
+ server.begin_sync(1)
707
+ _publish(server, GI_B, GROUP_B)
708
+ assert server._gather_error is None
709
+
710
+ def test_default_timeout_is_the_documented_one(self, server_factory):
711
+ server = server_factory()
712
+ assert server._stall_timeout == trainer_mod.DEFAULT_STALL_TIMEOUT_S
713
+
714
+
715
+class TestConcurrentPublishAndFree:
716
+ def test_the_lifecycle_survives_concurrent_publishes_and_signals(
717
+ self, server_factory
718
+ ):
719
+ """Every group published must end up freed and out of the cache, with
720
+ end_sync returning cleanly — the property whose violation is a hang."""
721
+ server = server_factory(gather_lookahead=2)
722
+ server.begin_sync(1)
723
+ groups = [(gi, (f"model.layers.{gi}.w",)) for gi in range(12)]
724
+
725
+ def _publisher():
726
+ for gi, names in groups:
727
+ _publish(server, gi, names)
728
+
729
+ def _signaler():
730
+ for gi, _names in groups:
731
+ while True:
732
+ with server._cache_cond:
733
+ published = gi in server._group_names
734
+ if published:
735
+ break
736
+ server.free_group(gi)
737
+
738
+ pub = threading.Thread(target=_publisher, daemon=True)
739
+ free = threading.Thread(target=_signaler, daemon=True)
740
+ pub.start()
741
+ free.start()
742
+ pub.join(timeout=10)
743
+ free.join(timeout=10)
744
+ assert not pub.is_alive() and not free.is_alive()
745
+
746
+ server.end_sync()
747
+ assert server._inflight_groups == []
748
+ assert server._cache == {}
749
+
750
+ def test_concurrent_signals_of_one_group_release_it_exactly_once(
751
+ self, server_factory
752
+ ):
753
+ """All live consumers signaling at once: the release must happen on the
754
+ last signal and only once, or the engine drops IPC refs twice (or
755
+ never)."""
756
+ server = server_factory()
757
+ server.begin_sync(8)
758
+ _publish(server, GI_A, GROUP_A)
759
+
760
+ barrier = threading.Barrier(8)
761
+
762
+ def _signal():
763
+ barrier.wait()
764
+ server.free_group(GI_A)
765
+
766
+ threads = [threading.Thread(target=_signal, daemon=True) for _ in range(8)]
767
+ for t in threads:
768
+ t.start()
769
+ for t in threads:
770
+ t.join(timeout=5)
771
+ assert not t.is_alive()
772
+
773
+ assert server._inflight_groups == []
774
+ assert server._freed_pending == [GI_A]
775
+ assert server._cache == {}
776
+
777
+
778
+class TestFakeServerAgreesWithTheRealOne:
779
+ """`test_sharded_rdt_trainer.py`'s `_FakeProducerServer` is a second,
780
+ independent model of this protocol that the engine-side tests assert on. Pin
781
+ the two against each other so the fake cannot silently drift from the
782
+ semantics the real server enforces."""
783
+
784
+ @staticmethod
785
+ def _fake():
786
+ from tests.distributed.test_sharded_rdt_trainer import _FakeProducerServer
787
+
788
+ return _FakeProducerServer(auto_free=False)
789
+
790
+ def test_the_barrier_and_early_signal_semantics_match(self, server_factory):
791
+ def publish_fake(srv, gi, names):
792
+ return srv.publish_group(gi, _entries(names))
793
+
794
+ for server, publish in (
795
+ (server_factory(), _publish),
796
+ (self._fake(), publish_fake),
797
+ ):
798
+ server.begin_sync(2)
799
+ # early signal, then publish, then the closing signal
800
+ server.free_group(GI_A)
801
+ assert publish(server, GI_A, GROUP_A) is None, "publish returns nothing"
802
+ assert server._inflight_groups == [GI_A]
803
+ server.free_group(GI_A)
804
+ assert server._inflight_groups == []
805
+ assert server.wait_freed() == [GI_A]
806
+ assert server.end_sync() == []
807
+
808
+
809
+class TestExportRing:
810
+ """The packed export path. Its one hard invariant is that the trainer packs
811
+ a group exactly the way ``publish_group`` rebuilds it — a layout
812
+ disagreement would serve the wrong bytes with nothing downstream to catch
813
+ it."""
814
+
815
+ @staticmethod
816
+ def _engine():
817
+ eng = ShardedRDTTrainerWeightTransferEngine.__new__(
818
+ ShardedRDTTrainerWeightTransferEngine
819
+ )
820
+ eng._export_ring = [None, None]
821
+ eng._export_ring_args = [None, None]
822
+ return eng
823
+
824
+ def _roundtrip(self, held, slot=0):
825
+ eng = self._engine()
826
+ storages, views, _refs = eng._pack_group_for_export(held, slot)
827
+ assert len(storages) == 1, "a packed group must export exactly one storage"
828
+ slot_t = eng._export_ring[slot]
829
+ return {
830
+ name: torch.as_strided(slot_t.view(getattr(torch, dt)), shape, stride, off)
831
+ for name, (_sid, dt, shape, stride, off) in views.items()
832
+ }
833
+
834
+ def test_packed_group_rebuilds_bit_identically(self):
835
+ held = [
836
+ ("bf16.2d", torch.randn(12, 7).bfloat16()),
837
+ ("fp32.1d", torch.randn(9)),
838
+ ("bf16.3d", torch.randn(4, 5, 6).bfloat16()),
839
+ ("bf16.transposed", torch.randn(6, 11).bfloat16().t()),
840
+ ("bf16.sliced", torch.randn(8, 32).bfloat16()[:, 3:19]),
841
+ ]
842
+ got = self._roundtrip(held)
843
+ for name, t in held:
844
+ assert got[name].dtype is t.dtype
845
+ assert tuple(got[name].shape) == tuple(t.shape)
846
+ assert torch.equal(got[name], t), name
847
+
848
+ def test_a_reused_slot_reallocates_and_still_rebuilds(self):
849
+ eng = self._engine()
850
+ eng._pack_group_for_export([("a", torch.randn(4, 4).bfloat16())], 0)
851
+ small = eng._export_ring[0].numel()
852
+ big = [("b", torch.randn(256, 128).bfloat16())]
853
+ _st, views, _rf = eng._pack_group_for_export(big, 0)
854
+ assert eng._export_ring[0].numel() > small
855
+ _name, (_sid, dt, shape, stride, off) = next(iter(views.items()))
856
+ rebuilt = torch.as_strided(
857
+ eng._export_ring[0].view(getattr(torch, dt)), shape, stride, off
858
+ )
859
+ assert torch.equal(rebuilt, big[0][1])
860
+
861
+ def test_slots_are_distinct_so_adjacent_groups_cannot_alias(self):
862
+ eng = self._engine()
863
+ eng._pack_group_for_export([("a", torch.randn(8, 8).bfloat16())], 0)
864
+ eng._pack_group_for_export([("b", torch.randn(8, 8).bfloat16())], 1)
865
+ assert eng._export_ring[0].data_ptr() != eng._export_ring[1].data_ptr()
866
+
867
+
868
+class TestExportRingSlotSafety:
869
+ """The export ring packs a group's small storages into one reused buffer and
870
+ ships per-name views into it, so a slot may only be rewritten once the group
871
+ that last used it is freed everywhere. The credit gate does not establish
872
+ that on its own: it bounds how many groups are unfreed, not the order they
873
+ free in, and barriers complete out of order (a consumer pre-frees a group it
874
+ pulls nothing from; a light group can finish before an earlier heavy one).
875
+ """
876
+
877
+ @pytest.fixture
878
+ def gather_engine(self, monkeypatch):
879
+ monkeypatch.setattr(torch.Tensor, "cuda", lambda self: self)
880
+ monkeypatch.setattr(
881
+ trainer_mod,
882
+ "reduce_tensor",
883
+ lambda base: (
884
+ None,
885
+ (base.numel(), None, None, None, None, None, -1, None),
886
+ ),
887
+ )
888
+ return _loop_engine
889
+
890
+ @staticmethod
891
+ def _await_published(server, gi, timeout=10):
892
+ deadline = time.monotonic() + timeout
893
+ while time.monotonic() < deadline:
894
+ with server._cache_cond:
895
+ if gi in server._group_names:
896
+ return True
897
+ time.sleep(0.002)
898
+ return False
899
+
900
+ def test_a_slot_is_never_rewritten_while_its_group_is_still_live(
901
+ self, server_factory, gather_engine
902
+ ):
903
+ """Free the later group of each pair first and hold the earlier one live
904
+ until the NEXT group is gathered -- the window a modular slot counter got
905
+ wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's)."""
906
+ server = server_factory(gather_lookahead=1)
907
+ n_groups = 8
908
+ engine = gather_engine(server, n_groups, lookahead=1)
909
+
910
+ violations: list[tuple[int, int]] = []
911
+ last_holder: dict[int, int] = {}
912
+ pending: dict[str, int] = {}
913
+ orig_pack = engine._pack_group_for_export
914
+ orig_pub = engine._publish_async
915
+
916
+ def _pack(held, slot_idx):
917
+ prev = last_holder.get(slot_idx)
918
+ if prev is not None and prev in engine._inflight:
919
+ violations.append((slot_idx, prev))
920
+ pending["slot"] = slot_idx
921
+ return orig_pack(held, slot_idx)
922
+
923
+ def _pub(group_idx, entries):
924
+ if "slot" in pending:
925
+ last_holder[pending.pop("slot")] = group_idx
926
+ return orig_pub(group_idx, entries)
927
+
928
+ engine._pack_group_for_export = _pack
929
+ engine._publish_async = _pub
930
+
931
+ def _consumer():
932
+ for base in range(0, n_groups, 2):
933
+ if not self._await_published(server, base + 1):
934
+ return
935
+ server.free_group(base + 1)
936
+ nxt = base + 2
937
+ if nxt < n_groups and not self._await_published(server, nxt):
938
+ return
939
+ server.free_group(base)
940
+
941
+ consumer = threading.Thread(target=_consumer, daemon=True)
942
+ consumer.start()
943
+ engine._run_gather_loop(update_future=None, live_count=1)
944
+ consumer.join(timeout=15)
945
+
946
+ assert last_holder, "no group was ring-packed; the test proves nothing"
947
+ assert not violations, f"slot rewritten while its group was live: {violations}"
948
+ assert engine._inflight == {}
949
+
950
+ def test_distinct_live_groups_never_share_a_slot(
951
+ self, server_factory, gather_engine
952
+ ):
953
+ server = server_factory(gather_lookahead=1)
954
+ engine = gather_engine(server, 6, lookahead=1)
955
+ seen: list[dict] = []
956
+ orig_pack = engine._pack_group_for_export
957
+
958
+ def _pack(held, slot_idx):
959
+ live = {
960
+ g: s for g, s in engine._slot_of_group.items() if g in engine._inflight
961
+ }
962
+ seen.append({"slot": slot_idx, "live": live})
963
+ return orig_pack(held, slot_idx)
964
+
965
+ engine._pack_group_for_export = _pack
966
+
967
+ def _consumer():
968
+ for gi in range(6):
969
+ if not self._await_published(server, gi):
970
+ return
971
+ server.free_group(gi)
972
+
973
+ consumer = threading.Thread(target=_consumer, daemon=True)
974
+ consumer.start()
975
+ engine._run_gather_loop(update_future=None, live_count=1)
976
+ consumer.join(timeout=15)
977
+
978
+ assert seen, "no group was ring-packed; the test proves nothing"
979
+ for rec in seen:
980
+ assert rec["slot"] not in rec["live"].values(), rec
981
+
982
+
983
+@pytest.fixture
984
+def sharing_server(server_factory):
985
+ """A producer whose serve slots are plain CPU tensors.
986
+
987
+ The pack is the only part of the serve path that touches a buffer, so
988
+ replacing the one allocation seam (``_new_serve_buffer``, which would
989
+ ``torch.empty(..., device="cuda")`` and register the result with NIXL) puts
990
+ the whole shared-pack path within reach of a CPU test.
991
+ """
992
+
993
+ def _make(**kwargs):
994
+ kwargs.setdefault("workers_per_replica", 2)
995
+ server = server_factory(**kwargs)
996
+ server._serve_device = torch.device("cpu")
997
+ server._new_serve_buffer = lambda n: torch.zeros(n, dtype=torch.uint8)
998
+ server.packs = 0
999
+ real_pack = server._pack_shared
1000
+
1001
+ def _counting_pack(gen, specs):
1002
+ server.packs += 1
1003
+ return real_pack(gen, specs)
1004
+
1005
+ server._pack_shared = _counting_pack
1006
+ return server
1007
+
1008
+ return _make
1009
+
1010
+
1011
+def _blob_bytes(server, name):
1012
+ """The cached tensor's bytes, as the pack should land them in the slot."""
1013
+ t = server._cache[name]
1014
+ return t.contiguous().view(torch.uint8).flatten()
1015
+
1016
+
1017
+class TestSharedServeSlots:
1018
+ """One registered slot serving several inference deployments.
1019
+
1020
+ Consumers whose ids differ by a multiple of ``workers_per_replica`` are the
1021
+ same worker of different deployments: same parallel config, same baked plan,
1022
+ same chunk sequence, byte-identical pack layout. They rendezvous per chunk,
1023
+ the last arrival packs, and every one of them returns that blob, so R
1024
+ deployments cost one ring and one pack rather than R of each.
1025
+ """
1026
+
1027
+ def test_one_deployment_never_shares(self, sharing_server):
1028
+ """With the width equal to the live count every group is a singleton, so
1029
+ each call packs for itself out of its own ring."""
1030
+ server = sharing_server(workers_per_replica=2)
1031
+ server.begin_sync(2, [0, 1])
1032
+ _publish(server, GI_A, GROUP_A)
1033
+ specs = [(GROUP_A[0], ())]
1034
+ blobs = [server.rdt_produce_weights_batched(specs, c, 0)[0] for c in (0, 1)]
1035
+
1036
+ assert server.packs == 2
1037
+ assert blobs[0].data_ptr() != blobs[1].data_ptr()
1038
+ assert len(server._serve_rings) == 2
1039
+
1040
+ def test_sharers_are_served_one_pack_out_of_one_buffer(self, sharing_server):
1041
+ server = sharing_server(workers_per_replica=2)
1042
+ server.begin_sync(4, [0, 1, 2, 3])
1043
+ _publish(server, GI_A, GROUP_A)
1044
+ specs = [(GROUP_A[0], ())]
1045
+
1046
+ with ThreadPoolExecutor(max_workers=2) as ex:
1047
+ futures = [
1048
+ ex.submit(server.rdt_produce_weights_batched, specs, c, 0)
1049
+ for c in (0, 2)
1050
+ ]
1051
+ blobs = [f.result(timeout=10)[0] for f in futures]
1052
+
1053
+ expected = _blob_bytes(server, GROUP_A[0])
1054
+ assert server.packs == 1
1055
+ assert blobs[0].data_ptr() == blobs[1].data_ptr()
1056
+ assert len(server._serve_rings) == 1
1057
+ assert torch.equal(blobs[0][: expected.numel()], expected)
1058
+
1059
+ def test_a_sharer_blocks_until_its_peer_arrives(self, sharing_server):
1060
+ """Without the wait the packer could repack a slot a peer had not read."""
1061
+ server = sharing_server(workers_per_replica=2, stall_timeout_s=30.0)
1062
+ server.begin_sync(4, [0, 1, 2, 3])
1063
+ _publish(server, GI_A, GROUP_A)
1064
+ specs = [(GROUP_A[0], ())]
1065
+
1066
+ with ThreadPoolExecutor(max_workers=2) as ex:
1067
+ first = ex.submit(server.rdt_produce_weights_batched, specs, 0, 0)
1068
+ with pytest.raises(FuturesTimeout):
1069
+ first.result(timeout=0.5)
1070
+ second = ex.submit(server.rdt_produce_weights_batched, specs, 2, 0)
1071
+ assert (
1072
+ first.result(timeout=10)[0].data_ptr()
1073
+ == second.result(timeout=10)[0].data_ptr()
1074
+ )
1075
+
1076
+ def test_the_slot_follows_the_consumers_issue_order(self, sharing_server):
1077
+ """The invariant slot safety rests on. The consumer drains pull i before
1078
+ issuing i+K, so choosing the slot as ``seq % K`` means a slot is only
1079
+ repacked once the read of its previous contents is over. Deriving it from
1080
+ a per-call counter on the producer instead (execution order) can hand a
1081
+ slot to a pull whose predecessor is still being read, because Ray may
1082
+ start a consumer's concurrent produce calls in any order.
1083
+ """
1084
+ server = sharing_server(workers_per_replica=2, num_rdt_buffers=2)
1085
+ server.begin_sync(4, [0, 1, 2, 3])
1086
+ names = ("a.w", "b.w", "c.w")
1087
+ for gi, name in enumerate(names):
1088
+ _publish(server, gi, (name,))
1089
+
1090
+ # Arrive OUT of issue order: seq 1 first, then 0, then 2.
1091
+ with ThreadPoolExecutor(max_workers=2) as ex:
1092
+ for seq in (1, 0, 2):
1093
+ futures = [
1094
+ ex.submit(
1095
+ server.rdt_produce_weights_batched,
1096
+ [(names[seq], ())],
1097
+ c,
1098
+ seq,
1099
+ )
1100
+ for c in (0, 2)
1101
+ ]
1102
+ [f.result(timeout=10) for f in futures]
1103
+
1104
+ slots = {seq: gen.slot for seq, gen in server._gens[0].items()}
1105
+ assert slots == {0: 0, 1: 1, 2: 0}, slots
1106
+
1107
+ def test_a_seqless_call_is_refused(self, sharing_server):
1108
+ """The slot comes from ``seq``, so a caller that sends none cannot be
1109
+ served safely."""
1110
+ server = sharing_server(workers_per_replica=2)
1111
+ server.begin_sync(1, [0])
1112
+ _publish(server, GI_A, GROUP_A)
1113
+
1114
+ with pytest.raises(ValueError, match="requires the consumer's issue"):
1115
+ server.rdt_produce_weights_batched([(GROUP_A[0], ())], 0)
1116
+
1117
+ def test_a_failed_pack_reaches_every_sharer(self, sharing_server):
1118
+ """Only the packer touches the buffer, so its failure has to be
1119
+ republished: a waiter must raise rather than read a half-written slot."""
1120
+ server = sharing_server(workers_per_replica=2)
1121
+ server.begin_sync(4, [0, 1, 2, 3])
1122
+ _publish(server, GI_A, GROUP_A)
1123
+ specs = [(GROUP_A[0], (("mutate_", (), ()),))] # not in ALLOWED_OPS
1124
+
1125
+ raised = []
1126
+ with ThreadPoolExecutor(max_workers=2) as ex:
1127
+ futures = [
1128
+ ex.submit(server.rdt_produce_weights_batched, specs, c, 0)
1129
+ for c in (0, 2)
1130
+ ]
1131
+ for future in futures:
1132
+ with pytest.raises((ValueError, RuntimeError)) as excinfo:
1133
+ future.result(timeout=10)
1134
+ raised.append(excinfo.value)
1135
+
1136
+ assert any(
1137
+ isinstance(e, ValueError) and "disallowed op" in str(e) for e in raised
1138
+ )
1139
+ assert any("packing this chunk failed" in str(e) for e in raised)
1140
+
1141
+ def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server):
1142
+ """A dead deployment must not be waited for: a survivor left alone in
1143
+ its group packs for itself instead of blocking on the corpse."""
1144
+ server = sharing_server(workers_per_replica=2, stall_timeout_s=30.0)
1145
+ server.begin_sync(3, [0, 1, 3]) # consumer 2 died
1146
+ _publish(server, GI_A, GROUP_A)
1147
+
1148
+ blob = server.rdt_produce_weights_batched([(GROUP_A[0], ())], 0, 0)[0]
1149
+
1150
+ expected = _blob_bytes(server, GROUP_A[0])
1151
+ assert torch.equal(blob[: expected.numel()], expected)
1152
+ assert server._sharers_of(0) == frozenset({0})
1153
+ assert server._sharers_of(1) == frozenset({1, 3})
1154
+
1155
+ def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server):
1156
+ """A consumer that dies INSIDE the sync window has no detector, exactly
1157
+ as for the free barrier; the rendezvous waits on the same channel."""
1158
+ server = sharing_server(workers_per_replica=2, stall_timeout_s=0.3)
1159
+ server.begin_sync(4, [0, 1, 2, 3])
1160
+ _publish(server, GI_A, GROUP_A)
1161
+
1162
+ with pytest.raises(RuntimeError, match="gather errored|RDT stall"):
1163
+ server.rdt_produce_weights_batched([(GROUP_A[0], ())], 0, 0)
1164
+ assert "RDT stall" in str(server._gather_error)
1165
+
1166
+
1167
+class TestSharedReservation:
1168
+ """``reserve_serve_buffer`` under sharing: one ring per GROUP, and the
1169
+ init-time check that the group's members pull the same chunks."""
1170
+
1171
+ def test_matching_plans_reserve_one_ring_between_them(self, sharing_server):
1172
+ server = sharing_server(workers_per_replica=2)
1173
+ for consumer_id in (0, 2):
1174
+ server.reserve_serve_buffer(consumer_id, 1024, "plan-a")
1175
+
1176
+ assert list(server._serve_rings) == [0]
1177
+ assert [s.numel() for s in server._serve_rings[0]] == [
1178
+ buffer_alloc_bytes(1024)
1179
+ ] * 2
1180
+
1181
+ def test_sharing_off_reserves_a_ring_per_consumer(self, sharing_server):
1182
+ server = sharing_server(workers_per_replica=0)
1183
+ for consumer_id in (0, 2):
1184
+ server.reserve_serve_buffer(consumer_id, 1024, "plan-a")
1185
+
1186
+ assert sorted(server._serve_rings) == [0, 2]
1187
+
1188
+ def test_mismatched_plans_fail_at_reservation(self, sharing_server):
1189
+ """Sharing rests on the sharers pulling the same chunks. Checked here,
1190
+ where it names both consumers, rather than at serve time, where it is a
1191
+ rendezvous nobody completes."""
1192
+ server = sharing_server(workers_per_replica=2)
1193
+ server.reserve_serve_buffer(0, 1024, "plan-a")
1194
+
1195
+ with pytest.raises(RuntimeError, match="do not pull the same chunks"):
1196
+ server.reserve_serve_buffer(2, 1024, "plan-b")
1197
+
1198
+ def test_different_groups_may_differ(self, sharing_server):
1199
+ """Only the members of ONE group have to agree; worker 0 and worker 1 of
1200
+ the same deployment hold different slices and never share a slot."""
1201
+ server = sharing_server(workers_per_replica=2)
1202
+ server.reserve_serve_buffer(0, 1024, "plan-a")
1203
+ server.reserve_serve_buffer(1, 1024, "plan-b")
1204
+
1205
+ assert sorted(server._serve_rings) == [0, 1]
1206
+
1207
+ def test_a_missing_digest_skips_the_check(self, sharing_server):
1208
+ """The digest is optional on the wire, so an engine that does not send
1209
+ one still reserves."""
1210
+ server = sharing_server(workers_per_replica=2)
1211
+ server.reserve_serve_buffer(0, 1024, None)
1212
+ server.reserve_serve_buffer(2, 1024, None)
1213
+
1214
+ assert list(server._serve_rings) == [0]
tests/distributed/test_sharded_rdt_trainer.py
new
+628
@@ -0,0 +1,628 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the sharded-RDT trainer-side engine.
4
+
5
+The trainer engine's own surface: init info and factory registration, the
6
+send_weights round trip against an in-process fake producer server, the
7
+ownership table it resolves and ships to the workers, and the producer's packed
8
+destination-view cache.
9
+
10
+The plan/replay half (including `RdtRouter`, which is consumer-side) lives in
11
+test_sharded_rdt_plan.py; the producer server's protocol in
12
+test_sharded_rdt_producer.py.
13
+"""
14
+
15
+import itertools
16
+from dataclasses import asdict
17
+
18
+import pytest
19
+import torch
20
+
21
+from tests.distributed.test_weight_transfer import RecordingClient
22
+from vllm.distributed.weight_transfer import (
23
+ ModuleSource,
24
+ ParamMeta,
25
+ WeightSource,
26
+ WeightTransferTrainerFactory,
27
+)
28
+from vllm.distributed.weight_transfer.base import layerwise_groups
29
+from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
30
+ ShardedRDTTrainerInitInfo,
31
+ ShardedRDTTrainerWeightTransferEngine,
32
+)
33
+
34
+
35
+class _ListSource(WeightSource):
36
+ """A WeightSource over an explicit ordered (name, cpu-tensor) list, so the
37
+ sharded-RDT group/order logic can be tested without a real model."""
38
+
39
+ def __init__(self, pairs):
40
+ self._pairs = list(pairs)
41
+
42
+ def metadata(self):
43
+ return [ParamMeta(n, t.dtype, tuple(t.shape)) for n, t in self._pairs]
44
+
45
+ def __iter__(self):
46
+ return iter(self._pairs)
47
+
48
+
49
+class _FakeProducerServer:
50
+ """In-process stand-in for the _RDTProducerServer Ray actor. Records the
51
+ engine->server call sequence and, by default, frees each group as soon as
52
+ it is published (simulating the consumers' free_group barrier) so the
53
+ gather loop's credit gate never parks. Mirrors the real server's
54
+ per-group barrier: publish/free are keyed by GROUP INDEX, signals count
55
+ to the ``begin_sync`` live total, publish_group returns nothing, and freed
56
+ groups flow back only through wait_freed / end_sync (see
57
+ test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne)."""
58
+
59
+ def __init__(self, auto_free=True):
60
+ self.order: list[str] = []
61
+ self.published: list[int] = []
62
+ self.live_count = 1
63
+ self._inflight_groups: list[int] = []
64
+ self.auto_free = auto_free
65
+ self.free_counts: dict[int, int] = {}
66
+ self._pending_freed: list[int] = []
67
+
68
+ # Alias for tests that inspect in-flight groups.
69
+ @property
70
+ def inflight(self):
71
+ return self._inflight_groups
72
+
73
+ def begin_sync(self, live_count, live_consumer_ids=None):
74
+ self.order.append("begin")
75
+ self.live_count = max(1, int(live_count))
76
+ self.live_ids = live_consumer_ids
77
+
78
+ def publish_group(self, group_idx, entries):
79
+ self.order.append("publish")
80
+ self.published.append(group_idx)
81
+ self._inflight_groups.append(group_idx)
82
+ if self.auto_free or self.free_counts.get(group_idx, 0) >= self.live_count:
83
+ self._inflight_groups.remove(group_idx)
84
+ self._pending_freed.append(group_idx)
85
+
86
+ def free_group(self, group_idx):
87
+ """Consumer back-edge; may arrive before the group's publish."""
88
+ self.free_counts[group_idx] = self.free_counts.get(group_idx, 0) + 1
89
+ if (
90
+ group_idx in self._inflight_groups
91
+ and self.free_counts[group_idx] >= self.live_count
92
+ ):
93
+ self._inflight_groups.remove(group_idx)
94
+ self._pending_freed.append(group_idx)
95
+
96
+ def wait_freed(self):
97
+ """The engine's credit gate. The real server blocks here; the fake
98
+ must already have a banked credit when the gate asks (auto_free, or a
99
+ test's own free_group) — anything else is the deadlock the real
100
+ watchdog would kill, so fail loudly. Not appended to ``order``: it is
101
+ pacing, not a lifecycle milestone."""
102
+ assert self._pending_freed, (
103
+ "wait_freed with nothing freed: the gather loop would deadlock"
104
+ )
105
+ freed = self._pending_freed
106
+ self._pending_freed = []
107
+ return freed
108
+
109
+ def end_sync(self):
110
+ self.order.append("end")
111
+ freed = self._pending_freed
112
+ self._pending_freed = []
113
+ return freed
114
+
115
+ def set_gather_error(self, message):
116
+ self.order.append("error")
117
+
118
+
119
+def _rdt_engine_with_fake_server(
120
+ source, *, is_sender, client, server, monkeypatch, fleet_owned=None
121
+):
122
+ """Build a ShardedRDTTrainerWeightTransferEngine wired to an in-process fake
123
+ server (no Ray, no CUDA IPC): bypass trainer_init's spawn, set the
124
+ group-major metadata, and route _rpc to the fake."""
125
+ import vllm.distributed.weight_transfer.sharded_rdt_trainer as mod
126
+
127
+ # reduce_tensor needs CUDA; the fake server never rebuilds, so stub it.
128
+ monkeypatch.setattr(mod, "reduce_tensor", lambda t: (None, ("fake",)))
129
+
130
+ init_info = ShardedRDTTrainerInitInfo(num_consumers=1, rank=0 if is_sender else 1)
131
+ engine = ShardedRDTTrainerWeightTransferEngine(
132
+ client=client, source=source, is_sender=is_sender, init_info=init_info
133
+ )
134
+ engine._meta = list(source.metadata())
135
+ names = [m.name for m in engine._meta]
136
+ engine._groups = layerwise_groups(names)
137
+ engine._server = server
138
+ engine._rpc = lambda method, *args: getattr(server, method)(*args)
139
+ # What trainer_init resolves from the source's held names + the fleet's
140
+ # all-gather. ``fleet_owned`` stands in for that all-gather so a partial-
141
+ # ownership rank can be tested without a real process group; the fleet must
142
+ # cover every name or the transpose rejects it (nothing would serve the rest).
143
+ if fleet_owned is None:
144
+ engine._resolve_ownership(1, 0)
145
+ else:
146
+ # Stands in for the (metadata digest, held-name bitmask) all-gather:
147
+ # each rank's group list becomes the mask its held names would set.
148
+ names = [m.name for m in engine._meta]
149
+ groups = engine._groups
150
+
151
+ def _fake_gather(_world, mine):
152
+ out = []
153
+ for owned in fleet_owned:
154
+ held = {n for gi in owned for n in groups[gi]}
155
+ mask = bytearray((len(names) + 7) // 8)
156
+ for i, n in enumerate(names):
157
+ if n in held:
158
+ mask[i >> 3] |= 1 << (i & 7)
159
+ out.append((mine[0], bytes(mask)))
160
+ return out
161
+
162
+ monkeypatch.setattr(engine, "_all_gather_owned", _fake_gather)
163
+ engine._resolve_ownership(len(fleet_owned), 0)
164
+ return engine
165
+
166
+
167
+def _rdt_source_two_layers():
168
+ return _ListSource(
169
+ [
170
+ ("embed.weight", torch.zeros(2)),
171
+ ("model.layers.0.w", torch.zeros(2)),
172
+ ("model.layers.1.w", torch.zeros(2)),
173
+ ("norm.weight", torch.zeros(2)),
174
+ ]
175
+ )
176
+
177
+
178
+class TestShardedRDTTrainerInitInfo:
179
+ def test_declares_backend(self):
180
+ assert ShardedRDTTrainerInitInfo.backend == "sharded_rdt"
181
+
182
+ def test_rank_is_keyword_only_and_drives_is_sender(self):
183
+ assert ShardedRDTTrainerInitInfo(num_consumers=4, rank=0).is_sender is True
184
+ assert ShardedRDTTrainerInitInfo(num_consumers=4, rank=1).is_sender is False
185
+ with pytest.raises(TypeError):
186
+ # rank is keyword-only.
187
+ ShardedRDTTrainerInitInfo(4, 0) # type: ignore[misc]
188
+
189
+ def test_registered_in_trainer_factory(self):
190
+ cls = WeightTransferTrainerFactory._registry["sharded_rdt"]()
191
+ assert cls is ShardedRDTTrainerWeightTransferEngine
192
+
193
+
194
+def test_sharded_rdt_trainer_init_requires_source():
195
+ with pytest.raises(ValueError, match="requires a WeightSource"):
196
+ ShardedRDTTrainerWeightTransferEngine.trainer_init(
197
+ ShardedRDTTrainerInitInfo(num_consumers=1, rank=0),
198
+ client=RecordingClient(),
199
+ source=None,
200
+ )
201
+
202
+
203
+def test_sharded_rdt_worker_init_info_is_group_major(monkeypatch):
204
+ source = _rdt_source_two_layers()
205
+ engine = _rdt_engine_with_fake_server(
206
+ source,
207
+ is_sender=True,
208
+ client=RecordingClient(),
209
+ server=_FakeProducerServer(),
210
+ monkeypatch=monkeypatch,
211
+ )
212
+ worker_init = engine._build_worker_init_info(["srv_rk0"])
213
+ # 4 params -> pre / layer0 / layer1 / post = 4 groups of length 1.
214
+ assert worker_init.names == [
215
+ "embed.weight",
216
+ "model.layers.0.w",
217
+ "model.layers.1.w",
218
+ "norm.weight",
219
+ ]
220
+ assert worker_init.group_lens == [1, 1, 1, 1]
221
+ assert worker_init.trainer_actor_names == ["srv_rk0"]
222
+ assert worker_init.produce_method_name == "rdt_produce_weights_batched"
223
+ assert sum(worker_init.group_lens) == len(worker_init.names)
224
+
225
+
226
+@pytest.mark.skipif(
227
+ not torch.cuda.is_available(),
228
+ reason="the gather loop's CUDA-IPC export needs a device",
229
+)
230
+def test_sharded_rdt_send_weights_drives_client_in_order(monkeypatch):
231
+ server = _FakeProducerServer(auto_free=True)
232
+ client = RecordingClient()
233
+ engine = _rdt_engine_with_fake_server(
234
+ _rdt_source_two_layers(),
235
+ is_sender=True,
236
+ client=client,
237
+ server=server,
238
+ monkeypatch=monkeypatch,
239
+ )
240
+ engine.send_weights()
241
+
242
+ assert client.order == ["start", "update", "finish"]
243
+ # begin, one publish per group (4), end.
244
+ assert server.order == ["begin", "publish", "publish", "publish", "publish", "end"]
245
+ assert len(server.published) == 4
246
+ # every group freed -> no engine-held refs remain.
247
+ assert engine._inflight == {}
248
+
249
+
250
+def test_sharded_rdt_send_weights_group_order_mismatch_raises(monkeypatch):
251
+ # Source whose iteration order disagrees with its metadata order.
252
+ class _BadSource(_ListSource):
253
+ def __iter__(self):
254
+ reordered = list(self._pairs)
255
+ reordered[0], reordered[1] = reordered[1], reordered[0]
256
+ return iter(reordered)
257
+
258
+ server = _FakeProducerServer()
259
+ engine = _rdt_engine_with_fake_server(
260
+ _BadSource(_rdt_source_two_layers()._pairs),
261
+ is_sender=True,
262
+ client=RecordingClient(),
263
+ server=server,
264
+ monkeypatch=monkeypatch,
265
+ )
266
+ with pytest.raises(RuntimeError, match="iteration order must match"):
267
+ engine.send_weights()
268
+ assert "error" in server.order # gather error propagated to the server
269
+
270
+
271
+@pytest.mark.skipif(
272
+ not torch.cuda.is_available(),
273
+ reason="the gather loop's CUDA-IPC export needs a device",
274
+)
275
+def test_sharded_rdt_non_sender_skips_client(monkeypatch):
276
+ class _RaisingClient(RecordingClient):
277
+ def start_weight_update(self):
278
+ raise AssertionError("non-sender must not touch the client")
279
+
280
+ def update_weights(self, update_info):
281
+ raise AssertionError("non-sender must not touch the client")
282
+
283
+ def finish_weight_update(self, weight_version: str | None = None):
284
+ raise AssertionError("non-sender must not touch the client")
285
+
286
+ server = _FakeProducerServer(auto_free=True)
287
+ client = _RaisingClient()
288
+ engine = _rdt_engine_with_fake_server(
289
+ _rdt_source_two_layers(),
290
+ is_sender=False,
291
+ client=client,
292
+ server=server,
293
+ monkeypatch=monkeypatch,
294
+ )
295
+ engine.send_weights() # gathers only; must not raise
296
+ assert client.order == []
297
+ assert server.order == ["begin", "publish", "publish", "publish", "publish", "end"]
298
+
299
+
300
+@pytest.mark.skipif(
301
+ not torch.cuda.is_available(),
302
+ reason="the gather loop's CUDA-IPC export needs a device",
303
+)
304
+def test_sharded_rdt_send_weights_surfaces_update_error(monkeypatch):
305
+ class _FailingUpdateClient(RecordingClient):
306
+ def update_weights(self, update_info):
307
+ self.order.append("update")
308
+ raise RuntimeError("inference side rejected update")
309
+
310
+ server = _FakeProducerServer(auto_free=True)
311
+ engine = _rdt_engine_with_fake_server(
312
+ _rdt_source_two_layers(),
313
+ is_sender=True,
314
+ client=_FailingUpdateClient(),
315
+ server=server,
316
+ monkeypatch=monkeypatch,
317
+ )
318
+ with pytest.raises(RuntimeError, match="inference side rejected update"):
319
+ engine.send_weights()
320
+
321
+
322
+class _OwnedSource(_ListSource):
323
+ """A source holding only some groups' names, like a pipeline-parallel rank.
324
+ Takes group indices for convenience and declares the names inside them."""
325
+
326
+ def __init__(self, pairs, owned_group_idx):
327
+ super().__init__(pairs)
328
+ self._owned = list(owned_group_idx)
329
+ groups = layerwise_groups([n for n, _ in pairs])
330
+ self._owned_names = [n for gi in self._owned for n in groups[gi]]
331
+
332
+ def held_names(self):
333
+ return list(self._owned_names)
334
+
335
+ def __iter__(self):
336
+ by_name = dict(self._pairs)
337
+ return iter([(n, by_name[n]) for n in self._owned_names])
338
+
339
+
340
+@pytest.mark.skipif(
341
+ not torch.cuda.is_available(),
342
+ reason="the gather loop's CUDA-IPC export needs a device",
343
+)
344
+def test_sharded_rdt_publishes_only_the_groups_it_holds(monkeypatch):
345
+ server = _FakeProducerServer(auto_free=True)
346
+ engine = _rdt_engine_with_fake_server(
347
+ _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]),
348
+ is_sender=False,
349
+ client=RecordingClient(),
350
+ server=server,
351
+ monkeypatch=monkeypatch,
352
+ fleet_owned=[[1, 2], [0, 3]], # this rank holds the layers, rank 1 the rest
353
+ )
354
+ engine.send_weights()
355
+
356
+ assert server.published == [1, 2]
357
+ assert server.order == ["begin", "publish", "publish", "end"]
358
+ assert engine._inflight == {}
359
+ # embed + norm on rank 1, the two layers here: two distinct owner sets,
360
+ # numbered by first appearance in metadata order.
361
+ assert engine._owner_sets == [[1], [0]]
362
+ assert engine._name_owner_class == [0, 1, 1, 0]
363
+
364
+
365
+def test_sharded_rdt_owned_group_order_mismatch_raises(monkeypatch):
366
+ class _MisorderedOwned(_OwnedSource):
367
+ def __iter__(self):
368
+ return iter(list(super().__iter__())[::-1])
369
+
370
+ server = _FakeProducerServer(auto_free=True)
371
+ engine = _rdt_engine_with_fake_server(
372
+ _MisorderedOwned(_rdt_source_two_layers()._pairs, [1, 2]),
373
+ is_sender=False,
374
+ client=RecordingClient(),
375
+ server=server,
376
+ monkeypatch=monkeypatch,
377
+ fleet_owned=[[1, 2], [0, 3]],
378
+ )
379
+ with pytest.raises(RuntimeError, match="iteration order must match"):
380
+ engine.send_weights()
381
+ assert "error" in server.order
382
+
383
+
384
+@pytest.mark.skipif(
385
+ not torch.cuda.is_available(),
386
+ reason="the gather loop's CUDA-IPC export needs a device",
387
+)
388
+def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch):
389
+ """The free barrier's target is one integer per sync: the live consumer
390
+ count, defaulting to the whole provisioned fleet."""
391
+ server = _FakeProducerServer(auto_free=True)
392
+ engine = _rdt_engine_with_fake_server(
393
+ _rdt_source_two_layers(),
394
+ is_sender=False,
395
+ client=RecordingClient(),
396
+ server=server,
397
+ monkeypatch=monkeypatch,
398
+ )
399
+ engine.send_weights()
400
+ assert server.live_count == 1 # num_consumers=1 in the harness
401
+ engine.send_weights(live_consumer_ids=[0])
402
+ assert server.live_count == 1
403
+
404
+
405
+class TestLiveCountPlumbing:
406
+ """``send_weights(live_consumer_ids)`` -> the barrier target. The
407
+ provisioned geometry is frozen; a degraded sync only lowers the target and
408
+ narrows each slot-sharing group's rendezvous."""
409
+
410
+ @staticmethod
411
+ def _engine(num_consumers, world=2, rank=0):
412
+ engine = ShardedRDTTrainerWeightTransferEngine.__new__(
413
+ ShardedRDTTrainerWeightTransferEngine
414
+ )
415
+ engine._init_info = ShardedRDTTrainerInitInfo(
416
+ rank=rank, num_consumers=num_consumers
417
+ )
418
+ engine.source = object() # send_weights asserts a source is present
419
+ received: list = []
420
+ engine._send_weights_inner = lambda count, ids: received.append((count, ids))
421
+ return engine, received
422
+
423
+ def test_none_counts_the_whole_provisioned_fleet(self):
424
+ engine, got = self._engine(num_consumers=8)
425
+ engine.send_weights(None)
426
+ assert got == [(8, list(range(8)))]
427
+
428
+ def test_a_live_set_counts_its_distinct_members(self):
429
+ engine, got = self._engine(num_consumers=8)
430
+ engine.send_weights([0, 1, 4, 5, 5])
431
+ assert got == [(4, [0, 1, 4, 5])]
432
+
433
+ def test_a_full_live_set_matches_the_provisioned_count(self):
434
+ engine, got = self._engine(num_consumers=8)
435
+ engine.send_weights(list(range(8)))
436
+ engine.send_weights(None)
437
+ assert got == [(8, list(range(8)))] * 2
438
+
439
+ def test_which_consumers_died_does_not_matter_only_how_many(self):
440
+ """The whole point of the barrier: no routed per-producer targets, so
441
+ the identity of the dead consumer is irrelevant to the producers."""
442
+ counts = []
443
+ for live in ([0, 1, 2, 3], [4, 5, 6, 7], [0, 2, 4, 6]):
444
+ engine, got = self._engine(num_consumers=8)
445
+ engine.send_weights(live)
446
+ counts += [c for c, _ids in got]
447
+ assert counts == [4, 4, 4]
448
+
449
+ def test_the_live_ids_travel_with_the_count(self):
450
+ """The count sizes the free barrier, the ids size the slot-sharing
451
+ rendezvous, so a producer that shares slots can tell WHICH consumers it
452
+ is still waiting for. They must describe the same set."""
453
+ engine, got = self._engine(num_consumers=8)
454
+ engine.send_weights([6, 0, 2, 2])
455
+ ((count, ids),) = got
456
+ assert ids == [0, 2, 6] and count == len(ids)
457
+
458
+
459
+def test_sharded_rdt_worker_init_info_carries_the_ownership_table(monkeypatch):
460
+ import json
461
+
462
+ server = _FakeProducerServer(auto_free=True)
463
+ engine = _rdt_engine_with_fake_server(
464
+ _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]),
465
+ is_sender=True,
466
+ client=RecordingClient(),
467
+ server=server,
468
+ monkeypatch=monkeypatch,
469
+ fleet_owned=[[0, 1], [2, 3]],
470
+ )
471
+ worker_init = engine._build_worker_init_info(["srv_rk0", "srv_rk1"])
472
+ assert worker_init.owner_sets == [[0], [1]]
473
+ assert worker_init.name_owner_class == [0, 0, 1, 1]
474
+ assert len(worker_init.name_owner_class) == len(worker_init.names)
475
+ # The payload crosses the control plane as JSON: nested lists must survive.
476
+ round_tripped = json.loads(json.dumps(asdict(worker_init)))
477
+ assert round_tripped["owner_sets"] == [[0], [1]]
478
+ assert round_tripped["name_owner_class"] == [0, 0, 1, 1]
479
+
480
+
481
+def test_weight_source_holds_everything_by_default():
482
+ """The contract's default: a source that says nothing holds the whole model."""
483
+ src = _rdt_source_two_layers()
484
+ assert src.held_names() is None
485
+ assert ModuleSource(torch.nn.Linear(2, 2)).held_names() is None
486
+
487
+
488
+def test_sharded_rdt_rejects_a_held_name_outside_metadata(monkeypatch):
489
+ class _BadHeld(_OwnedSource):
490
+ def held_names(self):
491
+ return ["embed.weight", "not.a.real.name"]
492
+
493
+ with pytest.raises(ValueError, match="not"):
494
+ _rdt_engine_with_fake_server(
495
+ _BadHeld(_rdt_source_two_layers()._pairs, [0]),
496
+ is_sender=False,
497
+ client=RecordingClient(),
498
+ server=_FakeProducerServer(),
499
+ monkeypatch=monkeypatch,
500
+ )
501
+
502
+
503
+def test_sharded_rdt_rejects_a_rank_holding_nothing(monkeypatch):
504
+ class _HoldsNothing(_ListSource):
505
+ def held_names(self):
506
+ return []
507
+
508
+ with pytest.raises(ValueError, match="empty"):
509
+ _rdt_engine_with_fake_server(
510
+ _HoldsNothing(_rdt_source_two_layers()._pairs),
511
+ is_sender=False,
512
+ client=RecordingClient(),
513
+ server=_FakeProducerServer(),
514
+ monkeypatch=monkeypatch,
515
+ )
516
+
517
+
518
+def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch):
519
+ """Every name must be held somewhere or it can never be served — caught
520
+ when the holdings are transposed, naming the orphan."""
521
+ with pytest.raises(ValueError, match="no trainer rank holds"):
522
+ _rdt_engine_with_fake_server(
523
+ _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]),
524
+ is_sender=False,
525
+ client=RecordingClient(),
526
+ server=_FakeProducerServer(),
527
+ monkeypatch=monkeypatch,
528
+ fleet_owned=[[0, 1], [2]], # nobody holds group 3
529
+ )
530
+
531
+
532
+def test_sharded_rdt_rejects_metadata_disagreement_across_ranks(monkeypatch):
533
+ """Only the sender's metadata reaches the consumers, so a rank describing
534
+ just its own share must fail loudly rather than silently drop the rest."""
535
+ engine = _rdt_engine_with_fake_server(
536
+ _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]),
537
+ is_sender=False,
538
+ client=RecordingClient(),
539
+ server=_FakeProducerServer(auto_free=True),
540
+ monkeypatch=monkeypatch,
541
+ fleet_owned=[[1, 2], [0, 3]], # a covering fleet, so construction succeeds
542
+ )
543
+ # Now rank 1 reports a DIFFERENT metadata digest for the same model.
544
+ monkeypatch.setattr(
545
+ engine,
546
+ "_all_gather_owned",
547
+ lambda w, mine: [mine, ("deadbeefdeadbeef", mine[1])],
548
+ )
549
+ with pytest.raises(ValueError, match="disagrees across trainer ranks"):
550
+ engine._resolve_ownership(2, 0)
551
+
552
+
553
+def _serve_ring_server(src_name, src):
554
+ """A producer server with one cached tensor and a pre-seeded serve ring, so a
555
+ pull needs no Ray and no NIXL registration. Returns (server, serve) where
556
+ ``serve(chain)`` packs one spec into the SAME ring slot every time — which is
557
+ what puts the destination-view cache, and only it, under test."""
558
+ from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
559
+ _RDTProducerServer,
560
+ )
561
+
562
+ srv = _RDTProducerServer(
563
+ num_rdt_buffers=2,
564
+ buffer_presize_gb=0.0,
565
+ gather_lookahead=2,
566
+ )
567
+ srv._cache[src_name] = src
568
+ srv._serve_rings[0] = [
569
+ torch.empty(1 << 16, dtype=torch.uint8, device="cuda") for _ in range(2)
570
+ ]
571
+
572
+ # The slot is ``seq % nring``, so stepping seq by nring keeps every pull in
573
+ # slot 0 while still making each one its own generation -- reusing a seq
574
+ # would rendezvous with the finished pack and return it unchanged.
575
+ seqs = itertools.count(0, srv._nring)
576
+
577
+ def serve(chain):
578
+ return srv.rdt_produce_weights_batched(
579
+ [(src_name, chain)], consumer_id=0, seq=next(seqs)
580
+ )[0]
581
+
582
+ return srv, serve
583
+
584
+
585
+@pytest.mark.skipif(
586
+ not torch.cuda.is_available(), reason="the producer server needs a CUDA device"
587
+)
588
+def test_serve_does_not_reuse_packed_views_of_another_shape():
589
+ """Two requests can share a name yet pack different slices of it.
590
+
591
+ The producer caches the destination views it carves into a serve ring slot.
592
+ Keyed by name alone, the second request is packed through the first's views.
593
+ Reachable when one name's copies split across owner-class chunks.
594
+ """
595
+ name = "model.layers.0.w"
596
+ src = torch.arange(64, dtype=torch.bfloat16, device="cuda").reshape(8, 8)
597
+ _srv, serve = _serve_ring_server(name, src)
598
+
599
+ serve((("narrow", (0, 0, 2), ()),)) # 2 rows
600
+ wide = src.narrow(0, 0, 6) # 6 rows, same name, same slot
601
+ blob = serve((("narrow", (0, 0, 6), ()),))
602
+
603
+ got = blob[: wide.numel() * wide.element_size()].view(wide.dtype).reshape(6, 8)
604
+ assert torch.equal(got, wide)
605
+
606
+
607
+@pytest.mark.skipif(
608
+ not torch.cuda.is_available(), reason="the producer server needs a CUDA device"
609
+)
610
+def test_serve_does_not_reuse_packed_views_of_another_dtype():
611
+ """The SILENT case of the same cache: two requests whose slices have the same
612
+ name and the same shape but different dtypes pack at identical offsets, so
613
+ reusing the stale views raises nothing — ``copy_`` just casts, and the blob
614
+ carries the wrong bytes with no check downstream.
615
+
616
+ ``view(dtype)`` rather than ``to(dtype)`` because the chain is replayed under
617
+ ``ALLOWED_OPS``, which rejects ``to``: same-itemsize dtypes keep the byte
618
+ length identical, which is exactly what makes the stale view fit.
619
+ """
620
+ name = "model.layers.0.w"
621
+ src = torch.arange(64, dtype=torch.bfloat16, device="cuda").reshape(8, 8)
622
+ _srv, serve = _serve_ring_server(name, src)
623
+
624
+ serve((("view", (torch.float16,), ()),)) # same shape and bytes, fp16
625
+ blob = serve(()) # same name and shape, bf16
626
+
627
+ got = blob[: src.numel() * src.element_size()].view(src.dtype).reshape(8, 8)
628
+ assert torch.equal(got, src), "packed through a cached view of the wrong dtype"
tests/distributed/test_weight_transfer.py
+158
@@ -32,8 +32,10 @@ from vllm.distributed.weight_transfer import (
32
)
33
from vllm.distributed.weight_transfer.base import (
34
TrainerInitInfo,
35
+ WeightTransferEngine,
36
WeightTransferInitRequest,
37
WeightTransferUpdateRequest,
38
+ layerwise_groups,
39
)
40
from vllm.distributed.weight_transfer.ipc_engine import (
41
IPCTrainerInitInfo,
@@ -1412,6 +1414,162 @@ class TestModuleSource:
1414
assert [m.shape for m in meta] == [tuple(t.shape) for _, t in pairs]
1415
1416
1417
+class TestWeightSourceGroupContract:
1418
+ """`groups()` / `iter_groups()` on the WeightSource ABC. Groups are what
1419
+ backends gather and free by, so the default must agree with
1420
+ `layerwise_groups` over `metadata()`, restricted to what this rank holds."""
1421
+
1422
+ class _Source(WeightSource):
1423
+ """Minimal source over an ordered (name, tensor) list, optionally holding
1424
+ only some names (in which case it iterates only their groups)."""
1425
+
1426
+ def __init__(self, names, held=None, reverse=False):
1427
+ self._pairs = [(n, torch.full((2,), float(i))) for i, n in enumerate(names)]
1428
+ self._held = held
1429
+ self._reverse = reverse
1430
+
1431
+ def metadata(self):
1432
+ return [ParamMeta(n, t.dtype, tuple(t.shape)) for n, t in self._pairs]
1433
+
1434
+ def held_names(self):
1435
+ return self._held
1436
+
1437
+ def __iter__(self):
1438
+ pairs = self._pairs
1439
+ if self._held is not None:
1440
+ keep = set(self._held)
1441
+ pairs = [(n, t) for n, t in pairs if n in keep]
1442
+ return iter(list(reversed(pairs)) if self._reverse else pairs)
1443
+
1444
+ def _source(self, names, held=None, reverse=False):
1445
+ return self._Source(names, held, reverse)
1446
+
1447
+ def test_groups_defaults_to_the_layerwise_partition(self):
1448
+ names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
1449
+ assert self._source(names).groups() == layerwise_groups(names)
1450
+
1451
+ def test_groups_keeps_only_groups_holding_a_held_name(self):
1452
+ names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
1453
+ held = ["model.layers.0.a", "model.layers.1.a"]
1454
+ assert self._source(names, held=held).groups() == [
1455
+ ["model.layers.0.a"],
1456
+ ["model.layers.1.a"],
1457
+ ]
1458
+
1459
+ def test_groups_keeps_a_partially_held_group_whole(self):
1460
+ """One held name selects the WHOLE group, unheld members included.
1461
+
1462
+ A group is the collective unit, so its membership has to be
1463
+ rank-uniform: a rank that narrowed it to what it holds would disagree
1464
+ with its peers about which names group *g* covers, and the gather would
1465
+ desynchronize. Selection is per group, but only whole groups; the
1466
+ per-name split is the backend's, via owner sets. This is the
1467
+ foreign-expert shape under expert parallelism — a rank holds some
1468
+ experts of a layer, not all of them.
1469
+ """
1470
+ names = ["model.layers.0.a", "model.layers.0.b", "model.layers.1.a"]
1471
+ held = ["model.layers.0.a"]
1472
+ assert self._source(names, held=held).groups() == [
1473
+ ["model.layers.0.a", "model.layers.0.b"],
1474
+ ]
1475
+
1476
+ def test_groups_order_follows_metadata_not_the_declaration(self):
1477
+ """The declaration is a SET of names; the groups it selects still come
1478
+ out in metadata order, so each pairs with the right ``iter_groups()``
1479
+ batch however the source listed them."""
1480
+ names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
1481
+ source = self._source(
1482
+ names, held=["model.layers.1.a", "model.layers.0.a", "model.layers.1.a"]
1483
+ )
1484
+ assert source.groups() == [["model.layers.0.a"], ["model.layers.1.a"]]
1485
+ assert [ns for ns, _ in source.iter_groups()] == [
1486
+ ["model.layers.0.a"],
1487
+ ["model.layers.1.a"],
1488
+ ]
1489
+
1490
+ def test_iter_groups_batches_the_stream_per_group(self):
1491
+ names = ["embed.w", "model.layers.0.a", "model.layers.0.b", "norm.w"]
1492
+ batches = list(self._source(names).iter_groups())
1493
+ assert [ns for ns, _ in batches] == [
1494
+ ["embed.w"],
1495
+ ["model.layers.0.a", "model.layers.0.b"],
1496
+ ["norm.w"],
1497
+ ]
1498
+ assert all(len(ns) == len(ts) for ns, ts in batches)
1499
+
1500
+ def test_iter_groups_yields_the_tensors_iteration_produced(self):
1501
+ names = ["model.layers.0.a", "model.layers.0.b"]
1502
+ (batch,) = list(self._source(names).iter_groups())
1503
+ _names, tensors = batch
1504
+ assert [float(t[0]) for t in tensors] == [0.0, 1.0]
1505
+
1506
+ def test_iter_groups_yields_only_held_groups(self):
1507
+ names = ["embed.w", "model.layers.0.a", "model.layers.1.a"]
1508
+ batches = list(self._source(names, held=["model.layers.1.a"]).iter_groups())
1509
+ assert [ns for ns, _ in batches] == [["model.layers.1.a"]]
1510
+
1511
+ def test_out_of_order_iteration_raises(self):
1512
+ """Materializing is usually a collective, so a rank that iterates out of
1513
+ order deadlocks its peers -- fail loudly instead."""
1514
+ source = self._source(["model.layers.0.a", "model.layers.0.b"], reverse=True)
1515
+ with pytest.raises(RuntimeError, match="iteration order must match"):
1516
+ list(source.iter_groups())
1517
+
1518
+ def test_a_source_may_override_iter_groups(self):
1519
+ """The extension point: materialize a whole group in one step instead of
1520
+ one generator resume per tensor."""
1521
+ calls = []
1522
+
1523
+ class _Batched(TestWeightSourceGroupContract._Source):
1524
+ def iter_groups(self):
1525
+ for group in self.groups():
1526
+ calls.append(len(group))
1527
+ yield group, [torch.zeros(2) for _ in group]
1528
+
1529
+ source = _Batched(["model.layers.0.a", "model.layers.0.b"])
1530
+ assert [ns for ns, _ in source.iter_groups()] == [
1531
+ ["model.layers.0.a", "model.layers.0.b"]
1532
+ ]
1533
+ assert calls == [2]
1534
+
1535
+
1536
+class TestDeferredProcessingContract:
1537
+ """`defers_processing` and `drain_pending` are two halves of one contract: a
1538
+ caller that takes over the update tail (running its own
1539
+ `finalize_layerwise_reload` instead of going through `finish_weight_update`)
1540
+ reads the flag and calls the method. Both must be answerable on any engine, or
1541
+ that caller ends up reaching through a getattr."""
1542
+
1543
+ def _engines(self):
1544
+ return {
1545
+ name: loader()
1546
+ for name, loader in WeightTransferEngineFactory._registry.items()
1547
+ }
1548
+
1549
+ def test_every_engine_declares_whether_it_defers(self):
1550
+ for name, cls in self._engines().items():
1551
+ assert isinstance(cls.defers_processing, bool), name
1552
+
1553
+ def test_every_engine_can_be_drained(self):
1554
+ """The default is a no-op, so a caller never has to check whether the
1555
+ method exists before calling it."""
1556
+ for name, cls in self._engines().items():
1557
+ assert callable(cls.drain_pending), name
1558
+
1559
+ def test_the_default_is_not_to_defer(self):
1560
+ assert WeightTransferEngine.defers_processing is False
1561
+
1562
+ def test_a_synchronous_engine_drains_as_a_no_op(self):
1563
+ engine = object.__new__(WeightTransferEngineFactory._registry["nccl"]())
1564
+ engine.drain_pending() # must not raise, and must not need any state
1565
+
1566
+ def test_the_rdt_engine_defers_and_overrides_the_drain(self):
1567
+ """The one engine the contract exists for."""
1568
+ cls = WeightTransferEngineFactory._registry["sharded_rdt"]()
1569
+ assert cls.defers_processing is True
1570
+ assert cls.drain_pending is not WeightTransferEngine.drain_pending
1571
+
1572
+
1573
class TestTrainerFactory:
1574
"""WeightTransferTrainerFactory registry mechanics."""
1575
tools/pre_commit/check_torch_cuda.py
+1
@@ -20,6 +20,7 @@ ALLOWED_FILES = {
20
"vllm/platforms/",
21
"vllm/device_allocator/",
22
"vllm/distributed/weight_transfer/ipc_engine.py",
23
+ "vllm/distributed/weight_transfer/sharded_rdt_engine.py",
24
"tests/distributed/test_packed_tensor.py",
25
}
26
vllm/config/weight_transfer.py
+1
-1
@@ -9,7 +9,7 @@ from vllm.config.utils import config
9
class WeightTransferConfig:
10
"""Configuration for weight transfer during RL training."""
11
12
- backend: Literal["nccl", "ipc", "sparse_nccl"] | str = "nccl"
12
+ backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt"] | str = "nccl"
13
"""The backend to use for weight transfer. Validated against the
14
`WeightTransferEngineFactory` registry at engine creation time.
15
"""
vllm/distributed/weight_transfer/base.py
+181
-6
@@ -3,7 +3,7 @@
3
"""Base class for weight transfer engines."""
4
5
from abc import ABC, abstractmethod
6
-from collections.abc import Iterator
6
+from collections.abc import Collection, Iterator
7
from dataclasses import dataclass, field
8
from typing import (
9
TYPE_CHECKING,
@@ -33,6 +33,76 @@ TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo")
33
# channel. The built-in `ModuleSource` uses `materialize_full_tensor`.
34
35
36
+def _stack_key(name: str) -> tuple[str, int] | None:
37
+ """``(prefix, index)`` of the OUTERMOST integer segment, or None if there is
38
+ none.
39
+
40
+ Outermost is what keeps a MoE layer whole:
41
+ ``model.layers.3.mlp.experts.7.w1`` keys on the layer, not the expert.
42
+ """
43
+ parts = name.split(".")
44
+ for i, part in enumerate(parts):
45
+ if part.isdigit():
46
+ return ".".join(parts[:i]), int(part)
47
+ return None
48
+
49
+
50
+def layerwise_groups(names: list[str]) -> list[list[str]]:
51
+ """Partition flat parameter names into one group per decoder layer, keyed on
52
+ the outermost index segment of each name.
53
+
54
+ This defines what a *group index* means for `WeightSource.groups` and
55
+ `WeightSource.iter_groups`: index *g* names the same group on every trainer
56
+ rank and every consumer, because it is derived from one rank's `metadata()`
57
+ order.
58
+
59
+ Keying on the index rather than a literal prefix needs no per-architecture
60
+ naming table: ``model.layers.0.``, ``model.language_model.layers.0.``,
61
+ ``transformer.h.0.``, ``backbone.layers.0.`` and a vision tower's
62
+ ``visual.blocks.0.`` all partition alike. Matching one fixed prefix does not,
63
+ and its failure is silent — every name lands in a single group holding the
64
+ whole model, which defeats the per-layer bound below.
65
+
66
+ Un-indexed names split by POSITION relative to the first indexed one: the pre
67
+ block (embeddings) and the post block (the final norm, `lm_head`, and any
68
+ inter-stack projector). Post lands last however early it arrived, which is
69
+ what a pipeline-parallel source needs — Megatron-Bridge streams the last
70
+ stage's output block *before* its layers.
71
+
72
+ Stacks come out in first-appearance order of their prefix and ascending index
73
+ within it, whatever order the source yielded them, so a source can normalize
74
+ an arbitrary export order by flattening this partition.
75
+
76
+ Backends that gather and free per group (sharded RDT) also use it as the unit
77
+ of transfer, which bounds their buffer sizes: without it a whole model becomes
78
+ one chunk.
79
+ """
80
+ pre: list[str] = []
81
+ post: list[str] = []
82
+ stacks: dict[tuple[str, int], list[str]] = {}
83
+ order: list[tuple[str, int]] = []
84
+ for name in names:
85
+ key = _stack_key(name)
86
+ if key is None:
87
+ (post if order else pre).append(name)
88
+ continue
89
+ if key not in stacks:
90
+ stacks[key] = []
91
+ order.append(key)
92
+ stacks[key].append(name)
93
+
94
+ prefix_rank: dict[str, int] = {}
95
+ for key in order:
96
+ prefix_rank.setdefault(key[0], len(prefix_rank))
97
+ order.sort(key=lambda key: (prefix_rank[key[0]], key[1]))
98
+
99
+ groups: list[list[str]] = [pre] if pre else []
100
+ groups += [stacks[key] for key in order]
101
+ if post:
102
+ groups.append(post)
103
+ return groups
104
+
105
+
106
def materialize_full_tensor(tensor: torch.Tensor) -> torch.Tensor:
107
"""Return a full, locally-materialized tensor ready to send.
108
@@ -66,10 +136,13 @@ class WeightSource(ABC):
136
case it should cache.
137
* iteration — yields fully-materialized `(name, tensor)` pairs, one at a
138
time. Materializing is typically a collective (FSDP `full_tensor()`, a
69
- Megatron export), so every trainer rank must iterate the same source in the
70
- same order in lockstep, or ranks deadlock. Under pipeline parallelism a
71
- rank may not own a parameter at all — iterating still drives the collective
72
- and the yielded tensor is only meaningful on the sender.
139
+ Megatron export), so the ranks that share a parameter must iterate it in
140
+ the same order in lockstep, or they deadlock.
141
+ * `held_names()` — which parameters this rank holds, for producers that are
142
+ split so each rank holds only part of the model. Defaults to all.
143
+ * `iter_groups()` — the same stream batched per gather group (see
144
+ `layerwise_groups`). Defaults to batching `__iter__`; override to
145
+ materialize a whole group in one step.
146
147
`iter(source)` must yield a *fresh* pass each round. Backends with custom
148
producer logic (Megatron export, RDT plans, MoE re-fusing) subclass this.
@@ -92,6 +165,79 @@ class WeightSource(ABC):
165
def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
166
raise NotImplementedError
167
168
+ def held_names(self) -> "Collection[str] | None":
169
+ """The parameters this rank holds, or None for all of them.
170
+
171
+ This is the whole ownership contract. Override it when producers are
172
+ split so each holds only part of the model — pipeline parallelism (a rank
173
+ holds some layers), expert parallelism (a rank holds some experts), or
174
+ any combination, including layouts that fit neither. A consumer routes
175
+ each name to a rank that holds it, so per-name is the granularity that
176
+ matters; the engine derives everything else from this.
177
+
178
+ Three requirements come with overriding it:
179
+
180
+ * `metadata()` must still describe the WHOLE model on every rank. The
181
+ group partition, the iteration checks and the consumers' pull plans are
182
+ all built from one rank's metadata, so a rank that reported only its own
183
+ share would leave the rest of the model silently un-transferred. The
184
+ sharded-RDT engine cross-checks this across ranks at init.
185
+ * Every name must be held by at least one rank, or it can never be
186
+ served. The engine raises at init naming the first orphan.
187
+ * Iteration must cover exactly `groups()` in metadata order, yielding a
188
+ real tensor for each held name and `None` for the rest. A group's
189
+ gather is a collective among the ranks that hold part of it, so the
190
+ name must still appear (to keep the order check aligned) while the
191
+ data is absent.
192
+
193
+ Returns:
194
+ The held parameter names, or None to hold every one.
195
+ """
196
+ return None
197
+
198
+ def groups(self) -> list[list[str]]:
199
+ """This rank's gather groups, in metadata order: `layerwise_groups` over
200
+ `metadata()`, restricted to the groups holding at least one held name.
201
+
202
+ A group with nothing held here is not iterated at all — its gather is a
203
+ collective among the ranks that do hold part of it.
204
+ """
205
+ groups = layerwise_groups([m.name for m in self.metadata()])
206
+ held = self.held_names()
207
+ if held is None:
208
+ return groups
209
+ held = set(held)
210
+ return [g for g in groups if any(n in held for n in g)]
211
+
212
+ def iter_groups(self) -> Iterator[tuple[list[str], list[torch.Tensor]]]:
213
+ """Yield one `(names, tensors)` batch per group from `groups()`.
214
+
215
+ The default drives `__iter__` and batches its output, checking as it goes
216
+ that the names arrive in metadata order — ranks sharing a parameter
217
+ materialize it with a collective, so a rank that iterates out of order
218
+ deadlocks its peers rather than returning wrong data.
219
+
220
+ Override when a backend can produce a whole group at once. Materializing
221
+ is usually a collective, and driving it per group instead of per tensor
222
+ turns ~37k generator resumes into ~95 on a per-expert MoE model (worth
223
+ ~0.9s per sync there). An override must yield the same batches in the same
224
+ order as this default.
225
+ """
226
+ it = iter(self)
227
+ for group in self.groups():
228
+ names: list[str] = []
229
+ tensors: list[torch.Tensor] = []
230
+ for expected in group:
231
+ name, tensor = next(it)
232
+ if name != expected:
233
+ raise RuntimeError(
234
+ f"WeightSource yielded {name!r} but expected "
235
+ f"{expected!r}; iteration order must match metadata()."
236
+ )
237
+ names.append(name)
238
+ tensors.append(tensor)
239
+ yield names, tensors
240
+
241
242
class ModuleSource(WeightSource):
243
"""`WeightSource` over `module.named_parameters()` — the common case.
@@ -207,6 +353,31 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
353
354
supports_draft_weight_update: bool = True
355
356
+ defers_processing: bool = False
357
+ """Whether `update_weights` returns before the weights are on the device.
358
+
359
+ An engine that pipelines its GPU post-processing onto background threads
360
+ cannot let `update_weights` synchronize the device — that would block on those
361
+ threads and serialize the pipeline. Such an engine sets this True, omits the
362
+ per-update sync, and guarantees completion in `finish_weight_update` instead.
363
+
364
+ Callers that go through `finish_weight_update` need do nothing: the engine
365
+ drains there. A caller that instead drives the tail itself — running its own
366
+ `finalize_layerwise_reload`, say — must read this flag and call
367
+ `drain_pending()` first, because with it set a returned `update_weights` means
368
+ "queued", not "applied".
369
+ """
370
+
371
+ def drain_pending(self) -> None:
372
+ """Block until every deferred update has been applied to the model.
373
+
374
+ The companion to `defers_processing`: a caller that has taken over the
375
+ update tail calls this to re-establish the guarantee that
376
+ `finish_weight_update` would otherwise have given it. Idempotent, and a
377
+ no-op by default — an engine that processes synchronously has nothing to
378
+ drain, so this is always safe to call.
379
+ """
380
+
381
def __init__(
382
self,
383
config: WeightTransferConfig,
@@ -385,7 +556,11 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
556
Symmetric to `WeightTransferEngine` but lives in the training process.
557
Constructed via the `trainer_init` factory classmethod; carries any
558
backend-specific state (NCCL communicators, IPC device info, transfer
388
- plans) on `self`.
559
+ plans) on `self`. Full-resync backends (NCCL, IPC) take a `WeightSource` at
560
+ `trainer_init` and replay it each round via the no-argument
561
+ `send_weights()`. Backends that push per-round deltas instead (e.g. sparse
562
+ patches) leave `source` as `None` and take their payload as a `send_weights`
563
+ argument.
564
565
Unlike the worker engine, the trainer side does not take a
566
`WeightTransferConfig`: the backend is selected from the init info's
vllm/distributed/weight_transfer/factory.py
+12
@@ -234,6 +234,12 @@ WeightTransferEngineFactory.register_engine(
234
"SparseNCCLWeightTransferEngine",
235
)
236
237
+WeightTransferEngineFactory.register_engine(
238
+ "sharded_rdt",
239
+ "vllm.distributed.weight_transfer.sharded_rdt_engine",
240
+ "ShardedRDTWeightTransferEngine",
241
+)
242
+
243
244
# Trainer-side engines, parallel to the worker registry above.
245
WeightTransferTrainerFactory.register_engine(
@@ -253,3 +259,9 @@ WeightTransferTrainerFactory.register_engine(
259
"vllm.distributed.weight_transfer.sparse_nccl_engine",
260
"SparseNCCLTrainerWeightTransferEngine",
261
)
262
+
263
+WeightTransferTrainerFactory.register_engine(
264
+ "sharded_rdt",
265
+ "vllm.distributed.weight_transfer.sharded_rdt_trainer",
266
+ "ShardedRDTTrainerWeightTransferEngine",
267
+)
vllm/distributed/weight_transfer/sharded_rdt_common.py
new
+339
@@ -0,0 +1,339 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared pieces of the sharded-RDT backend: the op-chain allowlist, buffer
4
+sizing and the minimum Ray version, plus ``RdtRouter`` — consumer-only, since
5
+routing is a consumer decision.
6
+
7
+The gather-group partition itself is ``base.layerwise_groups``: it defines what a
8
+group index means for any ``WeightSource``, not just this transport.
9
+"""
10
+
11
+from collections.abc import Callable
12
+from typing import Any
13
+
14
+import torch
15
+
16
+# The op-chain contract, with two enforcers. The consumer's ``FakeRDTTensor``
17
+# intercepts exactly these methods during the bake and records the string name;
18
+# anything else reaches ``__torch_dispatch__`` and raises, so a loader needing
19
+# real data fails at init rather than transferring wrong bytes. The producer
20
+# replays a chain as ``getattr(tensor, op)(...)`` and refuses anything outside
21
+# ``ALLOWED_OPS``, so a spoofed consumer cannot invoke arbitrary methods.
22
+#
23
+# Every entry must be a pure view; ``to`` in particular must never be added, as
24
+# it is the dtype/device escape the bake exists to reject. Both sets derive from
25
+# this one table so they cannot disagree.
26
+SUPPORTED_OPS: dict[Callable, str] = {
27
+ torch.Tensor.narrow: "narrow",
28
+ torch.Tensor.view: "view",
29
+ torch.Tensor.reshape: "reshape",
30
+ torch.Tensor.__getitem__: "__getitem__",
31
+ torch.Tensor.unsqueeze: "unsqueeze",
32
+ torch.Tensor.squeeze: "squeeze",
33
+ torch.Tensor.transpose: "transpose",
34
+ torch.Tensor.t: "t",
35
+ torch.Tensor.permute: "permute",
36
+ torch.Tensor.flatten: "flatten",
37
+ torch.Tensor.contiguous: "contiguous",
38
+ torch.Tensor.chunk: "chunk",
39
+ # Multi-return like chunk: one child per output, trailing __getitem__(i).
40
+ torch.Tensor.unbind: "unbind",
41
+}
42
+
43
+ALLOWED_OPS = frozenset(SUPPORTED_OPS.values())
44
+
45
+# Ray Direct Transport landed in halves: the actor opt-in
46
+# ``enable_tensor_transport`` in 2.49, and ``ray.experimental``'s
47
+# ``register_nixl_memory`` / ``set_target_for_ref`` — which this transport calls
48
+# directly — in 2.55. The floor is the tested version rather than the earliest
49
+# one that could import; 2.55 has never been run against.
50
+RDT_MIN_RAY_VERSION = "2.56.0"
51
+
52
+
53
+def check_ray_rdt_version() -> None:
54
+ """Refuse an installed Ray older than the one this backend is tested on.
55
+
56
+ vLLM does not depend on Ray, so there is no pin to carry this. Without the
57
+ check the failure is an opaque option-validation error out of
58
+ ``.options(enable_tensor_transport=True)`` (below 2.49) or an ImportError
59
+ raised deep in the first pull, long after init reported success (below 2.55).
60
+
61
+ Raises:
62
+ ValueError: the installed Ray predates ``RDT_MIN_RAY_VERSION``.
63
+ """
64
+ import importlib.metadata
65
+
66
+ from packaging import version
67
+
68
+ required = version.parse(RDT_MIN_RAY_VERSION)
69
+ current = version.parse(importlib.metadata.version("ray"))
70
+ if current < required:
71
+ raise ValueError(
72
+ f"The 'sharded_rdt' weight transfer backend requires Ray "
73
+ f">= {required}, the version it is tested against; Ray Direct "
74
+ f"Transport (ray.experimental.register_nixl_memory / "
75
+ f"set_target_for_ref) needs at least 2.55. Found {current}. "
76
+ f"Run `pip install -U 'ray>={required}'`."
77
+ )
78
+
79
+
80
+def assign_producer_indices(
81
+ num_producers: int, num_consumers: int, consumer_idx: int
82
+) -> list[int]:
83
+ """Producers (global indices) that consumer ``consumer_idx`` binds."""
84
+ p = max(1, num_producers)
85
+ c = max(1, num_consumers)
86
+ if p >= c:
87
+ return list(range(consumer_idx * p // c, (consumer_idx + 1) * p // c))
88
+ return [consumer_idx * p // c]
89
+
90
+
91
+class RdtRouter:
92
+ """Decides which producer serves each weight name, and — once bound — sends.
93
+
94
+ Ownership is per NAME. ``owner_sets`` holds the few distinct producer sets
95
+ that occur, and ``name_owner_class[i]`` indexes it for ``names[i]``; a name
96
+ determines both its owner set and its gather group, so neither appears in the
97
+ routing API. Empty tables mean every producer holds everything.
98
+
99
+ Any placement is expressible this way: a pipeline stage is a set of names
100
+ sharing one owner set, an expert is a name whose owner set is a single rank,
101
+ and a group produced by two stages is just two classes inside one group.
102
+
103
+ Both engines derive the same tables from the same wire data, so they agree on
104
+ who serves what. Disagreement is not a wrong answer but a hang or a loud
105
+ misroute: a pull sent to a producer that never gathered the name trips its
106
+ served-names guard.
107
+
108
+ Routing is consumer-only. The trainer publishes what it holds and answers
109
+ whatever arrives; it never asks who serves what.
110
+ """
111
+
112
+ def __init__(
113
+ self,
114
+ num_producers: int,
115
+ num_consumers: int,
116
+ owner_sets: list[list[int]] | None = None,
117
+ name_owner_class: list[int] | None = None,
118
+ names: list[str] | None = None,
119
+ group_lens: list[int] | None = None,
120
+ workers_per_replica: int = 0,
121
+ ) -> None:
122
+ """Build the routing tables from the wire data both engines receive.
123
+
124
+ The four table arguments travel together and default together: with all
125
+ of them empty the router degrades to "every producer holds everything".
126
+
127
+ Args:
128
+ num_producers: Trainer ranks. Owner indices are positions in this
129
+ range, and ``validate`` rejects any that fall outside it.
130
+ num_consumers: Inference workers across the whole fleet. Fixes
131
+ the id space; the block carve uses ``workers_per_replica``
132
+ of it.
133
+ owner_sets: The distinct producer sets that occur, one row per
134
+ owner class; each row is deduplicated and sorted here. Empty
135
+ means a single class owning every producer.
136
+ name_owner_class: Parallel to ``names`` — ``name_owner_class[i]``
137
+ indexes ``owner_sets`` for ``names[i]``. A name with no entry
138
+ falls back to class 0.
139
+ names: Every parameter name, in group-major order: concatenating
140
+ the gather groups reproduces this list exactly. The other two
141
+ tables are keyed by this order.
142
+ group_lens: Length of each gather group, consecutive over
143
+ ``names``, so they sum to ``len(names)``. Fixes name -> group
144
+ and the per-group owner union the free barrier fans out to.
145
+ workers_per_replica: Consumers per inference DEPLOYMENT. The block
146
+ carve spreads this many consumers over an owner set and every
147
+ deployment reuses that carve, so the same worker of each
148
+ deployment resolves one producer (see ``producer_for``). 0
149
+ means one deployment: carve over the whole fleet.
150
+
151
+ Raises:
152
+ ValueError: ``workers_per_replica`` does not divide
153
+ ``num_consumers``, so the deployments are not uniform and two
154
+ workers of one deployment would share a block index.
155
+ """
156
+ self.num_producers = max(1, num_producers)
157
+ self.num_consumers = max(1, num_consumers)
158
+ # Width of the block carve: one deployment's worth of consumers, so
159
+ # deployments overlay rather than spread. See ``producer_for``.
160
+ self._block_consumers = min(
161
+ self.num_consumers, max(1, workers_per_replica or self.num_consumers)
162
+ )
163
+ if self.num_consumers % self._block_consumers:
164
+ raise ValueError(
165
+ f"workers_per_replica={workers_per_replica} does not divide "
166
+ f"num_consumers={self.num_consumers}; the inference deployments "
167
+ f"must be uniform for the per-deployment block carve to be "
168
+ f"well defined."
169
+ )
170
+ self._owner_sets = (
171
+ [sorted(set(row)) for row in owner_sets]
172
+ if owner_sets
173
+ else [list(range(self.num_producers))]
174
+ )
175
+ names = list(names or [])
176
+ classes = list(name_owner_class or [])
177
+ self._class_of = {
178
+ n: (classes[i] if i < len(classes) else 0) for i, n in enumerate(names)
179
+ }
180
+ # name -> gather group, and the per-group owner union the free barrier
181
+ # fans out to. Both are pure functions of the tables, so precompute once.
182
+ self._group_of: dict[str, int] = {}
183
+ self._group_owners: list[list[int]] = []
184
+ pos = 0
185
+ for gi, glen in enumerate(group_lens or []):
186
+ union: set[int] = set()
187
+ for n in names[pos : pos + glen]:
188
+ self._group_of[n] = gi
189
+ c = self._class_of.get(n, 0)
190
+ # Bounds-checked: a bad class must surface from validate() as a
191
+ # named ValueError, not as an IndexError from this precompute.
192
+ if 0 <= c < len(self._owner_sets):
193
+ union.update(self._owner_sets[c])
194
+ pos += glen
195
+ self._group_owners.append(sorted(union))
196
+ # Bound at init on the consumer; the rule methods never touch these.
197
+ self._actors: list[Any] = []
198
+ self._produce_methods: list[Any] = []
199
+ self.consumer_id = 0
200
+
201
+ # ---------------- rules (pure; safe to use unbound) ----------------
202
+
203
+ def class_of(self, name: str) -> int:
204
+ """The name's owner class — the planner's bucketing key, since all names
205
+ of a chunk must share one producer."""
206
+ return self._class_of.get(name, 0)
207
+
208
+ def owners(self, name: str) -> list[int]:
209
+ """Every producer holding ``name``."""
210
+ return list(self._owner_sets[self.class_of(name)])
211
+
212
+ def group_owners(self, group_idx: int) -> list[int]:
213
+ """Every producer holding any name of ``group_idx``.
214
+
215
+ The one group-keyed rule, because the free barrier is per group: a
216
+ consumer signals ``free_group(gi)`` at each of these and the producer
217
+ counts signals against its live-consumer total. Names route; groups free.
218
+ """
219
+ if not self._group_owners:
220
+ return list(range(self.num_producers))
221
+ return list(self._group_owners[group_idx])
222
+
223
+ def producer_for(self, consumer_id: int, name: str) -> int:
224
+ """The single producer ``consumer_id`` pulls ``name`` from.
225
+
226
+ Blocks consumers across the name's owner set with the same rule that
227
+ binds producers globally, then rotates by the name's GROUP index so a
228
+ consumer spreads its groups over its block instead of hammering one NIC.
229
+ Rotating per group (not per name) keeps every name of a chunk on one
230
+ producer, which is what lets a chunk be a single pull.
231
+
232
+ The carve is over ONE DEPLOYMENT (this consumer's index within its own
233
+ deployment, over ``workers_per_replica`` of them), so the same worker of
234
+ every deployment resolves the same producer for every name. Carving over
235
+ the whole fleet instead spreads the multi-owner names of different
236
+ deployments onto different producers while single-owner names still
237
+ route by ownership, so a producer serves several distinct worker
238
+ indices, each needing its own serve ring, and none of them can share a
239
+ slot. With one deployment the width is the fleet, so this is the plain
240
+ block rule.
241
+ """
242
+ own = self.owners(name)
243
+ if not own:
244
+ raise ValueError(f"{name!r} has no owner")
245
+ block = assign_producer_indices(
246
+ len(own), self._block_consumers, consumer_id % self._block_consumers
247
+ )
248
+ return own[block[self._group_of.get(name, 0) % len(block)]]
249
+
250
+ def validate(self) -> None:
251
+ """Check the ownership tables can be served.
252
+
253
+ Raises:
254
+ ValueError: an owner set is empty or out of range, or a class index
255
+ does not resolve.
256
+ """
257
+ for c, row in enumerate(self._owner_sets):
258
+ if not row:
259
+ raise ValueError(f"owner set {c} is empty: no producer holds it")
260
+ bad = [p for p in row if not 0 <= p < self.num_producers]
261
+ if bad:
262
+ raise ValueError(f"owner set {c} out of range: {bad}")
263
+ for name, c in self._class_of.items():
264
+ if not 0 <= c < len(self._owner_sets):
265
+ raise ValueError(
266
+ f"{name!r} has owner class {c}, but only "
267
+ f"{len(self._owner_sets)} owner set(s) were shipped"
268
+ )
269
+
270
+ # ---------------- send surface (needs the binding) ----------------
271
+
272
+ def bind(self, actors: list, produce_methods: list, consumer_id: int) -> None:
273
+ """Attach this consumer's producer handles. Rebuilt wholesale per init,
274
+ never appended to: every owner index is a position in these lists, so a
275
+ rejoining engine that re-inits must not shift them."""
276
+ if len(actors) != len(produce_methods):
277
+ raise ValueError("actors and produce_methods must be parallel")
278
+ self._actors = list(actors)
279
+ self._produce_methods = list(produce_methods)
280
+ self.consumer_id = consumer_id
281
+
282
+ def _bound(self) -> None:
283
+ if not self._produce_methods:
284
+ raise RuntimeError("RdtRouter has no producers bound; call bind() first.")
285
+
286
+ def pull(self, owner: int, keys: list, seq: int) -> Any:
287
+ """Issue one packed pull to ``owner``.
288
+
289
+ The consumer id keys the producer's serve ring; without it every worker
290
+ is served out of ring 0 and concurrent pulls overwrite each other's blob.
291
+ ``seq`` is this call's index in the stream to ``owner`` (see ``_Chunk``):
292
+ it selects the serve slot in ISSUE order, so a slot is never repacked
293
+ while the read of its previous contents is still in flight.
294
+ """
295
+ self._bound()
296
+ return self._produce_methods[owner].remote(
297
+ keys, consumer_id=self.consumer_id, seq=seq
298
+ )
299
+
300
+ def free_group(self, group_idx: int) -> list:
301
+ """Signal the group done at every producer holding any of its names."""
302
+ self._bound()
303
+ return [
304
+ self._actors[p].free_group.remote(group_idx)
305
+ for p in self.group_owners(group_idx)
306
+ ]
307
+
308
+ def reserve_serve_buffers(
309
+ self, bytes_by_producer: list[int], plan_digests: list[str] | None = None
310
+ ) -> list:
311
+ """Ask each producer to pre-register a serve ring sized to the most this
312
+ consumer will pull from it.
313
+
314
+ ``plan_digests[p]`` describes the chunks this consumer pulls from
315
+ producer ``p``, in pull order. A producer that shares one serve ring
316
+ across the consumers of several deployments compares it across them, so
317
+ a fleet whose deployments are not identical fails at init instead of
318
+ stalling mid-sync. A producer that shares nothing ignores it.
319
+ """
320
+ self._bound()
321
+ digests: list = (
322
+ list(plan_digests) if plan_digests else [None] * len(bytes_by_producer)
323
+ )
324
+ return [
325
+ self._actors[p].reserve_serve_buffer.remote(
326
+ self.consumer_id, nb, digests[p]
327
+ )
328
+ for p, nb in enumerate(bytes_by_producer)
329
+ if nb > 0
330
+ ]
331
+
332
+
333
+def buffer_alloc_bytes(nbytes: int, presize: int = 0) -> int:
334
+ """Size a NIXL buffer / ring slot for ``nbytes``: the max of the request, an
335
+ optional ``presize`` floor, and a coarse 256MB round-up, so the buffer is
336
+ allocated ONCE and never regrows. Regrowth is a correctness hazard, not just a
337
+ perf one -- see ``buffer_presize_gb``. Shared by the consumer's receive buffers
338
+ and the producer's serve rings."""
339
+ return max(nbytes, presize, -(-nbytes // (256 << 20)) * (256 << 20))
vllm/distributed/weight_transfer/sharded_rdt_engine.py
new
+1479
@@ -0,0 +1,1479 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Sharded Ray Direct Transport (RDT) weight transfer engine (consumer side).
4
+
5
+Pulls only the slice each vLLM worker consumes under tensor/expert parallelism,
6
+not the full HF-format tensor.
7
+
8
+Two phases. BAKE, once at ``init_transfer_engine``: drive ``model.load_weights``
9
+against ``FakeRDTTensor`` placeholders and record, per leaf module, how each
10
+destination slice is fetched (an op chain) and where it lands (an ``as_strided``
11
+descriptor). REPLAY, every sync: no ``load_weights``, no ``FakeRDTTensor``
12
+dispatch, no discovery — pull the recorded slices in packed chunks over a ring
13
+of receive buffers, scatter them into freshly materialized params, then quant
14
+and kernel-copy. A live name with no recorded plan fails the plan build: there
15
+is no fallback load.
16
+
17
+Weights arrive in checkpoint format: the engine drives layerwise reload itself,
18
+in ``start_weight_update`` / ``finish_weight_update``.
19
+
20
+Data flow
21
+---------
22
+One thing at four resolutions, over three lifetimes. ``FetchKey`` --
23
+``(name, op_chain)``, "which slice of which trainer tensor" -- is the atom;
24
+everything else is bookkeeping around it.
25
+
26
+ BAKE once at init, kept for the engine's life
27
+ FakeRDTTensor intercepts the model's loaders; each copy_ records a
28
+ _Scatter (sharded_rdt_fake): src FetchKey, owning layer, the destination
29
+ as_strided region, and the produced dtype/nbytes.
30
+ -> _name_to_plan: name -> that module's scatter list
31
+
32
+ PLAN once, cached, _build_call_plan
33
+ _Chunk one packed pull: its scatters, deduped keys, the byte-exact
34
+ pack_layout, which producer serves it, and what to run after
35
+ (materialize / quant / free). One per (group, owner class).
36
+ _CallPlan all chunks + pre_free.
37
+
38
+ RUN per chunk, per sync, _run_chunk_pipeline
39
+ _Chunk -> _PendingPull issued, not yet landed: Ray ref, buffer views,
40
+ ring slot. [RPC thread]
41
+ -> _ProcItem chunk + results + slot. [hand-off to the
42
+ background scatter thread]
43
+
44
+The RUN pair stays split on purpose: it is the thread boundary, and only
45
+``targets`` should outlive the get -- carrying the Ray ref and the whole-buffer
46
+blob into the queue would keep both alive for the scatter's lifetime.
47
+
48
+See docs/training/weight_transfer/sharded_rdt.md for the design and the measured
49
+results behind the choices here.
50
+"""
51
+
52
+import time
53
+from dataclasses import dataclass, field
54
+from math import prod
55
+from typing import TYPE_CHECKING, Any, cast
56
+
57
+import torch
58
+
59
+from vllm.config.weight_transfer import WeightTransferConfig
60
+from vllm.distributed.weight_transfer.base import (
61
+ WeightTransferEngine,
62
+ WeightTransferInitInfo,
63
+ WeightTransferUpdateInfo,
64
+)
65
+from vllm.distributed.weight_transfer.sharded_rdt_common import (
66
+ RdtRouter,
67
+ buffer_alloc_bytes,
68
+ check_ray_rdt_version,
69
+)
70
+from vllm.distributed.weight_transfer.sharded_rdt_fake import (
71
+ BakeSink,
72
+ FakeRDTTensor,
73
+ FetchKey,
74
+ _Scatter,
75
+)
76
+from vllm.logger import init_logger
77
+
78
+if TYPE_CHECKING:
79
+ from vllm.config import VllmConfig
80
+
81
+logger = init_logger(__name__)
82
+
83
+
84
+@dataclass
85
+class _Chunk:
86
+ """One packed pull plus its post-processing (see Data flow).
87
+
88
+ A module's copies span chunks when its experts sit in several owner classes,
89
+ so ``materialize`` fires on its FIRST chunk and ``quant`` on its LAST --
90
+ materialize-once by construction, not by a runtime counter. ``pack_layout``
91
+ mirrors the producer's rule byte-exactly.
92
+ """
93
+
94
+ scatters: "list[_Scatter]"
95
+ keys: "list[FetchKey]"
96
+ pack_layout: "list[tuple[int, torch.dtype, int, tuple[int, ...]]]" # off,dt,n,shape
97
+ pack_bytes: int
98
+ materialize: "list[Any]"
99
+ quant: "list[Any]"
100
+ free: "list[int]"
101
+ # Trainer rank serving this chunk, resolved at plan time. One per chunk:
102
+ # every name in it shares an owner class, so one producer holds them all.
103
+ owner: int
104
+ # This chunk's index among the chunks this worker pulls from ``owner``, in
105
+ # issue order. It picks the producer's serve slot (``seq % ring_depth``),
106
+ # which is what makes slot reuse safe: the pipeline drains pull i before
107
+ # issuing i+K, so the slot seq+K takes was last used by seq, whose read is
108
+ # finished. Choosing the slot in EXECUTION order instead -- as a per-call
109
+ # counter on the producer would -- can hand a slot to a pull whose
110
+ # predecessor is still being read, because Ray may start a consumer's
111
+ # concurrent produce calls in any order. Static, since the plan is.
112
+ seq: int = 0
113
+
114
+
115
+@dataclass
116
+class _CallPlan:
117
+ """The static plan for one sync (see Data flow). Pure, so it is built once
118
+ and reused; runtime is then execution only.
119
+
120
+ ``pre_free`` = groups with NO chunk on this worker, signaled at sync start
121
+ (owners tolerate a signal preceding its publish). With the last-chunk signals
122
+ this keeps the completeness invariant consumer-local: every group is signaled
123
+ exactly once.
124
+ """
125
+
126
+ chunks: "list[_Chunk]"
127
+ pre_free: "list[int]"
128
+
129
+
130
+@dataclass
131
+class _PendingPull:
132
+ """A dispatched pull whose blocking ``ray.get`` has not run (see Data flow).
133
+
134
+ ``targets``/``blob`` must stay strongly referenced until it completes:
135
+ ``set_target_for_ref`` stores WEAKREFS, so dropping them silently reroutes the
136
+ transfer into a fallback buffer."""
137
+
138
+ ref: "Any"
139
+ keys: "list[FetchKey]"
140
+ targets: "list[torch.Tensor]"
141
+ blob: "torch.Tensor"
142
+ slot: int
143
+
144
+
145
+@dataclass
146
+class _ProcItem:
147
+ """A landed pull handed to the background scatter thread (see Data flow).
148
+
149
+ ``results`` alias the ring buffer ``slot``, held as strong refs so they
150
+ outlive the RPC-thread frame until the scatter consumes them.
151
+ """
152
+
153
+ chunk: "_Chunk"
154
+ results: "dict[FetchKey, torch.Tensor]"
155
+ slot: int
156
+
157
+
158
+@dataclass
159
+class ShardedRDTWeightTransferInitInfo(WeightTransferInitInfo):
160
+ """Initialization info for the sharded RDT backend."""
161
+
162
+ trainer_actor_names: list[str] = field(default_factory=list)
163
+ """Names of all trainer Ray actors exposing the producer method (set via
164
+ ``.options(name=...)``), ordered by trainer rank. ``RdtRouter`` picks one of
165
+ them per pull, out of the name's owner set. Every actor is bound regardless,
166
+ because ``free_group`` fans out to every owner. Must be non-empty; a
167
+ single-producer trainer passes a one-element list."""
168
+
169
+ trainer_actor_namespace: str | None = None
170
+ """Optional Ray namespace the trainer actor(s) live in."""
171
+
172
+ produce_method_name: str = "rdt_produce_weights_batched"
173
+ """Name of the trainer-side producer method. It and the rest of the serve
174
+ surface (``free_group``, ``reserve_serve_buffer``) are documented where they
175
+ are implemented, on ``_RDTProducerServer`` in ``sharded_rdt_trainer.py``."""
176
+
177
+ names: list[str] = field(default_factory=list)
178
+ """The trainer's complete, flat param name list. The bake drives
179
+ ``model.load_weights`` over all of them once and keys the plan by source
180
+ name."""
181
+
182
+ dtype_names: list[str] = field(default_factory=list)
183
+ """Dtype name (e.g. 'bfloat16') for each entry of ``names``."""
184
+
185
+ shapes: list[list[int]] = field(default_factory=list)
186
+ """Full HF shape for each entry of ``names``."""
187
+
188
+ group_lens: list[int] = field(default_factory=list)
189
+ """Partition of ``names`` into gather groups, in the SAME order the trainers
190
+ gather and publish them (group-major; ``sum(group_lens) == len(names)``, and
191
+ ``names`` must be ordered to match). Required: the engine pre-builds the whole
192
+ static chunk/signal plan from it at init, and fires ``free_group`` at every
193
+ owner as each group's last chunk completes."""
194
+
195
+ owner_sets: list[list[int]] = field(default_factory=list)
196
+ """The distinct producer sets that occur, each a sorted list of trainer ranks
197
+ (indices into ``trainer_actor_names``). Indexed by ``name_owner_class``.
198
+ Empty means every producer holds every name."""
199
+
200
+ name_owner_class: list[int] = field(default_factory=list)
201
+ """Per-name index into ``owner_sets``, parallel to ``names``: which producers
202
+ hold each name. Empty means every producer holds every name.
203
+
204
+ This one table expresses every layout the trainer can have — pipeline stages
205
+ (the names of these groups have this owner set), expert parallelism (this
206
+ expert name has this one-rank owner set), and combinations of the two — so
207
+ the engine cuts each group's baked copies into one chunk per distinct class
208
+ present and routes each chunk to that class's owner. Derived by the trainer
209
+ from ``WeightSource.held_names()``."""
210
+
211
+ num_consumers: int = 0
212
+ """Total consumer count across the fleet, for M:N routing. Authoritative when
213
+ > 0; at 0 the engine infers it from ``parallel_config``, which is correct for
214
+ the supported serving modes but worth setting explicitly under M:N. Each
215
+ worker's distinct index comes from ``_global_worker_index``."""
216
+
217
+ num_rdt_buffers: int = 2
218
+ """Depth of the consumer receive-buffer ring. Must match the
219
+ producer's — ``_run_chunk_pipeline``'s slot-safety argument rests on it. 2 =
220
+ double buffer: chunk i+1's serve overlaps chunk i's RDMA, and scatter(i-1)
221
+ overlaps RDMA(i) in the other slot. Keep depth x chunk_bytes under the
222
+ fabric's address-translation reach (~2-3 GB/flow on the reference 8xB200 RoCE
223
+ cluster, where K=3 measurably hurt)."""
224
+
225
+ buffer_presize_gb: float = 0.0
226
+ """Pre-size each packed receive-buffer slot to this many GiB
227
+ (0 = size to the first chunk + coarse 256MB round-up). Set it to cover the
228
+ model's largest atomic chunk (e.g. an untied lm_head). Sizing buffers ONCE
229
+ matters beyond perf -- see the doc's "Sizing buffers once matters beyond
230
+ throughput"."""
231
+
232
+ replica_rank: int = 0
233
+ """This inference engine's ordinal in the fleet (0..``num_replicas``-1).
234
+
235
+ Multi-engine deployments run several INDEPENDENT inference engines, each with
236
+ its own self-contained parallel config, so every engine's
237
+ ``_global_worker_index`` restarts at 0 and would collide across engines. The
238
+ driver gives each engine a distinct ``replica_rank`` (with identical
239
+ ``num_replicas``) so the engine offsets its consumers into a globally distinct
240
+ range for the M:N block assignment.Default 0/1 (single
241
+ engine) needs no override."""
242
+
243
+ num_replicas: int = 1
244
+ """Number of independent inference engines in the fleet. Default 1 => the
245
+ per-replica offset is 0 and consumer identity is exactly
246
+ ``_global_worker_index`` (preserves single-engine and single-DP-deployment
247
+ behavior). When > 1, ``workers_per_replica = num_consumers // num_replicas``
248
+ and this engine's consumers occupy
249
+ ``replica_rank * workers_per_replica + _global_worker_index()``. Assumes a
250
+ uniform fleet (every replica has the same worker count). Set alongside
251
+ ``replica_rank``, by the driver."""
252
+
253
+
254
+@dataclass
255
+class ShardedRDTWeightTransferUpdateInfo(WeightTransferUpdateInfo):
256
+ """Update info for the sharded RDT backend: intentionally EMPTY.
257
+
258
+ The chunk/free plan is a pure function of the baked plan and the driver's
259
+ gather-group partition, both fixed for the engine's lifetime, so it is built
260
+ once at ``init_transfer_engine`` from ``ShardedRDTWeightTransferInitInfo``'s
261
+ ``names`` + ``group_lens``. ONE ``update_weights`` per sync then just re-runs
262
+ that plan; there is nothing per-sync to carry.
263
+ """
264
+
265
+
266
+class ShardedRDTWeightTransferEngine(
267
+ WeightTransferEngine[
268
+ ShardedRDTWeightTransferInitInfo,
269
+ ShardedRDTWeightTransferUpdateInfo,
270
+ ]
271
+):
272
+ """Pull-based RDT/NIXL backend that transports only the slice each worker
273
+ consumes.
274
+
275
+ Requires ``distributed_executor_backend="ray"``, ``nixl`` in the shared env,
276
+ a named trainer actor exposing a ``@ray.method(tensor_transport="nixl")``
277
+ producer, and weight loaders that stay inside ``SUPPORTED_OPS`` — anything
278
+ needing real data (``.to``, ``.item``, arithmetic, bool-mask indexing)
279
+ raises during the bake.
280
+
281
+ The plan is baked once at ``init_transfer_engine`` into one scatter list
282
+ per fully-loaded leaf module, indexed by source name; every
283
+ ``update_weights`` replays the modules its gathered names cover.
284
+ """
285
+
286
+ init_info_cls = ShardedRDTWeightTransferInitInfo
287
+ update_info_cls = ShardedRDTWeightTransferUpdateInfo
288
+ # receive_weights pulls synchronously but defers GPU post-processing to a
289
+ # background thread so it overlaps the next chunk's pull, so ``update_weights``
290
+ # skips the base's device sync and ``finish_weight_update`` drains the
291
+ # deferred work before finalize.
292
+ defers_processing = True
293
+ # The baked replay plan is a function of one concrete model's parameter
294
+ # layout, so a separate draft model cannot reuse it.
295
+ supports_draft_weight_update = False
296
+
297
+ def __init__(
298
+ self,
299
+ config: WeightTransferConfig,
300
+ vllm_config: "VllmConfig",
301
+ device: torch.device,
302
+ model: torch.nn.Module,
303
+ ) -> None:
304
+ super().__init__(config, vllm_config, device, model)
305
+ # This engine's data plane is Ray's: the consumer resolves each trainer's
306
+ # serve actor with ``ray.get_actor`` and pulls over NIXL through the
307
+ # workers' tensor transport, neither of which exists unless the workers
308
+ # ARE Ray actors. Caught here, at construction, because the alternative
309
+ # is an opaque failure during the first handshake. A single-worker
310
+ # deployment defaults to "uni", so "ray" has to be asked for explicitly.
311
+ backend = vllm_config.parallel_config.distributed_executor_backend
312
+ # A custom Executor subclass is a deliberate override we cannot judge;
313
+ # only reject the named backends that are definitely not Ray.
314
+ if isinstance(backend, str) and backend != "ray":
315
+ raise ValueError(
316
+ f"The 'sharded_rdt' weight transfer backend requires "
317
+ f"distributed_executor_backend='ray', got {backend!r}. Pass "
318
+ f"--distributed-executor-backend ray (and keep "
319
+ f"VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1, the default, so the "
320
+ f"workers are RayExecutorV2 actors with tensor transport)."
321
+ )
322
+ # Ownership table, routing rules, and (once bound) every producer's
323
+ # handle — pulls go to one owner per chunk, but free_group fans out to
324
+ # every owner of the group, so all handles are bound.
325
+ self._router: RdtRouter | None = None
326
+ # Driver-supplied total consumer count (init_info.num_consumers); 0 => infer
327
+ # from parallel_config. See _num_consumers.
328
+ self._num_consumers_override: int = 0
329
+ # Source name -> the plan for the module consuming it. Several names of a
330
+ # fused module share one entry; replay dedups. A live name absent here
331
+ # fails the plan build.
332
+ self._name_to_plan: dict[str, list[_Scatter]] = {}
333
+ # name -> (dtype_name, shape) for every init name; the bake builds its
334
+ # lazies from this.
335
+ self._name_meta: dict[str, tuple[str, list[int]]] = {}
336
+ # Names whose copy_ fired during the bake. Unbaked names not here never
337
+ # move data (e.g. non-local EP experts), so receive_weights skips them.
338
+ self._live_names: set[str] = set()
339
+
340
+ # ---- Consumer-side pre-registered receive buffers ----------------------
341
+ # One persistent 1-D buffer per dtype per ring slot, STORAGE-registered with
342
+ # NIXL once, so every per-slice view hits Ray's registration cache and the
343
+ # recv path does zero registration in steady state. See the doc's "Sizing
344
+ # buffers once matters beyond throughput" for why they rarely grow.
345
+ #
346
+ # Ring depth is both the slot count and the in-flight pull count, one
347
+ # quantity from ``num_rdt_buffers``. Decoupling them with a spare slot
348
+ # measurably slows the sync.
349
+ self._ring_depth = 1
350
+ self._dest_buffers: list[dict[torch.dtype, torch.Tensor]] = [{}]
351
+ self._slot_read_done: list[Any] = [] # one torch.cuda.Event per slot
352
+ # Generation counters guarding the events (see _ensure_proc_worker).
353
+ self._slot_queued: list[int] = []
354
+ self._slot_done: list[int] = []
355
+ self._slot_cv: Any = None
356
+ self._pull_slot = 0
357
+ # (id(chunk), slot) -> (buffer data_ptr, per-key views).
358
+ # Lives on the engine, not on the _Chunk: the plan is static and shared
359
+ # across syncs, so a per-slot runtime cache has no business inside it.
360
+ self._targets_cache: dict[tuple[int, int], tuple[int, list[torch.Tensor]]] = {}
361
+ self._buffer_presize = 0 # bytes; set from init_info
362
+ self._pending_frees: list[Any] = [] # free_group signal refs, drained per sync
363
+ # The STATIC plan: built once — at init from init_info.group_lens, else
364
+ # lazily on the first update_weights — and reused. Non-None means
365
+ # update_weights ignores its per-sync names.
366
+ self._cached_plan: _CallPlan | None = None
367
+ # Completed syncs. The chunk pipeline runs SERIAL during sync 0, when both
368
+ # sides still register buffers, and pipelines from sync 1. See the doc.
369
+ self._completed_syncs = 0
370
+
371
+ # ---- Background post-processing thread (pull/process pipelining) -------
372
+ # The RPC thread hands pulled buffer views to this worker, which runs
373
+ # materialize/scatter/quant/kernel_copy on its own stream while the next
374
+ # pull proceeds. drain_pending() joins it before finalize.
375
+ self._proc_queue: Any | None = None
376
+ self._proc_thread: Any | None = None
377
+ self._proc_stream: Any | None = None
378
+ # Dedicated quant thread (pass 2): own queue + stream so quant never
379
+ # sits between two items' scatters on the scatter thread (measured to
380
+ # stall the RPC thread's slot handshake ~0.5-0.75s/iter).
381
+ self._quant_queue: Any | None = None
382
+ self._quant_thread: Any | None = None
383
+ self._quant_stream: Any | None = None
384
+ self._proc_error: BaseException | None = None
385
+
386
+ def init_transfer_engine(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
387
+ """Configure the ring, bind the producers, bake the replay plan, and
388
+ pre-register every NIXL buffer -- in that order, because each step depends
389
+ on the previous one.
390
+
391
+ The bake drives ``model.load_weights`` and the pre-registration blocks on
392
+ RPCs to the producers, so this is a heavyweight one-off; every later
393
+ ``update_weights`` is pure replay.
394
+ """
395
+ # Read the attribute directly: if a vLLM bump renames it, this guard
396
+ # must fail loudly at init rather than silently disappear.
397
+ if self.parallel_config.enable_eplb:
398
+ raise RuntimeError(
399
+ "sharded_rdt does not support EPLB (enable_eplb=true): dynamic "
400
+ "expert rearrangement invalidates the baked replay plan, which "
401
+ "records each expert's destination slot once at init — after a "
402
+ "rearrangement the replay would silently load weights into the "
403
+ "wrong expert slots. Disable EPLB or use another weight-transfer "
404
+ "backend."
405
+ )
406
+ # The worker is often a different install from whatever spawned it, so
407
+ # check here too and not only where the actor options are set.
408
+ check_ray_rdt_version()
409
+ self._configure_ring(init_info)
410
+ self._resolve_producers(init_info)
411
+
412
+ # A pure dry run: the trainer's gather cache is empty at init, so nothing
413
+ # can (or does) get pulled -- we only record how each slice is fetched and
414
+ # where it lands, then restore the model.
415
+ self._bake(init_info)
416
+ self._build_static_plan(init_info)
417
+ # Register ALL NIXL memory now, while the fabric is idle: dma-buf GPUDirect
418
+ # registration concurrent with in-flight RDMA intermittently fails
419
+ # (ibv_reg_mr 'Bad address'), which bites under M:N fan-in. Sizes come from
420
+ # the static plan, so this is exact.
421
+ self._preregister_at_init()
422
+ # Start the background post-processing worker (pull/process pipelining).
423
+ self._ensure_proc_worker()
424
+
425
+ def _configure_ring(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
426
+ """Ring depth K.
427
+
428
+ Must run before ``_ensure_proc_worker`` creates the per-slot events and
429
+ counters, and before any buffer is grown (both happen on the first pull).
430
+ """
431
+ self._num_consumers_override = int(init_info.num_consumers or 0)
432
+ k = max(1, int(init_info.num_rdt_buffers))
433
+ self._ring_depth = k
434
+ self._dest_buffers = [{} for _ in range(k)]
435
+ self._targets_cache = {}
436
+ self._buffer_presize = int(float(init_info.buffer_presize_gb) * (1 << 30))
437
+ logger.info(
438
+ "[RDT-RING] active_pulls=%d slots=%d presize=%.2fGiB",
439
+ k,
440
+ k,
441
+ self._buffer_presize / (1 << 30),
442
+ )
443
+
444
+ def _resolve_producers(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
445
+ """Work out this worker's consumer identity, build the router, and bind
446
+ EVERY producer actor.
447
+
448
+ Pull routing is M:N: each chunk goes to ONE producer holding every name
449
+ in it (see ``RdtRouter``), so one producer serves a whole pull. All
450
+ producers are bound regardless, because the per-group ``free_group``
451
+ signal fans out to every owner of a group, including producers this
452
+ worker never pulls from.
453
+ """
454
+ try:
455
+ import ray
456
+ except ImportError as e:
457
+ raise RuntimeError(
458
+ "Ray is required for the 'sharded_rdt' weight transfer "
459
+ "backend. Install Ray and run workers as Ray actors "
460
+ "(distributed_executor_backend='ray')."
461
+ ) from e
462
+
463
+ producer_names = list(init_info.trainer_actor_names)
464
+ if not producer_names:
465
+ raise RuntimeError(
466
+ "Sharded RDT engine requires a trainer producer: set "
467
+ "init_info.trainer_actor_names."
468
+ )
469
+
470
+ consumer_id = self._resolve_consumer_id(init_info)
471
+ name_owner_class = list(init_info.name_owner_class or [])
472
+ if name_owner_class and len(name_owner_class) != len(init_info.names):
473
+ raise RuntimeError(
474
+ f"Sharded RDT engine: {len(name_owner_class)} owner-class entries "
475
+ f"for {len(init_info.names)} names."
476
+ )
477
+ router = RdtRouter(
478
+ len(producer_names),
479
+ self._num_consumers(),
480
+ list(init_info.owner_sets) or None,
481
+ name_owner_class or None,
482
+ list(init_info.names),
483
+ list(init_info.group_lens),
484
+ # Carve the block per deployment so every replica's worker w pulls
485
+ # each name from the same producer, which is what lets that producer
486
+ # serve all of them from one slot.
487
+ workers_per_replica=self._workers_per_replica(init_info),
488
+ )
489
+ router.validate()
490
+
491
+ # A fresh router per init, never appended to: every owner index is a
492
+ # position in these lists, so a rejoining engine's second
493
+ # init_transfer_engine must not shift them.
494
+ actors: list[Any] = []
495
+ methods: list[Any] = []
496
+
497
+ for chosen_name in producer_names:
498
+ try:
499
+ actor = ray.get_actor(
500
+ chosen_name,
501
+ namespace=init_info.trainer_actor_namespace,
502
+ )
503
+ except ValueError as e:
504
+ raise RuntimeError(
505
+ f"Sharded RDT engine could not find trainer actor "
506
+ f"{chosen_name!r} (namespace="
507
+ f"{init_info.trainer_actor_namespace!r})."
508
+ ) from e
509
+ # ``ray.get_actor`` resolves by NAME, in a process that never
510
+ # imported the producer's class, and ``enable_tensor_transport``
511
+ # lives in that class's metadata: the creating process passes
512
+ # ``meta.enable_tensor_transport``, and Ray infers it from the
513
+ # class's ``@ray.method(tensor_transport=...)`` decorators. A
514
+ # name-resolved handle therefore always reports False and the
515
+ # dispatch guard rejects the pull, even though the trainer set the
516
+ # option. Inherent to resolving by name across processes, not a
517
+ # version bug; it fails on Ray 2.56.0.
518
+ #
519
+ # Forcing it skips no validation. Ray's next guard,
520
+ # ``actor_has_tensor_transport``, asks the LIVE actor whether it can
521
+ # build a NIXL agent and still runs, and ``_spawn_server`` hardcodes
522
+ # the actor option, so a misconfigured producer cannot reach here.
523
+ actor._ray_enable_tensor_transport = True
524
+ actors.append(actor)
525
+ methods.append(getattr(actor, init_info.produce_method_name))
526
+ router.bind(actors, methods, consumer_id)
527
+ self._router = router
528
+ logger.info(
529
+ "Sharded RDT engine (consumer %d) bound to all %d producers "
530
+ "(batched method %r, %d owner class(es))",
531
+ consumer_id,
532
+ len(producer_names),
533
+ init_info.produce_method_name,
534
+ len(init_info.owner_sets) or 1,
535
+ )
536
+
537
+ def _workers_per_replica(self, init_info: ShardedRDTWeightTransferInitInfo) -> int:
538
+ """Consumers per inference deployment, assuming a uniform fleet.
539
+
540
+ Read by this worker's consumer id and by the router's block carve, which
541
+ must not disagree, so it is derived once here.
542
+ """
543
+ return max(1, self._num_consumers() // max(1, int(init_info.num_replicas or 1)))
544
+
545
+ def _resolve_consumer_id(self, init_info: ShardedRDTWeightTransferInitInfo) -> int:
546
+ """This worker's DISTINCT index in 0..C-1 across the whole fleet.
547
+
548
+ Within one engine that is ``_global_worker_index()``. But a fleet of
549
+ INDEPENDENT engines (each with its own parallel config) restarts that
550
+ index at 0 per engine, so each engine offsets into its own range using
551
+ ``replica_rank``: with a uniform fleet,
552
+ ``workers_per_replica = C // num_replicas``. ``num_replicas`` defaults to
553
+ 1 (offset 0), preserving single-engine and single-DP-deployment behaviour.
554
+ """
555
+ replica_rank = max(0, int(init_info.replica_rank or 0))
556
+ return (
557
+ replica_rank * self._workers_per_replica(init_info)
558
+ + self._global_worker_index()
559
+ )
560
+
561
+ def _build_static_plan(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
562
+ """Build the chunk/free plan once. It never changes across syncs, so
563
+ ``update_weights`` needs no per-sync names."""
564
+ if not init_info.group_lens:
565
+ raise ValueError(
566
+ "Sharded RDT engine requires init_info.group_lens (the gather-"
567
+ "group partition of init_info.names)."
568
+ )
569
+ if sum(init_info.group_lens) != len(init_info.names):
570
+ raise ValueError(
571
+ f"init_info.group_lens sums to {sum(init_info.group_lens)} "
572
+ f"but {len(init_info.names)} names were given."
573
+ )
574
+ self._cached_plan = self._build_call_plan(init_info.names, init_info.group_lens)
575
+ logger.info(
576
+ "[RDT-PLAN] pre-built static call plan at init: %d chunks",
577
+ len(self._cached_plan.chunks),
578
+ )
579
+
580
+ def _preregister_at_init(self) -> None:
581
+ """Register every NIXL buffer this worker will use at init, before any
582
+ transfer runs, so nothing registers during the sync-0 RDMA churn.
583
+
584
+ Both sides are sized from the static plan: receive buffers are
585
+ ``ring_depth`` slots at the largest chunk's ``pack_bytes``, and each bound
586
+ producer is asked to pre-register a serve ring at the max bytes this
587
+ consumer will pull from it. A no-op when this worker has no chunks."""
588
+ plan = self._cached_plan
589
+ if plan is None or not plan.chunks:
590
+ return
591
+ from ray.experimental import register_nixl_memory
592
+
593
+ # (a) consumer receive buffers — one per ring slot, at the largest chunk.
594
+ max_pack = max(c.pack_bytes for c in plan.chunks)
595
+ alloc = buffer_alloc_bytes(max_pack, self._buffer_presize)
596
+ for slot in range(self._ring_depth):
597
+ buffer = self._dest_buffers[slot].get(torch.uint8)
598
+ if buffer is None or buffer.numel() < alloc:
599
+ buffer = torch.empty(alloc, dtype=torch.uint8, device=self.device)
600
+ register_nixl_memory(buffer)
601
+ self._dest_buffers[slot][torch.uint8] = buffer
602
+
603
+ # (b) producer serve rings — max bytes this consumer pulls from each
604
+ # bound producer, plus the per-producer plan digest a producer sharing
605
+ # one ring across deployments checks them against.
606
+ assert self._router is not None
607
+ serve_bytes = [0] * self._router.num_producers
608
+ per_owner_keys: list[list] = [[] for _ in range(self._router.num_producers)]
609
+ for c in plan.chunks:
610
+ serve_bytes[c.owner] = max(serve_bytes[c.owner], c.pack_bytes)
611
+ per_owner_keys[c.owner].append(c.keys)
612
+ digests = [_plan_digest(keys) for keys in per_owner_keys]
613
+ import ray
614
+
615
+ refs = self._router.reserve_serve_buffers(serve_bytes, digests)
616
+ if refs:
617
+ ray.get(refs) # block until every serve ring is registered
618
+ logger.info(
619
+ "[RDT-PLAN] pre-registered %d receive slots (%.0f MiB each) + serve "
620
+ "rings on %d producer(s) %s",
621
+ self._ring_depth,
622
+ alloc / (1 << 20),
623
+ len(refs),
624
+ [nb // (1 << 20) for nb in serve_bytes],
625
+ )
626
+
627
+ def _global_worker_index(self) -> int:
628
+ """This worker's stable, distinct global index across the inference fleet:
629
+ ``data_parallel_index * world_size + rank`` over the TP*PP world.
630
+
631
+ ``data_parallel_index``, not ``data_parallel_rank``: vLLM resets the
632
+ latter to 0 in a dense worker but keeps the former as the distinct global
633
+ DP rank. Same formula as the sibling ``nccl_engine``, so dense-via-TP and
634
+ MoE-via-DP+EP both yield distinct 0..C-1."""
635
+ pc = self.parallel_config
636
+ return pc.data_parallel_index * pc.world_size + pc.rank # world_size = TP*PP
637
+
638
+ def _num_consumers(self) -> int:
639
+ """Total inference-worker count. Prefers the driver-supplied
640
+ ``init_info.num_consumers`` (authoritative -- the driver knows the whole
641
+ fleet); else ``world_size_across_dp``, the same stride
642
+ ``_global_worker_index`` indexes with, so the two agree at any pp."""
643
+ if self._num_consumers_override > 0:
644
+ return self._num_consumers_override
645
+ return self.parallel_config.world_size_across_dp
646
+
647
+ def start_weight_update(self) -> None:
648
+ """Put the model's params on meta so layerwise reload streams them in
649
+ as each layer's slices land. Baked replay uses checkpoint format."""
650
+ from vllm.model_executor.model_loader.reload import (
651
+ initialize_layerwise_reload,
652
+ )
653
+
654
+ initialize_layerwise_reload(self.model)
655
+
656
+ def finish_weight_update(self) -> None:
657
+ """Drain the deferred pull/process pipeline (so every layer is fully
658
+ loaded) before finalizing the layerwise reload."""
659
+ from vllm.model_executor.model_loader.reload import (
660
+ finalize_layerwise_reload,
661
+ )
662
+
663
+ self.drain_pending()
664
+ finalize_layerwise_reload(self.model, self.model_config)
665
+
666
+ def update_weights(self, update_info: dict[str, Any]) -> None:
667
+ """Receive one update. Unlike the base, does NOT issue a per-update
668
+ device sync: post-processing is deferred to background threads and a
669
+ sync here would block on them and serialize the pull/process pipeline.
670
+ Completion is guaranteed by ``drain_pending`` in
671
+ ``finish_weight_update``."""
672
+ self.receive_weights(self.parse_update_info(update_info))
673
+
674
+ def receive_weights(
675
+ self,
676
+ update_info: ShardedRDTWeightTransferUpdateInfo,
677
+ ) -> None:
678
+ """Pull + replay the baked leaf modules the sync covers.
679
+
680
+ The chunk/free plan is STATIC across syncs — a pure function of the baked
681
+ plan and the driver's group partition — so it was built once at init and
682
+ every sync just re-runs the pipeline over its self-describing chunks, with
683
+ no per-sync bookkeeping and an empty ``update_info``.
684
+
685
+ Assumes each baked module's source names fall within one gather group,
686
+ which the per-layer / pre / post partition guarantees (a leaf module's
687
+ sources all live in one decoder layer). A module that did span groups
688
+ would be planned once per group and could pull a name whose group the
689
+ pipeline already freed, which parks the pull until the producer's stall
690
+ watchdog fires.
691
+ """
692
+ del update_info # the plan is static; nothing arrives per sync
693
+ if self._router is None:
694
+ raise RuntimeError(
695
+ "Sharded RDT engine not initialized. Call init_transfer_engine() first."
696
+ )
697
+ # Surface any error the background thread hit on a prior item promptly.
698
+ self._raise_proc_error()
699
+ if self._cached_plan is None:
700
+ raise RuntimeError(
701
+ "Sharded RDT engine has no call plan: init_info.group_lens must "
702
+ "be supplied at init_transfer_engine()."
703
+ )
704
+ self._run_chunk_pipeline(self._cached_plan)
705
+
706
+ def _build_fake_weights(
707
+ self,
708
+ names: list[str],
709
+ sink: "BakeSink",
710
+ device: torch.device,
711
+ ) -> list[tuple[str, torch.Tensor]]:
712
+ """Zero-storage lazies for ``names``, dtype/shape from the init metadata,
713
+ all feeding the bake's recording sink."""
714
+ return [
715
+ (
716
+ name,
717
+ FakeRDTTensor(
718
+ name=name,
719
+ shape=torch.Size(self._name_meta[name][1]),
720
+ dtype=_dtype_from_name(self._name_meta[name][0]),
721
+ device=device,
722
+ sink=sink,
723
+ ),
724
+ )
725
+ for name in names
726
+ ]
727
+
728
+ def _issue_pull(self, chunk: "_Chunk", slot: int) -> "_PendingPull":
729
+ """Reserve ``slot``, lay the targets out in its buffer, dispatch the
730
+ produce RPC and point the transfer at the buffer — WITHOUT the blocking
731
+ ``ray.get`` (that is ``_complete_pull``). The chunked pipeline issues
732
+ chunk i+1 before completing chunk i, so the producer serves the next
733
+ chunk while the in-flight RDMA streams.
734
+
735
+ Slot-reuse guard, both stages required: a generation wait (the CUDA event
736
+ binds only to its LAST record, so synchronizing before the background
737
+ thread recorded this item's event passes silently — observed as
738
+ nondeterministic weight corruption), then the event synchronize. It must
739
+ precede ``set_target_for_ref``, not just the get: the transfer may start
740
+ any time after the metadata push. See the doc's "slot generation
741
+ handshake".
742
+ """
743
+ from ray.experimental import register_nixl_memory, set_target_for_ref
744
+
745
+ if self._slot_read_done:
746
+ with self._slot_cv:
747
+ while self._slot_done[slot] < self._slot_queued[slot]:
748
+ self._slot_cv.wait(timeout=1.0)
749
+ self._slot_read_done[slot].synchronize()
750
+
751
+ # Every slice is byte-packed into ONE uint8 buffer, 16B-aligned in
752
+ # keys order, under one NIXL descriptor. ``chunk.pack_layout`` mirrors the
753
+ # producer's rule byte-exactly; ``targets`` carves the dtype views back out.
754
+ keys = chunk.keys
755
+ cur = chunk.pack_bytes
756
+ buffers = self._dest_buffers[slot]
757
+ buffer = buffers.get(torch.uint8)
758
+ if buffer is None or buffer.numel() < cur:
759
+ # Size ONCE with headroom: Ray's desc cache is keyed by data_ptr and
760
+ # outlives its tensor, so a regrowth can false-hit a recycled pointer
761
+ # and skip registering the new extent -> NIXL_ERR_NOT_FOUND.
762
+ alloc = buffer_alloc_bytes(cur, self._buffer_presize)
763
+ buffer = torch.empty(alloc, dtype=torch.uint8, device=self.device)
764
+ register_nixl_memory(buffer)
765
+ buffers[torch.uint8] = buffer
766
+ targets = self._pull_targets(chunk, slot, buffer)
767
+ # One produce RPC serves the whole packed chunk, since the group's owner
768
+ # holds every slice. The buffer view must stay strongly referenced through
769
+ # the get: set_target_for_ref stores WEAKREFS, and a dropped target
770
+ # reroutes the transfer into a fallback buffer.
771
+ blob = buffer[:cur]
772
+ assert self._router is not None
773
+ ref = self._router.pull(chunk.owner, keys, chunk.seq)
774
+ set_target_for_ref(ref, [blob])
775
+ return _PendingPull(
776
+ ref=ref,
777
+ keys=keys,
778
+ targets=targets,
779
+ blob=blob,
780
+ slot=slot,
781
+ )
782
+
783
+ def _pull_targets(
784
+ self, chunk: "_Chunk", slot: int, buffer: torch.Tensor
785
+ ) -> "list[torch.Tensor]":
786
+ """Per-key dtype views into ``slot``'s buffer.
787
+
788
+ The packed layout is static, so the views are built once per (chunk, slot)
789
+ rather than once per pull -- rebuilding them cost ~1150 Python ops per pull
790
+ at 235B. Keyed on the buffer pointer as well, so a regrow invalidates
791
+ instead of handing back views into a freed buffer.
792
+ """
793
+ cached = self._targets_cache.get((id(chunk), slot))
794
+ if cached is not None and cached[0] == buffer.data_ptr():
795
+ return cached[1]
796
+ targets = [
797
+ buffer[off : off + n * dt.itemsize].view(dt).reshape(shape)
798
+ for off, dt, n, shape in chunk.pack_layout
799
+ ]
800
+ self._targets_cache[(id(chunk), slot)] = (buffer.data_ptr(), targets)
801
+ return targets
802
+
803
+ def _complete_pull(self, pending: "_PendingPull") -> "dict[FetchKey, torch.Tensor]":
804
+ """Blocking half of a pull: the NIXL read lands during this ``ray.get``."""
805
+ import ray
806
+
807
+ ray.get(pending.ref)
808
+ return dict(zip(pending.keys, pending.targets))
809
+
810
+ # ---------------- Bake (dry run, at init) / replay ----------------
811
+
812
+ def _bake(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
813
+ """Bake the replay plan once, as a self-driven meta dry run.
814
+
815
+ Puts the params on meta, then drives ``model.load_weights`` over
816
+ ``init_info.names`` through the model's ORIGINAL loaders (the stamps
817
+ bypass ``online_process_loader``, so ``_layerwise_process`` is never in
818
+ the path). Nothing materializes or pulls; the fake's ``copy_`` records
819
+ the source op chain and the meta destination's geometry. Afterwards one
820
+ scatter list per fully-loaded leaf module (copied numel == loadable
821
+ size) is indexed by source name; a partial or unrecordable module fails
822
+ the plan build. The model is restored.
823
+
824
+ This leans on layerwise internals a public API should expose first-class:
825
+ a currently-loading hook instead of monkeypatched stamps, a dry-run mode
826
+ instead of bypassing ``online_process_loader``, and an
827
+ ``abort_layerwise_reload`` instead of ``_restore_after_dry_run``.
828
+ """
829
+ from vllm.model_executor.model_loader.reload.layerwise import (
830
+ initialize_layerwise_reload,
831
+ )
832
+ from vllm.model_executor.model_loader.reload.utils import get_layer_size
833
+
834
+ names, dtype_names, shapes = (
835
+ init_info.names,
836
+ init_info.dtype_names,
837
+ init_info.shapes,
838
+ )
839
+ self._name_meta = {n: (d, s) for n, d, s in zip(names, dtype_names, shapes)}
840
+ if not names:
841
+ return
842
+
843
+ model = self.model
844
+ recorder = BakeSink()
845
+
846
+ _t0 = time.perf_counter()
847
+ with torch.device(self.device):
848
+ # Meta-restore params + save kernel tensors (we bypass the loader
849
+ # wrapping it installs, below).
850
+ initialize_layerwise_reload(model)
851
+ # Stamp the *original* loaders (bypassing online_process_loader), so
852
+ # the single load pass runs the loaders on meta and records via the
853
+ # fake's copy_ — with no inline _layerwise_process, no deferral.
854
+ self._install_recording_stamps(model, recorder)
855
+ model.load_weights(self._build_fake_weights(names, recorder, self.device))
856
+ # Keep only fully-loaded modules (copied numel >= loadable size, the
857
+ # test online_process_loader uses): a partial module leaves unwritten
858
+ # regions that finalize would init, so baking it scatters garbage.
859
+ for module, recorded in recorder.copies_by_layer.items():
860
+ if not recorded or any(c is None for c in recorded):
861
+ continue # unrecordable copy_ -> slow path
862
+ # Guard above guarantees every entry is a real _Scatter.
863
+ copies = cast("list[_Scatter]", recorded)
864
+ copied = sum(prod(c.shape) for c in copies)
865
+ if copied < get_layer_size(module):
866
+ continue # partial -> slow path
867
+ for c in copies:
868
+ # Every name of the module shares ONE list, so identity
869
+ # dedups them back to one module at plan time.
870
+ self._name_to_plan[c.src[0]] = copies
871
+ self._restore_after_dry_run(model)
872
+
873
+ # Names whose copy_ fired during the bake. Names not in here no-op for
874
+ # this worker (e.g. foreign-EP experts) and are skipped; live names that
875
+ # did not bake fail the plan build.
876
+ self._live_names = set(recorder.copied_names)
877
+
878
+ n_groups = len({id(g) for g in self._name_to_plan.values()})
879
+ logger.info(
880
+ "Sharded RDT dry-run baked %d/%d names into %d leaf modules "
881
+ "(%d live) in %.3fs",
882
+ len(self._name_to_plan),
883
+ len(names),
884
+ n_groups,
885
+ len(self._live_names),
886
+ time.perf_counter() - _t0,
887
+ )
888
+
889
+ def _install_recording_stamps(
890
+ self, model: torch.nn.Module, recorder: "BakeSink"
891
+ ) -> None:
892
+ """Wrap each loadable param's ``weight_loader`` to stamp
893
+ ``recorder.current = (leaf_module, param_name)`` before delegating to the
894
+ original loader, so the fake's ``copy_`` can attribute each recorded copy.
895
+ ``functools.wraps`` keeps the loader's real signature (so vLLM's
896
+ ``_layerwise_process`` ``param`` redirect still works if a stamp leaks),
897
+ and ``_rdt_stamp_inner`` tags it so ``_restore_after_dry_run`` can unwrap it.
898
+ """
899
+ import functools
900
+
901
+ from vllm.model_executor.model_loader.reload.layerwise import (
902
+ _get_original_loader,
903
+ )
904
+ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
905
+ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
906
+
907
+ def _make_stamp(layer, name, inner, added=False):
908
+ @functools.wraps(inner) # keep ``inner``'s signature (incl. ``param``)
909
+ def stamp(*args, **kwargs):
910
+ recorder.current = (layer, name)
911
+ try:
912
+ return inner(*args, **kwargs)
913
+ finally:
914
+ recorder.current = None
915
+
916
+ # Tag so _restore_after_dry_run can detect and unwrap leaked stamps,
917
+ # and so a second bake doesn't double-wrap.
918
+ stamp._rdt_stamp_inner = inner # type: ignore[attr-defined]
919
+ stamp._rdt_stamp_added = added # type: ignore[attr-defined]
920
+ return stamp
921
+
922
+ for module in model.modules():
923
+ for name, tensor in get_layer_tensors(module).items():
924
+ if getattr(tensor, "weight_loader", None) is None:
925
+ # A param with NO loader (e.g. GLM's router bias, a plain
926
+ # nn.Parameter) is still loaded by the model's load_weights
927
+ # through the getattr(param, "weight_loader",
928
+ # default_weight_loader) fallback — unstamped, its bake copy
929
+ # would be unattributable, failing the module's coverage
930
+ # gate and the plan build. Stamp the same default loader
931
+ # the fallback would pick; the restore deletes it again.
932
+ tensor.weight_loader = _make_stamp(
933
+ module, name, default_weight_loader, added=True
934
+ )
935
+ continue
936
+ # Bypass online_process_loader: stamp the *original* loader.
937
+ original = _get_original_loader(tensor)
938
+ tensor.weight_loader = _make_stamp(module, name, original)
939
+
940
+ def _restore_after_dry_run(self, model: torch.nn.Module) -> None:
941
+ """Restore each layerwise layer's saved kernel tensors without pulling
942
+ (a real ``finalize_layerwise_reload`` would materialize/load) and reset
943
+ its info. Also unwrap any recording ``stamp`` left on the params, since a
944
+ leaked stamp would sit under the next sync's ``online_process_loader`` and
945
+ silently break ``_layerwise_process``'s ``param`` redirect.
946
+ """
947
+ from vllm.model_executor.model_loader.reload.layerwise import (
948
+ LAYERWISE_INFO,
949
+ _place_kernel_tensors,
950
+ )
951
+ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
952
+
953
+ for layer in model.modules():
954
+ info = LAYERWISE_INFO.get(layer)
955
+ if info is not None and info.can_load():
956
+ if info.kernel_tensors is not None:
957
+ _place_kernel_tensors(layer, info)
958
+ info.reset()
959
+ # Unwrap any recording stamps left on the (now-restored) params so they
960
+ # never leak into a later update_weights. ``_rdt_stamp_inner`` is set by
961
+ # ``_install_recording_stamps``; unwrap repeatedly in case of nesting.
962
+ for module in model.modules():
963
+ for _name, tensor in get_layer_tensors(module).items():
964
+ loader = getattr(tensor, "weight_loader", None)
965
+ added = False
966
+ while loader is not None and hasattr(loader, "_rdt_stamp_inner"):
967
+ added = added or getattr(loader, "_rdt_stamp_added", False)
968
+ loader = loader._rdt_stamp_inner
969
+ tensor.weight_loader = loader
970
+ if added:
971
+ # The stamp was ATTACHED to a param that had no loader
972
+ # (see _install_recording_stamps); leave none behind.
973
+ del tensor.weight_loader
974
+ if hasattr(model, "_original_do_torchao_reload"):
975
+ model._do_torchao_reload = model._original_do_torchao_reload
976
+
977
+ # ---------------- Background post-processing (pull/process pipeline) -------
978
+
979
+ def _ensure_proc_worker(self) -> None:
980
+ """Lazily create the per-slot events, the background CUDA stream, the work
981
+ queue, and the single processing thread. Idempotent."""
982
+ if self._proc_thread is not None:
983
+ return
984
+ import queue
985
+ import threading
986
+
987
+ self._slot_read_done = [torch.cuda.Event() for _ in range(self._ring_depth)]
988
+ # Generation handshake for the events above: the RPC thread counts items
989
+ # queued per slot, this thread counts records, and a pull may synchronize()
990
+ # only once done[slot] has caught up with queued[slot]. Without it the
991
+ # synchronize can bind to a stale record and pass, letting the next RDMA
992
+ # overwrite a slot under a pending scatter. See the doc's "The slot
993
+ # generation handshake".
994
+ self._slot_queued = [0] * self._ring_depth
995
+ self._slot_done = [0] * self._ring_depth
996
+ self._slot_cv = threading.Condition()
997
+ self._proc_stream = torch.cuda.Stream(device=self.device)
998
+ self._proc_queue = queue.Queue()
999
+ self._proc_error = None
1000
+ t = threading.Thread(
1001
+ target=self._proc_worker_loop, name="rdt-postprocess", daemon=True
1002
+ )
1003
+ self._proc_thread = t
1004
+ t.start()
1005
+ self._quant_stream = torch.cuda.Stream(device=self.device)
1006
+ self._quant_queue = queue.Queue()
1007
+ qt = threading.Thread(
1008
+ target=self._quant_worker_loop, name="rdt-quant", daemon=True
1009
+ )
1010
+ self._quant_thread = qt
1011
+ qt.start()
1012
+
1013
+ def _proc_worker_loop(self) -> None:
1014
+ """Single persistent thread: run each queued item's process phase on the
1015
+ background stream. Exits on the ``None`` sentinel (shutdown). An item that
1016
+ raises is recorded in ``_proc_error`` and re-raised on the RPC thread /
1017
+ at drain, so a failed sync fails loudly rather than corrupting silently.
1018
+ """
1019
+ torch.cuda.set_device(self.device)
1020
+ q = self._proc_queue
1021
+ assert q is not None
1022
+ while True:
1023
+ item = q.get()
1024
+ try:
1025
+ if item is None:
1026
+ return
1027
+ # _process_item publishes the slot generation internally (in a
1028
+ # finally around its scatter pass), so an error here still
1029
+ # unblocks the RPC thread's generation wait.
1030
+ self._process_item(item)
1031
+ except BaseException as e: # noqa: BLE001 - surfaced on the RPC thread
1032
+ self._proc_error = e
1033
+ logger.exception("RDT background post-processing failed")
1034
+ finally:
1035
+ q.task_done()
1036
+
1037
+ def _mark_slot_done(self, slot: int) -> None:
1038
+ """Publish that a queued item's read-done event has been recorded (or the
1039
+ item failed) so a pull waiting to reuse ``slot`` can proceed to its
1040
+ CUDA-event synchronize."""
1041
+ with self._slot_cv:
1042
+ self._slot_done[slot] += 1
1043
+ self._slot_cv.notify_all()
1044
+
1045
+ def _raise_proc_error(self) -> None:
1046
+ """Re-raise (once) any error captured by the background thread."""
1047
+ if self._proc_error is not None:
1048
+ err = self._proc_error
1049
+ self._proc_error = None
1050
+ raise RuntimeError("RDT background post-processing failed") from err
1051
+
1052
+ def drain_pending(self) -> None:
1053
+ """Block until the background thread has processed every queued item and
1054
+ its stream work is complete, then re-raise any error it hit. Called from
1055
+ the worker's ``finish_weight_update`` before ``finalize_layerwise_reload``
1056
+ so every baked layer is fully loaded (and ``info.reset()``-ed) first."""
1057
+ if self._proc_queue is not None:
1058
+ self._proc_queue.join() # every put() item task_done()'d
1059
+ # The scatter thread feeds the quant thread, so join it SECOND (all
1060
+ # completed-group batches have been put by now), then sync both streams
1061
+ # so finalize sees fully-materialized, quanted, reset layers.
1062
+ if self._quant_queue is not None:
1063
+ self._quant_queue.join()
1064
+ if self._proc_stream is not None:
1065
+ self._proc_stream.synchronize()
1066
+ if self._quant_stream is not None:
1067
+ self._quant_stream.synchronize()
1068
+ self._raise_proc_error()
1069
+ # Ensure every fired free_group signal has EXECUTED on
1070
+ # the producer before the sync ends: ``begin_sync`` resets the producer's
1071
+ # per-group signal counts, so a signal landing after the next sync
1072
+ # started would credit a group it does not belong to, over-crediting the
1073
+ # gather-lookahead backpressure (extra gather resident -> trainer OOM
1074
+ # risk).
1075
+ if self._pending_frees:
1076
+ import ray
1077
+
1078
+ try:
1079
+ ray.get(self._pending_frees)
1080
+ finally:
1081
+ self._pending_frees.clear()
1082
+ # One sync iteration fully drained; buffers/registrations are at (or
1083
+ # nearer) high-water — the chunk pipeline may issue ahead from now on.
1084
+ self._completed_syncs += 1
1085
+
1086
+ def _dispatch_item(self, item: "_ProcItem") -> None:
1087
+ """Hand one chunk item to the background scatter thread.
1088
+
1089
+ Counts the item against its slot BEFORE dispatch: the next pull into
1090
+ that slot must wait until the background thread has processed (and
1091
+ RECORDED the read-done event for) every item ever queued on it.
1092
+ """
1093
+ with self._slot_cv:
1094
+ self._slot_queued[item.slot] += 1
1095
+ assert self._proc_queue is not None
1096
+ self._proc_queue.put(item)
1097
+
1098
+ def _chunk_module_scatters(
1099
+ self, modules: "list[list[_Scatter]]"
1100
+ ) -> "list[tuple[int, list[_Scatter]]]":
1101
+ """Cut the modules' copies into one chunk per distinct owner class
1102
+ present, ascending by class index, as ``(class_idx, scatters)`` pairs.
1103
+
1104
+ A chunk is one packed pull, so every name in it must share a producer —
1105
+ which is exactly what an owner class is. The cut is a pure function of
1106
+ the bake and the ownership table. Copy order within a chunk is bake
1107
+ order, and a module's copies may span chunks (materialize/quant fire on
1108
+ its first/last chunk; see ``_build_call_plan``)."""
1109
+ assert self._router is not None
1110
+ by_class: dict[int, list[_Scatter]] = {}
1111
+ for copies in modules:
1112
+ for c in copies:
1113
+ by_class.setdefault(self._router.class_of(c.src[0]), []).append(c)
1114
+ return [(ci, by_class[ci]) for ci in sorted(by_class)]
1115
+
1116
+ def _build_call_plan(self, names: list[str], group_lens: list[int]) -> "_CallPlan":
1117
+ """Build the STATIC plan for one whole-sync call.
1118
+
1119
+ Pure — no pulls, no engine state touched — so the result is cached and
1120
+ reused every sync. Three passes:
1121
+ 1. Split ``names`` into gather groups, one chunk per owner class
1122
+ present in this worker's baked copies, recording each group's last
1123
+ chunk for its ``free_group`` signal (or ``pre_free`` when this worker
1124
+ pulls nothing from it). Every group is signaled exactly once by
1125
+ construction. The stream has no per-group call boundaries, so group
1126
+ L+1's first chunk issues while L's still stream.
1127
+ 2. Per leaf module, find its FIRST and LAST chunk — materialize on the
1128
+ first, quant/kernel/reset on the last, correct by construction
1129
+ instead of by runtime counters.
1130
+ 3. Assemble ``_Chunk``s: dedup keys and precompute the packed layout.
1131
+ """
1132
+ router = self._router
1133
+ assert router is not None, "init_transfer_engine() must run before planning"
1134
+
1135
+ # --- pass 1: gather groups -> per-owner-class scatter chunks -----------
1136
+ raw_chunks: list[list[_Scatter]] = []
1137
+ free_at: dict[int, list[int]] = {} # chunk idx -> groups to signal after it
1138
+ pre_free: list[int] = [] # groups with no chunk here (signal at start)
1139
+ unbaked: list[str] = []
1140
+ pos = 0
1141
+ for gi, glen in enumerate(group_lens):
1142
+ gnames = names[pos : pos + glen]
1143
+ pos += glen
1144
+ modules: list[list[_Scatter]] = []
1145
+ seen: set[int] = set()
1146
+ for n in gnames:
1147
+ mod = self._name_to_plan.get(n)
1148
+ if mod is None:
1149
+ if n in self._live_names:
1150
+ unbaked.append(n)
1151
+ elif id(mod) not in seen:
1152
+ seen.add(id(mod))
1153
+ modules.append(mod)
1154
+ if not modules:
1155
+ # Nothing to pull for this group on this worker; its owners
1156
+ # still published it, so signal it done at sync start (they
1157
+ # tolerate a signal that arrives before the publish).
1158
+ pre_free.append(gi)
1159
+ continue
1160
+ for _ci, scatters in self._chunk_module_scatters(modules):
1161
+ raw_chunks.append(scatters)
1162
+ free_at.setdefault(len(raw_chunks) - 1, []).append(gi)
1163
+ if unbaked:
1164
+ # A live name with no baked plan cannot be loaded: there is no
1165
+ # fallback, and its pull would target groups the pipeline has
1166
+ # already freed. Fail here, where the names are known.
1167
+ raise RuntimeError(
1168
+ f"Sharded RDT: {len(unbaked)} live name(s) did not bake, "
1169
+ f"first: {unbaked[:3]}. Fix the loader so these names bake "
1170
+ "(fully load their module in one pass), or exclude them from "
1171
+ "the synced set."
1172
+ )
1173
+
1174
+ # --- pass 2: per-module first/last chunk -> materialize/quant ----------
1175
+ first_at: dict[int, int] = {}
1176
+ last_at: dict[int, int] = {}
1177
+ layer_by_id: dict[int, Any] = {}
1178
+ for ci, scatters in enumerate(raw_chunks):
1179
+ for sc in scatters:
1180
+ lid = id(sc.layer)
1181
+ layer_by_id[lid] = sc.layer
1182
+ first_at.setdefault(lid, ci)
1183
+ last_at[lid] = ci
1184
+ materialize_at: dict[int, list[Any]] = {}
1185
+ quant_at: dict[int, list[Any]] = {}
1186
+ for lid, ci in first_at.items():
1187
+ materialize_at.setdefault(ci, []).append(layer_by_id[lid])
1188
+ for lid, ci in last_at.items():
1189
+ quant_at.setdefault(ci, []).append(layer_by_id[lid])
1190
+
1191
+ # --- pass 3: assemble _Chunks (dedup keys + precompute pack layout) ----
1192
+ from collections import defaultdict
1193
+
1194
+ chunks: list[_Chunk] = []
1195
+ per_owner_seq: dict[int, int] = defaultdict(int)
1196
+ for ci, scatters in enumerate(raw_chunks):
1197
+ keys: list[FetchKey] = []
1198
+ kmeta: dict[FetchKey, tuple[torch.dtype, tuple[int, ...]]] = {}
1199
+ for sc in scatters:
1200
+ if sc.src not in kmeta:
1201
+ kmeta[sc.src] = (sc.dtype, sc.shape)
1202
+ keys.append(sc.src)
1203
+ pack_layout: list[tuple[int, torch.dtype, int, tuple[int, ...]]] = []
1204
+ cur = 0
1205
+ for k in keys:
1206
+ dt, shape = kmeta[k]
1207
+ numel = prod(shape) or 1 # type: ignore[arg-type]
1208
+ off = (cur + 15) & ~15
1209
+ pack_layout.append((off, dt, numel, shape))
1210
+ cur = off + numel * dt.itemsize
1211
+ owner = router.producer_for(router.consumer_id, scatters[0].src[0])
1212
+ chunks.append(
1213
+ _Chunk(
1214
+ seq=per_owner_seq[owner],
1215
+ scatters=scatters,
1216
+ keys=keys,
1217
+ pack_layout=pack_layout,
1218
+ pack_bytes=cur,
1219
+ materialize=materialize_at.get(ci, []),
1220
+ quant=quant_at.get(ci, []),
1221
+ free=free_at.get(ci, []),
1222
+ # Every name of a chunk shares an owner class and a group,
1223
+ # so any of them resolves the same producer.
1224
+ owner=owner,
1225
+ )
1226
+ )
1227
+ per_owner_seq[owner] += 1
1228
+ return _CallPlan(chunks=chunks, pre_free=pre_free)
1229
+
1230
+ def _signal_group_done(self, group_idx: int) -> None:
1231
+ """Fire-and-forget ``free_group`` signal at EVERY owner of the group.
1232
+
1233
+ The per-group barrier: each owner counts one signal per live consumer
1234
+ and frees the group (releasing its lookahead credit) on the last one, so
1235
+ every owner must hear from every consumer — including owners this worker
1236
+ pulls nothing from. Refs are held and drained in ``drain_pending`` so
1237
+ every signal has EXECUTED before the sync ends: ``begin_sync`` resets
1238
+ the counters, and a straggler landing in the next sync would credit a
1239
+ group it does not belong to."""
1240
+ assert self._router is not None
1241
+ self._pending_frees.extend(self._router.free_group(group_idx))
1242
+
1243
+ def _run_chunk_pipeline(self, plan: "_CallPlan") -> None:
1244
+ """Pipelined chunk pulls over the ring of receive slots.
1245
+
1246
+ Issues up to ``ring_depth`` produce RPCs ahead of the blocking gets, so
1247
+ while chunk i's RDMA streams the producer serves i+1 into its own ring
1248
+ slot and the background thread scatters i-1 out of another. Reads stay
1249
+ serialized on the shared NIC — the bandwidth floor, not a loss.
1250
+
1251
+ Slot safety rests on two arguments spelled out in the doc: the producer's
1252
+ ring is no shallower than this one and drain-before-issue orders its
1253
+ reuse; the consumer's slots are held by ``_issue_pull``'s generation
1254
+ handshake.
1255
+ """
1256
+ from collections import deque
1257
+
1258
+ inflight: deque[tuple[_PendingPull, _Chunk]] = deque()
1259
+
1260
+ # Gather groups with no chunk on this worker (plan.pre_free): signal
1261
+ # them done before the pipeline — their owners still published them, and
1262
+ # a signal that arrives before its publish is counted (fire-and-forget).
1263
+ for gi in plan.pre_free:
1264
+ self._signal_group_done(gi)
1265
+
1266
+ def drain_one() -> None:
1267
+ pending, chunk = inflight.popleft()
1268
+ results = self._complete_pull(pending)
1269
+ self._dispatch_item(
1270
+ _ProcItem(chunk=chunk, results=results, slot=pending.slot)
1271
+ )
1272
+ # Each gather group whose LAST chunk this is: its read is done, so
1273
+ # this consumer is finished with the group -> signal every owner
1274
+ # (fire-and-forget, off the critical path).
1275
+ for gi in chunk.free:
1276
+ self._signal_group_done(gi)
1277
+
1278
+ for chunk in plan.chunks:
1279
+ if not chunk.keys:
1280
+ # _chunk_module_scatters never emits empty chunks; keep the
1281
+ # signals safe anyway if this ever changes.
1282
+ for gi in chunk.free:
1283
+ self._signal_group_done(gi)
1284
+ continue
1285
+ # Drain BEFORE issue once the ring is full: frees this chunk's slot
1286
+ # (generation-wise) and guarantees the producer-slot invariant above.
1287
+ # Sync 0 runs SERIAL (max 1 in flight): both sides still grow and
1288
+ # register buffers, and a producer registration mid-flight churns its
1289
+ # agent-metadata version under an in-flight pull (see __init__).
1290
+ depth = 1 if self._completed_syncs == 0 else self._ring_depth
1291
+ if len(inflight) >= depth:
1292
+ drain_one()
1293
+ slot = self._pull_slot
1294
+ self._pull_slot = (slot + 1) % self._ring_depth
1295
+ inflight.append((self._issue_pull(chunk, slot), chunk))
1296
+ while inflight:
1297
+ drain_one()
1298
+
1299
+ def _process_item(self, item: "_ProcItem") -> None:
1300
+ """Scatter-thread half: materialize this chunk's first-seen modules,
1301
+ scatter its slices on the process stream, publish the slot, then hand the
1302
+ modules it COMPLETES to the quant thread.
1303
+
1304
+ Mirrors ``_layerwise_process`` minus the loader replay. Once every scatter
1305
+ reading ``item.slot`` is enqueued, records the slot's read-done event so
1306
+ the RPC thread can block on it before overwriting the slot.
1307
+ """
1308
+ from vllm.model_executor.model_loader.reload.layerwise import (
1309
+ LAYERWISE_INFO,
1310
+ )
1311
+ from vllm.model_executor.model_loader.reload.meta import materialize_layer
1312
+
1313
+ results = item.results
1314
+ chunk = item.chunk
1315
+ with (
1316
+ torch.cuda.device(self.device),
1317
+ torch.cuda.stream(self._proc_stream),
1318
+ torch.device(self.device),
1319
+ ):
1320
+ # PASS 1 — slot readers: materialize the modules whose FIRST scatter
1321
+ # is in this chunk (empty HF params, once per module by construction)
1322
+ # then scatter this chunk's copies. The scatter copies are the ONLY
1323
+ # reads of the receive buffer; quant and kernel-copy operate on the
1324
+ # scattered params. Releasing the slot right after the scatters lets
1325
+ # the NEXT chunk's RDMA overwrite the buffer while quant still runs.
1326
+ try:
1327
+ for layer in chunk.materialize:
1328
+ info = LAYERWISE_INFO.get(layer)
1329
+ if info is None or not info.can_load():
1330
+ raise RuntimeError(
1331
+ f"Baked replay: layer {type(layer).__name__} "
1332
+ "was not set up for reload this sync "
1333
+ "(start_weight_update must run before "
1334
+ "update_weights)."
1335
+ )
1336
+ materialize_layer(layer, info)
1337
+ for sc in chunk.scatters:
1338
+ param = getattr(sc.layer, sc.param_name)
1339
+ # Through .data, exactly like default_weight_loader: params
1340
+ # that keep requires_grad=True (plain nn.Parameters like
1341
+ # GLM's router bias) reject an in-place copy through an
1342
+ # autograd view of the leaf.
1343
+ dst = param.data.as_strided(sc.shape, sc.stride, sc.offset)
1344
+ with torch._C.DisableTorchFunctionSubclass():
1345
+ dst.copy_(results[sc.src])
1346
+ # All reads of this slot's buffer are now enqueued on the process
1347
+ # stream; record + publish so the RPC thread can reuse the slot.
1348
+ self._slot_read_done[item.slot].record(self._proc_stream)
1349
+ finally:
1350
+ # Publish even on error so the RPC thread's generation wait
1351
+ # unblocks (the failure itself surfaces via _raise_proc_error).
1352
+ self._mark_slot_done(item.slot)
1353
+
1354
+ # PASS 2 — param readers only: quant / kernel-copy / reset for the
1355
+ # modules whose LAST scatter is in this chunk. Handed to the
1356
+ # DEDICATED quant thread (own CUDA stream, event-chained after this
1357
+ # chunk's scatters) so it never delays this thread's next pass-1.
1358
+ # Running it in order here stalls the RPC thread's slot handshake by
1359
+ # ~0.5-0.75s/iter, since every group's quant delays the next publish.
1360
+ if chunk.quant:
1361
+ ready = torch.cuda.Event()
1362
+ ready.record(self._proc_stream)
1363
+ if self._quant_queue is None:
1364
+ self._run_quant(chunk.quant, ready)
1365
+ else:
1366
+ self._quant_queue.put((chunk.quant, ready))
1367
+
1368
+ def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None:
1369
+ """Quant/kernel-copy/reset the given COMPLETED leaf modules, exactly as
1370
+ _layerwise_process. Runs on the quant thread's own stream, ordered after
1371
+ the modules' scatters via ``ready``; touches only the scattered params
1372
+ (never a receive slot), so it can overlap subsequent chunks' RDMA and
1373
+ scatters. ``info.reset()`` is what makes finalize skip the layer —
1374
+ drain_pending joins the quant queue before finalize runs."""
1375
+ from vllm.model_executor.layers.quantization.base_config import (
1376
+ QuantizeMethodBase,
1377
+ )
1378
+ from vllm.model_executor.model_loader.reload.layerwise import (
1379
+ LAYERWISE_INFO,
1380
+ _copy_and_restore_kernel_tensors,
1381
+ )
1382
+
1383
+ stream = self._quant_stream or self._proc_stream
1384
+ assert stream is not None # created in _ensure_proc_worker before use
1385
+ with (
1386
+ torch.cuda.device(self.device),
1387
+ torch.cuda.stream(stream),
1388
+ torch.device(self.device),
1389
+ ):
1390
+ stream.wait_event(ready)
1391
+ for layer in layers:
1392
+ info = LAYERWISE_INFO.get(layer)
1393
+ assert info is not None # completed leaf module is set up for reload
1394
+ quant_method = getattr(layer, "quant_method", None)
1395
+ if isinstance(quant_method, QuantizeMethodBase):
1396
+ if hasattr(layer, "_already_called_process_weights_after_loading"):
1397
+ delattr(layer, "_already_called_process_weights_after_loading")
1398
+ quant_method.process_weights_after_loading(layer)
1399
+ # Copy into persistent kernel storage (preserves cudagraph refs).
1400
+ if info.kernel_tensors is not None:
1401
+ _copy_and_restore_kernel_tensors(layer, info)
1402
+ # Reset so finalize_layerwise_reload skips this (loaded) layer.
1403
+ info.reset()
1404
+
1405
+ def _quant_worker_loop(self) -> None:
1406
+ """Dedicated quant thread: drains (completed_modules, scatter-done event)
1407
+ batches. Errors surface via _proc_error like the scatter thread's."""
1408
+ torch.cuda.set_device(self.device)
1409
+ q = self._quant_queue
1410
+ assert q is not None
1411
+ while True:
1412
+ batch = q.get()
1413
+ try:
1414
+ if batch is None:
1415
+ return
1416
+ self._run_quant(*batch)
1417
+ except BaseException as e: # noqa: BLE001 - surfaced on the RPC thread
1418
+ self._proc_error = e
1419
+ logger.exception("RDT quant thread failed")
1420
+ finally:
1421
+ q.task_done()
1422
+
1423
+ def shutdown(self) -> None:
1424
+ # Stop the background post-processing thread (drain, then sentinel + join)
1425
+ # before dropping the state it touches.
1426
+ if self._proc_thread is not None:
1427
+ try:
1428
+ self.drain_pending()
1429
+ except Exception:
1430
+ logger.exception("RDT drain during shutdown failed")
1431
+ assert self._proc_queue is not None
1432
+ self._proc_queue.put(None) # sentinel
1433
+ self._proc_thread.join(timeout=30)
1434
+ self._proc_thread = None
1435
+ self._proc_queue = None
1436
+ self._proc_stream = None
1437
+ if self._quant_thread is not None:
1438
+ assert self._quant_queue is not None
1439
+ self._quant_queue.put(None) # sentinel
1440
+ self._quant_thread.join(timeout=30)
1441
+ self._quant_thread = None
1442
+ self._quant_queue = None
1443
+ self._quant_stream = None
1444
+ self._slot_read_done = []
1445
+ self._router = None
1446
+ # Drop strong references to baked modules so the model can be freed.
1447
+ self._name_to_plan.clear()
1448
+ self._name_meta.clear()
1449
+ self._live_names.clear()
1450
+ # Drop the cached plan (holds _Scatter refs to the baked layers).
1451
+ self._cached_plan = None
1452
+ # Release the receive buffers (their NIXL registration is pinned for the
1453
+ # process lifetime; freeing the tensors just drops our strong refs).
1454
+ self._dest_buffers = [{} for _ in range(self._ring_depth)]
1455
+
1456
+
1457
+def _plan_digest(keys_per_chunk: list) -> str:
1458
+ """Digest of the chunks one consumer pulls from one producer, in pull order.
1459
+
1460
+ Two consumers a producer serves out of ONE shared serve ring must agree on
1461
+ this, since sharing rests on their plans being identical; the producer
1462
+ compares it at init. Over the whole (name, op-chain) list rather than the
1463
+ names, because the chains decide the bytes each pull returns.
1464
+ """
1465
+ import hashlib
1466
+
1467
+ h = hashlib.sha256()
1468
+ for keys in keys_per_chunk:
1469
+ h.update(repr(keys).encode())
1470
+ h.update(b"\n")
1471
+ return h.hexdigest()[:32]
1472
+
1473
+
1474
+def _dtype_from_name(name: str) -> torch.dtype:
1475
+ """Resolve a string like 'bfloat16' to torch.bfloat16."""
1476
+ dtype = getattr(torch, name, None)
1477
+ if not isinstance(dtype, torch.dtype):
1478
+ raise ValueError(f"Unknown torch dtype name: {name!r}")
1479
+ return dtype
vllm/distributed/weight_transfer/sharded_rdt_fake.py
new
+300
@@ -0,0 +1,300 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Op-chain recording for the sharded-RDT backend.
4
+
5
+The consumer asks the trainer for the exact slice a worker consumes, described as
6
+an op chain replayed on the trainer's live tensor. ``FakeRDTTensor`` builds the
7
+chain by intercepting the model's own weight loaders; ``copy_`` is its data sink,
8
+``BakeSink`` — the dry run records how each slice would be fetched and where it
9
+lands.
10
+
11
+Chains are built against meta tensors, so nothing here touches Ray or a GPU.
12
+"""
13
+
14
+from collections import defaultdict
15
+from collections.abc import Callable
16
+from dataclasses import dataclass, field
17
+from typing import Any
18
+
19
+import torch
20
+
21
+from vllm.distributed.weight_transfer.sharded_rdt_common import SUPPORTED_OPS
22
+
23
+# ("op_name", positional_args, sorted_kwargs_items). Entries must be hashable so
24
+# the chain can serve as a FetchKey.
25
+OpSpec = tuple[str, tuple[Any, ...], tuple[tuple[str, Any], ...]]
26
+OpChain = tuple[OpSpec, ...]
27
+FetchKey = tuple[str, OpChain]
28
+
29
+
30
+def _freeze_kwargs(kwargs: dict[str, Any]) -> tuple[tuple[str, Any], ...]:
31
+ """Sort kwargs into a tuple of items for hashable storage in OpSpec."""
32
+ return tuple(sorted(kwargs.items()))
33
+
34
+
35
+class _UnsupportedFakeOp(NotImplementedError):
36
+ """Raised when a weight loader does something to a FakeRDTTensor that cannot
37
+ be expressed as a slice request.
38
+
39
+ Surfaced as NotImplementedError so callers can distinguish "this backend
40
+ can't handle this loader" from genuine bugs.
41
+ """
42
+
43
+
44
+@dataclass
45
+class _Scatter:
46
+ """One recorded scatter: pull ``src``, copy it into ``layer``'s
47
+ ``param_name`` at the recorded strided region. The bake's output unit; see
48
+ the engine module's Data flow.
49
+
50
+ Self-contained so replay needs no lookups. Destination geometry is read off
51
+ the meta view and rebuilt each sync as ``as_strided(shape, stride, offset)``
52
+ -- ``layer`` is resolved to a param at replay time, never baked, since every
53
+ sync re-materializes fresh tensors.
54
+
55
+ ``shape`` and ``dtype`` also size the slice ON THE WIRE, so ``dtype`` is the
56
+ PRODUCED dtype (the fake's after its op chain), not the source name's:
57
+ ``view(dtype)`` is allowlisted, and taking the source's would size the slice
58
+ with the wrong itemsize and carve the packed blob differently on the two
59
+ sides.
60
+ """
61
+
62
+ layer: Any
63
+ param_name: str
64
+ src: FetchKey
65
+ offset: int
66
+ shape: tuple[int, ...]
67
+ stride: tuple[int, ...]
68
+ dtype: torch.dtype
69
+
70
+
71
+def _meta_copy_(dest: torch.Tensor, src: "FakeRDTTensor") -> torch.Tensor:
72
+ """Fire ``dest.copy_`` from a zero-storage meta source of ``src``'s geometry.
73
+
74
+ Moves no data but still counts against the layer's loaded numel, which drives
75
+ ``_layerwise_process``; skipping it would leave the layer looking unloaded
76
+ forever.
77
+
78
+ A non-meta ``dest`` is one the layerwise reload never moved there (the
79
+ ``SKIP_LOAD_TENSORS`` set, e.g. GLM's router bias). Torch forbids
80
+ ``real.copy_(meta)``, and skipping the count is consistent -- ``get_layer_size``
81
+ excludes the same set. The caller still RECORDS the copy, so the replay writes
82
+ the real param.
83
+ """
84
+ if not dest.is_meta:
85
+ return dest
86
+ meta_src = torch.empty(src.shape, dtype=src.dtype, device="meta")
87
+ with torch._C.DisableTorchFunctionSubclass():
88
+ return dest.copy_(meta_src)
89
+
90
+
91
+@dataclass
92
+class BakeSink:
93
+ """``copy_`` sink for the dry-run bake: record how each slice would be fetched
94
+ and where it would land, and move nothing.
95
+
96
+ ``_install_recording_stamps`` sets ``current = (leaf_module, param_name)``
97
+ around each loader so ``accept_copy`` can attribute the copy. An unstamped
98
+ ``copy_`` cannot be attributed and stays unrecorded, so its module fails the
99
+ coverage gate and takes the plain load. ``copies_by_layer`` is keyed by module
100
+ object, so iterating it yields each leaf module once.
101
+ """
102
+
103
+ copies_by_layer: "dict[Any, list[_Scatter | None]]" = field(
104
+ default_factory=lambda: defaultdict(list)
105
+ )
106
+ current: "tuple[Any, str] | None" = None
107
+ # Source names whose ``copy_`` fired during the bake. Names not here never
108
+ # moved data (e.g. experts owned by another EP rank) and are skipped; names
109
+ # here but unbaked fail the plan build.
110
+ copied_names: "set[str]" = field(default_factory=set)
111
+
112
+ def accept_copy(self, dest: torch.Tensor, src: "FakeRDTTensor") -> torch.Tensor:
113
+ self.copied_names.add(src._name)
114
+ if self.current is not None:
115
+ layer, param_name = self.current
116
+ if dest.numel() != src.numel():
117
+ # The packed layout sizes each slice from the DESTINATION shape,
118
+ # while the producer packs what the replayed chain produces. A
119
+ # broadcasting copy_ makes those two lengths differ, and neither
120
+ # side can notice: the consumer's buffer view is exactly
121
+ # prod(dest.shape) elements, so it reshapes cleanly over bytes
122
+ # the producer laid out at different offsets. Refuse at bake.
123
+ raise _UnsupportedFakeOp(
124
+ f"FakeRDTTensor: broadcasting copy_ into {param_name!r} of "
125
+ f"{type(layer).__name__} ({src.numel()} elements from "
126
+ f"{src._name!r} into {dest.numel()}). The sharded RDT "
127
+ "backend sends exactly the slice the loader asks for, so "
128
+ "the source and destination must have equal numel."
129
+ )
130
+ self.copies_by_layer[layer].append(
131
+ _Scatter(
132
+ layer=layer,
133
+ param_name=param_name,
134
+ src=src._key(),
135
+ offset=dest.storage_offset(),
136
+ shape=tuple(dest.shape),
137
+ stride=tuple(dest.stride()),
138
+ dtype=src.dtype,
139
+ )
140
+ )
141
+ return _meta_copy_(dest, src)
142
+
143
+
144
+class FakeRDTTensor(torch.Tensor):
145
+ """Zero-storage tensor that records how to fetch a weight slice.
146
+
147
+ ``_make_wrapper_subclass`` gives it shape/dtype/device without storage. Every
148
+ op in ``SUPPORTED_OPS`` returns a child with the spec appended; ``copy_``
149
+ delegates to the installed sink. Anything else reaches ``__torch_dispatch__``
150
+ and raises ``_UnsupportedFakeOp``, so failures are loud rather than silently
151
+ fetching the wrong bytes.
152
+ """
153
+
154
+ # Declared at class scope so mypy can infer attribute types on
155
+ # instances built via ``_make_wrapper_subclass`` (where ``__new__``
156
+ # returns a tensor it can't annotate as ``self``).
157
+ _name: str
158
+ _ops: OpChain
159
+ # Handles this fake's copy_. ``None`` only on bare construction.
160
+ _sink: "BakeSink | None"
161
+
162
+ @staticmethod
163
+ def __new__(
164
+ cls,
165
+ name: str,
166
+ shape: torch.Size,
167
+ dtype: torch.dtype,
168
+ device: torch.device,
169
+ ops: OpChain = (),
170
+ sink: "BakeSink | None" = None,
171
+ ) -> "FakeRDTTensor":
172
+ t = torch.Tensor._make_wrapper_subclass(
173
+ cls,
174
+ shape,
175
+ dtype=dtype,
176
+ device=device,
177
+ requires_grad=False,
178
+ )
179
+ t._name = name
180
+ t._ops = tuple(ops)
181
+ t._sink = sink
182
+ return t
183
+
184
+ def _key(self) -> FetchKey:
185
+ return (self._name, self._ops)
186
+
187
+ def _make_child(
188
+ self,
189
+ new_shape: torch.Size,
190
+ new_dtype: torch.dtype,
191
+ *new_ops: OpSpec,
192
+ ) -> "FakeRDTTensor":
193
+ """Append one or more ops to the chain and return a fresh child.
194
+
195
+ Variadic so multi-return ops (e.g. chunk) can append both the base
196
+ op and an indexing op in a single call.
197
+ """
198
+ return FakeRDTTensor(
199
+ name=self._name,
200
+ shape=new_shape,
201
+ dtype=new_dtype,
202
+ device=self.device,
203
+ ops=self._ops + new_ops,
204
+ sink=self._sink,
205
+ )
206
+
207
+ def _meta(self) -> torch.Tensor:
208
+ """A zero-storage meta tensor of this fake's shape/dtype, so PyTorch
209
+ itself computes post-op geometry rather than us reimplementing it."""
210
+ return torch.empty(self.shape, dtype=self.dtype, device="meta")
211
+
212
+ @classmethod
213
+ def __torch_function__(
214
+ cls,
215
+ func,
216
+ types,
217
+ args=(),
218
+ kwargs=None,
219
+ ):
220
+ kwargs = kwargs or {}
221
+
222
+ # Checked before the allowlist: copy_ must not be recorded into the chain.
223
+ if func is torch.Tensor.copy_:
224
+ dest = args[0]
225
+ src = args[1] if len(args) > 1 else kwargs.get("src")
226
+ if isinstance(src, cls) and src._sink is not None:
227
+ return src._sink.accept_copy(dest, src)
228
+
229
+ # Allowlisted slice/view/shape ops: append to chain and return child.
230
+ op_name = SUPPORTED_OPS.get(func)
231
+ if op_name is not None:
232
+ self_ = args[0]
233
+ if isinstance(self_, cls):
234
+ rest = tuple(args[1:])
235
+ return cls._intercept(self_, func, op_name, rest, kwargs)
236
+
237
+ # Metadata reads are answered by the wrapper subclass and stop here; ops
238
+ # that need data fall through to __torch_dispatch__, which raises.
239
+ with torch._C.DisableTorchFunctionSubclass():
240
+ return func(*args, **kwargs)
241
+
242
+ @classmethod
243
+ def _intercept(
244
+ cls,
245
+ self_: "FakeRDTTensor",
246
+ func: Callable,
247
+ op_name: str,
248
+ args: tuple,
249
+ kwargs: dict,
250
+ ):
251
+ """Append the op and return a child, or a tuple of children for
252
+ multi-return ops. Each child's geometry comes from running the op on a
253
+ meta tensor."""
254
+ meta = self_._meta()
255
+ with torch._C.DisableTorchFunctionSubclass():
256
+ meta_result = func(meta, *args, **kwargs)
257
+
258
+ base_op: OpSpec = (op_name, tuple(args), _freeze_kwargs(kwargs))
259
+
260
+ # Single-tensor result: one child carrying base_op alone.
261
+ if isinstance(meta_result, torch.Tensor):
262
+ return self_._make_child(meta_result.shape, meta_result.dtype, base_op)
263
+
264
+ # Multi-return: one child per output, each carrying base_op plus a
265
+ # trailing __getitem__(i) so the replay can index the result.
266
+ if isinstance(meta_result, tuple | list):
267
+ return tuple(
268
+ self_._make_child(
269
+ m.shape,
270
+ m.dtype,
271
+ base_op,
272
+ ("__getitem__", (i,), ()),
273
+ )
274
+ for i, m in enumerate(meta_result)
275
+ )
276
+
277
+ # Non-tensor result (e.g. .item() snuck into the allowlist).
278
+ raise _UnsupportedFakeOp(
279
+ f"FakeRDTTensor: {op_name!r} returned a non-tensor "
280
+ f"({type(meta_result).__name__}); cannot defer."
281
+ )
282
+
283
+ @classmethod
284
+ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
285
+ kwargs = kwargs or {}
286
+
287
+ # A fake at the aten level means an unsupported loader op. Name the op and
288
+ # the chain; materializing instead would mask a correctness bug.
289
+ for arg in (*args, *kwargs.values()):
290
+ if isinstance(arg, cls):
291
+ raise _UnsupportedFakeOp(
292
+ f"FakeRDTTensor: unsupported op {func} reached "
293
+ f"__torch_dispatch__ on fake {arg._name!r} "
294
+ f"(chain={arg._ops}). Supported ops are: "
295
+ f"{sorted(SUPPORTED_OPS.values())}, plus copy_. "
296
+ "Loaders that need .to(), .float(), .item(), arithmetic, "
297
+ "bool-mask indexing, or .data access are not supported by "
298
+ "the sharded RDT backend."
299
+ )
300
+ return func(*args, **kwargs)
vllm/distributed/weight_transfer/sharded_rdt_trainer.py
new
+1521
@@ -0,0 +1,1521 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Trainer-side engine for the sharded-RDT (pull-based NIXL) backend.
4
+
5
+RDT is pull-based, so unlike NCCL this engine broadcasts nothing. It owns a
6
+per-rank producer server (an internal Ray actor exposing the NIXL serve surface
7
+the worker engine dials by name), and on each ``send_weights`` gathers this
8
+rank's weights group-by-group from the ``WeightSource``, shares each group into
9
+the server over CUDA IPC, and — on the sender — drives the inference-side
10
+start/update/finish handshake, whose single empty ``update_weights`` unblocks the
11
+workers to pull.
12
+
13
+All serve-side state lives on the server actor, so trainer processes need no
14
+mixin, no named actors and no special actor options.
15
+
16
+See docs/training/weight_transfer/sharded_rdt.md for the publish -> serve ->
17
+free_group -> release lifecycle and the ownership model.
18
+"""
19
+
20
+import contextlib
21
+import threading
22
+import time
23
+import uuid
24
+from collections.abc import Callable, Collection
25
+from concurrent.futures import ThreadPoolExecutor
26
+from dataclasses import asdict, dataclass, field
27
+from typing import Any, ClassVar
28
+
29
+import ray
30
+import torch
31
+from torch.multiprocessing.reductions import (
32
+ StorageWeakRef,
33
+ rebuild_cuda_tensor,
34
+ reduce_tensor,
35
+)
36
+from typing_extensions import Self
37
+
38
+from vllm.distributed.weight_transfer.base import (
39
+ ParamMeta,
40
+ TrainerInitInfo,
41
+ TrainerWeightTransferEngine,
42
+ VLLMWeightSyncClient,
43
+ WeightSource,
44
+ layerwise_groups,
45
+)
46
+from vllm.distributed.weight_transfer.sharded_rdt_common import (
47
+ ALLOWED_OPS,
48
+ buffer_alloc_bytes,
49
+ check_ray_rdt_version,
50
+)
51
+from vllm.logger import init_logger
52
+
53
+logger = init_logger(__name__)
54
+
55
+# Gathered-but-unfreed groups the gather loop may run ahead by before it stops
56
+# gathering. A gathered group is published immediately (publish_group never
57
+# blocks), so at most ``lookahead + 1`` groups are resident on the trainer — the
58
+# memory bound larger models size against, pinned by the tests.
59
+#
60
+# At 1, group N+1 is gathered and serveable while the consumers pull N, so the
61
+# boundary's free-barrier latency hides behind live pulls. Raise it only if one
62
+# group's gather is slower than its pulls.
63
+DEFAULT_GATHER_LOOKAHEAD = 1
64
+
65
+# Seconds with no publish, produce or free before the
66
+# producer declares the sync dead. A consumer that dies mid-sync never signals
67
+# ``free_group``, so the group is never released and the waits below block
68
+# forever — which stops this rank iterating its WeightSource, a collective, and
69
+# wedges every other trainer rank with no exception anywhere. This converts that
70
+# into one real error.
71
+#
72
+# A liveness backstop, not a latency target: nominal gaps are sub-second, so this
73
+# is ~100x margin. It rides the init info because the producer is a Ray actor and
74
+# does not inherit the trainer's environment.
75
+DEFAULT_STALL_TIMEOUT_S = 300.0
76
+# How often a blocked waiter re-checks the progress stamp.
77
+_STALL_POLL_S = 5.0
78
+
79
+# Tensors whose STORAGE is at most this are packed into a ring slot and exported
80
+# as one CUDA-IPC handle; larger storages are exported directly. Batching the
81
+# small ones is what keeps sync time steady -- a per-tensor handle for each of
82
+# ~37k names dominated the variance -- and `_pack_group_for_export` is also the
83
+# only publish path that orders the pack against the sidecar's IPC read, so the
84
+# packing is not optional.
85
+_EXPORT_RING_MAX_BYTES = 64 << 20
86
+
87
+# The actor method the worker engine dials for the NIXL pull; fixed by contract.
88
+PRODUCE_METHOD_NAME = "rdt_produce_weights_batched"
89
+
90
+
91
+# Cross-deployment serve-slot sharing.
92
+#
93
+# Consumers whose ids differ by a multiple of ``workers_per_replica`` are the
94
+# same worker of different inference deployments, so they share a baked plan and
95
+# pull byte-identical chunks. A serve ring each would cost one full ring per
96
+# deployment of the producer's GPU and repack the same bytes once per deployment.
97
+#
98
+# NIXL reads are one-sided, so R readers can read ONE registered slot; what the
99
+# producer cannot see is when a reader has FINISHED. The release edge is
100
+# therefore the consumer's ISSUE order, sent as ``seq``: the pipeline drains pull
101
+# i before issuing i+K, so slot ``seq % K`` was last packed for ``seq - K``,
102
+# whose read is over. A counter on THIS side would follow execution order
103
+# instead, which is wrong -- Ray may start a consumer's K concurrent produce
104
+# calls in any order, so a call that executes K-before another can be a pull
105
+# still being read, and its slot gets repacked underneath the reader. Silent,
106
+# showing up only as a logprob drift.
107
+#
108
+# Hence one rendezvous per generation, keyed by ``seq``: the group's live sharers
109
+# meet there, the LAST to arrive packs, and all return that one blob. It carries
110
+# the release edge too, since ``seq`` is packed only once every live sharer has
111
+# arrived, so each has drained ``seq - K``. Mismatched plans would rendezvous on
112
+# differing bytes, so they are rejected at init by ``reserve_serve_buffer``'s
113
+# ``plan_digest``. A sharer that dies mid-sync stalls its group, exactly as a
114
+# dead consumer already stalls the per-group free barrier; the next
115
+# ``begin_sync`` narrows the live set.
116
+#
117
+# LIMITATION: the rendezvous synchronizes deployments that a pull-based transfer
118
+# would otherwise leave independent -- per chunk every sharer waits for the
119
+# slowest, so skew between their DP pause/resume brackets lands in every
120
+# deployment's sync time. Nothing about P2P requires this. It is the price of
121
+# serving out of GPU memory, which is too scarce for a ring per consumer, so
122
+# slots must be reused, and reuse needs a release edge.
123
+#
124
+# TODO: explore staging packed chunks in HOST memory. Host buffers are cheap
125
+# enough for a ring per consumer, which would drop the rendezvous, the
126
+# cross-replica coupling and the slot-reuse hazard together. The cost is a D2H
127
+# bounce on the serve path, which has been the bottleneck before, so it needs
128
+# measuring rather than assuming.
129
+
130
+
131
+@dataclass
132
+class _SharedPack:
133
+ """One generation of one slot-sharing group: the rendezvous state for a
134
+ single chunk, identified by the consumers' issue index."""
135
+
136
+ group: int
137
+ slot: int
138
+ """``seq % ring_depth``. Fixed at creation, so slot reuse follows the
139
+ consumers' issue order rather than this side's execution order."""
140
+ arrived: set[int] = field(default_factory=set)
141
+ blob: Any = None
142
+ packing: bool = False
143
+ done: bool = False
144
+ error: BaseException | None = None
145
+
146
+
147
+@dataclass
148
+class ShardedRDTTrainerInitInfo(TrainerInitInfo):
149
+ """Trainer init info for the sharded-RDT backend.
150
+
151
+ Identical on every rank except ``rank`` (rank 0 is the sender). Carries only
152
+ the must-agree wire params; the sender forwards them verbatim onto the
153
+ worker-side init info so the two cannot drift. Server-actor names are
154
+ generated per rank and all-gathered by the engine, not supplied here.
155
+ """
156
+
157
+ backend: ClassVar[str] = "sharded_rdt"
158
+
159
+ num_consumers: int
160
+ """Total inference-worker (consumer) count across the whole fleet
161
+ (DP*TP*PP*PCP), for the M:N block assignment / free ref-count."""
162
+ workers_per_replica: int = 0
163
+ """Consumers per inference DEPLOYMENT (``num_consumers // num_replicas``).
164
+
165
+ Fixes the slot-sharing groups: consumers whose ids differ by a multiple of
166
+ this are the same worker of different deployments, so they bake identical
167
+ plans and pull byte-identical chunks, and ONE registered serve slot can
168
+ serve all of them. 0 disables sharing, and so does the single-deployment
169
+ value ``num_consumers``, which makes every group a singleton and the serve
170
+ path identical to the unshared one.
171
+ """
172
+ trainer_actor_namespace: str | None = None
173
+ """Ray namespace the engine spawns its serve actors in. The inference
174
+ workers (which run in their own EngineCore subprocess with its own
175
+ ``ray.init``) resolve those actors by name, so this must be the namespace
176
+ they can see. Forwarded to the worker-side init info."""
177
+ num_rdt_buffers: int = 2
178
+ """Serve/receive ring depth K (must match the worker)."""
179
+ buffer_presize_gb: float = 0.0
180
+ """Serve-buffer pre-size floor in GiB (avoids NIXL desc-cache churn)."""
181
+ gather_lookahead: int = DEFAULT_GATHER_LOOKAHEAD
182
+ """Gathered-but-unfreed groups the gather loop may run ahead by. Bounds
183
+ trainer-resident memory at ``gather_lookahead + 1`` groups."""
184
+ stall_timeout_s: float = DEFAULT_STALL_TIMEOUT_S
185
+ """Seconds of no publish/serve/free progress before the producer fails the
186
+ sync (see ``DEFAULT_STALL_TIMEOUT_S``)."""
187
+
188
+
189
+class _RDTProducerServer:
190
+ """Per-rank NIXL serve surface: an internal Ray actor sharing the trainer
191
+ rank's GPU over CUDA IPC. Holds the gather cache, per-consumer serve rings,
192
+ the per-group free barrier and the packed serve. The engine feeds it with
193
+ ``publish_group``; workers pull with ``rdt_produce_weights_batched`` and
194
+ signal with ``free_group``.
195
+
196
+ A plain class — the engine wraps it with ``ray.remote(...)`` at spawn so the
197
+ actor options live in one place.
198
+ """
199
+
200
+ def __init__(
201
+ self,
202
+ *,
203
+ num_rdt_buffers: int,
204
+ buffer_presize_gb: float,
205
+ gather_lookahead: int,
206
+ served_names: list[str] | None = None,
207
+ stall_timeout_s: float = DEFAULT_STALL_TIMEOUT_S,
208
+ workers_per_replica: int = 0,
209
+ ) -> None:
210
+ import gc
211
+
212
+ self._device_index = torch.accelerator.current_device_index()
213
+ # name -> rebuilt CUDA-IPC tensor (or view); guarded by _cache_cond.
214
+ self._cache: dict[str, torch.Tensor] = {}
215
+ self._cache_cond = threading.Condition()
216
+ self._gather_error: BaseException | None = None
217
+
218
+ # Monotonic stamp of the last forward step of the
219
+ # publish -> serve -> free loop. The waits below poll it instead of
220
+ # blocking forever, so a dead consumer fails the sync with a real error.
221
+ self._stall_timeout = float(stall_timeout_s)
222
+ self._last_progress = time.monotonic()
223
+
224
+ # Names this producer publishes, so a misrouted pull fails loudly
225
+ # instead of blocking forever in the cache wait. None = serve anything.
226
+ self._served_names = set(served_names) if served_names is not None else None
227
+
228
+ # Every live consumer signals free_group(gi) at every
229
+ # owner, once per sync; the group frees when the count reaches
230
+ # begin_sync's live total — one uniform target, no routed per-producer
231
+ # ones. Signals may precede the publish, which completes them.
232
+ self._live_count = 1
233
+ self._free_counts: dict[int, int] = {}
234
+ # gi -> the names this producer published for it (what release drops).
235
+ self._group_names: dict[int, list[str]] = {}
236
+
237
+ # Published-but-unfreed groups, plus freed ones not
238
+ # yet handed back. The memory gate lives in the ENGINE's gather loop; this
239
+ # side only accounts — free_group moves a group from _inflight_groups to
240
+ # _freed_pending, which the engine collects via wait_freed / end_sync to
241
+ # drop its storage refs. _lookahead is engine-enforced, unused here.
242
+ self._lookahead = max(0, gather_lookahead)
243
+ self._inflight_groups: list[int] = []
244
+ self._freed_pending: list[int] = []
245
+
246
+ # Ring of packed serve buffers, one ring per SLOT-SHARING GROUP (one
247
+ # consumer per group unless several deployments share); the slot within
248
+ # it is the consumer's ``seq % nring``.
249
+ self._nring = max(1, num_rdt_buffers)
250
+ self._serve_rings: dict[int, list[torch.Tensor | None]] = {}
251
+ self._serve_lock = threading.Lock()
252
+ # registerMem on a shared NIXL agent is not concurrency-safe; serialize.
253
+ self._reg_lock = threading.Lock()
254
+ self._buffer_presize = int(buffer_presize_gb * (1 << 30))
255
+ self._serve_device = torch.device("cuda", self._device_index)
256
+
257
+ # Slot-sharing state, all guarded by _cache_cond, so the rendezvous
258
+ # waits ride the stall watchdog like every other wait here.
259
+ # ``_share_width`` = consumers per deployment; 0 means every group is a
260
+ # singleton and this degenerates to the unshared serve path.
261
+ self._share_width = max(0, int(workers_per_replica))
262
+ self._live_ids: set[int] | None = None
263
+ self._sharing_active = False
264
+ self._sharers: dict[int, frozenset] = {} # group -> live members
265
+ self._gens: dict[int, dict] = {} # group -> {seq: _SharedPack}
266
+ # group -> {cid: plan digest}. Init-time state, not cleared per sync: it
267
+ # is what proves the sharers of a group pull the same chunks.
268
+ self._plan_digests: dict[int, dict[int, str]] = {}
269
+
270
+ # (sharing group, ring idx, packed layout) ->
271
+ # (buffer data_ptr, destination views). Keyed on the buffer pointer too, so
272
+ # a ring regrow invalidates rather than writing into a freed buffer. See
273
+ # the serve path for why the layout, not the spec names, is the key.
274
+ self._pack_dsts: dict[tuple, tuple[int, list[torch.Tensor]]] = {}
275
+ # [RDT-EXPORT-RING] handle -> rebuilt base. The trainer reuses buffers
276
+ # across groups, so the same handle arrives on many publishes; without
277
+ # this we reopen it every group. Cleared per sync so no IPC mapping into
278
+ # trainer memory outlives the sync that made it.
279
+ self._base_cache: dict[tuple, torch.Tensor] = {}
280
+
281
+ # Freeze the static post-init object graph so gen-2 GC never stops the
282
+ # world mid-serve, which shows up as serve stragglers.
283
+ gc.collect()
284
+ gc.freeze()
285
+
286
+ def ping(self) -> int:
287
+ return self._device_index
288
+
289
+ # ---------------- stall watchdog ----------------
290
+
291
+ def _note_progress_locked(self) -> None:
292
+ """Record that the credit loop moved. Caller holds ``_cache_cond``."""
293
+ self._last_progress = time.monotonic()
294
+
295
+ def _wait_for(self, blocked: Callable[[], bool], what: str) -> None:
296
+ """``_cache_cond.wait()`` with a liveness bound.
297
+
298
+ Waits while ``blocked()`` holds, returning early on a gather error. If
299
+ nothing on this producer progresses for ``_stall_timeout``, self-fires
300
+ ``set_gather_error``, which every waiter here already checks — so the rank
301
+ unwinds through one path and the driver gets a real exception.
302
+
303
+ The progress stamp is global to the producer, not per-waiter: a merely
304
+ slow waiter is kept alive by its peers' progress. Caller holds
305
+ ``_cache_cond``.
306
+ """
307
+ while blocked():
308
+ if self._gather_error is not None:
309
+ return
310
+ self._cache_cond.wait(_STALL_POLL_S)
311
+ if not blocked() or self._gather_error is not None:
312
+ return
313
+ stalled = time.monotonic() - self._last_progress
314
+ if stalled >= self._stall_timeout:
315
+ msg = (
316
+ f"RDT stall: no progress for {stalled:.0f}s while waiting for "
317
+ f"{what} (timeout {self._stall_timeout:.0f}s). A consumer most "
318
+ f"likely died mid-sync: {len(self._inflight_groups)} group(s) "
319
+ f"published and unfreed."
320
+ )
321
+ logger.error("[rdt-stall] %s", msg)
322
+ # Same channel a gather failure uses: wakes every waiter here, and
323
+ # each raises rather than returning a half-served result.
324
+ self._gather_error = RuntimeError(msg)
325
+ self._cache_cond.notify_all()
326
+ return
327
+
328
+ def warmup_nixl(self) -> None:
329
+ """Create this server's NIXL agent now, while the rank's GPU is quiet.
330
+
331
+ Called at spawn, before the server-name all-gather, so no rank can be
332
+ spinning in a collective. Creating the agent lazily instead deadlocks on
333
+ EFA-class fabrics (see "warmup_nixl breaks a startup deadlock" in the
334
+ doc). The warmup buffer stays registered so the agent's CUDA-HMEM path
335
+ stays initialized.
336
+ """
337
+ from ray.experimental import register_nixl_memory
338
+
339
+ self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda")
340
+ with self._reg_lock:
341
+ register_nixl_memory(self._nixl_warmup_buf)
342
+
343
+ # ---------------- engine-facing (per sync) ----------------
344
+
345
+ def begin_sync(
346
+ self, live_count: int, live_consumer_ids: list | None = None
347
+ ) -> None:
348
+ """Reset per-sync free/backpressure state and set this sync's barrier
349
+ target.
350
+
351
+ ``live_count`` is how many consumers take part in THIS sync. Required, no
352
+ default: a forgotten argument silently targeting 1 would free groups after
353
+ the FIRST signal while others still pull — use-after-free, not an error.
354
+ The driver awaits the previous sync's finish before the next begins, so
355
+ nothing is in flight; a straggler signal would otherwise credit the wrong
356
+ sync, which is why the consumer drains its signals before finishing.
357
+
358
+ ``live_consumer_ids`` is that same live set enumerated, which is what
359
+ sizes each slot-sharing group's rendezvous; a count cannot, since a group
360
+ is a specific set of ids. ``None`` leaves every group a singleton.
361
+
362
+ The packed-destination cache deliberately survives: the layout repeats
363
+ every sync. So does ``_plan_digests``, which is init-time state.
364
+ """
365
+ with self._cache_cond:
366
+ self._gather_error = None
367
+ self._live_count = max(1, int(live_count))
368
+ self._live_ids = (
369
+ set(int(c) for c in live_consumer_ids)
370
+ if live_consumer_ids is not None
371
+ else None
372
+ )
373
+ # Sharing is on only when the geometry shows several deployments;
374
+ # one deployment takes the singleton path.
375
+ self._sharing_active = (
376
+ self._share_width > 0
377
+ and self._live_ids is not None
378
+ and len(self._live_ids) > self._share_width
379
+ )
380
+ self._sharers.clear()
381
+ self._gens.clear()
382
+ self._inflight_groups.clear()
383
+ self._freed_pending.clear()
384
+ self._free_counts.clear()
385
+ self._group_names.clear()
386
+ self._base_cache.clear()
387
+ self._note_progress_locked()
388
+
389
+ def _release_group_locked(self, group_idx: int) -> None:
390
+ """Drop a freed group's cache entries and queue it for the engine
391
+ (whose gather-credit gate blocks in ``wait_freed`` on exactly this).
392
+
393
+ Shared by the last ``free_group`` and by ``publish_group`` completing an
394
+ early-signaled group. Caller must hold ``_cache_cond``.
395
+ """
396
+ for name in self._group_names.pop(group_idx, ()):
397
+ self._cache.pop(name, None)
398
+ self._free_counts.pop(group_idx, None)
399
+ if group_idx in self._inflight_groups:
400
+ self._inflight_groups.remove(group_idx)
401
+ self._freed_pending.append(group_idx)
402
+
403
+ def publish_group(self, group_idx: int, entries: tuple) -> None:
404
+ """Rebuild one gather group's CUDA-IPC tensors into the serve cache.
405
+
406
+ NEVER blocks: a gathered group is serveable immediately. The memory bound
407
+ lives in the engine's gather loop, which stops GATHERING (not publishing)
408
+ past ``gather_lookahead`` unfreed groups.
409
+
410
+ ``entries`` is ``(storages, views)``: one CUDA-IPC export per storage,
411
+ plus per-name ``(sid, dtype_name, shape, stride, storage_offset)`` rebuilt
412
+ here as ``as_strided`` views.
413
+
414
+ Signals can arrive BEFORE their publish — a consumer pulling nothing of a
415
+ group signals it as its plan starts — so a group whose barrier is already
416
+ satisfied is released here. Freed groups reach the engine only through
417
+ ``wait_freed`` / ``end_sync``: a freed notice riding an unharvested async
418
+ publish result would wedge the loop.
419
+ """
420
+ storages, views = entries
421
+ bases: dict[int, torch.Tensor] = {}
422
+ for sid, reduce_args in storages.items():
423
+ cached = self._base_cache.get(reduce_args)
424
+ if cached is not None:
425
+ bases[sid] = cached
426
+ continue
427
+ list_args = list(reduce_args)
428
+ # Index 6 of reduce_tensor's args is the exporter's device index;
429
+ # rebuild on this server's device (same physical GPU as the rank).
430
+ list_args[6] = self._device_index
431
+ bases[sid] = self._base_cache[reduce_args] = rebuild_cuda_tensor(*list_args)
432
+ rebuilt: dict[str, torch.Tensor] = {}
433
+ for name, (sid, dtype_name, shape, stride, storage_offset) in views.items():
434
+ typed = bases[sid].view(getattr(torch, dtype_name))
435
+ rebuilt[name] = torch.as_strided(typed, shape, stride, storage_offset)
436
+ del bases
437
+
438
+ with self._cache_cond:
439
+ self._cache.update(rebuilt)
440
+ self._inflight_groups.append(group_idx)
441
+ self._group_names[group_idx] = list(rebuilt)
442
+ if self._free_counts.get(group_idx, 0) >= self._live_count:
443
+ self._release_group_locked(group_idx)
444
+ self._note_progress_locked()
445
+ self._cache_cond.notify_all()
446
+
447
+ def wait_freed(self) -> list[int]:
448
+ """Block until at least one published group has been freed; return and
449
+ clear the freed backlog. The engine's gather-credit gate calls this once
450
+ its loop is ``gather_lookahead`` groups ahead, so this wait is where the
451
+ trainer paces to the consumers' pull rate.
452
+
453
+ Raises rather than returning empty when the sync errors or stalls: the
454
+ engine is blocked here inside its gather loop, and an empty return would
455
+ spin it straight back in.
456
+ """
457
+ with self._cache_cond:
458
+ self._wait_for(lambda: not self._freed_pending, "a freed-group credit")
459
+ if not self._freed_pending:
460
+ raise RuntimeError(
461
+ "gather errored while waiting for a freed-group credit: "
462
+ f"{self._gather_error!r}"
463
+ )
464
+ freed = self._freed_pending
465
+ self._freed_pending = []
466
+ return freed
467
+
468
+ def end_sync(self) -> list[int]:
469
+ """Block until every published group has been freed by its consumers;
470
+ return the remaining freed keys so the engine drops its last refs."""
471
+ with self._cache_cond:
472
+ self._wait_for(
473
+ lambda: bool(self._inflight_groups),
474
+ "every published group to be freed",
475
+ )
476
+ freed = self._freed_pending
477
+ self._freed_pending = []
478
+ self._base_cache.clear()
479
+ return freed
480
+
481
+ def set_gather_error(self, message: str) -> None:
482
+ """Record a trainer-side gather failure so blocked serves / publishes
483
+ stop waiting and surface it."""
484
+ with self._cache_cond:
485
+ self._gather_error = RuntimeError(message)
486
+ self._cache_cond.notify_all()
487
+
488
+ # ---------------- consumer-facing (called by name over Ray) ----------------
489
+
490
+ def free_group(self, group_idx: int) -> None:
491
+ """Consumer back-edge: one consumer is done with group ``group_idx``,
492
+ either because its last chunk landed or because it had nothing to pull.
493
+
494
+ The per-group barrier counts one signal per live consumer against the
495
+ ``begin_sync`` count; every consumer signals every owner, so the target is
496
+ the same integer everywhere. The last signal drops the cache entries and
497
+ queues the group as gather credit for ``wait_freed``. A signal arriving
498
+ before its publish is only counted — ``publish_group`` completes it.
499
+ """
500
+ gi = int(group_idx)
501
+ with self._cache_cond:
502
+ count = self._free_counts.get(gi, 0) + 1
503
+ self._free_counts[gi] = count
504
+ if gi in self._group_names and count >= self._live_count:
505
+ self._release_group_locked(gi)
506
+ self._note_progress_locked()
507
+ self._cache_cond.notify_all()
508
+
509
+ def _new_serve_buffer(self, nbytes: int) -> torch.Tensor:
510
+ """Allocate + NIXL-register one serve slot: the single allocation seam,
511
+ so registration cannot be skipped on either of the two paths that make
512
+ buffers (the init-time reservation and the serve-path backstop)."""
513
+ from ray.experimental import register_nixl_memory
514
+
515
+ t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device)
516
+ with self._reg_lock:
517
+ register_nixl_memory(t)
518
+ return t
519
+
520
+ def reserve_serve_buffer(
521
+ self, consumer_id: int, nbytes: int, plan_digest: str | None = None
522
+ ) -> None:
523
+ """Pre-allocate + NIXL-register this consumer's serve ring before any
524
+ pull, while the fabric is idle (avoids registration races during the
525
+ sync-0 RDMA churn under M:N fan-in). Idempotent; grows only if needed.
526
+
527
+ The ring is keyed by SLOT-SHARING GROUP, so the sharers of one group
528
+ reserve ONE ring between them. They all call this, same group and same
529
+ size, so the whole body runs under ``_serve_lock``, or two concurrent
530
+ first-callers would each allocate a ring and one would be dropped while
531
+ still registered.
532
+
533
+ ``plan_digest`` is the caller's ordered chunk plan for THIS producer.
534
+ Sharers must agree on it, and this is where a disagreement is caught
535
+ rather than at serve time, where it is a rendezvous nobody completes.
536
+
537
+ Raises:
538
+ RuntimeError: two consumers of one sharing group disagree on the
539
+ chunks they pull from this producer.
540
+ """
541
+ sg = self._share_group(consumer_id)
542
+ if plan_digest is not None:
543
+ with self._cache_cond:
544
+ seen = self._plan_digests.setdefault(sg, {})
545
+ clash = next(
546
+ ((c, d) for c, d in seen.items() if d != plan_digest), None
547
+ )
548
+ if clash is None:
549
+ seen[consumer_id] = plan_digest
550
+ if clash is not None:
551
+ raise RuntimeError(
552
+ f"consumers {clash[0]} and {consumer_id} are in slot-sharing "
553
+ f"group {sg} but do not pull the same chunks from this "
554
+ f"producer ({clash[1]} != {plan_digest}). Sharing a serve "
555
+ f"slot requires the deployments to be identical; check that "
556
+ f"workers_per_replica matches the inference fleet's "
557
+ f"geometry."
558
+ )
559
+ # This is where the producer's NIXL serve memory is committed: one ring
560
+ # of K slots per SHARING GROUP that pulls from this rank -- with a single
561
+ # deployment, one ring per consumer, since every group is a singleton.
562
+ # Routing only sends a consumer here for names this rank owns, so the
563
+ # count is the fan-in, not the fleet.
564
+ #
565
+ # Every slot is sized by the LARGEST chunk that consumer pulls from us,
566
+ # rounded up to a 256 MB multiple, so the whole ring pays for the worst
567
+ # chunk in the plan while most chunks are far smaller.
568
+ #
569
+ # TODO: that over-allocates. Sizing slots per chunk instead of per plan,
570
+ # or capping chunk bytes so no chunk is an outlier, would shrink it; so
571
+ # would dropping the 256 MB round-up, which is slack once a buffer never
572
+ # regrows. Row-splitting the vocab matrix shrinks the outlier itself.
573
+ # The host-staging TODO above would remove these rings altogether.
574
+ alloc = buffer_alloc_bytes(nbytes, self._buffer_presize)
575
+ with self._serve_lock:
576
+ rings = self._serve_rings.setdefault(sg, [None] * self._nring)
577
+ for i in range(self._nring):
578
+ slot = rings[i]
579
+ if slot is None or slot.numel() < alloc:
580
+ rings[i] = self._new_serve_buffer(alloc)
581
+
582
+ def _share_group(self, consumer_id: int) -> int:
583
+ """The slot-sharing group ``consumer_id`` belongs to: its index within
584
+ its own deployment, since that is what fixes the plan. Without a width
585
+ the group is the consumer itself and nothing is shared."""
586
+ return consumer_id % self._share_width if self._share_width > 0 else consumer_id
587
+
588
+ def _sharers_of(self, sg: int) -> frozenset:
589
+ """The LIVE consumers of group ``sg``: the rendezvous width. Derived from
590
+ ``begin_sync``'s live set, so a degraded sync narrows the barrier instead
591
+ of waiting forever on a dead deployment. Caller holds ``_cache_cond``."""
592
+ cached = self._sharers.get(sg)
593
+ if cached is None:
594
+ if not self._sharing_active or self._live_ids is None:
595
+ cached = frozenset((sg,))
596
+ else:
597
+ cached = frozenset(
598
+ c for c in self._live_ids if c % self._share_width == sg
599
+ ) or frozenset((sg,))
600
+ self._sharers[sg] = cached
601
+ return cached
602
+
603
+ def _join_share_group(self, consumer_id: int, seq: int) -> tuple:
604
+ """Register this call's arrival in its group's rendezvous for ``seq`` and
605
+ return ``(generation, is_packer)``.
606
+
607
+ Exactly one caller per generation gets ``is_packer=True``, whichever
608
+ completes the arrival set, and it is the only one that touches the GPU.
609
+ The slot is ``seq % nring``: fixed by the consumers' issue order, so it
610
+ needs no release accounting on this side.
611
+ """
612
+ sg = self._share_group(consumer_id)
613
+ with self._cache_cond:
614
+ gens = self._gens.setdefault(sg, {})
615
+ gen = gens.get(seq)
616
+ if gen is None:
617
+ gen = gens[seq] = _SharedPack(group=sg, slot=seq % self._nring)
618
+ gen.arrived.add(consumer_id)
619
+ if gen.packing or gen.done or not gen.arrived >= self._sharers_of(sg):
620
+ return gen, False
621
+ gen.packing = True
622
+ return gen, True
623
+
624
+ def _await_shared_pack(self, gen: _SharedPack) -> torch.Tensor:
625
+ """Block until this generation's packer published its blob, and return
626
+ it. A failed pack is re-raised here, so every sharer of a bad pack fails
627
+ instead of reading a half-written slot."""
628
+ with self._cache_cond:
629
+ self._wait_for(
630
+ lambda: not gen.done and gen.error is None,
631
+ "a co-replica's shared serve pack",
632
+ )
633
+ if gen.error is not None:
634
+ raise RuntimeError(
635
+ f"the sharer packing this chunk failed: {gen.error!r}"
636
+ ) from gen.error
637
+ if not gen.done:
638
+ raise RuntimeError(
639
+ f"gather errored while awaiting a shared pack: "
640
+ f"{self._gather_error!r}"
641
+ )
642
+ return gen.blob
643
+
644
+ def _serve_slot(self, sg: int, idx: int, need: int) -> torch.Tensor:
645
+ """Group ``sg``'s ring slot ``idx``, grown if this chunk outgrew the
646
+ reservation. Growing registers memory while the fabric is busy, the
647
+ hazard ``reserve_serve_buffer`` exists to avoid, so it is a backstop: the
648
+ reservation is sized from the same static plan as this pack."""
649
+ with self._serve_lock:
650
+ buffer = self._serve_rings.setdefault(sg, [None] * self._nring)[idx]
651
+ if buffer is not None and buffer.numel() >= need:
652
+ return buffer
653
+ buffer = self._new_serve_buffer(buffer_alloc_bytes(need, self._buffer_presize))
654
+ with self._serve_lock:
655
+ self._serve_rings.setdefault(sg, [None] * self._nring)[idx] = buffer
656
+ return buffer
657
+
658
+ def _fail_shared_pack(self, gen: _SharedPack, exc: BaseException) -> None:
659
+ """Publish a pack failure so waiting sharers raise instead of hanging on
660
+ a generation that will never complete."""
661
+ with self._cache_cond:
662
+ gen.error = exc
663
+ self._cache_cond.notify_all()
664
+
665
+ @ray.method(tensor_transport="nixl")
666
+ def rdt_produce_weights_batched(
667
+ self, specs: list, consumer_id: int = 0, seq: int = -1
668
+ ):
669
+ """Serve one batched slice request over NIXL.
670
+
671
+ Waits until the specs' names are cached, then rendezvouses with the
672
+ other live sharers of this consumer's slot-sharing group. The last to
673
+ arrive replays each spec's op chain (pure views into cached tensors,
674
+ guarded by ALLOWED_OPS) and byte-packs the slices 16B-aligned into the
675
+ group's ring slot ``seq % nring``, mirroring the consumer's identical
676
+ layout; every sharer then returns that one packed blob, so R deployments
677
+ cost one pack and one slot instead of R. With a single deployment every
678
+ group is a singleton, so the arriving call is always its own packer and
679
+ this is the unshared serve path exactly.
680
+
681
+ ``seq`` is the caller's index in its pull stream to this producer, which
682
+ is what fixes the slot. Callers that want a single slice pass one spec
683
+ and read the blob back with that slice's dtype/shape.
684
+
685
+ Raises:
686
+ ValueError: ``seq`` was not supplied.
687
+ """
688
+ needed = sorted({n for n, _ in specs})
689
+ if self._served_names is not None:
690
+ unserved = [n for n in needed if n not in self._served_names]
691
+ if unserved:
692
+ # Without this the cache wait below would block forever: this
693
+ # producer never gathers these names, so they can never arrive.
694
+ raise RuntimeError(
695
+ f"pull routed to the wrong producer: {unserved[:3]} "
696
+ f"({len(unserved)} names) are not served here"
697
+ )
698
+ with self._cache_cond:
699
+ self._wait_for(
700
+ lambda: not all(n in self._cache for n in needed),
701
+ f"{len(needed)} name(s) to be published",
702
+ )
703
+ if not all(n in self._cache for n in needed):
704
+ # _wait_for only gives up on a gather error (its own stall
705
+ # included), so the names are never coming.
706
+ raise RuntimeError(
707
+ f"gather errored before {needed}: {self._gather_error!r}"
708
+ )
709
+
710
+ if seq < 0:
711
+ # The slot is derived from it, so a caller that does not send one
712
+ # cannot be served safely.
713
+ raise ValueError(
714
+ "rdt_produce_weights_batched requires the consumer's issue index `seq`"
715
+ )
716
+ gen, is_packer = self._join_share_group(consumer_id, seq)
717
+ if not is_packer:
718
+ blob = self._await_shared_pack(gen)
719
+ else:
720
+ try:
721
+ blob = self._pack_shared(gen, specs)
722
+ except BaseException as e:
723
+ self._fail_shared_pack(gen, e)
724
+ raise
725
+ # A served pull is the third of the four progress signals the stall
726
+ # watchdog reads (publish / produce / free / begin_sync): a long sync whose
727
+ # consumers pull steadily but slowly must never trip it.
728
+ with self._cache_cond:
729
+ self._note_progress_locked()
730
+ return [blob]
731
+
732
+ def _pack_shared(self, gen: _SharedPack, specs: list) -> torch.Tensor:
733
+ """Replay the op chains into this generation's slot and publish the blob
734
+ to the sharers waiting on it.
735
+
736
+ Runs on exactly one call per generation and OUTSIDE ``_cache_cond``: the
737
+ pack is GPU work and must not block publishes, frees or other groups.
738
+ """
739
+ sliced: list = [] # (byte_off, tensor)
740
+ pack_cur = 0
741
+ for name, chain in specs:
742
+ t = self._cache[name]
743
+ for op, args, kw in chain:
744
+ if op not in ALLOWED_OPS:
745
+ raise ValueError(f"{name!r}: disallowed op {op!r}")
746
+ t = getattr(t, op)(*args, **dict(kw))
747
+ off = (pack_cur + 15) & ~15
748
+ pack_cur = off + t.numel() * t.element_size()
749
+ sliced.append((off, t))
750
+
751
+ buffer = self._serve_slot(gen.group, gen.slot, pack_cur)
752
+
753
+ # The destination views are a pure function of the packed
754
+ # layout, which is byte-identical every sync, so build them once per
755
+ # (sharing group, ring slot, layout) and reuse — rebuilding per call cost
756
+ # 5.2ms of a 7.5ms 384-spec 235B group.
757
+ #
758
+ # The key is the LAYOUT, not the spec names: a name can appear in two
759
+ # requests with different op chains (one name's copies can split across
760
+ # owner-class chunks), and serving the second through the first's views
761
+ # would write the wrong bytes with nothing downstream to catch it.
762
+ dst_key = (
763
+ gen.group,
764
+ gen.slot,
765
+ tuple((off, t.dtype, t.shape) for off, t in sliced),
766
+ )
767
+ cached = self._pack_dsts.get(dst_key)
768
+ if cached is None or cached[0] != buffer.data_ptr():
769
+ dsts = []
770
+ for off, t in sliced:
771
+ nb = t.numel() * t.element_size()
772
+ dsts.append(buffer[off : off + nb].view(t.dtype).reshape(t.shape))
773
+ self._pack_dsts[dst_key] = (buffer.data_ptr(), dsts)
774
+ else:
775
+ dsts = cached[1]
776
+ torch._foreach_copy_(dsts, [t for _off, t in sliced])
777
+
778
+ blob = buffer[:pack_cur]
779
+ with self._cache_cond:
780
+ gen.blob = blob
781
+ gen.done = True
782
+ self._note_progress_locked()
783
+ self._cache_cond.notify_all()
784
+ return blob
785
+
786
+ def shutdown(self) -> None:
787
+ with self._cache_cond:
788
+ self._cache.clear()
789
+ with self._cache_cond:
790
+ # Generations hold blob views into the rings, so they go first.
791
+ self._gens.clear()
792
+ self._sharers.clear()
793
+ with self._serve_lock:
794
+ self._serve_rings.clear()
795
+ # Must go with the rings: these are views INTO them, and the
796
+ # data_ptr guard that normally invalidates them cannot tell a freed
797
+ # buffer from a new one recycled at the same address.
798
+ self._pack_dsts.clear()
799
+
800
+
801
+class ShardedRDTTrainerWeightTransferEngine(
802
+ TrainerWeightTransferEngine[ShardedRDTTrainerInitInfo]
803
+):
804
+ """Trainer-side engine for the pull-based sharded-RDT backend.
805
+
806
+ Lives on every trainer rank. Owns a per-rank `_RDTProducerServer` actor
807
+ (the NIXL serve surface). `send_weights` gathers this rank's weights
808
+ group-by-group from the `WeightSource`, shares each group into the server
809
+ over CUDA IPC, and — on the sender — drives the inference-side handshake so
810
+ the workers pull. Non-sender ranks only gather (staying in the collective).
811
+ """
812
+
813
+ init_info_cls = ShardedRDTTrainerInitInfo
814
+
815
+ def __init__(
816
+ self,
817
+ *,
818
+ client: VLLMWeightSyncClient,
819
+ source: WeightSource,
820
+ is_sender: bool = True,
821
+ init_info: ShardedRDTTrainerInitInfo,
822
+ ) -> None:
823
+ super().__init__(client=client, source=source, is_sender=is_sender)
824
+ self._init_info = init_info
825
+ self._server: Any = None # the _RDTProducerServer actor handle
826
+ self._server_name: str | None = None
827
+ # Group-major metadata / partition, computed at trainer_init.
828
+ self._meta: list[ParamMeta] = []
829
+ self._groups: list[list[str]] = []
830
+ # Ownership resolved at trainer_init from the fleet's held names: the
831
+ # distinct owner sets, the per-name index into them, this rank's groups
832
+ # and its held set. Routing itself is consumer-side; the trainer only
833
+ # ships the table.
834
+ self._owner_sets: list[list[int]] = []
835
+ self._name_owner_class: list[int] = []
836
+ self._owned_idx: list[int] = []
837
+ self._held_names: set[str] | None = None
838
+ # Strong refs to gathered tensors we've shared into the server, keyed by
839
+ # group index. CUDA-IPC exports must outlive the importer, so we hold
840
+ # them until the server reports the group freed. See send_weights.
841
+ self._inflight: dict[int, dict[str, torch.Tensor]] = {}
842
+ # [RDT-EXPORT-RING] Reusable packed slots, one per resident group. See
843
+ # _pack_group_for_export. Persistent across syncs: freeing and
844
+ # reallocating them each sync fragmented the caching allocator badly
845
+ # enough to grow reserved memory ~1.5 GB per sync.
846
+ self._export_ring: list[torch.Tensor | None] = []
847
+ self._export_ring_args: list[tuple | None] = []
848
+ # Which live group holds which slot, and the slots free to hand out. A
849
+ # modular counter is unsafe here: the credit gate bounds the NUMBER of
850
+ # unfreed groups, not their order, so it can hand a later group the slot
851
+ # a still-live group is being served from.
852
+ self._slot_of_group: dict[int, int] = {}
853
+ self._free_slots: list[int] = []
854
+ # Set by trainer_init from the server-name all-gather; retained so a
855
+ # restarted consumer can be re-initialized without another collective.
856
+ self._server_names: list[str] | None = None
857
+
858
+ def _rpc(self, method: str, *args: Any) -> Any:
859
+ """Call one of the server actor's methods and block for the result.
860
+ The single seam through which the engine talks to its server, so tests
861
+ can inject a local (non-Ray) fake server."""
862
+ import ray
863
+
864
+ return ray.get(getattr(self._server, method).remote(*args))
865
+
866
+ @staticmethod
867
+ def _contig_stride(shape) -> tuple:
868
+ stride: list[int] = []
869
+ acc = 1
870
+ for s in reversed(list(shape)):
871
+ stride.append(acc)
872
+ acc *= int(s)
873
+ return tuple(reversed(stride))
874
+
875
+ def _pack_group_for_export(self, held: list, slot_idx: int) -> tuple:
876
+ """Copy ``held`` into ring slot ``slot_idx``; return ``(storages, views,
877
+ refs)`` in the same shape the per-storage path returns, but with ONE
878
+ storage. Views are built exactly as ``publish_group`` rebuilds them, so
879
+ the two sides cannot disagree about layout.
880
+
881
+ Slot safety is not the credit gate alone: that bounds how MANY groups
882
+ are unfreed, not the order they free in, and barriers do complete out of
883
+ order. The caller takes ``slot_idx`` from a pool keyed by live group
884
+ (``_slot_of_group`` / ``_free_slots``); ``_drop_inflight`` returns a slot
885
+ only once its group is freed everywhere.
886
+ """
887
+ offsets: list[int] = []
888
+ cur = 0
889
+ for _name, t in held:
890
+ cur = (cur + 15) & ~15 # 16B keeps every element size we ship aligned
891
+ offsets.append(cur)
892
+ cur += t.numel() * t.element_size()
893
+ need = max(16, (cur + 15) & ~15)
894
+
895
+ slot = self._export_ring[slot_idx]
896
+ device = held[0][1].device
897
+ if slot is None or slot.numel() < need or slot.device != device:
898
+ self._export_ring[slot_idx] = self._export_ring_args[slot_idx] = None
899
+ slot = torch.empty(need, dtype=torch.uint8, device=device)
900
+ self._export_ring[slot_idx] = slot
901
+
902
+ ust = slot.untyped_storage()
903
+ sid = ust.data_ptr()
904
+ reduce_args = self._export_ring_args[slot_idx]
905
+ if reduce_args is None:
906
+ base = torch.empty(0, dtype=torch.uint8, device=device)
907
+ base.set_(ust, 0, (ust.nbytes(),))
908
+ _rebuild, reduce_args = reduce_tensor(base)
909
+ self._export_ring_args[slot_idx] = reduce_args
910
+
911
+ views: dict[str, tuple] = {}
912
+ refs: dict[str, torch.Tensor] = {}
913
+ for (name, t), off in zip(held, offsets):
914
+ shape, esz = tuple(t.shape), t.element_size()
915
+ stride = self._contig_stride(shape)
916
+ dst = torch.as_strided(slot.view(t.dtype), shape, stride, off // esz)
917
+ dst.copy_(t)
918
+ refs[name] = dst
919
+ views[name] = (
920
+ sid,
921
+ str(t.dtype).split(".")[-1],
922
+ list(shape),
923
+ list(stride),
924
+ off // esz,
925
+ )
926
+ # The producer sidecar reads this slot from another process over CUDA
927
+ # IPC, while these copies are merely enqueued on this rank's stream;
928
+ # nothing orders the two. Without this sync it packs bytes the copy has
929
+ # not written yet -- silent corruption, and only for ring-packed tensors.
930
+ if slot.is_cuda:
931
+ torch.cuda.current_stream(slot.device).synchronize()
932
+ return {sid: reduce_args}, views, refs
933
+
934
+ def _publish_async(self, group_idx: int, entries):
935
+ """Fire publish_group WITHOUT waiting on the RPC (the gather loop
936
+ overlaps the publish's server-side rebuild with the next group's
937
+ gather) and return a handle that ``_await_publish`` resolves. Ray actor
938
+ handle in production; a plain (non-Ray) fake server runs inline."""
939
+ method = self._server.publish_group
940
+ remote = getattr(method, "remote", None)
941
+ if remote is not None:
942
+ return remote(group_idx, entries)
943
+ return method(group_idx, entries)
944
+
945
+ def _await_publish(self, ref) -> None:
946
+ """Resolve one async publish so a server-side rebuild error surfaces at
947
+ window depth instead of at end_sync. Publishes carry nothing back —
948
+ freed groups flow only through wait_freed/end_sync (one channel; see
949
+ publish_group)."""
950
+ import ray
951
+
952
+ if isinstance(ref, ray.ObjectRef):
953
+ ray.get(ref)
954
+
955
+ # ---------------- construction ----------------
956
+
957
+ @classmethod
958
+ def trainer_init(
959
+ cls,
960
+ init_info: ShardedRDTTrainerInitInfo,
961
+ *,
962
+ client: VLLMWeightSyncClient,
963
+ source: WeightSource | None = None,
964
+ ) -> Self:
965
+ if source is None:
966
+ raise ValueError(
967
+ "Sharded RDT trainer weight transfer requires a WeightSource."
968
+ )
969
+ engine = cls(
970
+ client=client,
971
+ source=source,
972
+ is_sender=init_info.is_sender,
973
+ init_info=init_info,
974
+ )
975
+
976
+ # The VMM-based expandable_segments allocator makes IPC storage opens ~9x
977
+ # slower on both export and rebuild, and CUDA-IPC publish is this engine's
978
+ # hot path. Frameworks enabling it should disable it around send_weights.
979
+ # Only the env var is visible; runtime changes are not introspectable.
980
+ import os as _os
981
+
982
+ if "expandable_segments:True" in _os.environ.get("PYTORCH_CUDA_ALLOC_CONF", ""):
983
+ logger.warning(
984
+ "Sharded-RDT trainer: PYTORCH_CUDA_ALLOC_CONF enables "
985
+ "expandable_segments; CUDA-IPC weight publishing will be "
986
+ "several times slower. Disable expandable segments around "
987
+ "weight sync (allocate gather buffers in classic segments)."
988
+ )
989
+
990
+ engine._meta = list(source.metadata())
991
+ names = [m.name for m in engine._meta]
992
+ engine._groups = layerwise_groups(names)
993
+ flat = [n for g in engine._groups for n in g]
994
+ if flat != names:
995
+ raise ValueError(
996
+ "Sharded RDT requires a WeightSource whose metadata order is "
997
+ "group-contiguous. Reorder the source so all names sharing a "
998
+ "layer index are adjacent."
999
+ )
1000
+
1001
+ world, rank = engine._world_and_rank()
1002
+ engine._resolve_ownership(world, rank)
1003
+ engine._spawn_server(sorted(engine._held_names or []))
1004
+
1005
+ # Every rank's server must exist before the sender's init RPC (the worker
1006
+ # init calls reserve_serve_buffer back on ALL producer servers). The
1007
+ # all-gather of server names doubles as that barrier.
1008
+ server_names = engine._all_gather_server_names(world, rank)
1009
+ # Retained so a RESTARTED consumer can be re-initialized without another
1010
+ # all-gather (see get_worker_init_payload). The uuid actor names stay valid
1011
+ # for the run, because producers never restart.
1012
+ engine._server_names = server_names
1013
+
1014
+ if engine.is_sender:
1015
+ worker_init = engine._build_worker_init_info(server_names)
1016
+ engine.client.init_weight_transfer_engine(asdict(worker_init))
1017
+ return engine
1018
+
1019
+ def get_worker_init_payload(self) -> dict:
1020
+ """The consumer-side init payload, rebuilt on demand. Pure — no collective.
1021
+
1022
+ Reads only retained state, so a restarted inference engine can rejoin at a
1023
+ sync boundary: the driver can ask any time, including mid-run with every
1024
+ rank in its own training step. A collective here would deadlock, since the
1025
+ ranks are not at a matching one.
1026
+
1027
+ Raises:
1028
+ RuntimeError: called before ``trainer_init`` cached the server names.
1029
+ """
1030
+ if self._server_names is None:
1031
+ raise RuntimeError(
1032
+ "get_worker_init_payload requires trainer_init to have run (the "
1033
+ "producer server names are gathered there)."
1034
+ )
1035
+ return asdict(self._build_worker_init_info(self._server_names))
1036
+
1037
+ def _world_and_rank(self) -> tuple[int, int]:
1038
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
1039
+ return torch.distributed.get_world_size(), torch.distributed.get_rank()
1040
+ return 1, self._init_info.rank
1041
+
1042
+ def _resolve_ownership(self, world: int, rank: int) -> None:
1043
+ """Resolve which rank holds which name, and this rank's publish plan.
1044
+
1045
+ A source may hold only part of the model — pipeline stages, expert
1046
+ parallelism, or any mix — so each rank declares its held names and the
1047
+ fleet all-gathers them. The consumers route per name, so the wire carries
1048
+ the transposed result: the distinct owner sets, and a per-name index into
1049
+ them.
1050
+
1051
+ The masks are positional over metadata order, which is why the metadata
1052
+ digest is checked first: a rank whose names disagree would transpose into
1053
+ the wrong owners entirely.
1054
+ """
1055
+ assert self.source is not None # guaranteed by trainer_init
1056
+ num_groups = len(self._groups)
1057
+ names = [m.name for m in self._meta]
1058
+ held = self.source.held_names()
1059
+ if held is None:
1060
+ held_set = set(names)
1061
+ else:
1062
+ held_set = {str(n) for n in held}
1063
+ unknown = sorted(held_set - set(names))
1064
+ if unknown:
1065
+ raise ValueError(
1066
+ f"WeightSource.held_names() lists {len(unknown)} name(s) not "
1067
+ f"in metadata(), e.g. {unknown[:3]}."
1068
+ )
1069
+ if not held_set:
1070
+ raise ValueError(
1071
+ "WeightSource.held_names() is empty; a rank with nothing to "
1072
+ "serve cannot take part in the gather"
1073
+ )
1074
+
1075
+ # One collective carries this rank's holdings as a bitmask over metadata
1076
+ # order, plus the digest of the metadata it indexes into. The digest is
1077
+ # what makes partial ownership safe: only the sender's metadata reaches
1078
+ # the consumers, so a rank describing just its own share would silently
1079
+ # serve the wrong weights.
1080
+ digest = self._meta_digest()
1081
+ mask = bytearray((len(names) + 7) // 8)
1082
+ for i, n in enumerate(names):
1083
+ if n in held_set:
1084
+ mask[i >> 3] |= 1 << (i & 7)
1085
+ per_rank = self._all_gather_owned(world, (digest, bytes(mask)))
1086
+ mismatched = [r for r, (d, *_rest) in enumerate(per_rank) if d != digest]
1087
+ if mismatched:
1088
+ raise ValueError(
1089
+ f"WeightSource.metadata() disagrees across trainer ranks "
1090
+ f"(rank {rank} digest {digest}, differing ranks {mismatched[:4]}). "
1091
+ "Every rank must describe the WHOLE model, even when it holds "
1092
+ "only some of it."
1093
+ )
1094
+ masks = [m for _d, m in per_rank]
1095
+
1096
+ # Transpose to per-name owners, then dedup into classes. Numbering by
1097
+ # FIRST APPEARANCE in metadata order keeps it a pure function of
1098
+ # rank-identical inputs, which is what lets a rejoining consumer rebuild
1099
+ # the identical table from get_worker_init_payload.
1100
+ owner_sets: list[list[int]] = []
1101
+ class_of_owners: dict[tuple[int, ...], int] = {}
1102
+ name_owner_class: list[int] = []
1103
+ for i, n in enumerate(names):
1104
+ owners = tuple(r for r, m in enumerate(masks) if m[i >> 3] & (1 << (i & 7)))
1105
+ if not owners:
1106
+ raise ValueError(
1107
+ f"no trainer rank holds {n!r}; every name in metadata() must "
1108
+ "be held by at least one rank or it can never be served."
1109
+ )
1110
+ ci = class_of_owners.get(owners)
1111
+ if ci is None:
1112
+ ci = len(owner_sets)
1113
+ class_of_owners[owners] = ci
1114
+ owner_sets.append(list(owners))
1115
+ name_owner_class.append(ci)
1116
+
1117
+ self._owner_sets = owner_sets
1118
+ self._name_owner_class = name_owner_class
1119
+ # Groups holding anything here: exactly what iteration must cover.
1120
+ self._owned_idx = [
1121
+ gi
1122
+ for gi in range(num_groups)
1123
+ if any(n in held_set for n in self._groups[gi])
1124
+ ]
1125
+ self._held_names = held_set
1126
+
1127
+ def _validate_held_yields(self, gi: int, names, tensors) -> None:
1128
+ """Check ``held_names()`` against what the source actually yields — the
1129
+ one place both sit side by side.
1130
+
1131
+ Without it, a source that claims a name but yields ``None`` for it dies
1132
+ by stall watchdog 300s later: consumers route pulls here, the pull
1133
+ passes the served-names guard, and the cache wait never completes. With
1134
+ it, that is an immediate error naming the weight. One set lookup per
1135
+ name per sync."""
1136
+ if self._held_names is None:
1137
+ return
1138
+ for name, tensor in zip(names, tensors):
1139
+ if (tensor is None) == (name in self._held_names):
1140
+ claim = "does not hold" if tensor is not None else "holds"
1141
+ have = "a real tensor" if tensor is not None else "None"
1142
+ raise RuntimeError(
1143
+ f"held_names() disagrees with the yielded tensors: group {gi} "
1144
+ f"name {name!r} yielded {have} but this rank {claim} it. A "
1145
+ "WeightSource must yield a real tensor for every held name "
1146
+ "and None for the rest (see WeightSource.held_names)."
1147
+ )
1148
+
1149
+ def _meta_digest(self) -> str:
1150
+ """Stable digest of this rank's metadata (name order + count)."""
1151
+ import hashlib
1152
+
1153
+ h = hashlib.sha256()
1154
+ h.update(f"{len(self._meta)}\n".encode())
1155
+ for m in self._meta:
1156
+ h.update(m.name.encode())
1157
+ h.update(b"\n")
1158
+ return h.hexdigest()[:16]
1159
+
1160
+ def _all_gather_owned(self, world: int, mine: tuple) -> list[tuple]:
1161
+ """All-gather each rank's (metadata digest, held-name bitmask)."""
1162
+ if world <= 1 or not (
1163
+ torch.distributed.is_available() and torch.distributed.is_initialized()
1164
+ ):
1165
+ return [mine]
1166
+ gathered: list[Any] = [None] * world
1167
+ torch.distributed.all_gather_object(gathered, mine)
1168
+ return gathered
1169
+
1170
+ def _all_gather_server_names(self, world: int, rank: int) -> list[str]:
1171
+ assert self._server_name is not None
1172
+ if world <= 1 or not (
1173
+ torch.distributed.is_available() and torch.distributed.is_initialized()
1174
+ ):
1175
+ return [self._server_name]
1176
+ gathered: list[str | None] = [None] * world
1177
+ torch.distributed.all_gather_object(gathered, self._server_name)
1178
+ return [n for n in gathered if n is not None]
1179
+
1180
+ def _spawn_server(self, served_names: list[str]) -> None:
1181
+ import ray
1182
+ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
1183
+
1184
+ # The trainer is usually a separate install from the workers, so it gets
1185
+ # its own check rather than trusting the consumer side's.
1186
+ check_ray_rdt_version()
1187
+
1188
+ ii = self._init_info
1189
+ self._server_name = f"vllm_rdt_producer_{uuid.uuid4().hex[:12]}_rk{ii.rank}"
1190
+ node_id = ray.get_runtime_context().get_node_id()
1191
+ # Pin the server to this rank's physical GPU: num_gpus=0 so Ray does not
1192
+ # allocate a second, CUDA_VISIBLE_DEVICES so CUDA IPC to the rank's
1193
+ # gathered tensors works. max_concurrency > 1 serves pulls while control
1194
+ # calls are in flight; enable_tensor_transport gives the NIXL serve.
1195
+ gpu_ids = ray.get_gpu_ids()
1196
+ # The server is the trainer rank's process twin: forward the env the
1197
+ # rank runs under (library paths etc.) so it imports torch/vllm the
1198
+ # same way, then pin it to the rank's physical GPU for CUDA IPC.
1199
+ import os
1200
+
1201
+ env_vars = {
1202
+ k: os.environ[k]
1203
+ for k in (
1204
+ "LD_LIBRARY_PATH",
1205
+ "LD_PRELOAD",
1206
+ "NCCL_CUMEM_ENABLE",
1207
+ "VLLM_NCCL_SO_PATH",
1208
+ "PATH",
1209
+ )
1210
+ if k in os.environ
1211
+ }
1212
+ if gpu_ids:
1213
+ # A num_gpus=0 actor sharing the rank's GPU: Ray would set
1214
+ # CUDA_VISIBLE_DEVICES="" and hide every GPU, so tell it not to touch
1215
+ # the var and pin the device ourselves.
1216
+ env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] = "1"
1217
+ env_vars["CUDA_VISIBLE_DEVICES"] = str(gpu_ids[0])
1218
+ runtime_env = {"env_vars": env_vars} if env_vars else {}
1219
+ server_cls = ray.remote(_RDTProducerServer).options(
1220
+ name=self._server_name,
1221
+ namespace=ii.trainer_actor_namespace,
1222
+ num_cpus=0,
1223
+ num_gpus=0,
1224
+ # Thread budget: under EP-local routing every consumer pulls from every
1225
+ # producer, and with issue-ahead each can park K produce calls in the
1226
+ # cache-wait — C*K blocked calls. A smaller pool queues the control
1227
+ # plane behind them and deadlocks: the parked produces wait for a
1228
+ # publish that can never get a thread (the sync-2 wedge at 235B tp8,
1229
+ # 16 parked vs 8 threads). Size for the worst case plus control-plane
1230
+ # slack.
1231
+ max_concurrency=ii.num_consumers * ii.num_rdt_buffers + 4,
1232
+ enable_tensor_transport=True,
1233
+ scheduling_strategy=NodeAffinitySchedulingStrategy(
1234
+ node_id=node_id, soft=False
1235
+ ),
1236
+ runtime_env=runtime_env,
1237
+ )
1238
+ self._server = server_cls.remote(
1239
+ served_names=served_names,
1240
+ num_rdt_buffers=ii.num_rdt_buffers,
1241
+ buffer_presize_gb=ii.buffer_presize_gb,
1242
+ gather_lookahead=ii.gather_lookahead,
1243
+ stall_timeout_s=ii.stall_timeout_s,
1244
+ workers_per_replica=ii.workers_per_replica,
1245
+ )
1246
+ ray.get(self._server.ping.remote())
1247
+ # Pre-barrier NIXL warmup: must complete before the server-name
1248
+ # all-gather so no rank can be inside a collective yet (see warmup_nixl).
1249
+ ray.get(self._server.warmup_nixl.remote())
1250
+
1251
+ def _build_worker_init_info(self, server_names: list[str]):
1252
+ from vllm.distributed.weight_transfer.sharded_rdt_engine import (
1253
+ ShardedRDTWeightTransferInitInfo,
1254
+ )
1255
+
1256
+ group_lens = [len(g) for g in self._groups]
1257
+ names = [m.name for m in self._meta]
1258
+ dtype_names = [str(m.dtype).split(".")[-1] for m in self._meta]
1259
+ shapes = [list(m.shape) for m in self._meta]
1260
+ return ShardedRDTWeightTransferInitInfo(
1261
+ trainer_actor_names=server_names,
1262
+ trainer_actor_namespace=self._init_info.trainer_actor_namespace,
1263
+ produce_method_name=PRODUCE_METHOD_NAME,
1264
+ names=names,
1265
+ dtype_names=dtype_names,
1266
+ shapes=shapes,
1267
+ group_lens=group_lens,
1268
+ owner_sets=self._owner_sets,
1269
+ name_owner_class=self._name_owner_class,
1270
+ num_consumers=self._init_info.num_consumers,
1271
+ num_rdt_buffers=self._init_info.num_rdt_buffers,
1272
+ buffer_presize_gb=self._init_info.buffer_presize_gb,
1273
+ )
1274
+
1275
+ # ---------------- per-round ----------------
1276
+
1277
+ def send_weights(self, live_consumer_ids: Collection[int] | None = None) -> None:
1278
+ """Gather this rank's weights and publish them for the consumers to pull.
1279
+
1280
+ ``live_consumer_ids`` restricts the sync to the consumers still alive;
1281
+ ``None`` serves the whole provisioned set. The provisioned geometry is
1282
+ FROZEN for the run — a degraded sync only lowers the live count handed to
1283
+ ``begin_sync``. Every rank must get the SAME live set, since they share
1284
+ the gather collectives, so the caller computes it once for all of them.
1285
+ """
1286
+ assert self.source is not None
1287
+ if live_consumer_ids is None:
1288
+ live_ids = list(range(self._init_info.num_consumers))
1289
+ live_count = len(live_ids)
1290
+ else:
1291
+ live_ids = sorted(set(int(c) for c in live_consumer_ids))
1292
+ live_count = len(live_ids)
1293
+ logger.warning(
1294
+ "[rdt-degraded] serving %d/%d live consumers; every group's "
1295
+ "free barrier counts to the live total",
1296
+ live_count,
1297
+ self._init_info.num_consumers,
1298
+ )
1299
+ self._send_weights_inner(live_count, live_ids)
1300
+
1301
+ def _send_weights_inner(self, live_count: int, live_ids: list[int]) -> None:
1302
+ if not self.is_sender:
1303
+ self._run_gather_loop(
1304
+ update_future=None, live_count=live_count, live_ids=live_ids
1305
+ )
1306
+ return
1307
+
1308
+ self.client.start_weight_update()
1309
+
1310
+ from vllm.distributed.weight_transfer.sharded_rdt_engine import (
1311
+ ShardedRDTWeightTransferUpdateInfo,
1312
+ )
1313
+
1314
+ empty_update = asdict(ShardedRDTWeightTransferUpdateInfo())
1315
+ with ThreadPoolExecutor(max_workers=1) as exe:
1316
+ # The workers block inside update_weights until they've pulled every
1317
+ # group, so it runs concurrently with the gather/publish loop.
1318
+ future = exe.submit(self.client.update_weights, empty_update)
1319
+ self._run_gather_loop(
1320
+ update_future=future, live_count=live_count, live_ids=live_ids
1321
+ )
1322
+ future.result() # surface inference-side errors
1323
+
1324
+ self.client.finish_weight_update()
1325
+
1326
+ def _run_gather_loop(
1327
+ self, update_future, live_count: int, live_ids: list[int] | None = None
1328
+ ) -> None:
1329
+ """Gather this rank's weights group-by-group and publish each into the
1330
+ server over CUDA IPC. A gathered group is published — serveable —
1331
+ immediately; the loop gates BEFORE the next gather while more than
1332
+ `gather_lookahead` groups are unfreed (the per-group free barrier: a
1333
+ credit releases when every live consumer has signaled the group). So
1334
+ the loop self-paces to the consumers' pull rate with at most
1335
+ `gather_lookahead + 1` groups resident. Runs on every rank; only the
1336
+ sender has an `update_future` to fail fast on."""
1337
+ assert self.source is not None # guaranteed by trainer_init
1338
+ # The live IDS ride along with the count: the free barrier needs only
1339
+ # the count, but the slot-sharing rendezvous needs to know WHICH
1340
+ # consumers are live.
1341
+ self._rpc("begin_sync", live_count, live_ids)
1342
+ # One generator resume per GROUP: `iter_groups` yields (names, tensors)
1343
+ # per owned group in metadata order. Every owner must reach a group in the
1344
+ # same order or their shared gather collective mismatches.
1345
+ groups = self.source.iter_groups()
1346
+ # Publishes fire without an inline ray.get and are harvested this window
1347
+ # deep, so the RPC and server-side rebuild overlap the NEXT group's
1348
+ # gather/export. The window also surfaces a server-side error at depth 2
1349
+ # rather than at end_sync.
1350
+ _PUBLISH_WINDOW = 2
1351
+ pending_publish: list = []
1352
+ # The memory bound: gate BEFORE gathering while more
1353
+ # than `bound` groups are unfreed, so at most bound + 1 are resident. The
1354
+ # count is exact — `_inflight` gains its entry right after each gather and
1355
+ # only wait_freed/end_sync shrinks it, so a group freed server-side but not
1356
+ # yet collected still holds its refs and still counts. The publish window
1357
+ # drains first, so every resident group is pullable while we wait.
1358
+ #
1359
+ # 0 is legal and means NO pipelining: the loop waits for group N to be
1360
+ # freed by every live consumer before gathering N+1, so exactly one group
1361
+ # is resident and the sync costs the full serialized
1362
+ # gather + publish + pull + free RTT per group. 1 (the default) overlaps
1363
+ # N+1's gather with N's pulls.
1364
+ bound = max(0, self._init_info.gather_lookahead)
1365
+
1366
+ # [RDT-EXPORT-RING] `bound + 1` slots is exactly the residency the credit
1367
+ # gate enforces. Small storages are packed into a slot and exported once
1368
+ # per slot instead of once per storage (518 -> 7 cudaIpcGetMemHandle
1369
+ # calls per rank per sync at 235B, each 0.5-1.5ms).
1370
+ nslots = bound + 1
1371
+ # data_ptr -> (storage weakref, reduce_args), this sync only. `storages`
1372
+ # is per-group, so a source that reuses a buffer across groups would
1373
+ # otherwise re-export it every group. The weakref is load-bearing: a
1374
+ # freed allocation can be reissued at the same address, and serving that
1375
+ # from a stale handle is silently wrong bytes.
1376
+ handle_cache: dict[int, tuple] = {}
1377
+ if len(getattr(self, "_export_ring", [])) != nslots:
1378
+ self._export_ring = [None] * nslots
1379
+ self._export_ring_args = [None] * nslots
1380
+ # Slots are handed out per group and returned by _drop_inflight, so the
1381
+ # pool starts full every sync.
1382
+ self._slot_of_group = {}
1383
+ self._free_slots = list(range(nslots))
1384
+
1385
+ try:
1386
+ for gi in self._owned_idx:
1387
+ group = self._groups[gi]
1388
+ while len(self._inflight) > bound:
1389
+ while pending_publish:
1390
+ self._await_publish(pending_publish.pop(0))
1391
+ self._drop_inflight(self._rpc("wait_freed"))
1392
+ names, tensors = next(groups)
1393
+ if list(names) != list(group):
1394
+ raise RuntimeError(
1395
+ f"WeightSource group yielded {len(names)} names starting "
1396
+ f"{names[:2]!r} but expected {len(group)} starting "
1397
+ f"{group[:2]!r}; iteration order must match metadata."
1398
+ )
1399
+ self._validate_held_yields(gi, names, tensors)
1400
+ # Share each unique STORAGE once and describe every name as an
1401
+ # as_strided view onto it. ``None`` means a name this rank does not
1402
+ # hold (a foreign expert); the source keeps it in the list so the
1403
+ # order check stays rank-uniform, and it is dropped here, before
1404
+ # the IPC export, matching the sidecar's served_names.
1405
+ storages: dict[int, tuple] = {}
1406
+ views: dict[str, tuple] = {}
1407
+ refs: dict[str, torch.Tensor] = {}
1408
+ small: list = []
1409
+ for name, tensor in zip(names, tensors):
1410
+ if tensor is None:
1411
+ continue
1412
+ tensor = tensor.detach()
1413
+ if not tensor.is_cuda:
1414
+ tensor = tensor.cuda()
1415
+ # Threshold on the STORAGE, not the tensor: a source may hand
1416
+ # out many small views of one large stacked buffer, and a
1417
+ # per-tensor test would pack the whole thing, costing +1
1418
+ # group of peak memory for a storage already coalesced.
1419
+ if tensor.untyped_storage().nbytes() <= _EXPORT_RING_MAX_BYTES:
1420
+ small.append((name, tensor))
1421
+ continue
1422
+ tensor = tensor.contiguous()
1423
+ refs[name] = tensor # keep the export alive
1424
+ ust = tensor.untyped_storage()
1425
+ sid = ust.data_ptr()
1426
+ if sid not in storages:
1427
+ cached = handle_cache.get(sid)
1428
+ if cached is not None and not cached[0].expired():
1429
+ reduce_args = cached[1]
1430
+ else:
1431
+ base = torch.empty(
1432
+ 0, dtype=torch.uint8, device=tensor.device
1433
+ )
1434
+ base.set_(ust, 0, (ust.nbytes(),))
1435
+ _rebuild, reduce_args = reduce_tensor(base)
1436
+ handle_cache[sid] = (StorageWeakRef(ust), reduce_args)
1437
+ storages[sid] = reduce_args
1438
+ views[name] = (
1439
+ sid,
1440
+ str(tensor.dtype).split(".")[-1],
1441
+ list(tensor.shape),
1442
+ list(tensor.stride()),
1443
+ tensor.storage_offset(),
1444
+ )
1445
+ if small:
1446
+ # Take a slot no live group holds. The gate above already
1447
+ # guarantees one exists, so this is a backstop -- but it
1448
+ # must drain publishes first, like the credit gate.
1449
+ while not self._free_slots:
1450
+ while pending_publish:
1451
+ self._await_publish(pending_publish.pop(0))
1452
+ self._drop_inflight(self._rpc("wait_freed"))
1453
+ slot_idx = self._free_slots.pop(0)
1454
+ self._slot_of_group[gi] = slot_idx
1455
+ st, vw, rf = self._pack_group_for_export(small, slot_idx)
1456
+ storages.update(st)
1457
+ views.update(vw)
1458
+ refs.update(rf)
1459
+ del small
1460
+ del tensors
1461
+ # The loop leaves its last `tensor` / `ust` / `base` bound
1462
+ # through the credit gate and into the next gather, and `base`
1463
+ # is a whole-storage view -- so one tensor and its entire
1464
+ # allocation stay alive past `_drop_inflight`, unseen by the
1465
+ # `_inflight` accounting the residency bound rests on.
1466
+ # Assignment, not `del`: unbound if no direct-path tensor.
1467
+ tensor = ust = base = None
1468
+ if not refs:
1469
+ # A group with nothing held here cannot occur (every group
1470
+ # carries replicated names), but publishing an empty group
1471
+ # would park a credit nobody's pull is waiting on; skip.
1472
+ continue
1473
+ # Hold our refs before publishing; drop them only when the
1474
+ # server reports the group freed (IPC export must outlive import).
1475
+ self._inflight[gi] = refs
1476
+ pending_publish.append(self._publish_async(gi, (storages, views)))
1477
+ while len(pending_publish) >= _PUBLISH_WINDOW:
1478
+ self._await_publish(pending_publish.pop(0))
1479
+ if update_future is not None and update_future.done():
1480
+ # update_weights returned/failed early — surface now instead
1481
+ # of blocking further gathers.
1482
+ update_future.result()
1483
+ while pending_publish:
1484
+ self._await_publish(pending_publish.pop(0))
1485
+ freed = self._rpc("end_sync")
1486
+ self._drop_inflight(freed)
1487
+ except BaseException as e:
1488
+ with contextlib.suppress(Exception):
1489
+ self._rpc("set_gather_error", repr(e))
1490
+ self._inflight.clear()
1491
+ raise
1492
+ finally:
1493
+ # Keep the slots, drop the handles: the sidecar's base cache is
1494
+ # cleared per sync, so a handle must not outlive the sync either.
1495
+ self._export_ring_args = [None] * len(getattr(self, "_export_ring", []))
1496
+
1497
+ def _drop_inflight(self, freed_keys: list) -> None:
1498
+ """Release a freed group's refs and return its export-ring slot."""
1499
+ free_slots = getattr(self, "_free_slots", None)
1500
+ slot_of_group = getattr(self, "_slot_of_group", None)
1501
+ for k in freed_keys:
1502
+ gi = int(k)
1503
+ self._inflight.pop(gi, None)
1504
+ if slot_of_group is None:
1505
+ continue
1506
+ slot = slot_of_group.pop(gi, None)
1507
+ if slot is not None and free_slots is not None and slot not in free_slots:
1508
+ free_slots.append(slot)
1509
+
1510
+ # ---------------- misc ----------------
1511
+
1512
+ def shutdown(self) -> None:
1513
+ if self._server is None:
1514
+ return
1515
+ import ray
1516
+
1517
+ with contextlib.suppress(Exception):
1518
+ ray.get(self._server.shutdown.remote())
1519
+ ray.kill(self._server)
1520
+ self._server = None
1521
+ self._inflight.clear()
vllm/v1/executor/ray_executor_v2.py
+14
@@ -361,6 +361,19 @@ class RayExecutorV2(MultiprocExecutor):
361
runtime_env = self._build_runtime_env()
362
resource_kwargs = self._get_actor_resource_kwargs()
363
364
+ # The sharded_rdt weight transfer backend pulls weights via Ray's tensor
365
+ # transport (@ray.method(tensor_transport=...)). Ray requires the
366
+ # *calling* actor (the vLLM worker) to opt in via enable_tensor_transport.
367
+ wt_cfg = self.vllm_config.weight_transfer_config
368
+ extra_actor_options: dict[str, object] = {}
369
+ if wt_cfg is not None and wt_cfg.backend == "sharded_rdt":
370
+ from vllm.distributed.weight_transfer.sharded_rdt_common import (
371
+ check_ray_rdt_version,
372
+ )
373
+
374
+ check_ray_rdt_version()
375
+ extra_actor_options["enable_tensor_transport"] = True
376
+
377
for bundle_idx in range(self.world_size):
378
bundle = bundle_assignments[bundle_idx]
379
is_driver_worker = self._is_driver_worker(bundle["rank"])
@@ -383,6 +396,7 @@ class RayExecutorV2(MultiprocExecutor):
396
**resource_kwargs,
397
scheduling_strategy=scheduling_strategy,
398
runtime_env=runtime_env,
399
+ **extra_actor_options,
400
)
401
.remote(
402
vllm_config=self.vllm_config,