iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[rl] Stateful Trainer Send: IPC [2/N] (#48981)

Browse Source 
Signed-off-by: haoaaron <ahao@anyscale.com>
This commit is contained in:
Aaron Hao
2026-07-30 21:59:08 +08:00
committed by GitHub
parent 59e831c09a
commit 30b4e7f479
8 changed files with 391 additions and 421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/rl/rlhf_http_ipc.py
+14 -37
View File
@@ -46,10 +46,12 @@ import torch
from openai import OpenAI
from transformers import AutoModelForCausalLM
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerSendWeightsArgs,
IPCWeightTransferEngine,
from vllm.distributed.weight_transfer import (
HTTPVLLMWeightSyncClient,
ModuleSource,
WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
BASE_URL = "http://localhost:8000"
MODEL_NAME = "facebook/opt-125m"
@@ -72,28 +74,6 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
return results
def init_weight_transfer_engine(base_url: str) -> None:
"""Initialize weight transfer via HTTP endpoint (no-op for IPC)."""
url = f"{base_url}/init_weight_transfer_engine"
payload = {"init_info": dict()}
response = requests.post(url, json=payload, timeout=60)
response.raise_for_status()
def start_weight_update(base_url: str) -> None:
"""Start a weight update via HTTP endpoint."""
url = f"{base_url}/start_weight_update"
response = requests.post(url, json={}, timeout=60)
response.raise_for_status()
def finish_weight_update(base_url: str) -> None:
"""Finish a weight update via HTTP endpoint."""
url = f"{base_url}/finish_weight_update"
response = requests.post(url, json={}, timeout=60)
response.raise_for_status()
def pause_generation(base_url: str) -> None:
"""Pause generation via HTTP endpoint."""
url = f"{base_url}/pause"
@@ -159,23 +139,20 @@ def main():
print("Initializing weight transfer (IPC backend)...")
# Initialize weight transfer on vLLM server (no-op for IPC, but still required)
init_weight_transfer_engine(BASE_URL)
# The trainer engine drives the inference side over HTTP. init for IPC is a
# no-op rendezvous; the same client carries start/update/finish.
engine = WeightTransferTrainerFactory.trainer_init(
init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
client=HTTPVLLMWeightSyncClient(BASE_URL),
source=ModuleSource(train_model),
)
# Pause generation before weight sync
pause_generation(BASE_URL)
# Start weight update, broadcast via IPC, then finish
start_weight_update(BASE_URL)
print("Broadcasting weights via CUDA IPC (HTTP)...")
trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
IPCWeightTransferEngine.trainer_send_weights(
iterator=train_model.named_parameters(),
trainer_args=trainer_args,
)
finish_weight_update(BASE_URL)
# One call drives start_weight_update / update_weights / finish_weight_update.
engine.send_weights()
# Resume generation after weight sync
resume_generation(BASE_URL)
examples/rl/rlhf_ipc.py
+18 -22
View File
@@ -29,10 +29,12 @@ from transformers import AutoModelForCausalLM
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerSendWeightsArgs,
IPCWeightTransferEngine,
from vllm.distributed.weight_transfer import (
ModuleSource,
RayVLLMWeightSyncClient,
WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
class MyLLM(LLM):
@@ -65,23 +67,19 @@ class TrainModel:
self.llm_handle = llm_handle
def init_weight_transfer(self):
# IPC backend doesn't need initialization info
ray.get(
self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict()))
"""Build the trainer-side IPC engine (no rendezvous needed for IPC)."""
self.engine = WeightTransferTrainerFactory.trainer_init(
init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
client=RayVLLMWeightSyncClient(self.llm_handle),
source=ModuleSource(self.train_model),
)
def broadcast_weights(
self, llm_handle: ray.actor.ActorHandle, packed: bool = False
):
"""Broadcast weights to the inference engine using IPC."""
self.llm_handle = llm_handle
trainer_args = IPCTrainerSendWeightsArgs(
send_mode="ray", llm_handle=llm_handle, packed=packed
)
IPCWeightTransferEngine.trainer_send_weights(
iterator=self.train_model.named_parameters(),
trainer_args=trainer_args,
)
def broadcast_weights(self):
"""Broadcast weights to the inference engine using IPC.
Drives start/update/finish on the inference side internally.
"""
self.engine.send_weights()
ray.init()
@@ -138,10 +136,8 @@ for output in outputs:
ray.get(llm.sleep.remote(level=0))
ray.get(train_model.init_weight_transfer.remote())
# Start weight update, sync weights, then finish
ray.get(llm.start_weight_update.remote())
ray.get(train_model.broadcast_weights.remote(llm))
ray.get(llm.finish_weight_update.remote())
# One call drives start_weight_update / update_weights / finish_weight_update.
ray.get(train_model.broadcast_weights.remote())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
examples/rl/rlhf_ipc_fsdp_ep.py
+49 -100
View File
@@ -14,8 +14,10 @@ Layout (4 GPUs, TP=1, DP=4, EP):
* A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
waits for initialization, and orchestrates generation / weight-sync.
Uses the built-in ``ray`` send_mode: each FSDP worker calls
``trainer_send_weights`` targeting its colocated LLM actor.
Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``)
and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only
rank 0 (the sender) ships the merged handles and drives the DP LLM actors through
its ``RayVLLMWeightSyncClient``.
This example was run on 4xH100.
"""
@@ -23,7 +25,6 @@ This example was run on 4xH100.
from __future__ import annotations
import os
from dataclasses import asdict
import ray
import torch
@@ -31,19 +32,26 @@ import torch.distributed as dist
from huggingface_hub import snapshot_download
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from torch.distributed._tensor import DTensor
from torch.distributed.fsdp import fully_shard
from transformers import AutoModelForCausalLM
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerSendWeightsArgs,
IPCWeightTransferEngine,
IPCWeightTransferInitInfo,
from vllm.distributed.weight_transfer import (
ModuleSource,
RayVLLMWeightSyncClient,
WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
from vllm.utils.network_utils import get_ip, get_open_port
# The inference side only needs the backend; the packed wire params are
# trainer-side and get propagated to the workers at the init handshake.
WEIGHT_TRANSFER_CONFIG = WeightTransferConfig(backend="ipc")
# Packed IPC transfer with a 1 GB buffer (matches the per-chunk buffer size).
PACKED = True
PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024
TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
@@ -125,72 +133,34 @@ class FSDPTrainWorker:
def get_weight_metadata(self):
return self.weight_names, self.weight_dtype_names, self.weight_shapes
def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True):
"""All-gather full params; all ranks create IPC handles, rank 0 sends.
def setup_engine(self, llm_handles):
"""Build the trainer IPC engine on every FSDP rank.
All ranks must call trainer_send_weights so they participate in the
all_gather_object collective inside _all_gather_and_merge_handles.
Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero).
Called on all ranks: rank 0 becomes the sender (drives the inference
side and performs the no-op IPC init handshake against all DP LLM
actors); the other ranks hold a null-client engine and only join the
IPC handle all-gather during send_weights.
"""
def _full_param_iter():
# HF's Qwen3MoeExperts (and other recent HF MoE impls) packs
# all experts into two fused 3-D tensors per layer:
# experts.gate_up_proj shape (E, 2*I, H)
# experts.down_proj shape (E, H, I)
# vLLM's Qwen3MoE load_weights still expects the older
# per-expert HF layout (experts.<i>.gate_proj.weight,
# experts.<i>.up_proj.weight, experts.<i>.down_proj.weight),
# so we un-fuse on the fly. Split order matches HF's forward:
# gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1)
# → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up.
params = self.model.state_dict()
for name in list(params.keys()):
param = params.pop(name)
if isinstance(param, DTensor):
tensor = param.full_tensor().detach().contiguous()
else:
tensor = param.detach().contiguous()
del param
if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3:
prefix = name[: -len(".gate_up_proj")]
num_experts, two_inter, _ = tensor.shape
inter = two_inter // 2
for i in range(num_experts):
expert = tensor[i]
yield (
f"{prefix}.{i}.gate_proj.weight",
expert[:inter].contiguous(),
)
yield (
f"{prefix}.{i}.up_proj.weight",
expert[inter:].contiguous(),
)
del tensor
elif name.endswith(".experts.down_proj") and tensor.dim() == 3:
prefix = name[: -len(".down_proj")]
num_experts = tensor.shape[0]
for i in range(num_experts):
yield (
f"{prefix}.{i}.down_proj.weight",
tensor[i].contiguous(),
)
del tensor
else:
yield name, tensor
trainer_args = IPCTrainerSendWeightsArgs(
send_mode="ray",
llm_handle=llm_handle,
packed=packed,
packed_buffer_size_bytes=1024 * 1024 * 1024, # 1 GB
)
IPCWeightTransferEngine.trainer_send_weights(
iterator=_full_param_iter(),
trainer_args=trainer_args,
self.engine = WeightTransferTrainerFactory.trainer_init(
init_info=IPCTrainerInitInfo(
rank=self.rank, # FSDP rank; sender is 0
packed=PACKED,
packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES,
),
client=RayVLLMWeightSyncClient(llm_handles),
source=ModuleSource(self.model),
)
def gather_and_broadcast_weights_ipc(self):
"""All-gather full params across FSDP ranks; rank 0 sends to vLLM.
Called on all ranks concurrently. `send_weights` gathers each param and
contributes to the IPC handle all-gather on every rank; only rank 0 (the
sender) drives start/update/finish on the inference side (the other
ranks' client RPCs no-op).
"""
self.engine.send_weights()
@ray.remote(num_cpus=1)
class DataParallelInferenceEngine:
@@ -224,7 +194,7 @@ class DataParallelInferenceEngine:
distributed_executor_backend="ray",
enable_expert_parallel=True,
gpu_memory_utilization=0.35,
weight_transfer_config=WeightTransferConfig(backend="ipc"),
weight_transfer_config=WEIGHT_TRANSFER_CONFIG,
enable_sleep_mode=True,
load_format="dummy",
dp_rank=r,
@@ -267,22 +237,6 @@ class DataParallelInferenceEngine:
rank_idx += 1
return ordered
def init_weight_transfer(self):
ray.get(
[
actor.init_weight_transfer_engine.remote(
dict(init_info=asdict(IPCWeightTransferInitInfo()))
)
for actor in self.llm_actors
]
)
def start_weight_update(self):
ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
def finish_weight_update(self):
ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
def sleep(self, level: int = 0):
ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
@@ -370,8 +324,10 @@ def main():
print("-" * 60)
# --- Weight transfer ---
print("[transfer] Initializing IPC weight transfer...")
ray.get(inference_engine.init_weight_transfer.remote())
# The rank-0 FSDP worker owns the trainer engine and drives the inference
# side (init/start/update/finish) for all DP actors via the Ray client.
print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...")
ray.get([w.setup_engine.remote(llm_actors) for w in fsdp_workers])
# Two-phase sleep/wake pattern:
# 1. sleep(level=1) — offload weights to CPU, discard KV cache
@@ -384,18 +340,11 @@ def main():
print("[sync] Waking weights (KV cache stays free)...")
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
print("[sync] Starting weight update...")
ray.get(inference_engine.start_weight_update.remote())
# All FSDP ranks participate in the IPC handle all-gather; rank 0's engine
# additionally drives start_weight_update / update_weights /
# finish_weight_update on the inference side.
print("[sync] Packed IPC transfer FSDP → vLLM...")
ray.get(
[
w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True)
for w in fsdp_workers
]
)
ray.get(inference_engine.finish_weight_update.remote())
ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers])
print("[sync] Weight transfer complete.")
print("[sync] Waking KV cache + scheduling...")
tests/distributed/test_weight_transfer.py
+68 -17
View File
@@ -27,10 +27,13 @@ from vllm.distributed.weight_transfer import (
WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.base import (
TrainerInitInfo,
WeightTransferInitRequest,
WeightTransferUpdateRequest,
)
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerInitInfo,
IPCTrainerWeightTransferEngine,
IPCWeightTransferEngine,
IPCWeightTransferInitInfo,
IPCWeightTransferUpdateInfo,
@@ -1071,7 +1074,7 @@ def inference_receive_ipc_tensor(
import torch
_set_ray_assigned_device()
device = _set_ray_assigned_device()
from vllm.config.parallel import ParallelConfig
from vllm.config.weight_transfer import WeightTransferConfig
@@ -1088,6 +1091,8 @@ def inference_receive_ipc_tensor(
for name, tensor in weights:
self.received.append((name, tensor.clone()))
# Trainer sends unpacked IPC handles; the worker learns packed=False from
# the init handshake below (IPCWeightTransferInitInfo defaults to False).
config = WeightTransferConfig(backend="ipc")
vllm_config = MagicMock()
parallel_config = MagicMock(spec=ParallelConfig)
@@ -1099,9 +1104,7 @@ def inference_receive_ipc_tensor(
vllm_config.model_config = MagicMock()
recorder = Recorder()
engine = IPCWeightTransferEngine(
config, vllm_config, _get_ray_assigned_device(), recorder
)
engine = IPCWeightTransferEngine(config, vllm_config, device, recorder)
# Transport-only test: bypass the set_current_vllm_config context that
# receive_weights enters, since vllm_config here is a mock.
import vllm.config as _vllm_config_mod
@@ -1211,6 +1214,7 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
torch.device("cuda:0"),
MagicMock(spec=torch.nn.Module),
)
# No init handshake here, so the engine keeps its default packed=False.
dummy_tensor = torch.ones(10, 10, device="cuda:0")
_, ipc_handle = reduce_tensor(dummy_tensor)
@@ -1264,8 +1268,8 @@ 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 trainer_init(cls, init_info, *, client, source):
return cls(client=client, source=source)
def send_weights(self):
pass
@@ -1369,18 +1373,17 @@ class TestModuleSource:
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_registry_has_ipc(self):
# IPC is the first backend migrated to the trainer engine; NCCL /
# sparse NCCL register in later PRs.
assert "ipc" in 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(),
MagicMock(backend="dummy"), # backend read from the init info
client=RecordingClient(),
source=ModuleSource(_module_with(("w", torch.zeros(2)))),
)
@@ -1395,20 +1398,26 @@ class TestTrainerFactory:
def test_unknown_backend_raises(self):
with pytest.raises(ValueError, match="Invalid weight transfer backend"):
WeightTransferTrainerFactory.trainer_init(
"nope",
WeightTransferConfig(backend="nope"),
MagicMock(),
MagicMock(backend="nope"),
client=RecordingClient(),
source=ModuleSource(_module_with(("w", torch.zeros(2)))),
)
def test_ipc_init_info_declares_backend(self):
assert IPCTrainerInitInfo.backend == "ipc"
def test_trainer_init_info_subclass_must_set_backend(self):
with pytest.raises(TypeError, match="class-level `backend`"):
class _NoBackend(TrainerInitInfo):
pass
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)))),
)
@@ -1417,10 +1426,52 @@ class TestTrainerEngineBase:
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
@pytest.mark.skipif(
torch.accelerator.device_count() < 1,
reason="Need at least 1 GPU (CUDA IPC handles).",
)
def test_ipc_trainer_send_weights_drives_client_in_order():
"""send_weights issues start -> update -> finish and ships per-round metadata;
the packed wire param rides the init info, not the per-round update_info."""
client = RecordingClient()
engine = IPCTrainerWeightTransferEngine(
client=client,
source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))),
packed=False,
)
engine.send_weights()
assert client.order == ["start", "update", "finish"]
assert client.last_update_info is not None
assert client.last_update_info["names"] == ["w"]
assert client.last_update_info["shapes"] == [[4]]
assert "packed" not in client.last_update_info
def test_ipc_trainer_init_ships_packed_to_worker():
"""trainer_init drives the inference-side init handshake and propagates the
must-agree `packed` flag to the worker."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU (CUDA IPC handles).")
client = RecordingClient()
engine = WeightTransferTrainerFactory.trainer_init(
init_info=IPCTrainerInitInfo(rank=0, packed=True), # backend from init info
client=client,
source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))),
)
assert isinstance(engine, IPCTrainerWeightTransferEngine)
assert engine.is_sender is True
assert engine.packed is True
assert client.order == ["init"]
assert client.last_init_info == {"packed": True}
vllm/distributed/weight_transfer/base.py
+37 -12
View File
@@ -5,7 +5,15 @@
from abc import ABC, abstractmethod
from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable
from typing import (
TYPE_CHECKING,
Any,
ClassVar,
Generic,
Protocol,
TypeVar,
runtime_checkable,
)
import torch
from typing_extensions import Self
@@ -18,7 +26,7 @@ from vllm.config.weight_transfer import WeightTransferConfig
TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
TConfig = TypeVar("TConfig", bound="WeightTransferConfig")
TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo")
# A trainer supplies its parameters as a `WeightSource` (defined below): a
# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()`
@@ -108,7 +116,7 @@ class WeightTransferInitInfo(ABC): # noqa: B024
@dataclass
class TrainerInitInfo(WeightTransferInitInfo):
class TrainerInitInfo:
"""Base trainer-side init info: which trainer rank drives the transfer.
`rank` is this trainer process's rank, provided **explicitly** by the
@@ -118,10 +126,25 @@ class TrainerInitInfo(WeightTransferInitInfo):
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
@@ -373,7 +396,7 @@ class VLLMWeightSyncClient(Protocol):
def finish_weight_update(self, weight_version: str | None = None) -> None: ...
class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
"""Trainer-side weight transfer engine.
Symmetric to `WeightTransferEngine` but lives in the training process.
@@ -382,6 +405,13 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
plans) on `self`. The `WeightSource` is required at `trainer_init`,
then replayed each round by the no-argument `send_weights()`.
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
@@ -392,22 +422,18 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
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]
# Subclasses should override this class attribute
init_info_cls: type[TTrainerInitInfo]
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.
@@ -418,8 +444,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
@abstractmethod
def trainer_init(
cls,
config: TConfig,
init_info: TInitInfo,
init_info: TTrainerInitInfo,
*,
client: "VLLMWeightSyncClient",
source: "WeightSource",
vllm/distributed/weight_transfer/factory.py
+21 -9
View File
@@ -18,9 +18,9 @@ if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.config.weight_transfer import WeightTransferConfig
from vllm.distributed.weight_transfer.base import (
TrainerInitInfo,
VLLMWeightSyncClient,
WeightSource,
WeightTransferInitInfo,
)
logger = init_logger(__name__)
@@ -164,9 +164,7 @@ class WeightTransferTrainerFactory:
@classmethod
def trainer_init(
cls,
backend: str,
config: "WeightTransferConfig",
init_info: "WeightTransferInitInfo",
init_info: "TrainerInitInfo",
*,
client: "VLLMWeightSyncClient",
source: "WeightSource",
@@ -176,16 +174,22 @@ class WeightTransferTrainerFactory:
Called on every trainer rank (multi-rank trainers construct on all
ranks; the sender is resolved inside the engine's ``trainer_init``).
The trainer side takes no `WeightTransferConfig` and no separate
`backend` argument: the backend is read from ``init_info.backend`` (a
`ClassVar` on each `TrainerInitInfo` subclass), and the static wire
params ride `init_info`.
Args:
backend: Backend name (must be registered).
config: Backend-specific weight transfer config.
init_info: Backend-specific trainer init info.
init_info: Backend-specific trainer init info. Its `backend`
selects the engine; it also carries the wire params (e.g.
`packed`).
client: Inference-side control-plane client.
source: `WeightSource` of `(name, tensor)` pairs to send each round.
Raises:
ValueError: If the backend is not registered.
ValueError: If `init_info.backend` is not registered.
"""
backend = init_info.backend
if backend not in cls._registry:
available = list(cls._registry.keys())
raise ValueError(
@@ -200,7 +204,6 @@ class WeightTransferTrainerFactory:
)
return engine_cls.trainer_init(
config=config,
init_info=init_info,
client=client,
source=source,
@@ -228,3 +231,12 @@ WeightTransferEngineFactory.register_engine(
"vllm.distributed.weight_transfer.sparse_nccl_engine",
"SparseNCCLWeightTransferEngine",
)
# Trainer-side engines. Backends register here as they migrate to the stateful
# trainer engine; NCCL / sparse NCCL keep their static trainer path until then.
WeightTransferTrainerFactory.register_engine(
"ipc",
"vllm.distributed.weight_transfer.ipc_engine",
"IPCTrainerWeightTransferEngine",
)
vllm/distributed/weight_transfer/ipc_engine.py
+183 -222
View File
@@ -3,19 +3,21 @@
"""IPC-based weight transfer engine using CUDA IPC for communication."""
import pickle
from collections.abc import Callable, Iterator
from dataclasses import asdict, dataclass
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, ClassVar
import pybase64 as base64
import ray
import requests
import torch
from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
from typing_extensions import Self
from vllm import envs
from vllm.config.weight_transfer import WeightTransferConfig
from vllm.distributed.weight_transfer.base import (
TrainerInitInfo,
TrainerWeightTransferEngine,
VLLMWeightSyncClient,
WeightSource,
WeightTransferEngine,
WeightTransferInitInfo,
WeightTransferUpdateInfo,
@@ -31,47 +33,35 @@ from vllm.distributed.weight_transfer.packed_tensor import (
@dataclass
class IPCTrainerSendWeightsArgs:
"""Arguments for IPC trainer_send_weights method."""
class IPCWeightTransferInitInfo(WeightTransferInitInfo):
"""Worker-side init info for IPC weight transfer. No rendezvous needed.
`packed` is a must-agree wire param: the trainer ships it here at the init
handshake so the worker decodes with the same setting the trainer encoded
with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does
not need the buffer size (producer-only)."""
send_mode: str | Callable[["IPCWeightTransferUpdateInfo"], None]
"""How to send updates to vLLM. Either a string ('ray' or 'http') for
built-in transports, or a callable that receives an
IPCWeightTransferUpdateInfo and performs the send."""
llm_handle: Any = None
"""Ray actor handle or list of handles (required for 'ray' send_mode)."""
url: str | None = None
"""Base URL for HTTP endpoint (required for 'http' send_mode)."""
packed: bool = False
"""Whether to use packed tensor transfer for bounded-memory chunking."""
packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
"""Size in bytes for each packed tensor buffer when packed=True."""
def __post_init__(self):
"""Validate that required arguments are provided for the selected mode."""
if callable(self.send_mode):
return
if self.send_mode == "ray" and self.llm_handle is None:
raise ValueError("llm_handle is required for 'ray' send_mode")
if self.send_mode == "http" and self.url is None:
raise ValueError("url is required for 'http' send_mode")
if self.send_mode not in ("ray", "http"):
raise ValueError(
f"send_mode must be 'ray', 'http', or a callable, "
f"got {self.send_mode!r}"
)
@dataclass
class IPCWeightTransferInitInfo(WeightTransferInitInfo):
"""Initialization info for IPC weight transfer backend. No init needed for IPC."""
class IPCTrainerInitInfo(TrainerInitInfo):
"""Trainer-side init info for IPC weight transfer. No rendezvous needed;
`rank` (from `TrainerInitInfo`) identifies this trainer process — rank 0
ships the merged IPC handles. All ranks still join the handle all-gather.
pass
`packed` / `packed_buffer_size_bytes` are the transfer's wire params. The
trainer propagates `packed` to the worker at `trainer_init` so the two sides
cannot disagree. `backend` is the factory dispatch key."""
backend: ClassVar[str] = "ipc"
packed: bool = False
packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
@dataclass
class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
"""Update info for IPC weight transfer backend."""
"""Per-round update info for the IPC weight transfer backend."""
names: list[str]
dtype_names: list[str]
@@ -85,8 +75,6 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
tensor_sizes: list[int] | None = None
"""Per-parameter sizes in bytes within the packed buffer.
Required when packed=True, unused otherwise."""
packed: bool = False
"""Whether this update uses packed tensor format."""
def __post_init__(self):
if self.ipc_handles_pickled is not None:
@@ -119,17 +107,14 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
f"`shapes` should be of the same size as `names`: "
f"got {len(self.shapes)} and {len(self.names)}"
)
if (
not self.packed
and isinstance(self.ipc_handles, list)
and len(self.ipc_handles) != num_params
):
# Unpacked transfers carry a per-parameter handle dict (a list); packed
# transfers carry a single dict for the whole buffer, so the list check
# only applies to the list case.
if isinstance(self.ipc_handles, list) and len(self.ipc_handles) != num_params:
raise ValueError(
f"`ipc_handles` should be of the same size as `names`: "
f"got {len(self.ipc_handles)} and {len(self.names)}"
)
if self.packed and self.tensor_sizes is None:
raise ValueError("`tensor_sizes` is required when packed=True")
class IPCWeightTransferEngine(
@@ -154,57 +139,21 @@ class IPCWeightTransferEngine(
device: torch.device,
model: torch.nn.Module,
) -> None:
"""
Initialize the IPC weight transfer engine.
Args:
config: The configuration for the weight transfer engine
vllm_config: The full vLLM config
device: The device this worker's model lives on
model: The local model instance which will receive the weights
"""
super().__init__(config, vllm_config, device, model)
def parse_update_info(
self, update_dict: dict[str, Any]
) -> IPCWeightTransferUpdateInfo:
"""Parse update dict, deserializing pickled IPC handles if present.
HTTP transport sends IPC handles as a base64-encoded pickle under the
key ``ipc_handles_pickled``. This method deserializes them back into
``ipc_handles`` before constructing the typed dataclass, keeping
serialization concerns out of the dataclass itself.
Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
payload is deserialized via ``pickle.loads``.
"""
pickled = update_dict.pop("ipc_handles_pickled", None)
if pickled is not None:
if update_dict.get("ipc_handles") is not None:
raise ValueError(
"Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
)
if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
raise ValueError(
"Refusing to deserialize `ipc_handles_pickled` without "
"VLLM_ALLOW_INSECURE_SERIALIZATION=1"
)
update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
return super().parse_update_info(update_dict)
# Set from the trainer-supplied init info at the handshake; defaults are
# only for the (unreachable) receive-before-init case.
self.packed = False
def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
"""
Initialize the weight transfer mechanism.
This is called once at the beginning of training.
No initialization needed for IPC backend.
Initialize the weight transfer mechanism. No data-plane rendezvous is
needed for IPC; this just records the trainer-supplied wire params so the
worker decodes exactly as the trainer encoded.
Args:
init_info: IPC initialization info (empty)
init_info: IPC initialization info (carries `packed`).
"""
pass
self.packed = init_info.packed
def start_weight_update(self) -> None:
"""Initialize layerwise reloading for the incoming checkpoint weights."""
@@ -226,6 +175,10 @@ class IPCWeightTransferEngine(
"""
Receive weights from the trainer via CUDA IPC handles and load them.
Whether the transfer is packed is read from `self.packed`, set at the
init handshake from the trainer's init info, so it is guaranteed to
match how the trainer encoded.
Args:
update_info: IPC update info containing parameter names, dtypes, shapes,
and IPC handles. Each IPC handle is a mapping between physical
@@ -238,8 +191,9 @@ class IPCWeightTransferEngine(
# rebuilt on the device the model lives on.
device_index = self.device.index
if update_info.packed:
assert update_info.tensor_sizes is not None
if self.packed:
if update_info.tensor_sizes is None:
raise ValueError("`tensor_sizes` is required when packed=True")
assert isinstance(update_info.ipc_handles, dict)
weights = packed_ipc_consumer(
ipc_handle=update_info.ipc_handles,
@@ -285,65 +239,108 @@ class IPCWeightTransferEngine(
pass
@staticmethod
def trainer_send_weights(
iterator: Iterator[tuple[str, torch.Tensor]],
trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs,
) -> None:
"""Send weights from trainer to inference workers via CUDA IPC.
def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
"""Removed. Use the stateful `IPCTrainerWeightTransferEngine` instead.
Supports two transport modes ('ray' and 'http') and two transfer
strategies:
- Non-packed (default): all weights in a single API call.
- Packed (packed=True): chunked transfer with bounded GPU memory.
For multi-GPU training, all ranks must call this method in
parallel. IPC handles are all-gathered across ranks and merged
so that each vLLM worker can find its own GPU UUID. Only rank 0
sends the payload to vLLM.
.. note::
This method calls ``update_weights`` internally. The caller must
call ``start_weight_update`` before and ``finish_weight_update``
after this method.
Args:
iterator: Iterator of (name, tensor) pairs. For multi-GPU,
each rank should yield the full tensor on its own GPU
(e.g. via FSDP full_tensor()).
trainer_args: IPCTrainerSendWeightsArgs or equivalent dict.
Transitional stub kept only to satisfy the (still abstract)
`WeightTransferEngine.trainer_send_weights`; that member is dropped from
the worker ABC once every backend has migrated to the trainer engine.
"""
args = (
IPCTrainerSendWeightsArgs(**trainer_args)
if isinstance(trainer_args, dict)
else trainer_args
raise NotImplementedError(
"The static IPC trainer path has been replaced by "
"IPCTrainerWeightTransferEngine. Build it via "
"WeightTransferTrainerFactory.trainer_init(IPCTrainerInitInfo(...), "
"client=..., source=...) and drive it with send_weights()."
)
device_index = torch.accelerator.current_device_index()
gpu_uuid = str(torch.cuda.get_device_properties(device_index).uuid)
if args.packed:
IPCWeightTransferEngine._send_packed(iterator, args, gpu_uuid)
else:
IPCWeightTransferEngine._send_unpacked(iterator, args, gpu_uuid)
@staticmethod
def _is_rank_zero() -> bool:
"""Return True if this is rank 0 or no distributed group exists."""
if not torch.distributed.is_initialized():
return True
return torch.distributed.get_rank() == 0
@staticmethod
class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]):
"""Trainer-side CUDA IPC weight transfer engine.
Called on every trainer rank. For multi-rank (e.g. FSDP) trainers all ranks
iterate the source (materializing each tensor) and contribute to the
IPC-handle all-gather; only the sender (rank 0) ships the merged handles to
the inference side. IPC transfer is straight-line (no concurrent broadcast
like NCCL): `update_weights` *is* the transfer, and it rides the client, so
it no-ops on non-senders.
`packed` / `packed_buffer_size_bytes` come from `IPCTrainerInitInfo`; the
sender propagates `packed` to the worker at `trainer_init`.
"""
init_info_cls = IPCTrainerInitInfo
def __init__(
self,
*,
client: VLLMWeightSyncClient,
source: WeightSource,
is_sender: bool = True,
packed: bool = False,
packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
) -> None:
super().__init__(client=client, source=source, is_sender=is_sender)
self.packed = packed
self.packed_buffer_size_bytes = packed_buffer_size_bytes
self.device_index = torch.accelerator.current_device_index()
self.gpu_uuid = str(torch.cuda.get_device_properties(self.device_index).uuid)
@classmethod
def trainer_init(
cls,
init_info: IPCTrainerInitInfo,
*,
client: VLLMWeightSyncClient,
source: WeightSource,
) -> Self:
engine = cls(
client=client,
source=source,
is_sender=init_info.is_sender,
packed=init_info.packed,
packed_buffer_size_bytes=init_info.packed_buffer_size_bytes,
)
# IPC needs no data-plane rendezvous. The sender ships the must-agree
# `packed` flag so the worker decodes exactly as this trainer encodes.
if engine.is_sender:
engine.client.init_weight_transfer_engine({"packed": init_info.packed})
return engine
def send_weights(self) -> None:
source = self.source
if self.is_sender:
self.client.start_weight_update()
# Unpacked returns strong refs to its IPC-shared copies; they must stay
# alive across `finish_weight_update` too.
weight_refs = self._send(source)
if self.is_sender:
self.client.finish_weight_update()
self._post_send_sync()
del weight_refs
# ---- data plane (runs on all ranks; only the sender ships) ----
def _send(self, source: WeightSource) -> list[torch.Tensor] | None:
if self.packed:
self._send_packed(source)
return None
return self._send_unpacked(source)
def _all_gather_and_merge_handles(
self,
handles: list[dict[str, tuple]],
) -> list[dict[str, tuple]]:
"""All-gather and merge IPC handle dicts across ranks in one call.
Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one
per parameter or one per chunk). A single all_gather_object
collects every rank's full list, then rank 0 merges per-index so
collects every rank's full list, then the sender merges per-index so
each dict maps every GPU UUID to its args.
Non-rank-0 returns a list of empty dicts.
No-op (returns handles unchanged) when no distributed group exists.
The all-gather runs over the *default* process group; this assumes the
default group is exactly the set of colocated trainer ranks and that the
sender is a member. No-op (returns handles unchanged) when no distributed
group exists.
"""
if (
not torch.distributed.is_initialized()
@@ -357,7 +354,7 @@ class IPCWeightTransferEngine(
torch.distributed.barrier()
torch.cuda.synchronize()
if torch.distributed.get_rank() == 0:
if self.is_sender:
merged: list[dict[str, tuple]] = []
for param_idx in range(len(handles)):
m: dict[str, tuple] = {}
@@ -378,25 +375,23 @@ class IPCWeightTransferEngine(
torch.distributed.barrier()
torch.cuda.ipc_collect()
@staticmethod
def _send_unpacked(
iterator: Iterator[tuple[str, torch.Tensor]],
args: IPCTrainerSendWeightsArgs,
gpu_uuid: str,
) -> None:
"""Send all weights in a single API call (non-packed mode)."""
def _send_unpacked(self, source: WeightSource) -> list[torch.Tensor]:
"""Iterate the source, build one IPC handle per param, all-gather the
handles across ranks, and (sender) ship them in one update call.
Returns the strong refs to every contiguous copy. reduce_tensor's args
do NOT keep storage alive, and non-contiguous inputs allocate fresh
storage in .contiguous(); the caller must keep these alive until the
post-send barrier (past `finish`) so the consumer's IPC views stay valid.
"""
names: list[str] = []
dtype_names: list[str] = []
shapes: list[list[int]] = []
ipc_handles: list[dict[str, tuple]] = []
# Hold strong refs to every contiguous copy until the send + post-send
# sync completes. reduce_tensor's returned args do NOT keep storage
# alive, and non-contiguous inputs allocate fresh storage in
# .contiguous() that would otherwise be GC'd before the consumer opens
# the IPC handle.
weight_refs: list[torch.Tensor] = []
for name, tensor in iterator:
for name, tensor in source:
# FSDP shards were gathered by the source; ensure contiguity here.
names.append(name)
dtype_names.append(str(tensor.dtype).split(".")[-1])
shapes.append(list(tensor.shape))
@@ -404,98 +399,64 @@ class IPCWeightTransferEngine(
weight = tensor.detach().contiguous()
weight_refs.append(weight)
_, ipc_args = reduce_tensor(weight)
ipc_handles.append({gpu_uuid: ipc_args})
ipc_handles.append({self.gpu_uuid: ipc_args})
ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
ipc_handles = self._all_gather_and_merge_handles(ipc_handles)
self._do_send(
names=names,
dtype_names=dtype_names,
shapes=shapes,
ipc_handles=ipc_handles,
)
return weight_refs
if IPCWeightTransferEngine._is_rank_zero():
IPCWeightTransferEngine._do_send(
args=args,
names=names,
dtype_names=dtype_names,
shapes=shapes,
ipc_handles=ipc_handles,
)
IPCWeightTransferEngine._post_send_sync()
@staticmethod
def _send_packed(
iterator: Iterator[tuple[str, torch.Tensor]],
args: IPCTrainerSendWeightsArgs,
gpu_uuid: str,
) -> None:
def _send_packed(self, source: WeightSource) -> None:
"""Send weights in bounded-memory chunks (packed mode)."""
post_iter_func: Callable = lambda item: item[1]
for chunk in packed_ipc_producer(
iterator=iterator,
gpu_uuid=gpu_uuid,
post_iter_func=post_iter_func,
buffer_size_bytes=args.packed_buffer_size_bytes,
iterator=iter(source),
gpu_uuid=self.gpu_uuid,
post_iter_func=lambda item: item[1],
buffer_size_bytes=self.packed_buffer_size_bytes,
):
ipc_handle = IPCWeightTransferEngine._all_gather_and_merge_handles(
[chunk.ipc_handle]
)[0]
ipc_handle = self._all_gather_and_merge_handles([chunk.ipc_handle])[0]
self._do_send(
names=chunk.names,
dtype_names=chunk.dtype_names,
shapes=chunk.shapes,
ipc_handles=ipc_handle,
tensor_sizes=chunk.tensor_sizes,
)
# Per-chunk barrier: the producer reuses a single IPC buffer across
# chunks, but only the sender waits for the consumers (via _do_send).
# Without syncing every rank here, non-sender ranks race ahead and
# overwrite their buffer while their colocated worker is still
# reading the current chunk, silently corrupting the transfer.
self._post_send_sync()
if IPCWeightTransferEngine._is_rank_zero():
IPCWeightTransferEngine._do_send(
args=args,
names=chunk.names,
dtype_names=chunk.dtype_names,
shapes=chunk.shapes,
ipc_handles=ipc_handle,
tensor_sizes=chunk.tensor_sizes,
packed=True,
)
IPCWeightTransferEngine._post_send_sync()
@staticmethod
def _do_send(
args: IPCTrainerSendWeightsArgs,
self,
names: list[str],
dtype_names: list[str],
shapes: list[list[int]],
ipc_handles: list[dict[str, tuple]] | dict[str, tuple],
tensor_sizes: list[int] | None = None,
packed: bool = False,
) -> None:
"""Send a single update payload via the configured transport."""
"""Build one update payload and ship it via the client. Only the sender
ships (non-sender ranks already contributed to the handle all-gather).
Emits raw `ipc_handles`; transports that cannot carry them natively
(HTTP/JSON) pickle them in their client (see `HTTPVLLMWeightSyncClient`).
"""
if not self.is_sender:
return
update_fields: dict[str, Any] = {
"names": names,
"dtype_names": dtype_names,
"shapes": shapes,
"packed": packed,
"ipc_handles": ipc_handles,
}
if tensor_sizes is not None:
update_fields["tensor_sizes"] = tensor_sizes
update_fields["ipc_handles"] = ipc_handles
update_info = IPCWeightTransferUpdateInfo(**update_fields)
if callable(args.send_mode):
args.send_mode(update_info)
elif args.send_mode == "ray":
handles = (
args.llm_handle
if isinstance(args.llm_handle, list)
else [args.llm_handle]
)
ray.get(
[
h.update_weights.remote(dict(update_info=asdict(update_info)))
for h in handles
]
)
elif args.send_mode == "http":
pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode(
"utf-8"
)
http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"}
http_fields["ipc_handles_pickled"] = pickled_handles
url = f"{args.url}/update_weights"
payload = {"update_info": http_fields}
response = requests.post(url, json=payload, timeout=300)
response.raise_for_status()
self.client.update_weights(asdict(update_info))
vllm/distributed/weight_transfer/packed_tensor.py
+1 -2
View File
@@ -10,8 +10,7 @@ from typing import Any
import torch
from torch.multiprocessing.reductions import reduce_tensor
# Default values for packed tensor configuration.
# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights.
# Default values for packed tensor transfer.
DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB
DEFAULT_PACKED_NUM_BUFFERS = 2