[rl] Stateful Trainer Send: New Abstractions [1/N] (#48042)

Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>

tests/distributed/test_weight_transfer.py

@@ -17,7 +17,19 @@ from torch.multiprocessing.reductions import reduce_tensor
 
 from vllm.config.parallel import ParallelConfig
 from vllm.config.weight_transfer import WeightTransferConfig
-from vllm.distributed.weight_transfer import WeightTransferEngineFactory
+from vllm.distributed.weight_transfer import (
+    HTTPVLLMWeightSyncClient,
+    ModuleSource,
+    RayVLLMWeightSyncClient,
+    TrainerWeightTransferEngine,
+    VLLMWeightSyncClient,
+    WeightTransferEngineFactory,
+    WeightTransferTrainerFactory,
+)
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
 from vllm.distributed.weight_transfer.ipc_engine import (
     IPCWeightTransferEngine,
     IPCWeightTransferInitInfo,
@@ -1214,3 +1226,194 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
 
     with pytest.raises(ValueError, match="IPC handle not found"):
         engine.receive_weights(update_info)
+
+
+class RecordingClient:
+    """A fake VLLMWeightSyncClient that records the order of calls."""
+
+    def __init__(self):
+        self.order: list[str] = []
+        self.last_init_info: dict | None = None
+        self.last_update_info: dict | None = None
+
+    def init_weight_transfer_engine(self, init_info: dict) -> None:
+        self.order.append("init")
+        self.last_init_info = init_info
+
+    def start_weight_update(self) -> None:
+        self.order.append("start")
+
+    def update_weights(self, update_info: dict) -> None:
+        self.order.append("update")
+        self.last_update_info = update_info
+
+    def finish_weight_update(self) -> None:
+        self.order.append("finish")
+
+
+def _module_with(*pairs):
+    """A tiny nn.Module exposing the given (name, tensor) pairs as parameters,
+    so trainer tests can build a ModuleSource without a real model."""
+    module = torch.nn.Module()
+    for name, tensor in pairs:
+        module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False))
+    return module
+
+
+class _DummyTrainerEngine(TrainerWeightTransferEngine):
+    """Minimal concrete trainer engine to exercise base-class + factory."""
+
+    @classmethod
+    def trainer_init(cls, config, init_info, *, client, source):
+        return cls(config, client=client, source=source)
+
+    def send_weights(self):
+        pass
+
+
+class TestTrainerClients:
+    """Structural protocol conformance for the built-in clients."""
+
+    def test_recording_client_is_protocol(self):
+        assert isinstance(RecordingClient(), VLLMWeightSyncClient)
+
+    def test_http_client_is_protocol(self):
+        assert isinstance(
+            HTTPVLLMWeightSyncClient("http://localhost:8000"), VLLMWeightSyncClient
+        )
+
+    def test_ray_client_is_protocol(self):
+        assert isinstance(RayVLLMWeightSyncClient(MagicMock()), VLLMWeightSyncClient)
+
+    def test_ray_client_sends_typed_requests(self, monkeypatch):
+        """Ray client must hand the actor typed Request objects, not raw dicts."""
+        import ray
+
+        monkeypatch.setattr(ray, "get", lambda refs: None)
+        handle = MagicMock()
+        client = RayVLLMWeightSyncClient(handle)
+
+        client.init_weight_transfer_engine({"master_addr": "x"})
+        (init_req,), _ = handle.init_weight_transfer_engine.remote.call_args
+        assert isinstance(init_req, WeightTransferInitRequest)
+        assert init_req.init_info == {"master_addr": "x"}
+
+        client.update_weights({"names": ["w"]})
+        (update_req,), _ = handle.update_weights.remote.call_args
+        assert isinstance(update_req, WeightTransferUpdateRequest)
+        assert update_req.update_info == {"names": ["w"]}
+
+    def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch):
+        """HTTP update_weights must encode raw ipc_handles as a base64 pickle."""
+        captured = {}
+
+        def fake_post(self, path, json=None):
+            captured["path"] = path
+            captured["json"] = json
+
+        monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
+        client = HTTPVLLMWeightSyncClient("http://localhost:8000")
+        client.update_weights({"names": ["w"], "ipc_handles": [{"gpu": ("args",)}]})
+        sent = captured["json"]["update_info"]
+        assert "ipc_handles" not in sent
+        assert "ipc_handles_pickled" in sent
+        assert pickle.loads(base64.b64decode(sent["ipc_handles_pickled"])) == [
+            {"gpu": ("args",)}
+        ]
+
+    def test_http_client_passes_through_nccl_update_info(self, monkeypatch):
+        """NCCL update_info has only JSON-native fields and passes unchanged."""
+        captured = {}
+
+        def fake_post(self, path, json=None):
+            captured["json"] = json
+
+        monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
+        client = HTTPVLLMWeightSyncClient("http://localhost:8000")
+        update_info = {"names": ["w"], "dtype_names": ["float32"], "shapes": [[4]]}
+        client.update_weights(update_info)
+        assert captured["json"]["update_info"] == update_info
+
+
+class TestModuleSource:
+    """`ModuleSource` metadata vs. materialized iteration (dense, no GPU)."""
+
+    def test_metadata_reads_shape_and_dtype(self):
+        source = ModuleSource(
+            _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3)))
+        )
+        meta = source.metadata()
+        assert [m.name for m in meta] == ["w", "b"]
+        assert [m.shape for m in meta] == [(2, 3), (3,)]
+        assert all(m.dtype == torch.float32 for m in meta)
+
+    def test_iteration_yields_materialized_tensors(self):
+        w = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        source = ModuleSource(_module_with(("w", w)))
+        pairs = list(source)
+        assert [name for name, _ in pairs] == ["w"]
+        assert torch.equal(pairs[0][1], w)
+
+    def test_source_is_reiterable(self):
+        source = ModuleSource(_module_with(("w", torch.zeros(2))))
+        assert [n for n, _ in source] == [n for n, _ in source] == ["w"]
+
+
+class TestTrainerFactory:
+    """WeightTransferTrainerFactory registry mechanics."""
+
+    def test_builtin_registry_has_no_trainer_backends_yet(self):
+        # Concrete backends register in the per-backend migration PRs.
+        assert WeightTransferTrainerFactory._registry == {}
+
+    def test_register_and_dispatch(self):
+        saved = dict(WeightTransferTrainerFactory._registry)
+        try:
+            WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine)
+            engine = WeightTransferTrainerFactory.trainer_init(
+                "dummy",
+                WeightTransferConfig(backend="dummy"),
+                MagicMock(),
+                client=RecordingClient(),
+                source=ModuleSource(_module_with(("w", torch.zeros(2)))),
+            )
+            assert isinstance(engine, _DummyTrainerEngine)
+            with pytest.raises(ValueError, match="already registered"):
+                WeightTransferTrainerFactory.register_engine(
+                    "dummy", _DummyTrainerEngine
+                )
+        finally:
+            WeightTransferTrainerFactory._registry = saved
+
+    def test_unknown_backend_raises(self):
+        with pytest.raises(ValueError, match="Invalid weight transfer backend"):
+            WeightTransferTrainerFactory.trainer_init(
+                "nope",
+                WeightTransferConfig(backend="nope"),
+                MagicMock(),
+                client=RecordingClient(),
+                source=ModuleSource(_module_with(("w", torch.zeros(2)))),
+            )
+
+
+class TestTrainerEngineBase:
+    """Base-class construction (no GPU)."""
+
+    def test_source_stored_and_sender_by_default(self):
+        engine = _DummyTrainerEngine(
+            WeightTransferConfig(backend="nccl"),
+            client=RecordingClient(),
+            source=ModuleSource(_module_with(("w", torch.zeros(2)))),
+        )
+        assert engine.is_sender is True
+        assert [name for name, _ in engine.source] == ["w"]
+
+    def test_shutdown_default_is_noop(self):
+        engine = _DummyTrainerEngine(
+            WeightTransferConfig(backend="nccl"),
+            client=RecordingClient(),
+            source=ModuleSource(_module_with(("w", torch.zeros(2)))),
+            is_sender=False,
+        )
+        assert engine.is_sender is False
+        engine.shutdown()  # must not raise

