vllm.distributed.weight_transfer.base

Base class for weight transfer engines.

Classes:

• ModuleSource –

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

• ParamMeta –

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

• Replicate –

  Every coordinate along this mesh axis holds the whole extent.

• Shard –

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

• TrainerInitInfo –

  Base trainer-side init info: which trainer rank drives the transfer.

• 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

• WeightTransferInitInfo –

  Base class for backend-specific initialization info.

• WeightTransferInitRequest –

  API-level weight transfer initialization request.

• WeightTransferUpdateInfo –

  Base class for backend-specific weight update info.

• WeightTransferUpdateRequest –

  API-level weight update request.

Functions:

• layerwise_groups –

  Partition flat parameter names into one group per decoder layer, keyed on

• materialize_full_tensor –

  Return a full, locally-materialized tensor ready to send.

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)

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

TrainerInitInfo dataclass

Base trainer-side init info: which trainer rank drives the transfer.

rank is this trainer process's rank, provided explicitly by the caller — the engine does not read it from a global process group, which is ambiguous once several groups (FSDP / TP / PP / EP) exist. Rank 0 is always the sender: only it opens the endpoint and drives the inference-side RPCs, while every rank still runs the trainer-side collectives. Backend subclasses add their own (positional) fields; rank is keyword-only so that ordering never conflicts.

Every concrete subclass sets a class-level backend string (the same key it registers under in WeightTransferTrainerFactory). The factory reads it to dispatch, so callers pass only the init info/ It is a ClassVar (a fixed per-backend constant), so it is not an __init__ field.

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

@dataclass
class TrainerInitInfo:
    """Base trainer-side init info: which trainer rank drives the transfer.

    `rank` is this trainer process's rank, provided **explicitly** by the
    caller — the engine does not read it from a global process group, which is
    ambiguous once several groups (FSDP / TP / PP / EP) exist. Rank 0 is always
    the sender: only it opens the endpoint and drives the inference-side RPCs,
    while every rank still runs the trainer-side collectives. Backend subclasses
    add their own (positional) fields; `rank` is keyword-only so that ordering
    never conflicts.

    Every concrete subclass sets a class-level `backend` string (the same key it
    registers under in `WeightTransferTrainerFactory`). The factory reads it to
    dispatch, so callers pass only the init info/ It is a `ClassVar`
    (a fixed per-backend constant), so it is not an ``__init__`` field.
    """

    backend: ClassVar[str]

    rank: int = field(kw_only=True)

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        if not getattr(cls, "backend", None):
            raise TypeError(
                f"{cls.__name__} must set a class-level `backend` string "
                "(the WeightTransferTrainerFactory registry key)."
            )

    @property
    def is_sender(self) -> bool:
        return self.rank == 0

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()

WeightTransferInitInfo dataclass

Bases: ABC

Base class for backend-specific initialization info.

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

@dataclass
class WeightTransferInitInfo(ABC):  # noqa: B024
    """Base class for backend-specific initialization info."""

    pass

WeightTransferInitRequest dataclass

API-level weight transfer initialization request.

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

@dataclass
class WeightTransferInitRequest:
    """API-level weight transfer initialization request."""

    init_info: dict[str, Any] = field(default_factory=dict)

WeightTransferUpdateInfo dataclass

Bases: ABC

Base class for backend-specific weight update info.

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

@dataclass
class WeightTransferUpdateInfo(ABC):  # noqa: B024
    """Base class for backend-specific weight update info."""

    pass

WeightTransferUpdateRequest dataclass

API-level weight update request.

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

@dataclass
class WeightTransferUpdateRequest:
    """API-level weight update request."""

    update_info: WeightTransferUpdatePayload = field(default_factory=dict)

_dtensor_layout(param)

