vllm.distributed.weight_transfer

Weight transfer engines for syncing model weights from trainers to inference workers.

Modules:

• base –

  Base class for weight transfer engines.

• clients –

  Built-in VLLMWeightSyncClient implementations.

• factory –

  Factory for weight transfer engines with lazy loading.

• ipc_engine –

  IPC-based weight transfer engine using CUDA IPC for communication.

• m2n_common –

  Shared helpers for the NCCL M2N (nccl_m2n) weight transfer backend.

• m2n_engine –

  Inference-side weight transfer engine built on NCCL M2N (nccl_m2n).

• m2n_layout –

  Destination-layout resolution for the NCCL M2N weight transfer backend.

• m2n_trainer –

  Trainer-side weight transfer engine for the NCCL M2N backend.

• modelexpress_engine –

  Re-export the weight-transfer backend maintained by ModelExpress.

• nccl_common –

  Shared NCCL initialization helpers for weight transfer engines.

• nccl_engine –

  NCCL-based (dense) weight transfer engine.

• packed_tensor –

  Packed tensor utilities for efficient weight transfer.

• sharded_rdt_common –

  Shared pieces of the sharded-RDT backend: the op-chain allowlist, buffer

• sharded_rdt_engine –

  Sharded Ray Direct Transport (RDT) weight transfer engine (consumer side).

• sharded_rdt_fake –

  Op-chain recording for the sharded-RDT backend.

• sharded_rdt_trainer –

  Trainer-side engine for the sharded-RDT (pull-based NIXL) backend.

• sparse_nccl_engine –

  Sparse NCCL weight transfer engine.

Classes:

• HTTPVLLMWeightSyncClient –

  Talks to a vLLM server over the RLHF HTTP routes.

• ModuleSource –

  WeightSource over module.named_parameters(): plain, FSDP or DTensor.

• ParamMeta –

  One output parameter: its name, wire dtype and full (HF) shape, plus

• RayVLLMWeightSyncClient –

  Talks to one or more vLLM AsyncLLM/LLM Ray actors.

• Replicate –

  Every coordinate along this mesh axis holds the whole extent.

• Shard –

  This mesh axis splits tensor dim dim the way torch.chunk does.

• TrainerWeightTransferEngine –

  Trainer-side weight transfer engine.

• VLLMWeightSyncClient –

  Trainer-side stub for the inference engine's weight-sync control plane.

• WeightSource –

  The trainer's weights, as metadata plus one or two iteration modes.

• WeightTransferEngine –

  Base class for weight transfer engines that handle transport of model weights

• WeightTransferEngineFactory –

  Factory for creating weight transfer engines with lazy loading.

• WeightTransferTrainerFactory –

  Factory for creating trainer-side weight transfer engines.

HTTPVLLMWeightSyncClient

Talks to a vLLM server over the RLHF HTTP routes.

Mirrors vllm/entrypoints/rl/online/api_router.py: /init_weight_transfer_engine, /start_weight_update, /update_weights, /finish_weight_update, /shutdown_weight_transfer_engine.

Source code in vllm/distributed/weight_transfer/clients.py

class HTTPVLLMWeightSyncClient:
    """Talks to a vLLM server over the RLHF HTTP routes.

    Mirrors `vllm/entrypoints/rl/online/api_router.py`:
    `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`,
    `/finish_weight_update`, `/shutdown_weight_transfer_engine`.
    """

    def __init__(self, base_url: str, timeout: float = 300) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def _post(self, path: str, json: dict[str, Any] | None = None) -> None:
        import requests

        response = requests.post(
            f"{self.base_url}/{path}", json=json, timeout=self.timeout
        )
        response.raise_for_status()

    def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None:
        self._post("init_weight_transfer_engine", {"init_info": init_info})

    def shutdown_weight_transfer_engine(self) -> None:
        self._post("shutdown_weight_transfer_engine")

    def start_weight_update(self) -> None:
        self._post("start_weight_update")

    def update_weights(self, update_info: WeightTransferUpdatePayload) -> None:
        self._post(
            "update_weights", {"update_info": _json_safe_update_payload(update_info)}
        )

    def finish_weight_update(self, weight_version: str | None = None) -> None:
        json = (
            {"weight_version": weight_version} if weight_version is not None else None
        )
        self._post("finish_weight_update", json)

ModuleSource

Bases: WeightSource

WeightSource over module.named_parameters(): plain, FSDP or DTensor.

A plain tensor is owned by every rank. A DTensor is owned by the global ranks of its device mesh, with the mesh's shape and the DTensor's placements. iter_full gathers each DTensor with full_tensor() (a collective among its owners); iter_local yields to_local(). A DTensor whose placements are not plain Shard / Replicate is declared replicated, and iter_local gathers it. metadata reads global shapes, never gathers, and is cached on first read, so shard the module first.

The mesh ranks are taken as trainer ranks, so the default process group's ranks must be the TrainerInitInfo.ranks.

Source code in vllm/distributed/weight_transfer/base.py

class ModuleSource(WeightSource):
    """`WeightSource` over `module.named_parameters()`: plain, FSDP or DTensor.

    A plain tensor is owned by every rank. A `DTensor` is owned by the global
    ranks of its device mesh, with the mesh's shape and the DTensor's
    placements. `iter_full` gathers each DTensor with `full_tensor()` (a
    collective among its owners); `iter_local` yields `to_local()`. A DTensor
    whose placements are not plain `Shard` / `Replicate` is declared replicated,
    and `iter_local` gathers it. `metadata` reads global shapes, never
    gathers, and is cached on first read, so shard the module first.

    The mesh ranks are taken as trainer ranks, so the default process group's
    ranks must be the `TrainerInitInfo.rank`s.
    """

    def __init__(self, module: torch.nn.Module) -> None:
        self._module = module

    @cached_property
    def metadata(self) -> list[ParamMeta]:
        metas = []
        for name, p in self._module.named_parameters():
            layout = _dtensor_layout(p)
            if layout is None:
                metas.append(ParamMeta(name, p.dtype, tuple(p.shape)))
                continue
            owners, mesh_shape, placements = layout
            metas.append(
                ParamMeta(
                    name,
                    p.dtype,
                    tuple(p.shape),
                    owners=owners,
                    mesh_shape=mesh_shape,
                    placements=placements,
                )
            )
        return metas

    def iter_full(self) -> Iterator[tuple[str, torch.Tensor]]:
        for name, param in self._module.named_parameters():
            if _is_member(param):
                yield name, materialize_full_tensor(param)

    def iter_local(self) -> Iterator[tuple[str, torch.Tensor]]:
        for name, param in self._module.named_parameters():
            if not _is_member(param):
                continue
            layout = _dtensor_layout(param)
            if layout is None:
                yield name, param
            elif layout[2] is None:
                yield name, materialize_full_tensor(param)
            else:
                yield name, param.to_local()  # type: ignore[attr-defined]