tools/pre_commit/check_forbidden_imports.py

@@ -39,6 +39,7 @@ CHECK_IMPORTS = {
             "vllm/distributed/device_communicators/shm_broadcast.py",
             "vllm/distributed/device_communicators/shm_object_storage.py",
             "vllm/distributed/weight_transfer/ipc_engine.py",
+            "vllm/distributed/weight_transfer/clients.py",
             "tests/distributed/test_weight_transfer.py",
             "vllm/utils/hashing.py",
             "tests/multimodal/media/test_base.py",

vllm/distributed/weight_transfer/__init__.py

@@ -5,10 +5,32 @@ Weight transfer engines for syncing model weights from trainers
 to inference workers.
 """
 
-from vllm.distributed.weight_transfer.base import WeightTransferEngine
-from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
+from vllm.distributed.weight_transfer.base import (
+    ModuleSource,
+    ParamMeta,
+    TrainerWeightTransferEngine,
+    VLLMWeightSyncClient,
+    WeightSource,
+    WeightTransferEngine,
+)
+from vllm.distributed.weight_transfer.clients import (
+    HTTPVLLMWeightSyncClient,
+    RayVLLMWeightSyncClient,
+)
+from vllm.distributed.weight_transfer.factory import (
+    WeightTransferEngineFactory,
+    WeightTransferTrainerFactory,
+)
 
 __all__ = [
     "WeightTransferEngine",
     "WeightTransferEngineFactory",
+    "TrainerWeightTransferEngine",
+    "WeightTransferTrainerFactory",
+    "VLLMWeightSyncClient",
+    "HTTPVLLMWeightSyncClient",
+    "RayVLLMWeightSyncClient",
+    "ParamMeta",
+    "WeightSource",
+    "ModuleSource",
 ]

vllm/distributed/weight_transfer/base.py

@@ -5,9 +5,10 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable
 
 import torch
+from typing_extensions import Self
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -17,6 +18,85 @@ from vllm.config.weight_transfer import WeightTransferConfig
 
 TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
 TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
+TConfig = TypeVar("TConfig", bound="WeightTransferConfig")
+
+# A trainer supplies its parameters as a `WeightSource` (defined below): a
+# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()`
+# channel. The built-in `ModuleSource` uses `materialize_full_tensor`.
+
+
+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. Trainer engines call
+    this at send time so the (potentially expensive) gather happens exactly
+    once — 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
+
+
+@dataclass(frozen=True)
+class ParamMeta:
+    """Name / wire dtype / full (HF) shape for one output parameter."""
+
+    name: str
+    dtype: torch.dtype
+    shape: tuple[int, ...]
+
+
+class WeightSource(ABC):
+    """A re-iterable source of the trainer's weights, handed to a trainer engine.
+
+    Two channels:
+
+    * `metadata()` — `(name, wire dtype, full shape)` for every parameter,
+      *without* transferring. Cheap when shapes are known locally (FSDP
+      `DTensor` global shape); may be expensive on first call for backends that
+      must materialize to learn shapes (e.g. a Megatron-Bridge export), in which
+      case it should cache.
+    * iteration — yields fully-materialized `(name, tensor)` pairs, one at a
+      time. Materializing is typically a collective (FSDP `full_tensor()`, a
+      Megatron export), so every trainer rank must iterate the same source in the
+      same order in lockstep, or ranks deadlock. Under pipeline parallelism a
+      rank may not own a parameter at all — iterating still drives the collective
+      and the yielded tensor is only meaningful on the sender.
+
+    `iter(source)` must yield a *fresh* pass each round. Backends with custom
+    producer logic (Megatron export, RDT plans, MoE re-fusing) subclass this.
+    """
+
+    @abstractmethod
+    def metadata(self) -> list[ParamMeta]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
+        raise NotImplementedError
+
+
+class ModuleSource(WeightSource):
+    """`WeightSource` over `module.named_parameters()` — the common case.
+
+    Handles both plain dense modules and FSDP-sharded ones with no special
+    casing: iteration all-gathers each `DTensor` via `full_tensor()` (a
+    collective) and passes regular tensors through. `metadata()` reads the
+    *global* `.shape` / `.dtype`, so it never triggers a gather.
+    """
+
+    def __init__(self, module: torch.nn.Module) -> None:
+        self._module = module
+
+    def metadata(self) -> list[ParamMeta]:
+        return [
+            ParamMeta(name, p.dtype, tuple(p.shape))
+            for name, p in self._module.named_parameters()
+        ]
+
+    def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
+        for name, param in self._module.named_parameters():
+            yield name, materialize_full_tensor(param)
 
 
 # Base protocols for backend-specific dataclasses
@@ -27,6 +107,26 @@ class WeightTransferInitInfo(ABC):  # noqa: B024
     pass
 
 
+@dataclass
+class TrainerInitInfo(WeightTransferInitInfo):
+    """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.
+    """
+
+    rank: int = field(kw_only=True)
+
+    @property
+    def is_sender(self) -> bool:
+        return self.rank == 0
+
+
 @dataclass
 class WeightTransferUpdateInfo(ABC):  # noqa: B024
     """Base class for backend-specific weight update info."""
@@ -243,3 +343,104 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
             >>> engine.trainer_send_weights(param_iter, trainer_args)
         """
         raise NotImplementedError
+
+
+@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: dict[str, Any]) -> None: ...
+
+    def finish_weight_update(self) -> None: ...
+
+
+class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
+    """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`. The `WeightSource` is required at `trainer_init`,
+    then replayed each round by the no-argument `send_weights()`.
+
+    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
+        config_cls: Type of backend-specific config
+    """
+
+    # Subclasses should override these class attributes
+    init_info_cls: type[TInitInfo]
+    config_cls: type[TConfig]
+
+    def __init__(
+        self,
+        config: TConfig,
+        *,
+        client: "VLLMWeightSyncClient",
+        source: "WeightSource",
+        is_sender: bool = True,
+    ) -> None:
+        self.config = config
+        self.is_sender = is_sender
+        # 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,
+        config: TConfig,
+        init_info: TInitInfo,
+        *,
+        client: "VLLMWeightSyncClient",
+        source: "WeightSource",
+    ) -> 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. `source` is stored on
+        `self.source`; after return, `send_weights()` is callable.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def send_weights(self) -> None:
+        """Push `self.source`'s 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."""

