[Distributed] Support pre-shared ncclUniqueId rendezvous for weight transfer (#53784)
Signed-off-by: Dharak Kharod <dharak.kharod@cohere.com> Co-authored-by: Cursor <cursoragent@cursor.com>
dharak-cohere committed
Sep 1, 2026 at 02:55 UTC
d0e695a91b67a8214c2e9ed77595186d2f2844b6
6 files changed
+673
-9
tests/distributed/test_weight_transfer_nccl_uid.py
new
+465
@@ -0,0 +1,465 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the pre-shared ncclUniqueId weight-transfer rendezvous.
4
+
5
+The rendezvous lets a peer that cannot join a TCPStore (e.g. a torch-free JAX
6
+trainer) initialize NCCL weight transfer from a unique id minted out of band.
7
+
8
+* Unit tests for the mutually-exclusive rendezvous-mode validation on the NCCL
9
+ init infos.
10
+* Two-GPU integration tests driving the real worker engine against a peer: one
11
+ via vLLM's own ``PyNcclCommunicator``, one via a genuinely torch-free trainer
12
+ (cupy buffer + raw ``ncclBroadcast``).
13
+
14
+The shared Ray helpers come from the sibling module (unchanged); the worker task
15
+is local so the sibling's TCP test and its worker stay untouched.
16
+"""
17
+
18
+from unittest.mock import MagicMock
19
+
20
+import pybase64 as base64
21
+import pytest
22
+import ray
23
+import torch
24
+
25
+from tests.distributed.test_weight_transfer import (
26
+ _init_ray_for_weight_transfer,
27
+ _set_ray_assigned_device,
28
+ create_mock_vllm_config,
29
+)
30
+from vllm.config.weight_transfer import WeightTransferConfig
31
+from vllm.distributed.weight_transfer.nccl_engine import (
32
+ NCCLTrainerInitInfo,
33
+ NCCLWeightTransferEngine,
34
+ NCCLWeightTransferInitInfo,
35
+)
36
+
37
+# Standard base64 of exactly 128 zero bytes: a well-formed (dummy) ncclUniqueId.
38
+VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode()
39
+
40
+
41
+class TestNCCLRendezvousValidation:
42
+ """Test the xor-of-rendezvous-modes validation on the NCCL init infos."""
43
+
44
+ def test_worker_init_info_is_keyword_only(self):
45
+ # kw_only so a stale positional call fails loudly rather than silently
46
+ # swapping fields once the rendezvous fields became optional.
47
+ with pytest.raises(TypeError):
48
+ NCCLWeightTransferInitInfo(1, 3, "127.0.0.1", 12345) # type: ignore[call-arg]
49
+
50
+ def test_trainer_init_info_accepts_positional_fields(self):
51
+ # NCCLTrainerInitInfo keeps positional subclass fields; only `rank`
52
+ # (from TrainerInitInfo) is keyword-only.
53
+ trainer_positional = NCCLTrainerInitInfo("127.0.0.1", 12345, 3, rank=0)
54
+ trainer_keyword = NCCLTrainerInitInfo(
55
+ master_address="127.0.0.1",
56
+ master_port=12345,
57
+ world_size=3,
58
+ rank=0,
59
+ )
60
+ assert trainer_positional == trainer_keyword
61
+
62
+ def test_tcp_mode_valid(self):
63
+ info = NCCLWeightTransferInitInfo(
64
+ rank_offset=1,
65
+ world_size=3,
66
+ master_address="127.0.0.1",
67
+ master_port=12345,
68
+ )
69
+ assert info.master_address == "127.0.0.1"
70
+ assert info.nccl_unique_id_b64 is None
71
+
72
+ def test_uid_mode_valid(self):
73
+ info = NCCLWeightTransferInitInfo(
74
+ rank_offset=1,
75
+ world_size=3,
76
+ nccl_unique_id_b64=VALID_UID_B64,
77
+ )
78
+ assert info.nccl_unique_id_b64 == VALID_UID_B64
79
+ assert info.master_address is None
80
+ assert info.master_port is None
81
+ assert VALID_UID_B64 not in repr(info)
82
+
83
+ def test_both_modes_raises(self):
84
+ with pytest.raises(ValueError, match="not both"):
85
+ NCCLWeightTransferInitInfo(
86
+ rank_offset=1,
87
+ world_size=3,
88
+ master_address="127.0.0.1",
89
+ master_port=12345,
90
+ nccl_unique_id_b64=VALID_UID_B64,
91
+ )
92
+
93
+ def test_neither_mode_raises(self):
94
+ with pytest.raises(ValueError, match="need nccl_unique_id_b64"):
95
+ NCCLWeightTransferInitInfo(rank_offset=1, world_size=3)
96
+
97
+ def test_half_specified_tcp_raises(self):
98
+ with pytest.raises(ValueError, match="master_port"):
99
+ NCCLWeightTransferInitInfo(
100
+ rank_offset=1,
101
+ world_size=3,
102
+ master_address="127.0.0.1",
103
+ )
104
+
105
+ def test_bad_base64_uid_raises(self):
106
+ with pytest.raises(ValueError, match="not valid standard base64"):
107
+ NCCLWeightTransferInitInfo(
108
+ rank_offset=1,
109
+ world_size=3,
110
+ nccl_unique_id_b64="!!!!",
111
+ )
112
+
113
+ def test_wrong_length_uid_raises(self):
114
+ short_uid = base64.b64encode(b"\x00" * 64).decode()
115
+ with pytest.raises(ValueError, match="128"):
116
+ NCCLWeightTransferInitInfo(
117
+ rank_offset=1,
118
+ world_size=3,
119
+ nccl_unique_id_b64=short_uid,
120
+ )
121
+
122
+ def test_parse_init_info_uid_valid(self):
123
+ engine = NCCLWeightTransferEngine(
124
+ WeightTransferConfig(backend="nccl"),
125
+ create_mock_vllm_config(),
126
+ torch.device("cuda"),
127
+ MagicMock(spec=torch.nn.Module),
128
+ )
129
+ init_info = engine.parse_init_info(
130
+ {
131
+ "nccl_unique_id_b64": VALID_UID_B64,
132
+ "rank_offset": 1,
133
+ "world_size": 3,
134
+ }
135
+ )
136
+ assert isinstance(init_info, NCCLWeightTransferInitInfo)
137
+ assert init_info.nccl_unique_id_b64 == VALID_UID_B64
138
+ assert init_info.master_address is None
139
+
140
+ def test_parse_init_info_both_modes_raises(self):
141
+ engine = NCCLWeightTransferEngine(
142
+ WeightTransferConfig(backend="nccl"),
143
+ create_mock_vllm_config(),
144
+ torch.device("cuda"),
145
+ MagicMock(spec=torch.nn.Module),
146
+ )
147
+ with pytest.raises(ValueError, match="not both"):
148
+ engine.parse_init_info(
149
+ {
150
+ "master_address": "127.0.0.1",
151
+ "master_port": 12345,
152
+ "nccl_unique_id_b64": VALID_UID_B64,
153
+ "rank_offset": 1,
154
+ "world_size": 3,
155
+ }
156
+ )
157
+
158
+ @pytest.mark.parametrize(
159
+ ("unique_id_bytes", "rank", "world_size", "match"),
160
+ [
161
+ (b"\x00" * 64, 0, 2, "128-byte"),
162
+ (b"\x00" * 128, -1, 2, "out of range"),
163
+ (b"\x00" * 128, 2, 2, "out of range"),
164
+ ],
165
+ )
166
+ def test_from_unique_id_bytes_validation(
167
+ self, unique_id_bytes, rank, world_size, match
168
+ ):
169
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
170
+
171
+ with pytest.raises(ValueError, match=match):
172
+ PyNcclCommunicator.from_unique_id_bytes(
173
+ unique_id_bytes,
174
+ rank=rank,
175
+ world_size=world_size,
176
+ device=0,
177
+ )
178
+
179
+
180
+@ray.remote(num_gpus=1)
181
+def inference_receive_tensor(
182
+ init_info_dict: dict,
183
+ tensor_shape: list[int],
184
+ tensor_dtype: str,
185
+) -> dict:
186
+ """Worker that joins via the given init dict and records the tensor it gets.
187
+
188
+ Local to this module (rather than the sibling's TCP-only worker) so it can
189
+ take a UID init payload -- the same dict boundary an external trainer POSTs.
190
+ """
191
+ import contextlib
192
+
193
+ import torch
194
+
195
+ _set_ray_assigned_device()
196
+
197
+ from vllm.config.parallel import ParallelConfig
198
+ from vllm.config.weight_transfer import WeightTransferConfig
199
+ from vllm.distributed.weight_transfer.nccl_engine import (
200
+ NCCLWeightTransferEngine,
201
+ NCCLWeightTransferInitInfo,
202
+ NCCLWeightTransferUpdateInfo,
203
+ )
204
+
205
+ class Recorder(torch.nn.Module):
206
+ def __init__(self):
207
+ super().__init__()
208
+ self.received = []
209
+
210
+ def load_weights(self, weights):
211
+ for name, tensor in weights:
212
+ self.received.append((name, tensor.clone()))
213
+
214
+ config = WeightTransferConfig(backend="nccl")
215
+ vllm_config = MagicMock()
216
+ parallel_config = MagicMock(spec=ParallelConfig)
217
+ parallel_config.rank = 0
218
+ parallel_config.world_size = 1
219
+ parallel_config.data_parallel_rank = 0
220
+ parallel_config.data_parallel_index = 0
221
+ vllm_config.parallel_config = parallel_config
222
+ vllm_config.model_config = MagicMock()
223
+
224
+ recorder = Recorder()
225
+ engine = NCCLWeightTransferEngine(
226
+ config, vllm_config, torch.device("cuda"), recorder
227
+ )
228
+ # Transport-only test: bypass the set_current_vllm_config context that
229
+ # receive_weights enters, since vllm_config here is a mock.
230
+ import vllm.config as _vllm_config_mod
231
+
232
+ _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
233
+
234
+ try:
235
+ engine.init_transfer_engine(NCCLWeightTransferInitInfo(**init_info_dict))
236
+ engine.receive_weights(
237
+ NCCLWeightTransferUpdateInfo(
238
+ names=["test.weight"],
239
+ dtype_names=[tensor_dtype],
240
+ shapes=[tensor_shape],
241
+ )
242
+ )
243
+ torch.accelerator.synchronize()
244
+
245
+ received_shape = None
246
+ received_sum = None
247
+ if len(recorder.received) == 1:
248
+ _name, tensor = recorder.received[0]
249
+ received_shape = list(tensor.shape)
250
+ received_sum = tensor.sum().item()
251
+
252
+ return {
253
+ "success": received_shape == tensor_shape
254
+ and abs(received_sum - torch.tensor(tensor_shape).prod().item()) < 0.01,
255
+ "received_shape": received_shape,
256
+ "received_sum": received_sum,
257
+ }
258
+ finally:
259
+ engine.shutdown()
260
+
261
+
262
+@ray.remote(num_gpus=1)
263
+def trainer_broadcast_tensor_uid(
264
+ uid_b64: str,
265
+ world_size: int,
266
+ tensor_shape: list[int],
267
+ tensor_dtype: str,
268
+) -> dict:
269
+ """Drive rank 0 via vLLM's own communicator from a pre-shared unique id.
270
+
271
+ A real torch-free trainer would mint the id and drive rank 0 itself; here we
272
+ reuse the real send protocol by injecting the communicator into a sender
273
+ engine, exercising the worker-side UID path end to end.
274
+ """
275
+ import pybase64 as base64
276
+ import torch
277
+
278
+ _set_ray_assigned_device()
279
+
280
+ from vllm.distributed.weight_transfer import ModuleSource
281
+ from vllm.distributed.weight_transfer.nccl_common import uid_init_process_group
282
+ from vllm.distributed.weight_transfer.nccl_engine import (
283
+ NCCLTrainerWeightTransferEngine,
284
+ )
285
+
286
+ class NoopClient:
287
+ def init_weight_transfer_engine(self, init_info):
288
+ pass
289
+
290
+ def start_weight_update(self):
291
+ pass
292
+
293
+ def update_weights(self, update_info):
294
+ pass
295
+
296
+ def finish_weight_update(self, weight_version=None):
297
+ pass
298
+
299
+ dtype = getattr(torch, tensor_dtype)
300
+ model = torch.nn.Module()
301
+ test_module = torch.nn.Module()
302
+ test_module.register_parameter(
303
+ "weight",
304
+ torch.nn.Parameter(
305
+ torch.ones(tensor_shape, dtype=dtype, device="cuda"), requires_grad=False
306
+ ),
307
+ )
308
+ model.add_module("test", test_module)
309
+
310
+ engine = NCCLTrainerWeightTransferEngine(
311
+ client=NoopClient(),
312
+ source=ModuleSource(model),
313
+ is_sender=True,
314
+ packed=False,
315
+ )
316
+ engine.model_update_group = uid_init_process_group(
317
+ base64.b64decode(uid_b64),
318
+ rank=0,
319
+ world_size=world_size,
320
+ device=torch.accelerator.current_device_index(),
321
+ )
322
+ try:
323
+ engine.send_weights()
324
+ torch.accelerator.synchronize()
325
+ return {"success": True}
326
+ finally:
327
+ engine.shutdown()
328
+
329
+
330
+@ray.remote(num_gpus=1)
331
+def trainer_broadcast_tensor_torch_free(
332
+ uid_b64: str,
333
+ world_size: int,
334
+ tensor_shape: list[int],
335
+ tensor_dtype: str,
336
+) -> dict:
337
+ """Trainer peer that uses no torch and no PyNcclCommunicator (a JAX stand-in).
338
+
339
+ Drives NCCL through the raw ctypes wrapper with a cupy send buffer, so it
340
+ exercises interop with a foreign NCCL peer rather than vLLM against itself.
341
+ ``PyNcclCommunicator.__init__`` runs a one-element warm-up all_reduce right
342
+ after ``ncclCommInitRank``; that is a collective, so this peer must issue the
343
+ matching all_reduce before the broadcast or both ranks deadlock.
344
+ """
345
+ import cupy as cp
346
+ import numpy as np
347
+ import pybase64 as base64
348
+
349
+ device = _set_ray_assigned_device()
350
+ cp.cuda.Device(device.index).use()
351
+
352
+ from vllm.distributed.device_communicators.pynccl_wrapper import (
353
+ NCCLLibrary,
354
+ buffer_type,
355
+ cudaStream_t,
356
+ ncclDataTypeEnum,
357
+ ncclRedOpTypeEnum,
358
+ )
359
+
360
+ assert tensor_dtype == "float32", "harness only wires float32"
361
+ # Fill on the host and copy up (a plain H2D memcpy) so we neither use torch
362
+ # nor make cupy JIT-compile a fill kernel (which needs NVRTC). All-ones lets
363
+ # the worker verify by sum == numel.
364
+ sendbuff = cp.asarray(np.ones(tensor_shape, dtype=np.float32))
365
+ warmup = cp.asarray(np.zeros(1, dtype=np.float32))
366
+ numel = int(sendbuff.size)
367
+ # The copies above run on the default stream; drain them before the NCCL
368
+ # collectives read the buffers on our own stream.
369
+ cp.cuda.Device(device.index).synchronize()
370
+ stream = cp.cuda.Stream()
371
+
372
+ nccl = NCCLLibrary()
373
+ comm = nccl.ncclCommInitRank(
374
+ world_size, nccl.unique_id_from_bytes(base64.b64decode(uid_b64)), 0
375
+ )
376
+ try:
377
+ # Mirror the worker's init-time warm-up all_reduce (rank 0 must join it).
378
+ nccl.ncclAllReduce(
379
+ buffer_type(warmup.data.ptr),
380
+ buffer_type(warmup.data.ptr),
381
+ 1,
382
+ ncclDataTypeEnum.ncclFloat32,
383
+ ncclRedOpTypeEnum.ncclSum,
384
+ comm,
385
+ cudaStream_t(stream.ptr),
386
+ )
387
+ nccl.ncclBroadcast(
388
+ buffer_type(sendbuff.data.ptr),
389
+ buffer_type(sendbuff.data.ptr),
390
+ numel,
391
+ ncclDataTypeEnum.ncclFloat32,
392
+ 0, # root
393
+ comm,
394
+ cudaStream_t(stream.ptr),
395
+ )
396
+ stream.synchronize()
397
+ return {"success": True}
398
+ finally:
399
+ # Abort (not the collective destroy) for an uncoordinated teardown.
400
+ nccl.ncclCommAbort(comm)
401
+
402
+
403
+def _run_uid_transfer(trainer_task) -> None:
404
+ """Mint a unique id, run the worker and the given trainer peer concurrently,
405
+ and assert the tensor round-trips. Both ranks must enter init together --
406
+ there is no store barrier on the UID path."""
407
+ from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
408
+
409
+ _init_ray_for_weight_transfer()
410
+
411
+ nccl = NCCLLibrary()
412
+ uid_b64 = base64.b64encode(bytes(nccl.ncclGetUniqueId().internal)).decode()
413
+
414
+ world_size = 2 # 1 trainer + 1 inference worker
415
+ tensor_shape = [100, 100]
416
+ tensor_dtype = "float32"
417
+
418
+ inference_future = inference_receive_tensor.remote(
419
+ {
420
+ "rank_offset": 1,
421
+ "world_size": world_size,
422
+ "nccl_unique_id_b64": uid_b64,
423
+ "packed": False,
424
+ },
425
+ tensor_shape,
426
+ tensor_dtype,
427
+ )
428
+ trainer_future = trainer_task.remote(
429
+ uid_b64, world_size, tensor_shape, tensor_dtype
430
+ )
431
+
432
+ try:
433
+ trainer_result, result = ray.get(
434
+ [trainer_future, inference_future], timeout=120
435
+ )
436
+ finally:
437
+ ray.shutdown()
438
+
439
+ assert trainer_result["success"], "Trainer should complete successfully"
440
+ assert result["success"], (
441
+ f"Weight transfer failed. "
442
+ f"Received shape: {result['received_shape']}, "
443
+ f"Received sum: {result['received_sum']}"
444
+ )
445
+
446
+
447
+@pytest.mark.skipif(
448
+ torch.accelerator.device_count() < 2,
449
+ reason="Need at least 2 GPUs to run NCCL weight transfer test.",
450
+)
451
+def test_nccl_weight_transfer_between_processes_uid():
452
+ """Weight transfer over a pre-shared ncclUniqueId (no TCPStore), with both
453
+ ranks on vLLM's ``PyNcclCommunicator``."""
454
+ _run_uid_transfer(trainer_broadcast_tensor_uid)
455
+
456
+
457
+@pytest.mark.skipif(
458
+ torch.accelerator.device_count() < 2,
459
+ reason="Need at least 2 GPUs to run NCCL weight transfer test.",
460
+)
461
+def test_nccl_weight_transfer_torch_free_trainer():
462
+ """UID weight transfer against a foreign peer that uses no torch and no vLLM
463
+ communicator, only a cupy buffer and raw ``ncclBroadcast``."""
464
+ pytest.importorskip("cupy")
465
+ _run_uid_transfer(trainer_broadcast_tensor_torch_free)
vllm/distributed/device_communicators/pynccl.py
+68
@@ -11,6 +11,7 @@ from torch.distributed import ProcessGroup, ReduceOp
11
12
import vllm.envs as envs
13
from vllm.distributed.device_communicators.pynccl_wrapper import (
14
+ NCCL_UNIQUE_ID_BYTES,
15
NCCLLibrary,
16
buffer_type,
17
cudaStream_t,
@@ -60,6 +61,9 @@ def register_nccl_symmetric_ops(pynccl_comm):
61
62
63
class PyNcclCommunicator:
64
+ # None for communicators built via `from_unique_id_bytes` (no process group).
65
+ group: ProcessGroup | StatelessProcessGroup | None
66
+
67
def __init__(
68
self,
69
group: ProcessGroup | StatelessProcessGroup,
@@ -128,6 +132,14 @@ class PyNcclCommunicator:
132
self.unique_id.internal[i] = byte
133
else:
134
self.unique_id = group.broadcast_obj(self.unique_id, src=0)
135
+ self._init_comm(device)
136
+
137
+ def _init_comm(self, device: int | str | torch.device) -> None:
138
+ """Create the communicator on `device` from the already-resolved
139
+ `self.unique_id` / `self.rank` / `self.world_size`, then run the
140
+ one-element warm-up all_reduce. Shared by `__init__` and
141
+ `from_unique_id_bytes` so the init handshake stays identical on both.
142
+ """
143
if isinstance(device, int):
144
device = torch.device(f"cuda:{device}")
145
elif isinstance(device, str):
@@ -148,6 +160,62 @@ class PyNcclCommunicator:
160
stream.synchronize()
161
del data
162
163
+ @classmethod
164
+ def from_unique_id_bytes(
165
+ cls,
166
+ unique_id_bytes: bytes,
167
+ rank: int,
168
+ world_size: int,
169
+ device: int | str | torch.device,
170
+ library_path: str | None = None,
171
+ ) -> "PyNcclCommunicator":
172
+ """Build a communicator from pre-shared ``ncclUniqueId`` bytes.
173
+
174
+ For peers that cannot join a ``StatelessProcessGroup`` / TCPStore (e.g. a
175
+ torch-free JAX trainer): every rank passes the same id, minted once via
176
+ ``ncclGetUniqueId`` and shared out of band. There is no barrier, so all
177
+ ranks must enter init concurrently or ``ncclCommInitRank`` hangs.
178
+
179
+ Warm-up handshake: immediately after ``ncclCommInitRank`` this issues a
180
+ one-element ``all_reduce`` (mirroring ``__init__``). It is a collective,
181
+ so every peer -- including a foreign, non-vLLM rank -- must issue a
182
+ matching one-element ``all_reduce`` before any other collective, or all
183
+ ranks deadlock.
184
+ """
185
+ if len(unique_id_bytes) != NCCL_UNIQUE_ID_BYTES:
186
+ raise ValueError(
187
+ f"expected a {NCCL_UNIQUE_ID_BYTES}-byte NCCL unique id, "
188
+ f"got {len(unique_id_bytes)} bytes"
189
+ )
190
+ if not 0 <= rank < world_size:
191
+ raise ValueError(f"rank {rank} out of range for world_size {world_size}")
192
+
193
+ self = cls.__new__(cls)
194
+ self.rank = rank
195
+ self.world_size = world_size
196
+ self.group = None
197
+
198
+ if self.world_size == 1 or envs.VLLM_DISABLE_PYNCCL:
199
+ self.available = False
200
+ self.disabled = True
201
+ return self
202
+ try:
203
+ self.nccl = NCCLLibrary(library_path)
204
+ except Exception as e:
205
+ # Unlike the TCPStore path, silently disabling here leaves the peer
206
+ # blocked in ncclCommInitRank until timeout. The caller explicitly
207
+ # asked to join from a unique id, so fail loudly instead.
208
+ raise RuntimeError(
209
+ "failed to load the NCCL library for unique-id rendezvous"
210
+ ) from e
211
+
212
+ self.available = True
213
+ self.disabled = False
214
+ self.nccl_version = self.nccl.ncclGetRawVersion()
215
+ self.unique_id = self.nccl.unique_id_from_bytes(unique_id_bytes)
216
+ self._init_comm(device)
217
+ return self
218
+
219
def destroy(self):
220
if self.available and not self.disabled:
221
# ncclCommAbort can block until all CUDA graphs that
vllm/distributed/device_communicators/pynccl_wrapper.py
+3
@@ -51,6 +51,9 @@ class ncclUniqueId(ctypes.Structure):
51
_fields_ = [("internal", ctypes.c_byte * 128)]
52
53
54
+NCCL_UNIQUE_ID_BYTES = ctypes.sizeof(ncclUniqueId)
55
+
56
+
57
# Mirror of NCCL's versioned ncclCommProperties_t (nccl_device/core.h,
58
# v2.31.2-1, 144 bytes). ncclCommQueryProperties fills fields gated by the
59
# version the caller declares in props.version, not by props.size, so never
vllm/distributed/weight_transfer/nccl_common.py
+124
-7
@@ -8,15 +8,19 @@ The dense (`NCCLWeightTransferEngine`) and sparse
8
sparse engine does not have to subclass the dense one.
9
"""
10
11
-from dataclasses import dataclass
11
+from dataclasses import asdict, dataclass, field
12
from typing import TYPE_CHECKING, Protocol
13
14
+import pybase64 as base64
15
import torch
16
17
if TYPE_CHECKING:
18
from vllm.config.parallel import ParallelConfig
19
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
20
21
+from vllm.distributed.device_communicators.pynccl_wrapper import (
22
+ NCCL_UNIQUE_ID_BYTES,
23
+)
24
from vllm.distributed.weight_transfer.base import WeightTransferInitInfo
25
from vllm.distributed.weight_transfer.packed_tensor import (
26
DEFAULT_PACKED_BUFFER_SIZE_BYTES,
@@ -24,21 +28,102 @@ from vllm.distributed.weight_transfer.packed_tensor import (
28
)
29
30
27
-@dataclass
31
+def decode_nccl_unique_id(
32
+ *,
33
+ master_address: str | None,
34
+ master_port: int | None,
35
+ nccl_unique_id_b64: str | None,
36
+ ctx: str,
37
+) -> bytes | None:
38
+ """Validate the rendezvous mode and decode a pre-shared unique id."""
39
+ has_uid = bool(nccl_unique_id_b64)
40
+ has_addr = master_address is not None
41
+ has_port = master_port is not None
42
+
43
+ if has_uid and (has_addr or has_port):
44
+ raise ValueError(
45
+ f"{ctx}: pass nccl_unique_id_b64 OR master_address/master_port, not both"
46
+ )
47
+ if not has_uid:
48
+ if not (has_addr and has_port):
49
+ raise ValueError(
50
+ f"{ctx}: need nccl_unique_id_b64, or both master_address and "
51
+ f"master_port (got master_address={master_address!r}, "
52
+ f"master_port={master_port!r})"
53
+ )
54
+ return None
55
+ try:
56
+ decoded = base64.b64decode(nccl_unique_id_b64, validate=True)
57
+ except (ValueError, TypeError) as e:
58
+ raise ValueError(
59
+ f"{ctx}: nccl_unique_id_b64 is not valid standard base64 "
60
+ f"(URL-safe base64 is not accepted)"
61
+ ) from e
62
+ if len(decoded) != NCCL_UNIQUE_ID_BYTES:
63
+ raise ValueError(
64
+ f"{ctx}: nccl_unique_id_b64 must decode to {NCCL_UNIQUE_ID_BYTES} "
65
+ f"bytes, got {len(decoded)}"
66
+ )
67
+ return decoded
68
+
69
+
70
+@dataclass(kw_only=True)
71
class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
29
- """Worker-side initialization info for NCCL-based weight transfer backends."""
72
+ """Worker-side initialization info for NCCL-based weight transfer backends.
73
31
- master_address: str
32
- master_port: int
74
+ Keyword-only (`kw_only`): adding the optional `nccl_unique_id_b64` field
75
+ means the rendezvous fields can no longer keep a fixed positional slot, so a
76
+ stale positional call fails loudly instead of silently swapping arguments.
77
+
78
+ Provide exactly one rendezvous mode:
79
+
80
+ * ``master_address`` + ``master_port`` -- TCPStore / ``StatelessProcessGroup``
81
+ rendezvous (requires torch on every rank, including the trainer), or
82
+ * ``nccl_unique_id_b64`` -- standard (RFC 4648, *not* URL-safe) base64 of the
83
+ 128 raw bytes from ``ncclGetUniqueId``, for torch-free trainers (e.g. JAX)
84
+ that mint the unique id out of band and share it (over HTTP, etc.). Note a
85
+ JAX peer must use ``base64.b64encode``, not ``urlsafe_b64encode``.
86
+
87
+ On the unique-id path all ranks must enter init concurrently (there is no
88
+ store barrier), and every peer must honor the warm-up handshake: the
89
+ worker's communicator issues a one-element ``all_reduce`` immediately after
90
+ ``ncclCommInitRank`` (see ``PyNcclCommunicator.from_unique_id_bytes``), so a
91
+ foreign peer must issue a matching one-element ``all_reduce`` before any
92
+ other collective or all ranks deadlock.
93
+ """
94
+
95
+ master_address: str | None = None
96
+ master_port: int | None = None
97
rank_offset: int
98
world_size: int
99
+ nccl_unique_id_b64: str | None = field(default=None, repr=False)
100
packed: bool = False
101
packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
102
packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
103
104
+ def __post_init__(self) -> None:
105
+ _ = self.nccl_unique_id_bytes
106
+
107
+ @property
108
+ def nccl_unique_id_bytes(self) -> bytes | None:
109
+ return decode_nccl_unique_id(
110
+ master_address=self.master_address,
111
+ master_port=self.master_port,
112
+ nccl_unique_id_b64=self.nccl_unique_id_b64,
113
+ ctx="NCCLWeightTransferInitInfo",
114
+ )
115
+
116
+
117
+def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict:
118
+ """Serialize a worker init info for `init_weight_transfer_engine`, dropping
119
+ the unset rendezvous field (the UID in TCP mode) so the wire payload carries
120
+ only the mode actually in use. Shared by the dense and sparse trainer
121
+ engines so the two cannot drift."""
122
+ return {key: value for key, value in asdict(init_info).items() if value is not None}
123
+
124
125
class NCCLRendezvous(Protocol):
41
- """The three fields `trainer_init` needs to open the trainer endpoint.
126
+ """The TCP rendezvous fields `trainer_init` needs.
127
128
Structural so each backend can keep its own trainer init info next to its
129
engine (dense `NCCLTrainerInitInfo`, `SparseNCCLTrainerInitInfo`) without
@@ -73,6 +158,27 @@ def stateless_init_process_group(
158
return PyNcclCommunicator(pg, device=device)
159
160
161
+def uid_init_process_group(
162
+ nccl_unique_id_bytes: bytes,
163
+ rank: int,
164
+ world_size: int,
165
+ device,
166
+) -> "PyNcclCommunicator":
167
+ """Join the NCCL group from pre-shared ``ncclUniqueId`` bytes.
168
+
169
+ The torch-free rendezvous alternative to `stateless_init_process_group`: no
170
+ TCPStore, and therefore no barrier -- every rank must enter concurrently.
171
+ """
172
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
173
+
174
+ return PyNcclCommunicator.from_unique_id_bytes(
175
+ nccl_unique_id_bytes,
176
+ rank=rank,
177
+ world_size=world_size,
178
+ device=device,
179
+ )
180
+
181
+
182
def worker_init_process_group(
183
init_info: NCCLWeightTransferInitInfo,
184
parallel_config: "ParallelConfig",
@@ -80,7 +186,7 @@ def worker_init_process_group(
186
"""Create the trainer<->worker NCCL group on an inference worker.
187
188
Computes a unique rank for this worker across all data-parallel groups and
83
- joins the stateless process group with the trainer.
189
+ joins the trainer via whichever rendezvous mode `init_info` carries.
190
"""
191
# Calculate the global rank in the trainer-worker process group.
192
# Must account for data parallel to get unique ranks across all workers.
@@ -93,6 +199,17 @@ def worker_init_process_group(
199
rank = worker_rank + init_info.rank_offset
200
201
device = torch.accelerator.current_device_index()
202
+ unique_id_bytes = init_info.nccl_unique_id_bytes
203
+ if unique_id_bytes is not None:
204
+ return uid_init_process_group(
205
+ unique_id_bytes,
206
+ rank,
207
+ init_info.world_size,
208
+ device,
209
+ )
210
+ # __post_init__ guarantees the TCP pair when there is no unique id.
211
+ assert init_info.master_address is not None
212
+ assert init_info.master_port is not None
213
return stateless_init_process_group(
214
init_info.master_address,
215
init_info.master_port,
vllm/distributed/weight_transfer/nccl_engine.py
+10
-1
@@ -26,6 +26,7 @@ from vllm.distributed.weight_transfer.base import (
26
)
27
from vllm.distributed.weight_transfer.nccl_common import (
28
NCCLWeightTransferInitInfo,
29
+ worker_init_payload,
30
worker_init_process_group,
31
)
32
from vllm.distributed.weight_transfer.nccl_common import (
@@ -57,6 +58,13 @@ class NCCLTrainerInitInfo(TrainerInitInfo):
58
`world_size` is the full trainer+worker NCCL group size. `rank` (from
59
`TrainerInitInfo`) identifies this trainer process; rank 0 is the sender.
60
61
+ The trainer joins over a TCPStore rendezvous (`master_address` +
62
+ `master_port`). Torch-free trainers (e.g. JAX) that cannot join a TCPStore
63
+ mint an `ncclUniqueId` themselves and drive rank 0 out of band; they ship a
64
+ `nccl_unique_id_b64` payload straight to the inference workers'
65
+ `init_weight_transfer_engine` (see `NCCLWeightTransferInitInfo`) rather than
66
+ going through this engine.
67
+
68
`packed` / buffer sizes are the transfer's wire params. The trainer
69
propagates them to the worker at `trainer_init` so the two sides cannot
70
disagree. Note this defaults to packed, unlike the worker-side
@@ -301,7 +309,8 @@ class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerIni
309
# open the trainer endpoint (rank 0); both sides must rendezvous together.
310
with ThreadPoolExecutor(max_workers=1) as exe:
311
future = exe.submit(
304
- engine.client.init_weight_transfer_engine, asdict(worker_init_info)
312
+ engine.client.init_weight_transfer_engine,
313
+ worker_init_payload(worker_init_info),
314
)
315
engine.model_update_group = open_trainer_endpoint(init_info)
316
future.result() # surface any inference-side init error
vllm/distributed/weight_transfer/sparse_nccl_engine.py
+3
-1
@@ -31,6 +31,7 @@ from vllm.distributed.weight_transfer.base import (
31
)
32
from vllm.distributed.weight_transfer.nccl_common import (
33
NCCLWeightTransferInitInfo,
34
+ worker_init_payload,
35
worker_init_process_group,
36
)
37
from vllm.distributed.weight_transfer.nccl_common import (
@@ -256,7 +257,8 @@ class SparseNCCLTrainerWeightTransferEngine(
257
# open the trainer endpoint (rank 0); both sides must rendezvous together.
258
with ThreadPoolExecutor(max_workers=1) as exe:
259
future = exe.submit(
259
- engine.client.init_weight_transfer_engine, asdict(worker_init_info)
260
+ engine.client.init_weight_transfer_engine,
261
+ worker_init_payload(worker_init_info),
262
)
263
# Open the trainer endpoint as NCCL rank 0 on the current device
264
# (the init info satisfies the helper's rendezvous protocol).