ParamMeta dataclass

One output parameter: its name, wire dtype and full (HF) shape, plus which trainer ranks own it and how it is laid out across them.

Metadata is identical on every trainer rank, ownership included, so every rank and every engine derives the same plan from it without a runtime exchange.

Attributes:

• owners (tuple[int, ...] | None) –

  Global trainer ranks (TrainerInitInfo.rank) that own this parameter, in mesh order. None means every rank.

• mesh_shape (tuple[int, ...] | None) –

  Row-major shape of the mesh over owners. Defaults to (len(owners),).

• placements (tuple[Placement, ...] | None) –

  One Replicate or Shard per mesh axis, describing what WeightSource.iter_local yields on each owner; Shard splits the way torch.chunk does, outer mesh axis first when two axes shard one dim. None means replicated; a source without iter_local leaves it unused.

Source code in vllm/distributed/weight_transfer/base.py

@dataclass(frozen=True)
class ParamMeta:
    """One output parameter: its name, wire dtype and full (HF) shape, plus
    which trainer ranks own it and how it is laid out across them.

    Metadata is identical on every trainer rank, ownership included, so every
    rank and every engine derives the same plan from it without a runtime
    exchange.

    Attributes:
        owners: Global trainer ranks (`TrainerInitInfo.rank`) that own this
            parameter, in mesh order. `None` means every rank.
        mesh_shape: Row-major shape of the mesh over `owners`. Defaults to
            `(len(owners),)`.
        placements: One `Replicate` or `Shard` per mesh axis, describing what
            `WeightSource.iter_local` yields on each owner; `Shard` splits the
            way `torch.chunk` does, outer mesh axis first when two axes shard
            one dim. `None` means replicated; a source without `iter_local`
            leaves it unused.

    """

    name: str
    dtype: torch.dtype
    shape: tuple[int, ...]
    owners: tuple[int, ...] | None = field(default=None, kw_only=True)
    mesh_shape: tuple[int, ...] | None = field(default=None, kw_only=True)
    placements: tuple[Placement, ...] | None = field(default=None, kw_only=True)

    def __post_init__(self) -> None:
        object.__setattr__(self, "shape", tuple(self.shape))
        if self.owners is None:
            if self.mesh_shape is not None or self.placements is not None:
                raise ValueError(
                    f"parameter {self.name!r} declares a mesh or placements "
                    "without owners"
                )
            return
        owners = tuple(self.owners)
        if (
            not owners
            or any(not isinstance(r, int) or isinstance(r, bool) for r in owners)
            or min(owners) < 0
            or len(set(owners)) != len(owners)
        ):
            raise ValueError(
                f"parameter {self.name!r} owners must be distinct non-negative "
                f"ranks, got {self.owners}"
            )
        mesh_shape = (len(owners),) if self.mesh_shape is None else self.mesh_shape
        mesh_shape = tuple(mesh_shape)
        if any(d <= 0 for d in mesh_shape) or math.prod(mesh_shape) != len(owners):
            raise ValueError(
                f"parameter {self.name!r} mesh_shape {mesh_shape} does not cover "
                f"its {len(owners)} owners"
            )
        placements = self.placements
        if placements is not None:
            placements = tuple(placements)
            if len(placements) != len(mesh_shape):
                raise ValueError(
                    f"parameter {self.name!r} has {len(placements)} placements "
                    f"for a {len(mesh_shape)}-D mesh"
                )
            for p in placements:
                if isinstance(p, Shard):
                    if not 0 <= p.dim < len(self.shape):
                        raise ValueError(
                            f"parameter {self.name!r} shards dim {p.dim} of a "
                            f"{len(self.shape)}-D tensor"
                        )
                elif not isinstance(p, Replicate):
                    raise TypeError(
                        f"parameter {self.name!r} placement {p!r} is neither "
                        "Replicate nor Shard"
                    )
            if all(isinstance(p, Replicate) for p in placements):
                placements = None
        object.__setattr__(self, "owners", owners)
        object.__setattr__(self, "mesh_shape", mesh_shape)
        object.__setattr__(self, "placements", placements)

RayVLLMWeightSyncClient

Talks to one or more vLLM AsyncLLM/LLM Ray actors.

Each call fans out to every handle and blocks on all of them, so a multi-actor (e.g. multi-DP) deployment is driven as one unit.

Source code in vllm/distributed/weight_transfer/clients.py

class RayVLLMWeightSyncClient:
    """Talks to one or more vLLM `AsyncLLM`/`LLM` Ray actors.

    Each call fans out to every handle and blocks on all of them, so a
    multi-actor (e.g. multi-DP) deployment is driven as one unit.
    """

    def __init__(self, handle: "ActorHandle | list[ActorHandle]") -> None:
        self.handles = handle if isinstance(handle, list) else [handle]

    def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None:
        import ray

        request = WeightTransferInitRequest(init_info=init_info)
        ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles])

    def shutdown_weight_transfer_engine(self) -> None:
        import ray

        ray.get([h.shutdown_weight_transfer_engine.remote() for h in self.handles])

    def start_weight_update(self) -> None:
        import ray

        ray.get([h.start_weight_update.remote() for h in self.handles])

    def update_weights(self, update_info: WeightTransferUpdatePayload) -> None:
        import ray

        request = WeightTransferUpdateRequest(update_info=update_info)
        ray.get([h.update_weights.remote(request) for h in self.handles])

    def finish_weight_update(self, weight_version: str | None = None) -> None:
        import ray

        ray.get([h.finish_weight_update.remote() for h in self.handles])
        if weight_version is not None:
            ray.get(
                [h.update_weight_version.remote(weight_version) for h in self.handles]
            )

Replicate dataclass

Every coordinate along this mesh axis holds the whole extent.

Source code in vllm/distributed/weight_transfer/base.py

@dataclass(frozen=True)
class Replicate:
    """Every coordinate along this mesh axis holds the whole extent."""

Shard dataclass

This mesh axis splits tensor dim dim the way torch.chunk does.

Source code in vllm/distributed/weight_transfer/base.py

@dataclass(frozen=True)
class Shard:
    """This mesh axis splits tensor dim `dim` the way `torch.chunk` does."""

    dim: int

TrainerWeightTransferEngine

Bases: ABC, Generic[TTrainerInitInfo]

Trainer-side weight transfer engine.

Symmetric to WeightTransferEngine but lives in the training process. Constructed via the trainer_init factory classmethod; carries any backend-specific state (NCCL communicators, IPC device info, transfer plans) on self. Full-resync backends (NCCL, IPC) take a WeightSource at trainer_init and replay it each round via the no-argument send_weights(). Backends that push per-round deltas instead (e.g. sparse patches) leave source as None and take their payload as a send_weights argument.

