[rl] Stateful Trainer Send: IPC [2/N] (#48981)
Signed-off-by: haoaaron <ahao@anyscale.com>
Aaron Hao committed
Jul 30, 2026 at 06:59 UTC
30b4e7f479674a9c4d8889d4857294d3bd5e6849
8 files changed
+392
-422
examples/rl/rlhf_http_ipc.py
+14
-37
@@ -46,10 +46,12 @@ import torch
46
from openai import OpenAI
47
from transformers import AutoModelForCausalLM
48
49
-from vllm.distributed.weight_transfer.ipc_engine import (
50
- IPCTrainerSendWeightsArgs,
51
- IPCWeightTransferEngine,
49
+from vllm.distributed.weight_transfer import (
50
+ HTTPVLLMWeightSyncClient,
51
+ ModuleSource,
52
+ WeightTransferTrainerFactory,
53
)
54
+from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
55
56
BASE_URL = "http://localhost:8000"
57
MODEL_NAME = "facebook/opt-125m"
@@ -72,28 +74,6 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
74
return results
75
76
75
-def init_weight_transfer_engine(base_url: str) -> None:
76
- """Initialize weight transfer via HTTP endpoint (no-op for IPC)."""
77
- url = f"{base_url}/init_weight_transfer_engine"
78
- payload = {"init_info": dict()}
79
- response = requests.post(url, json=payload, timeout=60)
80
- response.raise_for_status()
81
-
82
-
83
-def start_weight_update(base_url: str) -> None:
84
- """Start a weight update via HTTP endpoint."""
85
- url = f"{base_url}/start_weight_update"
86
- response = requests.post(url, json={}, timeout=60)
87
- response.raise_for_status()
88
-
89
-
90
-def finish_weight_update(base_url: str) -> None:
91
- """Finish a weight update via HTTP endpoint."""
92
- url = f"{base_url}/finish_weight_update"
93
- response = requests.post(url, json={}, timeout=60)
94
- response.raise_for_status()
95
-
96
-
77
def pause_generation(base_url: str) -> None:
78
"""Pause generation via HTTP endpoint."""
79
url = f"{base_url}/pause"
@@ -159,23 +139,20 @@ def main():
139
140
print("Initializing weight transfer (IPC backend)...")
141
162
- # Initialize weight transfer on vLLM server (no-op for IPC, but still required)
163
- init_weight_transfer_engine(BASE_URL)
142
+ # The trainer engine drives the inference side over HTTP. init for IPC is a
143
+ # no-op rendezvous; the same client carries start/update/finish.
144
+ engine = WeightTransferTrainerFactory.trainer_init(
145
+ init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
146
+ client=HTTPVLLMWeightSyncClient(BASE_URL),
147
+ source=ModuleSource(train_model),
148
+ )
149
150
# Pause generation before weight sync
151
pause_generation(BASE_URL)
152
168
- # Start weight update, broadcast via IPC, then finish
169
- start_weight_update(BASE_URL)
170
-
153
print("Broadcasting weights via CUDA IPC (HTTP)...")
172
- trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
173
- IPCWeightTransferEngine.trainer_send_weights(
174
- iterator=train_model.named_parameters(),
175
- trainer_args=trainer_args,
176
- )
177
-
178
- finish_weight_update(BASE_URL)
154
+ # One call drives start_weight_update / update_weights / finish_weight_update.
155
+ engine.send_weights()
156
157
# Resume generation after weight sync
158
resume_generation(BASE_URL)
examples/rl/rlhf_ipc.py
+18
-22
@@ -29,10 +29,12 @@ from transformers import AutoModelForCausalLM
29
30
from vllm import LLM, SamplingParams
31
from vllm.config import WeightTransferConfig
32
-from vllm.distributed.weight_transfer.ipc_engine import (
33
- IPCTrainerSendWeightsArgs,
34
- IPCWeightTransferEngine,
32
+from vllm.distributed.weight_transfer import (
33
+ ModuleSource,
34
+ RayVLLMWeightSyncClient,
35
+ WeightTransferTrainerFactory,
36
)
37
+from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
38
39
40
class MyLLM(LLM):
@@ -65,23 +67,19 @@ class TrainModel:
67
self.llm_handle = llm_handle
68
69
def init_weight_transfer(self):
68
- # IPC backend doesn't need initialization info
69
- ray.get(
70
- self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict()))
70
+ """Build the trainer-side IPC engine (no rendezvous needed for IPC)."""
71
+ self.engine = WeightTransferTrainerFactory.trainer_init(
72
+ init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
73
+ client=RayVLLMWeightSyncClient(self.llm_handle),
74
+ source=ModuleSource(self.train_model),
75
)
76
73
- def broadcast_weights(
74
- self, llm_handle: ray.actor.ActorHandle, packed: bool = False
75
- ):
76
- """Broadcast weights to the inference engine using IPC."""
77
- self.llm_handle = llm_handle
78
- trainer_args = IPCTrainerSendWeightsArgs(
79
- send_mode="ray", llm_handle=llm_handle, packed=packed
80
- )
81
- IPCWeightTransferEngine.trainer_send_weights(
82
- iterator=self.train_model.named_parameters(),
83
- trainer_args=trainer_args,
84
- )
77
+ def broadcast_weights(self):
78
+ """Broadcast weights to the inference engine using IPC.
79
+
80
+ Drives start/update/finish on the inference side internally.
81
+ """
82
+ self.engine.send_weights()
83
84
85
ray.init()
@@ -138,10 +136,8 @@ for output in outputs:
136
ray.get(llm.sleep.remote(level=0))
137
138
ray.get(train_model.init_weight_transfer.remote())
141
-# Start weight update, sync weights, then finish
142
-ray.get(llm.start_weight_update.remote())
143
-ray.get(train_model.broadcast_weights.remote(llm))
144
-ray.get(llm.finish_weight_update.remote())
139
+# One call drives start_weight_update / update_weights / finish_weight_update.
140
+ray.get(train_model.broadcast_weights.remote())
141
142
ray.get(llm.wake_up.remote(tags=["scheduling"]))
143
examples/rl/rlhf_ipc_fsdp_ep.py
+49
-100
@@ -14,8 +14,10 @@ Layout (4 GPUs, TP=1, DP=4, EP):
14
* A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
15
waits for initialization, and orchestrates generation / weight-sync.
16
17
-Uses the built-in ``ray`` send_mode: each FSDP worker calls
18
-``trainer_send_weights`` targeting its colocated LLM actor.
17
+Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``)
18
+and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only
19
+rank 0 (the sender) ships the merged handles and drives the DP LLM actors through
20
+its ``RayVLLMWeightSyncClient``.
21
22
This example was run on 4xH100.
23
"""
@@ -23,7 +25,6 @@ This example was run on 4xH100.
25
from __future__ import annotations
26
27
import os
26
-from dataclasses import asdict
28
29
import ray
30
import torch
@@ -31,19 +32,26 @@ import torch.distributed as dist
32
from huggingface_hub import snapshot_download
33
from ray.util.placement_group import placement_group
34
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
34
-from torch.distributed._tensor import DTensor
35
from torch.distributed.fsdp import fully_shard
36
from transformers import AutoModelForCausalLM
37
38
from vllm import LLM, SamplingParams
39
from vllm.config import WeightTransferConfig
40
-from vllm.distributed.weight_transfer.ipc_engine import (
41
- IPCTrainerSendWeightsArgs,
42
- IPCWeightTransferEngine,
43
- IPCWeightTransferInitInfo,
40
+from vllm.distributed.weight_transfer import (
41
+ ModuleSource,
42
+ RayVLLMWeightSyncClient,
43
+ WeightTransferTrainerFactory,
44
)
45
+from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
46
from vllm.utils.network_utils import get_ip, get_open_port
47
48
+# The inference side only needs the backend; the packed wire params are
49
+# trainer-side and get propagated to the workers at the init handshake.
50
+WEIGHT_TRANSFER_CONFIG = WeightTransferConfig(backend="ipc")
51
+# Packed IPC transfer with a 1 GB buffer (matches the per-chunk buffer size).
52
+PACKED = True
53
+PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024
54
+
55
TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
56
VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
57
@@ -125,72 +133,34 @@ class FSDPTrainWorker:
133
def get_weight_metadata(self):
134
return self.weight_names, self.weight_dtype_names, self.weight_shapes
135
128
- def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True):
129
- """All-gather full params; all ranks create IPC handles, rank 0 sends.
136
+ def setup_engine(self, llm_handles):
137
+ """Build the trainer IPC engine on every FSDP rank.
138
131
- All ranks must call trainer_send_weights so they participate in the
132
- all_gather_object collective inside _all_gather_and_merge_handles.
133
- Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero).
139
+ Called on all ranks: rank 0 becomes the sender (drives the inference
140
+ side and performs the no-op IPC init handshake against all DP LLM
141
+ actors); the other ranks hold a null-client engine and only join the
142
+ IPC handle all-gather during send_weights.
143
"""
135
-
136
- def _full_param_iter():
137
- # HF's Qwen3MoeExperts (and other recent HF MoE impls) packs
138
- # all experts into two fused 3-D tensors per layer:
139
- # experts.gate_up_proj shape (E, 2*I, H)
140
- # experts.down_proj shape (E, H, I)
141
- # vLLM's Qwen3MoE load_weights still expects the older
142
- # per-expert HF layout (experts.<i>.gate_proj.weight,
143
- # experts.<i>.up_proj.weight, experts.<i>.down_proj.weight),
144
- # so we un-fuse on the fly. Split order matches HF's forward:
145
- # gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1)
146
- # → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up.
147
- params = self.model.state_dict()
148
- for name in list(params.keys()):
149
- param = params.pop(name)
150
- if isinstance(param, DTensor):
151
- tensor = param.full_tensor().detach().contiguous()
152
- else:
153
- tensor = param.detach().contiguous()
154
- del param
155
-
156
- if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3:
157
- prefix = name[: -len(".gate_up_proj")]
158
- num_experts, two_inter, _ = tensor.shape
159
- inter = two_inter // 2
160
- for i in range(num_experts):
161
- expert = tensor[i]
162
- yield (
163
- f"{prefix}.{i}.gate_proj.weight",
164
- expert[:inter].contiguous(),
165
- )
166
- yield (
167
- f"{prefix}.{i}.up_proj.weight",
168
- expert[inter:].contiguous(),
169
- )
170
- del tensor
171
- elif name.endswith(".experts.down_proj") and tensor.dim() == 3:
172
- prefix = name[: -len(".down_proj")]
173
- num_experts = tensor.shape[0]
174
- for i in range(num_experts):
175
- yield (
176
- f"{prefix}.{i}.down_proj.weight",
177
- tensor[i].contiguous(),
178
- )
179
- del tensor
180
- else:
181
- yield name, tensor
182
-
183
- trainer_args = IPCTrainerSendWeightsArgs(
184
- send_mode="ray",
185
- llm_handle=llm_handle,
186
- packed=packed,
187
- packed_buffer_size_bytes=1024 * 1024 * 1024, # 1 GB
188
- )
189
- IPCWeightTransferEngine.trainer_send_weights(
190
- iterator=_full_param_iter(),
191
- trainer_args=trainer_args,
144
+ self.engine = WeightTransferTrainerFactory.trainer_init(
145
+ init_info=IPCTrainerInitInfo(
146
+ rank=self.rank, # FSDP rank; sender is 0
147
+ packed=PACKED,
148
+ packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES,
149
+ ),
150
+ client=RayVLLMWeightSyncClient(llm_handles),
151
+ source=ModuleSource(self.model),
152
)
153
154
+ def gather_and_broadcast_weights_ipc(self):
155
+ """All-gather full params across FSDP ranks; rank 0 sends to vLLM.
156
+
157
+ Called on all ranks concurrently. `send_weights` gathers each param and
158
+ contributes to the IPC handle all-gather on every rank; only rank 0 (the
159
+ sender) drives start/update/finish on the inference side (the other
160
+ ranks' client RPCs no-op).
161
+ """
162
+ self.engine.send_weights()
163
+
164
165
@ray.remote(num_cpus=1)
166
class DataParallelInferenceEngine:
@@ -224,7 +194,7 @@ class DataParallelInferenceEngine:
194
distributed_executor_backend="ray",
195
enable_expert_parallel=True,
196
gpu_memory_utilization=0.35,
227
- weight_transfer_config=WeightTransferConfig(backend="ipc"),
197
+ weight_transfer_config=WEIGHT_TRANSFER_CONFIG,
198
enable_sleep_mode=True,
199
load_format="dummy",
200
dp_rank=r,
@@ -267,22 +237,6 @@ class DataParallelInferenceEngine:
237
rank_idx += 1
238
return ordered
239
270
- def init_weight_transfer(self):
271
- ray.get(
272
- [
273
- actor.init_weight_transfer_engine.remote(
274
- dict(init_info=asdict(IPCWeightTransferInitInfo()))
275
- )
276
- for actor in self.llm_actors
277
- ]
278
- )
279
-
280
- def start_weight_update(self):
281
- ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
282
-
283
- def finish_weight_update(self):
284
- ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
285
-
240
def sleep(self, level: int = 0):
241
ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
242
@@ -370,8 +324,10 @@ def main():
324
print("-" * 60)
325
326
# --- Weight transfer ---
373
- print("[transfer] Initializing IPC weight transfer...")
374
- ray.get(inference_engine.init_weight_transfer.remote())
327
+ # The rank-0 FSDP worker owns the trainer engine and drives the inference
328
+ # side (init/start/update/finish) for all DP actors via the Ray client.
329
+ print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...")
330
+ ray.get([w.setup_engine.remote(llm_actors) for w in fsdp_workers])
331
332
# Two-phase sleep/wake pattern:
333
# 1. sleep(level=1) — offload weights to CPU, discard KV cache
@@ -384,18 +340,11 @@ def main():
340
print("[sync] Waking weights (KV cache stays free)...")
341
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
342
387
- print("[sync] Starting weight update...")
388
- ray.get(inference_engine.start_weight_update.remote())
389
-
343
+ # All FSDP ranks participate in the IPC handle all-gather; rank 0's engine
344
+ # additionally drives start_weight_update / update_weights /
345
+ # finish_weight_update on the inference side.
346
print("[sync] Packed IPC transfer FSDP → vLLM...")
391
- ray.get(
392
- [
393
- w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True)
394
- for w in fsdp_workers
395
- ]
396
- )
397
-
398
- ray.get(inference_engine.finish_weight_update.remote())
347
+ ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers])
348
print("[sync] Weight transfer complete.")
349
350
print("[sync] Waking KV cache + scheduling...")
tests/distributed/test_weight_transfer.py
+68
-17
@@ -27,10 +27,13 @@ from vllm.distributed.weight_transfer import (
27
WeightTransferTrainerFactory,
28
)
29
from vllm.distributed.weight_transfer.base import (
30
+ TrainerInitInfo,
31
WeightTransferInitRequest,
32
WeightTransferUpdateRequest,
33
)
34
from vllm.distributed.weight_transfer.ipc_engine import (
35
+ IPCTrainerInitInfo,
36
+ IPCTrainerWeightTransferEngine,
37
IPCWeightTransferEngine,
38
IPCWeightTransferInitInfo,
39
IPCWeightTransferUpdateInfo,
@@ -1071,7 +1074,7 @@ def inference_receive_ipc_tensor(
1074
1075
import torch
1076
1074
- _set_ray_assigned_device()
1077
+ device = _set_ray_assigned_device()
1078
1079
from vllm.config.parallel import ParallelConfig
1080
from vllm.config.weight_transfer import WeightTransferConfig
@@ -1088,6 +1091,8 @@ def inference_receive_ipc_tensor(
1091
for name, tensor in weights:
1092
self.received.append((name, tensor.clone()))
1093
1094
+ # Trainer sends unpacked IPC handles; the worker learns packed=False from
1095
+ # the init handshake below (IPCWeightTransferInitInfo defaults to False).
1096
config = WeightTransferConfig(backend="ipc")
1097
vllm_config = MagicMock()
1098
parallel_config = MagicMock(spec=ParallelConfig)
@@ -1099,9 +1104,7 @@ def inference_receive_ipc_tensor(
1104
vllm_config.model_config = MagicMock()
1105
1106
recorder = Recorder()
1102
- engine = IPCWeightTransferEngine(
1103
- config, vllm_config, _get_ray_assigned_device(), recorder
1104
- )
1107
+ engine = IPCWeightTransferEngine(config, vllm_config, device, recorder)
1108
# Transport-only test: bypass the set_current_vllm_config context that
1109
# receive_weights enters, since vllm_config here is a mock.
1110
import vllm.config as _vllm_config_mod
@@ -1211,6 +1214,7 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
1214
torch.device("cuda:0"),
1215
MagicMock(spec=torch.nn.Module),
1216
)
1217
+ # No init handshake here, so the engine keeps its default packed=False.
1218
1219
dummy_tensor = torch.ones(10, 10, device="cuda:0")
1220
_, ipc_handle = reduce_tensor(dummy_tensor)
@@ -1264,8 +1268,8 @@ class _DummyTrainerEngine(TrainerWeightTransferEngine):
1268
"""Minimal concrete trainer engine to exercise base-class + factory."""
1269
1270
@classmethod
1267
- def trainer_init(cls, config, init_info, *, client, source):
1268
- return cls(config, client=client, source=source)
1271
+ def trainer_init(cls, init_info, *, client, source):
1272
+ return cls(client=client, source=source)
1273
1274
def send_weights(self):
1275
pass
@@ -1369,18 +1373,17 @@ class TestModuleSource:
1373
class TestTrainerFactory:
1374
"""WeightTransferTrainerFactory registry mechanics."""
1375
1372
- def test_builtin_registry_has_no_trainer_backends_yet(self):
1373
- # Concrete backends register in the per-backend migration PRs.
1374
- assert WeightTransferTrainerFactory._registry == {}
1376
+ def test_registry_has_ipc(self):
1377
+ # IPC is the first backend migrated to the trainer engine; NCCL /
1378
+ # sparse NCCL register in later PRs.
1379
+ assert "ipc" in WeightTransferTrainerFactory._registry
1380
1381
def test_register_and_dispatch(self):
1382
saved = dict(WeightTransferTrainerFactory._registry)
1383
try:
1384
WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine)
1385
engine = WeightTransferTrainerFactory.trainer_init(
1381
- "dummy",
1382
- WeightTransferConfig(backend="dummy"),
1383
- MagicMock(),
1386
+ MagicMock(backend="dummy"), # backend read from the init info
1387
client=RecordingClient(),
1388
source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1389
)
@@ -1395,20 +1398,26 @@ class TestTrainerFactory:
1398
def test_unknown_backend_raises(self):
1399
with pytest.raises(ValueError, match="Invalid weight transfer backend"):
1400
WeightTransferTrainerFactory.trainer_init(
1398
- "nope",
1399
- WeightTransferConfig(backend="nope"),
1400
- MagicMock(),
1401
+ MagicMock(backend="nope"),
1402
client=RecordingClient(),
1403
source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1404
)
1405
1406
+ def test_ipc_init_info_declares_backend(self):
1407
+ assert IPCTrainerInitInfo.backend == "ipc"
1408
+
1409
+ def test_trainer_init_info_subclass_must_set_backend(self):
1410
+ with pytest.raises(TypeError, match="class-level `backend`"):
1411
+
1412
+ class _NoBackend(TrainerInitInfo):
1413
+ pass
1414
+
1415
1416
class TestTrainerEngineBase:
1417
"""Base-class construction (no GPU)."""
1418
1419
def test_source_stored_and_sender_by_default(self):
1420
engine = _DummyTrainerEngine(
1411
- WeightTransferConfig(backend="nccl"),
1421
client=RecordingClient(),
1422
source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1423
)
@@ -1417,10 +1426,52 @@ class TestTrainerEngineBase:
1426
1427
def test_shutdown_default_is_noop(self):
1428
engine = _DummyTrainerEngine(
1420
- WeightTransferConfig(backend="nccl"),
1429
client=RecordingClient(),
1430
source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1431
is_sender=False,
1432
)
1433
assert engine.is_sender is False
1434
engine.shutdown() # must not raise
1435
+
1436
+
1437
+@pytest.mark.skipif(
1438
+ torch.accelerator.device_count() < 1,
1439
+ reason="Need at least 1 GPU (CUDA IPC handles).",
1440
+)
1441
+def test_ipc_trainer_send_weights_drives_client_in_order():
1442
+ """send_weights issues start -> update -> finish and ships per-round metadata;
1443
+ the packed wire param rides the init info, not the per-round update_info."""
1444
+ client = RecordingClient()
1445
+ engine = IPCTrainerWeightTransferEngine(
1446
+ client=client,
1447
+ source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))),
1448
+ packed=False,
1449
+ )
1450
+
1451
+ engine.send_weights()
1452
+
1453
+ assert client.order == ["start", "update", "finish"]
1454
+ assert client.last_update_info is not None
1455
+ assert client.last_update_info["names"] == ["w"]
1456
+ assert client.last_update_info["shapes"] == [[4]]
1457
+ assert "packed" not in client.last_update_info
1458
+
1459
+
1460
+def test_ipc_trainer_init_ships_packed_to_worker():
1461
+ """trainer_init drives the inference-side init handshake and propagates the
1462
+ must-agree `packed` flag to the worker."""
1463
+ if torch.accelerator.device_count() < 1:
1464
+ pytest.skip("Need at least 1 GPU (CUDA IPC handles).")
1465
+
1466
+ client = RecordingClient()
1467
+ engine = WeightTransferTrainerFactory.trainer_init(
1468
+ init_info=IPCTrainerInitInfo(rank=0, packed=True), # backend from init info
1469
+ client=client,
1470
+ source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))),
1471
+ )
1472
+
1473
+ assert isinstance(engine, IPCTrainerWeightTransferEngine)
1474
+ assert engine.is_sender is True
1475
+ assert engine.packed is True
1476
+ assert client.order == ["init"]
1477
+ assert client.last_init_info == {"packed": True}
vllm/distributed/weight_transfer/base.py
+37
-12
@@ -5,7 +5,15 @@
5
from abc import ABC, abstractmethod
6
from collections.abc import Iterator
7
from dataclasses import dataclass, field
8
-from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable
8
+from typing import (
9
+ TYPE_CHECKING,
10
+ Any,
11
+ ClassVar,
12
+ Generic,
13
+ Protocol,
14
+ TypeVar,
15
+ runtime_checkable,
16
+)
17
18
import torch
19
from typing_extensions import Self
@@ -18,7 +26,7 @@ from vllm.config.weight_transfer import WeightTransferConfig
26
27
TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
28
TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
21
-TConfig = TypeVar("TConfig", bound="WeightTransferConfig")
29
+TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo")
30
31
# A trainer supplies its parameters as a `WeightSource` (defined below): a
32
# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()`
@@ -108,7 +116,7 @@ class WeightTransferInitInfo(ABC): # noqa: B024
116
117
118
@dataclass
111
-class TrainerInitInfo(WeightTransferInitInfo):
119
+class TrainerInitInfo:
120
"""Base trainer-side init info: which trainer rank drives the transfer.
121
122
`rank` is this trainer process's rank, provided **explicitly** by the
@@ -118,10 +126,25 @@ class TrainerInitInfo(WeightTransferInitInfo):
126
while every rank still runs the trainer-side collectives. Backend subclasses
127
add their own (positional) fields; `rank` is keyword-only so that ordering
128
never conflicts.
129
+
130
+ Every concrete subclass sets a class-level `backend` string (the same key it
131
+ registers under in `WeightTransferTrainerFactory`). The factory reads it to
132
+ dispatch, so callers pass only the init info/ It is a `ClassVar`
133
+ (a fixed per-backend constant), so it is not an ``__init__`` field.
134
"""
135
136
+ backend: ClassVar[str]
137
+
138
rank: int = field(kw_only=True)
139
140
+ def __init_subclass__(cls, **kwargs: Any) -> None:
141
+ super().__init_subclass__(**kwargs)
142
+ if not getattr(cls, "backend", None):
143
+ raise TypeError(
144
+ f"{cls.__name__} must set a class-level `backend` string "
145
+ "(the WeightTransferTrainerFactory registry key)."
146
+ )
147
+
148
@property
149
def is_sender(self) -> bool:
150
return self.rank == 0
@@ -373,7 +396,7 @@ class VLLMWeightSyncClient(Protocol):
396
def finish_weight_update(self, weight_version: str | None = None) -> None: ...
397
398
376
-class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
399
+class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]):
400
"""Trainer-side weight transfer engine.
401
402
Symmetric to `WeightTransferEngine` but lives in the training process.
@@ -382,6 +405,13 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
405
plans) on `self`. The `WeightSource` is required at `trainer_init`,
406
then replayed each round by the no-argument `send_weights()`.
407
408
+ Unlike the worker engine, the trainer side does not take a
409
+ `WeightTransferConfig`: the backend is selected from the init info's
410
+ `backend` `ClassVar` (so callers pass only the init info), and the static
411
+ wire params (packed, buffer sizes) ride the backend-specific
412
+ `TrainerInitInfo`, which the sender also propagates to the worker at the init
413
+ handshake.
414
+
415
Multi-rank trainers: `trainer_init` and `send_weights` are
416
called on *every* trainer rank. Rank 0 is the sender, resolved once at
417
`trainer_init` into `is_sender`. Non-sender ranks still run every
@@ -392,22 +422,18 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
422
423
Subclasses should define:
424
init_info_cls: Type of backend-specific trainer init info
395
- config_cls: Type of backend-specific config
425
"""
426
398
- # Subclasses should override these class attributes
399
- init_info_cls: type[TInitInfo]
400
- config_cls: type[TConfig]
427
+ # Subclasses should override this class attribute
428
+ init_info_cls: type[TTrainerInitInfo]
429
430
def __init__(
431
self,
404
- config: TConfig,
432
*,
433
client: "VLLMWeightSyncClient",
434
source: "WeightSource",
435
is_sender: bool = True,
436
) -> None:
410
- self.config = config
437
self.is_sender = is_sender
438
# The real client is held on every rank; each engine only *calls* it when
439
# `is_sender`, so non-sender ranks never touch the wire.
@@ -418,8 +444,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
444
@abstractmethod
445
def trainer_init(
446
cls,
421
- config: TConfig,
422
- init_info: TInitInfo,
447
+ init_info: TTrainerInitInfo,
448
*,
449
client: "VLLMWeightSyncClient",
450
source: "WeightSource",
vllm/distributed/weight_transfer/factory.py
+21
-9
@@ -18,9 +18,9 @@ if TYPE_CHECKING:
18
from vllm.config import VllmConfig
19
from vllm.config.weight_transfer import WeightTransferConfig
20
from vllm.distributed.weight_transfer.base import (
21
+ TrainerInitInfo,
22
VLLMWeightSyncClient,
23
WeightSource,
23
- WeightTransferInitInfo,
24
)
25
26
logger = init_logger(__name__)
@@ -164,9 +164,7 @@ class WeightTransferTrainerFactory:
164
@classmethod
165
def trainer_init(
166
cls,
167
- backend: str,
168
- config: "WeightTransferConfig",
169
- init_info: "WeightTransferInitInfo",
167
+ init_info: "TrainerInitInfo",
168
*,
169
client: "VLLMWeightSyncClient",
170
source: "WeightSource",
@@ -176,16 +174,22 @@ class WeightTransferTrainerFactory:
174
Called on every trainer rank (multi-rank trainers construct on all
175
ranks; the sender is resolved inside the engine's ``trainer_init``).
176
177
+ The trainer side takes no `WeightTransferConfig` and no separate
178
+ `backend` argument: the backend is read from ``init_info.backend`` (a
179
+ `ClassVar` on each `TrainerInitInfo` subclass), and the static wire
180
+ params ride `init_info`.
181
+
182
Args:
180
- backend: Backend name (must be registered).
181
- config: Backend-specific weight transfer config.
182
- init_info: Backend-specific trainer init info.
183
+ init_info: Backend-specific trainer init info. Its `backend`
184
+ selects the engine; it also carries the wire params (e.g.
185
+ `packed`).
186
client: Inference-side control-plane client.
187
source: `WeightSource` of `(name, tensor)` pairs to send each round.
188
189
Raises:
187
- ValueError: If the backend is not registered.
190
+ ValueError: If `init_info.backend` is not registered.
191
"""
192
+ backend = init_info.backend
193
if backend not in cls._registry:
194
available = list(cls._registry.keys())
195
raise ValueError(
@@ -200,7 +204,6 @@ class WeightTransferTrainerFactory:
204
)
205
206
return engine_cls.trainer_init(
203
- config=config,
207
init_info=init_info,
208
client=client,
209
source=source,
@@ -228,3 +231,12 @@ WeightTransferEngineFactory.register_engine(
231
"vllm.distributed.weight_transfer.sparse_nccl_engine",
232
"SparseNCCLWeightTransferEngine",
233
)
234
+
235
+
236
+# Trainer-side engines. Backends register here as they migrate to the stateful
237
+# trainer engine; NCCL / sparse NCCL keep their static trainer path until then.
238
+WeightTransferTrainerFactory.register_engine(
239
+ "ipc",
240
+ "vllm.distributed.weight_transfer.ipc_engine",
241
+ "IPCTrainerWeightTransferEngine",
242
+)
vllm/distributed/weight_transfer/ipc_engine.py
+184
-223
@@ -3,19 +3,21 @@
3
"""IPC-based weight transfer engine using CUDA IPC for communication."""
4
5
import pickle
6
-from collections.abc import Callable, Iterator
6
from dataclasses import asdict, dataclass
8
-from typing import TYPE_CHECKING, Any
7
+from typing import TYPE_CHECKING, Any, ClassVar
8
9
import pybase64 as base64
11
-import ray
12
-import requests
10
import torch
11
from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
12
+from typing_extensions import Self
13
14
from vllm import envs
15
from vllm.config.weight_transfer import WeightTransferConfig
16
from vllm.distributed.weight_transfer.base import (
17
+ TrainerInitInfo,
18
+ TrainerWeightTransferEngine,
19
+ VLLMWeightSyncClient,
20
+ WeightSource,
21
WeightTransferEngine,
22
WeightTransferInitInfo,
23
WeightTransferUpdateInfo,
@@ -31,47 +33,35 @@ from vllm.distributed.weight_transfer.packed_tensor import (
33
34
35
@dataclass
34
-class IPCTrainerSendWeightsArgs:
35
- """Arguments for IPC trainer_send_weights method."""
36
-
37
- send_mode: str | Callable[["IPCWeightTransferUpdateInfo"], None]
38
- """How to send updates to vLLM. Either a string ('ray' or 'http') for
39
- built-in transports, or a callable that receives an
40
- IPCWeightTransferUpdateInfo and performs the send."""
41
- llm_handle: Any = None
42
- """Ray actor handle or list of handles (required for 'ray' send_mode)."""
43
- url: str | None = None
44
- """Base URL for HTTP endpoint (required for 'http' send_mode)."""
45
- packed: bool = False
46
- """Whether to use packed tensor transfer for bounded-memory chunking."""
47
- packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
48
- """Size in bytes for each packed tensor buffer when packed=True."""
36
+class IPCWeightTransferInitInfo(WeightTransferInitInfo):
37
+ """Worker-side init info for IPC weight transfer. No rendezvous needed.
38
50
- def __post_init__(self):
51
- """Validate that required arguments are provided for the selected mode."""
52
- if callable(self.send_mode):
53
- return
54
- if self.send_mode == "ray" and self.llm_handle is None:
55
- raise ValueError("llm_handle is required for 'ray' send_mode")
56
- if self.send_mode == "http" and self.url is None:
57
- raise ValueError("url is required for 'http' send_mode")
58
- if self.send_mode not in ("ray", "http"):
59
- raise ValueError(
60
- f"send_mode must be 'ray', 'http', or a callable, "
61
- f"got {self.send_mode!r}"
62
- )
39
+ `packed` is a must-agree wire param: the trainer ships it here at the init
40
+ handshake so the worker decodes with the same setting the trainer encoded
41
+ with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does
42
+ not need the buffer size (producer-only)."""
43
+
44
+ packed: bool = False
45
46
47
@dataclass
66
-class IPCWeightTransferInitInfo(WeightTransferInitInfo):
67
- """Initialization info for IPC weight transfer backend. No init needed for IPC."""
48
+class IPCTrainerInitInfo(TrainerInitInfo):
49
+ """Trainer-side init info for IPC weight transfer. No rendezvous needed;
50
+ `rank` (from `TrainerInitInfo`) identifies this trainer process — rank 0
51
+ ships the merged IPC handles. All ranks still join the handle all-gather.
52
69
- pass
53
+ `packed` / `packed_buffer_size_bytes` are the transfer's wire params. The
54
+ trainer propagates `packed` to the worker at `trainer_init` so the two sides
55
+ cannot disagree. `backend` is the factory dispatch key."""
56
+
57
+ backend: ClassVar[str] = "ipc"
58
+ packed: bool = False
59
+ packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
60
61
62
@dataclass
63
class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
74
- """Update info for IPC weight transfer backend."""
64
+ """Per-round update info for the IPC weight transfer backend."""
65
66
names: list[str]
67
dtype_names: list[str]
@@ -85,8 +75,6 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
75
tensor_sizes: list[int] | None = None
76
"""Per-parameter sizes in bytes within the packed buffer.
77
Required when packed=True, unused otherwise."""
88
- packed: bool = False
89
- """Whether this update uses packed tensor format."""
78
79
def __post_init__(self):
80
if self.ipc_handles_pickled is not None:
@@ -119,17 +107,14 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
107
f"`shapes` should be of the same size as `names`: "
108
f"got {len(self.shapes)} and {len(self.names)}"
109
)
122
- if (
123
- not self.packed
124
- and isinstance(self.ipc_handles, list)
125
- and len(self.ipc_handles) != num_params
126
- ):
110
+ # Unpacked transfers carry a per-parameter handle dict (a list); packed
111
+ # transfers carry a single dict for the whole buffer, so the list check
112
+ # only applies to the list case.
113
+ if isinstance(self.ipc_handles, list) and len(self.ipc_handles) != num_params:
114
raise ValueError(
115
f"`ipc_handles` should be of the same size as `names`: "
116
f"got {len(self.ipc_handles)} and {len(self.names)}"
117
)
131
- if self.packed and self.tensor_sizes is None:
132
- raise ValueError("`tensor_sizes` is required when packed=True")
118
119
120
class IPCWeightTransferEngine(
@@ -154,57 +139,21 @@ class IPCWeightTransferEngine(
139
device: torch.device,
140
model: torch.nn.Module,
141
) -> None:
157
- """
158
- Initialize the IPC weight transfer engine.
159
-
160
- Args:
161
- config: The configuration for the weight transfer engine
162
- vllm_config: The full vLLM config
163
- device: The device this worker's model lives on
164
- model: The local model instance which will receive the weights
165
- """
142
super().__init__(config, vllm_config, device, model)
167
-
168
- def parse_update_info(
169
- self, update_dict: dict[str, Any]
170
- ) -> IPCWeightTransferUpdateInfo:
171
- """Parse update dict, deserializing pickled IPC handles if present.
172
-
173
- HTTP transport sends IPC handles as a base64-encoded pickle under the
174
- key ``ipc_handles_pickled``. This method deserializes them back into
175
- ``ipc_handles`` before constructing the typed dataclass, keeping
176
- serialization concerns out of the dataclass itself.
177
-
178
- Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
179
- payload is deserialized via ``pickle.loads``.
180
- """
181
- pickled = update_dict.pop("ipc_handles_pickled", None)
182
- if pickled is not None:
183
- if update_dict.get("ipc_handles") is not None:
184
- raise ValueError(
185
- "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
186
- )
187
-
188
- if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
189
- raise ValueError(
190
- "Refusing to deserialize `ipc_handles_pickled` without "
191
- "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
192
- )
193
-
194
- update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
195
-
196
- return super().parse_update_info(update_dict)
143
+ # Set from the trainer-supplied init info at the handshake; defaults are
144
+ # only for the (unreachable) receive-before-init case.
145
+ self.packed = False
146
147
def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
148
"""
200
- Initialize the weight transfer mechanism.
201
- This is called once at the beginning of training.
202
- No initialization needed for IPC backend.
149
+ Initialize the weight transfer mechanism. No data-plane rendezvous is
150
+ needed for IPC; this just records the trainer-supplied wire params so the
151
+ worker decodes exactly as the trainer encoded.
152
153
Args:
205
- init_info: IPC initialization info (empty)
154
+ init_info: IPC initialization info (carries `packed`).
155
"""
207
- pass
156
+ self.packed = init_info.packed
157
158
def start_weight_update(self) -> None:
159
"""Initialize layerwise reloading for the incoming checkpoint weights."""
@@ -226,6 +175,10 @@ class IPCWeightTransferEngine(
175
"""
176
Receive weights from the trainer via CUDA IPC handles and load them.
177
178
+ Whether the transfer is packed is read from `self.packed`, set at the
179
+ init handshake from the trainer's init info, so it is guaranteed to
180
+ match how the trainer encoded.
181
+
182
Args:
183
update_info: IPC update info containing parameter names, dtypes, shapes,
184
and IPC handles. Each IPC handle is a mapping between physical
@@ -238,8 +191,9 @@ class IPCWeightTransferEngine(
191
# rebuilt on the device the model lives on.
192
device_index = self.device.index
193
241
- if update_info.packed:
242
- assert update_info.tensor_sizes is not None
194
+ if self.packed:
195
+ if update_info.tensor_sizes is None:
196
+ raise ValueError("`tensor_sizes` is required when packed=True")
197
assert isinstance(update_info.ipc_handles, dict)
198
weights = packed_ipc_consumer(
199
ipc_handle=update_info.ipc_handles,
@@ -285,65 +239,108 @@ class IPCWeightTransferEngine(
239
pass
240
241
@staticmethod
288
- def trainer_send_weights(
289
- iterator: Iterator[tuple[str, torch.Tensor]],
290
- trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs,
291
- ) -> None:
292
- """Send weights from trainer to inference workers via CUDA IPC.
242
+ def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
243
+ """Removed. Use the stateful `IPCTrainerWeightTransferEngine` instead.
244
294
- Supports two transport modes ('ray' and 'http') and two transfer
295
- strategies:
296
- - Non-packed (default): all weights in a single API call.
297
- - Packed (packed=True): chunked transfer with bounded GPU memory.
245
+ Transitional stub kept only to satisfy the (still abstract)
246
+ `WeightTransferEngine.trainer_send_weights`; that member is dropped from
247
+ the worker ABC once every backend has migrated to the trainer engine.
248
+ """
249
+ raise NotImplementedError(
250
+ "The static IPC trainer path has been replaced by "
251
+ "IPCTrainerWeightTransferEngine. Build it via "
252
+ "WeightTransferTrainerFactory.trainer_init(IPCTrainerInitInfo(...), "
253
+ "client=..., source=...) and drive it with send_weights()."
254
+ )
255
299
- For multi-GPU training, all ranks must call this method in
300
- parallel. IPC handles are all-gathered across ranks and merged
301
- so that each vLLM worker can find its own GPU UUID. Only rank 0
302
- sends the payload to vLLM.
256
304
- .. note::
305
- This method calls ``update_weights`` internally. The caller must
306
- call ``start_weight_update`` before and ``finish_weight_update``
307
- after this method.
257
+class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]):
258
+ """Trainer-side CUDA IPC weight transfer engine.
259
309
- Args:
310
- iterator: Iterator of (name, tensor) pairs. For multi-GPU,
311
- each rank should yield the full tensor on its own GPU
312
- (e.g. via FSDP full_tensor()).
313
- trainer_args: IPCTrainerSendWeightsArgs or equivalent dict.
314
- """
315
- args = (
316
- IPCTrainerSendWeightsArgs(**trainer_args)
317
- if isinstance(trainer_args, dict)
318
- else trainer_args
319
- )
320
- device_index = torch.accelerator.current_device_index()
321
- gpu_uuid = str(torch.cuda.get_device_properties(device_index).uuid)
322
- if args.packed:
323
- IPCWeightTransferEngine._send_packed(iterator, args, gpu_uuid)
324
- else:
325
- IPCWeightTransferEngine._send_unpacked(iterator, args, gpu_uuid)
260
+ Called on every trainer rank. For multi-rank (e.g. FSDP) trainers all ranks
261
+ iterate the source (materializing each tensor) and contribute to the
262
+ IPC-handle all-gather; only the sender (rank 0) ships the merged handles to
263
+ the inference side. IPC transfer is straight-line (no concurrent broadcast
264
+ like NCCL): `update_weights` *is* the transfer, and it rides the client, so
265
+ it no-ops on non-senders.
266
327
- @staticmethod
328
- def _is_rank_zero() -> bool:
329
- """Return True if this is rank 0 or no distributed group exists."""
330
- if not torch.distributed.is_initialized():
331
- return True
332
- return torch.distributed.get_rank() == 0
267
+ `packed` / `packed_buffer_size_bytes` come from `IPCTrainerInitInfo`; the
268
+ sender propagates `packed` to the worker at `trainer_init`.
269
+ """
270
+
271
+ init_info_cls = IPCTrainerInitInfo
272
+
273
+ def __init__(
274
+ self,
275
+ *,
276
+ client: VLLMWeightSyncClient,
277
+ source: WeightSource,
278
+ is_sender: bool = True,
279
+ packed: bool = False,
280
+ packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
281
+ ) -> None:
282
+ super().__init__(client=client, source=source, is_sender=is_sender)
283
+ self.packed = packed
284
+ self.packed_buffer_size_bytes = packed_buffer_size_bytes
285
+ self.device_index = torch.accelerator.current_device_index()
286
+ self.gpu_uuid = str(torch.cuda.get_device_properties(self.device_index).uuid)
287
+
288
+ @classmethod
289
+ def trainer_init(
290
+ cls,
291
+ init_info: IPCTrainerInitInfo,
292
+ *,
293
+ client: VLLMWeightSyncClient,
294
+ source: WeightSource,
295
+ ) -> Self:
296
+ engine = cls(
297
+ client=client,
298
+ source=source,
299
+ is_sender=init_info.is_sender,
300
+ packed=init_info.packed,
301
+ packed_buffer_size_bytes=init_info.packed_buffer_size_bytes,
302
+ )
303
+ # IPC needs no data-plane rendezvous. The sender ships the must-agree
304
+ # `packed` flag so the worker decodes exactly as this trainer encodes.
305
+ if engine.is_sender:
306
+ engine.client.init_weight_transfer_engine({"packed": init_info.packed})
307
+ return engine
308
+
309
+ def send_weights(self) -> None:
310
+ source = self.source
311
+ if self.is_sender:
312
+ self.client.start_weight_update()
313
+ # Unpacked returns strong refs to its IPC-shared copies; they must stay
314
+ # alive across `finish_weight_update` too.
315
+ weight_refs = self._send(source)
316
+ if self.is_sender:
317
+ self.client.finish_weight_update()
318
+ self._post_send_sync()
319
+ del weight_refs
320
+
321
+ # ---- data plane (runs on all ranks; only the sender ships) ----
322
+
323
+ def _send(self, source: WeightSource) -> list[torch.Tensor] | None:
324
+ if self.packed:
325
+ self._send_packed(source)
326
+ return None
327
+ return self._send_unpacked(source)
328
334
- @staticmethod
329
def _all_gather_and_merge_handles(
330
+ self,
331
handles: list[dict[str, tuple]],
332
) -> list[dict[str, tuple]]:
333
"""All-gather and merge IPC handle dicts across ranks in one call.
334
335
Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one
336
per parameter or one per chunk). A single all_gather_object
342
- collects every rank's full list, then rank 0 merges per-index so
337
+ collects every rank's full list, then the sender merges per-index so
338
each dict maps every GPU UUID to its args.
339
345
- Non-rank-0 returns a list of empty dicts.
346
- No-op (returns handles unchanged) when no distributed group exists.
340
+ The all-gather runs over the *default* process group; this assumes the
341
+ default group is exactly the set of colocated trainer ranks and that the
342
+ sender is a member. No-op (returns handles unchanged) when no distributed
343
+ group exists.
344
"""
345
if (
346
not torch.distributed.is_initialized()
@@ -357,7 +354,7 @@ class IPCWeightTransferEngine(
354
torch.distributed.barrier()
355
torch.cuda.synchronize()
356
360
- if torch.distributed.get_rank() == 0:
357
+ if self.is_sender:
358
merged: list[dict[str, tuple]] = []
359
for param_idx in range(len(handles)):
360
m: dict[str, tuple] = {}
@@ -378,25 +375,23 @@ class IPCWeightTransferEngine(
375
torch.distributed.barrier()
376
torch.cuda.ipc_collect()
377
381
- @staticmethod
382
- def _send_unpacked(
383
- iterator: Iterator[tuple[str, torch.Tensor]],
384
- args: IPCTrainerSendWeightsArgs,
385
- gpu_uuid: str,
386
- ) -> None:
387
- """Send all weights in a single API call (non-packed mode)."""
378
+ def _send_unpacked(self, source: WeightSource) -> list[torch.Tensor]:
379
+ """Iterate the source, build one IPC handle per param, all-gather the
380
+ handles across ranks, and (sender) ship them in one update call.
381
+
382
+ Returns the strong refs to every contiguous copy. reduce_tensor's args
383
+ do NOT keep storage alive, and non-contiguous inputs allocate fresh
384
+ storage in .contiguous(); the caller must keep these alive until the
385
+ post-send barrier (past `finish`) so the consumer's IPC views stay valid.
386
+ """
387
names: list[str] = []
388
dtype_names: list[str] = []
389
shapes: list[list[int]] = []
390
ipc_handles: list[dict[str, tuple]] = []
392
- # Hold strong refs to every contiguous copy until the send + post-send
393
- # sync completes. reduce_tensor's returned args do NOT keep storage
394
- # alive, and non-contiguous inputs allocate fresh storage in
395
- # .contiguous() that would otherwise be GC'd before the consumer opens
396
- # the IPC handle.
391
weight_refs: list[torch.Tensor] = []
392
399
- for name, tensor in iterator:
393
+ for name, tensor in source:
394
+ # FSDP shards were gathered by the source; ensure contiguity here.
395
names.append(name)
396
dtype_names.append(str(tensor.dtype).split(".")[-1])
397
shapes.append(list(tensor.shape))
@@ -404,98 +399,64 @@ class IPCWeightTransferEngine(
399
weight = tensor.detach().contiguous()
400
weight_refs.append(weight)
401
_, ipc_args = reduce_tensor(weight)
407
- ipc_handles.append({gpu_uuid: ipc_args})
408
-
409
- ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
410
-
411
- if IPCWeightTransferEngine._is_rank_zero():
412
- IPCWeightTransferEngine._do_send(
413
- args=args,
414
- names=names,
415
- dtype_names=dtype_names,
416
- shapes=shapes,
417
- ipc_handles=ipc_handles,
418
- )
419
-
420
- IPCWeightTransferEngine._post_send_sync()
402
+ ipc_handles.append({self.gpu_uuid: ipc_args})
403
+
404
+ ipc_handles = self._all_gather_and_merge_handles(ipc_handles)
405
+ self._do_send(
406
+ names=names,
407
+ dtype_names=dtype_names,
408
+ shapes=shapes,
409
+ ipc_handles=ipc_handles,
410
+ )
411
+ return weight_refs
412
422
- @staticmethod
423
- def _send_packed(
424
- iterator: Iterator[tuple[str, torch.Tensor]],
425
- args: IPCTrainerSendWeightsArgs,
426
- gpu_uuid: str,
427
- ) -> None:
413
+ def _send_packed(self, source: WeightSource) -> None:
414
"""Send weights in bounded-memory chunks (packed mode)."""
429
- post_iter_func: Callable = lambda item: item[1]
430
-
415
for chunk in packed_ipc_producer(
432
- iterator=iterator,
433
- gpu_uuid=gpu_uuid,
434
- post_iter_func=post_iter_func,
435
- buffer_size_bytes=args.packed_buffer_size_bytes,
416
+ iterator=iter(source),
417
+ gpu_uuid=self.gpu_uuid,
418
+ post_iter_func=lambda item: item[1],
419
+ buffer_size_bytes=self.packed_buffer_size_bytes,
420
):
437
- ipc_handle = IPCWeightTransferEngine._all_gather_and_merge_handles(
438
- [chunk.ipc_handle]
439
- )[0]
440
-
441
- if IPCWeightTransferEngine._is_rank_zero():
442
- IPCWeightTransferEngine._do_send(
443
- args=args,
444
- names=chunk.names,
445
- dtype_names=chunk.dtype_names,
446
- shapes=chunk.shapes,
447
- ipc_handles=ipc_handle,
448
- tensor_sizes=chunk.tensor_sizes,
449
- packed=True,
450
- )
451
-
452
- IPCWeightTransferEngine._post_send_sync()
421
+ ipc_handle = self._all_gather_and_merge_handles([chunk.ipc_handle])[0]
422
+ self._do_send(
423
+ names=chunk.names,
424
+ dtype_names=chunk.dtype_names,
425
+ shapes=chunk.shapes,
426
+ ipc_handles=ipc_handle,
427
+ tensor_sizes=chunk.tensor_sizes,
428
+ )
429
+ # Per-chunk barrier: the producer reuses a single IPC buffer across
430
+ # chunks, but only the sender waits for the consumers (via _do_send).
431
+ # Without syncing every rank here, non-sender ranks race ahead and
432
+ # overwrite their buffer while their colocated worker is still
433
+ # reading the current chunk, silently corrupting the transfer.
434
+ self._post_send_sync()
435
454
- @staticmethod
436
def _do_send(
456
- args: IPCTrainerSendWeightsArgs,
437
+ self,
438
names: list[str],
439
dtype_names: list[str],
440
shapes: list[list[int]],
441
ipc_handles: list[dict[str, tuple]] | dict[str, tuple],
442
tensor_sizes: list[int] | None = None,
462
- packed: bool = False,
443
) -> None:
464
- """Send a single update payload via the configured transport."""
444
+ """Build one update payload and ship it via the client. Only the sender
445
+ ships (non-sender ranks already contributed to the handle all-gather).
446
+
447
+ Emits raw `ipc_handles`; transports that cannot carry them natively
448
+ (HTTP/JSON) pickle them in their client (see `HTTPVLLMWeightSyncClient`).
449
+ """
450
+ if not self.is_sender:
451
+ return
452
update_fields: dict[str, Any] = {
453
"names": names,
454
"dtype_names": dtype_names,
455
"shapes": shapes,
469
- "packed": packed,
456
+ "ipc_handles": ipc_handles,
457
}
458
if tensor_sizes is not None:
459
update_fields["tensor_sizes"] = tensor_sizes
460
474
- update_fields["ipc_handles"] = ipc_handles
461
update_info = IPCWeightTransferUpdateInfo(**update_fields)
476
-
477
- if callable(args.send_mode):
478
- args.send_mode(update_info)
479
- elif args.send_mode == "ray":
480
- handles = (
481
- args.llm_handle
482
- if isinstance(args.llm_handle, list)
483
- else [args.llm_handle]
484
- )
485
- ray.get(
486
- [
487
- h.update_weights.remote(dict(update_info=asdict(update_info)))
488
- for h in handles
489
- ]
490
- )
491
- elif args.send_mode == "http":
492
- pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode(
493
- "utf-8"
494
- )
495
- http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"}
496
- http_fields["ipc_handles_pickled"] = pickled_handles
497
-
498
- url = f"{args.url}/update_weights"
499
- payload = {"update_info": http_fields}
500
- response = requests.post(url, json=payload, timeout=300)
501
- response.raise_for_status()
462
+ self.client.update_weights(asdict(update_info))
vllm/distributed/weight_transfer/packed_tensor.py
+1
-2
@@ -10,8 +10,7 @@ from typing import Any
10
import torch
11
from torch.multiprocessing.reductions import reduce_tensor
12
13
-# Default values for packed tensor configuration.
14
-# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights.
13
+# Default values for packed tensor transfer.
14
DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB
15
DEFAULT_PACKED_NUM_BUFFERS = 2
16