@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 262 lines 9.37 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Shared NCCL initialization helpers for weight transfer engines.
4
5 The dense (`NCCLWeightTransferEngine`) and sparse
6 (`SparseNCCLWeightTransferEngine`) backends are independent engines that share
7 *only* their process-group initialization. That common logic lives here so the
8 sparse engine does not have to subclass the dense one.
9 """
10
11 from dataclasses import 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,
27 DEFAULT_PACKED_NUM_BUFFERS,
28 )
29 from vllm.utils.nccl import unpinned_nccl_env
30
31
32 def decode_nccl_unique_id(
33 *,
34 master_address: str | None,
35 master_port: int | None,
36 nccl_unique_id_b64: str | None,
37 ctx: str,
38 ) -> bytes | None:
39 """Validate the rendezvous mode and decode a pre-shared unique id."""
40 has_uid = bool(nccl_unique_id_b64)
41 has_addr = master_address is not None
42 has_port = master_port is not None
43
44 if has_uid and (has_addr or has_port):
45 raise ValueError(
46 f"{ctx}: pass nccl_unique_id_b64 OR master_address/master_port, not both"
47 )
48 if not has_uid:
49 if not (has_addr and has_port):
50 raise ValueError(
51 f"{ctx}: need nccl_unique_id_b64, or both master_address and "
52 f"master_port (got master_address={master_address!r}, "
53 f"master_port={master_port!r})"
54 )
55 return None
56 try:
57 decoded = base64.b64decode(nccl_unique_id_b64, validate=True)
58 except (ValueError, TypeError) as e:
59 raise ValueError(
60 f"{ctx}: nccl_unique_id_b64 is not valid standard base64 "
61 f"(URL-safe base64 is not accepted)"
62 ) from e
63 if len(decoded) != NCCL_UNIQUE_ID_BYTES:
64 raise ValueError(
65 f"{ctx}: nccl_unique_id_b64 must decode to {NCCL_UNIQUE_ID_BYTES} "
66 f"bytes, got {len(decoded)}"
67 )
68 return decoded
69
70
71 @dataclass(kw_only=True)
72 class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
73 """Worker-side initialization info for NCCL-based weight transfer backends.
74
75 Keyword-only (`kw_only`): adding the optional `nccl_unique_id_b64` field
76 means the rendezvous fields can no longer keep a fixed positional slot, so a
77 stale positional call fails loudly instead of silently swapping arguments.
78
79 Provide exactly one rendezvous mode:
80
81 * ``master_address`` + ``master_port`` -- TCPStore / ``StatelessProcessGroup``
82 rendezvous (requires torch on every rank, including the trainer), or
83 * ``nccl_unique_id_b64`` -- standard (RFC 4648, *not* URL-safe) base64 of the
84 128 raw bytes from ``ncclGetUniqueId``, for torch-free trainers (e.g. JAX)
85 that mint the unique id out of band and share it (over HTTP, etc.). Note a
86 JAX peer must use ``base64.b64encode``, not ``urlsafe_b64encode``.
87
88 On the unique-id path all ranks must enter init concurrently (there is no
89 store barrier), and every peer must honor the warm-up handshake: the
90 worker's communicator issues a one-element ``all_reduce`` immediately after
91 ``ncclCommInitRank`` (see ``PyNcclCommunicator.from_unique_id_bytes``), so a
92 foreign peer must issue a matching one-element ``all_reduce`` before any
93 other collective or all ranks deadlock.
94 """
95
96 master_address: str | None = None
97 master_port: int | None = None
98 rank_offset: int
99 world_size: int
100 nccl_unique_id_b64: str | None = field(default=None, repr=False)
101 packed: bool = False
102 packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
103 packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
104
105 def __post_init__(self) -> None:
106 _ = self.nccl_unique_id_bytes
107
108 @property
109 def nccl_unique_id_bytes(self) -> bytes | None:
110 return decode_nccl_unique_id(
111 master_address=self.master_address,
112 master_port=self.master_port,
113 nccl_unique_id_b64=self.nccl_unique_id_b64,
114 ctx="NCCLWeightTransferInitInfo",
115 )
116
117
118 def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict:
119 """Serialize a worker init info for `init_weight_transfer_engine`, dropping
120 the unset rendezvous field (the UID in TCP mode) so the wire payload carries
121 only the mode actually in use. Shared by the dense and sparse trainer
122 engines so the two cannot drift."""
123 return {key: value for key, value in asdict(init_info).items() if value is not None}
124
125
126 class NCCLRendezvous(Protocol):
127 """The TCP rendezvous fields `trainer_init` needs.
128
129 Structural so each backend can keep its own trainer init info next to its
130 engine (dense `NCCLTrainerInitInfo`, `SparseNCCLTrainerInitInfo`) without
131 this shared module importing either.
132 """
133
134 master_address: str
135 master_port: int
136 world_size: int
137
138
139 def stateless_init_process_group(
140 master_address: str,
141 master_port: int,
142 rank: int,
143 world_size: int,
144 device,
145 ) -> "PyNcclCommunicator":
146 """VLLM provides `StatelessProcessGroup` to create a process group
147 without considering the global process group in torch.distributed.
148 It is recommended to create `StatelessProcessGroup`, and then initialize
149 the data-plane communication (NCCL) between external (train processes)
150 and vLLM workers.
151 """
152 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
153 from vllm.distributed.utils import StatelessProcessGroup
154
155 pg = StatelessProcessGroup.create(
156 host=master_address, port=master_port, rank=rank, world_size=world_size
157 )
158 with unpinned_nccl_env():
159 return PyNcclCommunicator(pg, device=device)
160
161
162 def uid_init_process_group(
163 nccl_unique_id_bytes: bytes,
164 rank: int,
165 world_size: int,
166 device,
167 ) -> "PyNcclCommunicator":
168 """Join the NCCL group from pre-shared ``ncclUniqueId`` bytes.
169
170 The torch-free rendezvous alternative to `stateless_init_process_group`: no
171 TCPStore, and therefore no barrier -- every rank must enter concurrently.
172 """
173 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
174
175 with unpinned_nccl_env():
176 return PyNcclCommunicator.from_unique_id_bytes(
177 nccl_unique_id_bytes,
178 rank=rank,
179 world_size=world_size,
180 device=device,
181 )
182
183
184 def worker_init_process_group(
185 init_info: NCCLWeightTransferInitInfo,
186 parallel_config: "ParallelConfig",
187 ) -> "PyNcclCommunicator":
188 """Create the trainer<->worker NCCL group on an inference worker.
189
190 Computes a unique rank for this worker across all data-parallel groups and
191 joins the trainer via whichever rendezvous mode `init_info` carries.
192 """
193 # Calculate the global rank in the trainer-worker process group.
194 # Must account for data parallel to get unique ranks across all workers.
195 dp_rank = parallel_config.data_parallel_index
196 world_size_per_dp = parallel_config.world_size # TP * PP
197 rank_within_dp = parallel_config.rank
198
199 # Unique rank across all DP groups
200 worker_rank = dp_rank * world_size_per_dp + rank_within_dp
201 rank = worker_rank + init_info.rank_offset
202
203 device = torch.accelerator.current_device_index()
204 unique_id_bytes = init_info.nccl_unique_id_bytes
205 if unique_id_bytes is not None:
206 return uid_init_process_group(
207 unique_id_bytes,
208 rank,
209 init_info.world_size,
210 device,
211 )
212 # __post_init__ guarantees the TCP pair when there is no unique id.
213 assert init_info.master_address is not None
214 assert init_info.master_port is not None
215 return stateless_init_process_group(
216 init_info.master_address,
217 init_info.master_port,
218 rank,
219 init_info.world_size,
220 device=device,
221 )
222
223
224 def trainer_init(
225 init_info: NCCLRendezvous | dict,
226 rank: int = 0,
227 ) -> "PyNcclCommunicator":
228 """Initialize NCCL process group for trainer-side weight transfer.
229
230 Uses the current CUDA device (torch.accelerator.current_device_index()).
231
232 Args:
233 init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or
234 worker NCCL init info), or a dict with keys:
235 - master_address: str
236 - master_port: int
237 - world_size: int
238 rank: This trainer process's rank in the group. The broadcast backends
239 have a single trainer at rank 0; multi-rank trainers (m2n) occupy
240 ranks `[0, num_trainer_ranks)` and the workers start after them.
241
242 Returns:
243 PyNcclCommunicator for weight transfer.
244
245 """
246 if isinstance(init_info, dict):
247 master_address = init_info["master_address"]
248 master_port = init_info["master_port"]
249 world_size = init_info["world_size"]
250 else:
251 master_address = init_info.master_address
252 master_port = init_info.master_port
253 world_size = init_info.world_size
254
255 device = torch.accelerator.current_device_index()
256 return stateless_init_process_group(
257 master_address,
258 master_port,
259 rank,
260 world_size,
261 device,
262 )