Unlike the worker engine, the trainer side does not take a WeightTransferConfig: the backend is selected from the init info's backend ClassVar (so callers pass only the init info), and the static wire params (packed, buffer sizes) ride the backend-specific TrainerInitInfo, which the sender also propagates to the worker at the init handshake.

Multi-rank trainers: trainer_init and send_weights are called on every trainer rank. Rank 0 is the sender, resolved once at trainer_init into is_sender. Non-sender ranks still run every collective (iterating the source, metadata export, IPC handle all-gather) so the group stays aligned, but each engine explicitly guards the control-plane RPCs and the transmit on self.is_sender, so only the sender touches the client.

Subclasses should define

init_info_cls: Type of backend-specific trainer init info

Methods:

• send_weights –

  Push weights to inference workers and drive the full update round

• shutdown –

  Tear down communicators / process groups. Default no-op.

• trainer_init –

  Rendezvous with the inference side and return a ready instance.

Source code in vllm/distributed/weight_transfer/base.py

class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
    """Trainer-side weight transfer engine.

    Symmetric to `WeightTransferEngine` but lives in the training process.
    Constructed via the `trainer_init` factory classmethod; carries any
    backend-specific state (NCCL communicators, IPC device info, transfer
    plans) on `self`. Full-resync backends (NCCL, IPC) take a `WeightSource` at
    `trainer_init` and replay it each round via the no-argument
    `send_weights()`. Backends that push per-round deltas instead (e.g. sparse
    patches) leave `source` as `None` and take their payload as a `send_weights`
    argument.

    Unlike the worker engine, the trainer side does not take a
    `WeightTransferConfig`: the backend is selected from the init info's
    `backend` `ClassVar` (so callers pass only the init info), and the static
    wire params (packed, buffer sizes) ride the backend-specific
    `TrainerInitInfo`, which the sender also propagates to the worker at the init
    handshake.

    Multi-rank trainers: `trainer_init` and `send_weights` are
    called on *every* trainer rank. Rank 0 is the sender, resolved once at
    `trainer_init` into `is_sender`. Non-sender ranks still run every
    collective (iterating the source, metadata export, IPC handle all-gather) so
    the group stays aligned, but each engine explicitly guards the control-plane
    RPCs and the transmit on `self.is_sender`, so only the sender touches the
    client.

    Subclasses should define:
        init_info_cls: Type of backend-specific trainer init info
    """

    # Subclasses should override this class attribute
    init_info_cls: type[TTrainerInitInfo]

    def __init__(
        self,
        *,
        client: "VLLMWeightSyncClient",
        source: "WeightSource | None" = None,
        is_sender: bool = True,
        rank: int = 0,
    ) -> None:
        self.is_sender = is_sender
        self.rank = rank
        # The real client is held on every rank; each engine only *calls* it when
        # `is_sender`, so non-sender ranks never touch the wire.
        self.client = client
        self.source = source

    @classmethod
    @abstractmethod
    def trainer_init(
        cls,
        init_info: TTrainerInitInfo,
        *,
        client: "VLLMWeightSyncClient",
        source: "WeightSource | None" = None,
    ) -> Self:
        """Rendezvous with the inference side and return a ready instance.

        Called on every trainer rank. The sender drives the full handshake via
        `client` (build the worker-side init info, call
        `client.init_weight_transfer_engine`, open the trainer-side endpoint);
        non-sender ranks skip the rendezvous and the RPC.
        """
        raise NotImplementedError

    @abstractmethod
    def send_weights(self) -> None:
        """Push weights to inference workers and drive the full update round
        trip: `start_weight_update`, `update_weights` (run concurrently with the
        trainer-side broadcast when the backend requires it), then
        `finish_weight_update`. Called on every trainer rank.
        """
        raise NotImplementedError

    def shutdown(self) -> None:
        """Tear down communicators / process groups. Default no-op."""

send_weights() abstractmethod

Push weights to inference workers and drive the full update round trip: start_weight_update, update_weights (run concurrently with the trainer-side broadcast when the backend requires it), then finish_weight_update. Called on every trainer rank.

Source code in vllm/distributed/weight_transfer/base.py

@abstractmethod
def send_weights(self) -> None:
    """Push weights to inference workers and drive the full update round
    trip: `start_weight_update`, `update_weights` (run concurrently with the
    trainer-side broadcast when the backend requires it), then
    `finish_weight_update`. Called on every trainer rank.
    """
    raise NotImplementedError

shutdown()

Tear down communicators / process groups. Default no-op.

Source code in vllm/distributed/weight_transfer/base.py

def shutdown(self) -> None:
    """Tear down communicators / process groups. Default no-op."""

trainer_init(init_info, *, client, source=None) abstractmethod classmethod

Rendezvous with the inference side and return a ready instance.

Called on every trainer rank. The sender drives the full handshake via client (build the worker-side init info, call client.init_weight_transfer_engine, open the trainer-side endpoint); non-sender ranks skip the rendezvous and the RPC.

Source code in vllm/distributed/weight_transfer/base.py

@classmethod
@abstractmethod
def trainer_init(
    cls,
    init_info: TTrainerInitInfo,
    *,
    client: "VLLMWeightSyncClient",
    source: "WeightSource | None" = None,
) -> Self:
    """Rendezvous with the inference side and return a ready instance.

    Called on every trainer rank. The sender drives the full handshake via
    `client` (build the worker-side init info, call
    `client.init_weight_transfer_engine`, open the trainer-side endpoint);
    non-sender ranks skip the rendezvous and the RPC.
    """
    raise NotImplementedError

VLLMWeightSyncClient

Bases: Protocol

Trainer-side stub for the inference engine's weight-sync control plane.

Mirrors the weight-sync methods that the inference engine exposes (EngineClient / the HTTP RLHF routes / Ray actors). A TrainerWeightTransferEngine drives the full handshake through this protocol so trainer code never has to know the transport.

All methods are synchronous and accept plain dicts (matching what the inference side already accepts). Concurrency that some backends need (e.g. NCCL must run update_weights concurrently with the trainer-side broadcast) is the engine's responsibility, not the client's, so the protocol stays a flat four-method surface that any wrapper can implement.

The protocol is structural (PEP 544), so user implementations need only define these four methods — no import or subclassing required.

Source code in vllm/distributed/weight_transfer/base.py

