| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Base class for weight transfer engines.""" |
| 4 | |
| 5 | from abc import ABC, abstractmethod |
| 6 | from collections.abc import Collection, Iterator |
| 7 | from dataclasses import dataclass, field |
| 8 | from typing import ( |
| 9 | TYPE_CHECKING, |
| 10 | Any, |
| 11 | ClassVar, |
| 12 | Generic, |
| 13 | Protocol, |
| 14 | TypeVar, |
| 15 | runtime_checkable, |
| 16 | ) |
| 17 | |
| 18 | import torch |
| 19 | from typing_extensions import Self |
| 20 | |
| 21 | if TYPE_CHECKING: |
| 22 | from vllm.config import VllmConfig |
| 23 | |
| 24 | from vllm.config.parallel import ParallelConfig |
| 25 | from vllm.config.weight_transfer import WeightTransferConfig |
| 26 | |
| 27 | TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo") |
| 28 | TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo") |
| 29 | TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo") |
| 30 | WeightTransferUpdatePayload = dict[str, Any] | list[dict[str, Any]] |
| 31 | |
| 32 | # A trainer supplies its parameters as a `WeightSource` (defined below): a |
| 33 | # re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()` |
| 34 | # channel. The built-in `ModuleSource` uses `materialize_full_tensor`. |
| 35 | |
| 36 | |
| 37 | def _stack_key(name: str) -> tuple[str, int] | None: |
| 38 | """``(prefix, index)`` of the OUTERMOST integer segment, or None if there is |
| 39 | none. |
| 40 | |
| 41 | Outermost is what keeps a MoE layer whole: |
| 42 | ``model.layers.3.mlp.experts.7.w1`` keys on the layer, not the expert. |
| 43 | """ |
| 44 | parts = name.split(".") |
| 45 | for i, part in enumerate(parts): |
| 46 | if part.isdigit(): |
| 47 | return ".".join(parts[:i]), int(part) |
| 48 | return None |
| 49 | |
| 50 | |
| 51 | def layerwise_groups(names: list[str]) -> list[list[str]]: |
| 52 | """Partition flat parameter names into one group per decoder layer, keyed on |
| 53 | the outermost index segment of each name. |
| 54 | |
| 55 | This defines what a *group index* means for `WeightSource.groups` and |
| 56 | `WeightSource.iter_groups`: index *g* names the same group on every trainer |
| 57 | rank and every consumer, because it is derived from one rank's `metadata()` |
| 58 | order. |
| 59 | |
| 60 | Keying on the index rather than a literal prefix needs no per-architecture |
| 61 | naming table: ``model.layers.0.``, ``model.language_model.layers.0.``, |
| 62 | ``transformer.h.0.``, ``backbone.layers.0.`` and a vision tower's |
| 63 | ``visual.blocks.0.`` all partition alike. Matching one fixed prefix does not, |
| 64 | and its failure is silent — every name lands in a single group holding the |
| 65 | whole model, which defeats the per-layer bound below. |
| 66 | |
| 67 | Un-indexed names split by POSITION relative to the first indexed one: the pre |
| 68 | block (embeddings) and the post block (the final norm, `lm_head`, and any |
| 69 | inter-stack projector). Post lands last however early it arrived, which is |
| 70 | what a pipeline-parallel source needs — Megatron-Bridge streams the last |
| 71 | stage's output block *before* its layers. |
| 72 | |
| 73 | Stacks come out in first-appearance order of their prefix and ascending index |
| 74 | within it, whatever order the source yielded them, so a source can normalize |
| 75 | an arbitrary export order by flattening this partition. |
| 76 | |
| 77 | Backends that gather and free per group (sharded RDT) also use it as the unit |
| 78 | of transfer, which bounds their buffer sizes: without it a whole model becomes |
| 79 | one chunk. |
| 80 | """ |
| 81 | pre: list[str] = [] |
| 82 | post: list[str] = [] |
| 83 | stacks: dict[tuple[str, int], list[str]] = {} |
| 84 | order: list[tuple[str, int]] = [] |
| 85 | for name in names: |
| 86 | key = _stack_key(name) |
| 87 | if key is None: |
| 88 | (post if order else pre).append(name) |
| 89 | continue |
| 90 | if key not in stacks: |
| 91 | stacks[key] = [] |
| 92 | order.append(key) |
| 93 | stacks[key].append(name) |
| 94 | |
| 95 | prefix_rank: dict[str, int] = {} |
| 96 | for key in order: |
| 97 | prefix_rank.setdefault(key[0], len(prefix_rank)) |
| 98 | order.sort(key=lambda key: (prefix_rank[key[0]], key[1])) |
| 99 | |
| 100 | groups: list[list[str]] = [pre] if pre else [] |
| 101 | groups += [stacks[key] for key in order] |
| 102 | if post: |
| 103 | groups.append(post) |
| 104 | return groups |
| 105 | |
| 106 | |
| 107 | def materialize_full_tensor(tensor: torch.Tensor) -> torch.Tensor: |
| 108 | """Return a full, locally-materialized tensor ready to send. |
| 109 | |
| 110 | FSDP shards (DTensors) expose `full_tensor()`, a collective all-gather; |
| 111 | regular tensors do not and are returned unchanged. Trainer engines call |
| 112 | this at send time so the (potentially expensive) gather happens exactly |
| 113 | once — reading `.shape`/`.dtype` for metadata does not trigger it. |
| 114 | """ |
| 115 | full_tensor = getattr(tensor, "full_tensor", None) |
| 116 | return full_tensor() if callable(full_tensor) else tensor |
| 117 | |
| 118 | |
| 119 | @dataclass(frozen=True) |
| 120 | class ParamMeta: |
| 121 | """Name / wire dtype / full (HF) shape for one output parameter.""" |
| 122 | |
| 123 | name: str |
| 124 | dtype: torch.dtype |
| 125 | shape: tuple[int, ...] |
| 126 | |
| 127 | |
| 128 | class WeightSource(ABC): |
| 129 | """A re-iterable source of the trainer's weights, handed to a trainer engine. |
| 130 | |
| 131 | Two channels: |
| 132 | |
| 133 | * `metadata()` — `(name, wire dtype, full shape)` for every parameter, |
| 134 | *without* transferring. Cheap when shapes are known locally (FSDP |
| 135 | `DTensor` global shape); may be expensive on first call for backends that |
| 136 | must materialize to learn shapes (e.g. a Megatron-Bridge export), in which |
| 137 | case it should cache. |
| 138 | * iteration — yields fully-materialized `(name, tensor)` pairs, one at a |
| 139 | time. Materializing is typically a collective (FSDP `full_tensor()`, a |
| 140 | Megatron export), so the ranks that share a parameter must iterate it in |
| 141 | the same order in lockstep, or they deadlock. |
| 142 | * `held_names()` — which parameters this rank holds, for producers that are |
| 143 | split so each rank holds only part of the model. Defaults to all. |
| 144 | * `iter_groups()` — the same stream batched per gather group (see |
| 145 | `layerwise_groups`). Defaults to batching `__iter__`; override to |
| 146 | materialize a whole group in one step. |
| 147 | |
| 148 | `iter(source)` must yield a *fresh* pass each round. Backends with custom |
| 149 | producer logic (Megatron export, RDT plans, MoE re-fusing) subclass this. |
| 150 | """ |
| 151 | |
| 152 | @abstractmethod |
| 153 | def metadata(self) -> list[ParamMeta]: |
| 154 | """Declare what iteration will yield, without transferring anything. |
| 155 | |
| 156 | Must agree with iteration element for element: the same parameters, in |
| 157 | the same order, with the same dtypes and shapes. Backends may read both |
| 158 | channels and trust that they match (dense NCCL sizes the worker's receive |
| 159 | buffers and its packed chunk boundaries from this, then sends the bytes |
| 160 | from iteration), so a source that disagrees between the two splits the |
| 161 | stream differently on each side. |
| 162 | """ |
| 163 | raise NotImplementedError |
| 164 | |
| 165 | @abstractmethod |
| 166 | def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: |
| 167 | raise NotImplementedError |
| 168 | |
| 169 | def held_names(self) -> "Collection[str] | None": |
| 170 | """The parameters this rank holds, or None for all of them. |
| 171 | |
| 172 | This is the whole ownership contract. Override it when producers are |
| 173 | split so each holds only part of the model — pipeline parallelism (a rank |
| 174 | holds some layers), expert parallelism (a rank holds some experts), or |
| 175 | any combination, including layouts that fit neither. A consumer routes |
| 176 | each name to a rank that holds it, so per-name is the granularity that |
| 177 | matters; the engine derives everything else from this. |
| 178 | |
| 179 | Three requirements come with overriding it: |
| 180 | |
| 181 | * `metadata()` must still describe the WHOLE model on every rank. The |
| 182 | group partition, the iteration checks and the consumers' pull plans are |
| 183 | all built from one rank's metadata, so a rank that reported only its own |
| 184 | share would leave the rest of the model silently un-transferred. The |
| 185 | sharded-RDT engine cross-checks this across ranks at init. |
| 186 | * Every name must be held by at least one rank, or it can never be |
| 187 | served. The engine raises at init naming the first orphan. |
| 188 | * Iteration must cover exactly `groups()` in metadata order, yielding a |
| 189 | real tensor for each held name and `None` for the rest. A group's |
| 190 | gather is a collective among the ranks that hold part of it, so the |
| 191 | name must still appear (to keep the order check aligned) while the |
| 192 | data is absent. |
| 193 | |
| 194 | Returns: |
| 195 | The held parameter names, or None to hold every one. |
| 196 | |
| 197 | """ |
| 198 | return None |
| 199 | |
| 200 | def groups(self) -> list[list[str]]: |
| 201 | """This rank's gather groups, in metadata order: `layerwise_groups` over |
| 202 | `metadata()`, restricted to the groups holding at least one held name. |
| 203 | |
| 204 | A group with nothing held here is not iterated at all — its gather is a |
| 205 | collective among the ranks that do hold part of it. |
| 206 | """ |
| 207 | groups = layerwise_groups([m.name for m in self.metadata()]) |
| 208 | held = self.held_names() |
| 209 | if held is None: |
| 210 | return groups |
| 211 | held = set(held) |
| 212 | return [g for g in groups if any(n in held for n in g)] |
| 213 | |
| 214 | def iter_groups(self) -> Iterator[tuple[list[str], list[torch.Tensor]]]: |
| 215 | """Yield one `(names, tensors)` batch per group from `groups()`. |
| 216 | |
| 217 | The default drives `__iter__` and batches its output, checking as it goes |
| 218 | that the names arrive in metadata order — ranks sharing a parameter |
| 219 | materialize it with a collective, so a rank that iterates out of order |
| 220 | deadlocks its peers rather than returning wrong data. |
| 221 | |
| 222 | Override when a backend can produce a whole group at once. Materializing |
| 223 | is usually a collective, and driving it per group instead of per tensor |
| 224 | turns ~37k generator resumes into ~95 on a per-expert MoE model (worth |
| 225 | ~0.9s per sync there). An override must yield the same batches in the same |
| 226 | order as this default. |
| 227 | """ |
| 228 | it = iter(self) |
| 229 | for group in self.groups(): |
| 230 | names: list[str] = [] |
| 231 | tensors: list[torch.Tensor] = [] |
| 232 | for expected in group: |
| 233 | name, tensor = next(it) |
| 234 | if name != expected: |
| 235 | raise RuntimeError( |
| 236 | f"WeightSource yielded {name!r} but expected " |
| 237 | f"{expected!r}; iteration order must match metadata()." |
| 238 | ) |
| 239 | names.append(name) |
| 240 | tensors.append(tensor) |
| 241 | yield names, tensors |
| 242 | |
| 243 | |
| 244 | class ModuleSource(WeightSource): |
| 245 | """`WeightSource` over `module.named_parameters()` — the common case. |
| 246 | |
| 247 | Handles both plain dense modules and FSDP-sharded ones with no special |
| 248 | casing: iteration all-gathers each `DTensor` via `full_tensor()` (a |
| 249 | collective) and passes regular tensors through. `metadata()` reads the |
| 250 | *global* `.shape` / `.dtype`, so it never triggers a gather. |
| 251 | """ |
| 252 | |
| 253 | def __init__(self, module: torch.nn.Module) -> None: |
| 254 | self._module = module |
| 255 | |
| 256 | def metadata(self) -> list[ParamMeta]: |
| 257 | return [ |
| 258 | ParamMeta(name, p.dtype, tuple(p.shape)) |
| 259 | for name, p in self._module.named_parameters() |
| 260 | ] |
| 261 | |
| 262 | def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: |
| 263 | for name, param in self._module.named_parameters(): |
| 264 | yield name, materialize_full_tensor(param) |
| 265 | |
| 266 | |
| 267 | # Base protocols for backend-specific dataclasses |
| 268 | @dataclass |
| 269 | class WeightTransferInitInfo(ABC): # noqa: B024 |
| 270 | """Base class for backend-specific initialization info.""" |
| 271 | |
| 272 | pass |
| 273 | |
| 274 | |
| 275 | @dataclass |
| 276 | class TrainerInitInfo: |
| 277 | """Base trainer-side init info: which trainer rank drives the transfer. |
| 278 | |
| 279 | `rank` is this trainer process's rank, provided **explicitly** by the |
| 280 | caller — the engine does not read it from a global process group, which is |
| 281 | ambiguous once several groups (FSDP / TP / PP / EP) exist. Rank 0 is always |
| 282 | the sender: only it opens the endpoint and drives the inference-side RPCs, |
| 283 | while every rank still runs the trainer-side collectives. Backend subclasses |
| 284 | add their own (positional) fields; `rank` is keyword-only so that ordering |
| 285 | never conflicts. |
| 286 | |
| 287 | Every concrete subclass sets a class-level `backend` string (the same key it |
| 288 | registers under in `WeightTransferTrainerFactory`). The factory reads it to |
| 289 | dispatch, so callers pass only the init info/ It is a `ClassVar` |
| 290 | (a fixed per-backend constant), so it is not an ``__init__`` field. |
| 291 | """ |
| 292 | |
| 293 | backend: ClassVar[str] |
| 294 | |
| 295 | rank: int = field(kw_only=True) |
| 296 | |
| 297 | def __init_subclass__(cls, **kwargs: Any) -> None: |
| 298 | super().__init_subclass__(**kwargs) |
| 299 | if not getattr(cls, "backend", None): |
| 300 | raise TypeError( |
| 301 | f"{cls.__name__} must set a class-level `backend` string " |
| 302 | "(the WeightTransferTrainerFactory registry key)." |
| 303 | ) |
| 304 | |
| 305 | @property |
| 306 | def is_sender(self) -> bool: |
| 307 | return self.rank == 0 |
| 308 | |
| 309 | |
| 310 | @dataclass |
| 311 | class WeightTransferUpdateInfo(ABC): # noqa: B024 |
| 312 | """Base class for backend-specific weight update info.""" |
| 313 | |
| 314 | pass |
| 315 | |
| 316 | |
| 317 | # API-level request classes (accept dicts for backend-agnostic serialization) |
| 318 | @dataclass |
| 319 | class WeightTransferInitRequest: |
| 320 | """API-level weight transfer initialization request.""" |
| 321 | |
| 322 | init_info: dict[str, Any] = field(default_factory=dict) |
| 323 | |
| 324 | |
| 325 | @dataclass |
| 326 | class WeightTransferUpdateRequest: |
| 327 | """API-level weight update request.""" |
| 328 | |
| 329 | update_info: WeightTransferUpdatePayload = field(default_factory=dict) |
| 330 | |
| 331 | |
| 332 | class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): |
| 333 | """Base class for weight transfer engines that handle transport of model weights |
| 334 | from a trainer to inference workers. |
| 335 | |
| 336 | This abstraction separates weight transfer transport logic from the worker |
| 337 | implementation, allowing different backends (NCCL, CUDA IPC, RDMA[TODO]) to be |
| 338 | plugged in. |
| 339 | |
| 340 | Each engine owns its full weight-update lifecycle: `start_weight_update`, |
| 341 | `update_weights`, and `finish_weight_update`. Layerwise reloading (used by |
| 342 | checkpoint-format engines) is opted into per engine by running it inside |
| 343 | `start_weight_update`/`finish_weight_update`. Engines that apply weights in |
| 344 | place (e.g. sparse patches) leave those methods as no-ops. |
| 345 | |
| 346 | Subclasses should define: |
| 347 | init_info_cls: Type of backend-specific initialization info |
| 348 | update_info_cls: Type of backend-specific update info |
| 349 | """ |
| 350 | |
| 351 | # Subclasses should override these class attributes |
| 352 | init_info_cls: type[TInitInfo] |
| 353 | update_info_cls: type[TUpdateInfo] |
| 354 | |
| 355 | supports_draft_weight_update: bool = True |
| 356 | |
| 357 | defers_processing: bool = False |
| 358 | """Whether `update_weights` returns before the weights are on the device. |
| 359 | |
| 360 | An engine that pipelines its GPU post-processing onto background threads |
| 361 | cannot let `update_weights` synchronize the device — that would block on those |
| 362 | threads and serialize the pipeline. Such an engine sets this True, omits the |
| 363 | per-update sync, and guarantees completion in `finish_weight_update` instead. |
| 364 | |
| 365 | Callers that go through `finish_weight_update` need do nothing: the engine |
| 366 | drains there. A caller that instead drives the tail itself — running its own |
| 367 | `finalize_layerwise_reload`, say — must read this flag and call |
| 368 | `drain_pending()` first, because with it set a returned `update_weights` means |
| 369 | "queued", not "applied". |
| 370 | """ |
| 371 | |
| 372 | def drain_pending(self) -> None: |
| 373 | """Block until every deferred update has been applied to the model. |
| 374 | |
| 375 | The companion to `defers_processing`: a caller that has taken over the |
| 376 | update tail calls this to re-establish the guarantee that |
| 377 | `finish_weight_update` would otherwise have given it. Idempotent, and a |
| 378 | no-op by default — an engine that processes synchronously has nothing to |
| 379 | drain, so this is always safe to call. |
| 380 | """ |
| 381 | |
| 382 | def __init__( |
| 383 | self, |
| 384 | config: WeightTransferConfig, |
| 385 | vllm_config: "VllmConfig", |
| 386 | device: torch.device, |
| 387 | model: torch.nn.Module, |
| 388 | ) -> None: |
| 389 | """Initialize the weight transfer engine. |
| 390 | |
| 391 | Args: |
| 392 | config: The configuration for the weight transfer engine |
| 393 | vllm_config: The full vLLM config (provides parallel/model config) |
| 394 | device: The device this worker's model lives on |
| 395 | model: The local model instance which will receive the weights |
| 396 | |
| 397 | """ |
| 398 | self.config = config |
| 399 | self.vllm_config = vllm_config |
| 400 | self.parallel_config: ParallelConfig = vllm_config.parallel_config |
| 401 | self.model_config = vllm_config.model_config |
| 402 | self.device = device |
| 403 | self.model = model |
| 404 | self._default_model_config = self.model_config |
| 405 | self._default_model = model |
| 406 | |
| 407 | def set_weight_update_target( |
| 408 | self, |
| 409 | model: torch.nn.Module, |
| 410 | model_config: Any, |
| 411 | ) -> None: |
| 412 | """Set the model that will receive the active weight update.""" |
| 413 | self.model = model |
| 414 | self.model_config = model_config |
| 415 | |
| 416 | def reset_weight_update_target(self) -> None: |
| 417 | """Restore weight updates to the engine's default target model.""" |
| 418 | self.model = self._default_model |
| 419 | self.model_config = self._default_model_config |
| 420 | |
| 421 | def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo: |
| 422 | """Construct typed init info from dict with validation. |
| 423 | |
| 424 | Args: |
| 425 | init_dict: Dictionary containing backend-specific initialization parameters |
| 426 | |
| 427 | Returns: |
| 428 | Typed backend-specific init info dataclass |
| 429 | |
| 430 | Raises: |
| 431 | ValueError: If init_dict is invalid for this backend |
| 432 | |
| 433 | """ |
| 434 | try: |
| 435 | return self.init_info_cls(**init_dict) |
| 436 | except TypeError as e: |
| 437 | raise ValueError( |
| 438 | f"Invalid init_info for {self.__class__.__name__}: {e}" |
| 439 | ) from e |
| 440 | |
| 441 | def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo: |
| 442 | """Construct typed update info from dict with validation. |
| 443 | |
| 444 | Args: |
| 445 | update_dict: Dictionary containing backend-specific update parameters |
| 446 | |
| 447 | Returns: |
| 448 | Typed backend-specific update info dataclass |
| 449 | |
| 450 | Raises: |
| 451 | ValueError: If update_dict is invalid for this backend |
| 452 | |
| 453 | """ |
| 454 | try: |
| 455 | return self.update_info_cls(**update_dict) |
| 456 | except TypeError as e: |
| 457 | raise ValueError( |
| 458 | f"Invalid update_info for {self.__class__.__name__}: {e}" |
| 459 | ) from e |
| 460 | |
| 461 | @abstractmethod |
| 462 | def init_transfer_engine(self, init_info: TInitInfo) -> None: |
| 463 | """Initialize the weight transfer mechanism. |
| 464 | This is called once at the beginning of training. |
| 465 | |
| 466 | Args: |
| 467 | init_info: Backend-specific initialization info |
| 468 | |
| 469 | """ |
| 470 | raise NotImplementedError |
| 471 | |
| 472 | @abstractmethod |
| 473 | def start_weight_update(self) -> None: |
| 474 | """Prepare the engine for a new weight update. |
| 475 | |
| 476 | Engines that receive weights in checkpoint format initialize layerwise reloading |
| 477 | here, else this is typically a no-op. |
| 478 | See: https://docs.vllm.ai/en/latest/training/layerwise/ for more details. |
| 479 | """ |
| 480 | raise NotImplementedError |
| 481 | |
| 482 | @abstractmethod |
| 483 | def finish_weight_update(self) -> None: |
| 484 | """Finalize the current weight update. |
| 485 | |
| 486 | Checkpoint-format engines finalize layerwise reloading here; engines |
| 487 | that apply weights in place leave this as a no-op. |
| 488 | """ |
| 489 | raise NotImplementedError |
| 490 | |
| 491 | def update_weights(self, update_info: dict[str, Any]) -> None: |
| 492 | """Receive one weight update chunk and load it into the model. |
| 493 | |
| 494 | Args: |
| 495 | update_info: Dictionary containing backend-specific update info |
| 496 | |
| 497 | """ |
| 498 | typed_update_info = self.parse_update_info(update_info) |
| 499 | self.receive_weights(typed_update_info) |
| 500 | # NCCL broadcast / IPC paths may be asynchronous. Synchronize here so the |
| 501 | # next step uses the new weights. |
| 502 | torch.accelerator.synchronize() |
| 503 | |
| 504 | @abstractmethod |
| 505 | def receive_weights(self, update_info: TUpdateInfo) -> None: |
| 506 | """Receive weights from the trainer and load them into the model. |
| 507 | |
| 508 | Args: |
| 509 | update_info: Backend-specific update info containing parameter metadata |
| 510 | and any backend-specific data |
| 511 | |
| 512 | """ |
| 513 | raise NotImplementedError |
| 514 | |
| 515 | @abstractmethod |
| 516 | def shutdown(self) -> None: |
| 517 | """Shutdown the weight transfer engine. |
| 518 | This should be called when the worker is shutting down. |
| 519 | """ |
| 520 | raise NotImplementedError |
| 521 | |
| 522 | |
| 523 | @runtime_checkable |
| 524 | class VLLMWeightSyncClient(Protocol): |
| 525 | """Trainer-side stub for the inference engine's weight-sync control plane. |
| 526 | |
| 527 | Mirrors the weight-sync methods that the inference engine exposes |
| 528 | (`EngineClient` / the HTTP RLHF routes / Ray actors). A |
| 529 | `TrainerWeightTransferEngine` drives the full handshake through this |
| 530 | protocol so trainer code never has to know the transport. |
| 531 | |
| 532 | All methods are synchronous and accept plain dicts (matching what the |
| 533 | inference side already accepts). Concurrency that some backends need |
| 534 | (e.g. NCCL must run `update_weights` concurrently with the trainer-side |
| 535 | broadcast) is the engine's responsibility, not the client's, so the |
| 536 | protocol stays a flat four-method surface that any wrapper can implement. |
| 537 | |
| 538 | The protocol is structural (PEP 544), so user implementations need only |
| 539 | define these four methods — no import or subclassing required. |
| 540 | """ |
| 541 | |
| 542 | def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ... |
| 543 | |
| 544 | def start_weight_update(self) -> None: ... |
| 545 | |
| 546 | def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: ... |
| 547 | |
| 548 | def finish_weight_update(self, weight_version: str | None = None) -> None: ... |
| 549 | |
| 550 | |
| 551 | class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): |
| 552 | """Trainer-side weight transfer engine. |
| 553 | |
| 554 | Symmetric to `WeightTransferEngine` but lives in the training process. |
| 555 | Constructed via the `trainer_init` factory classmethod; carries any |
| 556 | backend-specific state (NCCL communicators, IPC device info, transfer |
| 557 | plans) on `self`. Full-resync backends (NCCL, IPC) take a `WeightSource` at |
| 558 | `trainer_init` and replay it each round via the no-argument |
| 559 | `send_weights()`. Backends that push per-round deltas instead (e.g. sparse |
| 560 | patches) leave `source` as `None` and take their payload as a `send_weights` |
| 561 | argument. |
| 562 | |
| 563 | Unlike the worker engine, the trainer side does not take a |
| 564 | `WeightTransferConfig`: the backend is selected from the init info's |
| 565 | `backend` `ClassVar` (so callers pass only the init info), and the static |
| 566 | wire params (packed, buffer sizes) ride the backend-specific |
| 567 | `TrainerInitInfo`, which the sender also propagates to the worker at the init |
| 568 | handshake. |
| 569 | |
| 570 | Multi-rank trainers: `trainer_init` and `send_weights` are |
| 571 | called on *every* trainer rank. Rank 0 is the sender, resolved once at |
| 572 | `trainer_init` into `is_sender`. Non-sender ranks still run every |
| 573 | collective (iterating the source, metadata export, IPC handle all-gather) so |
| 574 | the group stays aligned, but each engine explicitly guards the control-plane |
| 575 | RPCs and the transmit on `self.is_sender`, so only the sender touches the |
| 576 | client. |
| 577 | |
| 578 | Subclasses should define: |
| 579 | init_info_cls: Type of backend-specific trainer init info |
| 580 | """ |
| 581 | |
| 582 | # Subclasses should override this class attribute |
| 583 | init_info_cls: type[TTrainerInitInfo] |
| 584 | |
| 585 | def __init__( |
| 586 | self, |
| 587 | *, |
| 588 | client: "VLLMWeightSyncClient", |
| 589 | source: "WeightSource | None" = None, |
| 590 | is_sender: bool = True, |
| 591 | ) -> None: |
| 592 | self.is_sender = is_sender |
| 593 | # The real client is held on every rank; each engine only *calls* it when |
| 594 | # `is_sender`, so non-sender ranks never touch the wire. |
| 595 | self.client = client |
| 596 | self.source = source |
| 597 | |
| 598 | @classmethod |
| 599 | @abstractmethod |
| 600 | def trainer_init( |
| 601 | cls, |
| 602 | init_info: TTrainerInitInfo, |
| 603 | *, |
| 604 | client: "VLLMWeightSyncClient", |
| 605 | source: "WeightSource | None" = None, |
| 606 | ) -> Self: |
| 607 | """Rendezvous with the inference side and return a ready instance. |
| 608 | |
| 609 | Called on every trainer rank. The sender drives the full handshake via |
| 610 | `client` (build the worker-side init info, call |
| 611 | `client.init_weight_transfer_engine`, open the trainer-side endpoint); |
| 612 | non-sender ranks skip the rendezvous and the RPC. |
| 613 | """ |
| 614 | raise NotImplementedError |
| 615 | |
| 616 | @abstractmethod |
| 617 | def send_weights(self) -> None: |
| 618 | """Push weights to inference workers and drive the full update round |
| 619 | trip: `start_weight_update`, `update_weights` (run concurrently with the |
| 620 | trainer-side broadcast when the backend requires it), then |
| 621 | `finish_weight_update`. Called on every trainer rank. |
| 622 | """ |
| 623 | raise NotImplementedError |
| 624 | |
| 625 | def shutdown(self) -> None: |
| 626 | """Tear down communicators / process groups. Default no-op.""" |