Weight sync refactor + move sparse nccl engine (#44353)
Signed-off-by: hao-aaron <ahao@anyscale.com> Signed-off-by: haoaaron <ahao@anyscale.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Aaron Hao committed
Jul 1, 2026 at 17:25 UTC
77a9c5ae28a3d054e6caf60c7e14082453b3ae47
32 files changed
+953
-1046
docs/training/layerwise.md
+1
-1
@@ -58,7 +58,7 @@ class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
58
59
### High Level Weight Transfer API
60
61
-The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
61
+The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Checkpoint-format weight transfer engines (e.g. the NCCL and IPC backends) run layerwise reloading automatically inside their `start_weight_update`/`finish_weight_update` lifecycle.
62
63
### Mid Level `reload_weights` API
64
docs/training/weight_transfer/README.md
+3
-2
@@ -17,6 +17,7 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
17
| ------- | --------- | -------- |
18
| [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
19
| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
20
+| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
21
22
## Configuration
23
@@ -41,7 +42,7 @@ vllm serve my-model \
42
--weight-transfer-config '{"backend": "nccl"}'
43
```
44
44
-The `backend` field accepts `"nccl"` (default) or `"ipc"`.
45
+The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
46
47
## API Endpoints
48
@@ -69,7 +70,7 @@ Both backends provide static methods that the trainer calls to send weights. The
70
EngineClass.trainer_init(init_info)
71
72
# 2. Start weight update on inference side
72
-llm.start_weight_update(is_checkpoint_format=True)
73
+llm.start_weight_update()
74
75
# 3. Send weights to inference workers
76
EngineClass.trainer_send_weights(
docs/training/weight_transfer/base.md
+28
-15
@@ -11,15 +11,23 @@ The `WeightTransferEngine` is a generic abstract class parameterized by two data
11
12
### Abstract Methods
13
14
-Subclasses must implement these four methods:
14
+Subclasses must implement these methods:
15
16
| Method | Side | Description |
17
| ------ | ---- | ----------- |
18
| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
19
-| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
19
+| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
20
+| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
21
+| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` |
22
| `shutdown()` | Inference | Clean up resources |
23
| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
24
25
+The base class provides two methods:
26
+
27
+1. `__init__` : Engines receive `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`) and `model` (`nn.Module`)
28
+2. `update_weights(update_info_dict)`: Thin wrapper for `receive_weights`: parses
29
+the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`.
30
+
31
### Request Classes
32
33
The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
@@ -81,7 +89,7 @@ class MyUpdateInfo(WeightTransferUpdateInfo):
89
### 2. Implement the Engine
90
91
```python
84
-from collections.abc import Callable, Iterator
92
+from collections.abc import Iterator
93
from typing import Any
94
import torch
95
@@ -93,18 +101,25 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
101
# Set up connection to trainer using init_info.endpoint, etc.
102
...
103
96
- def receive_weights(
97
- self,
98
- update_info: MyUpdateInfo,
99
- load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
100
- ) -> None:
101
- # Receive each weight and call load_weights incrementally
104
+ def start_weight_update(self) -> None:
105
+ # Checkpoint-format engines: run initialize_layerwise_reload(self.model).
106
+ # In-place engines: no-op
107
+ ...
108
+
109
+ def finish_weight_update(self) -> None:
110
+ # Checkpoint-format engines: run finalize_layerwise_reload(...).
111
+ # In-place engines: no-op
112
+ ...
113
+
114
+ def receive_weights(self, update_info: MyUpdateInfo) -> None:
115
+ weights = []
116
for name, dtype_name, shape in zip(
117
update_info.names, update_info.dtype_names, update_info.shapes
118
):
119
dtype = getattr(torch, dtype_name)
120
weight = self._fetch_weight(name, shape, dtype)
107
- load_weights([(name, weight)])
121
+ weights.append((name, weight))
122
+ self.model.load_weights(weights)
123
124
def shutdown(self) -> None:
125
# Clean up resources
@@ -121,9 +136,6 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
136
...
137
```
138
124
-!!! important
125
- The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models.
126
-
139
### 3. Register with the Factory
140
141
```python
@@ -147,7 +159,7 @@ Once registered, users can select your backend via `WeightTransferConfig(backend
159
160
## WeightTransferEngineFactory
161
150
-The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) 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.
162
+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.
163
164
```python
165
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
@@ -155,7 +167,8 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
167
# Create an engine from config
168
engine = WeightTransferEngineFactory.create_engine(
169
config=weight_transfer_config,
158
- parallel_config=parallel_config,
170
+ vllm_config=vllm_config,
171
+ device=device,
172
model=model,
173
)
174
```
docs/training/weight_transfer/ipc.md
+2
-2
@@ -55,7 +55,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
55
llm_handle=llm_actor_handle,
56
)
57
# start
58
-ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))
58
+ray.get(llm_actor_handle.start_weight_update.remote())
59
# send weights
60
IPCWeightTransferEngine.trainer_send_weights(
61
iterator=model.named_parameters(),
@@ -80,7 +80,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
80
# start
81
base_url = "http://localhost:8000"
82
url = f"{base_url}/start_weight_update"
83
-response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60)
83
+response = requests.post(url, json={}, timeout=60)
84
response.raise_for_status()
85
# send weights
86
IPCWeightTransferEngine.trainer_send_weights(
docs/training/weight_transfer/nccl.md
+14
-11
@@ -11,7 +11,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
11
## How It Works
12
13
1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
14
-2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
14
+2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
15
3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
16
17
## Initialization
@@ -93,7 +93,7 @@ remaining three steps are:
93
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
94
95
# 1. Start the weight update
96
-llm.start_weight_update(is_checkpoint_format=True)
96
+llm.start_weight_update()
97
98
# 2. Receive weights (can be called multiple times for chunked transfers)
99
llm.update_weights(
@@ -116,19 +116,22 @@ must match the order in which the trainer iterates over its parameters.
116
117
`start_weight_update` must be called before `update_weights`, and
118
`finish_weight_update` must be called after all weight chunks have been
119
-transferred. The `is_checkpoint_format` flag controls whether layerwise reload
120
-processing is applied (`True` for checkpoint-format weights, `False` for
121
-pre-processed kernel-format weights).
119
+transferred. The NCCL engine receives checkpoint-format weights and applies
120
+layerwise reload processing automatically inside `start_weight_update` /
121
+`finish_weight_update`.
122
123
-Sparse NCCL patches still use `update_kind="sparse_flat"` inside
124
-`update_info`, but they should be wrapped in
125
-`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
126
-directly to runtime/kernel-format parameters. The current sparse MVP requires
127
-`TP=1` and `PP=1`.
123
+## Sparse NCCL
124
+
125
+Sparse, flat-index weight patches use a separate backend,
126
+`WeightTransferConfig(backend="sparse_nccl")`, implemented by
127
+`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group
128
+initialization with the dense engine; patches are applied directly in place to
129
+existing parameters (no layerwise reload). The current sparse MVP requires
130
+`TP=1` and `PP=1`. See the example below.
131
132
## Examples
133
134
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
132
-- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
135
+- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
136
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
137
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
examples/rl/rlhf_async_new_apis.py
+1
-1
@@ -306,7 +306,7 @@ gen_futures = [
306
307
ray.get(llm.pause_after_n_tokens.remote())
308
309
-ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
309
+ray.get(llm.start_weight_update.remote())
310
311
inference_handle = llm.update_weights.remote(
312
WeightTransferUpdateRequest(
examples/rl/rlhf_http_ipc.py
+3
-7
@@ -80,14 +80,10 @@ def init_weight_transfer_engine(base_url: str) -> None:
80
response.raise_for_status()
81
82
83
-def start_weight_update(
84
- base_url: str,
85
- is_checkpoint_format: bool = True,
86
-) -> None:
83
+def start_weight_update(base_url: str) -> None:
84
"""Start a weight update via HTTP endpoint."""
85
url = f"{base_url}/start_weight_update"
89
- payload = {"is_checkpoint_format": is_checkpoint_format}
90
- response = requests.post(url, json=payload, timeout=60)
86
+ response = requests.post(url, json={}, timeout=60)
87
response.raise_for_status()
88
89
@@ -170,7 +166,7 @@ def main():
166
pause_generation(BASE_URL)
167
168
# Start weight update, broadcast via IPC, then finish
173
- start_weight_update(BASE_URL, is_checkpoint_format=False)
169
+ start_weight_update(BASE_URL)
170
171
print("Broadcasting weights via CUDA IPC (HTTP)...")
172
trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
examples/rl/rlhf_http_nccl.py
+3
-7
@@ -83,14 +83,10 @@ def init_weight_transfer_engine(
83
response.raise_for_status()
84
85
86
-def start_weight_update(
87
- base_url: str,
88
- is_checkpoint_format: bool = True,
89
-) -> None:
86
+def start_weight_update(base_url: str) -> None:
87
"""Start a weight update via HTTP endpoint."""
88
url = f"{base_url}/start_weight_update"
92
- payload = {"is_checkpoint_format": is_checkpoint_format}
93
- response = requests.post(url, json=payload, timeout=60)
89
+ response = requests.post(url, json={}, timeout=60)
90
response.raise_for_status()
91
92
@@ -223,7 +219,7 @@ def main():
219
shapes.append(list(p.shape))
220
221
# Start weight update
226
- start_weight_update(BASE_URL, is_checkpoint_format=True)
222
+ start_weight_update(BASE_URL)
223
224
# Start the update_weights call in a separate thread since it will block
225
# waiting for NCCL broadcasts
examples/rl/rlhf_ipc.py
+1
-1
@@ -139,7 +139,7 @@ ray.get(llm.sleep.remote(level=0))
139
140
ray.get(train_model.init_weight_transfer.remote())
141
# Start weight update, sync weights, then finish
142
-ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
142
+ray.get(llm.start_weight_update.remote())
143
ray.get(train_model.broadcast_weights.remote(llm))
144
ray.get(llm.finish_weight_update.remote())
145
examples/rl/rlhf_ipc_fsdp_ep.py
+3
-10
@@ -277,15 +277,8 @@ class DataParallelInferenceEngine:
277
]
278
)
279
280
- def start_weight_update(self, is_checkpoint_format: bool = True):
281
- ray.get(
282
- [
283
- actor.start_weight_update.remote(
284
- is_checkpoint_format=is_checkpoint_format
285
- )
286
- for actor in self.llm_actors
287
- ]
288
- )
280
+ def start_weight_update(self):
281
+ ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
282
283
def finish_weight_update(self):
284
ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
@@ -392,7 +385,7 @@ def main():
385
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
386
387
print("[sync] Starting weight update...")
395
- ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
388
+ ray.get(inference_engine.start_weight_update.remote())
389
390
print("[sync] Packed IPC transfer FSDP → vLLM...")
391
ray.get(
examples/rl/rlhf_nccl.py
+1
-1
@@ -202,7 +202,7 @@ ray.get([train_handle, inference_handle])
202
names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
203
204
# Start weight update
205
-ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
205
+ray.get(llm.start_weight_update.remote())
206
207
# Issue update_weights call with NCCL-specific update info
208
# packed=True enables efficient batched tensor broadcasting
examples/rl/rlhf_nccl_fsdp_ep.py
+1
-1
@@ -299,7 +299,7 @@ async def main():
299
print(f"[sync] Got metadata for {len(names)} parameters.")
300
301
print("[sync] Starting weight update...")
302
- await engine.start_weight_update(is_checkpoint_format=True)
302
+ await engine.start_weight_update()
303
304
print("[sync] Broadcasting weights from FSDP → vLLM...")
305
broadcast_handles = [
examples/rl/rlhf_sparse_nccl.py
+11
-8
@@ -44,11 +44,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
44
45
from vllm import LLM, SamplingParams
46
from vllm.config import WeightTransferConfig
47
-from vllm.distributed.weight_transfer.base import SparseWeightPatch
47
from vllm.distributed.weight_transfer.nccl_engine import (
48
NCCLTrainerSendWeightsArgs,
49
NCCLWeightTransferEngine,
50
)
51
+from vllm.distributed.weight_transfer.sparse_nccl_engine import (
52
+ SparseNCCLWeightTransferEngine,
53
+ SparseWeightPatch,
54
+)
55
from vllm.utils.network_utils import get_ip, get_open_port
56
57
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
@@ -244,7 +247,6 @@ class TrainModel:
247
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
248
shapes=[list(self.patched_param.shape)],
249
num_updates_list=[flat_indices.numel()],
247
- update_kind="sparse_flat",
250
)
251
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
252
@@ -271,7 +273,7 @@ class TrainModel:
273
raise RuntimeError("Sparse patch has not been prepared")
274
275
start = time.perf_counter()
274
- NCCLWeightTransferEngine.trainer_send_sparse_weights(
276
+ SparseNCCLWeightTransferEngine.trainer_send_weights(
277
iter(self.pending_sparse_patches),
278
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
279
)
@@ -282,6 +284,7 @@ class TrainModel:
284
285
def launch_llm(
286
scheduling_inference: PlacementGroupSchedulingStrategy,
287
+ backend: str = "nccl",
288
):
289
return ray.remote(
290
num_cpus=0,
@@ -293,7 +296,7 @@ def launch_llm(
296
tensor_parallel_size=1,
297
distributed_executor_backend="ray",
298
gpu_memory_utilization=0.7,
296
- weight_transfer_config=WeightTransferConfig(backend="nccl"),
299
+ weight_transfer_config=WeightTransferConfig(backend=backend),
300
)
301
302
@@ -332,7 +335,7 @@ def run_dense_phase(
335
scheduling_inference: PlacementGroupSchedulingStrategy,
336
) -> dict[str, object]:
337
ray.get(train_model.reset_model.remote())
335
- llm = launch_llm(scheduling_inference)
338
+ llm = launch_llm(scheduling_inference, backend="nccl")
339
try:
340
dense_before = collect_vllm_generations(llm)
341
@@ -351,7 +354,7 @@ def run_dense_phase(
354
)
355
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
356
ray.get([trainer_init, inference_init])
354
- ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
357
+ ray.get(llm.start_weight_update.remote())
358
359
dense_update_info, dense_payload_bytes = ray.get(
360
train_model.get_dense_update_info.remote()
@@ -391,7 +394,7 @@ def run_sparse_phase(
394
scheduling_inference: PlacementGroupSchedulingStrategy,
395
) -> dict[str, object]:
396
ray.get(train_model.reset_model.remote())
394
- llm = launch_llm(scheduling_inference)
397
+ llm = launch_llm(scheduling_inference, backend="sparse_nccl")
398
try:
399
sparse_before = collect_vllm_generations(llm)
400
@@ -410,7 +413,7 @@ def run_sparse_phase(
413
)
414
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
415
ray.get([trainer_init, inference_init])
413
- ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
416
+ ray.get(llm.start_weight_update.remote())
417
418
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
419
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
tests/distributed/test_weight_transfer.py
+266
-253
@@ -18,7 +18,6 @@ from torch.multiprocessing.reductions import reduce_tensor
18
from vllm.config.parallel import ParallelConfig
19
from vllm.config.weight_transfer import WeightTransferConfig
20
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
21
-from vllm.distributed.weight_transfer.base import SparseWeightPatch
21
from vllm.distributed.weight_transfer.ipc_engine import (
22
IPCWeightTransferEngine,
23
IPCWeightTransferInitInfo,
@@ -29,6 +28,11 @@ from vllm.distributed.weight_transfer.nccl_engine import (
28
NCCLWeightTransferInitInfo,
29
NCCLWeightTransferUpdateInfo,
30
)
31
+from vllm.distributed.weight_transfer.sparse_nccl_engine import (
32
+ SparseNCCLWeightTransferEngine,
33
+ SparseNCCLWeightTransferUpdateInfo,
34
+ SparseWeightPatch,
35
+)
36
from vllm.platforms import current_platform
37
from vllm.utils.network_utils import get_open_port
38
@@ -80,6 +84,18 @@ def create_mock_parallel_config(
84
return config
85
86
87
+def create_mock_vllm_config(
88
+ rank: int = 0,
89
+ world_size: int = 1,
90
+ dp_rank: int = 0,
91
+) -> MagicMock:
92
+ """Create a mock VllmConfig exposing parallel_config and model_config."""
93
+ vllm_config = MagicMock()
94
+ vllm_config.parallel_config = create_mock_parallel_config(rank, world_size, dp_rank)
95
+ vllm_config.model_config = MagicMock()
96
+ return vllm_config
97
+
98
+
99
# --- Unit Tests: NCCLWeightTransferUpdateInfo Validation ---
100
101
@@ -87,7 +103,6 @@ class TestNCCLWeightTransferUpdateInfoValidation:
103
"""Test NCCLWeightTransferUpdateInfo dataclass validation."""
104
105
def test_valid_update_info(self):
90
- """Test creating valid NCCLWeightTransferUpdateInfo."""
106
info = NCCLWeightTransferUpdateInfo(
107
names=["layer.weight", "layer.bias"],
108
dtype_names=["float32", "float32"],
@@ -98,7 +113,6 @@ class TestNCCLWeightTransferUpdateInfoValidation:
113
assert info.shapes == [[10, 10], [10]]
114
115
def test_mismatched_dtype_names_raises(self):
101
- """Test that mismatched dtype_names length raises ValueError."""
116
with pytest.raises(ValueError, match="dtype_names"):
117
NCCLWeightTransferUpdateInfo(
118
names=["layer.weight", "layer.bias"],
@@ -107,7 +121,6 @@ class TestNCCLWeightTransferUpdateInfoValidation:
121
)
122
123
def test_mismatched_shapes_raises(self):
110
- """Test that mismatched shapes length raises ValueError."""
124
with pytest.raises(ValueError, match="shapes"):
125
NCCLWeightTransferUpdateInfo(
126
names=["layer.weight", "layer.bias"],
@@ -116,73 +129,59 @@ class TestNCCLWeightTransferUpdateInfoValidation:
129
)
130
131
def test_empty_lists_valid(self):
119
- """Test that empty lists are valid."""
120
- info = NCCLWeightTransferUpdateInfo(
121
- names=[],
122
- dtype_names=[],
123
- shapes=[],
124
- )
132
+ info = NCCLWeightTransferUpdateInfo(names=[], dtype_names=[], shapes=[])
133
assert len(info.names) == 0
134
135
+
136
+# --- Unit Tests: SparseNCCLWeightTransferUpdateInfo Validation ---
137
+
138
+
139
+class TestSparseNCCLWeightTransferUpdateInfoValidation:
140
+ """Test SparseNCCLWeightTransferUpdateInfo dataclass validation."""
141
+
142
def test_valid_sparse_update_info(self):
128
- """Test creating valid sparse NCCL update info."""
129
- info = NCCLWeightTransferUpdateInfo(
143
+ info = SparseNCCLWeightTransferUpdateInfo(
144
names=["layer.weight", "layer.bias"],
145
dtype_names=["float32", "bfloat16"],
146
shapes=[[10, 10], [10]],
147
num_updates_list=[4, 2],
134
- update_kind="sparse_flat",
148
)
136
- assert info.update_kind == "sparse_flat"
149
assert info.num_updates_list == [4, 2]
150
139
- def test_sparse_update_requires_num_updates_list(self):
140
- with pytest.raises(ValueError, match="`num_updates_list` is required"):
141
- NCCLWeightTransferUpdateInfo(
142
- names=["layer.weight"],
151
+ def test_mismatched_dtype_names_raises(self):
152
+ with pytest.raises(ValueError, match="dtype_names"):
153
+ SparseNCCLWeightTransferUpdateInfo(
154
+ names=["layer.weight", "layer.bias"],
155
dtype_names=["float32"],
144
- shapes=[[10, 10]],
145
- update_kind="sparse_flat",
156
+ shapes=[[10, 10], [10]],
157
+ num_updates_list=[4, 2],
158
)
159
148
- def test_sparse_update_rejects_empty_num_updates_list(self):
160
+ def test_rejects_empty_num_updates_list(self):
161
with pytest.raises(ValueError, match="cannot be empty"):
150
- NCCLWeightTransferUpdateInfo(
162
+ SparseNCCLWeightTransferUpdateInfo(
163
names=[],
164
dtype_names=[],
165
shapes=[],
166
num_updates_list=[],
155
- update_kind="sparse_flat",
167
)
168
158
- def test_sparse_update_rejects_packed(self):
159
- with pytest.raises(ValueError, match="cannot be combined with `packed=True`"):
160
- NCCLWeightTransferUpdateInfo(
161
- names=["layer.weight"],
162
- dtype_names=["float32"],
163
- shapes=[[10, 10]],
164
- num_updates_list=[3],
165
- update_kind="sparse_flat",
166
- packed=True,
167
- )
168
-
169
- def test_sparse_update_rejects_mismatched_num_updates(self):
169
+ def test_rejects_mismatched_num_updates(self):
170
with pytest.raises(ValueError, match="`num_updates_list`"):
171
- NCCLWeightTransferUpdateInfo(
171
+ SparseNCCLWeightTransferUpdateInfo(
172
names=["layer.weight", "layer.bias"],
173
dtype_names=["float32", "float32"],
174
shapes=[[10, 10], [10]],
175
num_updates_list=[3],
176
- update_kind="sparse_flat",
176
)
177
179
- def test_dense_update_rejects_sparse_metadata(self):
180
- with pytest.raises(ValueError, match="Sparse metadata"):
181
- NCCLWeightTransferUpdateInfo(
178
+ def test_rejects_negative_num_updates(self):
179
+ with pytest.raises(ValueError, match="non-negative"):
180
+ SparseNCCLWeightTransferUpdateInfo(
181
names=["layer.weight"],
182
dtype_names=["float32"],
183
shapes=[[10, 10]],
185
- num_updates_list=[3],
184
+ num_updates_list=[-1],
185
)
186
187
@@ -192,14 +191,17 @@ class TestNCCLWeightTransferUpdateInfoValidation:
191
class TestNCCLEngineParsing:
192
"""Test NCCLWeightTransferEngine parsing methods."""
193
195
- def test_parse_init_info_valid(self):
196
- """Test parsing valid init info dict."""
194
+ def _make_engine(self):
195
config = WeightTransferConfig(backend="nccl")
198
- parallel_config = create_mock_parallel_config()
199
- engine = NCCLWeightTransferEngine(
200
- config, parallel_config, MagicMock(spec=torch.nn.Module)
196
+ return NCCLWeightTransferEngine(
197
+ config,
198
+ create_mock_vllm_config(),
199
+ "cuda",
200
+ MagicMock(spec=torch.nn.Module),
201
)
202
203
+ def test_parse_init_info_valid(self):
204
+ engine = self._make_engine()
205
init_info = engine.parse_init_info(
206
{
207
"master_address": "127.0.0.1",
@@ -208,7 +210,6 @@ class TestNCCLEngineParsing:
210
"world_size": 3,
211
}
212
)
211
-
213
assert isinstance(init_info, NCCLWeightTransferInitInfo)
214
assert init_info.master_address == "127.0.0.1"
215
assert init_info.master_port == 12345
@@ -216,29 +217,12 @@ class TestNCCLEngineParsing:
217
assert init_info.world_size == 3
218
219
def test_parse_init_info_missing_field_raises(self):
219
- """Test parsing init info with missing required field."""
220
- config = WeightTransferConfig(backend="nccl")
221
- parallel_config = create_mock_parallel_config()
222
- engine = NCCLWeightTransferEngine(
223
- config, parallel_config, MagicMock(spec=torch.nn.Module)
224
- )
225
-
220
+ engine = self._make_engine()
221
with pytest.raises(ValueError, match="Invalid init_info"):
227
- engine.parse_init_info(
228
- {
229
- "master_address": "127.0.0.1",
230
- # Missing master_port, rank_offset, world_size
231
- }
232
- )
222
+ engine.parse_init_info({"master_address": "127.0.0.1"})
223
224
def test_parse_update_info_valid(self):
235
- """Test parsing valid update info dict."""
236
- config = WeightTransferConfig(backend="nccl")
237
- parallel_config = create_mock_parallel_config()
238
- engine = NCCLWeightTransferEngine(
239
- config, parallel_config, MagicMock(spec=torch.nn.Module)
240
- )
241
-
225
+ engine = self._make_engine()
226
update_info = engine.parse_update_info(
227
{
228
"names": ["w1", "w2"],
@@ -246,7 +230,6 @@ class TestNCCLEngineParsing:
230
"shapes": [[100, 100], [50]],
231
}
232
)
249
-
233
assert isinstance(update_info, NCCLWeightTransferUpdateInfo)
234
assert update_info.names == ["w1", "w2"]
235
assert update_info.dtype_names == ["float32", "bfloat16"]
@@ -260,40 +243,140 @@ class TestEngineRegistry:
243
"""Test weight transfer engine registry."""
244
245
def test_create_engine_nccl(self):
263
- """Test factory creates NCCL engine."""
246
config = WeightTransferConfig(backend="nccl")
265
- parallel_config = create_mock_parallel_config()
247
engine = WeightTransferEngineFactory.create_engine(
267
- config, parallel_config, MagicMock(spec=torch.nn.Module)
248
+ config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
249
)
250
assert isinstance(engine, NCCLWeightTransferEngine)
251
252
def test_create_engine_ipc(self):
272
- """Test factory creates IPC engine."""
253
config = WeightTransferConfig(backend="ipc")
274
- parallel_config = create_mock_parallel_config()
254
engine = WeightTransferEngineFactory.create_engine(
276
- config, parallel_config, MagicMock(spec=torch.nn.Module)
255
+ config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
256
)
257
assert isinstance(engine, IPCWeightTransferEngine)
258
259
+ def test_create_engine_sparse_nccl(self):
260
+ config = WeightTransferConfig(backend="sparse_nccl")
261
+ engine = WeightTransferEngineFactory.create_engine(
262
+ config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
263
+ )
264
+ assert isinstance(engine, SparseNCCLWeightTransferEngine)
265
+
266
def test_create_engine_invalid_backend(self):
281
- """Test factory raises for invalid backend."""
267
config = WeightTransferConfig(backend="invalid")
283
- parallel_config = create_mock_parallel_config()
268
with pytest.raises(ValueError, match="Invalid weight transfer backend"):
269
WeightTransferEngineFactory.create_engine(
286
- config, parallel_config, MagicMock(spec=torch.nn.Module)
270
+ config,
271
+ create_mock_vllm_config(),
272
+ "cuda",
273
+ MagicMock(spec=torch.nn.Module),
274
)
275
276
def test_register_duplicate_raises(self):
290
- """Test registering duplicate engine name raises."""
277
with pytest.raises(ValueError, match="already registered"):
278
WeightTransferEngineFactory.register_engine(
279
"nccl", NCCLWeightTransferEngine
280
)
281
282
283
+# --- Unit Tests: Sparse patch application (CPU) ---
284
+
285
+
286
+class TestSparseNCCLPatchApplication:
287
+ """Test SparseNCCLWeightTransferEngine._apply_patch on a real param."""
288
+
289
+ def _make_engine(self, model):
290
+ config = WeightTransferConfig(backend="sparse_nccl")
291
+ return SparseNCCLWeightTransferEngine(
292
+ config, create_mock_vllm_config(), "cpu", model
293
+ )
294
+
295
+ def _make_model(self, numel: int = 8):
296
+ model = torch.nn.Module()
297
+ model.register_parameter(
298
+ "w", torch.nn.Parameter(torch.zeros(numel), requires_grad=False)
299
+ )
300
+
301
+ def get_parameter(name):
302
+ assert name == "w"
303
+ return model.w
304
+
305
+ model.get_parameter = get_parameter
306
+ return model
307
+
308
+ def test_apply_patch_updates_only_selected_entries(self):
309
+ model = self._make_model(8)
310
+ engine = self._make_engine(model)
311
+ engine._apply_patch(
312
+ SparseWeightPatch(
313
+ name="w",
314
+ indices=torch.tensor([1, 3], dtype=torch.int32),
315
+ values=torch.tensor([5.0, 7.0], dtype=torch.float32),
316
+ )
317
+ )
318
+ expected = torch.zeros(8)
319
+ expected[1] = 5.0
320
+ expected[3] = 7.0
321
+ assert torch.equal(model.w.data, expected)
322
+
323
+ def test_apply_patch_rejects_mismatched_lengths(self):
324
+ model = self._make_model(8)
325
+ engine = self._make_engine(model)
326
+ with pytest.raises(ValueError, match="matching lengths"):
327
+ engine._apply_patch(
328
+ SparseWeightPatch(
329
+ name="w",
330
+ indices=torch.tensor([1, 3], dtype=torch.int32),
331
+ values=torch.tensor([5.0], dtype=torch.float32),
332
+ )
333
+ )
334
+
335
+ def test_apply_patch_rejects_non_int32_indices(self):
336
+ model = self._make_model(8)
337
+ engine = self._make_engine(model)
338
+ with pytest.raises(ValueError, match="int32 indices"):
339
+ engine._apply_patch(
340
+ SparseWeightPatch(
341
+ name="w",
342
+ indices=torch.tensor([1], dtype=torch.int64),
343
+ values=torch.tensor([5.0], dtype=torch.float32),
344
+ )
345
+ )
346
+
347
+ def test_apply_patch_rejects_dtype_mismatch(self):
348
+ model = self._make_model(8)
349
+ engine = self._make_engine(model)
350
+ with pytest.raises(ValueError, match="does not match"):
351
+ engine._apply_patch(
352
+ SparseWeightPatch(
353
+ name="w",
354
+ indices=torch.tensor([1], dtype=torch.int32),
355
+ values=torch.tensor([5.0], dtype=torch.bfloat16),
356
+ )
357
+ )
358
+
359
+ def test_apply_patch_rejects_non_contiguous_param(self):
360
+ model = torch.nn.Module()
361
+ model.register_parameter(
362
+ "w",
363
+ torch.nn.Parameter(
364
+ torch.arange(12, dtype=torch.float32).view(3, 4).t(),
365
+ requires_grad=False,
366
+ ),
367
+ )
368
+ model.get_parameter = lambda name: model.w
369
+ engine = self._make_engine(model)
370
+ with pytest.raises(NotImplementedError, match="contiguous params"):
371
+ engine._apply_patch(
372
+ SparseWeightPatch(
373
+ name="w",
374
+ indices=torch.tensor([1], dtype=torch.int32),
375
+ values=torch.tensor([1.0], dtype=torch.float32),
376
+ )
377
+ )
378
+
379
+
380
# --- Test receive_weights without init raises ---
381
382
@@ -303,42 +386,37 @@ def test_nccl_receive_weights_without_init_raises():
386
pytest.skip("Need at least 1 GPU for this test")
387
388
config = WeightTransferConfig(backend="nccl")
306
- parallel_config = create_mock_parallel_config()
389
engine = NCCLWeightTransferEngine(
308
- config, parallel_config, MagicMock(spec=torch.nn.Module)
390
+ config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
391
)
392
393
update_info = NCCLWeightTransferUpdateInfo(
312
- names=["w"],
313
- dtype_names=["float32"],
314
- shapes=[[10]],
394
+ names=["w"], dtype_names=["float32"], shapes=[[10]]
395
)
396
397
with pytest.raises(RuntimeError, match="not initialized"):
318
- engine.receive_weights(update_info, lambda x: None)
398
+ engine.receive_weights(update_info)
399
400
321
-def test_nccl_receive_sparse_weights_without_init_raises():
401
+def test_sparse_nccl_receive_weights_without_init_raises():
402
"""Test that sparse receive raises if init_transfer_engine wasn't called."""
403
if torch.accelerator.device_count() < 1:
404
pytest.skip("Need at least 1 GPU for this test")
405
326
- config = WeightTransferConfig(backend="nccl")
327
- parallel_config = create_mock_parallel_config()
328
- engine = NCCLWeightTransferEngine(
329
- config, parallel_config, MagicMock(spec=torch.nn.Module)
406
+ config = WeightTransferConfig(backend="sparse_nccl")
407
+ engine = SparseNCCLWeightTransferEngine(
408
+ config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
409
)
410
332
- update_info = NCCLWeightTransferUpdateInfo(
411
+ update_info = SparseNCCLWeightTransferUpdateInfo(
412
names=["w"],
413
dtype_names=["float32"],
414
shapes=[[10]],
415
num_updates_list=[2],
337
- update_kind="sparse_flat",
416
)
417
418
with pytest.raises(RuntimeError, match="not initialized"):
341
- engine.receive_sparse_weights(update_info, lambda x: None)
419
+ engine.receive_weights(update_info)
420
421
422
# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
@@ -387,6 +465,7 @@ def inference_receive_tensor(
465
tensor_dtype: str,
466
) -> dict:
467
"""Inference task that receives tensor via NCCLWeightTransferEngine."""
468
+ import contextlib
469
from unittest.mock import MagicMock
470
471
import torch
@@ -401,17 +480,32 @@ def inference_receive_tensor(
480
NCCLWeightTransferUpdateInfo,
481
)
482
404
- # Create engine with mock parallel config
483
+ class Recorder(torch.nn.Module):
484
+ def __init__(self):
485
+ super().__init__()
486
+ self.received = []
487
+
488
+ def load_weights(self, weights):
489
+ for name, tensor in weights:
490
+ self.received.append((name, tensor.clone()))
491
+
492
config = WeightTransferConfig(backend="nccl")
493
+ vllm_config = MagicMock()
494
parallel_config = MagicMock(spec=ParallelConfig)
495
parallel_config.rank = 0
496
parallel_config.world_size = 1
497
parallel_config.data_parallel_rank = 0
498
parallel_config.data_parallel_index = 0
499
+ vllm_config.parallel_config = parallel_config
500
+ vllm_config.model_config = MagicMock()
501
412
- engine = NCCLWeightTransferEngine(
413
- config, parallel_config, MagicMock(spec=torch.nn.Module)
414
- )
502
+ recorder = Recorder()
503
+ engine = NCCLWeightTransferEngine(config, vllm_config, "cuda", recorder)
504
+ # Transport-only test: bypass the set_current_vllm_config context that
505
+ # receive_weights enters, since vllm_config here is a mock.
506
+ import vllm.config as _vllm_config_mod
507
+
508
+ _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
509
510
# Initialize the engine (joins as rank 1)
511
init_info = NCCLWeightTransferInitInfo(
@@ -422,20 +516,12 @@ def inference_receive_tensor(
516
)
517
engine.init_transfer_engine(init_info)
518
425
- # Receive weights with a no-op load_weights that captures the tensor
426
- received_tensors = []
427
-
428
- def noop_load_weights(weights: list[tuple[str, torch.Tensor]]):
429
- for name, tensor in weights:
430
- # Clone tensor to keep it after engine cleans up
431
- received_tensors.append((name, tensor.clone()))
432
-
519
update_info = NCCLWeightTransferUpdateInfo(
520
names=["test.weight"],
521
dtype_names=[tensor_dtype],
522
shapes=[tensor_shape],
523
)
438
- engine.receive_weights(update_info, noop_load_weights)
524
+ engine.receive_weights(update_info)
525
torch.accelerator.synchronize()
526
527
# Verify we received the tensor
@@ -443,11 +529,10 @@ def inference_receive_tensor(
529
received_shape = None
530
received_sum = None
531
446
- if len(received_tensors) == 1:
447
- name, tensor = received_tensors[0]
532
+ if len(recorder.received) == 1:
533
+ name, tensor = recorder.received[0]
534
received_shape = list(tensor.shape)
535
received_sum = tensor.sum().item()
450
- # Check shape matches and values are all 1s (trainer sends ones)
536
if received_shape == tensor_shape:
537
expected_sum = 1.0 * torch.tensor(tensor_shape).prod().item()
538
if abs(received_sum - expected_sum) < 0.01:
@@ -478,11 +563,9 @@ def test_nccl_weight_transfer_between_processes():
563
master_port = get_open_port()
564
world_size = 2 # 1 trainer + 1 inference worker
565
481
- # Tensor to transfer: 100x100 ones
566
tensor_shape = [100, 100]
567
tensor_dtype = "float32"
568
485
- # Start both tasks concurrently - Ray assigns GPUs automatically
569
inference_future = inference_receive_tensor.remote(
570
master_address, master_port, world_size, tensor_shape, tensor_dtype
571
)
@@ -490,7 +573,6 @@ def test_nccl_weight_transfer_between_processes():
573
master_address, master_port, world_size, tensor_shape, tensor_dtype
574
)
575
493
- # Wait for both to complete
576
trainer_result, result = ray.get([trainer_future, inference_future])
577
578
assert trainer_result, "Trainer should complete successfully"
@@ -514,10 +596,12 @@ def trainer_broadcast_sparse_tensor(
596
597
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
598
from vllm.distributed.utils import StatelessProcessGroup
517
- from vllm.distributed.weight_transfer.base import SparseWeightPatch
599
from vllm.distributed.weight_transfer.nccl_engine import (
600
NCCLTrainerSendWeightsArgs,
520
- NCCLWeightTransferEngine,
601
+ )
602
+ from vllm.distributed.weight_transfer.sparse_nccl_engine import (
603
+ SparseNCCLWeightTransferEngine,
604
+ SparseWeightPatch,
605
)
606
607
pg = StatelessProcessGroup.create(
@@ -533,7 +617,7 @@ def trainer_broadcast_sparse_tensor(
617
indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device),
618
values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device),
619
)
536
- NCCLWeightTransferEngine.trainer_send_sparse_weights(
620
+ SparseNCCLWeightTransferEngine.trainer_send_weights(
621
iter([patch]),
622
NCCLTrainerSendWeightsArgs(group=comm),
623
)
@@ -547,7 +631,7 @@ def inference_receive_sparse_tensor(
631
master_port: int,
632
world_size: int,
633
) -> dict:
550
- """Inference task that receives sparse patches via NCCLWeightTransferEngine."""
634
+ """Inference task that receives sparse patches via the sparse engine."""
635
from unittest.mock import MagicMock
636
637
import torch
@@ -556,22 +640,40 @@ def inference_receive_sparse_tensor(
640
641
from vllm.config.parallel import ParallelConfig
642
from vllm.config.weight_transfer import WeightTransferConfig
559
- from vllm.distributed.weight_transfer.nccl_engine import (
560
- NCCLWeightTransferEngine,
561
- NCCLWeightTransferInitInfo,
562
- NCCLWeightTransferUpdateInfo,
643
+ from vllm.distributed.weight_transfer.sparse_nccl_engine import (
644
+ SparseNCCLWeightTransferEngine,
645
+ SparseNCCLWeightTransferUpdateInfo,
646
)
647
565
- config = WeightTransferConfig(backend="nccl")
648
+ config = WeightTransferConfig(backend="sparse_nccl")
649
+ vllm_config = MagicMock()
650
parallel_config = MagicMock(spec=ParallelConfig)
651
parallel_config.rank = 0
652
parallel_config.world_size = 1
653
parallel_config.data_parallel_rank = 0
654
parallel_config.data_parallel_index = 0
655
+ vllm_config.parallel_config = parallel_config
656
+ vllm_config.model_config = MagicMock()
657
572
- engine = NCCLWeightTransferEngine(
573
- config, parallel_config, MagicMock(spec=torch.nn.Module)
658
+ # Real module holding the target parameter the patch will modify.
659
+ model = torch.nn.Module()
660
+ model.register_parameter(
661
+ "w", torch.nn.Parameter(torch.zeros(30, device="cuda"), requires_grad=False)
662
+ )
663
+ model.get_parameter = lambda name: model.w
664
+
665
+ update_info = SparseNCCLWeightTransferUpdateInfo(
666
+ names=["w"],
667
+ dtype_names=["float32"],
668
+ shapes=[[30]],
669
+ num_updates_list=[3],
670
)
671
+
672
+ engine = SparseNCCLWeightTransferEngine(config, vllm_config, "cuda", model)
673
+ from vllm.distributed.weight_transfer.nccl_common import (
674
+ NCCLWeightTransferInitInfo,
675
+ )
676
+
677
engine.init_transfer_engine(
678
NCCLWeightTransferInitInfo(
679
master_address=master_address,
@@ -580,32 +682,18 @@ def inference_receive_sparse_tensor(
682
world_size=world_size,
683
)
684
)
583
-
584
- target = torch.zeros(30, dtype=torch.float32, device=device)
585
-
586
- def apply_sparse_patches(patches: list[SparseWeightPatch]):
587
- for patch in patches:
588
- target.index_copy_(0, patch.indices.to(torch.long), patch.values)
589
-
590
- update_info = NCCLWeightTransferUpdateInfo(
591
- names=["test.weight"],
592
- dtype_names=["float32"],
593
- shapes=[[30]],
594
- num_updates_list=[3],
595
- update_kind="sparse_flat",
596
- )
597
- engine.receive_sparse_weights(update_info, apply_sparse_patches)
685
+ engine.receive_weights(update_info)
686
torch.accelerator.synchronize()
687
688
expected = torch.zeros(30, dtype=torch.float32, device=device)
689
expected[[1, 7, 25]] = torch.tensor(
690
[10.0, 20.0, 30.0], dtype=torch.float32, device=device
691
)
604
- success = torch.equal(target, expected)
692
+ success = torch.equal(model.w.data, expected)
693
engine.shutdown()
694
return {
695
"success": success,
608
- "selected_values": target[[1, 7, 25]].cpu().tolist(),
696
+ "selected_values": model.w.data[[1, 7, 25]].cpu().tolist(),
697
}
698
699
@@ -644,11 +732,9 @@ class TestIPCWeightTransferUpdateInfoValidation:
732
"""Test IPCWeightTransferUpdateInfo dataclass validation."""
733
734
def test_valid_update_info(self):
647
- """Test creating valid IPCWeightTransferUpdateInfo."""
735
if torch.accelerator.device_count() < 1:
736
pytest.skip("Need at least 1 GPU for this test")
737
651
- # Create a dummy tensor and IPC handle
738
dummy_tensor = torch.ones(10, 10, device="cuda:0")
739
_, ipc_handle = reduce_tensor(dummy_tensor)
740
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
@@ -666,7 +752,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
752
assert len(info.ipc_handles) == 1
753
754
def test_mismatched_dtype_names_raises(self):
669
- """Test that mismatched dtype_names length raises ValueError."""
755
if torch.accelerator.device_count() < 1:
756
pytest.skip("Need at least 1 GPU for this test")
757
@@ -684,7 +769,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
769
)
770
771
def test_mismatched_shapes_raises(self):
687
- """Test that mismatched shapes length raises ValueError."""
772
if torch.accelerator.device_count() < 1:
773
pytest.skip("Need at least 1 GPU for this test")
774
@@ -702,7 +786,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
786
)
787
788
def test_mismatched_ipc_handles_raises(self):
705
- """Test that mismatched ipc_handles length raises ValueError."""
789
if torch.accelerator.device_count() < 1:
790
pytest.skip("Need at least 1 GPU for this test")
791
@@ -719,44 +802,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
802
ipc_handles=ipc_handles,
803
)
804
722
- def test_sparse_update_kind_rejected(self):
723
- """Test that IPC backend rejects sparse update metadata."""
724
- if torch.accelerator.device_count() < 1:
725
- pytest.skip("Need at least 1 GPU for this test")
726
-
727
- dummy_tensor = torch.ones(10, 10, device="cuda:0")
728
- ipc_handle = reduce_tensor(dummy_tensor)
729
- gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
730
- ipc_handles = [{gpu_uuid: ipc_handle}]
731
-
732
- with pytest.raises(NotImplementedError, match="dense updates"):
733
- IPCWeightTransferUpdateInfo(
734
- names=["layer.weight"],
735
- dtype_names=["float32"],
736
- shapes=[[10, 10]],
737
- num_updates_list=[1],
738
- ipc_handles=ipc_handles,
739
- update_kind="sparse_flat",
740
- )
741
-
742
- def test_sparse_methods_not_supported(self):
743
- """Test that IPC engine inherits sparse rejection from the base class."""
744
- config = WeightTransferConfig(backend="ipc")
745
- parallel_config = create_mock_parallel_config()
746
- engine = IPCWeightTransferEngine(
747
- config, parallel_config, MagicMock(spec=torch.nn.Module)
748
- )
749
-
750
- with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
751
- engine.receive_sparse_weights(MagicMock(), lambda _: None)
752
- with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
753
- engine.trainer_send_sparse_weights(
754
- iter([]),
755
- {"mode": "http", "url": "http://localhost:8000"},
756
- )
757
-
805
def test_valid_update_info_from_pickled(self, monkeypatch):
759
- """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
806
if torch.accelerator.device_count() < 1:
807
pytest.skip("Need at least 1 GPU for this test")
808
@@ -779,7 +825,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
825
assert info.ipc_handles_pickled is None
826
827
def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
782
- """Test that pickled handles are rejected unless env flag is enabled."""
828
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
829
830
with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
@@ -791,7 +836,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
836
)
837
838
def test_both_handles_and_pickled_raises(self):
794
- """Test that providing both ipc_handles and ipc_handles_pickled raises."""
839
if torch.accelerator.device_count() < 1:
840
pytest.skip("Need at least 1 GPU for this test")
841
@@ -812,7 +856,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
856
)
857
858
def test_neither_handles_nor_pickled_raises(self):
815
- """Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
859
with pytest.raises(ValueError, match="must be provided"):
860
IPCWeightTransferUpdateInfo(
861
names=["layer.weight"],
@@ -821,7 +864,6 @@ class TestIPCWeightTransferUpdateInfoValidation:
864
)
865
866
def test_empty_lists_valid(self):
824
- """Test that empty lists are valid."""
867
info = IPCWeightTransferUpdateInfo(
868
names=[],
869
dtype_names=[],
@@ -837,18 +879,21 @@ class TestIPCWeightTransferUpdateInfoValidation:
879
class TestIPCEngineParsing:
880
"""Test IPCWeightTransferEngine parsing methods."""
881
882
+ def _make_engine(self):
883
+ config = WeightTransferConfig(backend="ipc")
884
+ return IPCWeightTransferEngine(
885
+ config,
886
+ create_mock_vllm_config(),
887
+ "cuda",
888
+ MagicMock(spec=torch.nn.Module),
889
+ )
890
+
891
def test_parse_update_info_valid(self):
841
- """Test parsing valid update info dict."""
892
if torch.accelerator.device_count() < 1:
893
pytest.skip("Need at least 1 GPU for this test")
894
845
- config = WeightTransferConfig(backend="ipc")
846
- parallel_config = create_mock_parallel_config()
847
- engine = IPCWeightTransferEngine(
848
- config, parallel_config, MagicMock(spec=torch.nn.Module)
849
- )
895
+ engine = self._make_engine()
896
851
- # Create dummy IPC handles
897
dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
898
dummy_tensor2 = torch.ones(50, device="cuda:0")
899
_, ipc_args1 = reduce_tensor(dummy_tensor1)
@@ -872,17 +917,12 @@ class TestIPCEngineParsing:
917
assert len(update_info.ipc_handles) == 2
918
919
def test_parse_update_info_pickled(self, monkeypatch):
875
- """Test parsing update info with pickled IPC handles (HTTP path)."""
920
if torch.accelerator.device_count() < 1:
921
pytest.skip("Need at least 1 GPU for this test")
922
923
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
924
881
- config = WeightTransferConfig(backend="ipc")
882
- parallel_config = create_mock_parallel_config()
883
- engine = IPCWeightTransferEngine(
884
- config, parallel_config, MagicMock(spec=torch.nn.Module)
885
- )
925
+ engine = self._make_engine()
926
927
dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
928
dummy_tensor2 = torch.ones(50, device="cuda:0")
@@ -909,12 +949,7 @@ class TestIPCEngineParsing:
949
assert gpu_uuid in update_info.ipc_handles[1]
950
951
def test_parse_update_info_ignores_none_pickled_handles(self):
912
- """Test Ray/asdict payloads with a null pickled field use ipc_handles."""
913
- config = WeightTransferConfig(backend="ipc")
914
- parallel_config = create_mock_parallel_config()
915
- engine = IPCWeightTransferEngine(
916
- config, parallel_config, MagicMock(spec=torch.nn.Module)
917
- )
952
+ engine = self._make_engine()
953
ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
954
955
update_info = engine.parse_update_info(
@@ -931,15 +966,10 @@ class TestIPCEngineParsing:
966
assert update_info.ipc_handles == ipc_handles
967
968
def test_parse_update_info_both_handles_and_pickled_raises(self):
934
- """Test that providing both ipc_handles and ipc_handles_pickled raises."""
969
if torch.accelerator.device_count() < 1:
970
pytest.skip("Need at least 1 GPU for this test")
971
938
- config = WeightTransferConfig(backend="ipc")
939
- parallel_config = create_mock_parallel_config()
940
- engine = IPCWeightTransferEngine(
941
- config, parallel_config, MagicMock(spec=torch.nn.Module)
942
- )
972
+ engine = self._make_engine()
973
974
dummy_tensor = torch.ones(10, 10, device="cuda:0")
975
_, ipc_handle = reduce_tensor(dummy_tensor)
@@ -981,9 +1011,6 @@ class TrainerActor:
1011
self.tensor = torch.ones(tensor_shape, dtype=dtype, device=device)
1012
self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer
1013
984
- # Create IPC handle (tensor must stay alive for IPC to work)
985
- # reduce_tensor returns (rebuild_func, args); we only send args
986
- # since the receiver imports rebuild_cuda_tensor directly.
1014
_, ipc_args = reduce_tensor(self.tensor)
1015
gpu_uuid = get_physical_gpu_id(device.index)
1016
@@ -1007,6 +1034,7 @@ def inference_receive_ipc_tensor(
1034
mode: str = "ray",
1035
) -> dict:
1036
"""Inference task that receives tensor via IPCWeightTransferEngine."""
1037
+ import contextlib
1038
import os
1039
1040
# Worker-side: ipc_handles_pickled is deserialized via pickle.
@@ -1025,30 +1053,36 @@ def inference_receive_ipc_tensor(
1053
IPCWeightTransferEngine,
1054
)
1055
1028
- # Create engine with mock parallel config
1056
+ class Recorder(torch.nn.Module):
1057
+ def __init__(self):
1058
+ super().__init__()
1059
+ self.received = []
1060
+
1061
+ def load_weights(self, weights):
1062
+ for name, tensor in weights:
1063
+ self.received.append((name, tensor.clone()))
1064
+
1065
config = WeightTransferConfig(backend="ipc")
1066
+ vllm_config = MagicMock()
1067
parallel_config = MagicMock(spec=ParallelConfig)
1068
parallel_config.rank = 0
1069
parallel_config.world_size = 1
1070
parallel_config.data_parallel_rank = 0
1071
parallel_config.data_parallel_index = 0
1072
+ vllm_config.parallel_config = parallel_config
1073
+ vllm_config.model_config = MagicMock()
1074
1036
- engine = IPCWeightTransferEngine(
1037
- config, parallel_config, MagicMock(spec=torch.nn.Module)
1038
- )
1075
+ recorder = Recorder()
1076
+ engine = IPCWeightTransferEngine(config, vllm_config, "cuda", recorder)
1077
+ # Transport-only test: bypass the set_current_vllm_config context that
1078
+ # receive_weights enters, since vllm_config here is a mock.
1079
+ import vllm.config as _vllm_config_mod
1080
+
1081
+ _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
1082
1040
- # Initialize the engine (no-op for IPC)
1083
init_info = IPCWeightTransferInitInfo()
1084
engine.init_transfer_engine(init_info)
1085
1044
- # Receive weights with a no-op load_weights that captures the tensor
1045
- received_tensors = []
1046
-
1047
- def noop_load_weights(weights: list[tuple[str, torch.Tensor]]):
1048
- for name, tensor in weights:
1049
- # Clone tensor to keep it after engine cleans up
1050
- received_tensors.append((name, tensor.clone()))
1051
-
1086
ipc_handles = [{ipc_handle_dict["gpu_uuid"]: ipc_handle_dict["ipc_handle"]}]
1087
1088
if mode == "ray":
@@ -1059,7 +1093,6 @@ def inference_receive_ipc_tensor(
1093
"ipc_handles": ipc_handles,
1094
}
1095
elif mode == "http":
1062
- # Simulate HTTP transport: pickle + base64 encode handles
1096
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
1097
update_dict = {
1098
"names": ["test.weight"],
@@ -1071,19 +1104,17 @@ def inference_receive_ipc_tensor(
1104
raise ValueError(f"Unknown mode: {mode}")
1105
1106
update_info = engine.parse_update_info(update_dict)
1074
- engine.receive_weights(update_info, noop_load_weights)
1107
+ engine.receive_weights(update_info)
1108
torch.accelerator.synchronize()
1109
1077
- # Verify we received the tensor
1110
success = False
1111
received_shape = None
1112
received_sum = None
1113
1082
- if len(received_tensors) == 1:
1083
- name, tensor = received_tensors[0]
1114
+ if len(recorder.received) == 1:
1115
+ name, tensor = recorder.received[0]
1116
received_shape = list(tensor.shape)
1117
received_sum = tensor.sum().item()
1086
- # Check shape matches and values are all 42s (trainer sends 42s)
1118
if received_shape == ipc_handle_dict["shape"]:
1119
expected_sum = 42.0 * torch.tensor(ipc_handle_dict["shape"]).prod().item()
1120
if abs(received_sum - expected_sum) < 0.01:
@@ -1104,23 +1135,12 @@ def inference_receive_ipc_tensor(
1135
)
1136
@pytest.mark.parametrize("mode", ["ray", "http"])
1137
def test_ipc_weight_transfer_between_processes(mode: str):
1107
- """Test IPC weight transfer from trainer to inference process using Ray.
1108
-
1109
- Parametrized over transport modes:
1110
- - 'ray': ipc_handles passed directly.
1111
- - 'http': ipc_handles pickled + base64-encoded, deserialized in
1112
- parse_update_info before constructing the dataclass.
1113
-
1114
- IPC requires same-GPU access, so we use a placement group to co-locate
1115
- the trainer actor and inference task on the same GPU.
1116
- """
1138
+ """Test IPC weight transfer from trainer to inference process using Ray."""
1139
from ray.util.placement_group import placement_group
1140
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
1141
1142
_init_ray_for_weight_transfer()
1143
1122
- # Create a placement group to ensure both processes are on the same GPU
1123
- # Use fractional GPUs so both tasks can share the same GPU bundle
1144
pg = placement_group([{"GPU": 1, "CPU": 2}])
1145
ray.get(pg.ready())
1146
@@ -1129,20 +1149,15 @@ def test_ipc_weight_transfer_between_processes(mode: str):
1149
placement_group_capture_child_tasks=True,
1150
)
1151
1132
- # Tensor to transfer: 100x100 filled with 42s
1152
tensor_shape = [100, 100]
1153
tensor_dtype = "float32"
1154
1136
- # Create trainer actor that holds the tensor and IPC handle (stays alive)
1155
trainer_actor = TrainerActor.options( # type: ignore[attr-defined]
1156
scheduling_strategy=scheduling_strategy
1157
).remote(tensor_shape, tensor_dtype)
1158
1141
- # Get IPC handle dict (tensor stays alive in trainer actor)
1159
ipc_handle_dict = ray.get(trainer_actor.get_ipc_handle_dict.remote())
1160
1144
- # Receive tensor in inference process using IPC handles (on same GPU)
1145
- # Trainer actor stays alive during this operation
1161
inference_result = ray.get(
1162
inference_receive_ipc_tensor.options(
1163
scheduling_strategy=scheduling_strategy
@@ -1162,12 +1177,10 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
1177
pytest.skip("Need at least 1 GPU for this test")
1178
1179
config = WeightTransferConfig(backend="ipc")
1165
- parallel_config = create_mock_parallel_config()
1180
engine = IPCWeightTransferEngine(
1167
- config, parallel_config, MagicMock(spec=torch.nn.Module)
1181
+ config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
1182
)
1183
1170
- # Create IPC handle with wrong GPU UUID
1184
dummy_tensor = torch.ones(10, 10, device="cuda:0")
1185
_, ipc_handle = reduce_tensor(dummy_tensor)
1186
wrong_uuid = "wrong-uuid-12345"
@@ -1181,4 +1194,4 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
1194
)
1195
1196
with pytest.raises(ValueError, match="IPC handle not found"):
1184
- engine.receive_weights(update_info, lambda x: None)
1197
+ engine.receive_weights(update_info)
tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
+19
-81
@@ -8,7 +8,6 @@ actual NCCL communication.
8
"""
9
10
import os
11
-from collections.abc import Callable
11
from dataclasses import dataclass
12
from unittest.mock import patch
13
@@ -48,7 +47,6 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
47
names: list[str] | None = None
48
dtype_names: list[str] | None = None
49
shapes: list[list[int]] | None = None
51
- num_updates_list: list[int] | None = None
50
51
52
class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
@@ -59,16 +57,20 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
57
58
# Class-level tracking for verification across processes
59
init_transfer_engine_called: bool = False
60
+ start_called: bool = False
61
receive_weights_called: bool = False
62
+ finish_called: bool = False
63
shutdown_called: bool = False
64
last_init_info: MockInitInfo | None = None
65
last_update_info: MockUpdateInfo | None = None
66
67
- def __init__(self, config, parallel_config, model):
68
- super().__init__(config, parallel_config, model)
67
+ def __init__(self, config, vllm_config, device, model):
68
+ super().__init__(config, vllm_config, device, model)
69
# Reset tracking on init
70
MockWeightTransferEngine.init_transfer_engine_called = False
71
+ MockWeightTransferEngine.start_called = False
72
MockWeightTransferEngine.receive_weights_called = False
73
+ MockWeightTransferEngine.finish_called = False
74
MockWeightTransferEngine.shutdown_called = False
75
MockWeightTransferEngine.last_init_info = None
76
MockWeightTransferEngine.last_update_info = None
@@ -77,37 +79,28 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
79
MockWeightTransferEngine.init_transfer_engine_called = True
80
MockWeightTransferEngine.last_init_info = init_info
81
80
- def receive_weights(
81
- self,
82
- update_info: MockUpdateInfo,
83
- load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
84
- ) -> None:
85
- MockWeightTransferEngine.receive_weights_called = True
86
- MockWeightTransferEngine.last_update_info = update_info
87
- # Simulate loading weights by calling load_weights with empty list
88
- # (In real implementation, this would receive and load actual weights)
89
- load_weights([])
90
-
91
- def receive_sparse_weights(
92
- self,
93
- update_info: MockUpdateInfo,
94
- apply_patches: Callable[[list], None],
95
- ) -> None:
82
+ def start_weight_update(self) -> None:
83
+ MockWeightTransferEngine.start_called = True
84
+
85
+ def finish_weight_update(self) -> None:
86
+ MockWeightTransferEngine.finish_called = True
87
+
88
+ def receive_weights(self, update_info: MockUpdateInfo) -> None:
89
MockWeightTransferEngine.receive_weights_called = True
90
MockWeightTransferEngine.last_update_info = update_info
98
- apply_patches([])
91
92
def shutdown(self) -> None:
93
MockWeightTransferEngine.shutdown_called = True
94
103
- def trainer_send_weights(self, *args, **kwargs):
95
+ @staticmethod
96
+ def trainer_send_weights(*args, **kwargs):
97
"""Mock method to simulate trainer sending weights."""
98
pass
99
100
108
-def mock_create_engine(config, parallel_config, model):
101
+def mock_create_engine(config, vllm_config, device, model):
102
"""Mock factory function that returns our mock engine."""
110
- return MockWeightTransferEngine(config, parallel_config, model)
103
+ return MockWeightTransferEngine(config, vllm_config, device, model)
104
105
106
# --- Tests ---
@@ -208,7 +201,7 @@ def test_update_weights_calls_engine():
201
llm.init_weight_transfer_engine(
202
WeightTransferInitRequest(init_info={"test_param": "init"})
203
)
211
- llm.start_weight_update(is_checkpoint_format=True)
204
+ llm.start_weight_update()
205
206
# Call update_weights
207
test_names = ["layer.weight", "layer.bias"]
@@ -243,61 +236,6 @@ def test_update_weights_calls_engine():
236
llm.finish_weight_update()
237
238
246
-@create_new_process_for_each_test()
247
-def test_update_weights_passes_sparse_metadata():
248
- """Test sparse update metadata is forwarded unchanged to the engine."""
249
- if torch.accelerator.device_count() < 1:
250
- pytest.skip("Need at least 1 GPU for this test")
251
-
252
- os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
253
- os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
254
-
255
- with patch(
256
- "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
257
- mock_create_engine,
258
- ):
259
- llm = LLM(
260
- model=MODEL_NAME,
261
- enforce_eager=True,
262
- load_format="dummy",
263
- tensor_parallel_size=1,
264
- weight_transfer_config=WeightTransferConfig(backend="nccl"),
265
- )
266
-
267
- llm.init_weight_transfer_engine(
268
- WeightTransferInitRequest(init_info={"test_param": "init"})
269
- )
270
- llm.start_weight_update(is_checkpoint_format=False)
271
-
272
- llm.update_weights(
273
- WeightTransferUpdateRequest(
274
- update_info={
275
- "names": ["layer.weight"],
276
- "dtype_names": ["bfloat16"],
277
- "shapes": [[100]],
278
- "num_updates_list": [3],
279
- "update_kind": "sparse_flat",
280
- }
281
- )
282
- )
283
-
284
- def check_sparse_update_called(self):
285
- engine = self.weight_transfer_engine
286
- if not engine.receive_weights_called:
287
- return None
288
- info = engine.last_update_info
289
- return (
290
- info.update_kind,
291
- info.num_updates_list,
292
- )
293
-
294
- results = llm.collective_rpc(check_sparse_update_called)
295
- for result in results:
296
- assert result == ("sparse_flat", [3])
297
-
298
- llm.finish_weight_update()
299
-
300
-
239
@create_new_process_for_each_test()
240
def test_full_weight_transfer_flow():
241
"""Test the complete weight transfer flow: init -> start -> update -> finish."""
@@ -327,7 +265,7 @@ def test_full_weight_transfer_flow():
265
)
266
267
# Step 2: Start weight update
330
- llm.start_weight_update(is_checkpoint_format=True)
268
+ llm.start_weight_update()
269
270
# Step 3: Update weights
271
llm.update_weights(
tests/v1/worker/test_gpu_model_runner.py
-69
@@ -7,7 +7,6 @@ from unittest.mock import Mock
7
import numpy as np
8
import pytest
9
import torch
10
-import torch.nn as nn
10
11
import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module
12
from vllm.config import (
@@ -23,7 +22,6 @@ from vllm.distributed.parallel_state import (
22
init_distributed_environment,
23
initialize_model_parallel,
24
)
26
-from vllm.distributed.weight_transfer.base import SparseWeightPatch
25
from vllm.lora.layers import LoRAMappingType
26
from vllm.lora.request import LoRARequest
27
from vllm.model_executor.layers.attention import Attention
@@ -864,73 +862,6 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler():
862
assert torch.equal(passed_draft_probs, expected_draft_probs)
863
864
867
-def test_apply_sparse_weight_patches_updates_only_selected_entries():
868
- class DummyModel(nn.Module):
869
- def __init__(self):
870
- super().__init__()
871
- self.weight = nn.Parameter(torch.zeros(6, dtype=torch.float32))
872
-
873
- runner = object.__new__(GPUModelRunner)
874
- runner.model = DummyModel()
875
-
876
- runner.apply_sparse_weight_patches(
877
- [
878
- SparseWeightPatch(
879
- name="weight",
880
- indices=torch.tensor([1, 4], dtype=torch.int32),
881
- values=torch.tensor([3.5, -2.0], dtype=torch.float32),
882
- )
883
- ]
884
- )
885
-
886
- expected = torch.tensor([0.0, 3.5, 0.0, 0.0, -2.0, 0.0], dtype=torch.float32)
887
- assert torch.equal(runner.get_model().weight.data, expected)
888
-
889
-
890
-def test_apply_sparse_weight_patches_rejects_mismatched_lengths():
891
- class DummyModel(nn.Module):
892
- def __init__(self):
893
- super().__init__()
894
- self.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
895
-
896
- runner = object.__new__(GPUModelRunner)
897
- runner.model = DummyModel()
898
-
899
- with pytest.raises(ValueError, match="matching lengths"):
900
- runner.apply_sparse_weight_patches(
901
- [
902
- SparseWeightPatch(
903
- name="weight",
904
- indices=torch.tensor([1, 2], dtype=torch.int32),
905
- values=torch.tensor([1.0], dtype=torch.float32),
906
- )
907
- ]
908
- )
909
-
910
-
911
-def test_apply_sparse_weight_patches_rejects_non_contiguous_param():
912
- class DummyModel(nn.Module):
913
- def __init__(self):
914
- super().__init__()
915
- self.weight = nn.Parameter(
916
- torch.arange(12, dtype=torch.float32).view(3, 4).t()
917
- )
918
-
919
- runner = object.__new__(GPUModelRunner)
920
- runner.model = DummyModel()
921
-
922
- with pytest.raises(NotImplementedError, match="contiguous params"):
923
- runner.apply_sparse_weight_patches(
924
- [
925
- SparseWeightPatch(
926
- name="weight",
927
- indices=torch.tensor([1], dtype=torch.int32),
928
- values=torch.tensor([1.0], dtype=torch.float32),
929
- )
930
- ]
931
- )
932
-
933
-
865
def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):
866
torch.set_default_dtype(torch.float16)
867
layer_0 = "model.layers.0.self_attn.attn"
tests/v1/worker/test_gpu_worker_weight_transfer.py
+72
-133
@@ -1,155 +1,94 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for GPUWorker weight-transfer pass-through behavior.
4
4
-from types import SimpleNamespace
5
-from unittest.mock import MagicMock
5
+The worker no longer contains transport, layerwise, or sparse logic: it only
6
+delegates to the configured weight transfer engine and tracks whether an update
7
+session is active. These tests verify that delegation and the session guard.
8
+"""
9
10
import pytest
8
-import torch
11
10
-from vllm.config.parallel import ParallelConfig
11
-from vllm.config.weight_transfer import WeightTransferConfig
12
-from vllm.distributed.weight_transfer.base import SparseWeightPatch
13
-from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine
12
from vllm.v1.worker.gpu_worker import Worker
13
14
17
-def _make_nccl_engine() -> NCCLWeightTransferEngine:
18
- parallel_config = MagicMock(spec=ParallelConfig)
19
- parallel_config.rank = 0
20
- parallel_config.world_size = 1
21
- parallel_config.data_parallel_rank = 0
22
- parallel_config.data_parallel_index = 0
23
- return NCCLWeightTransferEngine(
24
- WeightTransferConfig(backend="nccl"),
25
- parallel_config,
26
- MagicMock(spec=torch.nn.Module),
27
- )
15
+class _RecordingEngine:
16
+ """Minimal stand-in for a weight transfer engine."""
17
18
+ def __init__(self, raise_on_update: bool = False):
19
+ self.raise_on_update = raise_on_update
20
+ self.started = False
21
+ self.finished = False
22
+ self.update_calls: list[dict] = []
23
30
-def test_update_weights_sparse_dispatches_to_sparse_receive(monkeypatch):
31
- monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
24
+ def start_weight_update(self) -> None:
25
+ self.started = True
26
+
27
+ def update_weights(self, update_info: dict) -> None:
28
+ self.update_calls.append(update_info)
29
+ if self.raise_on_update:
30
+ raise ValueError("boom")
31
+
32
+ def finish_weight_update(self) -> None:
33
+ self.finished = True
34
33
- worker = object.__new__(Worker)
34
- worker.device = "cpu"
35
- worker.parallel_config = SimpleNamespace(world_size=1)
36
- worker.weight_transfer_engine = _make_nccl_engine()
37
- worker._weight_update_active = True
38
- worker._is_checkpoint_format = False
39
-
40
- applied_patches = []
41
-
42
- def apply_sparse_weight_patches(patches):
43
- applied_patches.extend(patches)
44
-
45
- worker.model_runner = SimpleNamespace(
46
- apply_sparse_weight_patches=apply_sparse_weight_patches,
47
- )
48
-
49
- received_kinds = []
50
-
51
- def receive_sparse_weights(update_info, apply_patches):
52
- received_kinds.append(update_info.update_kind)
53
- apply_patches(
54
- [
55
- SparseWeightPatch(
56
- name="layer.weight",
57
- indices=torch.tensor([1], dtype=torch.int32),
58
- values=torch.tensor([2.0], dtype=torch.float32),
59
- )
60
- ]
61
- )
62
-
63
- worker.weight_transfer_engine.receive_sparse_weights = receive_sparse_weights
64
-
65
- Worker.update_weights(
66
- worker,
67
- {
68
- "names": ["layer.weight"],
69
- "dtype_names": ["float32"],
70
- "shapes": [[4]],
71
- "num_updates_list": [1],
72
- "update_kind": "sparse_flat",
73
- },
74
- )
75
-
76
- assert received_kinds == ["sparse_flat"]
77
- assert len(applied_patches) == 1
78
- assert torch.equal(applied_patches[0].indices, torch.tensor([1], dtype=torch.int32))
79
-
80
-
81
-def test_update_weights_sparse_rejects_tp_or_pp(monkeypatch):
82
- monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
35
36
+def _make_worker(engine: _RecordingEngine | None) -> Worker:
37
worker = object.__new__(Worker)
85
- worker.device = "cpu"
86
- worker.parallel_config = SimpleNamespace(world_size=2)
87
- worker.weight_transfer_engine = _make_nccl_engine()
88
- worker._weight_update_active = True
89
- worker._is_checkpoint_format = False
90
- worker.model_runner = SimpleNamespace(apply_sparse_weight_patches=lambda _: None)
91
-
92
- with pytest.raises(NotImplementedError, match="TP=1 and PP=1"):
93
- Worker.update_weights(
94
- worker,
95
- {
96
- "names": ["layer.weight"],
97
- "dtype_names": ["float32"],
98
- "shapes": [[4]],
99
- "num_updates_list": [1],
100
- "update_kind": "sparse_flat",
101
- },
102
- )
103
- assert worker._weight_update_active is False
104
- assert worker._is_checkpoint_format is True
38
+ worker.weight_transfer_engine = engine
39
+ worker._weight_update_active = False
40
+ return worker
41
42
107
-def test_update_weights_sparse_rejects_checkpoint_format(monkeypatch):
108
- monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
43
+def test_start_update_finish_delegates_to_engine():
44
+ engine = _RecordingEngine()
45
+ worker = _make_worker(engine)
46
110
- worker = object.__new__(Worker)
111
- worker.device = "cpu"
112
- worker.parallel_config = SimpleNamespace(world_size=1)
113
- worker.weight_transfer_engine = _make_nccl_engine()
114
- worker._weight_update_active = True
115
- worker._is_checkpoint_format = True
116
- worker.model_runner = SimpleNamespace(model=MagicMock())
117
-
118
- with pytest.raises(ValueError, match="start_weight_update"):
119
- Worker.update_weights(
120
- worker,
121
- {
122
- "names": ["layer.weight"],
123
- "dtype_names": ["float32"],
124
- "shapes": [[4]],
125
- "num_updates_list": [1],
126
- "update_kind": "sparse_flat",
127
- },
128
- )
47
+ Worker.start_weight_update(worker)
48
+ assert engine.started is True
49
+ assert worker._weight_update_active is True
50
+
51
+ Worker.update_weights(worker, {"names": ["w"]})
52
+ assert engine.update_calls == [{"names": ["w"]}]
53
+ assert worker._weight_update_active is True
54
+
55
+ Worker.finish_weight_update(worker)
56
+ assert engine.finished is True
57
assert worker._weight_update_active is False
130
- assert worker._is_checkpoint_format is True
58
59
133
-def test_update_weights_resets_state_when_update_info_is_invalid(monkeypatch):
134
- monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
60
+def test_double_start_raises():
61
+ worker = _make_worker(_RecordingEngine())
62
+ Worker.start_weight_update(worker)
63
+ with pytest.raises(RuntimeError, match="already"):
64
+ Worker.start_weight_update(worker)
65
136
- worker = object.__new__(Worker)
137
- worker.device = "cpu"
138
- worker.parallel_config = SimpleNamespace(world_size=1)
139
- worker.weight_transfer_engine = _make_nccl_engine()
140
- worker._weight_update_active = True
141
- worker._is_checkpoint_format = False
142
-
143
- with pytest.raises(ValueError, match="cannot be empty"):
144
- Worker.update_weights(
145
- worker,
146
- {
147
- "names": [],
148
- "dtype_names": [],
149
- "shapes": [],
150
- "num_updates_list": [],
151
- "update_kind": "sparse_flat",
152
- },
153
- )
66
+
67
+def test_update_without_start_raises():
68
+ worker = _make_worker(_RecordingEngine())
69
+ with pytest.raises(RuntimeError, match="start_weight_update must be called"):
70
+ Worker.update_weights(worker, {"names": ["w"]})
71
+
72
+
73
+def test_finish_without_start_raises():
74
+ worker = _make_worker(_RecordingEngine())
75
+ with pytest.raises(RuntimeError, match="without a matching"):
76
+ Worker.finish_weight_update(worker)
77
+
78
+
79
+def test_update_resets_active_on_error():
80
+ engine = _RecordingEngine(raise_on_update=True)
81
+ worker = _make_worker(engine)
82
+ Worker.start_weight_update(worker)
83
+
84
+ with pytest.raises(ValueError, match="boom"):
85
+ Worker.update_weights(worker, {"names": ["w"]})
86
+
87
+ # A failed update ends the session so the next start is clean.
88
assert worker._weight_update_active is False
155
- assert worker._is_checkpoint_format is True
89
+
90
+
91
+def test_missing_engine_raises():
92
+ worker = _make_worker(None)
93
+ with pytest.raises(RuntimeError, match="Weight transfer not configured"):
94
+ Worker.start_weight_update(worker)
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"] | str = "nccl"
12
+ backend: Literal["nccl", "ipc", "sparse_nccl"] | 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
+60
-74
@@ -3,12 +3,15 @@
3
"""Base class for weight transfer engines."""
4
5
from abc import ABC, abstractmethod
6
-from collections.abc import Callable, Iterator
7
-from dataclasses import KW_ONLY, dataclass, field
8
-from typing import Any, Generic, Literal, TypeVar
6
+from collections.abc import Iterator
7
+from dataclasses import dataclass, field
8
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
9
10
import torch
11
12
+if TYPE_CHECKING:
13
+ from vllm.config import VllmConfig
14
+
15
from vllm.config.parallel import ParallelConfig
16
from vllm.config.weight_transfer import WeightTransferConfig
17
@@ -28,44 +31,7 @@ class WeightTransferInitInfo(ABC): # noqa: B024
31
class WeightTransferUpdateInfo(ABC): # noqa: B024
32
"""Base class for backend-specific weight update info."""
33
31
- _: KW_ONLY
32
- update_kind: Literal["dense", "sparse_flat"] = "dense"
33
- """Weight update format."""
34
- num_updates_list: list[int] | None = None
35
- """Number of sparse entries to receive for each parameter in ``names``."""
36
-
37
- def __post_init__(self) -> None:
38
- if self.update_kind not in ("dense", "sparse_flat"):
39
- raise ValueError(f"Unsupported update_kind: {self.update_kind}")
40
- if self.update_kind == "dense":
41
- if self.num_updates_list is not None:
42
- raise ValueError(
43
- "Sparse metadata is only supported for `update_kind='sparse_flat'`"
44
- )
45
- return
46
-
47
- if self.num_updates_list is None:
48
- raise ValueError("`num_updates_list` is required for sparse updates")
49
- if len(self.num_updates_list) == 0:
50
- raise ValueError("`num_updates_list` cannot be empty for sparse updates")
51
- if any(num_updates < 0 for num_updates in self.num_updates_list):
52
- raise ValueError("Sparse `num_updates_list` entries must be non-negative")
53
-
54
- names = getattr(self, "names", None)
55
- if names is not None and len(self.num_updates_list) != len(names):
56
- raise ValueError(
57
- f"`num_updates_list` should be of the same size as `names`: "
58
- f"got {len(self.num_updates_list)} and {len(names)}"
59
- )
60
-
61
-
62
-@dataclass
63
-class SparseWeightPatch:
64
- """A sparse in-place patch for one existing parameter."""
65
-
66
- name: str
67
- indices: torch.Tensor
68
- values: torch.Tensor
34
+ pass
35
36
37
# API-level request classes (accept dicts for backend-agnostic serialization)
@@ -89,9 +55,15 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
55
from a trainer to inference workers.
56
57
This abstraction separates weight transfer transport logic from the worker
92
- implementation, allowing different backends (NCCL, CUDA IPC[TODO], RDMA[TODO]) to be
58
+ implementation, allowing different backends (NCCL, CUDA IPC, RDMA[TODO]) to be
59
plugged in.
60
61
+ Each engine owns its full weight-update lifecycle: `start_weight_update`,
62
+ `update_weights`, and `finish_weight_update`. Layerwise reloading (used by
63
+ checkpoint-format engines) is opted into per engine by running it inside
64
+ `start_weight_update`/`finish_weight_update`. Engines that apply weights in
65
+ place (e.g. sparse patches) leave those methods as no-ops.
66
+
67
Subclasses should define:
68
init_info_cls: Type of backend-specific initialization info
69
update_info_cls: Type of backend-specific update info
@@ -104,7 +76,8 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
76
def __init__(
77
self,
78
config: WeightTransferConfig,
107
- parallel_config: ParallelConfig,
79
+ vllm_config: "VllmConfig",
80
+ device: torch.device,
81
model: torch.nn.Module,
82
) -> None:
83
"""
@@ -112,11 +85,15 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
85
86
Args:
87
config: The configuration for the weight transfer engine
115
- parallel_config: The configuration for the parallel setup
88
+ vllm_config: The full vLLM config (provides parallel/model config)
89
+ device: The device this worker's model lives on
90
model: The local model instance which will receive the weights
91
"""
92
self.config = config
119
- self.parallel_config = parallel_config
93
+ self.vllm_config = vllm_config
94
+ self.parallel_config: ParallelConfig = vllm_config.parallel_config
95
+ self.model_config = vllm_config.model_config
96
+ self.device = device
97
self.model = model
98
99
def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
@@ -171,32 +148,50 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
148
raise NotImplementedError
149
150
@abstractmethod
174
- def receive_weights(
175
- self,
176
- update_info: TUpdateInfo,
177
- load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
178
- ) -> None:
151
+ def start_weight_update(self) -> None:
152
+ """
153
+ Prepare the engine for a new weight update.
154
+
155
+ Engines that receive weights in checkpoint format initialize layerwise reloading
156
+ here, else this is typically a no-op.
157
+ See: https://docs.vllm.ai/en/latest/training/layerwise/ for more details.
158
+ """
159
+ raise NotImplementedError
160
+
161
+ @abstractmethod
162
+ def finish_weight_update(self) -> None:
163
"""
180
- Receive weights from the trainer and load them incrementally.
164
+ Finalize the current weight update.
165
+
166
+ Checkpoint-format engines finalize layerwise reloading here; engines
167
+ that apply weights in place leave this as a no-op.
168
+ """
169
+ raise NotImplementedError
170
+
171
+ def update_weights(self, update_info: dict[str, Any]) -> None:
172
+ """
173
+ Receive one weight update chunk and load it into the model.
174
+
175
+ Args:
176
+ update_info: Dictionary containing backend-specific update info
177
+ """
178
+ typed_update_info = self.parse_update_info(update_info)
179
+ self.receive_weights(typed_update_info)
180
+ # NCCL broadcast / IPC paths may be asynchronous. Synchronize here so the
181
+ # next step uses the new weights.
182
+ torch.accelerator.synchronize()
183
+
184
+ @abstractmethod
185
+ def receive_weights(self, update_info: TUpdateInfo) -> None:
186
+ """
187
+ Receive weights from the trainer and load them into the model.
188
189
Args:
190
update_info: Backend-specific update info containing parameter metadata
191
and any backend-specific data
185
- load_weights: Callable that loads weights into the model. Called
186
- incrementally for each weight to avoid OOM.
192
"""
193
raise NotImplementedError
194
190
- def receive_sparse_weights(
191
- self,
192
- update_info: TUpdateInfo,
193
- apply_patches: Callable[[list[SparseWeightPatch]], None],
194
- ) -> None:
195
- """Receive sparse weight patches from the trainer."""
196
- raise NotImplementedError(
197
- f"{self.__class__.__name__} does not support sparse weight updates"
198
- )
199
-
195
@abstractmethod
196
def shutdown(self) -> None:
197
"""
@@ -208,7 +203,7 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
203
@staticmethod
204
@abstractmethod
205
def trainer_send_weights(
211
- iterator: Iterator[tuple[str, torch.Tensor]],
206
+ iterator: Iterator[Any],
207
trainer_args: dict[str, Any] | Any,
208
) -> None:
209
"""
@@ -218,8 +213,7 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
213
to send weights to all inference workers.
214
215
Args:
221
- iterator: Iterator of model parameters. Returns (name, tensor) tuples.
222
- The tensors should be on the appropriate device for the backend.
216
+ iterator: Iterator of backend-specific items to send.
217
trainer_args: Dictionary containing backend-specific arguments needed
218
to send weights. The structure depends on the backend:
219
- NCCL: Contains 'group', 'src', 'packed', etc.
@@ -231,11 +225,3 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
225
>>> engine.trainer_send_weights(param_iter, trainer_args)
226
"""
227
raise NotImplementedError
234
-
235
- @staticmethod
236
- def trainer_send_sparse_weights(
237
- _iterator: Iterator[SparseWeightPatch],
238
- _trainer_args: dict[str, Any] | Any,
239
- ) -> None:
240
- """Send sparse weight patches from trainer to inference workers."""
241
- raise NotImplementedError("Sparse weight updates are not supported")
vllm/distributed/weight_transfer/factory.py
+12
-4
@@ -12,7 +12,7 @@ from vllm.logger import init_logger
12
if TYPE_CHECKING:
13
import torch
14
15
- from vllm.config.parallel import ParallelConfig
15
+ from vllm.config import VllmConfig
16
from vllm.config.weight_transfer import WeightTransferConfig
17
18
logger = init_logger(__name__)
@@ -76,14 +76,16 @@ class WeightTransferEngineFactory:
76
def create_engine(
77
cls,
78
config: "WeightTransferConfig",
79
- parallel_config: "ParallelConfig",
79
+ vllm_config: "VllmConfig",
80
+ device: "torch.device",
81
model: "torch.nn.Module",
82
) -> WeightTransferEngine:
83
"""Create a weight transfer engine instance.
84
85
Args:
86
config: Weight transfer configuration containing the backend name
86
- parallel_config: Parallel configuration for the engine
87
+ vllm_config: The full vLLM config (provides parallel/model config)
88
+ device: The device this worker's model lives on
89
model: The local model instance which will receive the weights
90
91
Returns:
@@ -106,7 +108,7 @@ class WeightTransferEngineFactory:
108
engine_cls.__name__,
109
)
110
109
- return engine_cls(config, parallel_config, model)
111
+ return engine_cls(config, vllm_config, device, model)
112
113
114
# Register built-in weight transfer engines here.
@@ -124,3 +126,9 @@ WeightTransferEngineFactory.register_engine(
126
"vllm.distributed.weight_transfer.ipc_engine",
127
"IPCWeightTransferEngine",
128
)
129
+
130
+WeightTransferEngineFactory.register_engine(
131
+ "sparse_nccl",
132
+ "vllm.distributed.weight_transfer.sparse_nccl_engine",
133
+ "SparseNCCLWeightTransferEngine",
134
+)
vllm/distributed/weight_transfer/ipc_engine.py
+34
-20
@@ -5,7 +5,7 @@
5
import pickle
6
from collections.abc import Callable, Iterator
7
from dataclasses import asdict, dataclass
8
-from typing import Any
8
+from typing import TYPE_CHECKING, Any
9
10
import pybase64 as base64
11
import ray
@@ -14,13 +14,15 @@ import torch
14
from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
15
16
from vllm import envs
17
-from vllm.config.parallel import ParallelConfig
17
from vllm.config.weight_transfer import WeightTransferConfig
18
from vllm.distributed.weight_transfer.base import (
19
WeightTransferEngine,
20
WeightTransferInitInfo,
21
WeightTransferUpdateInfo,
22
)
23
+
24
+if TYPE_CHECKING:
25
+ from vllm.config import VllmConfig
26
from vllm.distributed.weight_transfer.packed_tensor import (
27
DEFAULT_PACKED_BUFFER_SIZE_BYTES,
28
packed_ipc_consumer,
@@ -87,10 +89,6 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
89
"""Whether this update uses packed tensor format."""
90
91
def __post_init__(self):
90
- super().__post_init__()
91
- if self.update_kind != "dense":
92
- raise NotImplementedError("IPC weight transfer only supports dense updates")
93
-
92
if self.ipc_handles_pickled is not None:
93
if self.ipc_handles is not None:
94
raise ValueError(
@@ -152,7 +150,8 @@ class IPCWeightTransferEngine(
150
def __init__(
151
self,
152
config: WeightTransferConfig,
155
- parallel_config: ParallelConfig,
153
+ vllm_config: "VllmConfig",
154
+ device: torch.device,
155
model: torch.nn.Module,
156
) -> None:
157
"""
@@ -160,10 +159,11 @@ class IPCWeightTransferEngine(
159
160
Args:
161
config: The configuration for the weight transfer engine
163
- parallel_config: The configuration for the parallel setup
162
+ vllm_config: The full vLLM config
163
+ device: The device this worker's model lives on
164
model: The local model instance which will receive the weights
165
"""
166
- super().__init__(config, parallel_config, model)
166
+ super().__init__(config, vllm_config, device, model)
167
168
def parse_update_info(
169
self, update_dict: dict[str, Any]
@@ -206,22 +206,37 @@ class IPCWeightTransferEngine(
206
"""
207
pass
208
209
- def receive_weights(
210
- self,
211
- update_info: IPCWeightTransferUpdateInfo,
212
- load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
213
- ) -> None:
209
+ def start_weight_update(self) -> None:
210
+ """Initialize layerwise reloading for the incoming checkpoint weights."""
211
+ from vllm.model_executor.model_loader.reload import (
212
+ initialize_layerwise_reload,
213
+ )
214
+
215
+ initialize_layerwise_reload(self.model)
216
+
217
+ def finish_weight_update(self) -> None:
218
+ """Finalize layerwise reloading after all weights have been received."""
219
+ from vllm.model_executor.model_loader.reload import (
220
+ finalize_layerwise_reload,
221
+ )
222
+
223
+ finalize_layerwise_reload(self.model, self.model_config)
224
+
225
+ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None:
226
"""
215
- Receive weights from the trainer via CUDA IPC handles.
227
+ Receive weights from the trainer via CUDA IPC handles and load them.
228
229
Args:
230
update_info: IPC update info containing parameter names, dtypes, shapes,
231
and IPC handles. Each IPC handle is a mapping between physical
232
GPU UUID and the rebuild_cuda_tensor args tuple.
221
- load_weights: Callable that loads weights into the model. Called
222
- incrementally for each weight to avoid OOM.
233
"""
224
- device_index = torch.accelerator.current_device_index()
234
+ # Use the worker's assigned device rather than the ambient current
235
+ # device: the receive path is no longer wrapped in
236
+ # `with torch.device(self.device)` by the caller, so the current device
237
+ # is not guaranteed to match self.device. The IPC tensors must be
238
+ # rebuilt on the device the model lives on.
239
+ device_index = self.device.index
240
241
if update_info.packed:
242
assert update_info.tensor_sizes is not None
@@ -234,7 +249,6 @@ class IPCWeightTransferEngine(
249
tensor_sizes=update_info.tensor_sizes,
250
device_index=device_index,
251
)
237
- load_weights(weights)
252
else:
253
assert isinstance(update_info.ipc_handles, list)
254
weights = []
@@ -260,7 +274,7 @@ class IPCWeightTransferEngine(
274
weight = rebuild_cuda_tensor(*list_args)
275
weights.append((name, weight))
276
263
- load_weights(weights)
277
+ self.model.load_weights(weights)
278
279
def shutdown(self) -> None:
280
pass
vllm/distributed/weight_transfer/nccl_common.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared NCCL initialization helpers for weight transfer engines.
4
+
5
+The dense (`NCCLWeightTransferEngine`) and sparse
6
+(`SparseNCCLWeightTransferEngine`) backends are independent engines that share
7
+*only* their process-group initialization. That common logic lives here so the
8
+sparse engine does not have to subclass the dense one.
9
+"""
10
+
11
+from dataclasses import dataclass
12
+from typing import TYPE_CHECKING
13
+
14
+import torch
15
+
16
+if TYPE_CHECKING:
17
+ from vllm.config.parallel import ParallelConfig
18
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
19
+
20
+from vllm.distributed.weight_transfer.base import WeightTransferInitInfo
21
+
22
+
23
+@dataclass
24
+class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
25
+ """Initialization info for NCCL-based weight transfer backends."""
26
+
27
+ master_address: str
28
+ master_port: int
29
+ rank_offset: int
30
+ world_size: int
31
+
32
+
33
+def stateless_init_process_group(
34
+ master_address: str,
35
+ master_port: int,
36
+ rank: int,
37
+ world_size: int,
38
+ device,
39
+) -> "PyNcclCommunicator":
40
+ """
41
+ vLLM provides `StatelessProcessGroup` to create a process group
42
+ without considering the global process group in torch.distributed.
43
+ It is recommended to create `StatelessProcessGroup`, and then initialize
44
+ the data-plane communication (NCCL) between external (train processes)
45
+ and vLLM workers.
46
+ """
47
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
48
+ from vllm.distributed.utils import StatelessProcessGroup
49
+
50
+ pg = StatelessProcessGroup.create(
51
+ host=master_address, port=master_port, rank=rank, world_size=world_size
52
+ )
53
+ return PyNcclCommunicator(pg, device=device)
54
+
55
+
56
+def worker_init_process_group(
57
+ init_info: NCCLWeightTransferInitInfo,
58
+ parallel_config: "ParallelConfig",
59
+) -> "PyNcclCommunicator":
60
+ """Create the trainer<->worker NCCL group on an inference worker.
61
+
62
+ Computes a unique rank for this worker across all data-parallel groups and
63
+ joins the stateless process group with the trainer.
64
+ """
65
+ # Calculate the global rank in the trainer-worker process group.
66
+ # Must account for data parallel to get unique ranks across all workers.
67
+ dp_rank = parallel_config.data_parallel_index
68
+ world_size_per_dp = parallel_config.world_size # TP * PP
69
+ rank_within_dp = parallel_config.rank
70
+
71
+ # Unique rank across all DP groups
72
+ worker_rank = dp_rank * world_size_per_dp + rank_within_dp
73
+ rank = worker_rank + init_info.rank_offset
74
+
75
+ device = torch.accelerator.current_device_index()
76
+ return stateless_init_process_group(
77
+ init_info.master_address,
78
+ init_info.master_port,
79
+ rank,
80
+ init_info.world_size,
81
+ device=device,
82
+ )
83
+
84
+
85
+def trainer_init(
86
+ init_info: NCCLWeightTransferInitInfo | dict,
87
+) -> "PyNcclCommunicator":
88
+ """
89
+ Initialize NCCL process group for trainer-side weight transfer.
90
+
91
+ The trainer is always rank 0 in the process group. Uses the current
92
+ CUDA device (torch.accelerator.current_device_index()).
93
+
94
+ Args:
95
+ init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys:
96
+ - master_address: str
97
+ - master_port: int
98
+ - world_size: int
99
+
100
+ Returns:
101
+ PyNcclCommunicator for weight transfer.
102
+ """
103
+ if isinstance(init_info, dict):
104
+ master_address = init_info["master_address"]
105
+ master_port = init_info["master_port"]
106
+ world_size = init_info["world_size"]
107
+ else:
108
+ master_address = init_info.master_address
109
+ master_port = init_info.master_port
110
+ world_size = init_info.world_size
111
+
112
+ # Trainer is always rank 0
113
+ device = torch.accelerator.current_device_index()
114
+ return stateless_init_process_group(
115
+ master_address,
116
+ master_port,
117
+ 0,
118
+ world_size,
119
+ device,
120
+ )
vllm/distributed/weight_transfer/nccl_engine.py
+48
-196
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""NCCL-based weight transfer engine."""
3
+"""NCCL-based (dense) weight transfer engine."""
4
5
from collections.abc import Callable, Iterator
6
from dataclasses import dataclass
@@ -9,31 +9,32 @@ from typing import TYPE_CHECKING, Any
9
import torch
10
11
if TYPE_CHECKING:
12
+ from vllm.config import VllmConfig
13
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
14
14
-from vllm.config.parallel import ParallelConfig
15
from vllm.config.weight_transfer import WeightTransferConfig
16
from vllm.distributed.weight_transfer.base import (
17
- SparseWeightPatch,
17
WeightTransferEngine,
19
- WeightTransferInitInfo,
18
WeightTransferUpdateInfo,
19
)
20
+from vllm.distributed.weight_transfer.nccl_common import (
21
+ NCCLWeightTransferInitInfo,
22
+ trainer_init,
23
+ worker_init_process_group,
24
+)
25
from vllm.distributed.weight_transfer.packed_tensor import (
26
DEFAULT_PACKED_BUFFER_SIZE_BYTES,
27
DEFAULT_PACKED_NUM_BUFFERS,
28
packed_nccl_broadcast_consumer,
29
)
30
28
-
29
-@dataclass
30
-class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
31
- """Initialization info for NCCL weight transfer backend."""
32
-
33
- master_address: str
34
- master_port: int
35
- rank_offset: int
36
- world_size: int
31
+# Re-exported for backward compatibility; canonical home is nccl_common.
32
+__all__ = [
33
+ "NCCLWeightTransferInitInfo",
34
+ "NCCLTrainerSendWeightsArgs",
35
+ "NCCLWeightTransferUpdateInfo",
36
+ "NCCLWeightTransferEngine",
37
+]
38
39
40
@dataclass
@@ -82,7 +83,6 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
83
84
def __post_init__(self):
85
"""Validate that all lists have the same length."""
85
- super().__post_init__()
86
num_params = len(self.names)
87
if len(self.dtype_names) != num_params:
88
raise ValueError(
@@ -94,13 +94,6 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
94
f"`shapes` should be of the same size as `names`: "
95
f"got {len(self.shapes)} and {len(self.names)}"
96
)
97
- if self.update_kind == "dense":
98
- return
99
-
100
- if self.packed:
101
- raise ValueError(
102
- "`update_kind='sparse_flat'` cannot be combined with `packed=True`"
103
- )
97
98
99
class NCCLWeightTransferEngine(
@@ -109,8 +102,10 @@ class NCCLWeightTransferEngine(
102
"""
103
Weight transfer engine using NCCL for communication between trainer and workers.
104
112
- This implementation uses NCCL broadcast operations to transfer weights from
113
- the trainer (rank 0) to all inference workers in a process group.
105
+ This implementation uses NCCL broadcast operations to transfer dense
106
+ checkpoint-format weights from the trainer (rank 0) to all inference workers
107
+ in a process group. Received weights are loaded via the model's
108
+ `load_weights` using the layerwise reload lifecycle.
109
"""
110
111
# Define backend-specific dataclass types
@@ -120,18 +115,11 @@ class NCCLWeightTransferEngine(
115
def __init__(
116
self,
117
config: WeightTransferConfig,
123
- parallel_config: ParallelConfig,
118
+ vllm_config: "VllmConfig",
119
+ device: torch.device,
120
model: torch.nn.Module,
121
) -> None:
126
- """
127
- Initialize the NCCL weight transfer engine.
128
-
129
- Args:
130
- config: The configuration for the weight transfer engine
131
- parallel_config: The configuration for the parallel setup
132
- model: The local model instance which will receive the weights
133
- """
134
- super().__init__(config, parallel_config, model)
122
+ super().__init__(config, vllm_config, device, model)
123
self.model_update_group: PyNcclCommunicator | None = None
124
125
def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
@@ -142,35 +130,29 @@ class NCCLWeightTransferEngine(
130
init_info: NCCL initialization info containing master address, port,
131
rank offset, and world size
132
"""
133
+ self.model_update_group = worker_init_process_group(
134
+ init_info, self.parallel_config
135
+ )
136
146
- # Calculate the global rank in the trainer-worker process group
147
- # Must account for data parallel to get unique ranks across all workers
148
- dp_rank = self.parallel_config.data_parallel_index
149
- world_size_per_dp = self.parallel_config.world_size # TP * PP
150
- rank_within_dp = self.parallel_config.rank
151
-
152
- # Unique rank across all DP groups
153
- worker_rank = dp_rank * world_size_per_dp + rank_within_dp
154
- rank = worker_rank + init_info.rank_offset
155
- # Create stateless process group
156
- device = torch.accelerator.current_device_index()
157
- self.model_update_group = (
158
- NCCLWeightTransferEngine._stateless_init_process_group(
159
- init_info.master_address,
160
- init_info.master_port,
161
- rank,
162
- init_info.world_size,
163
- device=device,
164
- )
137
+ def start_weight_update(self) -> None:
138
+ """Initialize layerwise reloading for the incoming checkpoint weights."""
139
+ from vllm.model_executor.model_loader.reload import (
140
+ initialize_layerwise_reload,
141
)
142
167
- def receive_weights(
168
- self,
169
- update_info: NCCLWeightTransferUpdateInfo,
170
- load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
171
- ) -> None:
143
+ initialize_layerwise_reload(self.model)
144
+
145
+ def finish_weight_update(self) -> None:
146
+ """Finalize layerwise reloading after all weights have been received."""
147
+ from vllm.model_executor.model_loader.reload import (
148
+ finalize_layerwise_reload,
149
+ )
150
+
151
+ finalize_layerwise_reload(self.model, self.model_config)
152
+
153
+ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None:
154
"""
173
- Receive weights from trainer via NCCL broadcast and load them incrementally.
155
+ Receive weights from trainer via NCCL broadcast.
156
157
If update_info.packed is True, uses packed tensor broadcasting for
158
efficient transfer of multiple weights in batches. Otherwise, uses simple
@@ -179,19 +161,12 @@ class NCCLWeightTransferEngine(
161
Args:
162
update_info: NCCL update info containing parameter names, dtypes, shapes,
163
and packed flag
182
- load_weights: Callable that loads weights into the model. Called
183
- incrementally for each batch of weights to avoid OOM.
164
"""
165
if self.model_update_group is None:
166
raise RuntimeError(
167
"NCCL weight transfer not initialized. "
168
"Call init_transfer_engine() first."
169
)
190
- if update_info.update_kind != "dense":
191
- raise ValueError(
192
- "Sparse updates must use `receive_sparse_weights`, not "
193
- "`receive_weights`"
194
- )
170
171
if update_info.packed:
172
# Build iterator of (name, (shape, dtype)) from update_info
@@ -206,9 +181,10 @@ class NCCLWeightTransferEngine(
181
iterator=state_dict_info_iterator(),
182
group=self.model_update_group,
183
src=0,
209
- post_unpack_func=load_weights,
184
+ post_unpack_func=self.model.load_weights,
185
buffer_size_bytes=update_info.packed_buffer_size_bytes,
186
num_buffers=update_info.packed_num_buffers,
187
+ device=self.device,
188
)
189
else:
190
# Use simple one-by-one broadcasting
@@ -216,49 +192,13 @@ class NCCLWeightTransferEngine(
192
update_info.names, update_info.dtype_names, update_info.shapes
193
):
194
dtype = getattr(torch, dtype_name)
219
- weight = torch.empty(shape, dtype=dtype, device="cuda")
195
+ weight = torch.empty(shape, dtype=dtype, device=self.device)
196
self.model_update_group.broadcast(
197
weight, src=0, stream=torch.cuda.current_stream()
198
)
223
- load_weights([(name, weight)])
199
+ self.model.load_weights([(name, weight)])
200
del weight
201
226
- def receive_sparse_weights(
227
- self,
228
- update_info: NCCLWeightTransferUpdateInfo,
229
- apply_patches: Callable[[list[SparseWeightPatch]], None],
230
- ) -> None:
231
- """Receive sparse flat-index patches from trainer via NCCL."""
232
- if self.model_update_group is None:
233
- raise RuntimeError(
234
- "NCCL weight transfer not initialized. "
235
- "Call init_transfer_engine() first."
236
- )
237
- if update_info.update_kind != "sparse_flat":
238
- raise ValueError("Sparse receive path requires `update_kind='sparse_flat'`")
239
- assert update_info.num_updates_list is not None
240
-
241
- for name, dtype_name, num_updates in zip(
242
- update_info.names,
243
- update_info.dtype_names,
244
- update_info.num_updates_list,
245
- ):
246
- dtype = getattr(torch, dtype_name)
247
- device = torch.accelerator.current_device_index()
248
- indices = torch.empty(num_updates, dtype=torch.int32, device=device)
249
- values = torch.empty(num_updates, dtype=dtype, device=device)
250
- self.model_update_group.broadcast(
251
- indices, src=0, stream=torch.cuda.current_stream()
252
- )
253
- self.model_update_group.broadcast(
254
- values, src=0, stream=torch.cuda.current_stream()
255
- )
256
- apply_patches(
257
- [SparseWeightPatch(name=name, indices=indices, values=values)]
258
- )
259
- del indices
260
- del values
261
-
202
def shutdown(self) -> None:
203
if self.model_update_group is not None:
204
# Clean up the communicator by removing the reference
@@ -269,7 +209,7 @@ class NCCLWeightTransferEngine(
209
iterator: Iterator[tuple[str, torch.Tensor]],
210
trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
211
) -> None:
272
- """Broadcast weights from trainer to vLLM workers.
212
+ """Broadcast dense weights from trainer to vLLM workers.
213
214
Args:
215
iterator: Iterator of model parameters. Returns (name, tensor) tuples
@@ -322,94 +262,6 @@ class NCCLWeightTransferEngine(
262
stream=args.stream or torch.cuda.current_stream(),
263
)
264
325
- @staticmethod
326
- def trainer_send_sparse_weights(
327
- iterator: Iterator[SparseWeightPatch],
328
- trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
329
- ) -> None:
330
- """Broadcast sparse flat-index patches from trainer to vLLM workers."""
331
- if isinstance(trainer_args, dict):
332
- args = NCCLTrainerSendWeightsArgs(**trainer_args)
333
- else:
334
- args = trainer_args
335
-
336
- if args.packed:
337
- raise ValueError(
338
- "Sparse NCCL updates cannot be combined with `packed=True`"
339
- )
340
-
341
- stream = args.stream or torch.cuda.current_stream()
342
- for patch in iterator:
343
- args.group.broadcast(patch.indices, src=args.src, stream=stream)
344
- args.group.broadcast(patch.values, src=args.src, stream=stream)
345
-
346
- @staticmethod
347
- def trainer_init(
348
- init_info: NCCLWeightTransferInitInfo | dict,
349
- ) -> "PyNcclCommunicator":
350
- """
351
- Initialize NCCL process group for trainer-side weight transfer.
352
-
353
- The trainer is always rank 0 in the process group. Uses the current
354
- CUDA device (torch.accelerator.current_device_index()).
355
-
356
- Args:
357
- init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys:
358
- - master_address: str
359
- - master_port: int
360
- - world_size: int
361
-
362
- Returns:
363
- PyNcclCommunicator for weight transfer.
364
-
365
- Example:
366
- >>> from vllm.distributed.weight_transfer.nccl_engine import (
367
- ... NCCLWeightTransferEngine,
368
- ... )
369
- >>> group = NCCLWeightTransferEngine.trainer_init(
370
- ... dict(
371
- ... master_address=master_address,
372
- ... master_port=master_port,
373
- ... world_size=world_size,
374
- ... ),
375
- ... )
376
- """
377
- if isinstance(init_info, dict):
378
- master_address = init_info["master_address"]
379
- master_port = init_info["master_port"]
380
- world_size = init_info["world_size"]
381
- else:
382
- # NCCLWeightTransferInitInfo object
383
- master_address = init_info.master_address
384
- master_port = init_info.master_port
385
- world_size = init_info.world_size
386
-
387
- # Trainer is always rank 0
388
- device = torch.accelerator.current_device_index()
389
- return NCCLWeightTransferEngine._stateless_init_process_group(
390
- master_address,
391
- master_port,
392
- 0,
393
- world_size,
394
- device,
395
- )
396
-
397
- @staticmethod
398
- def _stateless_init_process_group(
399
- master_address, master_port, rank, world_size, device
400
- ):
401
- """
402
- vLLM provides `StatelessProcessGroup` to create a process group
403
- without considering the global process group in torch.distributed.
404
- It is recommended to create `StatelessProcessGroup`, and then initialize
405
- the data-plane communication (NCCL) between external (train processes)
406
- and vLLM workers.
407
- """
408
- from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
409
- from vllm.distributed.utils import StatelessProcessGroup
410
-
411
- pg = StatelessProcessGroup.create(
412
- host=master_address, port=master_port, rank=rank, world_size=world_size
413
- )
414
- pynccl = PyNcclCommunicator(pg, device=device)
415
- return pynccl
265
+ # Trainer-side process-group setup. Delegates to the shared helper so the
266
+ # sparse engine can reuse the exact same rendezvous without subclassing.
267
+ trainer_init = staticmethod(trainer_init)
vllm/distributed/weight_transfer/packed_tensor.py
+6
-3
@@ -186,6 +186,7 @@ def packed_nccl_broadcast_consumer(
186
post_unpack_func: Callable[[list[tuple[str, torch.Tensor]]], None],
187
buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
188
num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
189
+ device: torch.device | str = "cuda",
190
) -> None:
191
"""Consume packed tensors and unpack them into a list of tensors.
192
@@ -199,6 +200,8 @@ def packed_nccl_broadcast_consumer(
200
Both producer and consumer must use the same value.
201
num_buffers: Number of buffers for double/triple buffering.
202
Both producer and consumer must use the same value.
203
+ device: Device for the receive buffers. Must match the device the NCCL
204
+ communicator was created on (the worker's assigned device).
205
206
"""
207
target_packed_tensor_size = buffer_size_bytes
@@ -211,7 +214,7 @@ def packed_nccl_broadcast_consumer(
214
]
215
packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
216
packed_tensors: list[torch.Tensor] = [
214
- torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers)
217
+ torch.empty(0, dtype=torch.uint8, device=device) for _ in range(num_buffers)
218
]
219
220
while True:
@@ -234,7 +237,7 @@ def packed_nccl_broadcast_consumer(
237
break
238
# Create a packed tensor and broadcast it
239
packed_tensors[buffer_idx] = torch.empty(
237
- packing_tensor_sizes[buffer_idx], dtype=torch.uint8, device="cuda"
240
+ packing_tensor_sizes[buffer_idx], dtype=torch.uint8, device=device
241
)
242
group.broadcast(packed_tensors[buffer_idx], src=src)
243
# Load the packed tensor into the model
@@ -259,7 +262,7 @@ def packed_nccl_broadcast_consumer(
262
packed_tensors[buffer_idx] = torch.empty(
263
packing_tensor_sizes[buffer_idx],
264
dtype=torch.uint8,
262
- device="cuda",
265
+ device=device,
266
)
267
group.broadcast(packed_tensors[buffer_idx], src=src)
268
# Load the packed tensor into the model
vllm/distributed/weight_transfer/sparse_nccl_engine.py
new
+223
@@ -0,0 +1,223 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Sparse NCCL weight transfer engine.
4
+
5
+A standalone engine (not a subclass of `NCCLWeightTransferEngine`) for applying
6
+sparse, flat-index weight patches in place. It shares only NCCL process-group
7
+initialization with the dense engine (via `nccl_common`); the update path
8
+applies index/value patches directly to existing model parameters and never runs
9
+layerwise reload.
10
+
11
+MVP limitations:
12
+* TP=1 and PP=1 only
13
+* uses runtime/kernel-format parameter names
14
+* not composable with checkpoint-format or packed updates
15
+"""
16
+
17
+from collections.abc import Iterator
18
+from dataclasses import dataclass
19
+from typing import TYPE_CHECKING, Any
20
+
21
+import torch
22
+
23
+if TYPE_CHECKING:
24
+ from vllm.config import VllmConfig
25
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
26
+
27
+from vllm.config.weight_transfer import WeightTransferConfig
28
+from vllm.distributed.weight_transfer.base import (
29
+ WeightTransferEngine,
30
+ WeightTransferUpdateInfo,
31
+)
32
+from vllm.distributed.weight_transfer.nccl_common import (
33
+ NCCLWeightTransferInitInfo,
34
+ trainer_init,
35
+ worker_init_process_group,
36
+)
37
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs
38
+
39
+__all__ = [
40
+ "SparseWeightPatch",
41
+ "SparseNCCLWeightTransferUpdateInfo",
42
+ "SparseNCCLWeightTransferEngine",
43
+]
44
+
45
+
46
+@dataclass
47
+class SparseWeightPatch:
48
+ """A sparse in-place patch for one existing parameter."""
49
+
50
+ name: str
51
+ indices: torch.Tensor
52
+ values: torch.Tensor
53
+
54
+
55
+@dataclass
56
+class SparseNCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
57
+ """Update info for the sparse NCCL weight transfer backend."""
58
+
59
+ names: list[str]
60
+ dtype_names: list[str]
61
+ shapes: list[list[int]]
62
+ num_updates_list: list[int]
63
+ """Number of sparse entries to receive for each parameter in ``names``."""
64
+
65
+ def __post_init__(self) -> None:
66
+ num_params = len(self.names)
67
+ if len(self.dtype_names) != num_params:
68
+ raise ValueError(
69
+ f"`dtype_names` should be of the same size as `names`: "
70
+ f"got {len(self.dtype_names)} and {len(self.names)}"
71
+ )
72
+ if len(self.shapes) != num_params:
73
+ raise ValueError(
74
+ f"`shapes` should be of the same size as `names`: "
75
+ f"got {len(self.shapes)} and {len(self.names)}"
76
+ )
77
+ if len(self.num_updates_list) == 0:
78
+ raise ValueError("`num_updates_list` cannot be empty for sparse updates")
79
+ if len(self.num_updates_list) != num_params:
80
+ raise ValueError(
81
+ f"`num_updates_list` should be of the same size as `names`: "
82
+ f"got {len(self.num_updates_list)} and {len(self.names)}"
83
+ )
84
+ if any(num_updates < 0 for num_updates in self.num_updates_list):
85
+ raise ValueError("Sparse `num_updates_list` entries must be non-negative")
86
+
87
+
88
+class SparseNCCLWeightTransferEngine(
89
+ WeightTransferEngine[NCCLWeightTransferInitInfo, SparseNCCLWeightTransferUpdateInfo]
90
+):
91
+ """
92
+ Sparse weight transfer engine using NCCL.
93
+
94
+ Receives flat-index (indices, values) patches broadcast from the trainer
95
+ (rank 0) and applies them in place to existing model parameters. Weights are
96
+ applied directly without layerwise reload, so `start_weight_update` and
97
+ `finish_weight_update` are no-ops.
98
+ """
99
+
100
+ init_info_cls = NCCLWeightTransferInitInfo
101
+ update_info_cls = SparseNCCLWeightTransferUpdateInfo
102
+
103
+ def __init__(
104
+ self,
105
+ config: WeightTransferConfig,
106
+ vllm_config: "VllmConfig",
107
+ device: torch.device,
108
+ model: torch.nn.Module,
109
+ ) -> None:
110
+ super().__init__(config, vllm_config, device, model)
111
+ self.model_update_group: PyNcclCommunicator | None = None
112
+
113
+ def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
114
+ """Initialize the NCCL process group with the trainer."""
115
+ self.model_update_group = worker_init_process_group(
116
+ init_info, self.parallel_config
117
+ )
118
+
119
+ def start_weight_update(self) -> None:
120
+ """No-op: sparse patches are applied in place, no layerwise reload."""
121
+ if self.parallel_config.world_size != 1:
122
+ raise NotImplementedError(
123
+ "Sparse weight updates currently require TP=1 and PP=1"
124
+ )
125
+
126
+ def finish_weight_update(self) -> None:
127
+ """No-op: sparse patches are applied in place, no layerwise reload."""
128
+ pass
129
+
130
+ def receive_weights(self, update_info: SparseNCCLWeightTransferUpdateInfo) -> None:
131
+ """Receive sparse flat-index patches from the trainer and apply them."""
132
+ if self.model_update_group is None:
133
+ raise RuntimeError(
134
+ "NCCL weight transfer not initialized. "
135
+ "Call init_transfer_engine() first."
136
+ )
137
+
138
+ # Use the worker's assigned device rather than the ambient current
139
+ # device: the receive path is no longer wrapped in
140
+ # `with torch.device(self.device)` by the caller, so the current device
141
+ # is not guaranteed to match self.device. The recv buffers must live on
142
+ # the same device as the NCCL communicator (created on self.device).
143
+ device = self.device
144
+ for name, dtype_name, num_updates in zip(
145
+ update_info.names,
146
+ update_info.dtype_names,
147
+ update_info.num_updates_list,
148
+ ):
149
+ dtype = getattr(torch, dtype_name)
150
+ indices = torch.empty(num_updates, dtype=torch.int32, device=device)
151
+ values = torch.empty(num_updates, dtype=dtype, device=device)
152
+ self.model_update_group.broadcast(
153
+ indices, src=0, stream=torch.cuda.current_stream()
154
+ )
155
+ self.model_update_group.broadcast(
156
+ values, src=0, stream=torch.cuda.current_stream()
157
+ )
158
+ self._apply_patch(
159
+ SparseWeightPatch(name=name, indices=indices, values=values)
160
+ )
161
+ del indices
162
+ del values
163
+
164
+ def _apply_patch(self, patch: SparseWeightPatch) -> None:
165
+ """Apply a single sparse flat-index patch to an existing model param."""
166
+ param = self.model.get_parameter(patch.name)
167
+ if not param.data.is_contiguous():
168
+ raise NotImplementedError(
169
+ "Sparse weight updates currently require contiguous params: "
170
+ f"{patch.name}"
171
+ )
172
+ if patch.indices.dtype != torch.int32:
173
+ raise ValueError(
174
+ f"Sparse weight updates currently require int32 indices: {patch.name}"
175
+ )
176
+ if patch.indices.ndim != 1 or patch.values.ndim != 1:
177
+ raise ValueError(
178
+ f"Sparse weight patches must be 1D flattened updates: {patch.name}"
179
+ )
180
+ if patch.indices.numel() != patch.values.numel():
181
+ raise ValueError(
182
+ f"`indices` and `values` must have matching lengths for {patch.name}"
183
+ )
184
+ if patch.values.dtype != param.dtype:
185
+ raise ValueError(
186
+ f"Sparse values dtype {patch.values.dtype} does not match "
187
+ f"parameter dtype {param.dtype} for {patch.name}"
188
+ )
189
+
190
+ flat_param = param.data.view(-1)
191
+ flat_param.index_copy_(
192
+ 0,
193
+ patch.indices.to(device=flat_param.device, dtype=torch.long),
194
+ patch.values.to(device=flat_param.device),
195
+ )
196
+
197
+ def shutdown(self) -> None:
198
+ if self.model_update_group is not None:
199
+ self.model_update_group = None
200
+
201
+ @staticmethod
202
+ def trainer_send_weights(
203
+ iterator: Iterator[SparseWeightPatch],
204
+ trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
205
+ ) -> None:
206
+ """Broadcast sparse flat-index patches from trainer to vLLM workers."""
207
+ if isinstance(trainer_args, dict):
208
+ args = NCCLTrainerSendWeightsArgs(**trainer_args)
209
+ else:
210
+ args = trainer_args
211
+
212
+ if args.packed:
213
+ raise ValueError(
214
+ "Sparse NCCL updates cannot be combined with `packed=True`"
215
+ )
216
+
217
+ stream = args.stream or torch.cuda.current_stream()
218
+ for patch in iterator:
219
+ args.group.broadcast(patch.indices, src=args.src, stream=stream)
220
+ args.group.broadcast(patch.values, src=args.src, stream=stream)
221
+
222
+ # Trainer-side process-group setup (shared with the dense engine).
223
+ trainer_init = staticmethod(trainer_init)
vllm/engine/protocol.py
+1
-1
@@ -244,7 +244,7 @@ class EngineClient(ABC):
244
"""Initialize weight transfer for RL training."""
245
raise NotImplementedError
246
247
- async def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
247
+ async def start_weight_update(self) -> None:
248
"""Start a new weight update."""
249
raise NotImplementedError
250
vllm/entrypoints/llm.py
+2
-5
@@ -873,12 +873,9 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
873
"init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
874
)
875
876
- def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
876
+ def start_weight_update(self) -> None:
877
"""Start a new weight update."""
878
- self.llm_engine.collective_rpc(
879
- "start_weight_update",
880
- kwargs={"is_checkpoint_format": is_checkpoint_format},
881
- )
878
+ self.llm_engine.collective_rpc("start_weight_update")
879
880
def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None:
881
"""
vllm/entrypoints/serve/dev/rlhf/api_router.py
+1
-8
@@ -129,14 +129,7 @@ async def init_weight_transfer_engine(raw_request: Request):
129
130
@router.post("/start_weight_update")
131
async def start_weight_update(raw_request: Request):
132
- try:
133
- body = await raw_request.json()
134
- except json.JSONDecodeError as e:
135
- raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904
136
- is_checkpoint_format = body.get("is_checkpoint_format", True)
137
- await engine_client(raw_request).start_weight_update(
138
- is_checkpoint_format=is_checkpoint_format
139
- )
132
+ await engine_client(raw_request).start_weight_update()
133
return JSONResponse(content={"message": "Weight update started"})
134
135
vllm/v1/engine/async_llm.py
+2
-5
@@ -1080,12 +1080,9 @@ class AsyncLLM(EngineClient):
1080
"init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
1081
)
1082
1083
- async def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
1083
+ async def start_weight_update(self) -> None:
1084
"""Start a new weight update."""
1085
- await self.collective_rpc(
1086
- "start_weight_update",
1087
- kwargs={"is_checkpoint_format": is_checkpoint_format},
1088
- )
1085
+ await self.collective_rpc("start_weight_update")
1086
1087
async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
1088
"""
vllm/v1/worker/gpu/model_runner.py
-6
@@ -377,12 +377,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
377
378
GPUModelRunnerV1.reload_weights(self, *args, **kwargs) # type: ignore[arg-type]
379
380
- def apply_sparse_weight_patches(self, *args, **kwargs) -> None:
381
- # TODO: Use full version instead of import when fully migrated to v2
382
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
383
-
384
- GPUModelRunnerV1.apply_sparse_weight_patches(self, *args, **kwargs) # type: ignore[arg-type]
385
-
380
def update_config(self, *args, **kwargs) -> None:
381
# TODO(Wentao): Use full version instead of import when fully migrated to v2
382
from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
vllm/v1/worker/gpu_model_runner.py
-39
@@ -49,7 +49,6 @@ from vllm.distributed.parallel_state import (
49
is_global_first_rank,
50
prepare_communication_buffer_for_model,
51
)
52
-from vllm.distributed.weight_transfer.base import SparseWeightPatch
52
from vllm.forward_context import (
53
BatchDescriptor,
54
set_forward_context,
@@ -3244,44 +3243,6 @@ class GPUModelRunner(
3243
return self.model.unwrap()
3244
return self.model
3245
3247
- def apply_sparse_weight_patches(self, patches: Iterable[SparseWeightPatch]) -> None:
3248
- """Apply sparse flat-index patches directly to existing model params."""
3249
- model = self.get_model()
3250
- for patch in patches:
3251
- param = model.get_parameter(patch.name)
3252
- if not param.data.is_contiguous():
3253
- raise NotImplementedError(
3254
- "Sparse weight updates currently require contiguous params: "
3255
- f"{patch.name}"
3256
- )
3257
-
3258
- if patch.indices.dtype != torch.int32:
3259
- raise ValueError(
3260
- "Sparse weight updates currently require int32 indices: "
3261
- f"{patch.name}"
3262
- )
3263
- if patch.indices.ndim != 1 or patch.values.ndim != 1:
3264
- raise ValueError(
3265
- f"Sparse weight patches must be 1D flattened updates: {patch.name}"
3266
- )
3267
- if patch.indices.numel() != patch.values.numel():
3268
- raise ValueError(
3269
- "`indices` and `values` must have matching lengths for "
3270
- f"{patch.name}"
3271
- )
3272
- if patch.values.dtype != param.dtype:
3273
- raise ValueError(
3274
- f"Sparse values dtype {patch.values.dtype} does not match "
3275
- f"parameter dtype {param.dtype} for {patch.name}"
3276
- )
3277
-
3278
- flat_param = param.data.view(-1)
3279
- flat_param.index_copy_(
3280
- 0,
3281
- patch.indices.to(device=flat_param.device, dtype=torch.long),
3282
- patch.values.to(device=flat_param.device),
3283
- )
3284
-
3246
def get_supported_generation_tasks(self) -> list[GenerationTask]:
3247
model = self.get_model()
3248
supported_tasks = list[GenerationTask]()
vllm/v1/worker/gpu_worker.py
+14
-81
@@ -155,7 +155,6 @@ class Worker(WorkerBase):
155
# is available, since the engine needs a reference to the model.
156
self.weight_transfer_engine: WeightTransferEngine | None = None
157
self._weight_update_active = False
158
- self._is_checkpoint_format = True
158
159
# Torch/CUDA profiler. Enabled and configured through profiler_config.
160
# Profiler wrapper is created lazily in profile() when start is called,
@@ -416,7 +415,8 @@ class Worker(WorkerBase):
415
if self.vllm_config.weight_transfer_config is not None:
416
self.weight_transfer_engine = WeightTransferEngineFactory.create_engine(
417
self.vllm_config.weight_transfer_config,
419
- self.vllm_config.parallel_config,
418
+ self.vllm_config,
419
+ self.device,
420
self.model_runner.get_model(),
421
)
422
@@ -1161,16 +1161,16 @@ class Worker(WorkerBase):
1161
typed_init_info = self.weight_transfer_engine.parse_init_info(init_info)
1162
self.weight_transfer_engine.init_transfer_engine(typed_init_info)
1163
1164
- def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
1164
+ def start_weight_update(self) -> None:
1165
"""
1166
Start a new weight update session.
1167
1168
- Args:
1169
- is_checkpoint_format: Whether incoming weights are in checkpoint
1170
- format (need layerwise processing) or kernel format (direct
1171
- copy / sparse patch application).
1168
+ Delegates engine-specific preparation (e.g. layerwise reload setup) to
1169
+ the configured weight transfer engine. The worker only tracks that a
1170
+ session is active.
1171
"""
1172
self._check_weight_transfer_engine()
1173
+ assert self.weight_transfer_engine is not None
1174
1175
if self._weight_update_active:
1176
raise RuntimeError(
@@ -1178,16 +1178,7 @@ class Worker(WorkerBase):
1178
"active. Call finish_weight_update first."
1179
)
1180
1181
- if is_checkpoint_format:
1182
- from vllm.model_executor.model_loader.reload import (
1183
- initialize_layerwise_reload,
1184
- )
1185
-
1186
- model = self.model_runner.model
1187
- with torch.device(self.device):
1188
- initialize_layerwise_reload(model)
1189
-
1190
- self._is_checkpoint_format = is_checkpoint_format
1181
+ self.weight_transfer_engine.start_weight_update()
1182
self._weight_update_active = True
1183
1184
def update_weights(self, update_info: dict) -> None:
@@ -1208,82 +1199,24 @@ class Worker(WorkerBase):
1199
"start_weight_update must be called before update_weights."
1200
)
1201
1211
- update_succeeded = False
1202
try:
1213
- # Parse dict into backend-specific typed dataclass
1214
- typed_update_info = self.weight_transfer_engine.parse_update_info(
1215
- update_info
1216
- )
1217
-
1218
- with torch.device(self.device):
1219
- if self._is_checkpoint_format:
1220
- if typed_update_info.update_kind != "dense":
1221
- raise ValueError(
1222
- "Sparse weight updates require "
1223
- "`start_weight_update(is_checkpoint_format=False)`."
1224
- )
1225
-
1226
- model = self.model_runner.model
1227
-
1228
- # Use layerwise reload pattern for checkpoint format weights
1229
- self.weight_transfer_engine.receive_weights(
1230
- typed_update_info,
1231
- load_weights=model.load_weights,
1232
- )
1233
- elif typed_update_info.update_kind == "sparse_flat":
1234
- if self.parallel_config.world_size != 1:
1235
- raise NotImplementedError(
1236
- "Sparse weight updates currently require TP=1 and PP=1"
1237
- )
1238
- self.weight_transfer_engine.receive_sparse_weights(
1239
- typed_update_info,
1240
- apply_patches=self.model_runner.apply_sparse_weight_patches,
1241
- )
1242
- else:
1243
- model = self.model_runner.model
1244
-
1245
- # Weights are already in kernel format, copy directly.
1246
- def load_weights_direct(
1247
- weights: list[tuple[str, torch.Tensor]],
1248
- ) -> None:
1249
- for name, weight in weights:
1250
- param = model.get_parameter(name)
1251
- param.copy_(weight)
1252
-
1253
- self.weight_transfer_engine.receive_weights(
1254
- typed_update_info,
1255
- load_weights=load_weights_direct,
1256
- )
1257
-
1258
- # NCCL broadcast/packed path are asynchronous.
1259
- # Sync here so the next step uses the new weights.
1260
- torch.accelerator.synchronize()
1261
- update_succeeded = True
1262
- finally:
1263
- if not update_succeeded:
1264
- self._weight_update_active = False
1265
- self._is_checkpoint_format = True
1203
+ self.weight_transfer_engine.update_weights(update_info)
1204
+ except BaseException:
1205
+ self._weight_update_active = False
1206
+ raise
1207
1208
def finish_weight_update(self) -> None:
1209
"""Finish the current weight update session."""
1210
self._check_weight_transfer_engine()
1211
+ assert self.weight_transfer_engine is not None
1212
1213
if not self._weight_update_active:
1214
raise RuntimeError(
1215
"finish_weight_update called without a matching start_weight_update."
1216
)
1217
1276
- if self._is_checkpoint_format:
1277
- from vllm.model_executor.model_loader.reload import (
1278
- finalize_layerwise_reload,
1279
- )
1280
-
1281
- model = self.model_runner.model
1282
- with torch.device(self.device):
1283
- finalize_layerwise_reload(model, self.model_config)
1284
-
1218
+ self.weight_transfer_engine.finish_weight_update()
1219
self._weight_update_active = False
1286
- self._is_checkpoint_format = True
1220
1221
def shutdown(self) -> None:
1222
gc.unfreeze()