@runtime_checkable
class VLLMWeightSyncClient(Protocol):
    """Trainer-side stub for the inference engine's weight-sync control plane.

    Mirrors the weight-sync methods that the inference engine exposes
    (`EngineClient` / the HTTP RLHF routes / Ray actors). A
    `TrainerWeightTransferEngine` drives the full handshake through this
    protocol so trainer code never has to know the transport.

    All methods are synchronous and accept plain dicts (matching what the
    inference side already accepts). Concurrency that some backends need
    (e.g. NCCL must run `update_weights` concurrently with the trainer-side
    broadcast) is the engine's responsibility, not the client's, so the
    protocol stays a flat four-method surface that any wrapper can implement.

    The protocol is structural (PEP 544), so user implementations need only
    define these four methods — no import or subclassing required.
    """

    def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ...

    def start_weight_update(self) -> None: ...

    def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: ...

    def finish_weight_update(self, weight_version: str | None = None) -> None: ...

WeightSource

Bases: ABC

The trainer's weights, as metadata plus one or two iteration modes.

• metadata declares every parameter of the whole model, on every rank: name, wire dtype, full shape, and which ranks own it in what layout (see ParamMeta). It transfers nothing.
• iter_full() yields each owned parameter as a full tensor.
• iter_local() yields each owned parameter as this rank's piece, shaped by ParamMeta.placements.

A source implements one mode or both, and nothing is derived from one mode to the other. Each engine consumes the mode it needs: broadcast, IPC and RDT consume iter_full, NCCL M2N consumes iter_local.

Both modes yield (name, tensor) pairs: exactly owned_names(rank), in metadata order. Either may run collectives among a parameter's owners (an FSDP all-gather, one per layer if it likes); every owner iterates in metadata order, so the collectives line up. Work is enqueued on the caller's current stream. Each call starts a fresh pass.

Methods:

• iter_full –

  Yield the full tensor of each name this rank owns.

• iter_local –

  Yield this rank's piece of each name it owns.

• owned_names –

  The names rank owns, in metadata order: exactly what it iterates.

Attributes:

• metadata (list[ParamMeta]) –

  Every parameter, declared without transferring anything.

Source code in vllm/distributed/weight_transfer/base.py

class WeightSource(ABC):
    """The trainer's weights, as metadata plus one or two iteration modes.

    * `metadata` declares every parameter of the whole model, on every rank:
      name, wire dtype, full shape, and which ranks own it in what layout (see
      `ParamMeta`). It transfers nothing.
    * `iter_full()` yields each owned parameter as a full tensor.
    * `iter_local()` yields each owned parameter as this rank's piece, shaped
      by `ParamMeta.placements`.

    A source implements one mode or both, and nothing is derived from one mode
    to the other. Each engine consumes the mode it needs: broadcast, IPC and
    RDT consume `iter_full`, NCCL M2N consumes `iter_local`.

    Both modes yield `(name, tensor)` pairs: exactly `owned_names(rank)`, in
    metadata order. Either may run collectives among a parameter's
    owners (an FSDP all-gather, one per layer if it likes); every owner
    iterates in metadata order, so the collectives line up. Work is enqueued
    on the caller's current stream. Each call starts a fresh pass.
    """

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        metadata = cls.__dict__.get("metadata")
        if metadata is not None and not isinstance(metadata, cached_property):
            raise TypeError(
                f"{cls.__name__}.metadata must be a functools.cached_property: "
                "engines read it repeatedly and it must not change."
            )

    @property
    @abstractmethod
    def metadata(self) -> list[ParamMeta]:
        """Every parameter, declared without transferring anything.

        Implement it as a `functools.cached_property`. Must be identical on
        every trainer rank, including ownership: engines plan from it without
        exchanging it.
        """
        raise NotImplementedError

    def owned_names(self, rank: int) -> list[str]:
        """The names `rank` owns, in metadata order: exactly what it iterates."""
        return [m.name for m in self.metadata if m.owners is None or rank in m.owners]

    def implements_iter_full(self) -> bool:
        return type(self).iter_full is not WeightSource.iter_full

    def implements_iter_local(self) -> bool:
        return type(self).iter_local is not WeightSource.iter_local

    def iter_full(self) -> Iterator[tuple[str, torch.Tensor]]:
        """Yield the full tensor of each name this rank owns."""
        raise NotImplementedError(f"{type(self).__name__} does not implement iter_full")

    def iter_local(self) -> Iterator[tuple[str, torch.Tensor]]:
        """Yield this rank's piece of each name it owns."""
        raise NotImplementedError(
            f"{type(self).__name__} does not implement iter_local"
        )

metadata abstractmethod property

Every parameter, declared without transferring anything.

Implement it as a functools.cached_property. Must be identical on every trainer rank, including ownership: engines plan from it without exchanging it.

iter_full()

Yield the full tensor of each name this rank owns.

Source code in vllm/distributed/weight_transfer/base.py

def iter_full(self) -> Iterator[tuple[str, torch.Tensor]]:
    """Yield the full tensor of each name this rank owns."""
    raise NotImplementedError(f"{type(self).__name__} does not implement iter_full")

iter_local()

Yield this rank's piece of each name it owns.

Source code in vllm/distributed/weight_transfer/base.py

def iter_local(self) -> Iterator[tuple[str, torch.Tensor]]:
    """Yield this rank's piece of each name it owns."""
    raise NotImplementedError(
        f"{type(self).__name__} does not implement iter_local"
    )

owned_names(rank)

The names rank owns, in metadata order: exactly what it iterates.

Source code in vllm/distributed/weight_transfer/base.py

def owned_names(self, rank: int) -> list[str]:
    """The names `rank` owns, in metadata order: exactly what it iterates."""
    return [m.name for m in self.metadata if m.owners is None or rank in m.owners]

WeightTransferEngine

Bases: ABC, Generic[TInitInfo, TUpdateInfo]

Base class for weight transfer engines that handle transport of model weights from a trainer to inference workers.

This abstraction separates weight transfer transport logic from the worker implementation, allowing different backends (NCCL, CUDA IPC, RDMA[TODO]) to be plugged in.

Each engine owns its full weight-update lifecycle: start_weight_update, update_weights, and finish_weight_update. Layerwise reloading (used by checkpoint-format engines) is opted into per engine by running it inside start_weight_update/finish_weight_update. Engines that apply weights in place (e.g. sparse patches) leave those methods as no-ops.

Subclasses should define

init_info_cls: Type of backend-specific initialization info update_info_cls: Type of backend-specific update info

Methods:

• __init__ –

  Initialize the weight transfer engine.

• drain_pending –

  Block until every deferred update has been applied to the model.

• finish_weight_update –

  Finalize the current weight update.

• init_transfer_engine –

  Initialize the weight transfer mechanism.

• parse_init_info –

  Construct typed init info from dict with validation.

• parse_update_info –

  Construct typed update info from dict with validation.

• receive_weights –

  Receive weights from the trainer and load them into the model.

• reset_weight_update_target –

  Restore weight updates to the engine's default target model.

• set_weight_update_target –

  Set the model that will receive the active weight update.

• shutdown –

  Release transfer resources without shutting down the inference worker.

