| 1 | # Base Classes and Custom Engines |
| 2 | |
| 3 | The weight transfer system is built from four abstractions, each independently |
| 4 | replaceable: |
| 5 | |
| 6 | | Abstraction | Side | Answers | |
| 7 | | ----------- | ---- | ------- | |
| 8 | | [`WeightSource`](#weightsource) | Trainer | *What* weights to send | |
| 9 | | [`VLLMWeightSyncClient`](#vllmweightsyncclient) | Trainer | *How to reach* the inference engine — the adapter for your RL stack's own vLLM wrapper | |
| 10 | | [`TrainerWeightTransferEngine`](#trainerweighttransferengine) | Trainer | *How to transmit* the bytes | |
| 11 | | [`WeightTransferEngine`](#weighttransferengine) | Inference | *How to receive* them and load them | |
| 12 | |
| 13 | The two engines are registered in two separate factories, |
| 14 | [`WeightTransferTrainerFactory`](#weighttransfertrainerfactory) and |
| 15 | [`WeightTransferEngineFactory`](#weighttransferenginefactory). They share backend |
| 16 | names by convention, but a trainer process never instantiates a worker engine or |
| 17 | vice versa, so the registries stay independent. |
| 18 | |
| 19 | ## Trainer Side |
| 20 | |
| 21 | ### WeightSource |
| 22 | |
| 23 | **This is the adapter for whatever shape your trainer's weights are in.** |
| 24 | `WeightSource` defines how you extract your weights for your specific framework. |
| 25 | |
| 26 | What an engine gets is always the same: HF-format parameter names, and tensors |
| 27 | already materialized to their full (unsharded) shape. Whatever gathering, |
| 28 | re-fusing, dequantizing or renaming that takes belongs inside the source. |
| 29 | |
| 30 | A `WeightSource` is **re-iterable** and has two required channels: |
| 31 | |
| 32 | - **`metadata() -> list[ParamMeta]`** — the name, wire dtype, and full shape of |
| 33 | every parameter, *without transferring anything*. Cheap when shapes are known |
| 34 | locally (an FSDP `DTensor` knows its global shape); may be expensive on the |
| 35 | first call for producers that must materialize to learn shapes (a |
| 36 | Megatron-Bridge export), in which case it should cache. |
| 37 | - **iteration** — yields fully-materialized `(name, tensor)` pairs, one at a |
| 38 | time. |
| 39 | |
| 40 | ```python |
| 41 | @dataclass(frozen=True) |
| 42 | class ParamMeta: |
| 43 | name: str |
| 44 | dtype: torch.dtype |
| 45 | shape: tuple[int, ...] |
| 46 | ``` |
| 47 | |
| 48 | !!! warning "The two channels must agree, element for element" |
| 49 | `metadata()` must declare exactly what iteration will yield: the same |
| 50 | parameters, in the same order, with the same dtypes and shapes. This is an |
| 51 | invariant of the ABC, not of any one backend — a source that reorders, omits, |
| 52 | or re-dtypes a parameter between the two channels is broken even if the |
| 53 | backend you happen to test against never notices. Backends are free to read |
| 54 | both channels and to trust that they match; dense NCCL does, and |
| 55 | [enforces it](nccl.md#the-two-weightsource-channels-must-agree). |
| 56 | |
| 57 | Materializing is typically a collective, so **every trainer rank must iterate the |
| 58 | same source in the same order, in lockstep**, or ranks deadlock. `metadata()` can |
| 59 | itself be a collective for custom producers, so it too runs on every rank — only |
| 60 | the sender ships the result. |
| 61 | |
| 62 | `iter(source)` must yield a *fresh* pass each round. |
| 63 | |
| 64 | #### ModuleSource |
| 65 | |
| 66 | `ModuleSource(module)` is the common case, over `module.named_parameters()`. It |
| 67 | handles plain and FSDP-sharded modules with no special casing: iteration |
| 68 | all-gathers each `DTensor` via `full_tensor()`, while `metadata()` reads the |
| 69 | *global* `.shape` / `.dtype` and so never triggers a gather. |
| 70 | |
| 71 | ```python |
| 72 | from vllm.distributed.weight_transfer import ModuleSource |
| 73 | |
| 74 | source = ModuleSource(model) |
| 75 | ``` |
| 76 | |
| 77 | #### Custom sources |
| 78 | |
| 79 | Subclass `WeightSource` when the weights need work to reach HF format — a |
| 80 | framework-specific export, a re-fusing step, a dtype cast. |
| 81 | |
| 82 | ```python |
| 83 | from vllm.distributed.weight_transfer import ParamMeta, WeightSource |
| 84 | |
| 85 | |
| 86 | class MegatronBridgeSource(WeightSource): |
| 87 | """Megatron model -> HF names, via a bridge that gathers TP/PP/EP internally |
| 88 | and returns full tensors on every rank.""" |
| 89 | |
| 90 | def __init__(self, bridge, module, dtype): |
| 91 | self._bridge, self._module, self._dtype = bridge, module, dtype |
| 92 | self._meta: list[ParamMeta] | None = None |
| 93 | |
| 94 | def _export(self): |
| 95 | return self._bridge.export_hf_weights(self._module) |
| 96 | |
| 97 | def metadata(self) -> list[ParamMeta]: |
| 98 | # Cache: for producers that must materialize to learn shapes, this is |
| 99 | # the expensive channel. Runs on every rank (it may be a collective). |
| 100 | if self._meta is None: |
| 101 | self._meta = [ |
| 102 | ParamMeta(name, self._dtype, tuple(t.shape)) |
| 103 | for name, t in self._export() |
| 104 | ] |
| 105 | return self._meta |
| 106 | |
| 107 | def __iter__(self): |
| 108 | # Must yield exactly what metadata() declared, in the same order. |
| 109 | for name, tensor in self._export(): |
| 110 | yield name, tensor.to(self._dtype).detach().contiguous() |
| 111 | ``` |
| 112 | |
| 113 | #### `held_names()`: partial ownership |
| 114 | |
| 115 | By default every rank is assumed to be able to produce every parameter — which is |
| 116 | what the source above does, since its bridge gathers across all parallelism before |
| 117 | yielding. That is simple and always correct, but it means the gather cost is paid |
| 118 | in full on every rank. |
| 119 | |
| 120 | Optionally, you can override `held_names()` when the ranks are split so each holds only part of the |
| 121 | model. It returns the |
| 122 | parameter names this rank holds (or `None`, the default, for all of them): |
| 123 | |
| 124 | ```python |
| 125 | def held_names(self): |
| 126 | # This pipeline stage's layers, and within them only this EP rank's experts. |
| 127 | return self._my_stage_names - self._foreign_expert_names |
| 128 | ``` |
| 129 | |
| 130 | This covers various trainer layouts — pipeline stages (a rank holds some layers), |
| 131 | expert parallelism (a rank holds some experts), both at once, or a shape that fits |
| 132 | neither. Backends that can route per parameter (see |
| 133 | [sharded RDT](sharded_rdt.md)) then pull each name from a rank that actually |
| 134 | holds it. |
| 135 | |
| 136 | Three requirements come with overriding it: |
| 137 | |
| 138 | - **`metadata()` must still describe the whole model on every rank.** Only the |
| 139 | sender's metadata reaches the inference side, so a rank that reported just its |
| 140 | own share would leave the rest silently un-transferred. Sharded RDT cross-checks |
| 141 | this across ranks at init. |
| 142 | - **Every name must be held by at least one rank**, or it can never be served. |
| 143 | The engine raises at init naming the first orphan. |
| 144 | - **Iteration yields `None` for a name this rank does not hold.** The name still |
| 145 | appears, in metadata order, so the order check stays aligned across ranks — only |
| 146 | the data is absent. Claiming a name and then yielding `None` for it is an error |
| 147 | the engine reports by name. |
| 148 | |
| 149 | !!! warning "Partial ownership only works with sharded RDT" |
| 150 | Only a backend that routes per parameter honours `held_names()`. |
| 151 | Broadcast backends ignore it and send every name from every rank, so |
| 152 | declaring partial ownership there changes nothing. |
| 153 | |
| 154 | #### Gather groups |
| 155 | |
| 156 | Some backends transfer a layer at a time rather than a model at a time, so they |
| 157 | partition `metadata()` into **gather groups**. `layerwise_groups` keys each name |
| 158 | on the **outermost index segment** it contains, so **a group is one decoder |
| 159 | layer**, with runs of un-indexed names (the embeddings, the final norm, |
| 160 | `lm_head`) forming groups of their own where they appear: |
| 161 | |
| 162 | ```text |
| 163 | group 0 model.embed_tokens.weight |
| 164 | group 1 model.layers.0.* <- one decoder layer |
| 165 | group 2 model.layers.1.* |
| 166 | ... |
| 167 | group N+1 model.norm.weight, lm_head.weight |
| 168 | ``` |
| 169 | |
| 170 | Keying on the index rather than a literal prefix means no per-architecture table: |
| 171 | `model.layers.0.`, `model.language_model.layers.0.` (recent Qwen text |
| 172 | checkpoints), `transformer.h.0.` (GPT-2, Falcon), `backbone.layers.0.` (Mamba) |
| 173 | and a vision tower's `visual.blocks.0.` all partition the same way. The index |
| 174 | taken is the *outermost* one, which keeps a MoE layer whole — a per-expert name |
| 175 | like `model.layers.3.mlp.experts.7.w1` keys on the layer, not the expert. |
| 176 | |
| 177 | Group index *g* means the same layer on every rank and every consumer, because |
| 178 | every side derives it from one rank's `metadata()`. That agreement is what lets a |
| 179 | backend bound its buffers to one layer and free a layer once everyone is done |
| 180 | with it. |
| 181 | |
| 182 | !!! warning "A leaf module's sources must all fall in one group" |
| 183 | The sharded-RDT engine frees a group as soon as its last chunk lands, so a |
| 184 | module split across groups would park a pull until the stall watchdog fires. |
| 185 | The default partition guarantees this; a `groups()` override must preserve |
| 186 | it. |
| 187 | |
| 188 | Two hooks follow from this, both with working defaults: |
| 189 | |
| 190 | - **`groups()`** — this rank's groups, in metadata order. The default is |
| 191 | `layerwise_groups(metadata())` restricted to the groups holding at least one |
| 192 | held name. A group with nothing held here is skipped entirely. |
| 193 | - **`iter_groups()`** — the same stream, batched one group at a time. The default |
| 194 | drives `__iter__` and batches its output, checking that names arrive in metadata |
| 195 | order. Override it when your framework can produce a whole group in one step: |
| 196 | materializing is usually a collective, and driving it per group rather than per |
| 197 | tensor turns ~37k generator resumes into ~95 on a per-expert MoE model. |
| 198 | |
| 199 | Because `metadata()` order defines the partition, **all names sharing a layer |
| 200 | index must be contiguous in it**. A source whose natural export order interleaves |
| 201 | layers — bucketing all the MoE experts together, say — has to reorder before |
| 202 | returning. |
| 203 | |
| 204 | ### VLLMWeightSyncClient |
| 205 | |
| 206 | **This is the adapter for however your RL stack reaches vLLM.** Many RL frameworks wrap |
| 207 | inference engines in their own abstractions, and each reaches vLLM its own way. |
| 208 | `VLLMWeightSyncClient` is the single seam where that bespoke shape is adapted, so |
| 209 | weight sync engines remain control plane agnostic. |
| 210 | |
| 211 | The contract is only this: **however the wrapper is shaped, it must bottom out in |
| 212 | the same four calls** — `init_weight_transfer_engine` once at setup, then |
| 213 | `start_weight_update` → one or more `update_weights` → `finish_weight_update` per |
| 214 | round. Everything a trainer engine needs from the inference side goes through them. |
| 215 | |
| 216 | ```python |
| 217 | class VLLMWeightSyncClient(Protocol): |
| 218 | def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ... |
| 219 | def start_weight_update(self) -> None: ... |
| 220 | def update_weights(self, update_info: dict[str, Any]) -> None: ... |
| 221 | def finish_weight_update(self, weight_version: str | None = None) -> None: ... |
| 222 | ``` |
| 223 | |
| 224 | It is a `@runtime_checkable` structural `Protocol` (PEP 544), which is what makes |
| 225 | adapting cheap: **any object with those four methods already satisfies it**. An existing |
| 226 | wrapper in your framework can usually become a client by gaining four forwarding |
| 227 | methods. |
| 228 | |
| 229 | Two implementations ship with vLLM: |
| 230 | |
| 231 | | Client | Talks to | |
| 232 | | ------ | -------- | |
| 233 | | `RayVLLMWeightSyncClient(handle)` | One or more `AsyncLLM`/`LLM` Ray actors. Accepts a list and fans each call out to every handle, blocking on all of them, so a multi-actor (e.g. multi-DP) deployment is driven as one unit | |
| 234 | | `HTTPVLLMWeightSyncClient(base_url, timeout=300)` | A vLLM server over the RLHF HTTP routes | |
| 235 | |
| 236 | Custom weight sync clients can be implemented like so: |
| 237 | |
| 238 | ```python |
| 239 | class MyFrameworkWeightSyncClient: |
| 240 | """Adapts one RL framework's rollout pool to the four weight-sync calls.""" |
| 241 | |
| 242 | def __init__(self, rollout_pool): |
| 243 | self.pool = rollout_pool # whatever your stack already has |
| 244 | |
| 245 | def init_weight_transfer_engine(self, init_info): |
| 246 | # Fan out to every replica and block: all of them receive weights. |
| 247 | self.pool.broadcast_rpc("init_weight_transfer_engine", init_info=init_info) |
| 248 | |
| 249 | def start_weight_update(self): |
| 250 | self.pool.broadcast_rpc("start_weight_update") |
| 251 | |
| 252 | def update_weights(self, update_info): |
| 253 | self.pool.broadcast_rpc("update_weights", update_info=update_info) |
| 254 | |
| 255 | def finish_weight_update(self, weight_version=None): |
| 256 | self.pool.broadcast_rpc("finish_weight_update") |
| 257 | if weight_version is not None: |
| 258 | self.pool.broadcast_rpc("update_weight_version", weight_version) |
| 259 | ``` |
| 260 | |
| 261 | Two things to get right in any adapter: |
| 262 | |
| 263 | - **Reach every replica, and block until all of them are done.** A weight update |
| 264 | is not a load-balanced request: every worker holding a copy of the model must |
| 265 | receive it. Returning before they all finish lets the trainer race ahead of |
| 266 | workers still loading. (Both built-in clients do this — Ray by fanning out over |
| 267 | its handles, HTTP because the server's DP client broadcasts internally.) |
| 268 | - **Raise on failure.** Trainer engines rely on exceptions to surface |
| 269 | inference-side errors; a client that swallows them turns a failed sync into |
| 270 | silently stale weights, or into a hang for backends whose transfer rendezvouses |
| 271 | with the worker. |
| 272 | |
| 273 | !!! note |
| 274 | HTTP cannot carry raw CUDA IPC handles, so `HTTPVLLMWeightSyncClient` pickles |
| 275 | and base64-encodes them into an `ipc_handles_pickled` field. The worker |
| 276 | deserializes it only when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`. Backends |
| 277 | whose payloads are JSON-native (NCCL) pass through untouched. |
| 278 | |
| 279 | ### TrainerWeightTransferEngine |
| 280 | |
| 281 | The trainer-side engine: it holds the transport state (NCCL communicators, IPC |
| 282 | device info, transfer plans), pulls weights from a `WeightSource`, and drives the |
| 283 | inference side through a `VLLMWeightSyncClient`. It is generic over its init info |
| 284 | type, constructed by the `trainer_init` classmethod factory, and driven by `send_weights()`. |
| 285 | |
| 286 | | Method | Description | |
| 287 | | ------ | ----------- | |
| 288 | | `trainer_init(init_info, *, client, source=None)` | Classmethod. Rendezvous with the inference side and return a ready instance | |
| 289 | | `send_weights()` | Push weights and drive the full update round trip | |
| 290 | | `shutdown()` | Tear down communicators / process groups. Default no-op | |
| 291 | |
| 292 | Both `trainer_init` and `send_weights` are called on **every** trainer rank. |
| 293 | `is_sender` is resolved once, at `trainer_init`, from `init_info.rank`. Each |
| 294 | engine holds the real client on every rank but guards the control-plane RPCs and |
| 295 | the transmit on `self.is_sender`, so only the sender touches the wire; non-sender |
| 296 | ranks still run every collective so the group stays aligned. |
| 297 | |
| 298 | The trainer side takes **no `WeightTransferConfig`**. The backend comes from the |
| 299 | init info's `backend` `ClassVar`, and the wire params ride the init info too. |
| 300 | |
| 301 | #### TrainerInitInfo |
| 302 | |
| 303 | The `init_info` passed to `trainer_init` above. It is how a caller configures a |
| 304 | transfer: it selects the backend, says which rank this process is, and carries the |
| 305 | wire params. Each backend subclasses it; the base class holds the one field every |
| 306 | backend needs. |
| 307 | |
| 308 | ```python |
| 309 | @dataclass |
| 310 | class TrainerInitInfo: |
| 311 | backend: ClassVar[str] # factory dispatch key |
| 312 | rank: int = field(kw_only=True) |
| 313 | |
| 314 | @property |
| 315 | def is_sender(self) -> bool: |
| 316 | return self.rank == 0 |
| 317 | ``` |
| 318 | |
| 319 | - **`rank`** is this trainer process's rank, supplied **explicitly**. The engine |
| 320 | does not read it from a global process group, which is ambiguous once several |
| 321 | groups (FSDP / TP / PP / EP) exist. **Rank 0 is always the sender** — this is |
| 322 | what `trainer_init` resolves into `is_sender`. It is keyword-only, so backend |
| 323 | subclasses can add positional fields freely. |
| 324 | - **`backend`** is a `ClassVar`, not an `__init__` field: it is a fixed |
| 325 | per-backend constant that the factory reads to dispatch, which is why callers |
| 326 | never pass a `backend=` argument. Every subclass must set it — |
| 327 | `__init_subclass__` raises otherwise. |
| 328 | |
| 329 | Subclasses also carry the transfer's **wire params** (`packed`, buffer sizes). |
| 330 | The sender propagates them to the worker inside `trainer_init`, so the two sides |
| 331 | cannot disagree. See [`NCCLTrainerInitInfo`](nccl.md#nccltrainerinitinfo) and |
| 332 | [`IPCTrainerInitInfo`](ipc.md#ipctrainerinitinfo) for the concrete fields. |
| 333 | |
| 334 | #### Full-Resync vs. Delta Backends |
| 335 | |
| 336 | `source` is optional, which splits the backends into two shapes: |
| 337 | |
| 338 | - **Full resync** (NCCL, IPC) — a stable `WeightSource` is fixed at |
| 339 | `trainer_init` and re-iterated each round; `send_weights()` takes no |
| 340 | arguments. These backends validate that `source` is non-null themselves. |
| 341 | - **Delta** (sparse NCCL) — the payload differs every round, so there is no |
| 342 | stable source. The engine takes no `source` and each round's payload is passed |
| 343 | straight to `send_weights(patches)`. |
| 344 | |
| 345 | #### Implementing a Custom Trainer Engine |
| 346 | |
| 347 | ```python |
| 348 | from dataclasses import dataclass |
| 349 | from typing import ClassVar |
| 350 | |
| 351 | from typing_extensions import Self |
| 352 | |
| 353 | from vllm.distributed.weight_transfer.base import ( |
| 354 | TrainerInitInfo, |
| 355 | TrainerWeightTransferEngine, |
| 356 | VLLMWeightSyncClient, |
| 357 | WeightSource, |
| 358 | ) |
| 359 | |
| 360 | |
| 361 | @dataclass |
| 362 | class MyTrainerInitInfo(TrainerInitInfo): |
| 363 | backend: ClassVar[str] = "my_backend" |
| 364 | |
| 365 | endpoint: str |
| 366 | chunk_size_bytes: int = 256 * 1024 * 1024 # a wire param: shipped to the worker |
| 367 | |
| 368 | |
| 369 | class MyTrainerWeightTransferEngine(TrainerWeightTransferEngine[MyTrainerInitInfo]): |
| 370 | init_info_cls = MyTrainerInitInfo |
| 371 | |
| 372 | def __init__(self, *, client, source, is_sender=True, chunk_size_bytes=0): |
| 373 | super().__init__(client=client, source=source, is_sender=is_sender) |
| 374 | self.chunk_size_bytes = chunk_size_bytes |
| 375 | |
| 376 | @classmethod |
| 377 | def trainer_init( |
| 378 | cls, |
| 379 | init_info: MyTrainerInitInfo, |
| 380 | *, |
| 381 | client: VLLMWeightSyncClient, |
| 382 | source: WeightSource | None = None, |
| 383 | ) -> Self: |
| 384 | if source is None: |
| 385 | raise ValueError("my_backend requires a WeightSource.") |
| 386 | engine = cls( |
| 387 | client=client, |
| 388 | source=source, |
| 389 | is_sender=init_info.is_sender, |
| 390 | chunk_size_bytes=init_info.chunk_size_bytes, |
| 391 | ) |
| 392 | if engine.is_sender: |
| 393 | # Ship the must-agree wire params so the worker decodes exactly as |
| 394 | # this trainer encodes, then open the trainer-side endpoint. |
| 395 | engine.client.init_weight_transfer_engine( |
| 396 | {"chunk_size_bytes": init_info.chunk_size_bytes} |
| 397 | ) |
| 398 | return engine |
| 399 | |
| 400 | def send_weights(self) -> None: |
| 401 | assert self.source is not None |
| 402 | meta = self.source.metadata() # every rank: may be a collective |
| 403 | if not self.is_sender: |
| 404 | for _ in self.source: # stay in the trainer-side collective |
| 405 | pass |
| 406 | return |
| 407 | |
| 408 | self.client.start_weight_update() |
| 409 | self.client.update_weights( |
| 410 | { |
| 411 | "names": [m.name for m in meta], |
| 412 | "dtype_names": [str(m.dtype).split(".")[-1] for m in meta], |
| 413 | "shapes": [list(m.shape) for m in meta], |
| 414 | } |
| 415 | ) |
| 416 | for name, tensor in self.source: |
| 417 | ... # transmit |
| 418 | self.client.finish_weight_update() |
| 419 | ``` |
| 420 | |
| 421 | Two things to get right, both of which have bitten the built-in backends: |
| 422 | |
| 423 | - **Drain before returning.** `send_weights` must not return with transfers |
| 424 | still in flight. Anything keeping a send buffer alive dies with the frame, and |
| 425 | the inference side's `finish_weight_update` post-processing can otherwise |
| 426 | finalize weights that have not landed. |
| 427 | - **Never join a control-plane thread on the error path.** If you run |
| 428 | `update_weights` on a side thread concurrently with a transmit (as NCCL does) |
| 429 | and the transmit raises, the worker is still blocked in the matching |
| 430 | collective and will never return. Shut the executor down without waiting, so |
| 431 | the real exception surfaces instead of hanging. |
| 432 | |
| 433 | ### WeightTransferTrainerFactory |
| 434 | |
| 435 | ```python |
| 436 | from vllm.distributed.weight_transfer import WeightTransferTrainerFactory |
| 437 | |
| 438 | # Lazy loading (recommended): the module is imported only when the backend is used |
| 439 | WeightTransferTrainerFactory.register_engine( |
| 440 | "my_backend", |
| 441 | "my_package.my_module", |
| 442 | "MyTrainerWeightTransferEngine", |
| 443 | ) |
| 444 | |
| 445 | # Or register the class directly |
| 446 | WeightTransferTrainerFactory.register_engine("my_backend", MyTrainerWeightTransferEngine) |
| 447 | |
| 448 | engine = WeightTransferTrainerFactory.trainer_init( |
| 449 | init_info=MyTrainerInitInfo(rank=0, endpoint="..."), # `backend` selects the engine |
| 450 | client=client, |
| 451 | source=source, |
| 452 | ) |
| 453 | ``` |
| 454 | |
| 455 | ## Inference Side |
| 456 | |
| 457 | ### WeightTransferEngine |
| 458 | |
| 459 | A generic abstract class parameterized by two dataclass types: |
| 460 | |
| 461 | - **`TInitInfo`** (extends `WeightTransferInitInfo`): backend-specific initialization parameters. |
| 462 | - **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): backend-specific weight update metadata. |
| 463 | |
| 464 | Subclasses must implement five methods: |
| 465 | |
| 466 | | Method | Description | |
| 467 | | ------ | ----------- | |
| 468 | | `init_transfer_engine(init_info)` | Initialize the communication channel on each inference worker, and record the trainer-supplied wire params | |
| 469 | | `start_weight_update()` | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines | |
| 470 | | `finish_weight_update()` | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines | |
| 471 | | `receive_weights(update_info)` | Receive weights and load them into `self.model` | |
| 472 | | `shutdown()` | Clean up resources | |
| 473 | |
| 474 | The base class provides: |
| 475 | |
| 476 | 1. `__init__`, taking `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`), and `model` (`nn.Module`). |
| 477 | 2. `update_weights(update_info_dict)`, a thin wrapper for `receive_weights`: it parses the dict into the typed dataclass, calls `receive_weights`, and synchronizes the device — unless the engine sets `defers_processing`, below. |
| 478 | 3. `parse_init_info` / `parse_update_info`, which convert API-level dicts into the typed dataclasses and raise `ValueError` on a bad payload. |
| 479 | 4. `set_weight_update_target` / `reset_weight_update_target`, used to retarget an update at the speculative draft model. |
| 480 | |
| 481 | !!! note "Read wire params from the handshake, not the payload" |
| 482 | Anything the two sides must agree on — `packed`, buffer geometry — arrives |
| 483 | on the **init info** and should be stored on `self` in |
| 484 | `init_transfer_engine`, then read from `self` in `receive_weights`. Per-round |
| 485 | update info carries only per-round metadata. This is what makes a |
| 486 | trainer/worker mismatch unrepresentable. |
| 487 | |
| 488 | !!! note "`defers_processing`: when a returned update means queued, not applied" |
| 489 | An engine that pipelines its GPU post-processing onto background threads |
| 490 | cannot let `update_weights` synchronize the device — that would block on |
| 491 | those threads and serialize the pipeline. Such an engine sets the class |
| 492 | attribute `defers_processing = True`, omits the per-update sync, and |
| 493 | guarantees completion in `finish_weight_update` instead. |
| 494 | |
| 495 | Callers that go through `finish_weight_update` need do nothing; the engine |
| 496 | drains there. A caller that drives the tail itself — running its own |
| 497 | `finalize_layerwise_reload`, say — must check the flag and call |
| 498 | `drain_pending()` first, because with it set a returned `update_weights` |
| 499 | means *queued*, not *applied*. `drain_pending()` is idempotent, and a no-op |
| 500 | on an engine that processes synchronously, so it is always safe to call. |
| 501 | |
| 502 | [Sharded RDT](sharded_rdt.md) is the built-in engine that sets it: it |
| 503 | scatters and quantizes on background threads with their own CUDA streams, so |
| 504 | its `drain_pending()` joins both queues and syncs both streams before |
| 505 | `finalize_layerwise_reload` runs. |
| 506 | |
| 507 | ### Request Classes |
| 508 | |
| 509 | The API-level request classes provide backend-agnostic serialization using plain dictionaries. |
| 510 | |
| 511 | ```python |
| 512 | from vllm.distributed.weight_transfer.base import ( |
| 513 | WeightTransferInitRequest, |
| 514 | WeightTransferUpdateRequest, |
| 515 | ) |
| 516 | |
| 517 | # Init request (dict is converted to backend-specific TInitInfo) |
| 518 | init_request = WeightTransferInitRequest( |
| 519 | init_info={"master_address": "10.0.0.1", "master_port": 29500, ...} |
| 520 | ) |
| 521 | |
| 522 | # Update request (dict is converted to backend-specific TUpdateInfo) |
| 523 | update_request = WeightTransferUpdateRequest( |
| 524 | update_info={"names": [...], "dtype_names": [...], "shapes": [...]} |
| 525 | ) |
| 526 | ``` |
| 527 | |
| 528 | Using a built-in client, you never construct these by hand — `RayVLLMWeightSyncClient` |
| 529 | wraps the dicts for you, and `HTTPVLLMWeightSyncClient` posts them as JSON. |
| 530 | |
| 531 | At the LLM/API layer, call `start_draft_weight_update()` instead of |
| 532 | `start_weight_update()` to target the speculative draft model; |
| 533 | `update_weights` / `finish_weight_update` are unchanged. Engines that cannot |
| 534 | support this set `supports_draft_weight_update = False`. |
| 535 | |
| 536 | ### Implementing a Custom Engine |
| 537 | |
| 538 | #### 1. Define Info Dataclasses |
| 539 | |
| 540 | ```python |
| 541 | from dataclasses import dataclass |
| 542 | from vllm.distributed.weight_transfer.base import ( |
| 543 | WeightTransferEngine, |
| 544 | WeightTransferInitInfo, |
| 545 | WeightTransferUpdateInfo, |
| 546 | ) |
| 547 | |
| 548 | @dataclass |
| 549 | class MyInitInfo(WeightTransferInitInfo): |
| 550 | endpoint: str |
| 551 | chunk_size_bytes: int = 256 * 1024 * 1024 # must-agree wire param |
| 552 | |
| 553 | @dataclass |
| 554 | class MyUpdateInfo(WeightTransferUpdateInfo): |
| 555 | names: list[str] |
| 556 | dtype_names: list[str] |
| 557 | shapes: list[list[int]] |
| 558 | # Per-round metadata only. |
| 559 | ``` |
| 560 | |
| 561 | #### 2. Implement the Engine |
| 562 | |
| 563 | ```python |
| 564 | class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]): |
| 565 | init_info_cls = MyInitInfo |
| 566 | update_info_cls = MyUpdateInfo |
| 567 | |
| 568 | def init_transfer_engine(self, init_info: MyInitInfo) -> None: |
| 569 | # Record the trainer's wire params, then set up the connection. |
| 570 | self.chunk_size_bytes = init_info.chunk_size_bytes |
| 571 | ... |
| 572 | |
| 573 | def start_weight_update(self) -> None: |
| 574 | # Checkpoint-format engines: run initialize_layerwise_reload(self.model). |
| 575 | # In-place engines: no-op |
| 576 | ... |
| 577 | |
| 578 | def finish_weight_update(self) -> None: |
| 579 | # Checkpoint-format engines: run finalize_layerwise_reload(...). |
| 580 | # In-place engines: no-op |
| 581 | ... |
| 582 | |
| 583 | def receive_weights(self, update_info: MyUpdateInfo) -> None: |
| 584 | weights = [] |
| 585 | for name, dtype_name, shape in zip( |
| 586 | update_info.names, update_info.dtype_names, update_info.shapes |
| 587 | ): |
| 588 | dtype = getattr(torch, dtype_name) |
| 589 | weight = self._fetch_weight(name, shape, dtype) |
| 590 | weights.append((name, weight)) |
| 591 | self.model.load_weights(weights) |
| 592 | |
| 593 | def shutdown(self) -> None: |
| 594 | # Clean up resources |
| 595 | ... |
| 596 | ``` |
| 597 | |
| 598 | #### 3. Register with the Factory |
| 599 | |
| 600 | ```python |
| 601 | from vllm.distributed.weight_transfer import WeightTransferEngineFactory |
| 602 | |
| 603 | # Option 1: Lazy loading (recommended for built-in engines) |
| 604 | WeightTransferEngineFactory.register_engine( |
| 605 | "my_backend", |
| 606 | "my_package.my_module", |
| 607 | "MyWeightTransferEngine", |
| 608 | ) |
| 609 | |
| 610 | # Option 2: Direct class registration |
| 611 | WeightTransferEngineFactory.register_engine( |
| 612 | "my_backend", |
| 613 | MyWeightTransferEngine, |
| 614 | ) |
| 615 | ``` |
| 616 | |
| 617 | Once registered, users select your backend via `WeightTransferConfig(backend="my_backend")`. |
| 618 | |
| 619 | ### WeightTransferEngineFactory |
| 620 | |
| 621 | The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, `sparse_nccl` and `sharded_rdt`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed. |
| 622 | |
| 623 | ```python |
| 624 | from vllm.distributed.weight_transfer import WeightTransferEngineFactory |
| 625 | |
| 626 | # Create an engine from config |
| 627 | engine = WeightTransferEngineFactory.create_engine( |
| 628 | config=weight_transfer_config, |
| 629 | vllm_config=vllm_config, |
| 630 | device=device, |
| 631 | model=model, |
| 632 | ) |
| 633 | ``` |
| 634 | |
| 635 | vLLM calls this for you during worker startup; you only need it directly when |
| 636 | embedding the engine in your own worker. |