(owners, mesh_shape, placements) of a DTensor, or None for a plain tensor. Placements are None when the layout is not plain Shard / Replicate (e.g. Partial, or FSDP+TP's strided shards).

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

def _dtensor_layout(
    param: torch.Tensor,
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[Placement, ...] | None] | None:
    """`(owners, mesh_shape, placements)` of a DTensor, or None for a plain
    tensor. Placements are None when the layout is not plain `Shard` /
    `Replicate` (e.g. `Partial`, or FSDP+TP's strided shards)."""
    device_mesh = getattr(param, "device_mesh", None)
    if device_mesh is None:
        return None
    from torch.distributed.tensor import Replicate as TorchReplicate
    from torch.distributed.tensor import Shard as TorchShard

    grid = device_mesh.mesh
    placements: list[Placement] = []
    for p in param.placements:  # type: ignore[attr-defined]
        if type(p) is TorchShard:
            placements.append(Shard(p.dim))
        elif type(p) is TorchReplicate:
            placements.append(Replicate())
        else:
            return tuple(grid.flatten().tolist()), tuple(grid.shape), None
    return tuple(grid.flatten().tolist()), tuple(grid.shape), tuple(placements)

_stack_key(name)

(prefix, index) of the OUTERMOST integer segment, or None if there is none.

Outermost is what keeps a MoE layer whole: model.layers.3.mlp.experts.7.w1 keys on the layer, not the expert.

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

def _stack_key(name: str) -> tuple[str, int] | None:
    """``(prefix, index)`` of the OUTERMOST integer segment, or None if there is
    none.

    Outermost is what keeps a MoE layer whole:
    ``model.layers.3.mlp.experts.7.w1`` keys on the layer, not the expert.
    """
    parts = name.split(".")
    for i, part in enumerate(parts):
        if part.isdigit():
            return ".".join(parts[:i]), int(part)
    return None

layerwise_groups(names)

Partition flat parameter names into one group per decoder layer, keyed on the outermost index segment of each name.

This defines what a group index means for backends that transfer a layer at a time: index g names the same group on every trainer rank and every consumer, because it is derived from metadata order, which every rank shares.

Keying on the index rather than a literal prefix needs no per-architecture naming table: model.layers.0., model.language_model.layers.0., transformer.h.0., backbone.layers.0. and a vision tower's visual.blocks.0. all partition alike. Matching one fixed prefix does not, and its failure is silent — every name lands in a single group holding the whole model, which defeats the per-layer bound below.

Un-indexed names split by POSITION relative to the first indexed one: the pre block (embeddings) and the post block (the final norm, lm_head, and any inter-stack projector). Post lands last however early it arrived, which is what a pipeline-parallel source needs — Megatron-Bridge streams the last stage's output block before its layers.

Stacks come out in first-appearance order of their prefix and ascending index within it, whatever order the source yielded them, so a source can normalize an arbitrary export order by flattening this partition.

Backends that gather and free per group (sharded RDT) also use it as the unit of transfer, which bounds their buffer sizes: without it a whole model becomes one chunk.

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

def layerwise_groups(names: list[str]) -> list[list[str]]:
    """Partition flat parameter names into one group per decoder layer, keyed on
    the outermost index segment of each name.

    This defines what a *group index* means for backends that transfer a layer
    at a time: index *g* names the same group on every trainer rank and every
    consumer, because it is derived from `metadata` order, which every rank
    shares.

    Keying on the index rather than a literal prefix needs no per-architecture
    naming table: ``model.layers.0.``, ``model.language_model.layers.0.``,
    ``transformer.h.0.``, ``backbone.layers.0.`` and a vision tower's
    ``visual.blocks.0.`` all partition alike. Matching one fixed prefix does not,
    and its failure is silent — every name lands in a single group holding the
    whole model, which defeats the per-layer bound below.

    Un-indexed names split by POSITION relative to the first indexed one: the pre
    block (embeddings) and the post block (the final norm, `lm_head`, and any
    inter-stack projector). Post lands last however early it arrived, which is
    what a pipeline-parallel source needs — Megatron-Bridge streams the last
    stage's output block *before* its layers.

    Stacks come out in first-appearance order of their prefix and ascending index
    within it, whatever order the source yielded them, so a source can normalize
    an arbitrary export order by flattening this partition.

    Backends that gather and free per group (sharded RDT) also use it as the unit
    of transfer, which bounds their buffer sizes: without it a whole model becomes
    one chunk.
    """
    pre: list[str] = []
    post: list[str] = []
    stacks: dict[tuple[str, int], list[str]] = {}
    order: list[tuple[str, int]] = []
    for name in names:
        key = _stack_key(name)
        if key is None:
            (post if order else pre).append(name)
            continue
        if key not in stacks:
            stacks[key] = []
            order.append(key)
        stacks[key].append(name)

    prefix_rank: dict[str, int] = {}
    for key in order:
        prefix_rank.setdefault(key[0], len(prefix_rank))
    order.sort(key=lambda key: (prefix_rank[key[0]], key[1]))

    groups: list[list[str]] = [pre] if pre else []
    groups += [stacks[key] for key in order]
    if post:
        groups.append(post)
    return groups

materialize_full_tensor(tensor)

Return a full, locally-materialized tensor ready to send.

FSDP shards (DTensors) expose full_tensor(), a collective all-gather; regular tensors do not and are returned unchanged. Sources call this at iteration time, so the gather happens once per round; reading .shape/.dtype for metadata does not trigger it.

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

def materialize_full_tensor(tensor: torch.Tensor) -> torch.Tensor:
    """Return a full, locally-materialized tensor ready to send.

    FSDP shards (DTensors) expose `full_tensor()`, a collective all-gather;
    regular tensors do not and are returned unchanged. Sources call this at
    iteration time, so the gather happens once per round; reading
    `.shape`/`.dtype` for metadata does not trigger it.
    """
    full_tensor = getattr(tensor, "full_tensor", None)
    return full_tensor() if callable(full_tensor) else tensor