• start_weight_update –

  Prepare the engine for a new weight update.

• update_weights –

  Receive one weight update chunk and load it into the model.

Attributes:

• defers_processing (bool) –

  Whether update_weights returns before the weights are on the device.

Source code in vllm/distributed/weight_transfer/base.py

class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
    """Base class for weight transfer engines that handle transport of model weights
    from a trainer to inference workers.

    This abstraction separates weight transfer transport logic from the worker
    implementation, allowing different backends (NCCL, CUDA IPC, RDMA[TODO]) to be
    plugged in.

    Each engine owns its full weight-update lifecycle: `start_weight_update`,
    `update_weights`, and `finish_weight_update`. Layerwise reloading (used by
    checkpoint-format engines) is opted into per engine by running it inside
    `start_weight_update`/`finish_weight_update`. Engines that apply weights in
    place (e.g. sparse patches) leave those methods as no-ops.

    Subclasses should define:
        init_info_cls: Type of backend-specific initialization info
        update_info_cls: Type of backend-specific update info
    """

    # Subclasses should override these class attributes
    init_info_cls: type[TInitInfo]
    update_info_cls: type[TUpdateInfo]

    supports_draft_weight_update: bool = True

    defers_processing: bool = False
    """Whether `update_weights` returns before the weights are on the device.

    An engine that pipelines its GPU post-processing onto background threads
    cannot let `update_weights` synchronize the device — that would block on those
    threads and serialize the pipeline. Such an engine sets this True, omits the
    per-update sync, and guarantees completion in `finish_weight_update` instead.

    Callers that go through `finish_weight_update` need do nothing: the engine
    drains there. A caller that instead drives the tail itself — running its own
    `finalize_layerwise_reload`, say — must read this flag and call
    `drain_pending()` first, because with it set a returned `update_weights` means
    "queued", not "applied".
    """

    def drain_pending(self) -> None:
        """Block until every deferred update has been applied to the model.

        The companion to `defers_processing`: a caller that has taken over the
        update tail calls this to re-establish the guarantee that
        `finish_weight_update` would otherwise have given it. Idempotent, and a
        no-op by default — an engine that processes synchronously has nothing to
        drain, so this is always safe to call.
        """

    def __init__(
        self,
        config: WeightTransferConfig,
        vllm_config: "VllmConfig",
        device: torch.device,
        model: torch.nn.Module,
    ) -> None:
        """Initialize the weight transfer engine.

        Args:
            config: The configuration for the weight transfer engine
            vllm_config: The full vLLM config (provides parallel/model config)
            device: The device this worker's model lives on
            model: The local model instance which will receive the weights

        """
        self.config = config
        self.vllm_config = vllm_config
        self.parallel_config: ParallelConfig = vllm_config.parallel_config
        self.model_config = vllm_config.model_config
        self.device = device
        self.model = model
        self._default_model_config = self.model_config
        self._default_model = model

    def set_weight_update_target(
        self,
        model: torch.nn.Module,
        model_config: Any,
    ) -> None:
        """Set the model that will receive the active weight update."""
        self.model = model
        self.model_config = model_config

    def reset_weight_update_target(self) -> None:
        """Restore weight updates to the engine's default target model."""
        self.model = self._default_model
        self.model_config = self._default_model_config

    def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
        """Construct typed init info from dict with validation.

        Args:
            init_dict: Dictionary containing backend-specific initialization parameters

        Returns:
            Typed backend-specific init info dataclass

        Raises:
            ValueError: If init_dict is invalid for this backend

        """
        try:
            return self.init_info_cls(**init_dict)
        except TypeError as e:
            raise ValueError(
                f"Invalid init_info for {self.__class__.__name__}: {e}"
            ) from e

    def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
        """Construct typed update info from dict with validation.

        Args:
            update_dict: Dictionary containing backend-specific update parameters

        Returns:
            Typed backend-specific update info dataclass

        Raises:
            ValueError: If update_dict is invalid for this backend

        """
        try:
            return self.update_info_cls(**update_dict)
        except TypeError as e:
            raise ValueError(
                f"Invalid update_info for {self.__class__.__name__}: {e}"
            ) from e

    @abstractmethod
    def init_transfer_engine(self, init_info: TInitInfo) -> None:
        """Initialize the weight transfer mechanism.
        This is called once at the beginning of training.

        Args:
            init_info: Backend-specific initialization info

        """
        raise NotImplementedError

    @abstractmethod
    def start_weight_update(self) -> None:
        """Prepare the engine for a new weight update.

        Engines that receive weights in checkpoint format initialize layerwise reloading
        here, else this is typically a no-op.
        See: https://docs.vllm.ai/en/latest/training/layerwise/ for more details.
        """
        raise NotImplementedError

    @abstractmethod
    def finish_weight_update(self) -> None:
        """Finalize the current weight update.

        Checkpoint-format engines finalize layerwise reloading here; engines
        that apply weights in place leave this as a no-op.
        """
        raise NotImplementedError

    def update_weights(self, update_info: dict[str, Any]) -> None:
        """Receive one weight update chunk and load it into the model.

        Args:
            update_info: Dictionary containing backend-specific update info

        """
        typed_update_info = self.parse_update_info(update_info)
        self.receive_weights(typed_update_info)
        # NCCL broadcast / IPC paths may be asynchronous. Synchronize here so the
        # next step uses the new weights.
        torch.accelerator.synchronize()

    @abstractmethod
    def receive_weights(self, update_info: TUpdateInfo) -> None:
        """Receive weights from the trainer and load them into the model.

        Args:
            update_info: Backend-specific update info containing parameter metadata
                        and any backend-specific data

        """
        raise NotImplementedError

    @abstractmethod
    def shutdown(self) -> None:
        """Release transfer resources without shutting down the inference worker."""
        raise NotImplementedError

defers_processing = False class-attribute instance-attribute

Whether update_weights returns before the weights are on the device.

An engine that pipelines its GPU post-processing onto background threads cannot let update_weights synchronize the device — that would block on those threads and serialize the pipeline. Such an engine sets this True, omits the per-update sync, and guarantees completion in finish_weight_update instead.

Callers that go through finish_weight_update need do nothing: the engine drains there. A caller that instead drives the tail itself — running its own finalize_layerwise_reload, say — must read this flag and call drain_pending() first, because with it set a returned update_weights means "queued", not "applied".

__init__(config, vllm_config, device, model)

Initialize the weight transfer engine.

Parameters:

• config

  (WeightTransferConfig) –

  The configuration for the weight transfer engine

• vllm_config

  (VllmConfig) –

  The full vLLM config (provides parallel/model config)

• device

  (device) –

  The device this worker's model lives on

• model

  (Module) –

  The local model instance which will receive the weights

Source code in vllm/distributed/weight_transfer/base.py