vllm/distributed/weight_transfer/clients.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Built-in `VLLMWeightSyncClient` implementations.
+
+These adapt the inference engine's weight-sync control plane to concrete
+transports. A `TrainerWeightTransferEngine` takes one of these (or any object
+with the same four methods — the protocol is structural) and drives the full
+handshake through it.
+
+Imports of `ray` / `requests` are deferred to call time so this module is
+importable without those packages installed.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
+
+if TYPE_CHECKING:
+    from ray.actor import ActorHandle
+
+
+def _json_safe_update_info(update_info: dict[str, Any]) -> dict[str, Any]:
+    """Make an update_info dict JSON-serializable for HTTP transport.
+
+    CUDA IPC handles (`ipc_handles`) are tuples of non-JSON-native objects, so
+    over HTTP they are pickled+base64-encoded into `ipc_handles_pickled` (which
+    the worker auto-deserializes when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`).
+    Other backends (NCCL) carry only JSON-native metadata and pass through
+    unchanged. Mirrors the old IPC `_do_send` HTTP branch.
+    """
+    ipc_handles = update_info.get("ipc_handles")
+    if ipc_handles is None:
+        return update_info
+
+    import pickle
+
+    import pybase64 as base64
+
+    out = {k: v for k, v in update_info.items() if k != "ipc_handles"}
+    out["ipc_handles_pickled"] = base64.b64encode(pickle.dumps(ipc_handles)).decode(
+        "utf-8"
+    )
+    return out
+
+
+class HTTPVLLMWeightSyncClient:
+    """Talks to a vLLM server over the RLHF HTTP routes.
+
+    Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`:
+    `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`,
+    `/finish_weight_update`.
+    """
+
+    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 start_weight_update(self) -> None:
+        self._post("start_weight_update")
+
+    def update_weights(self, update_info: dict[str, Any]) -> None:
+        self._post(
+            "update_weights", {"update_info": _json_safe_update_info(update_info)}
+        )
+
+    def finish_weight_update(self) -> None:
+        self._post("finish_weight_update")
+
+
+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 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: dict[str, Any]) -> 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) -> None:
+        import ray
+
+        ray.get([h.finish_weight_update.remote() for h in self.handles])

