Skip to content

vllm.distributed.weight_transfer.nccl_common

Shared NCCL initialization helpers for weight transfer engines.

The dense (NCCLWeightTransferEngine) and sparse (SparseNCCLWeightTransferEngine) backends are independent engines that share only their process-group initialization. That common logic lives here so the sparse engine does not have to subclass the dense one.

Classes:

Functions:

NCCLRendezvous

Bases: Protocol

The TCP rendezvous fields trainer_init needs.

Structural so each backend can keep its own trainer init info next to its engine (dense NCCLTrainerInitInfo, SparseNCCLTrainerInitInfo) without this shared module importing either.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
126
127
128
129
130
131
132
133
134
135
136
class NCCLRendezvous(Protocol):
    """The TCP rendezvous fields `trainer_init` needs.

    Structural so each backend can keep its own trainer init info next to its
    engine (dense `NCCLTrainerInitInfo`, `SparseNCCLTrainerInitInfo`) without
    this shared module importing either.
    """

    master_address: str
    master_port: int
    world_size: int

NCCLWeightTransferInitInfo dataclass

Bases: WeightTransferInitInfo

Worker-side initialization info for NCCL-based weight transfer backends.

Keyword-only (kw_only): adding the optional nccl_unique_id_b64 field means the rendezvous fields can no longer keep a fixed positional slot, so a stale positional call fails loudly instead of silently swapping arguments.

Provide exactly one rendezvous mode:

  • master_address + master_port -- TCPStore / StatelessProcessGroup rendezvous (requires torch on every rank, including the trainer), or
  • nccl_unique_id_b64 -- standard (RFC 4648, not URL-safe) base64 of the 128 raw bytes from ncclGetUniqueId, for torch-free trainers (e.g. JAX) that mint the unique id out of band and share it (over HTTP, etc.). Note a JAX peer must use base64.b64encode, not urlsafe_b64encode.

On the unique-id path all ranks must enter init concurrently (there is no store barrier), and every peer must honor the warm-up handshake: the worker's communicator issues a one-element all_reduce immediately after ncclCommInitRank (see PyNcclCommunicator.from_unique_id_bytes), so a foreign peer must issue a matching one-element all_reduce before any other collective or all ranks deadlock.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@dataclass(kw_only=True)
class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
    """Worker-side initialization info for NCCL-based weight transfer backends.

    Keyword-only (`kw_only`): adding the optional `nccl_unique_id_b64` field
    means the rendezvous fields can no longer keep a fixed positional slot, so a
    stale positional call fails loudly instead of silently swapping arguments.

    Provide exactly one rendezvous mode:

    * ``master_address`` + ``master_port`` -- TCPStore / ``StatelessProcessGroup``
      rendezvous (requires torch on every rank, including the trainer), or
    * ``nccl_unique_id_b64`` -- standard (RFC 4648, *not* URL-safe) base64 of the
      128 raw bytes from ``ncclGetUniqueId``, for torch-free trainers (e.g. JAX)
      that mint the unique id out of band and share it (over HTTP, etc.). Note a
      JAX peer must use ``base64.b64encode``, not ``urlsafe_b64encode``.

    On the unique-id path all ranks must enter init concurrently (there is no
    store barrier), and every peer must honor the warm-up handshake: the
    worker's communicator issues a one-element ``all_reduce`` immediately after
    ``ncclCommInitRank`` (see ``PyNcclCommunicator.from_unique_id_bytes``), so a
    foreign peer must issue a matching one-element ``all_reduce`` before any
    other collective or all ranks deadlock.
    """

    master_address: str | None = None
    master_port: int | None = None
    rank_offset: int
    world_size: int
    nccl_unique_id_b64: str | None = field(default=None, repr=False)
    packed: bool = False
    packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
    packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS

    def __post_init__(self) -> None:
        _ = self.nccl_unique_id_bytes

    @property
    def nccl_unique_id_bytes(self) -> bytes | None:
        return decode_nccl_unique_id(
            master_address=self.master_address,
            master_port=self.master_port,
            nccl_unique_id_b64=self.nccl_unique_id_b64,
            ctx="NCCLWeightTransferInitInfo",
        )