def __init__(
    self,
    config: WeightTransferConfig,
    vllm_config: "VllmConfig",
    device: torch.device,
    model: torch.nn.Module,
) -> None:
    """Initialize the weight transfer engine.

    Args:
        config: The configuration for the weight transfer engine
        vllm_config: The full vLLM config (provides parallel/model config)
        device: The device this worker's model lives on
        model: The local model instance which will receive the weights

    """
    self.config = config
    self.vllm_config = vllm_config
    self.parallel_config: ParallelConfig = vllm_config.parallel_config
    self.model_config = vllm_config.model_config
    self.device = device
    self.model = model
    self._default_model_config = self.model_config
    self._default_model = model

drain_pending()

Block until every deferred update has been applied to the model.

The companion to defers_processing: a caller that has taken over the update tail calls this to re-establish the guarantee that finish_weight_update would otherwise have given it. Idempotent, and a no-op by default — an engine that processes synchronously has nothing to drain, so this is always safe to call.

Source code in vllm/distributed/weight_transfer/base.py

def drain_pending(self) -> None:
    """Block until every deferred update has been applied to the model.

    The companion to `defers_processing`: a caller that has taken over the
    update tail calls this to re-establish the guarantee that
    `finish_weight_update` would otherwise have given it. Idempotent, and a
    no-op by default — an engine that processes synchronously has nothing to
    drain, so this is always safe to call.
    """

finish_weight_update() abstractmethod

Finalize the current weight update.

Checkpoint-format engines finalize layerwise reloading here; engines that apply weights in place leave this as a no-op.

Source code in vllm/distributed/weight_transfer/base.py

@abstractmethod
def finish_weight_update(self) -> None:
    """Finalize the current weight update.

    Checkpoint-format engines finalize layerwise reloading here; engines
    that apply weights in place leave this as a no-op.
    """
    raise NotImplementedError

init_transfer_engine(init_info) abstractmethod

Initialize the weight transfer mechanism. This is called once at the beginning of training.

Parameters:

• init_info

  (TInitInfo) –

  Backend-specific initialization info

Source code in vllm/distributed/weight_transfer/base.py

@abstractmethod
def init_transfer_engine(self, init_info: TInitInfo) -> None:
    """Initialize the weight transfer mechanism.
    This is called once at the beginning of training.

    Args:
        init_info: Backend-specific initialization info

    """
    raise NotImplementedError

parse_init_info(init_dict)

Construct typed init info from dict with validation.

Parameters:

• init_dict

  (dict[str, Any]) –

  Dictionary containing backend-specific initialization parameters

Returns:

• TInitInfo –

  Typed backend-specific init info dataclass

Raises:

• ValueError –

  If init_dict is invalid for this backend

Source code in vllm/distributed/weight_transfer/base.py

def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
    """Construct typed init info from dict with validation.

    Args:
        init_dict: Dictionary containing backend-specific initialization parameters

    Returns:
        Typed backend-specific init info dataclass

    Raises:
        ValueError: If init_dict is invalid for this backend

    """
    try:
        return self.init_info_cls(**init_dict)
    except TypeError as e:
        raise ValueError(
            f"Invalid init_info for {self.__class__.__name__}: {e}"
        ) from e

parse_update_info(update_dict)

Construct typed update info from dict with validation.

Parameters:

• update_dict

  (dict[str, Any]) –

  Dictionary containing backend-specific update parameters

Returns:

• TUpdateInfo –

  Typed backend-specific update info dataclass

Raises:

• ValueError –

  If update_dict is invalid for this backend

Source code in vllm/distributed/weight_transfer/base.py

def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
    """Construct typed update info from dict with validation.

    Args:
        update_dict: Dictionary containing backend-specific update parameters

    Returns:
        Typed backend-specific update info dataclass

    Raises:
        ValueError: If update_dict is invalid for this backend

    """
    try:
        return self.update_info_cls(**update_dict)
    except TypeError as e:
        raise ValueError(
            f"Invalid update_info for {self.__class__.__name__}: {e}"
        ) from e

receive_weights(update_info) abstractmethod

Receive weights from the trainer and load them into the model.

Parameters:

• update_info

  (TUpdateInfo) –

  Backend-specific update info containing parameter metadata and any backend-specific data

Source code in vllm/distributed/weight_transfer/base.py

@abstractmethod
def receive_weights(self, update_info: TUpdateInfo) -> None:
    """Receive weights from the trainer and load them into the model.

    Args:
        update_info: Backend-specific update info containing parameter metadata
                    and any backend-specific data

    """
    raise NotImplementedError

reset_weight_update_target()

Restore weight updates to the engine's default target model.

Source code in vllm/distributed/weight_transfer/base.py

def reset_weight_update_target(self) -> None:
    """Restore weight updates to the engine's default target model."""
    self.model = self._default_model
    self.model_config = self._default_model_config

set_weight_update_target(model, model_config)

Set the model that will receive the active weight update.

Source code in vllm/distributed/weight_transfer/base.py

def set_weight_update_target(
    self,
    model: torch.nn.Module,
    model_config: Any,
) -> None:
    """Set the model that will receive the active weight update."""
    self.model = model
    self.model_config = model_config

shutdown() abstractmethod

Release transfer resources without shutting down the inference worker.

Source code in vllm/distributed/weight_transfer/base.py

@abstractmethod
def shutdown(self) -> None:
    """Release transfer resources without shutting down the inference worker."""
    raise NotImplementedError

start_weight_update() abstractmethod

Prepare the engine for a new weight update.

Engines that receive weights in checkpoint format initialize layerwise reloading here, else this is typically a no-op. See: https://docs.vllm.ai/en/latest/training/layerwise/ for more details.

Source code in vllm/distributed/weight_transfer/base.py

@abstractmethod
def start_weight_update(self) -> None:
    """Prepare the engine for a new weight update.

    Engines that receive weights in checkpoint format initialize layerwise reloading
    here, else this is typically a no-op.
    See: https://docs.vllm.ai/en/latest/training/layerwise/ for more details.
    """
    raise NotImplementedError

update_weights(update_info)

Receive one weight update chunk and load it into the model.

Parameters:

• update_info

  (dict[str, Any]) –

  Dictionary containing backend-specific update info

Source code in vllm/distributed/weight_transfer/base.py

def update_weights(self, update_info: dict[str, Any]) -> None:
    """Receive one weight update chunk and load it into the model.

    Args:
        update_info: Dictionary containing backend-specific update info

    """
    typed_update_info = self.parse_update_info(update_info)
    self.receive_weights(typed_update_info)
    # NCCL broadcast / IPC paths may be asynchronous. Synchronize here so the
    # next step uses the new weights.
    torch.accelerator.synchronize()

WeightTransferEngineFactory

Factory for creating weight transfer engines with lazy loading.

