| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Trainer-side weight transfer engine for the NCCL M2N backend. |
| 4 | |
| 5 | Symmetric to `M2NWeightTransferEngine` but in the training process. Every |
| 6 | trainer rank joins the shared communicator and runs every reshard, sending its |
| 7 | own local shard; only rank 0 touches the inference control plane. |
| 8 | |
| 9 | Workers join the communicator during init and run reshard from inside |
| 10 | `update_weights`, so those two RPCs overlap with work on this side. Both run on |
| 11 | a helper thread and are joined afterwards, which keeps that overlap inside the |
| 12 | engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather than |
| 13 | in every caller. |
| 14 | """ |
| 15 | |
| 16 | from concurrent.futures import Future, ThreadPoolExecutor |
| 17 | from dataclasses import dataclass, field |
| 18 | from typing import TYPE_CHECKING, Any, ClassVar, cast |
| 19 | |
| 20 | import torch |
| 21 | from typing_extensions import Self |
| 22 | |
| 23 | from vllm.distributed.weight_transfer.base import ( |
| 24 | TrainerInitInfo, |
| 25 | TrainerWeightTransferEngine, |
| 26 | VLLMWeightSyncClient, |
| 27 | WeightSource, |
| 28 | ) |
| 29 | from vllm.distributed.weight_transfer.m2n_common import ( |
| 30 | M2NMesh, |
| 31 | M2NParamMeta, |
| 32 | Placements, |
| 33 | check_transferable, |
| 34 | comm_ptr, |
| 35 | import_m2n, |
| 36 | publish_destination_placements, |
| 37 | resolve_layout, |
| 38 | to_mesh, |
| 39 | to_placements, |
| 40 | validate_layout, |
| 41 | ) |
| 42 | from vllm.distributed.weight_transfer.m2n_source import M2NWeightSource |
| 43 | from vllm.distributed.weight_transfer.nccl_common import ( |
| 44 | decode_nccl_unique_id, |
| 45 | trainer_init, |
| 46 | uid_init_process_group, |
| 47 | ) |
| 48 | from vllm.logger import init_logger |
| 49 | |
| 50 | if TYPE_CHECKING: |
| 51 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 52 | |
| 53 | logger = init_logger(__name__) |
| 54 | |
| 55 | __all__ = ["M2NTrainerInitInfo", "M2NTrainerWeightTransferEngine"] |
| 56 | |
| 57 | |
| 58 | def _dtype_name(dtype: torch.dtype) -> str: |
| 59 | """Short name for a dtype, e.g. `bfloat16`, for the worker-side init dict.""" |
| 60 | return str(dtype).split(".")[-1] |
| 61 | |
| 62 | |
| 63 | @dataclass(kw_only=True) |
| 64 | class M2NTrainerInitInfo(TrainerInitInfo): |
| 65 | """Trainer-side init info for nccl_m2n. |
| 66 | |
| 67 | `rank` (from `TrainerInitInfo`) is this trainer process's rank; it is also |
| 68 | its rank in the shared communicator, since the trainer occupies |
| 69 | `[0, num_trainer_ranks)`. Rank 0 drives the control plane. |
| 70 | """ |
| 71 | |
| 72 | backend: ClassVar[str] = "nccl_m2n" |
| 73 | |
| 74 | master_address: str | None = None |
| 75 | master_port: int | None = None |
| 76 | world_size: int |
| 77 | """Trainer ranks + all inference workers.""" |
| 78 | num_trainer_ranks: int = 1 |
| 79 | dst_mesh_dims: tuple[int, int] | None = None |
| 80 | """How the inference ranks are laid out: axis 0 replicates and axis 1 |
| 81 | shards. The trainer declares it so both sides describe the destination |
| 82 | identically; defaults to a flat `(num_workers, 1)`, which is all a |
| 83 | replicated destination needs.""" |
| 84 | max_cta: int | None = None |
| 85 | nccl_unique_id_b64: str | None = field(default=None, repr=False) |
| 86 | |
| 87 | @property |
| 88 | def destination_mesh_dims(self) -> tuple[int, int]: |
| 89 | """`dst_mesh_dims`, or a flat mesh over every inference worker.""" |
| 90 | if self.dst_mesh_dims is not None: |
| 91 | return self.dst_mesh_dims |
| 92 | return (self.world_size - self.num_trainer_ranks, 1) |
| 93 | |
| 94 | def __post_init__(self) -> None: |
| 95 | """Reject rank counts or destination meshes that cannot form a group.""" |
| 96 | _ = self.nccl_unique_id_bytes |
| 97 | if not 0 < self.num_trainer_ranks < self.world_size: |
| 98 | raise ValueError( |
| 99 | f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at " |
| 100 | f"least one worker in world_size {self.world_size}" |
| 101 | ) |
| 102 | num_workers = self.world_size - self.num_trainer_ranks |
| 103 | if self.dst_mesh_dims is not None: |
| 104 | dims = tuple(self.dst_mesh_dims) |
| 105 | if len(dims) != 2 or dims[0] * dims[1] != num_workers: |
| 106 | raise ValueError( |
| 107 | f"`dst_mesh_dims` {dims} must be two dims covering the " |
| 108 | f"{num_workers} inference workers" |
| 109 | ) |
| 110 | if not 0 <= self.rank < self.num_trainer_ranks: |
| 111 | raise ValueError( |
| 112 | f"trainer `rank` ({self.rank}) must be below " |
| 113 | f"`num_trainer_ranks` ({self.num_trainer_ranks})" |
| 114 | ) |
| 115 | |
| 116 | @property |
| 117 | def nccl_unique_id_bytes(self) -> bytes | None: |
| 118 | return decode_nccl_unique_id( |
| 119 | master_address=self.master_address, |
| 120 | master_port=self.master_port, |
| 121 | nccl_unique_id_b64=self.nccl_unique_id_b64, |
| 122 | ctx="M2NTrainerInitInfo", |
| 123 | ) |
| 124 | |
| 125 | |
| 126 | class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitInfo]): |
| 127 | """Trainer-side engine: sends every rank's local shard, once per parameter. |
| 128 | |
| 129 | Called on every trainer rank. All ranks join the communicator and run every |
| 130 | reshard; only rank 0 touches the control plane. |
| 131 | """ |
| 132 | |
| 133 | init_info_cls = M2NTrainerInitInfo |
| 134 | |
| 135 | def __init__( |
| 136 | self, |
| 137 | *, |
| 138 | client: VLLMWeightSyncClient, |
| 139 | source: M2NWeightSource, |
| 140 | is_controller: bool, |
| 141 | num_trainer_ranks: int, |
| 142 | ) -> None: |
| 143 | """Hold the client and source; `trainer_init` fills in the group.""" |
| 144 | super().__init__(client=client, source=source, is_sender=is_controller) |
| 145 | self.source: M2NWeightSource = source |
| 146 | self.is_controller = is_controller |
| 147 | self.num_trainer_ranks = num_trainer_ranks |
| 148 | self.group: PyNcclCommunicator | None = None |
| 149 | self._m2n: Any = None |
| 150 | self._handle: Any = None |
| 151 | self._metas: list[M2NParamMeta] = [] |
| 152 | self._src_mesh: M2NMesh | None = None |
| 153 | self._dst_mesh: M2NMesh | None = None |
| 154 | self._dst_placements: list[Placements | None] = [] |
| 155 | self._executor: ThreadPoolExecutor | None = None |
| 156 | |
| 157 | @classmethod |
| 158 | def trainer_init( |
| 159 | cls, |
| 160 | init_info: M2NTrainerInitInfo, |
| 161 | *, |
| 162 | client: VLLMWeightSyncClient, |
| 163 | source: WeightSource | None = None, |
| 164 | ) -> Self: |
| 165 | """Build the engine and rendezvous with the inference side. |
| 166 | |
| 167 | Runs on *every* trainer rank. Rank 0 additionally ships the transfer |
| 168 | plan to the workers and drives the control plane; the other ranks only |
| 169 | build local state and join the communicator. Every rank participates |
| 170 | in every reshard and sends its local shard. |
| 171 | |
| 172 | The ordering here is not arbitrary -- see the comments inline. In |
| 173 | short: validate before anything blocks; start the init RPC without |
| 174 | waiting (workers join the communicator inside that RPC, so it cannot |
| 175 | return until this side has joined too); then join the NCCL |
| 176 | communicator; then create the M2N handle; and only then wait for the RPC |
| 177 | to finish. |
| 178 | """ |
| 179 | # The base WeightSource carries name/dtype/shape but not layout, which |
| 180 | # is not enough to plan a reshard. Fail here rather than at the first |
| 181 | # missing attribute deep inside the send loop. |
| 182 | if not isinstance(source, M2NWeightSource): |
| 183 | raise TypeError( |
| 184 | "nccl_m2n needs per-parameter layouts, so its source must be an " |
| 185 | f"M2NWeightSource (e.g. DTensorModuleSource); got " |
| 186 | f"{type(source).__name__}" |
| 187 | ) |
| 188 | |
| 189 | # Trainer-side weight transfer engine for this rank. `__init__` only |
| 190 | # holds the client and source; the communicator, meshes, and M2N handle |
| 191 | # are filled in below. |
| 192 | engine = cls( |
| 193 | client=client, |
| 194 | source=source, |
| 195 | is_controller=init_info.rank == 0, |
| 196 | num_trainer_ranks=init_info.num_trainer_ranks, |
| 197 | ) |
| 198 | engine._m2n = import_m2n() |
| 199 | |
| 200 | engine._src_mesh = source.mesh() |
| 201 | if engine._src_mesh.size != init_info.num_trainer_ranks: |
| 202 | raise ValueError( |
| 203 | f"source mesh covers {engine._src_mesh.size} ranks, but there " |
| 204 | f"are {init_info.num_trainer_ranks} trainer ranks" |
| 205 | ) |
| 206 | |
| 207 | # Validate the whole plan before any rendezvous. A layout m2n cannot |
| 208 | # express is a mismatched collective at send time, which hangs rather |
| 209 | # than raises -- so every precondition is checked while an exception can |
| 210 | # still propagate to the caller. |
| 211 | engine._metas = cast(list[M2NParamMeta], source.metadata()) |
| 212 | for meta in engine._metas: |
| 213 | check_transferable(meta.name, meta.dtype, meta.shape) |
| 214 | mesh, placements = resolve_layout(engine._src_mesh, meta.placements) |
| 215 | validate_layout(mesh, placements, meta.shape, "source") |
| 216 | |
| 217 | rendezvous: Future | None = None |
| 218 | if engine.is_controller: |
| 219 | engine._executor = ThreadPoolExecutor( |
| 220 | max_workers=1, thread_name_prefix="m2n-weight-sync" |
| 221 | ) |
| 222 | # Only trainer rank 0 starts the control-plane RPC. Workers join |
| 223 | # the communicator inside it, so submit it first, join below, and |
| 224 | # wait for it at the end. |
| 225 | rendezvous = engine._executor.submit( |
| 226 | client.init_weight_transfer_engine, |
| 227 | engine._worker_init_info(init_info), |
| 228 | ) |
| 229 | |
| 230 | # Every trainer rank joins -- reshard is collective over the whole |
| 231 | # communicator. Trainer ranks take |
| 232 | # [0, num_trainer_ranks) and the workers follow, which is the |
| 233 | # contiguous-interval layout m2n requires of both meshes. |
| 234 | unique_id_bytes = init_info.nccl_unique_id_bytes |
| 235 | if unique_id_bytes is not None: |
| 236 | engine.group = uid_init_process_group( |
| 237 | unique_id_bytes, |
| 238 | rank=init_info.rank, |
| 239 | world_size=init_info.world_size, |
| 240 | device=torch.accelerator.current_device_index(), |
| 241 | ) |
| 242 | else: |
| 243 | assert init_info.master_address is not None |
| 244 | assert init_info.master_port is not None |
| 245 | engine.group = trainer_init( |
| 246 | { |
| 247 | "master_address": init_info.master_address, |
| 248 | "master_port": init_info.master_port, |
| 249 | "world_size": init_info.world_size, |
| 250 | }, |
| 251 | rank=init_info.rank, |
| 252 | ) |
| 253 | |
| 254 | # The trainer declares the inference topology; the worker uses exactly |
| 255 | # this, so neither side has to infer the other's factorization. |
| 256 | engine._dst_mesh = M2NMesh( |
| 257 | init_info.destination_mesh_dims, init_info.num_trainer_ranks |
| 258 | ) |
| 259 | engine._dst_placements = engine._receive_destination_plan( |
| 260 | init_info.num_trainer_ranks |
| 261 | ) |
| 262 | logger.info( |
| 263 | "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), " |
| 264 | "inference ranks [%d, %d)", |
| 265 | len(engine._metas), |
| 266 | init_info.num_trainer_ranks, |
| 267 | init_info.num_trainer_ranks, |
| 268 | init_info.world_size, |
| 269 | ) |
| 270 | |
| 271 | # Created after the group so a failed rendezvous does not leave a live |
| 272 | # handle behind; shutdown() releases it once caller streams are synced. |
| 273 | engine._handle = engine._m2n.Handle.create( |
| 274 | engine._m2n.Config(max_cta=init_info.max_cta) |
| 275 | ) |
| 276 | |
| 277 | # Now safe to collect: both sides have joined, so the RPC has returned |
| 278 | # or raised. Doing this last means a worker-side init failure surfaces |
| 279 | # here, as an exception from trainer_init, rather than at the first send. |
| 280 | if rendezvous is not None: |
| 281 | rendezvous.result() |
| 282 | return engine |
| 283 | |
| 284 | def _receive_destination_plan( |
| 285 | self, first_worker_rank: int |
| 286 | ) -> list[Placements | None]: |
| 287 | """Receive and validate the inference-side per-parameter placements.""" |
| 288 | if self.group is None or self._dst_mesh is None: |
| 289 | raise RuntimeError( |
| 290 | "destination plan requires the communicator and destination mesh" |
| 291 | ) |
| 292 | |
| 293 | # The first inference worker publishes one placement per parameter. |
| 294 | # Every trainer rank receives the complete plan in this single, |
| 295 | # initialization-time broadcast. |
| 296 | placements = publish_destination_placements( |
| 297 | self.group, first_worker_rank, None, len(self._metas) |
| 298 | ) |
| 299 | # The send loop matches placements to metadata by position, so reject |
| 300 | # an incomplete or oversized plan before entering any M2N collective. |
| 301 | if len(placements) != len(self._metas): |
| 302 | raise RuntimeError( |
| 303 | f"inference side planned {len(placements)} parameters, but " |
| 304 | f"this trainer announced {len(self._metas)}" |
| 305 | ) |
| 306 | # Check every worker-selected destination against the corresponding |
| 307 | # global parameter shape before the first transfer. |
| 308 | for meta, placement in zip(self._metas, placements): |
| 309 | mesh, resolved = resolve_layout(self._dst_mesh, placement) |
| 310 | validate_layout(mesh, resolved, meta.shape, "destination") |
| 311 | return placements |
| 312 | |
| 313 | def _worker_init_info(self, init_info: M2NTrainerInitInfo) -> dict[str, Any]: |
| 314 | """Handshake payload: rendezvous, both meshes, and the transfer plan.""" |
| 315 | assert self._src_mesh is not None |
| 316 | payload: dict[str, Any] = { |
| 317 | "master_address": init_info.master_address, |
| 318 | "master_port": init_info.master_port, |
| 319 | "nccl_unique_id_b64": init_info.nccl_unique_id_b64, |
| 320 | "rank_offset": init_info.num_trainer_ranks, |
| 321 | "world_size": init_info.world_size, |
| 322 | "src_mesh_dims": list(self._src_mesh.dims), |
| 323 | "dst_mesh_dims": list(init_info.destination_mesh_dims), |
| 324 | "names": [m.name for m in self._metas], |
| 325 | "dtype_names": [_dtype_name(m.dtype) for m in self._metas], |
| 326 | "shapes": [list(m.shape) for m in self._metas], |
| 327 | "src_placements": [ |
| 328 | None if m.placements is None else list(m.placements) |
| 329 | for m in self._metas |
| 330 | ], |
| 331 | "max_cta": init_info.max_cta, |
| 332 | } |
| 333 | return {key: value for key, value in payload.items() if value is not None} |
| 334 | |
| 335 | def send_weights(self) -> None: |
| 336 | """Drive one update round: start, reshard concurrently, then finish.""" |
| 337 | if self._handle is None or self.group is None: |
| 338 | raise RuntimeError("nccl_m2n trainer engine is not initialized") |
| 339 | |
| 340 | update: Future | None = None |
| 341 | if self.is_controller: |
| 342 | self.client.start_weight_update() |
| 343 | # The workers reshard from inside `update_weights`, concurrently |
| 344 | # with the sends below. |
| 345 | assert self._executor is not None |
| 346 | update = self._executor.submit( |
| 347 | self.client.update_weights, |
| 348 | {"names": [m.name for m in self._metas]}, |
| 349 | ) |
| 350 | |
| 351 | try: |
| 352 | self._send() |
| 353 | finally: |
| 354 | if update is not None: |
| 355 | update.result() |
| 356 | |
| 357 | if self.is_controller: |
| 358 | self.client.finish_weight_update() |
| 359 | |
| 360 | def _send(self) -> None: |
| 361 | """Reshard each local shard into its worker-planned destination.""" |
| 362 | assert self.group is not None |
| 363 | assert self._src_mesh is not None |
| 364 | assert self._dst_mesh is not None |
| 365 | m2n = self._m2n |
| 366 | comm = comm_ptr(self.group) |
| 367 | stream = torch.cuda.current_stream() |
| 368 | |
| 369 | for index, (name, tensor) in enumerate(self.source): |
| 370 | meta = self._metas[index] |
| 371 | if name != meta.name: |
| 372 | raise RuntimeError( |
| 373 | f"weight source yielded '{name}' at position {index}, but " |
| 374 | f"its metadata declared '{meta.name}'; the source must " |
| 375 | "iterate in a stable order" |
| 376 | ) |
| 377 | shard = tensor.detach().contiguous() |
| 378 | src_mesh, src_placements = resolve_layout(self._src_mesh, meta.placements) |
| 379 | dst_mesh, dst_placements = resolve_layout( |
| 380 | self._dst_mesh, self._dst_placements[index] |
| 381 | ) |
| 382 | m2n.reshard( |
| 383 | shard, |
| 384 | None, |
| 385 | comm, |
| 386 | stream, |
| 387 | src_mesh=to_mesh(m2n, src_mesh), |
| 388 | src_placements=to_placements(m2n, src_placements), |
| 389 | dst_mesh=to_mesh(m2n, dst_mesh), |
| 390 | dst_placements=to_placements(m2n, dst_placements), |
| 391 | handle=self._handle, |
| 392 | ) |
| 393 | # A custom source may yield temporary storage allocated on another |
| 394 | # stream. M2N retains only its device pointer for asynchronous work, |
| 395 | # so tell the allocator not to reuse that storage until this stream |
| 396 | # has finished with it. For storage allocated on this stream, |
| 397 | # record_stream is effectively a no-op. |
| 398 | shard.record_stream(stream) |
| 399 | |
| 400 | # As in the other weight-transfer backends, wait before returning |
| 401 | # rather than relying on same-stream ordering. The caller may then |
| 402 | # safely mutate parameters or continue work on another stream. |
| 403 | stream.synchronize() |
| 404 | |
| 405 | def shutdown(self) -> None: |
| 406 | """Destroy the m2n handle, join the helper thread, and drop the group.""" |
| 407 | if self._handle is not None: |
| 408 | torch.accelerator.synchronize() |
| 409 | self._handle.destroy() |
| 410 | self._handle = None |
| 411 | if self._executor is not None: |
| 412 | self._executor.shutdown() |
| 413 | self._executor = None |
| 414 | self.group = None |