@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2ce6f3cf67
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 14:00 UTC

[Feat][RL][2/2] Native Weight Syncing API: IPC (#34171)

Signed-off-by: hao-aaron <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Signed-off-by: ahao-anyscale <ahao@anyscale.com>

Aaron Hao committed Feb 27, 2026 at 12:45 UTC 2ce6f3cf67934ebe199188c9a1f83ff1c2d8ba96
14 files changed +1189 -45
.buildkite/test-amd.yaml
+2 -1
@@ -278,7 +278,8 @@ steps:
278 - popd
279 # NEW rlhf examples
280 - pushd ../examples/offline_inference/new_weight_syncing
281 - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
281 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
282 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
283 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
284 - popd
285
.buildkite/test_areas/distributed.yaml
+2 -1
@@ -103,7 +103,8 @@ steps:
103 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
104 # NEW rlhf examples
105 - cd new_weight_syncing
106 - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
106 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
107 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
108
109 - label: Distributed Tests (8 GPUs)(H100)
110 timeout_in_minutes: 10
examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
+6 -2
@@ -42,6 +42,7 @@ from vllm.distributed.weight_transfer.base import (
42 WeightTransferUpdateRequest,
43 )
44 from vllm.distributed.weight_transfer.nccl_engine import (
45 + NCCLTrainerSendWeightsArgs,
46 NCCLWeightTransferEngine,
47 NCCLWeightTransferInitInfo,
48 NCCLWeightTransferUpdateInfo,
@@ -152,11 +153,14 @@ class TrainModel:
153
154 def broadcast_weights(self, packed: bool = True):
155 """Broadcast weights to the inference engine."""
155 - NCCLWeightTransferEngine.trainer_send_weights(
156 - iterator=self.model.named_parameters(),
156 + trainer_args = NCCLTrainerSendWeightsArgs(
157 group=self.model_update_group,
158 packed=packed,
159 )
160 + NCCLWeightTransferEngine.trainer_send_weights(
161 + iterator=self.model.named_parameters(),
162 + trainer_args=trainer_args,
163 + )
164
165 @torch.inference_mode()
166 def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
examples/offline_inference/new_weight_syncing/rlhf_ipc.py new
+149
@@ -0,0 +1,149 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray,
5 +with IPC-based weight syncing APIs
6 +
7 +The script colocates the training and inference workloads onto the same GPU using Ray.
8 +
9 +The example performs the following steps:
10 +
11 +* Request a placement group of 1 GPU.
12 +* Place the inference model on the above GPU using the placement group.
13 +* Place and load the training model on the same GPU using the placement group.
14 +* Generate text from a list of prompts using the inference engine.
15 +* Update the weights of the training model and broadcast the updated weights
16 + to the inference engine by using CUDA IPC handles. Note that
17 + for demonstration purposes we simply zero out the weights.
18 +
19 +This example assumes a single-node cluster with a single GPU,
20 +but can be extended to multiple GPUs.
21 +"""
22 +
23 +import os
24 +
25 +import ray
26 +from ray.util.placement_group import placement_group
27 +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
28 +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,
35 +)
36 +
37 +
38 +class MyLLM(LLM):
39 + """Configure the vLLM worker for Ray placement group execution."""
40 +
41 + def __init__(self, *args, **kwargs):
42 + # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
43 + # so that vLLM can manage its own device placement within the worker.
44 + os.environ.pop("CUDA_VISIBLE_DEVICES", None)
45 + # Each worker uses 0.4 GPU so that two instances fit on the same GPU.
46 + os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
47 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
48 + # needed for ipc handle serialization
49 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
50 + super().__init__(*args, **kwargs)
51 +
52 +
53 +# Load the OPT-125M model onto GPU 0 for the training workload.
54 +
55 +MODEL_NAME = "facebook/opt-125m"
56 +
57 +
58 +@ray.remote
59 +class TrainModel:
60 + def __init__(self, llm_handle: ray.actor.ActorHandle):
61 + self.train_model = AutoModelForCausalLM.from_pretrained(
62 + MODEL_NAME,
63 + )
64 + self.train_model.to("cuda:0")
65 + self.llm_handle = llm_handle
66 +
67 + 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()))
71 + )
72 +
73 + def broadcast_weights(self, llm_handle: ray.actor.ActorHandle):
74 + """Broadcast weights to the inference engine using IPC."""
75 + self.llm_handle = llm_handle
76 + trainer_args = IPCTrainerSendWeightsArgs(mode="ray", llm_handle=llm_handle)
77 + IPCWeightTransferEngine.trainer_send_weights(
78 + iterator=self.train_model.named_parameters(),
79 + trainer_args=trainer_args,
80 + )
81 +
82 +
83 +ray.init()
84 +
85 +pg_colocate = placement_group([{"GPU": 1, "CPU": 0}])
86 +ray.get(pg_colocate.ready())
87 +
88 +
89 +llm = ray.remote(
90 + num_cpus=0,
91 + num_gpus=0,
92 + scheduling_strategy=PlacementGroupSchedulingStrategy(
93 + placement_group=pg_colocate,
94 + placement_group_capture_child_tasks=True,
95 + ),
96 +)(MyLLM).remote(
97 + model=MODEL_NAME,
98 + enforce_eager=True,
99 + tensor_parallel_size=1,
100 + distributed_executor_backend="ray",
101 + gpu_memory_utilization=0.7,
102 + weight_transfer_config=WeightTransferConfig(backend="ipc"),
103 + load_format="dummy",
104 +)
105 +
106 +train_model = TrainModel.options(
107 + num_gpus=0.1,
108 + num_cpus=0,
109 + scheduling_strategy=PlacementGroupSchedulingStrategy(
110 + placement_group=pg_colocate, placement_group_capture_child_tasks=True
111 + ),
112 +).remote(llm)
113 +
114 +
115 +# Generate text from the prompts.
116 +prompts = [
117 + "Hello, my name is",
118 + "The president of the United States is",
119 + "The capital of France is",
120 + "The future of AI is",
121 +]
122 +
123 +sampling_params = SamplingParams(temperature=0)
124 +
125 +outputs = ray.get(llm.generate.remote(prompts, sampling_params))
126 +
127 +print("-" * 50)
128 +for output in outputs:
129 + prompt = output.prompt
130 + generated_text = output.outputs[0].text
131 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
132 + print("-" * 50)
133 +
134 +ray.get(llm.sleep.remote(level=0))
135 +
136 +ray.get(train_model.init_weight_transfer.remote())
137 +# Synchronize the updated weights to the inference engine using batched API.
138 +ray.get(train_model.broadcast_weights.remote(llm))
139 +
140 +ray.get(llm.wake_up.remote(tags=["scheduling"]))
141 +
142 +# Generate text with the updated model.
143 +outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
144 +print("-" * 50)
145 +for output in outputs_updated:
146 + prompt = output.prompt
147 + generated_text = output.outputs[0].text
148 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
149 + print("-" * 50)
examples/offline_inference/new_weight_syncing/rlhf_nccl.py renamed
+10 -2
@@ -36,6 +36,7 @@ from transformers import AutoModelForCausalLM
36 from vllm import LLM, SamplingParams
37 from vllm.config import WeightTransferConfig
38 from vllm.distributed.weight_transfer.nccl_engine import (
39 + NCCLTrainerSendWeightsArgs,
40 NCCLWeightTransferEngine,
41 )
42 from vllm.utils.network_utils import get_ip, get_open_port
@@ -90,11 +91,14 @@ class TrainModel:
91
92 def broadcast_weights(self, packed: bool = True):
93 """Broadcast weights to the inference engine."""
93 - NCCLWeightTransferEngine.trainer_send_weights(
94 - iterator=self.model.named_parameters(),
94 + trainer_args = NCCLTrainerSendWeightsArgs(
95 group=self.model_update_group,
96 packed=packed,
97 )
98 + NCCLWeightTransferEngine.trainer_send_weights(
99 + iterator=self.model.named_parameters(),
100 + trainer_args=trainer_args,
101 + )
102
103
104 # Initialize Ray and set the visible devices. The vLLM engine will
@@ -156,6 +160,8 @@ for output in outputs:
160 print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
161 print("-" * 50)
162
163 +ray.get(llm.sleep.remote(level=0))
164 +
165 # Set up the communication channel between the training process and the
166 # inference engine.
167 master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
@@ -197,6 +203,8 @@ inference_handle = llm.update_weights.remote(
203 train_handle = train_model.broadcast_weights.remote(packed=True)
204 ray.get([train_handle, inference_handle])
205
206 +ray.get(llm.wake_up.remote(tags=["scheduling"]))
207 +
208 # Generate text with the updated model. The output is expected to be normal
209 # because the weights are updated.
210 outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
examples/online_serving/new_weight_syncing/rlhf_http_ipc.py new
+181
@@ -0,0 +1,181 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
5 +via HTTP API, with IPC-based weight syncing APIs.
6 +
7 +Unlike rlhf_nccl.py which uses NCCL and can use separate GPUs, this script
8 +uses CUDA IPC which requires the training model and vLLM server to be on the
9 +same GPU. Memory must be carefully managed to fit both models.
10 +
11 +Unlike rlhf.py which creates a vLLM instance programmatically, this script
12 +assumes you have already started a vLLM server using `vllm serve`. It uses:
13 +- OpenAI-compatible API for inference requests
14 +- HTTP endpoints for weight transfer control plane
15 +- CUDA IPC for actual weight data transfer
16 +
17 +Prerequisites:
18 + Start a vLLM server with weight transfer enabled and reduced GPU memory
19 + utilization to leave room for the training model:
20 +
21 + $ VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
22 + vllm serve facebook/opt-125m --enforce-eager \
23 + --weight-transfer-config '{"backend": "ipc"}' \
24 + --load-format dummy \
25 + --gpu-memory-utilization 0.5
26 +
27 + Then run this script:
28 +
29 + $ python rlhf_http_ipc.py
30 +
31 +The example performs the following steps:
32 +
33 +* Load the training model on GPU 0 (same GPU as the vLLM server).
34 +* Generate text using the vLLM server via OpenAI-compatible API. The output
35 + is expected to be nonsense because the server is initialized with dummy weights.
36 +* Initialize weight transfer via HTTP endpoint (no-op for IPC).
37 +* Broadcast the real weights from the training model to the vLLM server
38 + using CUDA IPC handles.
39 +* Generate text again to show normal output after the weight update.
40 +"""
41 +
42 +import os
43 +
44 +import requests
45 +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,
52 +)
53 +
54 +BASE_URL = "http://localhost:8000"
55 +MODEL_NAME = "facebook/opt-125m"
56 +
57 +# Enable insecure serialization for IPC handle serialization
58 +os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
59 +
60 +
61 +def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
62 + """Generate completions using the OpenAI-compatible API."""
63 + results = []
64 + for prompt in prompts:
65 + response = client.completions.create(
66 + model=model,
67 + prompt=prompt,
68 + max_tokens=32,
69 + temperature=0,
70 + )
71 + results.append(response.choices[0].text)
72 + return results
73 +
74 +
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 pause_generation(base_url: str) -> None:
84 + """Pause generation via HTTP endpoint."""
85 + url = f"{base_url}/pause"
86 + response = requests.post(url, timeout=60)
87 + response.raise_for_status()
88 +
89 +
90 +def resume_generation(base_url: str) -> None:
91 + """Resume generation via HTTP endpoint."""
92 + url = f"{base_url}/resume"
93 + response = requests.post(url, timeout=60)
94 + response.raise_for_status()
95 +
96 +
97 +def get_world_size(base_url: str) -> int:
98 + """Get world size from the vLLM server."""
99 + url = f"{base_url}/get_world_size"
100 + response = requests.get(url, timeout=10)
101 + response.raise_for_status()
102 + return response.json()["world_size"]
103 +
104 +
105 +def main():
106 + # IPC requires the training model to be on the same GPU as the vLLM server
107 + # The server should be started on GPU 0 with reduced memory utilization
108 + device = "cuda:0"
109 + torch.cuda.set_device(device)
110 +
111 + # Load the training model on the same GPU as the server
112 + # Use bfloat16 to reduce memory footprint
113 + print(f"Loading training model: {MODEL_NAME} on {device}")
114 + print(
115 + "Note: Ensure the vLLM server was started with --gpu-memory-utilization 0.5 "
116 + "or lower to leave room for the training model."
117 + )
118 + train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
119 + train_model.to(device)
120 + train_model.eval() # Set to eval mode to save memory
121 +
122 + # Create OpenAI client pointing to the vLLM server
123 + client = OpenAI(
124 + base_url=f"{BASE_URL}/v1",
125 + api_key="EMPTY", # vLLM doesn't require an API key by default
126 + )
127 +
128 + # Test prompts
129 + prompts = [
130 + "Hello, my name is",
131 + "The president of the United States is",
132 + "The capital of France is",
133 + "The future of AI is",
134 + ]
135 +
136 + # Generate text before weight update. The output is expected to be nonsense
137 + # because the server is initialized with dummy weights.
138 + print("-" * 50)
139 + print("Generating text BEFORE weight update (expect nonsense):")
140 + print("-" * 50)
141 + outputs = generate_completions(client, MODEL_NAME, prompts)
142 + for prompt, generated_text in zip(prompts, outputs):
143 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
144 + print("-" * 50)
145 +
146 + print("Initializing weight transfer (IPC backend)...")
147 +
148 + # Initialize weight transfer on vLLM server (no-op for IPC, but still required)
149 + init_weight_transfer_engine(BASE_URL)
150 +
151 + # Pause generation before weight sync
152 + pause_generation(BASE_URL)
153 +
154 + # Broadcast weights via IPC handles using HTTP mode
155 + print("Broadcasting weights via CUDA IPC (HTTP)...")
156 + trainer_args = IPCTrainerSendWeightsArgs(mode="http", url=BASE_URL)
157 + IPCWeightTransferEngine.trainer_send_weights(
158 + iterator=train_model.named_parameters(),
159 + trainer_args=trainer_args,
160 + )
161 +
162 + # Resume generation after weight sync
163 + resume_generation(BASE_URL)
164 +
165 + # Generate text after weight update. The output is expected to be normal
166 + # because the real weights are now loaded.
167 + print("-" * 50)
168 + print("Generating text AFTER weight update:")
169 + print("-" * 50)
170 + outputs_updated = generate_completions(client, MODEL_NAME, prompts)
171 + for prompt, generated_text in zip(prompts, outputs_updated):
172 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
173 + print("-" * 50)
174 +
175 + # Note: The training model and IPC handles remain in memory.
176 + # In a real RLHF training loop, you would update the training model
177 + # and create new IPC handles for each weight update.
178 +
179 +
180 +if __name__ == "__main__":
181 + main()
examples/online_serving/new_weight_syncing/rlhf_http_nccl.py renamed
+6 -2
@@ -39,6 +39,7 @@ from openai import OpenAI
39 from transformers import AutoModelForCausalLM
40
41 from vllm.distributed.weight_transfer.nccl_engine import (
42 + NCCLTrainerSendWeightsArgs,
43 NCCLWeightTransferEngine,
44 )
45 from vllm.utils.network_utils import get_ip, get_open_port
@@ -214,11 +215,14 @@ def main():
215
216 # Broadcast all weights from trainer to vLLM workers
217 print("Broadcasting weights via NCCL...")
217 - NCCLWeightTransferEngine.trainer_send_weights(
218 - iterator=train_model.named_parameters(),
218 + trainer_args = NCCLTrainerSendWeightsArgs(
219 group=model_update_group,
220 packed=True,
221 )
222 + NCCLWeightTransferEngine.trainer_send_weights(
223 + iterator=train_model.named_parameters(),
224 + trainer_args=trainer_args,
225 + )
226
227 # Wait for update_weights to complete
228 update_thread.join()
tests/distributed/test_weight_transfer.py
+453 -2
@@ -3,18 +3,26 @@
3 """Tests for weight transfer engine backends.
4
5 Unit tests for engine classes (parsing, validation, registry).
6 -Integration test for NCCL weight transfer between processes using Ray.
6 +Integration tests for NCCL and IPC weight transfer between processes using Ray.
7 """
8
9 +import base64
10 +import pickle
11 from unittest.mock import MagicMock
12
13 import pytest
14 import ray
15 import torch
16 +from torch.multiprocessing.reductions import reduce_tensor
17
18 from vllm.config.parallel import ParallelConfig
19 from vllm.config.weight_transfer import WeightTransferConfig
20 from vllm.distributed.weight_transfer import WeightTransferEngineFactory
21 +from vllm.distributed.weight_transfer.ipc_engine import (
22 + IPCWeightTransferEngine,
23 + IPCWeightTransferInitInfo,
24 + IPCWeightTransferUpdateInfo,
25 +)
26 from vllm.distributed.weight_transfer.nccl_engine import (
27 NCCLWeightTransferEngine,
28 NCCLWeightTransferInitInfo,
@@ -155,9 +163,29 @@ class TestEngineRegistry:
163 engine = WeightTransferEngineFactory.create_engine(config, parallel_config)
164 assert isinstance(engine, NCCLWeightTransferEngine)
165
166 + def test_create_engine_ipc(self):
167 + """Test factory creates IPC engine."""
168 + config = WeightTransferConfig(backend="ipc")
169 + parallel_config = create_mock_parallel_config()
170 + engine = WeightTransferEngineFactory.create_engine(config, parallel_config)
171 + assert isinstance(engine, IPCWeightTransferEngine)
172 +
173 def test_create_engine_invalid_backend(self):
174 """Test factory raises for invalid backend."""
160 - config = WeightTransferConfig(backend="invalid")
175 + # Pydantic validates Literal types at construction, so we can't create
176 + # a config with an invalid backend. Instead, we test by directly
177 + # accessing the registry or using model_construct to bypass validation.
178 + from pydantic import ValidationError
179 +
180 + # Test that Pydantic prevents invalid backend at construction
181 + with pytest.raises(ValidationError):
182 + WeightTransferConfig(backend="invalid")
183 +
184 + # Test factory error by creating a config with valid backend but
185 + # then manually modifying the backend attribute (bypassing validation)
186 + config = WeightTransferConfig(backend="nccl")
187 + # Use object.__setattr__ to bypass Pydantic validation
188 + object.__setattr__(config, "backend", "invalid")
189 parallel_config = create_mock_parallel_config()
190 with pytest.raises(ValueError, match="Invalid weight transfer backend"):
191 WeightTransferEngineFactory.create_engine(config, parallel_config)
@@ -344,3 +372,426 @@ def test_nccl_weight_transfer_between_processes():
372 f"Received shape: {result['received_shape']}, "
373 f"Received sum: {result['received_sum']}"
374 )
375 +
376 +
377 +# --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
378 +
379 +
380 +class TestIPCWeightTransferUpdateInfoValidation:
381 + """Test IPCWeightTransferUpdateInfo dataclass validation."""
382 +
383 + def test_valid_update_info(self):
384 + """Test creating valid IPCWeightTransferUpdateInfo."""
385 + if torch.cuda.device_count() < 1:
386 + pytest.skip("Need at least 1 GPU for this test")
387 +
388 + # Create a dummy tensor and IPC handle
389 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
390 + ipc_handle = reduce_tensor(dummy_tensor)
391 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
392 + ipc_handles = [{gpu_uuid: ipc_handle}]
393 +
394 + info = IPCWeightTransferUpdateInfo(
395 + names=["layer.weight"],
396 + dtype_names=["float32"],
397 + shapes=[[10, 10]],
398 + ipc_handles=ipc_handles,
399 + )
400 + assert info.names == ["layer.weight"]
401 + assert info.dtype_names == ["float32"]
402 + assert info.shapes == [[10, 10]]
403 + assert len(info.ipc_handles) == 1
404 +
405 + def test_mismatched_dtype_names_raises(self):
406 + """Test that mismatched dtype_names length raises ValueError."""
407 + if torch.cuda.device_count() < 1:
408 + pytest.skip("Need at least 1 GPU for this test")
409 +
410 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
411 + ipc_handle = reduce_tensor(dummy_tensor)
412 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
413 + ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
414 +
415 + with pytest.raises(ValueError, match="dtype_names"):
416 + IPCWeightTransferUpdateInfo(
417 + names=["layer.weight", "layer.bias"],
418 + dtype_names=["float32"], # Only one dtype
419 + shapes=[[10, 10], [10]],
420 + ipc_handles=ipc_handles,
421 + )
422 +
423 + def test_mismatched_shapes_raises(self):
424 + """Test that mismatched shapes length raises ValueError."""
425 + if torch.cuda.device_count() < 1:
426 + pytest.skip("Need at least 1 GPU for this test")
427 +
428 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
429 + ipc_handle = reduce_tensor(dummy_tensor)
430 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
431 + ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
432 +
433 + with pytest.raises(ValueError, match="shapes"):
434 + IPCWeightTransferUpdateInfo(
435 + names=["layer.weight", "layer.bias"],
436 + dtype_names=["float32", "float32"],
437 + shapes=[[10, 10]], # Only one shape
438 + ipc_handles=ipc_handles,
439 + )
440 +
441 + def test_mismatched_ipc_handles_raises(self):
442 + """Test that mismatched ipc_handles length raises ValueError."""
443 + if torch.cuda.device_count() < 1:
444 + pytest.skip("Need at least 1 GPU for this test")
445 +
446 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
447 + ipc_handle = reduce_tensor(dummy_tensor)
448 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
449 + ipc_handles = [{gpu_uuid: ipc_handle}] # Only one handle
450 +
451 + with pytest.raises(ValueError, match="ipc_handles"):
452 + IPCWeightTransferUpdateInfo(
453 + names=["layer.weight", "layer.bias"],
454 + dtype_names=["float32", "float32"],
455 + shapes=[[10, 10], [10]],
456 + ipc_handles=ipc_handles,
457 + )
458 +
459 + def test_valid_update_info_from_pickled(self):
460 + """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
461 + if torch.cuda.device_count() < 1:
462 + pytest.skip("Need at least 1 GPU for this test")
463 +
464 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
465 + ipc_handle = reduce_tensor(dummy_tensor)
466 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
467 + ipc_handles = [{gpu_uuid: ipc_handle}]
468 +
469 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
470 +
471 + info = IPCWeightTransferUpdateInfo(
472 + names=["layer.weight"],
473 + dtype_names=["float32"],
474 + shapes=[[10, 10]],
475 + ipc_handles_pickled=pickled,
476 + )
477 + assert info.ipc_handles == ipc_handles
478 + assert info.ipc_handles_pickled is None
479 +
480 + def test_both_handles_and_pickled_raises(self):
481 + """Test that providing both ipc_handles and ipc_handles_pickled raises."""
482 + if torch.cuda.device_count() < 1:
483 + pytest.skip("Need at least 1 GPU for this test")
484 +
485 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
486 + ipc_handle = reduce_tensor(dummy_tensor)
487 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
488 + ipc_handles = [{gpu_uuid: ipc_handle}]
489 +
490 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
491 +
492 + with pytest.raises(ValueError, match="Cannot specify both"):
493 + IPCWeightTransferUpdateInfo(
494 + names=["layer.weight"],
495 + dtype_names=["float32"],
496 + shapes=[[10, 10]],
497 + ipc_handles=ipc_handles,
498 + ipc_handles_pickled=pickled,
499 + )
500 +
501 + def test_neither_handles_nor_pickled_raises(self):
502 + """Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
503 + with pytest.raises(ValueError, match="must be provided"):
504 + IPCWeightTransferUpdateInfo(
505 + names=["layer.weight"],
506 + dtype_names=["float32"],
507 + shapes=[[10, 10]],
508 + )
509 +
510 + def test_empty_lists_valid(self):
511 + """Test that empty lists are valid."""
512 + info = IPCWeightTransferUpdateInfo(
513 + names=[],
514 + dtype_names=[],
515 + shapes=[],
516 + ipc_handles=[],
517 + )
518 + assert len(info.names) == 0
519 +
520 +
521 +# --- Unit Tests: IPC Engine Parsing ---
522 +
523 +
524 +class TestIPCEngineParsing:
525 + """Test IPCWeightTransferEngine parsing methods."""
526 +
527 + def test_parse_update_info_valid(self):
528 + """Test parsing valid update info dict."""
529 + if torch.cuda.device_count() < 1:
530 + pytest.skip("Need at least 1 GPU for this test")
531 +
532 + config = WeightTransferConfig(backend="ipc")
533 + parallel_config = create_mock_parallel_config()
534 + engine = IPCWeightTransferEngine(config, parallel_config)
535 +
536 + # Create dummy IPC handles
537 + dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
538 + dummy_tensor2 = torch.ones(50, device="cuda:0")
539 + ipc_handle1 = reduce_tensor(dummy_tensor1)
540 + ipc_handle2 = reduce_tensor(dummy_tensor2)
541 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
542 + ipc_handles = [{gpu_uuid: ipc_handle1}, {gpu_uuid: ipc_handle2}]
543 +
544 + update_info = engine.parse_update_info(
545 + {
546 + "names": ["w1", "w2"],
547 + "dtype_names": ["float32", "bfloat16"],
548 + "shapes": [[100, 100], [50]],
549 + "ipc_handles": ipc_handles,
550 + }
551 + )
552 +
553 + assert isinstance(update_info, IPCWeightTransferUpdateInfo)
554 + assert update_info.names == ["w1", "w2"]
555 + assert update_info.dtype_names == ["float32", "bfloat16"]
556 + assert update_info.shapes == [[100, 100], [50]]
557 + assert len(update_info.ipc_handles) == 2
558 +
559 + def test_parse_update_info_pickled(self):
560 + """Test parsing update info with pickled IPC handles (HTTP path)."""
561 + if torch.cuda.device_count() < 1:
562 + pytest.skip("Need at least 1 GPU for this test")
563 +
564 + config = WeightTransferConfig(backend="ipc")
565 + parallel_config = create_mock_parallel_config()
566 + engine = IPCWeightTransferEngine(config, parallel_config)
567 +
568 + dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
569 + dummy_tensor2 = torch.ones(50, device="cuda:0")
570 + ipc_handle1 = reduce_tensor(dummy_tensor1)
571 + ipc_handle2 = reduce_tensor(dummy_tensor2)
572 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
573 + ipc_handles = [{gpu_uuid: ipc_handle1}, {gpu_uuid: ipc_handle2}]
574 +
575 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
576 +
577 + update_info = engine.parse_update_info(
578 + {
579 + "names": ["w1", "w2"],
580 + "dtype_names": ["float32", "bfloat16"],
581 + "shapes": [[100, 100], [50]],
582 + "ipc_handles_pickled": pickled,
583 + }
584 + )
585 +
586 + assert isinstance(update_info, IPCWeightTransferUpdateInfo)
587 + assert update_info.names == ["w1", "w2"]
588 + assert len(update_info.ipc_handles) == 2
589 + assert update_info.ipc_handles_pickled is None
590 + assert gpu_uuid in update_info.ipc_handles[0]
591 + assert gpu_uuid in update_info.ipc_handles[1]
592 +
593 +
594 +# --- Integration Test: IPC Weight Transfer Between Ray Tasks ---
595 +
596 +
597 +def get_physical_gpu_id(device_index: int = 0) -> str:
598 + """Get physical GPU UUID for a device."""
599 + props = torch.cuda.get_device_properties(device_index)
600 + return str(props.uuid)
601 +
602 +
603 +@ray.remote(num_gpus=0.5)
604 +class TrainerActor:
605 + """Trainer actor that creates and holds CUDA IPC handles."""
606 +
607 + def __init__(self, tensor_shape: list[int], tensor_dtype: str):
608 + # Create tensor on GPU and keep it alive
609 + dtype = getattr(torch, tensor_dtype)
610 + self.tensor = torch.ones(tensor_shape, dtype=dtype, device="cuda:0")
611 + self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer
612 +
613 + # Create IPC handle (tensor must stay alive for IPC to work)
614 + ipc_handle = reduce_tensor(self.tensor)
615 + gpu_uuid = get_physical_gpu_id(0)
616 +
617 + torch.cuda.synchronize()
618 +
619 + self.ipc_handle_dict = {
620 + "ipc_handle": ipc_handle,
621 + "gpu_uuid": gpu_uuid,
622 + "shape": tensor_shape,
623 + "dtype": tensor_dtype,
624 + }
625 +
626 + def get_ipc_handle_dict(self) -> dict:
627 + """Return IPC handle dict. Tensor stays alive in this actor."""
628 + return self.ipc_handle_dict
629 +
630 +
631 +@ray.remote(num_gpus=0.5)
632 +def inference_receive_ipc_tensor(
633 + ipc_handle_dict: dict,
634 + mode: str = "ray",
635 +) -> dict:
636 + """Inference task that receives tensor via IPCWeightTransferEngine."""
637 + from unittest.mock import MagicMock
638 +
639 + import torch
640 +
641 + from vllm.config.parallel import ParallelConfig
642 + from vllm.config.weight_transfer import WeightTransferConfig
643 + from vllm.distributed.weight_transfer.ipc_engine import (
644 + IPCWeightTransferEngine,
645 + )
646 +
647 + # Create engine with mock parallel config
648 + config = WeightTransferConfig(backend="ipc")
649 + parallel_config = MagicMock(spec=ParallelConfig)
650 + parallel_config.rank = 0
651 + parallel_config.world_size = 1
652 + parallel_config.data_parallel_rank = 0
653 +
654 + engine = IPCWeightTransferEngine(config, parallel_config)
655 +
656 + # Initialize the engine (no-op for IPC)
657 + init_info = IPCWeightTransferInitInfo()
658 + engine.init_transfer_engine(init_info)
659 +
660 + # Receive weights with a no-op load_weights that captures the tensor
661 + received_tensors = []
662 +
663 + def noop_load_weights(weights: list[tuple[str, torch.Tensor]]):
664 + for name, tensor in weights:
665 + # Clone tensor to keep it after engine cleans up
666 + received_tensors.append((name, tensor.clone()))
667 +
668 + # Build update dict and go through parse_update_info (exercises __post_init__)
669 + ipc_handles = [{ipc_handle_dict["gpu_uuid"]: ipc_handle_dict["ipc_handle"]}]
670 +
671 + if mode == "ray":
672 + update_dict: dict = {
673 + "names": ["test.weight"],
674 + "dtype_names": [ipc_handle_dict["dtype"]],
675 + "shapes": [ipc_handle_dict["shape"]],
676 + "ipc_handles": ipc_handles,
677 + }
678 + elif mode == "http":
679 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
680 + update_dict = {
681 + "names": ["test.weight"],
682 + "dtype_names": [ipc_handle_dict["dtype"]],
683 + "shapes": [ipc_handle_dict["shape"]],
684 + "ipc_handles_pickled": pickled,
685 + }
686 + else:
687 + raise ValueError(f"Unknown mode: {mode}")
688 +
689 + update_info = engine.parse_update_info(update_dict)
690 + engine.receive_weights(update_info, noop_load_weights)
691 + torch.cuda.synchronize()
692 +
693 + # Verify we received the tensor
694 + success = False
695 + received_shape = None
696 + received_sum = None
697 +
698 + if len(received_tensors) == 1:
699 + name, tensor = received_tensors[0]
700 + received_shape = list(tensor.shape)
701 + received_sum = tensor.sum().item()
702 + # Check shape matches and values are all 42s (trainer sends 42s)
703 + if received_shape == ipc_handle_dict["shape"]:
704 + expected_sum = 42.0 * torch.tensor(ipc_handle_dict["shape"]).prod().item()
705 + if abs(received_sum - expected_sum) < 0.01:
706 + success = True
707 +
708 + engine.shutdown()
709 +
710 + return {
711 + "success": success,
712 + "received_shape": received_shape,
713 + "received_sum": received_sum,
714 + }
715 +
716 +
717 +@pytest.mark.skipif(
718 + torch.cuda.device_count() < 1,
719 + reason="Need at least 1 GPU to run IPC weight transfer test.",
720 +)
721 +@pytest.mark.parametrize("mode", ["ray", "http"])
722 +def test_ipc_weight_transfer_between_processes(mode: str):
723 + """Test IPC weight transfer from trainer to inference process using Ray.
724 +
725 + Parametrized over transport modes:
726 + - 'ray': ipc_handles passed directly.
727 + - 'http': ipc_handles pickled + base64-encoded, unpickled via __post_init__.
728 +
729 + IPC requires same-GPU access, so we use a placement group to co-locate
730 + the trainer actor and inference task on the same GPU.
731 + """
732 + from ray.util.placement_group import placement_group
733 + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
734 +
735 + ray.init(ignore_reinit_error=True)
736 +
737 + # Create a placement group to ensure both processes are on the same GPU
738 + # Use fractional GPUs so both tasks can share the same GPU bundle
739 + pg = placement_group([{"GPU": 1, "CPU": 2}])
740 + ray.get(pg.ready())
741 +
742 + scheduling_strategy = PlacementGroupSchedulingStrategy(
743 + placement_group=pg,
744 + placement_group_capture_child_tasks=True,
745 + )
746 +
747 + # Tensor to transfer: 100x100 filled with 42s
748 + tensor_shape = [100, 100]
749 + tensor_dtype = "float32"
750 +
751 + # Create trainer actor that holds the tensor and IPC handle (stays alive)
752 + trainer_actor = TrainerActor.options( # type: ignore[attr-defined]
753 + scheduling_strategy=scheduling_strategy
754 + ).remote(tensor_shape, tensor_dtype)
755 +
756 + # Get IPC handle dict (tensor stays alive in trainer actor)
757 + ipc_handle_dict = ray.get(trainer_actor.get_ipc_handle_dict.remote())
758 +
759 + # Receive tensor in inference process using IPC handles (on same GPU)
760 + # Trainer actor stays alive during this operation
761 + inference_result = ray.get(
762 + inference_receive_ipc_tensor.options(
763 + scheduling_strategy=scheduling_strategy
764 + ).remote(ipc_handle_dict, mode=mode)
765 + )
766 +
767 + assert inference_result["success"], (
768 + f"IPC weight transfer failed (mode={mode}). "
769 + f"Received shape: {inference_result['received_shape']}, "
770 + f"Received sum: {inference_result['received_sum']}"
771 + )
772 +
773 +
774 +def test_ipc_receive_weights_missing_gpu_uuid_raises():
775 + """Test that receive_weights raises if GPU UUID not found in IPC handles."""
776 + if torch.cuda.device_count() < 1:
777 + pytest.skip("Need at least 1 GPU for this test")
778 +
779 + config = WeightTransferConfig(backend="ipc")
780 + parallel_config = create_mock_parallel_config()
781 + engine = IPCWeightTransferEngine(config, parallel_config)
782 +
783 + # Create IPC handle with wrong GPU UUID
784 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
785 + ipc_handle = reduce_tensor(dummy_tensor)
786 + wrong_uuid = "wrong-uuid-12345"
787 + ipc_handles = [{wrong_uuid: ipc_handle}]
788 +
789 + update_info = IPCWeightTransferUpdateInfo(
790 + names=["w"],
791 + dtype_names=["float32"],
792 + shapes=[[10, 10]],
793 + ipc_handles=ipc_handles,
794 + )
795 +
796 + with pytest.raises(ValueError, match="IPC handle not found"):
797 + engine.receive_weights(update_info, lambda x: None)
tools/pre_commit/check_forbidden_imports.py
+2
@@ -37,6 +37,8 @@ CHECK_IMPORTS = {
37 "vllm/distributed/device_communicators/all_reduce_utils.py",
38 "vllm/distributed/device_communicators/shm_broadcast.py",
39 "vllm/distributed/device_communicators/shm_object_storage.py",
40 + "vllm/distributed/weight_transfer/ipc_engine.py",
41 + "tests/distributed/test_weight_transfer.py",
42 "vllm/utils/hashing.py",
43 "tests/multimodal/media/test_base.py",
44 "tests/tokenizers_/test_hf.py",
vllm/config/weight_transfer.py
+1 -1
@@ -9,5 +9,5 @@ from vllm.config.utils import config
9 class WeightTransferConfig:
10 """Configuration for weight transfer during RL training."""
11
12 - backend: Literal["nccl"] = "nccl"
12 + backend: Literal["nccl", "ipc"] = "nccl"
13 """The backend to use for weight transfer."""
vllm/distributed/weight_transfer/base.py
+28 -1
@@ -3,7 +3,7 @@
3 """Base class for weight transfer engines."""
4
5 from abc import ABC, abstractmethod
6 -from collections.abc import Callable
6 +from collections.abc import Callable, Iterator
7 from dataclasses import KW_ONLY, dataclass, field
8 from typing import Any, Generic, TypeVar
9
@@ -156,3 +156,30 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
156 This should be called when the worker is shutting down.
157 """
158 raise NotImplementedError
159 +
160 + @staticmethod
161 + @abstractmethod
162 + def trainer_send_weights(
163 + iterator: Iterator[tuple[str, torch.Tensor]],
164 + trainer_args: dict[str, Any] | Any,
165 + ) -> None:
166 + """
167 + Send weights from trainer to inference workers.
168 +
169 + This is a static method that can be called from the trainer process
170 + to send weights to all inference workers.
171 +
172 + Args:
173 + iterator: Iterator of model parameters. Returns (name, tensor) tuples.
174 + The tensors should be on the appropriate device for the backend.
175 + trainer_args: Dictionary containing backend-specific arguments needed
176 + to send weights. The structure depends on the backend:
177 + - NCCL: Contains 'group', 'src', 'packed', etc.
178 + - IPC: Contains 'mode' ('http' or 'ray'),
179 + 'llm_handle' (for Ray), 'url' (for HTTP), etc.
180 +
181 + Example:
182 + >>> param_iter = ((n, p) for n, p in model.named_parameters())
183 + >>> engine.trainer_send_weights(param_iter, trainer_args)
184 + """
185 + raise NotImplementedError
vllm/distributed/weight_transfer/factory.py
+6
@@ -114,3 +114,9 @@ WeightTransferEngineFactory.register_engine(
114 "vllm.distributed.weight_transfer.nccl_engine",
115 "NCCLWeightTransferEngine",
116 )
117 +
118 +WeightTransferEngineFactory.register_engine(
119 + "ipc",
120 + "vllm.distributed.weight_transfer.ipc_engine",
121 + "IPCWeightTransferEngine",
122 +)
vllm/distributed/weight_transfer/ipc_engine.py new
+291
@@ -0,0 +1,291 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""IPC-based weight transfer engine using CUDA IPC for communication."""
4 +
5 +import base64
6 +import pickle
7 +from collections.abc import Callable, Iterator
8 +from dataclasses import asdict, dataclass
9 +from typing import Any
10 +
11 +import requests
12 +import torch
13 +from torch.multiprocessing.reductions import reduce_tensor
14 +
15 +from vllm.config.parallel import ParallelConfig
16 +from vllm.config.weight_transfer import WeightTransferConfig
17 +from vllm.distributed.weight_transfer.base import (
18 + WeightTransferEngine,
19 + WeightTransferInitInfo,
20 + WeightTransferUpdateInfo,
21 +)
22 +
23 +
24 +@dataclass
25 +class IPCTrainerSendWeightsArgs:
26 + """Arguments for IPC trainer_send_weights method."""
27 +
28 + mode: str
29 + """Transport mode: 'http' or 'ray'."""
30 + llm_handle: Any = None
31 + """Ray ObjectRef to LLM handle (required for 'ray' mode)."""
32 + url: str | None = None
33 + """Base URL for HTTP endpoint (required for 'http' mode)."""
34 +
35 + def __post_init__(self):
36 + """Validate that required arguments are provided for the selected mode."""
37 + if self.mode == "ray" and self.llm_handle is None:
38 + raise ValueError("llm_handle is required for 'ray' mode")
39 + if self.mode == "http" and self.url is None:
40 + raise ValueError("url is required for 'http' mode")
41 + if self.mode not in ("ray", "http"):
42 + raise ValueError(f"mode must be 'ray' or 'http', got {self.mode}")
43 +
44 +
45 +@dataclass
46 +class IPCWeightTransferInitInfo(WeightTransferInitInfo):
47 + """Initialization info for IPC weight transfer backend. No init needed for IPC."""
48 +
49 + pass
50 +
51 +
52 +@dataclass
53 +class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
54 + """Update info for IPC weight transfer backend.
55 +
56 + Accepts IPC handles either directly via ``ipc_handles`` (Ray transport)
57 + or as a base64-encoded pickle via ``ipc_handles_pickled`` (HTTP transport).
58 + Exactly one of the two must be provided; if ``ipc_handles_pickled`` is set
59 + it is unpickled into ``ipc_handles`` during ``__post_init__``.
60 + """
61 +
62 + names: list[str]
63 + dtype_names: list[str]
64 + shapes: list[list[int]]
65 + ipc_handles: list[dict[str, tuple[Callable, tuple]]] | None = None
66 + """IPC handles mapping physical GPU UUID to (func, args) tuple.
67 + Each handle is a dictionary mapping GPU UUID strings to IPC handle tuples."""
68 + ipc_handles_pickled: str | None = None
69 + """Base64-encoded pickled IPC handles, used for HTTP transport."""
70 +
71 + def __post_init__(self):
72 + if self.ipc_handles_pickled is not None:
73 + if self.ipc_handles is not None:
74 + raise ValueError(
75 + "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
76 + )
77 + self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled))
78 + self.ipc_handles_pickled = None
79 +
80 + if self.ipc_handles is None:
81 + raise ValueError(
82 + "Either `ipc_handles` or `ipc_handles_pickled` must be provided"
83 + )
84 +
85 + num_params = len(self.names)
86 + if len(self.dtype_names) != num_params:
87 + raise ValueError(
88 + f"`dtype_names` should be of the same size as `names`: "
89 + f"got {len(self.dtype_names)} and {len(self.names)}"
90 + )
91 + if len(self.shapes) != num_params:
92 + raise ValueError(
93 + f"`shapes` should be of the same size as `names`: "
94 + f"got {len(self.shapes)} and {len(self.names)}"
95 + )
96 + if len(self.ipc_handles) != num_params:
97 + raise ValueError(
98 + f"`ipc_handles` should be of the same size as `names`: "
99 + f"got {len(self.ipc_handles)} and {len(self.names)}"
100 + )
101 +
102 +
103 +class IPCWeightTransferEngine(
104 + WeightTransferEngine[IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo]
105 +):
106 + """
107 + Weight transfer engine using CUDA IPC for communication between trainer and workers.
108 +
109 + This implementation uses CUDA IPC to transfer weights from the trainer (rank 0)
110 + to all inference workers in a process group. IPC handles are used to share
111 + memory between processes on the same node.
112 + """
113 +
114 + # Define backend-specific dataclass types
115 + init_info_cls = IPCWeightTransferInitInfo
116 + update_info_cls = IPCWeightTransferUpdateInfo
117 +
118 + def __init__(
119 + self, config: WeightTransferConfig, parallel_config: ParallelConfig
120 + ) -> None:
121 + """
122 + Initialize the IPC weight transfer engine.
123 +
124 + Args:
125 + config: The configuration for the weight transfer engine
126 + parallel_config: The configuration for the parallel setup
127 + """
128 + super().__init__(config, parallel_config)
129 +
130 + def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
131 + """
132 + Initialize the weight transfer mechanism.
133 + This is called once at the beginning of training.
134 + No initialization needed for IPC backend.
135 +
136 + Args:
137 + init_info: IPC initialization info (empty)
138 + """
139 + pass
140 +
141 + def receive_weights(
142 + self,
143 + update_info: IPCWeightTransferUpdateInfo,
144 + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
145 + ) -> None:
146 + """
147 + Receive weights from the trainer via CUDA IPC handles.
148 +
149 + Args:
150 + update_info: IPC update info containing parameter names, dtypes, shapes,
151 + and IPC handles. Each IPC handle is a mapping between physical
152 + GPU UUID and the IPC handle tuple (func, args).
153 + load_weights: Callable that loads weights into the model. Called
154 + incrementally for each weight to avoid OOM.
155 + """
156 + assert update_info.ipc_handles is not None
157 + weights = []
158 + for name, _dtype_name, _shape, ipc_handle in zip(
159 + update_info.names,
160 + update_info.dtype_names,
161 + update_info.shapes,
162 + update_info.ipc_handles,
163 + ):
164 + device_index = torch.cuda.current_device()
165 + props = torch.cuda.get_device_properties(device_index)
166 + physical_gpu_id = str(props.uuid)
167 +
168 + if physical_gpu_id not in ipc_handle:
169 + raise ValueError(
170 + f"IPC handle not found for GPU UUID {physical_gpu_id}. "
171 + f"Available UUIDs: {list(ipc_handle.keys())}"
172 + )
173 +
174 + handle = ipc_handle[physical_gpu_id]
175 +
176 + func, args = handle
177 + list_args = list(args) # type: ignore
178 + # Index 6 is the device_index parameter in torch's
179 + # IPC handle tuple (rebuild_cuda_tensor). Update it
180 + # to the current device since the logical index can
181 + # differ between sender and receiver.
182 + list_args[6] = device_index
183 + weight = func(*list_args) # type: ignore
184 + weights.append((name, weight))
185 +
186 + load_weights(weights)
187 +
188 + def shutdown(self) -> None:
189 + """
190 + Shutdown the weight transfer engine.
191 + """
192 + pass
193 +
194 + @staticmethod
195 + def trainer_send_weights(
196 + iterator: Iterator[tuple[str, torch.Tensor]],
197 + trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs,
198 + ) -> None:
199 + """
200 + Send weights from trainer to inference workers via CUDA IPC.
201 +
202 + Supports two modes:
203 + - 'ray': Sends weights via Ray RPC to a Ray-based LLM handle
204 + - 'http': Sends weights via HTTP POST to a vLLM HTTP server
205 +
206 + Args:
207 + iterator: Iterator of model parameters. Returns (name, tensor) tuples.
208 + Tensors should be on the same GPU as the inference workers.
209 + trainer_args: Dictionary containing IPC-specific arguments.
210 + Should contain keys from IPCTrainerSendWeightsArgs:
211 + - mode: 'ray' or 'http'
212 + - llm_handle: Ray ObjectRef (for 'ray' mode)
213 + - url: Base URL string (for 'http' mode)
214 +
215 + Example (Ray mode):
216 + >>> from vllm.distributed.weight_transfer.ipc_engine import (
217 + ... IPCWeightTransferEngine,
218 + ... IPCTrainerSendWeightsArgs,
219 + ... )
220 + >>> param_iter = ((n, p) for n, p in model.named_parameters())
221 + >>> args = IPCTrainerSendWeightsArgs(mode="ray", llm_handle=llm_handle)
222 + >>> IPCWeightTransferEngine.trainer_send_weights(param_iter, asdict(args))
223 +
224 + Example (HTTP mode):
225 + >>> args = IPCTrainerSendWeightsArgs(
226 + ... mode="http", url="http://localhost:8000"
227 + ... )
228 + >>> IPCWeightTransferEngine.trainer_send_weights(param_iter, asdict(args))
229 + """
230 + # Parse trainer args - accept either dict or dataclass instance
231 + if isinstance(trainer_args, dict):
232 + args = IPCTrainerSendWeightsArgs(**trainer_args)
233 + else:
234 + args = trainer_args
235 +
236 + # Get physical GPU UUID
237 + device_index = torch.cuda.current_device()
238 + props = torch.cuda.get_device_properties(device_index)
239 + gpu_uuid = str(props.uuid)
240 +
241 + # Collect weight metadata and create IPC handles
242 + names = []
243 + dtype_names = []
244 + shapes = []
245 + ipc_handles = []
246 +
247 + for name, tensor in iterator:
248 + names.append(name)
249 + dtype_names.append(str(tensor.dtype).split(".")[-1])
250 + shapes.append(list(tensor.shape))
251 +
252 + # Create IPC handle for this weight tensor
253 + # The tensor must remain in memory for IPC to work
254 + weight = tensor.detach().contiguous()
255 + ipc_handle = reduce_tensor(weight)
256 + ipc_handles.append({gpu_uuid: ipc_handle})
257 +
258 + # Send weights based on mode
259 + if args.mode == "ray":
260 + # Ray mode: send via Ray RPC
261 + import ray
262 +
263 + update_info = asdict(
264 + IPCWeightTransferUpdateInfo(
265 + names=names,
266 + dtype_names=dtype_names,
267 + shapes=shapes,
268 + ipc_handles=ipc_handles,
269 + )
270 + )
271 + ray.get(
272 + args.llm_handle.update_weights.remote(dict(update_info=update_info))
273 + )
274 + elif args.mode == "http":
275 + # HTTP mode: send via HTTP POST with pickled handles
276 + # Pickle and base64 encode IPC handles for HTTP transmission
277 + pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode(
278 + "utf-8"
279 + )
280 +
281 + url = f"{args.url}/update_weights"
282 + payload = {
283 + "update_info": {
284 + "names": names,
285 + "dtype_names": dtype_names,
286 + "shapes": shapes,
287 + "ipc_handles_pickled": pickled_handles,
288 + }
289 + }
290 + response = requests.post(url, json=payload, timeout=300)
291 + response.raise_for_status()
vllm/distributed/weight_transfer/nccl_engine.py
+52 -33
@@ -35,6 +35,32 @@ class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
35 world_size: int
36
37
38 +@dataclass
39 +class NCCLTrainerSendWeightsArgs:
40 + """Arguments for NCCL trainer_send_weights method."""
41 +
42 + group: Any
43 + """Process group (PyNcclCommunicator) for NCCL communication."""
44 + src: int = 0
45 + """Source rank (default 0, trainer is typically rank 0)."""
46 + post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor] | None = None
47 + """Optional function to apply to each (name, tensor) pair before broadcasting.
48 + If None, extracts just the tensor."""
49 + packed: bool = False
50 + """Whether to use packed tensor broadcasting for efficiency.
51 + When True, multiple tensors are batched together before broadcasting
52 + to reduce NCCL communication overhead."""
53 + stream: torch.cuda.Stream | None = None
54 + """CUDA stream to use for broadcasting if packed is False.
55 + If packed is True, new streams will be created for each buffer."""
56 + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
57 + """Size in bytes for each packed tensor buffer.
58 + Must match the value used in NCCLWeightTransferUpdateInfo."""
59 + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
60 + """Number of buffers for double/triple buffering during packed transfer.
61 + Must match the value used in NCCLWeightTransferUpdateInfo."""
62 +
63 +
64 @dataclass
65 class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
66 """Update info for NCCL weight transfer backend."""
@@ -47,7 +73,7 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
73 When True, multiple tensors are batched together before broadcasting
74 to reduce NCCL communication overhead."""
75 packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
50 - """Size in bytes for each packed tensor buffer. Default is 1GB.
76 + """Size in bytes for each packed tensor buffer.
77 Both producer and consumer must use the same value."""
78 packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
79 """Number of buffers for double/triple buffering during packed transfer.
@@ -186,47 +212,38 @@ class NCCLWeightTransferEngine(
212 @staticmethod
213 def trainer_send_weights(
214 iterator: Iterator[tuple[str, torch.Tensor]],
189 - group: Any,
190 - src: int = 0,
191 - post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor]
192 - | None = None,
193 - packed: bool = False,
194 - stream: torch.cuda.Stream | None = None,
195 - packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
196 - packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
215 + trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
216 ) -> None:
217 """Broadcast weights from trainer to vLLM workers.
218
219 Args:
220 iterator: Iterator of model parameters. Returns (name, tensor) tuples
202 - group: Process group (PyNcclCommunicator)
203 - src: Source rank (default 0, trainer is typically rank 0)
204 - post_iter_func: Optional function to apply to each (name, tensor) pair
205 - before broadcasting. If None, extracts just the tensor.
206 - packed: Whether to use packed tensor broadcasting for efficiency.
207 - When True, multiple tensors are batched together before
208 - broadcasting to reduce NCCL communication overhead.
209 - stream: CUDA stream to use for broadcasting if packed is False.
210 - If packed is True, new streams will be created for each buffer.
211 - packed_buffer_size_bytes: Size in bytes for each packed tensor buffer.
212 - Must match the value used in NCCLWeightTransferUpdateInfo.
213 - packed_num_buffers: Number of buffers for double/triple buffering.
214 - Must match the value used in NCCLWeightTransferUpdateInfo.
221 + trainer_args: Dictionary or NCCLTrainerSendWeightsArgs instance containing
222 + NCCL-specific arguments. If a dict, should contain keys from
223 + NCCLTrainerSendWeightsArgs.
224
225 Example:
226 >>> from vllm.distributed.weight_transfer.nccl_engine import (
227 ... NCCLWeightTransferEngine,
228 + ... NCCLTrainerSendWeightsArgs,
229 ... )
230 >>> param_iter = ((n, p) for n, p in model.named_parameters())
221 - >>> NCCLWeightTransferEngine.trainer_send_weights(
222 - ... param_iter, group, packed=True
223 - ... )
231 + >>> args = NCCLTrainerSendWeightsArgs(group=group, packed=True)
232 + >>> NCCLWeightTransferEngine.trainer_send_weights(param_iter, args)
233 """
225 - if post_iter_func is None:
234 + # Parse trainer args - accept either dict or dataclass instance
235 + if isinstance(trainer_args, dict):
236 + args = NCCLTrainerSendWeightsArgs(**trainer_args)
237 + else:
238 + args = trainer_args
239 +
240 + if args.post_iter_func is None:
241 # Default: extract just the tensor from (name, tensor) tuple
242 post_iter_func = lambda x: x[1]
243 + else:
244 + post_iter_func = args.post_iter_func
245
229 - if packed:
246 + if args.packed:
247 # Use packed tensor broadcasting for efficiency
248 from vllm.distributed.weight_transfer.packed_tensor import (
249 packed_broadcast_producer,
@@ -234,18 +251,20 @@ class NCCLWeightTransferEngine(
251
252 packed_broadcast_producer(
253 iterator=iterator,
237 - group=group,
238 - src=src,
254 + group=args.group,
255 + src=args.src,
256 post_iter_func=post_iter_func,
240 - buffer_size_bytes=packed_buffer_size_bytes,
241 - num_buffers=packed_num_buffers,
257 + buffer_size_bytes=args.packed_buffer_size_bytes,
258 + num_buffers=args.packed_num_buffers,
259 )
260 else:
261 # Use simple one-by-one broadcasting
262 for item in iterator:
263 tensor = post_iter_func(item)
247 - group.broadcast(
248 - tensor, src=src, stream=stream or torch.cuda.current_stream()
264 + args.group.broadcast(
265 + tensor,
266 + src=args.src,
267 + stream=args.stream or torch.cuda.current_stream(),
268 )
269
270 @staticmethod