This factory implements a registry pattern that supports: - Lazy loading: Engine modules are only imported when actually needed - Extensibility: Custom engines can be registered at runtime - Centralized registration: All built-in engines registered in one place

Methods:

• create_engine –

  Create a weight transfer engine instance.

• register_engine –

  Register an engine with lazy-loading or direct class reference.

Source code in vllm/distributed/weight_transfer/factory.py

class WeightTransferEngineFactory:
    """Factory for creating weight transfer engines with lazy loading.

    This factory implements a registry pattern that supports:
    - Lazy loading: Engine modules are only imported when actually needed
    - Extensibility: Custom engines can be registered at runtime
    - Centralized registration: All built-in engines registered in one place
    """

    _registry: dict[str, Callable[[], type[WeightTransferEngine]]] = {}

    @classmethod
    def register_engine(
        cls,
        name: str,
        module_path_or_cls: str | type[WeightTransferEngine],
        class_name: str | None = None,
    ) -> None:
        """Register an engine with lazy-loading or direct class reference.

        Supports two calling conventions:
        1. Lazy loading: register_engine(name, module_path, class_name)
        2. Direct class: register_engine(name, engine_cls)

        Args:
            name: The name to register the engine under (e.g., "nccl")
            module_path_or_cls: Either a module path string for lazy loading,
                or the engine class directly
            class_name: Name of the engine class (required if module_path is string)

        Raises:
            ValueError: If an engine with the same name is already registered

        """
        if name in cls._registry:
            raise ValueError(f"Weight transfer engine '{name}' is already registered.")

        if isinstance(module_path_or_cls, str):
            # Lazy loading path
            module_path = module_path_or_cls
            if class_name is None:
                raise ValueError(
                    "class_name is required when registering with module path"
                )

            def loader() -> type[WeightTransferEngine]:
                module = importlib.import_module(module_path)
                return getattr(module, class_name)

            cls._registry[name] = loader
        else:
            # Direct class registration
            engine_cls = module_path_or_cls
            cls._registry[name] = lambda: engine_cls

    @classmethod
    def create_engine(
        cls,
        config: "WeightTransferConfig",
        vllm_config: "VllmConfig",
        device: "torch.device",
        model: "torch.nn.Module",
    ) -> WeightTransferEngine:
        """Create a weight transfer engine instance.

        Args:
            config: Weight transfer configuration containing the backend name
            vllm_config: The full vLLM config (provides parallel/model config)
            device: The device this worker's model lives on
            model: The local model instance which will receive the weights

        Returns:
            An initialized weight transfer engine instance

        Raises:
            ValueError: If the backend is not registered

        """
        backend = config.backend
        if backend not in cls._registry:
            available = list(cls._registry.keys())
            raise ValueError(
                f"Invalid weight transfer backend: {backend}. "
                f"Available engines: {available}"
            )
        engine_cls = cls._registry[backend]()

        logger.info(
            "Creating weight transfer engine: %s",
            engine_cls.__name__,
        )

        return engine_cls(config, vllm_config, device, model)

create_engine(config, vllm_config, device, model) classmethod

Create a weight transfer engine instance.

Parameters:

• config

  (WeightTransferConfig) –

  Weight transfer configuration containing the backend name

• vllm_config

  (VllmConfig) –

  The full vLLM config (provides parallel/model config)

• device

  (device) –

  The device this worker's model lives on

• model

  (Module) –

  The local model instance which will receive the weights

Returns:

• WeightTransferEngine –

  An initialized weight transfer engine instance

Raises:

• ValueError –

  If the backend is not registered

Source code in vllm/distributed/weight_transfer/factory.py

@classmethod
def create_engine(
    cls,
    config: "WeightTransferConfig",
    vllm_config: "VllmConfig",
    device: "torch.device",
    model: "torch.nn.Module",
) -> WeightTransferEngine:
    """Create a weight transfer engine instance.

    Args:
        config: Weight transfer configuration containing the backend name
        vllm_config: The full vLLM config (provides parallel/model config)
        device: The device this worker's model lives on
        model: The local model instance which will receive the weights

    Returns:
        An initialized weight transfer engine instance

    Raises:
        ValueError: If the backend is not registered

    """
    backend = config.backend
    if backend not in cls._registry:
        available = list(cls._registry.keys())
        raise ValueError(
            f"Invalid weight transfer backend: {backend}. "
            f"Available engines: {available}"
        )
    engine_cls = cls._registry[backend]()

    logger.info(
        "Creating weight transfer engine: %s",
        engine_cls.__name__,
    )

    return engine_cls(config, vllm_config, device, model)

register_engine(name, module_path_or_cls, class_name=None) classmethod

Register an engine with lazy-loading or direct class reference.

Supports two calling conventions: 1. Lazy loading: register_engine(name, module_path, class_name) 2. Direct class: register_engine(name, engine_cls)

Parameters:

• name

  (str) –

  The name to register the engine under (e.g., "nccl")

• module_path_or_cls

  (str | type[WeightTransferEngine]) –

  Either a module path string for lazy loading, or the engine class directly

• class_name

  (str | None, default: None ) –

  Name of the engine class (required if module_path is string)

Raises:

• ValueError –

  If an engine with the same name is already registered

Source code in vllm/distributed/weight_transfer/factory.py

@classmethod
def register_engine(
    cls,
    name: str,
    module_path_or_cls: str | type[WeightTransferEngine],
    class_name: str | None = None,
) -> None:
    """Register an engine with lazy-loading or direct class reference.

    Supports two calling conventions:
    1. Lazy loading: register_engine(name, module_path, class_name)
    2. Direct class: register_engine(name, engine_cls)

    Args:
        name: The name to register the engine under (e.g., "nccl")
        module_path_or_cls: Either a module path string for lazy loading,
            or the engine class directly
        class_name: Name of the engine class (required if module_path is string)

    Raises:
        ValueError: If an engine with the same name is already registered

    """
    if name in cls._registry:
        raise ValueError(f"Weight transfer engine '{name}' is already registered.")

    if isinstance(module_path_or_cls, str):
        # Lazy loading path
        module_path = module_path_or_cls
        if class_name is None:
            raise ValueError(
                "class_name is required when registering with module path"
            )

        def loader() -> type[WeightTransferEngine]:
            module = importlib.import_module(module_path)
            return getattr(module, class_name)

        cls._registry[name] = loader
    else:
        # Direct class registration
        engine_cls = module_path_or_cls
        cls._registry[name] = lambda: engine_cls

WeightTransferTrainerFactory

Factory for creating trainer-side weight transfer engines.

Parallel to WeightTransferEngineFactory, with its own lazy-import registry. The trainer-side and worker-side registries are kept separate: they share backend names by convention, but the trainer process never instantiates a worker engine and vice versa, so unifying them would only couple the import graphs.

Methods:

• register_engine –

  Register a trainer engine. Same conventions as

• trainer_init –

  Build and rendezvous a ready-to-send trainer engine.

Source code in vllm/distributed/weight_transfer/factory.py

class WeightTransferTrainerFactory:
    """Factory for creating trainer-side weight transfer engines.

    Parallel to `WeightTransferEngineFactory`, with its own lazy-import
    registry. The trainer-side and worker-side registries are kept separate:
    they share backend names by convention, but the trainer process never
    instantiates a worker engine and vice versa, so unifying them would only
    couple the import graphs.
    """

    _registry: dict[str, Callable[[], type[TrainerWeightTransferEngine]]] = {}

    @classmethod
    def register_engine(
        cls,
        name: str,
        module_path_or_cls: "str | type[TrainerWeightTransferEngine]",
        class_name: str | None = None,
    ) -> None:
        """Register a trainer engine. Same conventions as
        `WeightTransferEngineFactory.register_engine`."""
        if name in cls._registry:
            raise ValueError(
                f"Weight transfer trainer engine '{name}' is already registered."
            )

        if isinstance(module_path_or_cls, str):
            module_path = module_path_or_cls
            if class_name is None:
                raise ValueError(
                    "class_name is required when registering with module path"
                )

            def loader() -> type[TrainerWeightTransferEngine]:
                module = importlib.import_module(module_path)
                return getattr(module, class_name)

            cls._registry[name] = loader
        else:
            engine_cls = module_path_or_cls
            cls._registry[name] = lambda: engine_cls

    @classmethod
    def trainer_init(
        cls,
        init_info: "TrainerInitInfo",
        *,
        client: "VLLMWeightSyncClient",
        source: "WeightSource | None" = None,
    ) -> TrainerWeightTransferEngine:
        """Build and rendezvous a ready-to-send trainer engine.

        Called on every trainer rank (multi-rank trainers construct on all
        ranks; the sender is resolved inside the engine's ``trainer_init``).

        The trainer side takes no `WeightTransferConfig` and no separate
        `backend` argument: the backend is read from ``init_info.backend`` (a
        `ClassVar` on each `TrainerInitInfo` subclass), and the static wire
        params ride `init_info`.

        Args:
            init_info: Backend-specific trainer init info. Its `backend`
                selects the engine; it also carries the wire params (e.g.
                `packed`).
            client: Inference-side control-plane client.
            source: `WeightSource` of `(name, tensor)` pairs to send each round,
                for full-resync backends (NCCL, IPC). Sparse backend
                omits it and passes its per-round payload to `send_weights`.

        Raises:
            ValueError: If `init_info.backend` is not registered.

        """
        backend = init_info.backend
        if backend not in cls._registry:
            available = list(cls._registry.keys())
            raise ValueError(
                f"Invalid weight transfer backend: {backend}. "
                f"Available trainer engines: {available}"
            )
        engine_cls = cls._registry[backend]()

        logger.info(
            "Creating weight transfer trainer engine: %s",
            engine_cls.__name__,
        )

        return engine_cls.trainer_init(
            init_info=init_info,
            client=client,
            source=source,
        )

register_engine(name, module_path_or_cls, class_name=None) classmethod

Register a trainer engine. Same conventions as WeightTransferEngineFactory.register_engine.

Source code in vllm/distributed/weight_transfer/factory.py

@classmethod
def register_engine(
    cls,
    name: str,
    module_path_or_cls: "str | type[TrainerWeightTransferEngine]",
    class_name: str | None = None,
) -> None:
    """Register a trainer engine. Same conventions as
    `WeightTransferEngineFactory.register_engine`."""
    if name in cls._registry:
        raise ValueError(
            f"Weight transfer trainer engine '{name}' is already registered."
        )

    if isinstance(module_path_or_cls, str):
        module_path = module_path_or_cls
        if class_name is None:
            raise ValueError(
                "class_name is required when registering with module path"
            )

        def loader() -> type[TrainerWeightTransferEngine]:
            module = importlib.import_module(module_path)
            return getattr(module, class_name)

        cls._registry[name] = loader
    else:
        engine_cls = module_path_or_cls
        cls._registry[name] = lambda: engine_cls

trainer_init(init_info, *, client, source=None) classmethod

Build and rendezvous a ready-to-send trainer engine.

Called on every trainer rank (multi-rank trainers construct on all ranks; the sender is resolved inside the engine's trainer_init).

The trainer side takes no WeightTransferConfig and no separate backend argument: the backend is read from init_info.backend (a ClassVar on each TrainerInitInfo subclass), and the static wire params ride init_info.

Parameters:

• init_info

  (TrainerInitInfo) –

  Backend-specific trainer init info. Its backend selects the engine; it also carries the wire params (e.g. packed).

• client

  (VLLMWeightSyncClient) –

  Inference-side control-plane client.

• source

  (WeightSource | None, default: None ) –

  WeightSource of (name, tensor) pairs to send each round, for full-resync backends (NCCL, IPC). Sparse backend omits it and passes its per-round payload to send_weights.

Raises:

• ValueError –

  If init_info.backend is not registered.

Source code in vllm/distributed/weight_transfer/factory.py

@classmethod
def trainer_init(
    cls,
    init_info: "TrainerInitInfo",
    *,
    client: "VLLMWeightSyncClient",
    source: "WeightSource | None" = None,
) -> TrainerWeightTransferEngine:
    """Build and rendezvous a ready-to-send trainer engine.

    Called on every trainer rank (multi-rank trainers construct on all
    ranks; the sender is resolved inside the engine's ``trainer_init``).

    The trainer side takes no `WeightTransferConfig` and no separate
    `backend` argument: the backend is read from ``init_info.backend`` (a
    `ClassVar` on each `TrainerInitInfo` subclass), and the static wire
    params ride `init_info`.

    Args:
        init_info: Backend-specific trainer init info. Its `backend`
            selects the engine; it also carries the wire params (e.g.
            `packed`).
        client: Inference-side control-plane client.
        source: `WeightSource` of `(name, tensor)` pairs to send each round,
            for full-resync backends (NCCL, IPC). Sparse backend
            omits it and passes its per-round payload to `send_weights`.

    Raises:
        ValueError: If `init_info.backend` is not registered.

    """
    backend = init_info.backend
    if backend not in cls._registry:
        available = list(cls._registry.keys())
        raise ValueError(
            f"Invalid weight transfer backend: {backend}. "
            f"Available trainer engines: {available}"
        )
    engine_cls = cls._registry[backend]()

    logger.info(
        "Creating weight transfer trainer engine: %s",
        engine_cls.__name__,
    )

    return engine_cls.trainer_init(
        init_info=init_info,
        client=client,
        source=source,
    )