decode_nccl_unique_id(*, master_address, master_port, nccl_unique_id_b64, ctx)

Validate the rendezvous mode and decode a pre-shared unique id.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def decode_nccl_unique_id(
    *,
    master_address: str | None,
    master_port: int | None,
    nccl_unique_id_b64: str | None,
    ctx: str,
) -> bytes | None:
    """Validate the rendezvous mode and decode a pre-shared unique id."""
    has_uid = bool(nccl_unique_id_b64)
    has_addr = master_address is not None
    has_port = master_port is not None

    if has_uid and (has_addr or has_port):
        raise ValueError(
            f"{ctx}: pass nccl_unique_id_b64 OR master_address/master_port, not both"
        )
    if not has_uid:
        if not (has_addr and has_port):
            raise ValueError(
                f"{ctx}: need nccl_unique_id_b64, or both master_address and "
                f"master_port (got master_address={master_address!r}, "
                f"master_port={master_port!r})"
            )
        return None
    try:
        decoded = base64.b64decode(nccl_unique_id_b64, validate=True)
    except (ValueError, TypeError) as e:
        raise ValueError(
            f"{ctx}: nccl_unique_id_b64 is not valid standard base64 "
            f"(URL-safe base64 is not accepted)"
        ) from e
    if len(decoded) != NCCL_UNIQUE_ID_BYTES:
        raise ValueError(
            f"{ctx}: nccl_unique_id_b64 must decode to {NCCL_UNIQUE_ID_BYTES} "
            f"bytes, got {len(decoded)}"
        )
    return decoded

stateless_init_process_group(master_address, master_port, rank, world_size, device)

VLLM provides StatelessProcessGroup to create a process group without considering the global process group in torch.distributed. It is recommended to create StatelessProcessGroup, and then initialize the data-plane communication (NCCL) between external (train processes) and vLLM workers.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def stateless_init_process_group(
    master_address: str,
    master_port: int,
    rank: int,
    world_size: int,
    device,
) -> "PyNcclCommunicator":
    """VLLM provides `StatelessProcessGroup` to create a process group
    without considering the global process group in torch.distributed.
    It is recommended to create `StatelessProcessGroup`, and then initialize
    the data-plane communication (NCCL) between external (train processes)
    and vLLM workers.
    """
    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
    from vllm.distributed.utils import StatelessProcessGroup

    pg = StatelessProcessGroup.create(
        host=master_address, port=master_port, rank=rank, world_size=world_size
    )
    with unpinned_nccl_env():
        return PyNcclCommunicator(pg, device=device)

trainer_init(init_info, rank=0)

Initialize NCCL process group for trainer-side weight transfer.

Uses the current CUDA device (torch.accelerator.current_device_index()).

Parameters:

  • init_info

    (NCCLRendezvous | dict) –

    Any object carrying the NCCLRendezvous fields (a trainer or worker NCCL init info), or a dict with keys: - master_address: str - master_port: int - world_size: int

  • rank

    (int, default: 0 ) –

    This trainer process's rank in the group. The broadcast backends have a single trainer at rank 0; multi-rank trainers (m2n) occupy ranks [0, num_trainer_ranks) and the workers start after them.

Returns:

Source code in vllm/distributed/weight_transfer/nccl_common.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def trainer_init(
    init_info: NCCLRendezvous | dict,
    rank: int = 0,
) -> "PyNcclCommunicator":
    """Initialize NCCL process group for trainer-side weight transfer.

    Uses the current CUDA device (torch.accelerator.current_device_index()).

    Args:
        init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or
            worker NCCL init info), or a dict with keys:
            - master_address: str
            - master_port: int
            - world_size: int
        rank: This trainer process's rank in the group. The broadcast backends
            have a single trainer at rank 0; multi-rank trainers (m2n) occupy
            ranks `[0, num_trainer_ranks)` and the workers start after them.

    Returns:
        PyNcclCommunicator for weight transfer.

    """
    if isinstance(init_info, dict):
        master_address = init_info["master_address"]
        master_port = init_info["master_port"]
        world_size = init_info["world_size"]
    else:
        master_address = init_info.master_address
        master_port = init_info.master_port
        world_size = init_info.world_size

    device = torch.accelerator.current_device_index()
    return stateless_init_process_group(
        master_address,
        master_port,
        rank,
        world_size,
        device,
    )

uid_init_process_group(nccl_unique_id_bytes, rank, world_size, device)

Join the NCCL group from pre-shared ncclUniqueId bytes.

The torch-free rendezvous alternative to stateless_init_process_group: no TCPStore, and therefore no barrier -- every rank must enter concurrently.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def uid_init_process_group(
    nccl_unique_id_bytes: bytes,
    rank: int,
    world_size: int,
    device,
) -> "PyNcclCommunicator":
    """Join the NCCL group from pre-shared ``ncclUniqueId`` bytes.

    The torch-free rendezvous alternative to `stateless_init_process_group`: no
    TCPStore, and therefore no barrier -- every rank must enter concurrently.
    """
    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

    with unpinned_nccl_env():
        return PyNcclCommunicator.from_unique_id_bytes(
            nccl_unique_id_bytes,
            rank=rank,
            world_size=world_size,
            device=device,
        )

worker_init_payload(init_info)

Serialize a worker init info for init_weight_transfer_engine, dropping the unset rendezvous field (the UID in TCP mode) so the wire payload carries only the mode actually in use. Shared by the dense and sparse trainer engines so the two cannot drift.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
118
119
120
121
122
123
def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict:
    """Serialize a worker init info for `init_weight_transfer_engine`, dropping
    the unset rendezvous field (the UID in TCP mode) so the wire payload carries
    only the mode actually in use. Shared by the dense and sparse trainer
    engines so the two cannot drift."""
    return {key: value for key, value in asdict(init_info).items() if value is not None}

worker_init_process_group(init_info, parallel_config)

Create the trainer<->worker NCCL group on an inference worker.

Computes a unique rank for this worker across all data-parallel groups and joins the trainer via whichever rendezvous mode init_info carries.

Source code in vllm/distributed/weight_transfer/nccl_common.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def worker_init_process_group(
    init_info: NCCLWeightTransferInitInfo,
    parallel_config: "ParallelConfig",
) -> "PyNcclCommunicator":
    """Create the trainer<->worker NCCL group on an inference worker.

    Computes a unique rank for this worker across all data-parallel groups and
    joins the trainer via whichever rendezvous mode `init_info` carries.
    """
    # Calculate the global rank in the trainer-worker process group.
    # Must account for data parallel to get unique ranks across all workers.
    dp_rank = parallel_config.data_parallel_index
    world_size_per_dp = parallel_config.world_size  # TP * PP
    rank_within_dp = parallel_config.rank

    # Unique rank across all DP groups
    worker_rank = dp_rank * world_size_per_dp + rank_within_dp
    rank = worker_rank + init_info.rank_offset

    device = torch.accelerator.current_device_index()
    unique_id_bytes = init_info.nccl_unique_id_bytes
    if unique_id_bytes is not None:
        return uid_init_process_group(
            unique_id_bytes,
            rank,
            init_info.world_size,
            device,
        )
    # __post_init__ guarantees the TCP pair when there is no unique id.
    assert init_info.master_address is not None
    assert init_info.master_port is not None
    return stateless_init_process_group(
        init_info.master_address,
        init_info.master_port,
        rank,
        init_info.world_size,
        device=device,
    )