vllm/distributed/weight_transfer/factory.py

@@ -6,7 +6,10 @@ import importlib
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
-from vllm.distributed.weight_transfer.base import WeightTransferEngine
+from vllm.distributed.weight_transfer.base import (
+    TrainerWeightTransferEngine,
+    WeightTransferEngine,
+)
 from vllm.logger import init_logger
 
 if TYPE_CHECKING:
@@ -14,6 +17,11 @@ if TYPE_CHECKING:
 
     from vllm.config import VllmConfig
     from vllm.config.weight_transfer import WeightTransferConfig
+    from vllm.distributed.weight_transfer.base import (
+        VLLMWeightSyncClient,
+        WeightSource,
+        WeightTransferInitInfo,
+    )
 
 logger = init_logger(__name__)
 
@@ -111,6 +119,94 @@ class WeightTransferEngineFactory:
         return engine_cls(config, vllm_config, device, model)
 
 
+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,
+        backend: str,
+        config: "WeightTransferConfig",
+        init_info: "WeightTransferInitInfo",
+        *,
+        client: "VLLMWeightSyncClient",
+        source: "WeightSource",
+    ) -> 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``).
+
+        Args:
+            backend: Backend name (must be registered).
+            config: Backend-specific weight transfer config.
+            init_info: Backend-specific trainer init info.
+            client: Inference-side control-plane client.
+            source: `WeightSource` of `(name, tensor)` pairs to send each round.
+
+        Raises:
+            ValueError: If the backend is not registered.
+        """
+        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(
+            config=config,
+            init_info=init_info,
+            client=client,
+            source=source,
+        )
+
+
 # Register built-in weight transfer engines here.
 # Registration should be centralized to ensure lazy loading -
 # engine modules are only imported when actually used.

vllm/v1/engine/async_llm.py

@@ -1067,17 +1067,8 @@ class AsyncLLM(EngineClient):
         Args:
             request: Weight transfer initialization request with backend-specific info
         """
-        from vllm.distributed.weight_transfer.base import (
-            WeightTransferInitRequest,
-        )
-
-        if isinstance(request, WeightTransferInitRequest):
-            init_info_dict = request.init_info
-        else:
-            raise TypeError(f"Expected WeightTransferInitRequest, got {type(request)}")
-
         await self.collective_rpc(
-            "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
+            "init_weight_transfer_engine", kwargs={"init_info": request.init_info}
         )
 
     async def start_weight_update(self) -> None:
@@ -1095,16 +1086,8 @@ class AsyncLLM(EngineClient):
         Args:
             request: Weight update request with backend-specific update info
         """
-
-        if isinstance(request, WeightTransferUpdateRequest):
-            update_info_dict = request.update_info
-        else:
-            raise TypeError(
-                f"Expected WeightTransferUpdateRequest, got {type(request)}"
-            )
-
         await self.collective_rpc(
-            "update_weights", kwargs={"update_info": update_info_dict}
+            "update_weights", kwargs={"update_info": request.update_info}
         )
 
     async def finish_weight_update(self) -> None: