[rl] Stateful Trainer Send: NCCL + Sparse NCCL [3/N] (#50902)
Signed-off-by: haoaaron <ahao@anyscale.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Aaron Hao committed
Aug 7, 2026 at 00:35 UTC
21ea5b4fa1062a379dd7e6795497ad6becd5a856
14 files changed
+1418
-779
examples/rl/rlhf_async_new_apis.py
+18
-77
@@ -41,7 +41,6 @@ causes unexpected behavior.
41
42
import asyncio
43
import uuid
44
-from dataclasses import asdict
44
45
import ray
46
import torch
@@ -50,16 +49,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
49
import vllm
50
from vllm import SamplingParams
51
from vllm.config import WeightTransferConfig
53
-from vllm.distributed.weight_transfer.base import (
54
- WeightTransferInitRequest,
55
- WeightTransferUpdateRequest,
56
-)
57
-from vllm.distributed.weight_transfer.nccl_engine import (
58
- NCCLTrainerSendWeightsArgs,
59
- NCCLWeightTransferEngine,
60
- NCCLWeightTransferInitInfo,
61
- NCCLWeightTransferUpdateInfo,
52
+from vllm.distributed.weight_transfer import (
53
+ ModuleSource,
54
+ RayVLLMWeightSyncClient,
55
+ WeightTransferTrainerFactory,
56
)
57
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
58
from vllm.platforms import current_platform
59
from vllm.utils.network_utils import get_ip, get_open_port
60
from vllm.v1.executor import Executor
@@ -141,40 +136,23 @@ class TrainModel:
136
self.port = get_open_port()
137
self.master_address = get_ip()
138
144
- def get_master_address_and_port(self):
145
- return self.master_address, self.port
146
-
147
- def get_weight_metadata(self):
148
- """Return weight names, dtypes, and shapes for weight transfer."""
149
- names = []
150
- dtype_names = []
151
- shapes = []
152
- for name, p in self.model.named_parameters():
153
- names.append(name)
154
- dtype_names.append(str(p.dtype).split(".")[-1])
155
- shapes.append(list(p.shape))
156
- return names, dtype_names, shapes
157
-
158
- def init_weight_transfer_group(self, world_size):
159
- """Initialize the NCCL process group for weight transfer."""
160
- self.model_update_group = NCCLWeightTransferEngine.trainer_init(
161
- dict(
139
+ def init_weight_transfer(self, world_size, llm_handle):
140
+ """Build the trainer-side weight-transfer engine and rendezvous."""
141
+ self.engine = WeightTransferTrainerFactory.trainer_init(
142
+ init_info=NCCLTrainerInitInfo(
143
master_address=self.master_address,
144
master_port=self.port,
145
world_size=world_size,
146
+ rank=0, # single-GPU trainer is the sole (sender) rank
147
+ packed=True,
148
),
149
+ client=RayVLLMWeightSyncClient(llm_handle),
150
+ source=ModuleSource(self.model),
151
)
152
168
- def broadcast_weights(self, packed: bool = True):
169
- """Broadcast weights to the inference engine."""
170
- trainer_args = NCCLTrainerSendWeightsArgs(
171
- group=self.model_update_group,
172
- packed=packed,
173
- )
174
- NCCLWeightTransferEngine.trainer_send_weights(
175
- iterator=self.model.named_parameters(),
176
- trainer_args=trainer_args,
177
- )
153
+ def broadcast_weights(self):
154
+ """Push weights to the inference engine (drives start/update/finish)."""
155
+ self.engine.send_weights()
156
157
@torch.inference_mode()
158
def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
@@ -263,32 +241,12 @@ batch_prompt_token_ids = [
241
242
# Set up the communication channel between the training process and the
243
# inference engine.
266
-master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
267
-
244
world_size = 2 # 1 trainer + 1 inference worker
269
-inference_handle = llm.init_weight_transfer_engine.remote(
270
- WeightTransferInitRequest(
271
- init_info=asdict(
272
- NCCLWeightTransferInitInfo(
273
- master_address=master_address,
274
- master_port=master_port,
275
- rank_offset=1,
276
- world_size=world_size,
277
- )
278
- )
279
- )
280
-)
281
-
282
-# Initialize weight transfer group on both the training actor and inference engine
283
-train_handle = train_model.init_weight_transfer_group.remote(world_size)
284
-ray.get([train_handle, inference_handle])
245
+ray.get(train_model.init_weight_transfer.remote(world_size, llm))
246
247
248
N_NEW_TOKENS = 100
249
289
-# Collect weight metadata once
290
-names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
291
-
250
# ── Phase 1: concurrent requests with weight sync ───────────────────
251
print(f"\n{'=' * 50}")
252
print(f"Prompts ({len(PROMPTS)}):")
@@ -306,24 +264,7 @@ gen_futures = [
264
265
ray.get(llm.pause_after_n_tokens.remote())
266
309
-ray.get(llm.start_weight_update.remote())
310
-
311
-inference_handle = llm.update_weights.remote(
312
- WeightTransferUpdateRequest(
313
- update_info=asdict(
314
- NCCLWeightTransferUpdateInfo(
315
- names=names,
316
- dtype_names=dtype_names,
317
- shapes=shapes,
318
- packed=True,
319
- )
320
- )
321
- )
322
-)
323
-train_handle = train_model.broadcast_weights.remote(packed=True)
324
-ray.get([train_handle, inference_handle])
325
-
326
-ray.get(llm.finish_weight_update.remote())
267
+ray.get(train_model.broadcast_weights.remote())
268
269
ray.get(llm.resume_generation.remote())
270
results = ray.get(gen_futures)
examples/rl/rlhf_http_nccl.py
+19
-113
@@ -38,10 +38,12 @@ import torch
38
from openai import OpenAI
39
from transformers import AutoModelForCausalLM
40
41
-from vllm.distributed.weight_transfer.nccl_engine import (
42
- NCCLTrainerSendWeightsArgs,
43
- NCCLWeightTransferEngine,
41
+from vllm.distributed.weight_transfer import (
42
+ HTTPVLLMWeightSyncClient,
43
+ ModuleSource,
44
+ WeightTransferTrainerFactory,
45
)
46
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
47
from vllm.utils.network_utils import get_ip, get_open_port
48
49
BASE_URL = "http://localhost:8000"
@@ -62,62 +64,6 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
64
return results
65
66
65
-def init_weight_transfer_engine(
66
- base_url: str,
67
- master_address: str,
68
- master_port: int,
69
- rank_offset: int,
70
- world_size: int,
71
-) -> None:
72
- """Initialize weight transfer via HTTP endpoint."""
73
- url = f"{base_url}/init_weight_transfer_engine"
74
- payload = {
75
- "init_info": dict(
76
- master_address=master_address,
77
- master_port=master_port,
78
- rank_offset=rank_offset,
79
- world_size=world_size,
80
- )
81
- }
82
- response = requests.post(url, json=payload, timeout=60)
83
- response.raise_for_status()
84
-
85
-
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"
89
- response = requests.post(url, json={}, timeout=60)
90
- response.raise_for_status()
91
-
92
-
93
-def update_weights(
94
- base_url: str,
95
- names: list[str],
96
- dtype_names: list[str],
97
- shapes: list[list[int]],
98
- packed: bool = False,
99
-) -> None:
100
- """Update weights via HTTP endpoint."""
101
- url = f"{base_url}/update_weights"
102
- payload = {
103
- "update_info": dict(
104
- names=names,
105
- dtype_names=dtype_names,
106
- shapes=shapes,
107
- packed=packed,
108
- )
109
- }
110
- response = requests.post(url, json=payload, timeout=300)
111
- response.raise_for_status()
112
-
113
-
114
-def finish_weight_update(base_url: str) -> None:
115
- """Finish a weight update via HTTP endpoint."""
116
- url = f"{base_url}/finish_weight_update"
117
- response = requests.post(url, json={}, timeout=60)
118
- response.raise_for_status()
119
-
120
-
67
def pause_generation(base_url: str) -> None:
68
"""Pause generation via HTTP endpoint."""
69
url = f"{base_url}/pause"
@@ -177,75 +123,35 @@ def main():
123
print("-" * 50)
124
125
# Set up the communication channel between the training process and the
180
- # vLLM server. The trainer is rank 0, vLLM worker(s) start at rank_offset.
126
+ # vLLM server. The trainer is rank 0, vLLM worker(s) start after it.
127
master_address = get_ip()
128
master_port = get_open_port()
183
- rank_offset = 1
129
130
print(f"Initializing weight transfer: master={master_address}:{master_port}")
131
187
- # Initialize weight transfer on vLLM server (this is async, server will
188
- # wait for NCCL connection)
189
- import threading
190
-
191
- init_thread = threading.Thread(
192
- target=init_weight_transfer_engine,
193
- args=(BASE_URL, master_address, master_port, rank_offset, world_size),
194
- )
195
- init_thread.start()
196
-
197
- # Initialize NCCL process group on trainer side
198
- model_update_group = NCCLWeightTransferEngine.trainer_init(
199
- dict(
132
+ # The trainer engine owns the full handshake. `trainer_init` kicks off the
133
+ # server's init_weight_transfer_engine (via the HTTP client) on a worker
134
+ # thread while opening the trainer's NCCL endpoint, so both ends rendezvous
135
+ # together — no manual threading needed here.
136
+ engine = WeightTransferTrainerFactory.trainer_init(
137
+ init_info=NCCLTrainerInitInfo(
138
master_address=master_address,
139
master_port=master_port,
140
world_size=world_size,
141
+ rank=0, # single-GPU trainer is the sole (sender) rank
142
+ packed=True,
143
),
144
+ client=HTTPVLLMWeightSyncClient(BASE_URL),
145
+ source=ModuleSource(train_model),
146
)
147
206
- # Wait for init_weight_transfer_engine to complete
207
- init_thread.join()
208
-
148
# Pause generation before weight sync
149
pause_generation(BASE_URL)
150
212
- # Collect weight metadata for the update request
213
- names = []
214
- dtype_names = []
215
- shapes = []
216
- for name, p in train_model.named_parameters():
217
- names.append(name)
218
- dtype_names.append(str(p.dtype).split(".")[-1])
219
- shapes.append(list(p.shape))
220
-
221
- # Start weight update
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
226
- # packed=True enables efficient batched tensor broadcasting
227
- update_thread = threading.Thread(
228
- target=update_weights,
229
- args=(BASE_URL, names, dtype_names, shapes, True), # packed=True
230
- )
231
- update_thread.start()
232
-
233
- # Broadcast all weights from trainer to vLLM workers
151
+ # One call drives start_weight_update / update_weights / finish_weight_update
152
+ # over HTTP, concurrent with the NCCL broadcast.
153
print("Broadcasting weights via NCCL...")
235
- trainer_args = NCCLTrainerSendWeightsArgs(
236
- group=model_update_group,
237
- packed=True,
238
- )
239
- NCCLWeightTransferEngine.trainer_send_weights(
240
- iterator=train_model.named_parameters(),
241
- trainer_args=trainer_args,
242
- )
243
-
244
- # Wait for update_weights to complete
245
- update_thread.join()
246
-
247
- # Finish weight update
248
- finish_weight_update(BASE_URL)
154
+ engine.send_weights()
155
156
# Resume generation after weight sync
157
resume_generation(BASE_URL)
examples/rl/rlhf_nccl.py
+33
-74
@@ -36,10 +36,12 @@ from transformers import AutoModelForCausalLM
36
37
from vllm import LLM, SamplingParams
38
from vllm.config import WeightTransferConfig
39
-from vllm.distributed.weight_transfer.nccl_engine import (
40
- NCCLTrainerSendWeightsArgs,
41
- NCCLWeightTransferEngine,
39
+from vllm.distributed.weight_transfer import (
40
+ ModuleSource,
41
+ RayVLLMWeightSyncClient,
42
+ WeightTransferTrainerFactory,
43
)
44
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
45
from vllm.platforms import current_platform
46
from vllm.utils.network_utils import get_ip, get_open_port
47
@@ -80,40 +82,33 @@ class TrainModel:
82
self.port = get_open_port()
83
self.master_address = get_ip()
84
83
- def get_master_address_and_port(self):
84
- return self.master_address, self.port
85
-
86
- def get_weight_metadata(self):
87
- """Return weight names, dtypes, and shapes for weight transfer."""
88
- names = []
89
- dtype_names = []
90
- shapes = []
91
- for name, p in self.model.named_parameters():
92
- names.append(name)
93
- dtype_names.append(str(p.dtype).split(".")[-1])
94
- shapes.append(list(p.shape))
95
- return names, dtype_names, shapes
96
-
97
- def init_weight_transfer_group(self, world_size):
98
- """Initialize the NCCL process group for weight transfer."""
99
- self.model_update_group = NCCLWeightTransferEngine.trainer_init(
100
- dict(
85
+ def init_weight_transfer(self, world_size, llm_handle):
86
+ """Build the trainer-side weight-transfer engine.
87
+
88
+ `trainer_init` drives the full handshake: it kicks off the inference
89
+ side's `init_weight_transfer_engine` (via the Ray client) on a worker
90
+ thread while opening the trainer's NCCL endpoint, so both ends
91
+ rendezvous together. After this returns, `send_weights()` is callable.
92
+ """
93
+ self.engine = WeightTransferTrainerFactory.trainer_init(
94
+ init_info=NCCLTrainerInitInfo(
95
master_address=self.master_address,
96
master_port=self.port,
97
world_size=world_size,
98
+ rank=0, # single-GPU trainer is the sole (sender) rank
99
+ packed=True,
100
),
101
+ client=RayVLLMWeightSyncClient(llm_handle),
102
+ source=ModuleSource(self.model),
103
)
104
107
- def broadcast_weights(self, packed: bool = True):
108
- """Broadcast weights to the inference engine."""
109
- trainer_args = NCCLTrainerSendWeightsArgs(
110
- group=self.model_update_group,
111
- packed=packed,
112
- )
113
- NCCLWeightTransferEngine.trainer_send_weights(
114
- iterator=self.model.named_parameters(),
115
- trainer_args=trainer_args,
116
- )
105
+ def broadcast_weights(self):
106
+ """Push the current weights to the inference engine.
107
+
108
+ Drives start/update/finish on the inference side and the NCCL
109
+ broadcast internally — one call.
110
+ """
111
+ self.engine.send_weights()
112
113
114
# Initialize Ray and set the visible devices. The vLLM engine will
@@ -178,51 +173,15 @@ for output in outputs:
173
ray.get(llm.sleep.remote(level=0))
174
175
# Set up the communication channel between the training process and the
181
-# inference engine.
182
-master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
183
-
176
+# inference engine. The trainer engine now owns the full handshake — the
177
+# driver only has to hand it the vLLM actor handle and the world size.
178
world_size = ray.get(llm.get_world_size.remote()) + 1 # +1 for the trainer
185
-inference_handle = llm.init_weight_transfer_engine.remote(
186
- dict(
187
- init_info=dict(
188
- master_address=master_address,
189
- master_port=master_port,
190
- rank_offset=1,
191
- world_size=world_size,
192
- )
193
- )
194
-)
195
-
196
-# Initialize weight transfer group on both the training actor and inference engine
197
-train_handle = train_model.init_weight_transfer_group.remote(world_size)
198
-ray.get([train_handle, inference_handle])
199
-
200
-# Synchronize the updated weights to the inference engine using batched API.
201
-# Collect all weight metadata from the training actor
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())
206
-
207
-# Issue update_weights call with NCCL-specific update info
208
-# packed=True enables efficient batched tensor broadcasting
209
-inference_handle = llm.update_weights.remote(
210
- dict(
211
- update_info=dict(
212
- names=names,
213
- dtype_names=dtype_names,
214
- shapes=shapes,
215
- packed=True,
216
- )
217
- )
218
-)
219
-
220
-# Broadcast all weights from trainer using the weight transfer API
221
-train_handle = train_model.broadcast_weights.remote(packed=True)
222
-ray.get([train_handle, inference_handle])
179
+ray.get(train_model.init_weight_transfer.remote(world_size, llm))
180
224
-# Finish weight update
225
-ray.get(llm.finish_weight_update.remote())
181
+# Synchronize the updated weights to the inference engine. The engine drives
182
+# start_weight_update / update_weights / finish_weight_update on the inference
183
+# side internally, concurrent with the NCCL broadcast.
184
+ray.get(train_model.broadcast_weights.remote())
185
186
ray.get(llm.wake_up.remote(tags=["scheduling"]))
187
examples/rl/rlhf_nccl_fsdp_ep.py
+241
-215
@@ -4,69 +4,79 @@
4
RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs).
5
6
8-GPU layout:
7
- Training — 4 GPUs, PyTorch FSDP2 (fully_shard)
8
- Inference — 4 GPUs, vLLM AsyncLLMEngine with expert parallelism +
7
+ Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors
8
+ Inference — 4 GPUs, a `vllm serve` HTTP server with expert parallelism +
9
data parallelism (TP=1, DP=4, enable_expert_parallel
10
→ EP_SIZE = TP×DP = 4)
11
12
-FSDP workers are Ray actors that form a single FSDP2 process group.
13
-Rank 0 gathers full parameters via DTensor.full_tensor() and broadcasts
14
-them to the vLLM inference engine through the NCCL weight-transfer API.
12
+The inference side is a standalone HTTP server (spawned by this script with
13
+`vllm serve`), so both the weight-sync control plane (HTTP) and the NCCL data
14
+plane run inside the rank-0 FSDP Ray actor. That lets the trainer use the
15
+unified `TrainerWeightTransferEngine.send_weights()` with an
16
+`HTTPVLLMWeightSyncClient` — one call drives start/update/finish on the server
17
+concurrently with the NCCL broadcast. Every FSDP rank builds an engine and calls
18
+`send_weights()`, so all 4 participate in the incremental `full_tensor()`
19
+all-gather; only rank 0 holds a communicator and broadcasts (it is the only
20
+trainer rank in the NCCL group).
21
16
-The inference engine uses AsyncLLMEngine which automatically spawns
17
-DP worker processes (no manual placement group needed). Weight sync
18
-uses pause_generation / resume_generation.
22
+GPU split (single node): the server takes GPUs 0-3 (CUDA_VISIBLE_DEVICES), and
23
+Ray (training) is restricted to GPUs 4-7.
24
25
Steps:
21
- 1. Launch 4 FSDP training workers.
22
- 2. Launch AsyncLLMEngine with EP+DP (dummy weights).
23
- 3. Generate from prompts → gibberish (random weights).
24
- 4. Pause generation, transfer weights from FSDP, resume.
26
+ 1. Launch the vLLM HTTP server (EP+DP, dummy weights) on GPUs 0-3.
27
+ 2. Launch 4 FSDP training workers (Ray) on GPUs 4-7.
28
+ 3. Generate from prompts over HTTP → gibberish (random weights).
29
+ 4. Pause generation, transfer weights FSDP → server over NCCL, resume.
30
5. Generate from prompts → sensible output (synced weights).
31
32
Assumes a single-node cluster with 8 GPUs.
33
"""
34
30
-import asyncio
35
+import json
36
import os
32
-import uuid
33
-from dataclasses import asdict
37
+import subprocess
38
+import sys
39
+import time
40
41
import ray
42
+import requests
43
import torch
44
import torch.distributed as dist
45
from huggingface_hub import snapshot_download
46
+from openai import OpenAI
47
from torch.distributed.fsdp import fully_shard
48
from transformers import AutoModelForCausalLM
49
42
-import vllm
43
-from vllm import SamplingParams
44
-from vllm.config import WeightTransferConfig
45
-from vllm.distributed.weight_transfer.base import (
46
- WeightTransferInitRequest,
47
- WeightTransferUpdateRequest,
48
-)
49
-from vllm.distributed.weight_transfer.nccl_engine import (
50
- NCCLTrainerSendWeightsArgs,
51
- NCCLWeightTransferEngine,
52
- NCCLWeightTransferInitInfo,
53
- NCCLWeightTransferUpdateInfo,
50
+from vllm.distributed.weight_transfer import (
51
+ HTTPVLLMWeightSyncClient,
52
+ ModuleSource,
53
+ WeightTransferTrainerFactory,
54
)
55
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
56
from vllm.utils.network_utils import get_ip, get_open_port
56
-from vllm.v1.executor import Executor
57
58
MODEL_NAME = "Qwen/Qwen3-30B-A3B"
59
+SERVED_MODEL_NAME = "policy"
60
61
FSDP_WORLD_SIZE = 4
62
INFERENCE_TP_SIZE = 1
63
INFERENCE_DP_SIZE = 4
64
65
+# Training (FSDP) GPUs are reserved through Ray; the inference server then runs
66
+# on the complementary GPUs (see main()). We do NOT hard-code the split via
67
+# CUDA_VISIBLE_DEVICES before ray.init(): that only restricts Ray when ray.init()
68
+# *starts* a local cluster, and is silently ignored when it connects to an
69
+# existing one (e.g. a shared/managed Ray cluster), causing training and the
70
+# server to collide on the same physical GPUs.
71
+SERVER_PORT = 8000
72
+BASE_URL = f"http://localhost:{SERVER_PORT}"
73
+
74
75
@ray.remote(num_gpus=1)
76
class FSDPTrainWorker:
77
"""
78
One FSDP2 training worker per GPU. Four of these form the FSDP group.
69
- Rank 0 additionally handles weight transfer to the vLLM engine.
79
+ Rank 0 additionally drives weight transfer to the vLLM server.
80
"""
81
82
def __init__(
@@ -78,6 +88,7 @@ class FSDPTrainWorker:
88
fsdp_master_port: int,
89
):
90
self.rank = rank
91
+ self.engine = None
92
93
os.environ["MASTER_ADDR"] = fsdp_master_addr
94
os.environ["MASTER_PORT"] = str(fsdp_master_port)
@@ -89,12 +100,6 @@ class FSDPTrainWorker:
100
model_name, torch_dtype=torch.bfloat16
101
)
102
92
- self.weight_names = [n for n, _ in model.named_parameters()]
93
- self.weight_dtype_names = [
94
- str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
95
- ]
96
- self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
97
-
103
for layer in model.model.layers:
104
fully_shard(layer)
105
fully_shard(model)
@@ -103,11 +108,14 @@ class FSDPTrainWorker:
108
109
self.transfer_port = None
110
self.transfer_master_address = None
106
- self.model_update_group = None
111
112
def get_rank(self):
113
return self.rank
114
115
+ def get_gpu_ids(self):
116
+ """Physical GPU id(s) Ray assigned to this worker (for server/train split)."""
117
+ return ray.get_gpu_ids()
118
+
119
# ---- weight-transfer setup (rank 0 only) ----
120
121
def setup_transfer_endpoint(self):
@@ -117,95 +125,130 @@ class FSDPTrainWorker:
125
self.transfer_master_address = get_ip()
126
return self.transfer_master_address, self.transfer_port
127
120
- def init_weight_transfer_group(self, transfer_world_size: int):
121
- """Join the weight-transfer NCCL group as rank 0 (the source)."""
122
- assert self.rank == 0
123
- self.model_update_group = NCCLWeightTransferEngine.trainer_init(
124
- dict(
125
- master_address=self.transfer_master_address,
126
- master_port=self.transfer_port,
128
+ def setup_engine(
129
+ self,
130
+ base_url: str,
131
+ transfer_master_address: str,
132
+ transfer_port: int,
133
+ transfer_world_size: int,
134
+ ):
135
+ """Build the trainer engine on every FSDP rank.
136
+
137
+ Called on all ranks with the shared rendezvous endpoint. Rank 0 is the
138
+ sender: `trainer_init` opens its rank-0 NCCL endpoint and, on a worker
139
+ thread, calls the server's `init_weight_transfer_engine` over HTTP so
140
+ both ends rendezvous together. The other ranks skip the rendezvous and
141
+ only join the FSDP all-gather during send_weights.
142
+ """
143
+ self.engine = WeightTransferTrainerFactory.trainer_init(
144
+ init_info=NCCLTrainerInitInfo(
145
+ master_address=transfer_master_address,
146
+ master_port=transfer_port,
147
world_size=transfer_world_size,
148
+ rank=self.rank, # FSDP rank; sender is rank 0
149
+ packed=True,
150
),
151
+ client=HTTPVLLMWeightSyncClient(base_url),
152
+ # Yields sharded DTensors; the engine reads global shape/dtype for
153
+ # metadata (no gather) and calls full_tensor() at broadcast time.
154
+ source=ModuleSource(self.model),
155
)
156
131
- def get_weight_metadata(self):
132
- """Return weight names, dtypes, and shapes captured before FSDP wrapping."""
133
- return self.weight_names, self.weight_dtype_names, self.weight_shapes
134
-
157
# ---- collective ops (ALL FSDP ranks must call concurrently) ----
158
137
- def gather_and_broadcast_weights(self, packed: bool = True):
138
- """
139
- All-gather full parameters and broadcast them to vLLM.
140
- Only rank 0 performs the actual NCCL broadcast; others just
141
- participate in the FSDP all-gather.
159
+ def gather_and_broadcast_weights(self):
160
+ """All-gather full parameters and broadcast them to the vLLM server.
161
143
- full_tensor() is a collective — all FSDP ranks must call it
144
- for each parameter in the same order. Rank 0 additionally
145
- feeds each gathered tensor to the weight-transfer engine.
162
+ Called on all FSDP ranks. `send_weights` gathers each param via
163
+ `full_tensor()` (a collective every rank must enter in the same order);
164
+ only rank 0 (the sender) drives the server-side update_weights
165
+ concurrently with the NCCL broadcast — the other ranks only gather.
166
"""
147
- if self.rank == 0:
148
-
149
- def _full_param_iter():
150
- for name, param in self.model.named_parameters():
151
- yield name, param.full_tensor()
152
-
153
- trainer_args = NCCLTrainerSendWeightsArgs(
154
- group=self.model_update_group,
155
- packed=packed,
156
- )
157
- NCCLWeightTransferEngine.trainer_send_weights(
158
- iterator=_full_param_iter(),
159
- trainer_args=trainer_args,
160
- )
161
- else:
162
- for _, param in self.model.named_parameters():
163
- param.full_tensor()
164
-
165
-
166
-def create_async_engine(**kwargs):
167
- """Create an AsyncLLMEngine directly (no subclass needed)."""
168
- engine_args = vllm.AsyncEngineArgs(**kwargs)
169
- vllm_config = engine_args.create_engine_config()
170
- executor_class = Executor.get_class(vllm_config)
171
- return vllm.AsyncLLMEngine(
172
- vllm_config=vllm_config,
173
- executor_class=executor_class,
174
- log_requests=engine_args.enable_log_requests,
175
- log_stats=not engine_args.disable_log_stats,
167
+ self.engine.send_weights()
168
+
169
+
170
+def start_vllm_server(server_gpus: str) -> subprocess.Popen:
171
+ """Spawn a `vllm serve` HTTP server (EP+DP) on `server_gpus` and wait for it."""
172
+ serve_args = [
173
+ "vllm",
174
+ "serve",
175
+ MODEL_NAME,
176
+ "--served-model-name",
177
+ SERVED_MODEL_NAME,
178
+ "--tensor-parallel-size",
179
+ str(INFERENCE_TP_SIZE),
180
+ "--data-parallel-size",
181
+ str(INFERENCE_DP_SIZE),
182
+ "--enable-expert-parallel",
183
+ "--enforce-eager",
184
+ "--load-format",
185
+ "dummy",
186
+ "--gpu-memory-utilization",
187
+ "0.7",
188
+ "--port",
189
+ str(SERVER_PORT),
190
+ "--weight-transfer-config",
191
+ json.dumps({"backend": "nccl"}),
192
+ ]
193
+ env = os.environ.copy()
194
+ env["CUDA_VISIBLE_DEVICES"] = server_gpus
195
+ env["VLLM_SERVER_DEV_MODE"] = "1" # exposes the weight-transfer endpoints
196
+ env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
197
+ print(f"[server] Launching: {' '.join(serve_args)} (GPUs {server_gpus})")
198
+ proc = subprocess.Popen(
199
+ serve_args,
200
+ env=env,
201
+ stdout=sys.stdout,
202
+ stderr=sys.stderr,
203
+ start_new_session=True,
204
)
205
178
-
179
-async def generate_batch(engine, prompts, sampling_params):
180
- """Generate completions for a batch of prompts."""
181
-
182
- async def gen_one(prompt):
183
- output = None
184
- async for request_output in engine.generate(
185
- {"prompt": prompt},
186
- sampling_params,
187
- request_id=str(uuid.uuid4()),
188
- ):
189
- output = request_output
190
- return output
191
-
192
- return await asyncio.gather(*[gen_one(p) for p in prompts])
206
+ # Wait for the server to come up (model load can take a while).
207
+ deadline = time.monotonic() + 1800
208
+ while True:
209
+ if proc.poll() is not None:
210
+ raise RuntimeError("vLLM server exited before becoming ready.")
211
+ try:
212
+ if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
213
+ break
214
+ except requests.RequestException:
215
+ pass
216
+ if time.monotonic() > deadline:
217
+ raise RuntimeError("vLLM server failed to start in time.")
218
+ time.sleep(2)
219
+ print("[server] Ready.")
220
+ return proc
221
+
222
+
223
+def generate_completions(client: OpenAI, prompts: list[str]) -> list[str]:
224
+ """Generate completions for a batch of prompts via the OpenAI HTTP API."""
225
+ results = []
226
+ for prompt in prompts:
227
+ response = client.completions.create(
228
+ model=SERVED_MODEL_NAME,
229
+ prompt=prompt,
230
+ max_tokens=32,
231
+ temperature=0,
232
+ )
233
+ results.append(response.choices[0].text)
234
+ return results
235
236
195
-async def main():
196
- ray.init()
197
-
237
+def main():
238
# Download model weights to local/shared disk once.
239
local_model_path = snapshot_download(MODEL_NAME)
240
print(f"[init] Model downloaded to {local_model_path}")
241
202
- # FSDP rendezvous address (single-node)
242
+ ray.init()
243
+
244
+ # FSDP rendezvous address (single-node).
245
fsdp_master_addr = get_ip()
246
fsdp_master_port = get_open_port()
247
206
- # Launch 4 FSDP training workers.
207
- # Ray allocates 1 GPU per worker; AsyncLLMEngine's internal DP
208
- # placement groups will land on the remaining 4 GPUs.
248
+ # Launch the FSDP training workers first so Ray reserves their GPUs, then
249
+ # place the inference server on the GPUs Ray did NOT use. This keeps the two
250
+ # on disjoint physical GPUs whether ray.init() started a fresh cluster or
251
+ # connected to an existing one.
252
fsdp_workers = [
253
FSDPTrainWorker.remote(
254
local_model_path,
@@ -219,126 +262,109 @@ async def main():
262
ray.get([w.get_rank.remote() for w in fsdp_workers])
263
print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.")
264
222
- # Launch vLLM with expert parallelism + data parallelism.
223
- # AsyncLLMEngine with data_parallel_backend="ray" creates its own
224
- # placement groups internally — no manual placement group needed.
225
- print("[engine] Creating AsyncLLMEngine...")
226
- engine = create_async_engine(
227
- model=local_model_path,
228
- enforce_eager=True,
229
- tensor_parallel_size=INFERENCE_TP_SIZE,
230
- data_parallel_size=INFERENCE_DP_SIZE,
231
- enable_expert_parallel=True,
232
- distributed_executor_backend="ray",
233
- data_parallel_backend="ray",
234
- weight_transfer_config=WeightTransferConfig(backend="nccl"),
235
- load_format="dummy",
236
- gpu_memory_utilization=0.7,
237
- )
238
- print("[engine] AsyncLLMEngine created.")
239
-
240
- prompts = [
241
- "Hello, my name is",
242
- "The president of the United States is",
243
- "The capital of France is",
244
- "The future of AI is",
265
+ # Discover the physical GPUs Ray assigned to training; run the server on the
266
+ # complementary GPUs.
267
+ training_gpus = {
268
+ int(g)
269
+ for ids in ray.get([w.get_gpu_ids.remote() for w in fsdp_workers])
270
+ for g in ids
271
+ }
272
+ num_gpus = int(ray.cluster_resources().get("GPU", 0))
273
+ num_server_gpus = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE
274
+ server_gpu_ids = [g for g in range(num_gpus) if g not in training_gpus][
275
+ :num_server_gpus
276
]
246
- sampling_params = SamplingParams(temperature=0)
247
-
248
- # Generate with dummy weights — expect gibberish.
249
- print("[generate] Starting generation with dummy weights...")
250
- outputs = await generate_batch(engine, prompts, sampling_params)
251
- print("[generate] Generation complete.")
252
-
253
- print("-" * 60)
254
- print("BEFORE weight sync (dummy weights):")
255
- print("-" * 60)
256
- for output in outputs:
257
- print(f"Prompt: {output.prompt!r}")
258
- print(f"Generated: {output.outputs[0].text!r}")
259
- print("-" * 60)
277
+ if len(server_gpu_ids) < num_server_gpus:
278
+ raise RuntimeError(
279
+ f"Need {num_server_gpus} free GPUs for the inference server but only "
280
+ f"found {server_gpu_ids} (training uses {sorted(training_gpus)} of "
281
+ f"{num_gpus} cluster GPUs)."
282
+ )
283
+ server_gpus = ",".join(str(g) for g in server_gpu_ids)
284
+ print(f"[init] Training GPUs {sorted(training_gpus)}; server GPUs [{server_gpus}].")
285
+
286
+ # Start the inference server on the complementary GPUs.
287
+ server_proc = start_vllm_server(server_gpus)
288
+ try:
289
+ client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
290
+
291
+ prompts = [
292
+ "Hello, my name is",
293
+ "The president of the United States is",
294
+ "The capital of France is",
295
+ "The future of AI is",
296
+ ]
297
261
- # --- Weight-transfer setup ---
262
- print("[transfer] Setting up weight-transfer endpoint...")
263
- transfer_addr, transfer_port = ray.get(
264
- fsdp_workers[0].setup_transfer_endpoint.remote()
265
- )
266
- print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}")
298
+ # Generate with dummy weights — expect gibberish.
299
+ print("[generate] Generating with dummy weights...")
300
+ outputs = generate_completions(client, prompts)
301
+ print("-" * 60)
302
+ print("BEFORE weight sync (dummy weights):")
303
+ print("-" * 60)
304
+ for prompt, text in zip(prompts, outputs):
305
+ print(f"Prompt: {prompt!r}")
306
+ print(f"Generated: {text!r}")
307
+ print("-" * 60)
308
+
309
+ # --- Weight-transfer setup ---
310
+ print("[transfer] Setting up weight-transfer endpoint...")
311
+ transfer_addr, transfer_port = ray.get(
312
+ fsdp_workers[0].setup_transfer_endpoint.remote()
313
+ )
314
+ print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}")
315
268
- transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1
269
- print(
270
- f"[transfer] World size: {transfer_world_size} "
271
- f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)"
272
- )
316
+ transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1
317
+ print(
318
+ f"[transfer] World size: {transfer_world_size} "
319
+ f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)"
320
+ )
321
274
- print("[transfer] Initializing NCCL groups...")
275
- train_handle = fsdp_workers[0].init_weight_transfer_group.remote(
276
- transfer_world_size
277
- )
278
- await engine.init_weight_transfer_engine(
279
- WeightTransferInitRequest(
280
- init_info=asdict(
281
- NCCLWeightTransferInitInfo(
282
- master_address=transfer_addr,
283
- master_port=transfer_port,
284
- rank_offset=1,
285
- world_size=transfer_world_size,
322
+ # Build the trainer engine on all FSDP ranks (rank 0 is the sender). The
323
+ # sender drives the server's init_weight_transfer_engine (HTTP) while
324
+ # opening the trainer NCCL endpoint, so both ends rendezvous together;
325
+ # the other ranks build a null-client engine that only gathers.
326
+ print("[transfer] Initializing NCCL groups (all FSDP ranks)...")
327
+ ray.get(
328
+ [
329
+ w.setup_engine.remote(
330
+ BASE_URL, transfer_addr, transfer_port, transfer_world_size
331
)
287
- )
332
+ for w in fsdp_workers
333
+ ]
334
)
289
- )
290
- ray.get(train_handle)
291
- print("[transfer] NCCL groups initialized.")
335
+ print("[transfer] NCCL groups initialized.")
336
293
- # --- Pause, transfer weights, resume ---
294
- print("[sync] Pausing generation...")
295
- await engine.pause_generation(mode="abort")
296
- print("[sync] Generation paused.")
337
+ # --- Pause, transfer weights, resume ---
338
+ print("[sync] Pausing generation...")
339
+ requests.post(f"{BASE_URL}/pause", timeout=60).raise_for_status()
340
298
- names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote())
299
- print(f"[sync] Got metadata for {len(names)} parameters.")
341
+ # All ranks participate in the FSDP all-gather; rank 0 additionally
342
+ # drives start/update/finish on the server and the NCCL broadcast.
343
+ print("[sync] Broadcasting weights from FSDP → vLLM...")
344
+ ray.get([w.gather_and_broadcast_weights.remote() for w in fsdp_workers])
345
+ print("[sync] Weight broadcast complete.")
346
301
- print("[sync] Starting weight update...")
302
- await engine.start_weight_update()
347
+ print("[sync] Resuming generation...")
348
+ requests.post(f"{BASE_URL}/resume", timeout=60).raise_for_status()
349
304
- print("[sync] Broadcasting weights from FSDP → vLLM...")
305
- broadcast_handles = [
306
- w.gather_and_broadcast_weights.remote(packed=True) for w in fsdp_workers
307
- ]
308
- await engine.update_weights(
309
- WeightTransferUpdateRequest(
310
- update_info=asdict(
311
- NCCLWeightTransferUpdateInfo(
312
- names=names,
313
- dtype_names=dtype_names,
314
- shapes=shapes,
315
- packed=True,
316
- )
317
- )
318
- )
319
- )
320
- ray.get(broadcast_handles)
321
-
322
- await engine.finish_weight_update()
323
- print("[sync] Weight broadcast complete.")
324
-
325
- print("[sync] Resuming generation...")
326
- await engine.resume_generation()
327
- print("[sync] Generation resumed.")
328
-
329
- # Generate with synced weights — expect sensible output.
330
- print("[generate] Starting generation with synced weights...")
331
- outputs_updated = await generate_batch(engine, prompts, sampling_params)
332
- print("[generate] Generation complete.")
333
-
334
- print("-" * 60)
335
- print("AFTER weight sync (real weights):")
336
- print("-" * 60)
337
- for output in outputs_updated:
338
- print(f"Prompt: {output.prompt!r}")
339
- print(f"Generated: {output.outputs[0].text!r}")
350
+ # Generate with synced weights — expect sensible output.
351
+ print("[generate] Generating with synced weights...")
352
+ outputs_updated = generate_completions(client, prompts)
353
+ print("-" * 60)
354
+ print("AFTER weight sync (real weights):")
355
print("-" * 60)
356
+ for prompt, text in zip(prompts, outputs_updated):
357
+ print(f"Prompt: {prompt!r}")
358
+ print(f"Generated: {text!r}")
359
+ print("-" * 60)
360
+ finally:
361
+ print("[server] Shutting down...")
362
+ server_proc.terminate()
363
+ try:
364
+ server_proc.wait(timeout=30)
365
+ except subprocess.TimeoutExpired:
366
+ server_proc.kill()
367
368
369
if __name__ == "__main__":
344
- asyncio.run(main())
370
+ main()
examples/rl/rlhf_sparse_nccl.py
+72
-116
@@ -44,12 +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.nccl_engine import (
48
- NCCLTrainerSendWeightsArgs,
49
- NCCLWeightTransferEngine,
47
+from vllm.distributed.weight_transfer import (
48
+ ModuleSource,
49
+ RayVLLMWeightSyncClient,
50
+ WeightTransferTrainerFactory,
51
)
52
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
53
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
52
- SparseNCCLWeightTransferEngine,
54
+ SparseNCCLTrainerInitInfo,
55
SparseWeightPatch,
56
)
57
from vllm.utils.network_utils import get_ip, get_open_port
@@ -87,7 +89,8 @@ class TrainModel:
89
self.model = None
90
self.patched_param = None
91
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
90
- self.model_update_group = None
92
+ self.dense_engine = None
93
+ self.sparse_engine = None
94
self.master_address = get_ip()
95
self.port = get_open_port()
96
self.reset_model()
@@ -112,34 +115,32 @@ class TrainModel:
115
self.port = get_open_port()
116
return self.master_address, self.port
117
115
- def init_weight_transfer_group(self, world_size: int) -> None:
116
- self.model_update_group = NCCLWeightTransferEngine.trainer_init(
117
- dict(
118
+ def init_dense_engine(self, world_size: int, llm_handle) -> None:
119
+ """Build the dense NCCL trainer engine"""
120
+ self.dense_engine = WeightTransferTrainerFactory.trainer_init(
121
+ init_info=NCCLTrainerInitInfo(
122
master_address=self.master_address,
123
master_port=self.port,
124
world_size=world_size,
121
- )
125
+ rank=0, # single-GPU trainer is the sole (sender) rank
126
+ packed=True,
127
+ ),
128
+ client=RayVLLMWeightSyncClient(llm_handle),
129
+ source=ModuleSource(self.model),
130
)
131
124
- def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
125
- names = []
126
- dtype_names = []
127
- shapes = []
128
- payload_bytes = 0
129
- for name, param in self.model.named_parameters():
130
- names.append(name)
131
- dtype_names.append(str(param.dtype).split(".")[-1])
132
- shapes.append(list(param.shape))
133
- payload_bytes += param.numel() * param.element_size()
134
-
135
- return (
136
- dict(
137
- names=names,
138
- dtype_names=dtype_names,
139
- shapes=shapes,
140
- packed=packed,
141
- ),
142
- payload_bytes,
132
+ def shutdown_engines(self) -> None:
133
+ """Drop the NCCL communicators before the vLLM engine they point at
134
+ goes away, so the next phase rendezvouses from a clean slate."""
135
+ for attr in ("dense_engine", "sparse_engine"):
136
+ engine = getattr(self, attr)
137
+ if engine is not None:
138
+ engine.shutdown()
139
+ setattr(self, attr, None)
140
+
141
+ def dense_payload_bytes(self) -> int:
142
+ return sum(
143
+ param.numel() * param.element_size() for param in self.model.parameters()
144
)
145
146
@torch.inference_mode()
@@ -173,7 +174,7 @@ class TrainModel:
174
self,
175
prompts: Sequence[str],
176
max_patch_rows: int = MAX_PATCH_ROWS,
176
- ) -> tuple[dict[str, object], list[int], str, int]:
177
+ ) -> tuple[list[int], str, int]:
178
selected_token_ids: list[int] = []
179
special_ids = set(self.tokenizer.all_special_ids)
180
for prompt in prompts:
@@ -220,63 +221,60 @@ class TrainModel:
221
flat_indices = (
222
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
223
)
223
- flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
224
+ flat_values = self.patched_param[row_ids].detach().reshape(-1).contiguous()
225
self.pending_sparse_patches = [
226
SparseWeightPatch(
227
name=PATCHED_PARAM_NAME,
228
indices=flat_indices.to(torch.int32),
229
values=flat_values,
230
+ full_shape=tuple(self.patched_param.shape),
231
)
232
]
233
patch_digest = hashlib.sha256(
234
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
233
- + self.pending_sparse_patches[0]
234
- .values.detach()
235
- .float()
236
- .cpu()
237
- .numpy()
238
- .tobytes()
235
+ + self.pending_sparse_patches[0].values.float().cpu().numpy().tobytes()
236
).hexdigest()
237
238
sparse_payload_bytes = (
239
flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
240
+ flat_values.numel() * flat_values.element_size()
241
)
245
- update_info = dict(
246
- names=[PATCHED_PARAM_NAME],
247
- dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
248
- shapes=[list(self.patched_param.shape)],
249
- num_updates_list=[flat_indices.numel()],
250
- )
251
- return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
242
+ return selected_token_ids, patch_digest, sparse_payload_bytes
243
253
- def broadcast_weights(self, packed: bool = False) -> float:
254
- if self.model_update_group is None:
255
- raise RuntimeError("Weight transfer group is not initialized")
244
+ def send_dense_weights(self) -> float:
245
+ """One call drives start/update/finish on the inference side
246
+ concurrently with the NCCL broadcast."""
247
+ if self.dense_engine is None:
248
+ raise RuntimeError("Dense engine is not initialized")
249
257
- trainer_args = NCCLTrainerSendWeightsArgs(
258
- group=self.model_update_group,
259
- packed=packed,
260
- )
250
start = time.perf_counter()
262
- NCCLWeightTransferEngine.trainer_send_weights(
263
- iterator=self.model.named_parameters(),
264
- trainer_args=trainer_args,
265
- )
251
+ self.dense_engine.send_weights()
252
torch.accelerator.synchronize()
253
return (time.perf_counter() - start) * 1000.0
254
269
- def broadcast_pending_sparse_patch(self) -> float:
270
- if self.model_update_group is None:
271
- raise RuntimeError("Weight transfer group is not initialized")
255
+ def init_sparse_engine(self, world_size: int, llm_handle) -> None:
256
+ """Build the sparse trainer engine"""
257
+ self.sparse_engine = WeightTransferTrainerFactory.trainer_init(
258
+ init_info=SparseNCCLTrainerInitInfo(
259
+ master_address=self.master_address,
260
+ master_port=self.port,
261
+ world_size=world_size,
262
+ rank=0, # single-GPU trainer is the sole (sender) rank
263
+ ),
264
+ client=RayVLLMWeightSyncClient(llm_handle),
265
+ )
266
+
267
+ def send_pending_sparse_patch(self) -> float:
268
+ """One call drives start/update/finish on the inference side
269
+ concurrently with the patch broadcasts. Sparse deltas differ each round,
270
+ so the patches are passed to `send_weights` rather than fixed at init."""
271
+ if self.sparse_engine is None:
272
+ raise RuntimeError("Sparse engine is not initialized")
273
if self.pending_sparse_patches is None:
274
raise RuntimeError("Sparse patch has not been prepared")
275
276
start = time.perf_counter()
276
- SparseNCCLWeightTransferEngine.trainer_send_weights(
277
- iter(self.pending_sparse_patches),
278
- NCCLTrainerSendWeightsArgs(group=self.model_update_group),
279
- )
277
+ self.sparse_engine.send_weights(self.pending_sparse_patches)
278
torch.accelerator.synchronize()
279
self.pending_sparse_patches = None
280
return (time.perf_counter() - start) * 1000.0
@@ -340,39 +338,17 @@ def run_dense_phase(
338
dense_before = collect_vllm_generations(llm)
339
340
ray.get(llm.sleep.remote(level=0))
343
- master_address, master_port = ray.get(train_model.create_rendezvous.remote())
341
+ ray.get(train_model.create_rendezvous.remote())
342
world_size = ray.get(llm.get_world_size.remote()) + 1
345
- inference_init = llm.init_weight_transfer_engine.remote(
346
- dict(
347
- init_info=dict(
348
- master_address=master_address,
349
- master_port=master_port,
350
- rank_offset=1,
351
- world_size=world_size,
352
- )
353
- )
354
- )
355
- trainer_init = train_model.init_weight_transfer_group.remote(world_size)
356
- ray.get([trainer_init, inference_init])
357
- ray.get(llm.start_weight_update.remote())
343
+ ray.get(train_model.init_dense_engine.remote(world_size, llm))
344
359
- dense_update_info, dense_payload_bytes = ray.get(
360
- train_model.get_dense_update_info.remote()
361
- )
362
- _, selected_token_ids, patch_digest, _ = ray.get(
345
+ dense_payload_bytes = ray.get(train_model.dense_payload_bytes.remote())
346
+ selected_token_ids, patch_digest, _ = ray.get(
347
train_model.prepare_sparse_patch.remote(PROMPTS)
348
)
349
366
- inference_update = llm.update_weights.remote(
367
- dict(update_info=dense_update_info)
368
- )
369
- dense_send_ms, _ = ray.get(
370
- [
371
- train_model.broadcast_weights.remote(packed=False),
372
- inference_update,
373
- ]
374
- )
375
- ray.get(llm.finish_weight_update.remote())
350
+ # One call drives start/update/finish + the NCCL broadcast.
351
+ dense_send_ms = ray.get(train_model.send_dense_weights.remote())
352
ray.get(llm.wake_up.remote(tags=["scheduling"]))
353
354
dense_after = collect_vllm_generations(llm)
@@ -386,6 +362,7 @@ def run_dense_phase(
362
"dense_send_ms": dense_send_ms,
363
}
364
finally:
365
+ ray.get(train_model.shutdown_engines.remote())
366
ray.kill(llm)
367
368
@@ -399,36 +376,14 @@ def run_sparse_phase(
376
sparse_before = collect_vllm_generations(llm)
377
378
ray.get(llm.sleep.remote(level=0))
402
- master_address, master_port = ray.get(train_model.create_rendezvous.remote())
379
+ ray.get(train_model.create_rendezvous.remote())
380
world_size = ray.get(llm.get_world_size.remote()) + 1
404
- inference_init = llm.init_weight_transfer_engine.remote(
405
- dict(
406
- init_info=dict(
407
- master_address=master_address,
408
- master_port=master_port,
409
- rank_offset=1,
410
- world_size=world_size,
411
- )
412
- )
413
- )
414
- trainer_init = train_model.init_weight_transfer_group.remote(world_size)
415
- ray.get([trainer_init, inference_init])
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))
420
- )
381
+ ray.get(train_model.init_sparse_engine.remote(world_size, llm))
382
422
- inference_update = llm.update_weights.remote(
423
- dict(update_info=sparse_update_info)
424
- )
425
- sparse_send_ms, _ = ray.get(
426
- [
427
- train_model.broadcast_pending_sparse_patch.remote(),
428
- inference_update,
429
- ]
383
+ selected_token_ids, patch_digest, sparse_payload_bytes = ray.get(
384
+ train_model.prepare_sparse_patch.remote(PROMPTS)
385
)
431
- ray.get(llm.finish_weight_update.remote())
386
+ sparse_send_ms = ray.get(train_model.send_pending_sparse_patch.remote())
387
ray.get(llm.wake_up.remote(tags=["scheduling"]))
388
389
sparse_after = collect_vllm_generations(llm)
@@ -442,6 +397,7 @@ def run_sparse_phase(
397
"sparse_send_ms": sparse_send_ms,
398
}
399
finally:
400
+ ray.get(train_model.shutdown_engines.remote())
401
ray.kill(llm)
402
403
tests/distributed/test_packed_tensor.py
-27
@@ -9,7 +9,6 @@ These utilities enable efficient batched tensor transfer over NCCL.
9
import pytest
10
import torch
11
12
-from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo
12
from vllm.distributed.weight_transfer.packed_tensor import (
13
pack_tensors,
14
packed_ipc_consumer,
@@ -68,32 +67,6 @@ def create_state_dict_info(
67
return {name: (tuple(tensor.shape), tensor.dtype) for name, tensor in params}
68
69
71
-# --- Unit Tests: NCCLWeightTransferUpdateInfo packed field ---
72
-
73
-
74
-class TestNCCLWeightTransferUpdateInfoPacked:
75
- """Test NCCLWeightTransferUpdateInfo dataclass packed field."""
76
-
77
- def test_packed_default_false(self):
78
- """Test that packed defaults to False."""
79
- info = NCCLWeightTransferUpdateInfo(
80
- names=["layer.weight"],
81
- dtype_names=["float32"],
82
- shapes=[[10, 10]],
83
- )
84
- assert info.packed is False
85
-
86
- def test_packed_can_be_set_true(self):
87
- """Test that packed can be set to True."""
88
- info = NCCLWeightTransferUpdateInfo(
89
- names=["layer.weight"],
90
- dtype_names=["float32"],
91
- shapes=[[10, 10]],
92
- packed=True,
93
- )
94
- assert info.packed is True
95
-
96
-
70
# --- Unit Tests: packed_nccl_broadcast_producer ---
71
72
tests/distributed/test_weight_transfer.py
+473
-20
@@ -7,6 +7,8 @@ Integration tests for NCCL and IPC weight transfer between processes using Ray.
7
"""
8
9
import pickle
10
+import threading
11
+import time
12
from unittest.mock import MagicMock
13
14
import pybase64 as base64
@@ -20,9 +22,11 @@ from vllm.config.weight_transfer import WeightTransferConfig
22
from vllm.distributed.weight_transfer import (
23
HTTPVLLMWeightSyncClient,
24
ModuleSource,
25
+ ParamMeta,
26
RayVLLMWeightSyncClient,
27
TrainerWeightTransferEngine,
28
VLLMWeightSyncClient,
29
+ WeightSource,
30
WeightTransferEngineFactory,
31
WeightTransferTrainerFactory,
32
)
@@ -39,11 +43,19 @@ from vllm.distributed.weight_transfer.ipc_engine import (
43
IPCWeightTransferUpdateInfo,
44
)
45
from vllm.distributed.weight_transfer.nccl_engine import (
46
+ NCCLTrainerInitInfo,
47
+ NCCLTrainerWeightTransferEngine,
48
NCCLWeightTransferEngine,
49
NCCLWeightTransferInitInfo,
50
NCCLWeightTransferUpdateInfo,
51
)
52
+from vllm.distributed.weight_transfer.packed_tensor import (
53
+ DEFAULT_PACKED_BUFFER_SIZE_BYTES,
54
+ DEFAULT_PACKED_NUM_BUFFERS,
55
+)
56
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
57
+ SparseNCCLTrainerInitInfo,
58
+ SparseNCCLTrainerWeightTransferEngine,
59
SparseNCCLWeightTransferEngine,
60
SparseNCCLWeightTransferUpdateInfo,
61
SparseWeightPatch,
@@ -535,11 +547,15 @@ def inference_receive_tensor(
547
_vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
548
549
# Initialize the engine (joins as rank 1)
550
+ # Trainer broadcasts a single tensor unpacked, so the worker must not
551
+ # expect the packed wire format (packed is a must-agree wire param shipped
552
+ # on the init info).
553
init_info = NCCLWeightTransferInitInfo(
554
master_address=master_address,
555
master_port=master_port,
556
rank_offset=1, # Trainer is rank 0, we become rank 1
557
world_size=world_size,
558
+ packed=False,
559
)
560
engine.init_transfer_engine(init_info)
561
@@ -616,39 +632,53 @@ def trainer_broadcast_sparse_tensor(
632
master_port: int,
633
world_size: int,
634
) -> bool:
619
- """Trainer task that broadcasts sparse patches via NCCL."""
635
+ """Trainer task that broadcasts sparse patches via the trainer engine.
636
+
637
+ The worker task drives its own init/receive directly (it is not an RPC
638
+ endpoint), so the engine gets a no-op control-plane client; the NCCL
639
+ rendezvous and the patch broadcasts are the real thing.
640
+ """
641
import torch
642
643
device = _set_ray_assigned_device()
644
624
- from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
625
- from vllm.distributed.utils import StatelessProcessGroup
626
- from vllm.distributed.weight_transfer.nccl_engine import (
627
- NCCLTrainerSendWeightsArgs,
628
- )
645
+ from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
646
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
630
- SparseNCCLWeightTransferEngine,
647
+ SparseNCCLTrainerInitInfo,
648
SparseWeightPatch,
649
)
650
634
- pg = StatelessProcessGroup.create(
635
- host=master_address,
636
- port=master_port,
637
- rank=0,
638
- world_size=world_size,
639
- )
640
- comm = PyNcclCommunicator(pg, device=device.index)
651
+ class NoopClient:
652
+ def init_weight_transfer_engine(self, init_info):
653
+ pass
654
+
655
+ def start_weight_update(self):
656
+ pass
657
+
658
+ def update_weights(self, update_info):
659
+ pass
660
+
661
+ def finish_weight_update(self):
662
+ pass
663
664
patch = SparseWeightPatch(
665
name="test.weight",
666
indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device),
667
values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device),
668
+ full_shape=(10, 10),
669
)
647
- SparseNCCLWeightTransferEngine.trainer_send_weights(
648
- iter([patch]),
649
- NCCLTrainerSendWeightsArgs(group=comm),
670
+ engine = WeightTransferTrainerFactory.trainer_init(
671
+ init_info=SparseNCCLTrainerInitInfo(
672
+ master_address=master_address,
673
+ master_port=master_port,
674
+ world_size=world_size,
675
+ rank=0,
676
+ ),
677
+ client=NoopClient(),
678
)
679
+ engine.send_weights([patch])
680
torch.accelerator.synchronize()
681
+ engine.shutdown()
682
return True
683
684
@@ -1369,14 +1399,26 @@ class TestModuleSource:
1399
source = ModuleSource(_module_with(("w", torch.zeros(2))))
1400
assert [n for n, _ in source] == [n for n, _ in source] == ["w"]
1401
1402
+ def test_metadata_agrees_with_iteration(self):
1403
+ """The two channels must line up element-for-element: engines declare
1404
+ the round from `metadata()` and then send what iteration yields."""
1405
+ source = ModuleSource(
1406
+ _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3)))
1407
+ )
1408
+ meta = source.metadata()
1409
+ pairs = list(source)
1410
+ assert [m.name for m in meta] == [name for name, _ in pairs]
1411
+ assert [m.dtype for m in meta] == [t.dtype for _, t in pairs]
1412
+ assert [m.shape for m in meta] == [tuple(t.shape) for _, t in pairs]
1413
+
1414
1415
class TestTrainerFactory:
1416
"""WeightTransferTrainerFactory registry mechanics."""
1417
1376
- def test_registry_has_ipc(self):
1377
- # IPC is the first backend migrated to the trainer engine; NCCL /
1378
- # sparse NCCL register in later PRs.
1418
+ def test_registry_has_all_backends(self):
1419
+ assert "nccl" in WeightTransferTrainerFactory._registry
1420
assert "ipc" in WeightTransferTrainerFactory._registry
1421
+ assert "sparse_nccl" in WeightTransferTrainerFactory._registry
1422
1423
def test_register_and_dispatch(self):
1424
saved = dict(WeightTransferTrainerFactory._registry)
@@ -1406,6 +1448,12 @@ class TestTrainerFactory:
1448
def test_ipc_init_info_declares_backend(self):
1449
assert IPCTrainerInitInfo.backend == "ipc"
1450
1451
+ def test_nccl_init_info_declares_backend(self):
1452
+ assert NCCLTrainerInitInfo.backend == "nccl"
1453
+
1454
+ def test_sparse_nccl_init_info_declares_backend(self):
1455
+ assert SparseNCCLTrainerInitInfo.backend == "sparse_nccl"
1456
+
1457
def test_trainer_init_info_subclass_must_set_backend(self):
1458
with pytest.raises(TypeError, match="class-level `backend`"):
1459
@@ -1475,3 +1523,408 @@ def test_ipc_trainer_init_ships_packed_to_worker():
1523
assert engine.packed is True
1524
assert client.order == ["init"]
1525
assert client.last_init_info == {"packed": True}
1526
+
1527
+
1528
+def test_nccl_trainer_init_ships_worker_init_info(monkeypatch):
1529
+ """The sender's trainer_init drives the inference-side init handshake with
1530
+ the worker-shaped init info (rank_offset=1) while opening its own endpoint,
1531
+ and propagates the must-agree wire params to the worker."""
1532
+ import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
1533
+
1534
+ # Bypass the real NCCL rendezvous.
1535
+ monkeypatch.setattr(
1536
+ nccl_engine_mod, "open_trainer_endpoint", lambda info: MagicMock()
1537
+ )
1538
+
1539
+ client = RecordingClient()
1540
+ engine = WeightTransferTrainerFactory.trainer_init(
1541
+ init_info=NCCLTrainerInitInfo(
1542
+ master_address="127.0.0.1",
1543
+ master_port=29500,
1544
+ world_size=3,
1545
+ rank=0,
1546
+ packed=True,
1547
+ packed_buffer_size_bytes=1024,
1548
+ packed_num_buffers=3,
1549
+ ),
1550
+ client=client,
1551
+ source=ModuleSource(_module_with(("w", torch.zeros(4)))),
1552
+ )
1553
+
1554
+ assert isinstance(engine, NCCLTrainerWeightTransferEngine)
1555
+ assert engine.is_sender is True
1556
+ assert engine.packed is True
1557
+ assert client.order == ["init"]
1558
+ assert client.last_init_info == {
1559
+ "master_address": "127.0.0.1",
1560
+ "master_port": 29500,
1561
+ "rank_offset": 1,
1562
+ "world_size": 3,
1563
+ "packed": True,
1564
+ "packed_buffer_size_bytes": 1024,
1565
+ "packed_num_buffers": 3,
1566
+ }
1567
+
1568
+
1569
+def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch):
1570
+ """The worker engine reads packed + buffer geometry from the
1571
+ trainer-supplied init info at the handshake, not from the config or the
1572
+ per-round update info."""
1573
+ import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
1574
+
1575
+ monkeypatch.setattr(
1576
+ nccl_engine_mod, "worker_init_process_group", lambda info, pc: MagicMock()
1577
+ )
1578
+
1579
+ engine = NCCLWeightTransferEngine(
1580
+ WeightTransferConfig(backend="nccl"),
1581
+ create_mock_vllm_config(),
1582
+ torch.device("cuda:0"),
1583
+ MagicMock(spec=torch.nn.Module),
1584
+ )
1585
+ assert engine.packed is False # pre-handshake default (legacy unpacked)
1586
+ engine.init_transfer_engine(
1587
+ NCCLWeightTransferInitInfo(
1588
+ master_address="127.0.0.1",
1589
+ master_port=29500,
1590
+ rank_offset=1,
1591
+ world_size=2,
1592
+ packed=True,
1593
+ packed_buffer_size_bytes=2048,
1594
+ packed_num_buffers=4,
1595
+ )
1596
+ )
1597
+
1598
+ assert engine.packed is True
1599
+ assert engine.packed_buffer_size_bytes == 2048
1600
+ assert engine.packed_num_buffers == 4
1601
+
1602
+
1603
+def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client():
1604
+ """Non-sender trainer ranks build an engine without opening an endpoint or
1605
+ touching the client; they only join the collectives in send_weights."""
1606
+ client = RecordingClient()
1607
+ engine = WeightTransferTrainerFactory.trainer_init(
1608
+ init_info=NCCLTrainerInitInfo(
1609
+ master_address="127.0.0.1",
1610
+ master_port=29500,
1611
+ world_size=3,
1612
+ rank=1,
1613
+ ),
1614
+ client=client,
1615
+ source=ModuleSource(_module_with(("w", torch.zeros(4)))),
1616
+ )
1617
+
1618
+ assert engine.is_sender is False
1619
+ assert engine.model_update_group is None
1620
+ assert client.order == []
1621
+
1622
+ # send_weights on a non-sender only iterates the source (packed mode needs
1623
+ # no CUDA stream on non-senders), never the client.
1624
+ engine.send_weights()
1625
+ assert client.order == []
1626
+
1627
+
1628
+@pytest.mark.skipif(
1629
+ torch.accelerator.device_count() < 1,
1630
+ reason="Need at least 1 GPU (NCCL broadcast / CUDA stream).",
1631
+)
1632
+def test_nccl_trainer_send_weights_drives_client_in_order():
1633
+ """send_weights issues start -> update -> finish and ships per-round
1634
+ metadata; the packed wire params ride the init handshake, not the
1635
+ per-round update_info."""
1636
+ client = RecordingClient()
1637
+ engine = NCCLTrainerWeightTransferEngine(
1638
+ client=client,
1639
+ source=ModuleSource(_module_with(("w", torch.zeros(4, device="cuda")))),
1640
+ packed=False,
1641
+ )
1642
+ # Bypass the real NCCL rendezvous; broadcast is a no-op.
1643
+ engine.model_update_group = MagicMock()
1644
+
1645
+ engine.send_weights()
1646
+
1647
+ assert client.order == ["start", "update", "finish"]
1648
+ assert client.last_update_info is not None
1649
+ assert client.last_update_info["names"] == ["w"]
1650
+ assert client.last_update_info["shapes"] == [[4]]
1651
+ assert "packed" not in client.last_update_info
1652
+
1653
+
1654
+class _ScriptedSource(WeightSource):
1655
+ """Declares `meta` but yields whatever `pairs` says — used to drive the
1656
+ metadata/iteration agreement checks."""
1657
+
1658
+ def __init__(self, meta, pairs):
1659
+ self._meta = meta
1660
+ self._pairs = pairs
1661
+
1662
+ def metadata(self):
1663
+ return list(self._meta)
1664
+
1665
+ def __iter__(self):
1666
+ yield from self._pairs
1667
+
1668
+
1669
+def _mock_group_engine(source, monkeypatch, **kwargs):
1670
+ """Unpacked trainer engine with a mocked group and stream (no GPU needed)."""
1671
+ engine = NCCLTrainerWeightTransferEngine(
1672
+ client=RecordingClient(), source=source, packed=False, **kwargs
1673
+ )
1674
+ engine.model_update_group = MagicMock()
1675
+ monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
1676
+ return engine
1677
+
1678
+
1679
+def test_nccl_trainer_init_requires_source():
1680
+ """NCCL is a full-resync backend: it cannot run without a WeightSource."""
1681
+ with pytest.raises(ValueError, match="requires a WeightSource"):
1682
+ NCCLTrainerWeightTransferEngine.trainer_init(
1683
+ NCCLTrainerInitInfo(
1684
+ master_address="127.0.0.1", master_port=29500, world_size=2, rank=0
1685
+ ),
1686
+ client=RecordingClient(),
1687
+ )
1688
+
1689
+
1690
+def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch):
1691
+ """The worker sizes its buffers (and cuts packed chunks) from metadata(), so
1692
+ iteration disagreeing with it must raise rather than corrupt the stream."""
1693
+ meta = [
1694
+ ParamMeta("w", torch.float32, (4,)),
1695
+ ParamMeta("b", torch.float32, (2,)),
1696
+ ]
1697
+ reordered = [("b", torch.zeros(2)), ("w", torch.zeros(4))]
1698
+ engine = _mock_group_engine(_ScriptedSource(meta, reordered), monkeypatch)
1699
+
1700
+ with pytest.raises(ValueError, match="disagrees with iteration at index 0"):
1701
+ engine.send_weights()
1702
+
1703
+
1704
+def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch):
1705
+ """A source that declares one wire dtype and materializes another would make
1706
+ the two sides disagree on every byte offset."""
1707
+ meta = [ParamMeta("w", torch.float32, (4,))]
1708
+ engine = _mock_group_engine(
1709
+ _ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]),
1710
+ monkeypatch,
1711
+ )
1712
+
1713
+ with pytest.raises(ValueError, match="disagrees with iteration"):
1714
+ engine.send_weights()
1715
+
1716
+
1717
+def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch):
1718
+ """Yielding fewer parameters than declared leaves the worker waiting."""
1719
+ meta = [
1720
+ ParamMeta("w", torch.float32, (4,)),
1721
+ ParamMeta("b", torch.float32, (2,)),
1722
+ ]
1723
+ engine = _mock_group_engine(
1724
+ _ScriptedSource(meta, [("w", torch.zeros(4))]), monkeypatch
1725
+ )
1726
+
1727
+ with pytest.raises(ValueError, match="yielded 1 parameters"):
1728
+ engine.send_weights()
1729
+
1730
+
1731
+def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch):
1732
+ """NCCL sends numel elements from data_ptr(), so a non-contiguous view must
1733
+ be linearized first or the worker receives unrelated memory."""
1734
+ base = torch.arange(6, dtype=torch.float32).reshape(2, 3)
1735
+ view = base.t() # non-contiguous
1736
+ meta = [ParamMeta("w", torch.float32, tuple(view.shape))]
1737
+ engine = _mock_group_engine(_ScriptedSource(meta, [("w", view)]), monkeypatch)
1738
+
1739
+ engine.send_weights()
1740
+
1741
+ sent = engine.model_update_group.broadcast.call_args.args[0]
1742
+ assert sent.is_contiguous()
1743
+ assert torch.equal(sent, view)
1744
+
1745
+
1746
+def test_nccl_trainer_send_weights_raises_instead_of_hanging(monkeypatch):
1747
+ """A failed broadcast must surface even while the inference-side
1748
+ update_weights is still blocked in its matching NCCL call.
1749
+
1750
+ Joining the RPC thread there would deadlock: the worker only returns once
1751
+ the broadcast it is waiting for arrives, which never happens.
1752
+ """
1753
+ rpc_entered = threading.Event()
1754
+ release_rpc = threading.Event()
1755
+
1756
+ class _BlockingClient(RecordingClient):
1757
+ def update_weights(self, update_info):
1758
+ rpc_entered.set()
1759
+ release_rpc.wait(timeout=30) # stands in for a wedged NCCL recv
1760
+ super().update_weights(update_info)
1761
+
1762
+ class _FailingSource(WeightSource):
1763
+ def metadata(self):
1764
+ return [ParamMeta("w", torch.float32, (4,))]
1765
+
1766
+ def __iter__(self):
1767
+ # Fail only once the RPC is provably in flight.
1768
+ rpc_entered.wait(timeout=60)
1769
+ raise RuntimeError("broadcast blew up")
1770
+
1771
+ engine = NCCLTrainerWeightTransferEngine(
1772
+ client=_BlockingClient(), source=_FailingSource(), packed=False
1773
+ )
1774
+ engine.model_update_group = MagicMock()
1775
+ monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
1776
+
1777
+ started = time.perf_counter()
1778
+ try:
1779
+ with pytest.raises(RuntimeError, match="broadcast blew up"):
1780
+ engine.send_weights()
1781
+ elapsed = time.perf_counter() - started
1782
+ assert rpc_entered.is_set(), "the RPC was never in flight"
1783
+ assert not release_rpc.is_set(), "send_weights waited for the wedged RPC"
1784
+ # Regression guard: joining the RPC thread would park here until the
1785
+ # client's own timeout expires instead of raising immediately.
1786
+ assert elapsed < 5.0, f"send_weights blocked {elapsed:.1f}s on the RPC"
1787
+ finally:
1788
+ release_rpc.set()
1789
+
1790
+
1791
+def _sparse_patch(device: str = "cpu") -> SparseWeightPatch:
1792
+ return SparseWeightPatch(
1793
+ name="w",
1794
+ indices=torch.tensor([1, 3], dtype=torch.int32, device=device),
1795
+ values=torch.tensor([1.0, 2.0], dtype=torch.float32, device=device),
1796
+ full_shape=(4, 4),
1797
+ )
1798
+
1799
+
1800
+def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch):
1801
+ """The sender's trainer_init drives the init handshake with the
1802
+ worker-shaped init info; sparse ships no packed wire params, so the worker
1803
+ keeps its unpacked defaults. Sparse takes no `source`."""
1804
+ import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod
1805
+
1806
+ monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock())
1807
+
1808
+ client = RecordingClient()
1809
+ engine = WeightTransferTrainerFactory.trainer_init(
1810
+ init_info=SparseNCCLTrainerInitInfo(
1811
+ master_address="127.0.0.1",
1812
+ master_port=29500,
1813
+ world_size=2,
1814
+ rank=0,
1815
+ ),
1816
+ client=client,
1817
+ )
1818
+
1819
+ assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine)
1820
+ assert client.order == ["init"]
1821
+ assert client.last_init_info == {
1822
+ "master_address": "127.0.0.1",
1823
+ "master_port": 29500,
1824
+ "rank_offset": 1,
1825
+ "world_size": 2,
1826
+ "packed": False,
1827
+ "packed_buffer_size_bytes": DEFAULT_PACKED_BUFFER_SIZE_BYTES,
1828
+ "packed_num_buffers": DEFAULT_PACKED_NUM_BUFFERS,
1829
+ }
1830
+
1831
+
1832
+def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
1833
+ """send_weights takes the round's patches and ships per-patch metadata
1834
+ (names / shapes / num_updates_list) + broadcasts indices + values each."""
1835
+ client = RecordingClient()
1836
+ engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1837
+ engine.model_update_group = MagicMock()
1838
+ # The group is a mock, so the stream is just a handle it is handed (and a
1839
+ # handle _post_send_sync can synchronize).
1840
+ monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
1841
+
1842
+ engine.send_weights([_sparse_patch()])
1843
+
1844
+ assert client.order == ["start", "update", "finish"]
1845
+ assert client.last_update_info is not None
1846
+ assert client.last_update_info["names"] == ["w"]
1847
+ assert client.last_update_info["shapes"] == [[4, 4]]
1848
+ assert client.last_update_info["num_updates_list"] == [2]
1849
+ # One broadcast for indices + one for values per patch.
1850
+ assert engine.model_update_group.broadcast.call_count == 2
1851
+
1852
+
1853
+def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
1854
+ """A round with no patches must not touch the client (an empty sparse
1855
+ update info is invalid by construction)."""
1856
+ client = RecordingClient()
1857
+ engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1858
+ engine.model_update_group = MagicMock()
1859
+
1860
+ engine.send_weights([])
1861
+ engine.send_weights() # no argument is also a no-op round
1862
+
1863
+ assert client.order == []
1864
+
1865
+
1866
+def test_sparse_nccl_trainer_send_weights_requires_full_shape():
1867
+ patch = _sparse_patch()
1868
+ patch.full_shape = None
1869
+ engine = SparseNCCLTrainerWeightTransferEngine(client=RecordingClient())
1870
+ engine.model_update_group = MagicMock()
1871
+
1872
+ with pytest.raises(ValueError, match="full_shape"):
1873
+ engine.send_weights([patch])
1874
+
1875
+
1876
+def test_sparse_nccl_trainer_rejects_source():
1877
+ """Sparse is a delta backend; a WeightSource would silently never be sent."""
1878
+ with pytest.raises(ValueError, match="takes no WeightSource"):
1879
+ SparseNCCLTrainerWeightTransferEngine(
1880
+ client=RecordingClient(),
1881
+ source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1882
+ )
1883
+
1884
+
1885
+def test_sparse_nccl_trainer_validates_patch_before_any_rpc():
1886
+ """Malformed patches must fail on the trainer, before start_weight_update:
1887
+ the worker's own checks only run once the broadcasts are already in flight,
1888
+ where a size mismatch wedges both sides instead of raising."""
1889
+ client = RecordingClient()
1890
+ engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1891
+ engine.model_update_group = MagicMock()
1892
+
1893
+ mismatched = SparseWeightPatch(
1894
+ name="w",
1895
+ indices=torch.tensor([1, 3], dtype=torch.int32),
1896
+ values=torch.tensor([1.0], dtype=torch.float32),
1897
+ full_shape=(4, 4),
1898
+ )
1899
+ with pytest.raises(ValueError, match="matching lengths"):
1900
+ engine.send_weights([mismatched])
1901
+
1902
+ wrong_index_dtype = SparseWeightPatch(
1903
+ name="w",
1904
+ indices=torch.tensor([1, 3], dtype=torch.int64),
1905
+ values=torch.tensor([1.0, 2.0], dtype=torch.float32),
1906
+ full_shape=(4, 4),
1907
+ )
1908
+ with pytest.raises(ValueError, match="int32 indices"):
1909
+ engine.send_weights([wrong_index_dtype])
1910
+
1911
+ assert client.order == []
1912
+
1913
+
1914
+def test_sparse_nccl_trainer_non_sender_skips_client():
1915
+ client = RecordingClient()
1916
+ engine = WeightTransferTrainerFactory.trainer_init(
1917
+ init_info=SparseNCCLTrainerInitInfo(
1918
+ master_address="127.0.0.1",
1919
+ master_port=29500,
1920
+ world_size=2,
1921
+ rank=1,
1922
+ ),
1923
+ client=client,
1924
+ )
1925
+
1926
+ assert engine.is_sender is False
1927
+ assert engine.model_update_group is None
1928
+ assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine)
1929
+ engine.send_weights([_sparse_patch()])
1930
+ assert client.order == []
vllm/distributed/weight_transfer/base.py
+10
-10
@@ -77,6 +77,7 @@ class WeightSource(ABC):
77
78
@abstractmethod
79
def metadata(self) -> list[ParamMeta]:
80
+ """Declare what iteration will yield, without transferring anything."""
81
raise NotImplementedError
82
83
@abstractmethod
@@ -402,8 +403,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
403
Symmetric to `WeightTransferEngine` but lives in the training process.
404
Constructed via the `trainer_init` factory classmethod; carries any
405
backend-specific state (NCCL communicators, IPC device info, transfer
405
- plans) on `self`. The `WeightSource` is required at `trainer_init`,
406
- then replayed each round by the no-argument `send_weights()`.
406
+ plans) on `self`.
407
408
Unlike the worker engine, the trainer side does not take a
409
`WeightTransferConfig`: the backend is selected from the init info's
@@ -431,7 +431,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
431
self,
432
*,
433
client: "VLLMWeightSyncClient",
434
- source: "WeightSource",
434
+ source: "WeightSource | None" = None,
435
is_sender: bool = True,
436
) -> None:
437
self.is_sender = is_sender
@@ -447,24 +447,24 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
447
init_info: TTrainerInitInfo,
448
*,
449
client: "VLLMWeightSyncClient",
450
- source: "WeightSource",
450
+ source: "WeightSource | None" = None,
451
) -> Self:
452
"""Rendezvous with the inference side and return a ready instance.
453
454
Called on every trainer rank. The sender drives the full handshake via
455
`client` (build the worker-side init info, call
456
`client.init_weight_transfer_engine`, open the trainer-side endpoint);
457
- non-sender ranks skip the rendezvous and the RPC. `source` is stored on
458
- `self.source`; after return, `send_weights()` is callable.
457
+ non-sender ranks skip the rendezvous and the RPC.
458
"""
459
raise NotImplementedError
460
461
@abstractmethod
462
def send_weights(self) -> None:
464
- """Push `self.source`'s weights to inference workers and drive the full
465
- update round trip: `start_weight_update`, `update_weights` (run
466
- concurrently with the trainer-side broadcast when the backend requires
467
- it), then `finish_weight_update`. Called on every trainer rank."""
463
+ """Push weights to inference workers and drive the full update round
464
+ trip: `start_weight_update`, `update_weights` (run concurrently with the
465
+ trainer-side broadcast when the backend requires it), then
466
+ `finish_weight_update`. Called on every trainer rank.
467
+ """
468
raise NotImplementedError
469
470
def shutdown(self) -> None:
vllm/distributed/weight_transfer/factory.py
+17
-4
@@ -167,7 +167,7 @@ class WeightTransferTrainerFactory:
167
init_info: "TrainerInitInfo",
168
*,
169
client: "VLLMWeightSyncClient",
170
- source: "WeightSource",
170
+ source: "WeightSource | None" = None,
171
) -> TrainerWeightTransferEngine:
172
"""Build and rendezvous a ready-to-send trainer engine.
173
@@ -184,7 +184,9 @@ class WeightTransferTrainerFactory:
184
selects the engine; it also carries the wire params (e.g.
185
`packed`).
186
client: Inference-side control-plane client.
187
- source: `WeightSource` of `(name, tensor)` pairs to send each round.
187
+ source: `WeightSource` of `(name, tensor)` pairs to send each round,
188
+ for full-resync backends (NCCL, IPC). Sparse backend
189
+ omits it and passes its per-round payload to `send_weights`.
190
191
Raises:
192
ValueError: If `init_info.backend` is not registered.
@@ -233,10 +235,21 @@ WeightTransferEngineFactory.register_engine(
235
)
236
237
236
-# Trainer-side engines. Backends register here as they migrate to the stateful
237
-# trainer engine; NCCL / sparse NCCL keep their static trainer path until then.
238
+# Trainer-side engines, parallel to the worker registry above.
239
+WeightTransferTrainerFactory.register_engine(
240
+ "nccl",
241
+ "vllm.distributed.weight_transfer.nccl_engine",
242
+ "NCCLTrainerWeightTransferEngine",
243
+)
244
+
245
WeightTransferTrainerFactory.register_engine(
246
"ipc",
247
"vllm.distributed.weight_transfer.ipc_engine",
248
"IPCTrainerWeightTransferEngine",
249
)
250
+
251
+WeightTransferTrainerFactory.register_engine(
252
+ "sparse_nccl",
253
+ "vllm.distributed.weight_transfer.sparse_nccl_engine",
254
+ "SparseNCCLTrainerWeightTransferEngine",
255
+)
vllm/distributed/weight_transfer/ipc_engine.py
+4
-1
@@ -291,8 +291,10 @@ class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitI
291
init_info: IPCTrainerInitInfo,
292
*,
293
client: VLLMWeightSyncClient,
294
- source: WeightSource,
294
+ source: WeightSource | None = None,
295
) -> Self:
296
+ if source is None:
297
+ raise ValueError("IPC trainer weight transfer requires a WeightSource.")
298
engine = cls(
299
client=client,
300
source=source,
@@ -307,6 +309,7 @@ class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitI
309
return engine
310
311
def send_weights(self) -> None:
312
+ assert self.source is not None
313
source = self.source
314
if self.is_sender:
315
self.client.start_weight_update()
vllm/distributed/weight_transfer/nccl_common.py
+25
-4
@@ -9,7 +9,7 @@ sparse engine does not have to subclass the dense one.
9
"""
10
11
from dataclasses import dataclass
12
-from typing import TYPE_CHECKING
12
+from typing import TYPE_CHECKING, Protocol
13
14
import torch
15
@@ -18,16 +18,36 @@ if TYPE_CHECKING:
18
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
19
20
from vllm.distributed.weight_transfer.base import WeightTransferInitInfo
21
+from vllm.distributed.weight_transfer.packed_tensor import (
22
+ DEFAULT_PACKED_BUFFER_SIZE_BYTES,
23
+ DEFAULT_PACKED_NUM_BUFFERS,
24
+)
25
26
27
@dataclass
28
class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
25
- """Initialization info for NCCL-based weight transfer backends."""
29
+ """Worker-side initialization info for NCCL-based weight transfer backends."""
30
31
master_address: str
32
master_port: int
33
rank_offset: int
34
world_size: int
35
+ packed: bool = False
36
+ packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
37
+ packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
38
+
39
+
40
+class NCCLRendezvous(Protocol):
41
+ """The three fields `trainer_init` needs to open the trainer endpoint.
42
+
43
+ Structural so each backend can keep its own trainer init info next to its
44
+ engine (dense `NCCLTrainerInitInfo`, `SparseNCCLTrainerInitInfo`) without
45
+ this shared module importing either.
46
+ """
47
+
48
+ master_address: str
49
+ master_port: int
50
+ world_size: int
51
52
53
def stateless_init_process_group(
@@ -83,7 +103,7 @@ def worker_init_process_group(
103
104
105
def trainer_init(
86
- init_info: NCCLWeightTransferInitInfo | dict,
106
+ init_info: NCCLRendezvous | dict,
107
) -> "PyNcclCommunicator":
108
"""
109
Initialize NCCL process group for trainer-side weight transfer.
@@ -92,7 +112,8 @@ def trainer_init(
112
CUDA device (torch.accelerator.current_device_index()).
113
114
Args:
95
- init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys:
115
+ init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or
116
+ worker NCCL init info), or a dict with keys:
117
- master_address: str
118
- master_port: int
119
- world_size: int
vllm/distributed/weight_transfer/nccl_engine.py
+292
-97
@@ -2,11 +2,13 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""NCCL-based (dense) weight transfer engine."""
4
5
-from collections.abc import Callable, Iterator
6
-from dataclasses import dataclass
7
-from typing import TYPE_CHECKING, Any
5
+from collections.abc import Iterator
6
+from concurrent.futures import ThreadPoolExecutor
7
+from dataclasses import asdict, dataclass
8
+from typing import TYPE_CHECKING, Any, ClassVar
9
10
import torch
11
+from typing_extensions import Self
12
13
if TYPE_CHECKING:
14
from vllm.config import VllmConfig
@@ -14,72 +16,76 @@ if TYPE_CHECKING:
16
17
from vllm.config.weight_transfer import WeightTransferConfig
18
from vllm.distributed.weight_transfer.base import (
19
+ ParamMeta,
20
+ TrainerInitInfo,
21
+ TrainerWeightTransferEngine,
22
+ VLLMWeightSyncClient,
23
+ WeightSource,
24
WeightTransferEngine,
25
WeightTransferUpdateInfo,
26
)
27
from vllm.distributed.weight_transfer.nccl_common import (
28
NCCLWeightTransferInitInfo,
22
- trainer_init,
29
worker_init_process_group,
30
)
31
+from vllm.distributed.weight_transfer.nccl_common import (
32
+ trainer_init as open_trainer_endpoint,
33
+)
34
from vllm.distributed.weight_transfer.packed_tensor import (
35
DEFAULT_PACKED_BUFFER_SIZE_BYTES,
36
DEFAULT_PACKED_NUM_BUFFERS,
37
packed_nccl_broadcast_consumer,
38
+ packed_nccl_broadcast_producer,
39
)
40
31
-# Re-exported for backward compatibility; canonical home is nccl_common.
41
+# NCCLWeightTransferInitInfo is re-exported here for convenience; its canonical
42
+# home is nccl_common, shared with the sparse backend.
43
__all__ = [
44
"NCCLWeightTransferInitInfo",
34
- "NCCLTrainerSendWeightsArgs",
45
+ "NCCLTrainerInitInfo",
46
"NCCLWeightTransferUpdateInfo",
47
"NCCLWeightTransferEngine",
48
+ "NCCLTrainerWeightTransferEngine",
49
]
50
51
52
@dataclass
41
-class NCCLTrainerSendWeightsArgs:
42
- """Arguments for NCCL trainer_send_weights method."""
43
-
44
- group: Any
45
- """Process group (PyNcclCommunicator) for NCCL communication."""
46
- src: int = 0
47
- """Source rank (default 0, trainer is typically rank 0)."""
48
- post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor] | None = None
49
- """Optional function to apply to each (name, tensor) pair before broadcasting.
50
- If None, extracts just the tensor."""
51
- packed: bool = False
52
- """Whether to use packed tensor broadcasting for efficiency.
53
- When True, multiple tensors are batched together before broadcasting
54
- to reduce NCCL communication overhead."""
55
- stream: torch.cuda.Stream | None = None
56
- """CUDA stream to use for broadcasting if packed is False.
57
- If packed is True, new streams will be created for each buffer."""
53
+class NCCLTrainerInitInfo(TrainerInitInfo):
54
+ """Trainer-side init info for the dense NCCL weight transfer backend.
55
+
56
+ The sender opens its endpoint as NCCL rank 0, so it needs no `rank_offset`.
57
+ `world_size` is the full trainer+worker NCCL group size. `rank` (from
58
+ `TrainerInitInfo`) identifies this trainer process; rank 0 is the sender.
59
+
60
+ `packed` / buffer sizes are the transfer's wire params. The trainer
61
+ propagates them to the worker at `trainer_init` so the two sides cannot
62
+ disagree. Note this defaults to packed, unlike the worker-side
63
+ `NCCLWeightTransferInitInfo`, whose default only applies when no trainer
64
+ ships a value. `backend` is the factory dispatch key."""
65
+
66
+ backend: ClassVar[str] = "nccl"
67
+
68
+ master_address: str
69
+ master_port: int
70
+ world_size: int
71
+ packed: bool = True
72
packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
59
- """Size in bytes for each packed tensor buffer.
60
- Must match the value used in NCCLWeightTransferUpdateInfo."""
73
packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
62
- """Number of buffers for double/triple buffering during packed transfer.
63
- Must match the value used in NCCLWeightTransferUpdateInfo."""
74
75
76
@dataclass
77
class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
68
- """Update info for NCCL weight transfer backend."""
78
+ """Per-round update info for the dense NCCL weight transfer backend.
79
+
80
+ Whether the transfer is packed (and the buffer geometry) is a must-agree
81
+ wire param carried on the init info (`NCCLTrainerInitInfo` /
82
+ `NCCLWeightTransferInitInfo`), not here; this carries only the per-round
83
+ parameter metadata.
84
+ """
85
86
names: list[str]
87
dtype_names: list[str]
88
shapes: list[list[int]]
73
- packed: bool = False
74
- """Whether to use packed tensor broadcasting for efficiency.
75
- When True, multiple tensors are batched together before broadcasting
76
- to reduce NCCL communication overhead."""
77
- packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
78
- """Size in bytes for each packed tensor buffer.
79
- Both producer and consumer must use the same value."""
80
- packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
81
- """Number of buffers for double/triple buffering during packed transfer.
82
- Both producer and consumer must use the same value."""
89
90
def __post_init__(self):
91
"""Validate that all lists have the same length."""
@@ -121,15 +127,25 @@ class NCCLWeightTransferEngine(
127
) -> None:
128
super().__init__(config, vllm_config, device, model)
129
self.model_update_group: PyNcclCommunicator | None = None
130
+ # Set from the trainer-supplied init info at the handshake; defaults are
131
+ # only for the (unreachable) receive-before-init case.
132
+ self.packed = False
133
+ self.packed_buffer_size_bytes = DEFAULT_PACKED_BUFFER_SIZE_BYTES
134
+ self.packed_num_buffers = DEFAULT_PACKED_NUM_BUFFERS
135
136
def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
137
"""
127
- Initialize NCCL process group with the trainer.
138
+ Initialize NCCL process group with the trainer and record the
139
+ trainer-supplied wire params so the worker decodes exactly as the
140
+ trainer encodes.
141
142
Args:
143
init_info: NCCL initialization info containing master address, port,
131
- rank offset, and world size
144
+ rank offset, world size, and the packed wire params
145
"""
146
+ self.packed = init_info.packed
147
+ self.packed_buffer_size_bytes = init_info.packed_buffer_size_bytes
148
+ self.packed_num_buffers = init_info.packed_num_buffers
149
self.model_update_group = worker_init_process_group(
150
init_info, self.parallel_config
151
)
@@ -154,13 +170,14 @@ class NCCLWeightTransferEngine(
170
"""
171
Receive weights from trainer via NCCL broadcast.
172
157
- If update_info.packed is True, uses packed tensor broadcasting for
158
- efficient transfer of multiple weights in batches. Otherwise, uses simple
159
- one-by-one broadcasting.
173
+ Whether to use packed broadcasting (and the buffer geometry) is read
174
+ from `self.packed` / `self.packed_*`, set at the init handshake from the
175
+ trainer's init info, so it is guaranteed to match how the trainer
176
+ encoded.
177
178
Args:
162
- update_info: NCCL update info containing parameter names, dtypes, shapes,
163
- and packed flag
179
+ update_info: NCCL update info containing parameter names, dtypes,
180
+ and shapes
181
"""
182
if self.model_update_group is None:
183
raise RuntimeError(
@@ -173,7 +190,7 @@ class NCCLWeightTransferEngine(
190
)
191
192
with disable_mtp_completeness_check():
176
- if update_info.packed:
193
+ if self.packed:
194
# Build iterator of (name, (shape, dtype)) from update_info
195
def state_dict_info_iterator():
196
for name, dtype_name, shape in zip(
@@ -187,8 +204,8 @@ class NCCLWeightTransferEngine(
204
group=self.model_update_group,
205
src=0,
206
post_unpack_func=self.model.load_weights,
190
- buffer_size_bytes=update_info.packed_buffer_size_bytes,
191
- num_buffers=update_info.packed_num_buffers,
207
+ buffer_size_bytes=self.packed_buffer_size_bytes,
208
+ num_buffers=self.packed_num_buffers,
209
device=self.device,
210
)
211
else:
@@ -210,63 +227,241 @@ class NCCLWeightTransferEngine(
227
self.model_update_group = None
228
229
@staticmethod
213
- def trainer_send_weights(
214
- iterator: Iterator[tuple[str, torch.Tensor]],
215
- trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
216
- ) -> None:
217
- """Broadcast dense weights from trainer to vLLM workers.
230
+ def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
231
+ """Removed. Use the stateful `NCCLTrainerWeightTransferEngine` instead.
232
219
- Args:
220
- iterator: Iterator of model parameters. Returns (name, tensor) tuples
221
- trainer_args: Dictionary or NCCLTrainerSendWeightsArgs instance containing
222
- NCCL-specific arguments. If a dict, should contain keys from
223
- NCCLTrainerSendWeightsArgs.
224
-
225
- Example:
226
- >>> from vllm.distributed.weight_transfer.nccl_engine import (
227
- ... NCCLWeightTransferEngine,
228
- ... NCCLTrainerSendWeightsArgs,
229
- ... )
230
- >>> param_iter = ((n, p) for n, p in model.named_parameters())
231
- >>> args = NCCLTrainerSendWeightsArgs(group=group, packed=True)
232
- >>> NCCLWeightTransferEngine.trainer_send_weights(param_iter, args)
233
+ Transitional stub kept only to satisfy the (still abstract)
234
+ `WeightTransferEngine.trainer_send_weights`; that member is dropped from
235
+ the worker ABC once every backend has migrated to the trainer engine.
236
"""
234
- # Parse trainer args - accept either dict or dataclass instance
235
- if isinstance(trainer_args, dict):
236
- args = NCCLTrainerSendWeightsArgs(**trainer_args)
237
- else:
238
- args = trainer_args
237
+ raise NotImplementedError(
238
+ "The static NCCL trainer path has been replaced by "
239
+ "NCCLTrainerWeightTransferEngine. Build it via "
240
+ "WeightTransferTrainerFactory.trainer_init(NCCLTrainerInitInfo(...), "
241
+ "client=..., source=...) and drive it with send_weights()."
242
+ )
243
240
- if args.post_iter_func is None:
241
- # Default: extract just the tensor from (name, tensor) tuple
242
- post_iter_func = lambda x: x[1]
243
- else:
244
- post_iter_func = args.post_iter_func
244
246
- if args.packed:
247
- # Use packed tensor broadcasting for efficiency
248
- from vllm.distributed.weight_transfer.packed_tensor import (
249
- packed_nccl_broadcast_producer,
245
+class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerInitInfo]):
246
+ """Trainer-side NCCL weight transfer engine.
247
+
248
+ On the sender (rank 0) holds the NCCL communicator and drives the full
249
+ update round trip: it runs the inference-side `update_weights` concurrently
250
+ with the trainer-side broadcast (both rendezvous inside the same NCCL
251
+ calls), then finishes the update. Non-sender trainer ranks hold no
252
+ communicator; they only iterate the source to stay in the trainer-side
253
+ collective (e.g. FSDP `full_tensor()`) and skip the client RPCs and the
254
+ broadcast (all guarded on `is_sender`).
255
+
256
+ `packed` / buffer sizes come from `NCCLTrainerInitInfo`; the sender
257
+ propagates them to the worker at `trainer_init` (on the worker-side init
258
+ info), so per-round payloads carry only parameter metadata.
259
+ """
260
+
261
+ init_info_cls = NCCLTrainerInitInfo
262
+
263
+ def __init__(
264
+ self,
265
+ *,
266
+ client: VLLMWeightSyncClient,
267
+ source: WeightSource,
268
+ is_sender: bool = True,
269
+ packed: bool = True,
270
+ packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
271
+ packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
272
+ ) -> None:
273
+ super().__init__(client=client, source=source, is_sender=is_sender)
274
+ self.packed = packed
275
+ self.packed_buffer_size_bytes = packed_buffer_size_bytes
276
+ self.packed_num_buffers = packed_num_buffers
277
+ self.model_update_group: PyNcclCommunicator | None = None
278
+
279
+ @classmethod
280
+ def trainer_init(
281
+ cls,
282
+ init_info: NCCLTrainerInitInfo,
283
+ *,
284
+ client: VLLMWeightSyncClient,
285
+ source: WeightSource | None = None,
286
+ ) -> Self:
287
+ if source is None:
288
+ raise ValueError("NCCL trainer weight transfer requires a WeightSource.")
289
+ engine = cls(
290
+ client=client,
291
+ source=source,
292
+ is_sender=init_info.is_sender,
293
+ packed=init_info.packed,
294
+ packed_buffer_size_bytes=init_info.packed_buffer_size_bytes,
295
+ packed_num_buffers=init_info.packed_num_buffers,
296
+ )
297
+ if not engine.is_sender:
298
+ # Non-sender trainer ranks aren't part of the transfer NCCL group and
299
+ # don't drive the inference side; they only participate in the
300
+ # trainer-side gather during send_weights.
301
+ return engine
302
+
303
+ # Workers sit at rank_offset 1, after the single trainer sender rank 0.
304
+ worker_init_info = NCCLWeightTransferInitInfo(
305
+ master_address=init_info.master_address,
306
+ master_port=init_info.master_port,
307
+ rank_offset=1,
308
+ world_size=init_info.world_size,
309
+ packed=init_info.packed,
310
+ packed_buffer_size_bytes=init_info.packed_buffer_size_bytes,
311
+ packed_num_buffers=init_info.packed_num_buffers,
312
+ )
313
+
314
+ # The inference workers block inside init_weight_transfer_engine waiting
315
+ # for the NCCL rendezvous, so we kick that off on a side thread while we
316
+ # open the trainer endpoint (rank 0); both sides must rendezvous together.
317
+ with ThreadPoolExecutor(max_workers=1) as exe:
318
+ future = exe.submit(
319
+ engine.client.init_weight_transfer_engine, asdict(worker_init_info)
320
)
321
+ engine.model_update_group = open_trainer_endpoint(init_info)
322
+ future.result() # surface any inference-side init error
323
+
324
+ return engine
325
+
326
+ def send_weights(self) -> None:
327
+ assert self.source is not None # guaranteed by trainer_init / __init__
328
+ source = self.source
329
+
330
+ # Metadata is declared without gathering. For Megatron it is itself a
331
+ # collective, so every rank runs it; only the sender ships it.
332
+ meta = source.metadata()
333
+
334
+ if not self.is_sender:
335
+ # Non-sender ranks only join the trainer-side gather collective.
336
+ self._broadcast(source, meta)
337
+ self._post_send_sync()
338
+ return
339
+
340
+ update_info = NCCLWeightTransferUpdateInfo(
341
+ names=[m.name for m in meta],
342
+ dtype_names=[str(m.dtype).split(".")[-1] for m in meta],
343
+ shapes=[list(m.shape) for m in meta],
344
+ )
345
346
+ self.client.start_weight_update()
347
+ # update_weights (workers receive) must run concurrently with the
348
+ # trainer-side broadcast — both rendezvous inside the same NCCL calls.
349
+ exe = ThreadPoolExecutor(max_workers=1)
350
+ try:
351
+ future = exe.submit(self.client.update_weights, asdict(update_info))
352
+ # Cheap best-effort: if update_weights already failed (e.g. a bad
353
+ # request rejected before any NCCL call), surface it now instead of
354
+ # hanging in broadcast waiting for a peer that will never arrive.
355
+ if future.done():
356
+ future.result()
357
+ self._broadcast(source, meta)
358
+ future.result() # surface inference-side errors
359
+ finally:
360
+ # Never wait for the RPC thread here. If the broadcast raised, the
361
+ # worker is still blocked in the matching NCCL call and will never
362
+ # return, so joining would turn the error into a permanent hang.
363
+ # Let the exception out instead; the transfer group is unusable
364
+ # either way and the caller has to tear it down. (The orphaned
365
+ # thread is still joined at interpreter exit, so a caller that wants
366
+ # to exit cleanly after such a failure must not wait on it.)
367
+ exe.shutdown(wait=False)
368
+ self.client.finish_weight_update()
369
+ self._post_send_sync()
370
+
371
+ def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None:
372
+ """Iterate the source (materializing each tensor — a collective on all
373
+ ranks) and, on the sender, broadcast from rank 0, packed or one-by-one.
374
+ Non-sender ranks only replay the iteration to stay in the collective."""
375
+ if not self.is_sender:
376
+ for _ in source:
377
+ pass
378
+ return
379
+
380
+ assert self.model_update_group is not None, (
381
+ "trainer_init() must be called before _broadcast()."
382
+ )
383
+ pairs = self._checked_iter(source, meta)
384
+ if self.packed:
385
packed_nccl_broadcast_producer(
253
- iterator=iterator,
254
- group=args.group,
255
- src=args.src,
256
- post_iter_func=post_iter_func,
257
- buffer_size_bytes=args.packed_buffer_size_bytes,
258
- num_buffers=args.packed_num_buffers,
386
+ iterator=pairs,
387
+ group=self.model_update_group,
388
+ src=0,
389
+ post_iter_func=lambda item: item[1],
390
+ buffer_size_bytes=self.packed_buffer_size_bytes,
391
+ num_buffers=self.packed_num_buffers,
392
)
393
else:
261
- # Use simple one-by-one broadcasting
262
- for item in iterator:
263
- tensor = post_iter_func(item)
264
- args.group.broadcast(
265
- tensor,
266
- src=args.src,
267
- stream=args.stream or torch.cuda.current_stream(),
394
+ stream = torch.cuda.current_stream()
395
+ for _name, tensor in pairs:
396
+ # NCCL sends `numel` elements straight from `data_ptr()`, so a
397
+ # non-contiguous view would ship whatever follows its base
398
+ # pointer. Keep the copy referenced until the broadcast is
399
+ # enqueued. (The packed path linearizes in `pack_tensors`.)
400
+ send = tensor if tensor.is_contiguous() else tensor.contiguous()
401
+ self.model_update_group.broadcast(send, src=0, stream=stream)
402
+
403
+ @staticmethod
404
+ def _checked_iter(
405
+ source: WeightSource, meta: list[ParamMeta]
406
+ ) -> Iterator[tuple[str, torch.Tensor]]:
407
+ """Yield the source's pairs, checking each against what the worker was
408
+ told to expect.
409
+
410
+ The worker sizes its receive buffers — and in packed mode cuts its chunk
411
+ boundaries — from the update info, which is built from `metadata()`. If
412
+ iteration disagrees with it, the two sides split the stream differently
413
+ and the transfer hangs in NCCL or loads garbage. Checking here costs one
414
+ comparison per parameter and turns that into an error naming the first
415
+ divergent parameter. Sender-only: under pipeline parallelism a
416
+ non-sender's yielded tensor is not meaningful.
417
+ """
418
+ sent = 0
419
+ for name, tensor in source:
420
+ if sent >= len(meta):
421
+ raise ValueError(
422
+ f"WeightSource yielded more parameters than metadata() "
423
+ f"declared ({len(meta)}); first extra is {name!r}."
424
+ )
425
+ expected = meta[sent]
426
+ if (
427
+ name != expected.name
428
+ or tensor.dtype != expected.dtype
429
+ or tuple(tensor.shape) != expected.shape
430
+ ):
431
+ raise ValueError(
432
+ "WeightSource metadata() disagrees with iteration at index "
433
+ f"{sent}: declared {expected.name!r} "
434
+ f"{expected.dtype} {tuple(expected.shape)}, got {name!r} "
435
+ f"{tensor.dtype} {tuple(tensor.shape)}. Both channels must "
436
+ "enumerate the same parameters in the same order."
437
)
438
+ sent += 1
439
+ yield name, tensor
440
+ if sent != len(meta):
441
+ raise ValueError(
442
+ f"WeightSource yielded {sent} parameters but metadata() "
443
+ f"declared {len(meta)}; the worker is waiting for the rest."
444
+ )
445
270
- # Trainer-side process-group setup. Delegates to the shared helper so the
271
- # sparse engine can reuse the exact same rendezvous without subclassing.
272
- trainer_init = staticmethod(trainer_init)
446
+ def _post_send_sync(self) -> None:
447
+ """Wait for this rank's transfer work to land before returning.
448
+
449
+ Broadcasts are only *enqueued* by `send_weights`: the unpacked path on
450
+ the current stream, the packed path on the producer's own streams (which
451
+ it drains itself). Waiting here lets a caller mutate parameters, or
452
+ start the next step on another stream, as soon as `send_weights`
453
+ returns, instead of silently depending on same-stream ordering. Every
454
+ rank waits: a non-sender's `full_tensor()` gathers feed the sender's
455
+ broadcast, so they must have landed before it may touch its shards.
456
+
457
+ Unlike IPC there is no cross-rank barrier here. Nothing in this backend
458
+ outlives the collective it travelled in (IPC's barrier keeps shared
459
+ buffers alive until every consumer has opened them), so a barrier would
460
+ only add a dependency on the default process group that this backend
461
+ otherwise does not have.
462
+ """
463
+ if torch.cuda.is_available():
464
+ torch.cuda.current_stream().synchronize()
465
+
466
+ def shutdown(self) -> None:
467
+ self.model_update_group = None
vllm/distributed/weight_transfer/packed_tensor.py
+8
@@ -177,6 +177,10 @@ def packed_nccl_broadcast_producer(
177
# Move to the next buffer
178
buffer_idx = (buffer_idx + 1) % num_buffers
179
180
+ # Drain every slot before returning.
181
+ for stream in streams:
182
+ stream.synchronize()
183
+
184
185
def packed_nccl_broadcast_consumer(
186
iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]],
@@ -279,6 +283,10 @@ def packed_nccl_broadcast_consumer(
283
)
284
break
285
286
+ # Drain every slot before returning.
287
+ for stream in streams:
288
+ stream.synchronize()
289
+
290
291
# ── IPC packed transfer ────────────────────────────────────────────────
292
vllm/distributed/weight_transfer/sparse_nccl_engine.py
+206
-21
@@ -14,11 +14,13 @@ MVP limitations:
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
17
+from collections.abc import Iterable
18
+from concurrent.futures import ThreadPoolExecutor
19
+from dataclasses import asdict, dataclass
20
+from typing import TYPE_CHECKING, Any, ClassVar
21
22
import torch
23
+from typing_extensions import Self
24
25
if TYPE_CHECKING:
26
from vllm.config import VllmConfig
@@ -26,20 +28,27 @@ if TYPE_CHECKING:
28
29
from vllm.config.weight_transfer import WeightTransferConfig
30
from vllm.distributed.weight_transfer.base import (
31
+ TrainerInitInfo,
32
+ TrainerWeightTransferEngine,
33
+ VLLMWeightSyncClient,
34
+ WeightSource,
35
WeightTransferEngine,
36
WeightTransferUpdateInfo,
37
)
38
from vllm.distributed.weight_transfer.nccl_common import (
39
NCCLWeightTransferInitInfo,
34
- trainer_init,
40
worker_init_process_group,
41
)
37
-from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs
42
+from vllm.distributed.weight_transfer.nccl_common import (
43
+ trainer_init as open_trainer_endpoint,
44
+)
45
46
__all__ = [
47
"SparseWeightPatch",
48
+ "SparseNCCLTrainerInitInfo",
49
"SparseNCCLWeightTransferUpdateInfo",
50
"SparseNCCLWeightTransferEngine",
51
+ "SparseNCCLTrainerWeightTransferEngine",
52
]
53
54
@@ -50,6 +59,25 @@ class SparseWeightPatch:
59
name: str
60
indices: torch.Tensor
61
values: torch.Tensor
62
+ full_shape: tuple[int, ...] | None = None
63
+ """Full shape of the patched parameter. Required when the patch is sent
64
+ via `SparseNCCLTrainerWeightTransferEngine` (it ships in the per-round
65
+ update info); the worker-side apply path does not read it."""
66
+
67
+
68
+@dataclass
69
+class SparseNCCLTrainerInitInfo(TrainerInitInfo):
70
+ """Trainer-side init info for the sparse NCCL weight transfer backend.
71
+
72
+ Same rendezvous shape as the dense NCCL backend (the sender opens its
73
+ endpoint as NCCL rank 0), but with no packed wire params: sparse transfers
74
+ are never packed. `backend` is the factory dispatch key."""
75
+
76
+ backend: ClassVar[str] = "sparse_nccl"
77
+
78
+ master_address: str
79
+ master_port: int
80
+ world_size: int
81
82
83
@dataclass
@@ -200,25 +228,182 @@ class SparseNCCLWeightTransferEngine(
228
self.model_update_group = None
229
230
@staticmethod
203
- def trainer_send_weights(
204
- iterator: Iterator[SparseWeightPatch],
205
- trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
231
+ def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
232
+ """Removed. Use the stateful `SparseNCCLTrainerWeightTransferEngine`.
233
+
234
+ Transitional stub kept only to satisfy the (still abstract)
235
+ `WeightTransferEngine.trainer_send_weights`; that member is dropped from
236
+ the worker ABC once every backend has migrated to the trainer engine.
237
+ """
238
+ raise NotImplementedError(
239
+ "The static sparse NCCL trainer path has been replaced by "
240
+ "SparseNCCLTrainerWeightTransferEngine. Build it via "
241
+ "WeightTransferTrainerFactory.trainer_init("
242
+ "SparseNCCLTrainerInitInfo(...), client=..., source=...) and drive "
243
+ "it with send_weights()."
244
+ )
245
+
246
+
247
+class SparseNCCLTrainerWeightTransferEngine(
248
+ TrainerWeightTransferEngine[SparseNCCLTrainerInitInfo]
249
+):
250
+ """Trainer-side sparse NCCL weight transfer engine.
251
+
252
+ Broadcasts flat-index (indices, values) patches from NCCL rank 0 while the
253
+ inference-side `update_weights` runs concurrently on a side thread (the
254
+ worker's recvs rendezvous inside the same NCCL broadcasts), then finishes
255
+ the update. Unlike the full-resync backends, sparse patches differ every
256
+ round (a fresh set of deltas from each optimizer step), so they are not a
257
+ stable `WeightSource`: the engine takes no `source`, and each round's
258
+ patches are passed straight to `send_weights(patches)`. A round with no
259
+ patches is a no-op.
260
+
261
+ The sparse backend assumes a single-rank trainer (matching its TP=1 / PP=1
262
+ MVP scope), so non-sender ranks skip `send_weights` entirely.
263
+ """
264
+
265
+ init_info_cls = SparseNCCLTrainerInitInfo
266
+
267
+ def __init__(
268
+ self,
269
+ *,
270
+ client: VLLMWeightSyncClient,
271
+ source: WeightSource | None = None,
272
+ is_sender: bool = True,
273
) -> None:
207
- """Broadcast sparse flat-index patches from trainer to vLLM workers."""
208
- if isinstance(trainer_args, dict):
209
- args = NCCLTrainerSendWeightsArgs(**trainer_args)
210
- else:
211
- args = trainer_args
274
+ # Sparse is a delta backend: it takes per-round patches via
275
+ # send_weights, so a `source` would silently never be sent. The
276
+ # parameter exists only to match the base/factory signature.
277
+ if source is not None:
278
+ raise ValueError(
279
+ "Sparse NCCL weight transfer takes no WeightSource; pass each "
280
+ "round's patches to send_weights(patches) instead."
281
+ )
282
+ super().__init__(client=client, source=source, is_sender=is_sender)
283
+ self.model_update_group: PyNcclCommunicator | None = None
284
+
285
+ @classmethod
286
+ def trainer_init(
287
+ cls,
288
+ init_info: SparseNCCLTrainerInitInfo,
289
+ *,
290
+ client: VLLMWeightSyncClient,
291
+ source: WeightSource | None = None,
292
+ ) -> Self:
293
+ engine = cls(client=client, source=source, is_sender=init_info.is_sender)
294
+ if not engine.is_sender:
295
+ return engine
296
213
- if args.packed:
297
+ # Workers sit at rank_offset 1, after the single trainer sender rank 0.
298
+ # Sparse transfers are never packed, so the worker keeps the unpacked
299
+ # defaults on its init info.
300
+ worker_init_info = NCCLWeightTransferInitInfo(
301
+ master_address=init_info.master_address,
302
+ master_port=init_info.master_port,
303
+ rank_offset=1,
304
+ world_size=init_info.world_size,
305
+ )
306
+
307
+ # The inference workers block inside init_weight_transfer_engine waiting
308
+ # for the NCCL rendezvous, so we kick that off on a side thread while we
309
+ # open the trainer endpoint (rank 0); both sides must rendezvous together.
310
+ with ThreadPoolExecutor(max_workers=1) as exe:
311
+ future = exe.submit(
312
+ engine.client.init_weight_transfer_engine, asdict(worker_init_info)
313
+ )
314
+ # Open the trainer endpoint as NCCL rank 0 on the current device
315
+ # (the init info satisfies the helper's rendezvous protocol).
316
+ engine.model_update_group = open_trainer_endpoint(init_info)
317
+ future.result() # surface any inference-side init error
318
+
319
+ return engine
320
+
321
+ def send_weights(self, patches: Iterable[SparseWeightPatch] | None = None) -> None:
322
+ """Broadcast this round's sparse patches. `patches` is the per-round
323
+ payload (sparse deltas differ every round), so it is passed here rather
324
+ than fixed at init. Every patch must set `full_shape`."""
325
+ if not self.is_sender:
326
+ return
327
+
328
+ patches = list(patches) if patches is not None else []
329
+ if not patches:
330
+ return
331
+
332
+ shapes = []
333
+ for patch in patches:
334
+ if patch.full_shape is None:
335
+ raise ValueError(
336
+ "SparseWeightPatch.full_shape must be set to send via the "
337
+ f"trainer engine: {patch.name}"
338
+ )
339
+ self._validate_patch(patch)
340
+ shapes.append(list(patch.full_shape))
341
+
342
+ update_info = SparseNCCLWeightTransferUpdateInfo(
343
+ names=[patch.name for patch in patches],
344
+ dtype_names=[str(patch.values.dtype).split(".")[-1] for patch in patches],
345
+ shapes=shapes,
346
+ num_updates_list=[patch.indices.numel() for patch in patches],
347
+ )
348
+
349
+ assert self.model_update_group is not None, (
350
+ "trainer_init() must be called before send_weights()."
351
+ )
352
+ self.client.start_weight_update()
353
+ # update_weights (workers receive) must run concurrently with the
354
+ # trainer-side broadcasts — both rendezvous inside the same NCCL calls.
355
+ exe = ThreadPoolExecutor(max_workers=1)
356
+ try:
357
+ future = exe.submit(self.client.update_weights, asdict(update_info))
358
+ # Cheap best-effort: if update_weights already failed (e.g. a bad
359
+ # request rejected before any NCCL call), surface it now instead of
360
+ # hanging in broadcast waiting for a peer that will never arrive.
361
+ if future.done():
362
+ future.result()
363
+ stream = torch.cuda.current_stream()
364
+ for patch in patches:
365
+ self.model_update_group.broadcast(patch.indices, src=0, stream=stream)
366
+ self.model_update_group.broadcast(patch.values, src=0, stream=stream)
367
+ future.result() # surface inference-side errors
368
+ finally:
369
+ # Never wait for the RPC thread here: if a broadcast raised, the
370
+ # worker is still blocked in the matching NCCL call and will never
371
+ # return, so joining would turn the error into a permanent hang.
372
+ # See NCCLTrainerWeightTransferEngine.send_weights.
373
+ exe.shutdown(wait=False)
374
+ self.client.finish_weight_update()
375
+ self._post_send_sync()
376
+
377
+ @staticmethod
378
+ def _validate_patch(patch: SparseWeightPatch) -> None:
379
+ """Reject a malformed patch before any NCCL call.
380
+
381
+ The worker checks the same invariants in `_apply_patch`, but by then the
382
+ broadcasts are already under way: a mismatch surfaces as a size mismatch
383
+ mid-transfer, which wedges both sides instead of raising. Checking here
384
+ keeps the failure on the trainer, before `start_weight_update`.
385
+ """
386
+ if patch.indices.dtype != torch.int32:
387
+ raise ValueError(
388
+ f"Sparse weight updates require int32 indices: {patch.name}"
389
+ )
390
+ if patch.indices.ndim != 1 or patch.values.ndim != 1:
391
raise ValueError(
215
- "Sparse NCCL updates cannot be combined with `packed=True`"
392
+ f"Sparse weight patches must be 1D flattened updates: {patch.name}"
393
+ )
394
+ if patch.indices.numel() != patch.values.numel():
395
+ raise ValueError(
396
+ f"`indices` and `values` must have matching lengths for {patch.name}"
397
)
398
218
- stream = args.stream or torch.cuda.current_stream()
219
- for patch in iterator:
220
- args.group.broadcast(patch.indices, src=args.src, stream=stream)
221
- args.group.broadcast(patch.values, src=args.src, stream=stream)
399
+ def _post_send_sync(self) -> None:
400
+ """Wait for the broadcasts to land before returning, so a caller may
401
+ rebuild or free the patch tensors as soon as `send_weights` returns
402
+ rather than relying on same-stream ordering. See
403
+ `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no
404
+ cross-rank barrier (and sparse is single-rank on the trainer anyway)."""
405
+ if torch.cuda.is_available():
406
+ torch.cuda.current_stream().synchronize()
407
223
- # Trainer-side process-group setup (shared with the dense engine).
224
- trainer_init = staticmethod(trainer_init)
408
+ def shutdown(self) -> None:
409
+ self.model_update_group = None