From fb1d8ccaf5bc46935031ef074950eaa8232cd15e Mon Sep 17 00:00:00 2001 From: Aaron Hao Date: Fri, 17 Jul 2026 00:11:06 -0700 Subject: [PATCH] [rl] Stateful Trainer Send: New Abstractions [1/N] (#48042) Signed-off-by: haoaaron Signed-off-by: Aaron Hao Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com> --- tests/distributed/test_weight_transfer.py | 205 ++++++++++++++++++- tools/pre_commit/check_forbidden_imports.py | 1 + vllm/distributed/weight_transfer/__init__.py | 26 ++- vllm/distributed/weight_transfer/base.py | 203 +++++++++++++++++- vllm/distributed/weight_transfer/clients.py | 114 +++++++++++ vllm/distributed/weight_transfer/factory.py | 98 ++++++++- vllm/v1/engine/async_llm.py | 21 +- 7 files changed, 644 insertions(+), 24 deletions(-) create mode 100644 vllm/distributed/weight_transfer/clients.py diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index f3423745ca5..b79aa1974d1 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/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 diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 365b2f5bb77..a2fc173f035 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/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", diff --git a/vllm/distributed/weight_transfer/__init__.py b/vllm/distributed/weight_transfer/__init__.py index af3322e0cbb..c78fd1f3cc5 100644 --- a/vllm/distributed/weight_transfer/__init__.py +++ b/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", ] diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 6dbd768d253..2e377e29253 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/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.""" diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py new file mode 100644 index 00000000000..4f54a6e291e --- /dev/null +++ b/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]) diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index a253363d736..4ea27c5ef58 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/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. diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 8bcd4ba89a4..93e02abf747 100644 --- a/vllm/v1/engine/async_llm.py +++ b/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: