@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c1858b7ec8
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 8, 20:15 UTC

[Feat][RL][1/2] Native Weight Syncing API: NCCL (#31943)

Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Co-authored-by: SumanthRH <sumanthrh99@gmail.com>

Aaron Hao committed Feb 5, 2026 at 09:13 UTC c1858b7ec8aa571dc0c0e00aded01019cca6a7e6
27 files changed +2974 -2
.buildkite/test-amd.yaml
+7
@@ -233,6 +233,7 @@ steps:
233 - tests/compile/fullgraph/test_basic_correctness.py
234 - examples/offline_inference/rlhf.py
235 - examples/offline_inference/rlhf_colocate.py
236 + - examples/offline_inference/new_weight_syncing/
237 - tests/examples/offline_inference/data_parallel.py
238 - tests/v1/distributed
239 - tests/v1/engine/test_engine_core_client.py
@@ -268,10 +269,16 @@ steps:
269 - pytest -v -s distributed/test_symm_mem_allreduce.py
270 # TODO: create a dedicated test section for multi-GPU example tests
271 # when we have multiple distributed example tests
272 + # OLD rlhf examples
273 - pushd ../examples/offline_inference
274 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
275 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
276 - popd
277 + # NEW rlhf examples
278 + - pushd ../examples/offline_inference/new_weight_syncing
279 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
280 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
281 + - popd
282
283 - label: Distributed Tests (8 GPUs) # 4min
284 timeout_in_minutes: 10
.buildkite/test-pipeline.yaml
+9
@@ -206,6 +206,7 @@ steps:
206 - tests/compile/fullgraph/test_basic_correctness.py
207 - examples/offline_inference/rlhf.py
208 - examples/offline_inference/rlhf_colocate.py
209 + - examples/offline_inference/new_weight_syncing/
210 - tests/examples/offline_inference/data_parallel.py
211 - tests/v1/distributed
212 - tests/v1/engine/test_engine_core_client.py
@@ -240,10 +241,16 @@ steps:
241 - pytest -v -s distributed/test_symm_mem_allreduce.py
242 # TODO: create a dedicated test section for multi-GPU example tests
243 # when we have multiple distributed example tests
244 + # OLD rlhf examples
245 - pushd ../examples/offline_inference
246 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
247 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
248 - popd
249 + # NEW rlhf examples
250 + - pushd ../examples/offline_inference/new_weight_syncing
251 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
252 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
253 + - popd
254
255 - label: Distributed Tests (8 GPUs) # 4min
256 timeout_in_minutes: 10
@@ -1146,6 +1153,8 @@ steps:
1153 - pytest -v -s distributed/test_shm_broadcast.py
1154 - pytest -v -s distributed/test_shm_buffer.py
1155 - pytest -v -s distributed/test_shm_storage.py
1156 + - pytest -v -s distributed/test_packed_tensor.py
1157 + - pytest -v -s distributed/test_weight_transfer.py
1158
1159 - label: 2 Node Tests (4 GPUs in total) # 16min
1160 timeout_in_minutes: 30
.buildkite/test_areas/distributed.yaml
+6
@@ -62,6 +62,7 @@ steps:
62 - tests/compile/fullgraph/test_basic_correctness.py
63 - examples/offline_inference/rlhf.py
64 - examples/offline_inference/rlhf_colocate.py
65 + - examples/offline_inference/new_weight_syncing/
66 - tests/examples/offline_inference/data_parallel.py
67 - tests/v1/distributed
68 - tests/v1/engine/test_engine_core_client.py
@@ -96,9 +97,14 @@ steps:
97 - pytest -v -s distributed/test_symm_mem_allreduce.py
98 # TODO: create a dedicated test section for multi-GPU example tests
99 # when we have multiple distributed example tests
100 + # OLD rlhf examples
101 - cd ../examples/offline_inference
102 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
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
107 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
108
109 - label: Distributed Tests (8 GPUs)(H100)
110 timeout_in_minutes: 10
examples/offline_inference/new_weight_syncing/rlhf.py new
+208
@@ -0,0 +1,208 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Demonstrates reinforcement learning using vLLM and Ray,
5 +with native weight syncing APIs at engine instance.
6 +
7 +The script separates training and inference workloads onto distinct GPUs
8 +so that Ray can manage process placement and inter-process communication.
9 +A Hugging Face Transformer model occupies one GPU for training, whereas a
10 +2x tensor-parallel vLLM inference engine occupies two GPUs.
11 +
12 +The example performs the following steps:
13 +* Load the training model on one gpu (scheduled via ray)
14 +* Initialize the inference model with dummy weights across
15 + two gpus using vLLM's tensor parallelism and Ray placement groups.
16 +* Generate gibberish from a list of prompts using the randomly initialized
17 + inference engine.
18 +* Update the weights of the training model and broadcast the updated weights
19 + to the inference engine by using a Ray collective RPC group.
20 +* Generating from the list of prompts after weight sync should result
21 + in sensible outputs.
22 +
23 +This example assumes a single-node cluster with three GPUs, but Ray
24 +supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
25 +workloads. Residual GPU activity interferes with vLLM memory profiling and
26 +causes unexpected behavior.
27 +"""
28 +
29 +import os
30 +
31 +import ray
32 +from ray.util.placement_group import placement_group
33 +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
34 +from transformers import AutoModelForCausalLM
35 +
36 +from vllm import LLM, SamplingParams
37 +from vllm.config import WeightTransferConfig
38 +from vllm.distributed.weight_transfer.nccl_engine import (
39 + NCCLWeightTransferEngine,
40 +)
41 +from vllm.utils.network_utils import get_ip, get_open_port
42 +
43 +MODEL_NAME = "facebook/opt-125m"
44 +# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128"
45 +
46 +
47 +class MyLLM(LLM):
48 + """Configure the vLLM worker for Ray placement group execution."""
49 +
50 + def __init__(self, *args, **kwargs):
51 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
52 + super().__init__(*args, **kwargs)
53 +
54 +
55 +@ray.remote(num_gpus=1)
56 +class TrainModel:
57 + """Ray actor that wraps the training model on a dedicated GPU."""
58 +
59 + def __init__(self, model_name: str):
60 + self.model = AutoModelForCausalLM.from_pretrained(
61 + model_name,
62 + ).to("cuda:0")
63 +
64 + self.port = get_open_port()
65 + self.master_address = get_ip()
66 +
67 + def get_master_address_and_port(self):
68 + return self.master_address, self.port
69 +
70 + def get_weight_metadata(self):
71 + """Return weight names, dtypes, and shapes for weight transfer."""
72 + names = []
73 + dtype_names = []
74 + shapes = []
75 + for name, p in self.model.named_parameters():
76 + names.append(name)
77 + dtype_names.append(str(p.dtype).split(".")[-1])
78 + shapes.append(list(p.shape))
79 + return names, dtype_names, shapes
80 +
81 + def init_weight_transfer_group(self, world_size):
82 + """Initialize the NCCL process group for weight transfer."""
83 + self.model_update_group = NCCLWeightTransferEngine.trainer_init(
84 + dict(
85 + master_address=self.master_address,
86 + master_port=self.port,
87 + world_size=world_size,
88 + ),
89 + )
90 +
91 + def broadcast_weights(self, packed: bool = True):
92 + """Broadcast weights to the inference engine."""
93 + NCCLWeightTransferEngine.trainer_send_weights(
94 + iterator=self.model.named_parameters(),
95 + group=self.model_update_group,
96 + packed=packed,
97 + )
98 +
99 +
100 +# Initialize Ray and set the visible devices. The vLLM engine will
101 +# be placed on GPUs 1 and 2.
102 +ray.init()
103 +
104 +# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
105 +# Learn more about Ray placement groups:
106 +# https://docs.ray.io/en/latest/placement-groups.html
107 +# Launch the training model actor. Ray's resource scheduler will allocate
108 +# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
109 +train_model = TrainModel.remote(MODEL_NAME)
110 +
111 +pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
112 +ray.get(pg_inference.ready())
113 +scheduling_inference = PlacementGroupSchedulingStrategy(
114 + placement_group=pg_inference,
115 + placement_group_capture_child_tasks=True,
116 + placement_group_bundle_index=0,
117 +)
118 +
119 +# Launch the vLLM inference engine. The `enforce_eager` flag reduces
120 +# start-up latency.
121 +# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
122 +# are now native to vLLM workers.
123 +llm = ray.remote(
124 + num_cpus=0,
125 + num_gpus=0,
126 + scheduling_strategy=scheduling_inference,
127 +)(MyLLM).remote(
128 + model=MODEL_NAME,
129 + enforce_eager=True,
130 + tensor_parallel_size=2,
131 + data_parallel_size=1,
132 + distributed_executor_backend="ray",
133 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
134 + load_format="dummy",
135 + quantization="fp8",
136 +)
137 +
138 +# Generate text from the prompts.
139 +prompts = [
140 + "Hello, my name is",
141 + "The president of the United States is",
142 + "The capital of France is",
143 + "The future of AI is",
144 +]
145 +
146 +sampling_params = SamplingParams(temperature=0)
147 +
148 +outputs = ray.get(llm.generate.remote(prompts, sampling_params))
149 +
150 +# Generate text with the initial model. The output is expected to be nonsense
151 +# because the weights are randomly initialized.
152 +print("-" * 50)
153 +for output in outputs:
154 + prompt = output.prompt
155 + generated_text = output.outputs[0].text
156 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
157 + print("-" * 50)
158 +
159 +# Set up the communication channel between the training process and the
160 +# inference engine.
161 +master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
162 +
163 +world_size = ray.get(llm.get_world_size.remote()) + 1 # +1 for the trainer
164 +inference_handle = llm.init_weight_transfer_engine.remote(
165 + dict(
166 + init_info=dict(
167 + master_address=master_address,
168 + master_port=master_port,
169 + rank_offset=1,
170 + world_size=world_size,
171 + )
172 + )
173 +)
174 +
175 +# Initialize weight transfer group on both the training actor and inference engine
176 +train_handle = train_model.init_weight_transfer_group.remote(world_size)
177 +ray.get([train_handle, inference_handle])
178 +
179 +# Synchronize the updated weights to the inference engine using batched API.
180 +# Collect all weight metadata from the training actor
181 +names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
182 +
183 +# Issue update_weights call with NCCL-specific update info
184 +# packed=True enables efficient batched tensor broadcasting
185 +inference_handle = llm.update_weights.remote(
186 + dict(
187 + update_info=dict(
188 + names=names,
189 + dtype_names=dtype_names,
190 + shapes=shapes,
191 + packed=True,
192 + )
193 + )
194 +)
195 +
196 +# Broadcast all weights from trainer using the weight transfer API
197 +train_handle = train_model.broadcast_weights.remote(packed=True)
198 +ray.get([train_handle, inference_handle])
199 +
200 +# Generate text with the updated model. The output is expected to be normal
201 +# because the weights are updated.
202 +outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
203 +print("-" * 50)
204 +for output in outputs_updated:
205 + prompt = output.prompt
206 + generated_text = output.outputs[0].text
207 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
208 + print("-" * 50)
examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py new
+283
@@ -0,0 +1,283 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Demonstrates async reinforcement learning using vLLM and Ray,
5 +with native weight syncing APIs at engine instance.
6 +
7 +The script separates training and inference workloads onto distinct GPUs
8 +so that Ray can manage process placement and inter-process communication.
9 +A Hugging Face Transformer model occupies one GPU for training, whereas a
10 +2x tensor-parallel vLLM inference engine occupies two GPUs.
11 +
12 +The example performs the following steps:
13 +* Load the training model on one gpu (scheduled via ray)
14 +* Initialize the inference model with dummy weights across
15 + two gpus using vLLM's tensor parallelism and Ray placement groups.
16 +* Generate gibberish from a list of prompts using the randomly initialized
17 + inference engine.
18 +* Pause generation once generation completes for one sequence
19 +* Update the weights of the training model and broadcast the updated weights
20 + to the inference engine by using a Ray collective RPC group.
21 +* Resume generation and print out the results
22 +
23 +This example assumes a single-node cluster with three GPUs, but Ray
24 +supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
25 +workloads. Residual GPU activity interferes with vLLM memory profiling and
26 +causes unexpected behavior.
27 +"""
28 +
29 +import os
30 +import uuid
31 +from dataclasses import asdict
32 +
33 +import ray
34 +import torch
35 +from ray.util.placement_group import placement_group
36 +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
37 +from transformers import AutoModelForCausalLM, AutoTokenizer
38 +
39 +import vllm
40 +from vllm import SamplingParams
41 +from vllm.config import WeightTransferConfig
42 +from vllm.distributed.weight_transfer.base import (
43 + WeightTransferInitRequest,
44 + WeightTransferUpdateRequest,
45 +)
46 +from vllm.distributed.weight_transfer.nccl_engine import (
47 + NCCLWeightTransferEngine,
48 + NCCLWeightTransferInitInfo,
49 + NCCLWeightTransferUpdateInfo,
50 +)
51 +from vllm.utils.network_utils import get_ip, get_open_port
52 +from vllm.v1.executor import Executor
53 +
54 +MODEL_NAME = "facebook/opt-125m"
55 +
56 +
57 +class MyLLM(vllm.AsyncLLMEngine):
58 + """Configure the vLLM worker for Ray placement group execution."""
59 +
60 + def __init__(self, **kwargs):
61 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
62 + engine_args = vllm.AsyncEngineArgs(**kwargs)
63 + vllm_config = engine_args.create_engine_config()
64 + executor_class = Executor.get_class(vllm_config)
65 + super().__init__(
66 + vllm_config=vllm_config,
67 + executor_class=executor_class,
68 + log_requests=engine_args.enable_log_requests,
69 + log_stats=not engine_args.disable_log_stats,
70 + )
71 +
72 + async def generate_with_retry(
73 + self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
74 + ) -> vllm.RequestOutput:
75 + finish_reason = "abort"
76 + while finish_reason == "abort":
77 + async for request_output in self.generate(
78 + {"prompt_token_ids": prompt_token_ids},
79 + sampling_params,
80 + request_id=str(uuid.uuid4()),
81 + ):
82 + output = request_output
83 + finish_reason = output.outputs[0].finish_reason
84 + if finish_reason == "abort":
85 + print(
86 + f"ABORT, prompt_token_ids: {prompt_token_ids}, "
87 + f"generated token_ids: {list(output.outputs[0].token_ids)}"
88 + )
89 + prompt_token_ids = prompt_token_ids + list(output.outputs[0].token_ids)
90 + return output
91 +
92 +
93 +@ray.remote(num_gpus=1)
94 +class TrainModel:
95 + """Ray actor that wraps the training model on a dedicated GPU."""
96 +
97 + def __init__(self, model_name: str):
98 + self.model = AutoModelForCausalLM.from_pretrained(
99 + model_name, dtype=torch.bfloat16
100 + ).to("cuda:0")
101 + self.port = get_open_port()
102 + self.master_address = get_ip()
103 +
104 + def get_master_address_and_port(self):
105 + return self.master_address, self.port
106 +
107 + def get_weight_metadata(self):
108 + """Return weight names, dtypes, and shapes for weight transfer."""
109 + names = []
110 + dtype_names = []
111 + shapes = []
112 + for name, p in self.model.named_parameters():
113 + names.append(name)
114 + dtype_names.append(str(p.dtype).split(".")[-1])
115 + shapes.append(list(p.shape))
116 + return names, dtype_names, shapes
117 +
118 + def init_weight_transfer_group(self, world_size):
119 + """Initialize the NCCL process group for weight transfer."""
120 + self.model_update_group = NCCLWeightTransferEngine.trainer_init(
121 + dict(
122 + master_address=self.master_address,
123 + master_port=self.port,
124 + world_size=world_size,
125 + ),
126 + )
127 +
128 + def broadcast_weights(self, packed: bool = True):
129 + """Broadcast weights to the inference engine."""
130 + NCCLWeightTransferEngine.trainer_send_weights(
131 + iterator=self.model.named_parameters(),
132 + group=self.model_update_group,
133 + packed=packed,
134 + )
135 +
136 +
137 +# Initialize Ray and set the visible devices. The vLLM engine will
138 +# be placed on GPUs 1 and 2.
139 +ray.init()
140 +
141 +# Launch the training model actor. Ray's resource scheduler will allocate
142 +# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
143 +train_model = TrainModel.remote(MODEL_NAME)
144 +
145 +# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
146 +# Learn more about Ray placement groups:
147 +# https://docs.ray.io/en/latest/placement-groups.html
148 +
149 +pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
150 +ray.get(pg_inference.ready())
151 +scheduling_inference = PlacementGroupSchedulingStrategy(
152 + placement_group=pg_inference,
153 + placement_group_capture_child_tasks=True,
154 + placement_group_bundle_index=0,
155 +)
156 +
157 +# Launch the vLLM inference engine. The `enforce_eager` flag reduces
158 +# start-up latency.
159 +# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
160 +# are now native to vLLM workers.
161 +llm = ray.remote(
162 + num_cpus=0,
163 + num_gpus=0,
164 + scheduling_strategy=scheduling_inference,
165 +)(MyLLM).remote(
166 + model=MODEL_NAME,
167 + enforce_eager=True,
168 + tensor_parallel_size=2,
169 + distributed_executor_backend="ray",
170 + load_format="dummy",
171 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
172 +)
173 +
174 +# Generate text from the prompts.
175 +prompts = [
176 + "My name is",
177 + "The president of the United States is",
178 + "The capital of France is",
179 + "The future of AI is",
180 +]
181 +
182 +# Tokenize prompts to token IDs
183 +tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
184 +prompt_token_ids_list = [
185 + tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts
186 +]
187 +
188 +sampling_params = [
189 + SamplingParams(temperature=0, max_tokens=2),
190 + SamplingParams(temperature=0, max_tokens=32),
191 + SamplingParams(temperature=0, max_tokens=32),
192 + SamplingParams(temperature=0, max_tokens=32),
193 +]
194 +
195 +# Set up the communication channel between the training process and the
196 +# inference engine.
197 +master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
198 +
199 +world_size = 3 # 1 trainer + 2 inference workers (tensor_parallel_size=2)
200 +inference_handle = llm.init_weight_transfer_engine.remote(
201 + WeightTransferInitRequest(
202 + init_info=asdict(
203 + NCCLWeightTransferInitInfo(
204 + master_address=master_address,
205 + master_port=master_port,
206 + rank_offset=1,
207 + world_size=world_size,
208 + )
209 + )
210 + )
211 +)
212 +
213 +# Initialize weight transfer group on both the training actor and inference engine
214 +train_handle = train_model.init_weight_transfer_group.remote(world_size)
215 +ray.get([train_handle, inference_handle])
216 +
217 +
218 +generation_futures = [
219 + llm.generate_with_retry.remote(prompt_token_ids, params)
220 + for prompt_token_ids, params in zip(prompt_token_ids_list, sampling_params)
221 +]
222 +
223 +finished, pending = ray.wait(generation_futures, num_returns=1)
224 +
225 +# Pause generation in preparation for weight sync
226 +ray.get(llm.pause_generation.remote(wait_for_inflight_requests=False))
227 +
228 +# Synchronize the updated weights to the inference engine using batched API.
229 +# Collect all weight metadata from the training actor
230 +names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
231 +
232 +# Issue update_weights call with NCCL-specific update info
233 +# packed=True enables efficient batched tensor broadcasting
234 +inference_handle = llm.update_weights.remote(
235 + WeightTransferUpdateRequest(
236 + update_info=asdict(
237 + NCCLWeightTransferUpdateInfo(
238 + names=names,
239 + dtype_names=dtype_names,
240 + shapes=shapes,
241 + packed=True,
242 + )
243 + )
244 + )
245 +)
246 +
247 +# Broadcast all weights from trainer using the weight transfer API
248 +train_handle = train_model.broadcast_weights.remote(packed=True)
249 +ray.get([train_handle, inference_handle])
250 +
251 +# Resume generation since weight sync is complete
252 +ray.get(llm.resume_generation.remote())
253 +
254 +# Get outputs separately - finished completed before pause, pending were paused/resumed
255 +finished_outputs = ray.get(finished)
256 +pending_outputs = ray.get(pending)
257 +
258 +# Requests that finished before the pause: all generation used original weights
259 +print("-" * 50)
260 +print("Requests that completed BEFORE weight change:")
261 +print("-" * 50)
262 +for output in finished_outputs:
263 + prompt_text = tokenizer.decode(output.prompt_token_ids)
264 + print(f"Prompt: {prompt_text!r}")
265 + print(f"Generated (with original weights): {output.outputs[0].text!r}")
266 + print("-" * 50)
267 +
268 +# Requests that were paused mid-generation: some text before, some after weight change
269 +print("Requests that were PAUSED and RESUMED after weight change:")
270 +print("-" * 50)
271 +for output in pending_outputs:
272 + # Decode the full prompt token IDs (original + generated before pause)
273 + full_prompt_text = tokenizer.decode(output.prompt_token_ids)
274 + # Find the original prompt by checking which one this output started with
275 + original_prompt = next(p for p in prompts if full_prompt_text.startswith(p))
276 + # output.prompt_token_ids contains original prompt + tokens generated before pause
277 + # output.outputs[0].text is what was generated after resuming with new weights
278 + text_before_pause = full_prompt_text[len(original_prompt) :]
279 + text_after_pause = output.outputs[0].text
280 + print(f"Original prompt: {original_prompt!r}")
281 + print(f"Generated before weight change: {text_before_pause!r}")
282 + print(f"Generated after weight change: {text_after_pause!r}")
283 + print("-" * 50)
examples/online_serving/rlhf_http.py new
+241
@@ -0,0 +1,241 @@
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 native weight syncing APIs.
6 +
7 +Unlike rlhf.py which creates a vLLM instance programmatically, this script
8 +assumes you have already started a vLLM server using `vllm serve`. It uses:
9 +- OpenAI-compatible API for inference requests
10 +- HTTP endpoints for weight transfer control plane
11 +- NCCL for actual weight data transfer
12 +
13 +Prerequisites:
14 + Start a vLLM server with weight transfer enabled:
15 +
16 + $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
17 + --enforce-eager \
18 + --weight-transfer-config '{"backend": "nccl"}' \
19 + --load-format dummy
20 +
21 + Then run this script:
22 +
23 + $ python rlhf_http.py
24 +
25 +The example performs the following steps:
26 +
27 +* Load the training model on GPU 0.
28 +* Generate text using the vLLM server via OpenAI-compatible API. The output
29 + is expected to be nonsense because the server is initialized with dummy weights.
30 +* Initialize weight transfer via HTTP endpoint.
31 +* Broadcast the real weights from the training model to the vLLM server
32 + using NCCL.
33 +* Generate text again to show normal output after the weight update.
34 +"""
35 +
36 +import requests
37 +import torch
38 +from openai import OpenAI
39 +from transformers import AutoModelForCausalLM
40 +
41 +from vllm.distributed.weight_transfer.nccl_engine import (
42 + NCCLWeightTransferEngine,
43 +)
44 +from vllm.utils.network_utils import get_ip, get_open_port
45 +
46 +BASE_URL = "http://localhost:8000"
47 +MODEL_NAME = "facebook/opt-125m"
48 +
49 +
50 +def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
51 + """Generate completions using the OpenAI-compatible API."""
52 + results = []
53 + for prompt in prompts:
54 + response = client.completions.create(
55 + model=model,
56 + prompt=prompt,
57 + max_tokens=32,
58 + temperature=0,
59 + )
60 + results.append(response.choices[0].text)
61 + return results
62 +
63 +
64 +def init_weight_transfer_engine(
65 + base_url: str,
66 + master_address: str,
67 + master_port: int,
68 + rank_offset: int,
69 + world_size: int,
70 +) -> None:
71 + """Initialize weight transfer via HTTP endpoint."""
72 + url = f"{base_url}/init_weight_transfer_engine"
73 + payload = {
74 + "init_info": dict(
75 + master_address=master_address,
76 + master_port=master_port,
77 + rank_offset=rank_offset,
78 + world_size=world_size,
79 + )
80 + }
81 + response = requests.post(url, json=payload, timeout=60)
82 + response.raise_for_status()
83 +
84 +
85 +def update_weights(
86 + base_url: str,
87 + names: list[str],
88 + dtype_names: list[str],
89 + shapes: list[list[int]],
90 + packed: bool = False,
91 +) -> None:
92 + """Update weights via HTTP endpoint."""
93 + url = f"{base_url}/update_weights"
94 + payload = {
95 + "update_info": dict(
96 + names=names,
97 + dtype_names=dtype_names,
98 + shapes=shapes,
99 + packed=packed,
100 + )
101 + }
102 + response = requests.post(url, json=payload, timeout=300)
103 + response.raise_for_status()
104 +
105 +
106 +def pause_generation(base_url: str) -> None:
107 + """Pause generation via HTTP endpoint."""
108 + url = f"{base_url}/pause"
109 + response = requests.post(url, timeout=60)
110 + response.raise_for_status()
111 +
112 +
113 +def resume_generation(base_url: str) -> None:
114 + """Resume generation via HTTP endpoint."""
115 + url = f"{base_url}/resume"
116 + response = requests.post(url, timeout=60)
117 + response.raise_for_status()
118 +
119 +
120 +def get_world_size(base_url: str) -> int:
121 + """Get world size from the vLLM server."""
122 + url = f"{base_url}/get_world_size"
123 + response = requests.get(url, timeout=10)
124 + response.raise_for_status()
125 + return response.json()["world_size"]
126 +
127 +
128 +def main():
129 + # Get the inference world size from the vLLM server
130 + inference_world_size = get_world_size(BASE_URL)
131 + world_size = inference_world_size + 1 # +1 for the trainer
132 + device = f"cuda:{inference_world_size}"
133 + torch.cuda.set_device(device)
134 +
135 + # Load the training model
136 + print(f"Loading training model: {MODEL_NAME}")
137 + train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
138 + train_model.to(device)
139 +
140 + # Create OpenAI client pointing to the vLLM server
141 + client = OpenAI(
142 + base_url=f"{BASE_URL}/v1",
143 + api_key="EMPTY", # vLLM doesn't require an API key by default
144 + )
145 +
146 + # Test prompts
147 + prompts = [
148 + "Hello, my name is",
149 + "The president of the United States is",
150 + "The capital of France is",
151 + "The future of AI is",
152 + ]
153 +
154 + # Generate text before weight update. The output is expected to be nonsense
155 + # because the server is initialized with dummy weights.
156 + print("-" * 50)
157 + print("Generating text BEFORE weight update (expect nonsense):")
158 + print("-" * 50)
159 + outputs = generate_completions(client, MODEL_NAME, prompts)
160 + for prompt, generated_text in zip(prompts, outputs):
161 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
162 + print("-" * 50)
163 +
164 + # Set up the communication channel between the training process and the
165 + # vLLM server. The trainer is rank 0, vLLM worker(s) start at rank_offset.
166 + master_address = get_ip()
167 + master_port = get_open_port()
168 + rank_offset = 1
169 +
170 + print(f"Initializing weight transfer: master={master_address}:{master_port}")
171 +
172 + # Initialize weight transfer on vLLM server (this is async, server will
173 + # wait for NCCL connection)
174 + import threading
175 +
176 + init_thread = threading.Thread(
177 + target=init_weight_transfer_engine,
178 + args=(BASE_URL, master_address, master_port, rank_offset, world_size),
179 + )
180 + init_thread.start()
181 +
182 + # Initialize NCCL process group on trainer side
183 + model_update_group = NCCLWeightTransferEngine.trainer_init(
184 + dict(
185 + master_address=master_address,
186 + master_port=master_port,
187 + world_size=world_size,
188 + ),
189 + )
190 +
191 + # Wait for init_weight_transfer_engine to complete
192 + init_thread.join()
193 +
194 + # Pause generation before weight sync
195 + pause_generation(BASE_URL)
196 +
197 + # Collect weight metadata for the update request
198 + names = []
199 + dtype_names = []
200 + shapes = []
201 + for name, p in train_model.named_parameters():
202 + names.append(name)
203 + dtype_names.append(str(p.dtype).split(".")[-1])
204 + shapes.append(list(p.shape))
205 +
206 + # Start the update_weights call in a separate thread since it will block
207 + # waiting for NCCL broadcasts
208 + # packed=True enables efficient batched tensor broadcasting
209 + update_thread = threading.Thread(
210 + target=update_weights,
211 + args=(BASE_URL, names, dtype_names, shapes, True), # packed=True
212 + )
213 + update_thread.start()
214 +
215 + # Broadcast all weights from trainer to vLLM workers
216 + print("Broadcasting weights via NCCL...")
217 + NCCLWeightTransferEngine.trainer_send_weights(
218 + iterator=train_model.named_parameters(),
219 + group=model_update_group,
220 + packed=True,
221 + )
222 +
223 + # Wait for update_weights to complete
224 + update_thread.join()
225 +
226 + # Resume generation after weight sync
227 + resume_generation(BASE_URL)
228 +
229 + # Generate text after weight update. The output is expected to be normal
230 + # because the real weights are now loaded.
231 + print("-" * 50)
232 + print("Generating text AFTER weight update:")
233 + print("-" * 50)
234 + outputs_updated = generate_completions(client, MODEL_NAME, prompts)
235 + for prompt, generated_text in zip(prompts, outputs_updated):
236 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
237 + print("-" * 50)
238 +
239 +
240 +if __name__ == "__main__":
241 + main()
tests/distributed/test_packed_tensor.py new
+443
@@ -0,0 +1,443 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for packed tensor broadcasting functionality.
4 +
5 +Unit tests for packed_broadcast_producer and packed_broadcast_consumer.
6 +These utilities enable efficient batched tensor transfer over NCCL.
7 +"""
8 +
9 +import pytest
10 +import torch
11 +
12 +from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo
13 +from vllm.distributed.weight_transfer.packed_tensor import (
14 + packed_broadcast_consumer,
15 + packed_broadcast_producer,
16 +)
17 +
18 +
19 +class MockCommunicationGroup:
20 + """Mock communication group for testing producer broadcast operations."""
21 +
22 + def __init__(self):
23 + self.broadcasted_tensors: list[torch.Tensor] = []
24 + self.broadcast_count = 0
25 + self.device = torch.device("cuda:0")
26 +
27 + def broadcast(self, tensor, src):
28 + """Mock broadcast that stores the tensor for later verification."""
29 + self.broadcasted_tensors.append(tensor.clone())
30 + self.broadcast_count += 1
31 +
32 +
33 +class MockConsumerCommunicationGroup:
34 + """Mock communication group for consumer that returns pre-stored tensors."""
35 +
36 + def __init__(self, tensors_to_return: list[torch.Tensor]):
37 + self.tensors_to_return = tensors_to_return
38 + self.current_index = 0
39 + self.device = torch.device("cuda:0")
40 +
41 + def broadcast(self, tensor, src):
42 + """Mock broadcast that fills the tensor with pre-stored data."""
43 + if self.current_index < len(self.tensors_to_return):
44 + tensor.copy_(self.tensors_to_return[self.current_index])
45 + self.current_index += 1
46 +
47 +
48 +def create_mock_model_params(
49 + num_layers: int = 3,
50 + dtype: torch.dtype = torch.float32,
51 +) -> list[tuple[str, torch.Tensor]]:
52 + """Create mock model parameters for testing."""
53 + params = []
54 + for i in range(num_layers):
55 + params.append((f"layer{i}.weight", torch.randn(10, 20, dtype=dtype)))
56 + params.append((f"layer{i}.bias", torch.randn(10, dtype=dtype)))
57 + return params
58 +
59 +
60 +def create_state_dict_info(
61 + params: list[tuple[str, torch.Tensor]],
62 +) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
63 + """Create state dict info (name -> (shape, dtype)) from params."""
64 + return {name: (tuple(tensor.shape), tensor.dtype) for name, tensor in params}
65 +
66 +
67 +# --- Unit Tests: NCCLWeightTransferUpdateInfo packed field ---
68 +
69 +
70 +class TestNCCLWeightTransferUpdateInfoPacked:
71 + """Test NCCLWeightTransferUpdateInfo dataclass packed field."""
72 +
73 + def test_packed_default_false(self):
74 + """Test that packed defaults to False."""
75 + info = NCCLWeightTransferUpdateInfo(
76 + names=["layer.weight"],
77 + dtype_names=["float32"],
78 + shapes=[[10, 10]],
79 + )
80 + assert info.packed is False
81 +
82 + def test_packed_can_be_set_true(self):
83 + """Test that packed can be set to True."""
84 + info = NCCLWeightTransferUpdateInfo(
85 + names=["layer.weight"],
86 + dtype_names=["float32"],
87 + shapes=[[10, 10]],
88 + packed=True,
89 + )
90 + assert info.packed is True
91 +
92 +
93 +# --- Unit Tests: packed_broadcast_producer ---
94 +
95 +
96 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
97 +class TestPackedBroadcastProducer:
98 + """Test packed_broadcast_producer function."""
99 +
100 + def test_producer_broadcasts_tensors(self):
101 + """Test that producer broadcasts all tensors."""
102 + params = create_mock_model_params()
103 + params_cuda = [(name, tensor.cuda()) for name, tensor in params]
104 +
105 + mock_group = MockCommunicationGroup()
106 +
107 + # Use a small target size to force multiple batches
108 + packed_broadcast_producer(
109 + iterator=iter(params_cuda),
110 + group=mock_group,
111 + src=0,
112 + post_iter_func=lambda x: x[1],
113 + buffer_size_bytes=500,
114 + )
115 +
116 + # Should have broadcasted some tensors
117 + assert mock_group.broadcast_count > 0
118 + assert len(mock_group.broadcasted_tensors) > 0
119 +
120 + def test_producer_single_large_tensor(self):
121 + """Test with a single tensor larger than target size."""
122 + # Create a large tensor
123 + large_tensor = torch.randn(1000, 1000, dtype=torch.float32).cuda()
124 + params = [("large_weight", large_tensor)]
125 +
126 + mock_group = MockCommunicationGroup()
127 +
128 + # Small target size to force the tensor to exceed it
129 + packed_broadcast_producer(
130 + iterator=iter(params),
131 + group=mock_group,
132 + src=0,
133 + post_iter_func=lambda x: x[1],
134 + buffer_size_bytes=100,
135 + )
136 +
137 + # Should still broadcast the tensor (at least 1 broadcast)
138 + assert mock_group.broadcast_count >= 1
139 + assert len(mock_group.broadcasted_tensors) >= 1
140 +
141 + # Verify the total broadcasted size matches the tensor
142 + expected_size = large_tensor.numel() * large_tensor.element_size()
143 + actual_size = sum(t.numel() for t in mock_group.broadcasted_tensors)
144 + assert actual_size == expected_size
145 +
146 + def test_producer_multiple_batches(self):
147 + """Test that tensors are properly batched when exceeding target size."""
148 + # Create many small tensors
149 + params = [
150 + (f"weight_{i}", torch.randn(10, 10, dtype=torch.float32).cuda())
151 + for i in range(20)
152 + ]
153 +
154 + mock_group = MockCommunicationGroup()
155 +
156 + # Small target size to force multiple batches
157 + packed_broadcast_producer(
158 + iterator=iter(params),
159 + group=mock_group,
160 + src=0,
161 + post_iter_func=lambda x: x[1],
162 + buffer_size_bytes=2000,
163 + )
164 +
165 + # Should have multiple broadcasts
166 + assert mock_group.broadcast_count > 1
167 +
168 + # Total size should match sum of all tensors
169 + expected_total = sum(t.numel() * t.element_size() for _, t in params)
170 + actual_total = sum(t.numel() for t in mock_group.broadcasted_tensors)
171 + assert actual_total == expected_total
172 +
173 + def test_producer_empty_iterator(self):
174 + """Test producer handles empty iterator gracefully."""
175 + mock_group = MockCommunicationGroup()
176 +
177 + packed_broadcast_producer(
178 + iterator=iter([]),
179 + group=mock_group,
180 + src=0,
181 + post_iter_func=lambda x: x[1],
182 + buffer_size_bytes=1000,
183 + )
184 +
185 + # No broadcasts for empty iterator
186 + assert mock_group.broadcast_count == 0
187 +
188 +
189 +# --- Unit Tests: packed_broadcast_consumer ---
190 +
191 +
192 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
193 +class TestPackedBroadcastConsumer:
194 + """Test packed_broadcast_consumer function."""
195 +
196 + def test_consumer_receives_tensors(self):
197 + """Test that consumer receives and unpacks tensors."""
198 + params = create_mock_model_params()
199 + params_cuda = [(name, tensor.cuda()) for name, tensor in params]
200 +
201 + buffer_size = 2000
202 +
203 + # First, run producer to get the broadcasted tensors
204 + producer_group = MockCommunicationGroup()
205 +
206 + packed_broadcast_producer(
207 + iterator=iter(params_cuda),
208 + group=producer_group,
209 + src=0,
210 + post_iter_func=lambda x: x[1],
211 + buffer_size_bytes=buffer_size,
212 + )
213 +
214 + # Now run consumer with the broadcasted tensors
215 + consumer_group = MockConsumerCommunicationGroup(
216 + producer_group.broadcasted_tensors
217 + )
218 +
219 + state_dict_info = create_state_dict_info(params_cuda)
220 +
221 + unpacked_tensors = {}
222 +
223 + def post_unpack_func(tensor_list):
224 + for name, tensor in tensor_list:
225 + unpacked_tensors[name] = tensor.clone()
226 +
227 + packed_broadcast_consumer(
228 + iterator=iter(state_dict_info.items()),
229 + group=consumer_group,
230 + src=0,
231 + post_unpack_func=post_unpack_func,
232 + buffer_size_bytes=buffer_size,
233 + )
234 +
235 + # Verify all parameters were unpacked
236 + assert len(unpacked_tensors) == len(params)
237 +
238 + # Verify each tensor matches the original
239 + for name, original_tensor in params_cuda:
240 + assert name in unpacked_tensors
241 + unpacked = unpacked_tensors[name]
242 + assert unpacked.shape == original_tensor.shape
243 + assert unpacked.dtype == original_tensor.dtype
244 + assert torch.allclose(unpacked, original_tensor, rtol=1e-5, atol=1e-7)
245 +
246 +
247 +# --- Integration Tests: Producer-Consumer Roundtrip ---
248 +
249 +
250 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
251 +class TestPackedBroadcastRoundtrip:
252 + """Test producer-consumer roundtrip behavior."""
253 +
254 + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
255 + def test_roundtrip_different_dtypes(self, dtype):
256 + """Test roundtrip with different data types."""
257 + params = create_mock_model_params(num_layers=2, dtype=dtype)
258 + params_cuda = [(name, tensor.cuda()) for name, tensor in params]
259 +
260 + buffer_size = 1000
261 + producer_group = MockCommunicationGroup()
262 +
263 + packed_broadcast_producer(
264 + iterator=iter(params_cuda),
265 + group=producer_group,
266 + src=0,
267 + post_iter_func=lambda x: x[1],
268 + buffer_size_bytes=buffer_size,
269 + )
270 +
271 + consumer_group = MockConsumerCommunicationGroup(
272 + producer_group.broadcasted_tensors
273 + )
274 +
275 + state_dict_info = create_state_dict_info(params_cuda)
276 + unpacked_tensors = {}
277 +
278 + def post_unpack_func(tensor_list):
279 + for name, tensor in tensor_list:
280 + unpacked_tensors[name] = tensor.clone()
281 +
282 + packed_broadcast_consumer(
283 + iterator=iter(state_dict_info.items()),
284 + group=consumer_group,
285 + src=0,
286 + post_unpack_func=post_unpack_func,
287 + buffer_size_bytes=buffer_size,
288 + )
289 +
290 + # Verify roundtrip preserves data
291 + for name, original_tensor in params_cuda:
292 + assert name in unpacked_tensors
293 + unpacked = unpacked_tensors[name]
294 + assert unpacked.dtype == dtype
295 + assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
296 +
297 + def test_roundtrip_mixed_dtypes(self):
298 + """Test roundtrip with mixed data types."""
299 + # Create params with mixed dtypes
300 + params = [
301 + ("layer1.weight", torch.randn(10, 20, dtype=torch.float32).cuda()),
302 + ("layer1.bias", torch.randn(10, dtype=torch.float16).cuda()),
303 + ("layer2.weight", torch.randn(20, 30, dtype=torch.bfloat16).cuda()),
304 + ]
305 +
306 + buffer_size = 500
307 + producer_group = MockCommunicationGroup()
308 +
309 + packed_broadcast_producer(
310 + iterator=iter(params),
311 + group=producer_group,
312 + src=0,
313 + post_iter_func=lambda x: x[1],
314 + buffer_size_bytes=buffer_size,
315 + )
316 +
317 + consumer_group = MockConsumerCommunicationGroup(
318 + producer_group.broadcasted_tensors
319 + )
320 +
321 + state_dict_info = create_state_dict_info(params)
322 + unpacked_tensors = {}
323 +
324 + def post_unpack_func(tensor_list):
325 + for name, tensor in tensor_list:
326 + unpacked_tensors[name] = tensor.clone()
327 +
328 + packed_broadcast_consumer(
329 + iterator=iter(state_dict_info.items()),
330 + group=consumer_group,
331 + src=0,
332 + post_unpack_func=post_unpack_func,
333 + buffer_size_bytes=buffer_size,
334 + )
335 +
336 + # Verify all params roundtrip correctly with correct dtypes
337 + for name, original_tensor in params:
338 + assert name in unpacked_tensors
339 + unpacked = unpacked_tensors[name]
340 + assert unpacked.shape == original_tensor.shape
341 + assert unpacked.dtype == original_tensor.dtype
342 + assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
343 +
344 + @pytest.mark.parametrize("target_size", [100, 1000, 10000, 100000])
345 + def test_roundtrip_different_batch_sizes(self, target_size):
346 + """Test roundtrip with different target batch sizes."""
347 + params = create_mock_model_params(num_layers=5)
348 + params_cuda = [(name, tensor.cuda()) for name, tensor in params]
349 +
350 + producer_group = MockCommunicationGroup()
351 +
352 + packed_broadcast_producer(
353 + iterator=iter(params_cuda),
354 + group=producer_group,
355 + src=0,
356 + post_iter_func=lambda x: x[1],
357 + buffer_size_bytes=target_size,
358 + )
359 +
360 + consumer_group = MockConsumerCommunicationGroup(
361 + producer_group.broadcasted_tensors
362 + )
363 +
364 + state_dict_info = create_state_dict_info(params_cuda)
365 + unpacked_tensors = {}
366 +
367 + def post_unpack_func(tensor_list):
368 + for name, tensor in tensor_list:
369 + unpacked_tensors[name] = tensor.clone()
370 +
371 + packed_broadcast_consumer(
372 + iterator=iter(state_dict_info.items()),
373 + group=consumer_group,
374 + src=0,
375 + post_unpack_func=post_unpack_func,
376 + buffer_size_bytes=target_size,
377 + )
378 +
379 + # Verify all params roundtrip correctly
380 + assert len(unpacked_tensors) == len(params)
381 + for name, original_tensor in params_cuda:
382 + assert name in unpacked_tensors
383 + assert torch.allclose(
384 + unpacked_tensors[name], original_tensor, rtol=1e-5, atol=1e-7
385 + )
386 +
387 + def test_roundtrip_non_contiguous_tensors(self):
388 + """Test roundtrip with non-contiguous tensors from the trainer."""
389 + # Create non-contiguous tensors (simulating trainer outputs)
390 + # Transposed tensors are non-contiguous
391 + weight1 = torch.randn(20, 10, dtype=torch.float32).cuda().T
392 + # Sliced tensors with step are non-contiguous
393 + weight2 = torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2]
394 + # Permuted tensors are non-contiguous
395 + weight3 = torch.randn(5, 10, 15, dtype=torch.bfloat16).cuda().permute(2, 0, 1)
396 +
397 + params = [
398 + ("layer1.weight", weight1),
399 + ("layer2.weight", weight2),
400 + ("layer3.weight", weight3),
401 + ]
402 +
403 + # Verify tensors are indeed non-contiguous
404 + for name, tensor in params:
405 + assert not tensor.is_contiguous(), f"{name} should be non-contiguous"
406 +
407 + buffer_size = 500
408 + producer_group = MockCommunicationGroup()
409 +
410 + packed_broadcast_producer(
411 + iterator=iter(params),
412 + group=producer_group,
413 + src=0,
414 + post_iter_func=lambda x: x[1],
415 + buffer_size_bytes=buffer_size,
416 + )
417 +
418 + consumer_group = MockConsumerCommunicationGroup(
419 + producer_group.broadcasted_tensors
420 + )
421 +
422 + state_dict_info = create_state_dict_info(params)
423 + unpacked_tensors = {}
424 +
425 + def post_unpack_func(tensor_list):
426 + for name, tensor in tensor_list:
427 + unpacked_tensors[name] = tensor.clone()
428 +
429 + packed_broadcast_consumer(
430 + iterator=iter(state_dict_info.items()),
431 + group=consumer_group,
432 + src=0,
433 + post_unpack_func=post_unpack_func,
434 + buffer_size_bytes=buffer_size,
435 + )
436 +
437 + # Verify all non-contiguous params roundtrip correctly
438 + for name, original_tensor in params:
439 + assert name in unpacked_tensors
440 + unpacked = unpacked_tensors[name]
441 + assert unpacked.shape == original_tensor.shape
442 + assert unpacked.dtype == original_tensor.dtype
443 + assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
tests/distributed/test_weight_transfer.py new
+346
@@ -0,0 +1,346 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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.
7 +"""
8 +
9 +from unittest.mock import MagicMock
10 +
11 +import pytest
12 +import ray
13 +import torch
14 +
15 +from vllm.config.parallel import ParallelConfig
16 +from vllm.config.weight_transfer import WeightTransferConfig
17 +from vllm.distributed.weight_transfer import WeightTransferEngineFactory
18 +from vllm.distributed.weight_transfer.nccl_engine import (
19 + NCCLWeightTransferEngine,
20 + NCCLWeightTransferInitInfo,
21 + NCCLWeightTransferUpdateInfo,
22 +)
23 +from vllm.utils.network_utils import get_open_port
24 +
25 +
26 +def create_mock_parallel_config(
27 + rank: int = 0,
28 + world_size: int = 1,
29 + dp_rank: int = 0,
30 +) -> ParallelConfig:
31 + """Create a mock ParallelConfig for testing."""
32 + config = MagicMock(spec=ParallelConfig)
33 + config.rank = rank
34 + config.world_size = world_size
35 + config.data_parallel_rank = dp_rank
36 + return config
37 +
38 +
39 +# --- Unit Tests: NCCLWeightTransferUpdateInfo Validation ---
40 +
41 +
42 +class TestNCCLWeightTransferUpdateInfoValidation:
43 + """Test NCCLWeightTransferUpdateInfo dataclass validation."""
44 +
45 + def test_valid_update_info(self):
46 + """Test creating valid NCCLWeightTransferUpdateInfo."""
47 + info = NCCLWeightTransferUpdateInfo(
48 + names=["layer.weight", "layer.bias"],
49 + dtype_names=["float32", "float32"],
50 + shapes=[[10, 10], [10]],
51 + )
52 + assert info.names == ["layer.weight", "layer.bias"]
53 + assert info.dtype_names == ["float32", "float32"]
54 + assert info.shapes == [[10, 10], [10]]
55 +
56 + def test_mismatched_dtype_names_raises(self):
57 + """Test that mismatched dtype_names length raises ValueError."""
58 + with pytest.raises(ValueError, match="dtype_names"):
59 + NCCLWeightTransferUpdateInfo(
60 + names=["layer.weight", "layer.bias"],
61 + dtype_names=["float32"], # Only one dtype
62 + shapes=[[10, 10], [10]],
63 + )
64 +
65 + def test_mismatched_shapes_raises(self):
66 + """Test that mismatched shapes length raises ValueError."""
67 + with pytest.raises(ValueError, match="shapes"):
68 + NCCLWeightTransferUpdateInfo(
69 + names=["layer.weight", "layer.bias"],
70 + dtype_names=["float32", "float32"],
71 + shapes=[[10, 10]], # Only one shape
72 + )
73 +
74 + def test_empty_lists_valid(self):
75 + """Test that empty lists are valid."""
76 + info = NCCLWeightTransferUpdateInfo(
77 + names=[],
78 + dtype_names=[],
79 + shapes=[],
80 + )
81 + assert len(info.names) == 0
82 +
83 +
84 +# --- Unit Tests: Engine Parsing ---
85 +
86 +
87 +class TestNCCLEngineParsing:
88 + """Test NCCLWeightTransferEngine parsing methods."""
89 +
90 + def test_parse_init_info_valid(self):
91 + """Test parsing valid init info dict."""
92 + config = WeightTransferConfig(backend="nccl")
93 + parallel_config = create_mock_parallel_config()
94 + engine = NCCLWeightTransferEngine(config, parallel_config)
95 +
96 + init_info = engine.parse_init_info(
97 + {
98 + "master_address": "127.0.0.1",
99 + "master_port": 12345,
100 + "rank_offset": 1,
101 + "world_size": 3,
102 + }
103 + )
104 +
105 + assert isinstance(init_info, NCCLWeightTransferInitInfo)
106 + assert init_info.master_address == "127.0.0.1"
107 + assert init_info.master_port == 12345
108 + assert init_info.rank_offset == 1
109 + assert init_info.world_size == 3
110 +
111 + def test_parse_init_info_missing_field_raises(self):
112 + """Test parsing init info with missing required field."""
113 + config = WeightTransferConfig(backend="nccl")
114 + parallel_config = create_mock_parallel_config()
115 + engine = NCCLWeightTransferEngine(config, parallel_config)
116 +
117 + with pytest.raises(ValueError, match="Invalid init_info"):
118 + engine.parse_init_info(
119 + {
120 + "master_address": "127.0.0.1",
121 + # Missing master_port, rank_offset, world_size
122 + }
123 + )
124 +
125 + def test_parse_update_info_valid(self):
126 + """Test parsing valid update info dict."""
127 + config = WeightTransferConfig(backend="nccl")
128 + parallel_config = create_mock_parallel_config()
129 + engine = NCCLWeightTransferEngine(config, parallel_config)
130 +
131 + update_info = engine.parse_update_info(
132 + {
133 + "names": ["w1", "w2"],
134 + "dtype_names": ["float32", "bfloat16"],
135 + "shapes": [[100, 100], [50]],
136 + }
137 + )
138 +
139 + assert isinstance(update_info, NCCLWeightTransferUpdateInfo)
140 + assert update_info.names == ["w1", "w2"]
141 + assert update_info.dtype_names == ["float32", "bfloat16"]
142 + assert update_info.shapes == [[100, 100], [50]]
143 +
144 +
145 +# --- Unit Tests: Engine Registry ---
146 +
147 +
148 +class TestEngineRegistry:
149 + """Test weight transfer engine registry."""
150 +
151 + def test_create_engine_nccl(self):
152 + """Test factory creates NCCL engine."""
153 + config = WeightTransferConfig(backend="nccl")
154 + parallel_config = create_mock_parallel_config()
155 + engine = WeightTransferEngineFactory.create_engine(config, parallel_config)
156 + assert isinstance(engine, NCCLWeightTransferEngine)
157 +
158 + def test_create_engine_invalid_backend(self):
159 + """Test factory raises for invalid backend."""
160 + config = WeightTransferConfig(backend="invalid")
161 + parallel_config = create_mock_parallel_config()
162 + with pytest.raises(ValueError, match="Invalid weight transfer backend"):
163 + WeightTransferEngineFactory.create_engine(config, parallel_config)
164 +
165 + def test_register_duplicate_raises(self):
166 + """Test registering duplicate engine name raises."""
167 + with pytest.raises(ValueError, match="already registered"):
168 + WeightTransferEngineFactory.register_engine(
169 + "nccl", NCCLWeightTransferEngine
170 + )
171 +
172 +
173 +# --- Test receive_weights without init raises ---
174 +
175 +
176 +def test_nccl_receive_weights_without_init_raises():
177 + """Test that receive_weights raises if init_transfer_engine wasn't called."""
178 + if torch.cuda.device_count() < 1:
179 + pytest.skip("Need at least 1 GPU for this test")
180 +
181 + config = WeightTransferConfig(backend="nccl")
182 + parallel_config = create_mock_parallel_config()
183 + engine = NCCLWeightTransferEngine(config, parallel_config)
184 +
185 + update_info = NCCLWeightTransferUpdateInfo(
186 + names=["w"],
187 + dtype_names=["float32"],
188 + shapes=[[10]],
189 + )
190 +
191 + with pytest.raises(RuntimeError, match="not initialized"):
192 + engine.receive_weights(update_info, lambda x: None)
193 +
194 +
195 +# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
196 +
197 +
198 +@ray.remote(num_gpus=1)
199 +def trainer_broadcast_tensor(
200 + master_address: str,
201 + master_port: int,
202 + world_size: int,
203 + tensor_shape: list[int],
204 + tensor_dtype: str,
205 +) -> bool:
206 + """Trainer task that broadcasts a tensor via NCCL."""
207 + import torch
208 +
209 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
210 + from vllm.distributed.utils import StatelessProcessGroup
211 +
212 + # Create process group as rank 0 (trainer)
213 + pg = StatelessProcessGroup.create(
214 + host=master_address,
215 + port=master_port,
216 + rank=0,
217 + world_size=world_size,
218 + )
219 + # Ray sets CUDA_VISIBLE_DEVICES, so device 0 is the assigned GPU
220 + comm = PyNcclCommunicator(pg, device=0)
221 +
222 + # Create and broadcast the tensor
223 + dtype = getattr(torch, tensor_dtype)
224 + tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device="cuda:0")
225 + comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream())
226 + torch.cuda.synchronize()
227 +
228 + return True
229 +
230 +
231 +@ray.remote(num_gpus=1)
232 +def inference_receive_tensor(
233 + master_address: str,
234 + master_port: int,
235 + world_size: int,
236 + tensor_shape: list[int],
237 + tensor_dtype: str,
238 +) -> dict:
239 + """Inference task that receives tensor via NCCLWeightTransferEngine."""
240 + from unittest.mock import MagicMock
241 +
242 + import torch
243 +
244 + from vllm.config.parallel import ParallelConfig
245 + from vllm.config.weight_transfer import WeightTransferConfig
246 + from vllm.distributed.weight_transfer.nccl_engine import (
247 + NCCLWeightTransferEngine,
248 + NCCLWeightTransferInitInfo,
249 + NCCLWeightTransferUpdateInfo,
250 + )
251 +
252 + # Create engine with mock parallel config
253 + config = WeightTransferConfig(backend="nccl")
254 + parallel_config = MagicMock(spec=ParallelConfig)
255 + parallel_config.rank = 0
256 + parallel_config.world_size = 1
257 + parallel_config.data_parallel_rank = 0
258 +
259 + engine = NCCLWeightTransferEngine(config, parallel_config)
260 +
261 + # Initialize the engine (joins as rank 1)
262 + init_info = NCCLWeightTransferInitInfo(
263 + master_address=master_address,
264 + master_port=master_port,
265 + rank_offset=1, # Trainer is rank 0, we become rank 1
266 + world_size=world_size,
267 + )
268 + engine.init_transfer_engine(init_info)
269 +
270 + # Receive weights with a no-op load_weights that captures the tensor
271 + received_tensors = []
272 +
273 + def noop_load_weights(weights: list[tuple[str, torch.Tensor]]):
274 + for name, tensor in weights:
275 + # Clone tensor to keep it after engine cleans up
276 + received_tensors.append((name, tensor.clone()))
277 +
278 + update_info = NCCLWeightTransferUpdateInfo(
279 + names=["test.weight"],
280 + dtype_names=[tensor_dtype],
281 + shapes=[tensor_shape],
282 + )
283 + engine.receive_weights(update_info, noop_load_weights)
284 + torch.cuda.synchronize()
285 +
286 + # Verify we received the tensor
287 + success = False
288 + received_shape = None
289 + received_sum = None
290 +
291 + if len(received_tensors) == 1:
292 + name, tensor = received_tensors[0]
293 + received_shape = list(tensor.shape)
294 + received_sum = tensor.sum().item()
295 + # Check shape matches and values are all 1s (trainer sends ones)
296 + if received_shape == tensor_shape:
297 + expected_sum = 1.0 * torch.tensor(tensor_shape).prod().item()
298 + if abs(received_sum - expected_sum) < 0.01:
299 + success = True
300 +
301 + engine.shutdown()
302 +
303 + return {
304 + "success": success,
305 + "received_shape": received_shape,
306 + "received_sum": received_sum,
307 + }
308 +
309 +
310 +@pytest.mark.skipif(
311 + torch.cuda.device_count() < 2,
312 + reason="Need at least 2 GPUs to run NCCL weight transfer test.",
313 +)
314 +def test_nccl_weight_transfer_between_processes():
315 + """Test NCCL weight transfer from trainer to inference process using Ray.
316 +
317 + This test verifies that the NCCLWeightTransferEngine can receive
318 + tensors broadcast by a trainer process via NCCL.
319 + """
320 + ray.init(ignore_reinit_error=True)
321 +
322 + master_address = "127.0.0.1"
323 + master_port = get_open_port()
324 + world_size = 2 # 1 trainer + 1 inference worker
325 +
326 + # Tensor to transfer: 100x100 ones
327 + tensor_shape = [100, 100]
328 + tensor_dtype = "float32"
329 +
330 + # Start both tasks concurrently - Ray assigns GPUs automatically
331 + inference_future = inference_receive_tensor.remote(
332 + master_address, master_port, world_size, tensor_shape, tensor_dtype
333 + )
334 + trainer_future = trainer_broadcast_tensor.remote(
335 + master_address, master_port, world_size, tensor_shape, tensor_dtype
336 + )
337 +
338 + # Wait for both to complete
339 + trainer_result, result = ray.get([trainer_future, inference_future])
340 +
341 + assert trainer_result, "Trainer should complete successfully"
342 + assert result["success"], (
343 + f"Weight transfer failed. "
344 + f"Received shape: {result['received_shape']}, "
345 + f"Received sum: {result['received_sum']}"
346 + )
tests/entrypoints/openai/test_openai_schema.py
+8
@@ -139,6 +139,14 @@ def test_openapi_stateless(case: schemathesis.Case):
139 # Skip responses API as it is meant to be stateful.
140 return
141
142 + # Skip weight transfer endpoints as they require special setup
143 + # (weight_transfer_config) and are meant to be stateful.
144 + if case.operation.path in (
145 + "/init_weight_transfer_engine",
146 + "/update_weights",
147 + ):
148 + return
149 +
150 timeout = {
151 # requires a longer timeout
152 ("POST", "/v1/chat/completions"): LONG_TIMEOUT_SECONDS,
tests/entrypoints/weight_transfer/__init__.py new
+3
@@ -0,0 +1,3 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
tests/entrypoints/weight_transfer/test_weight_transfer_llm.py new
+300
@@ -0,0 +1,300 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for weight transfer APIs via LLM class.
4 +
5 +These tests use a mock weight transfer engine to verify that the API
6 +calls the correct methods with the right arguments, without requiring
7 +actual NCCL communication.
8 +"""
9 +
10 +import os
11 +from collections.abc import Callable
12 +from dataclasses import dataclass
13 +from unittest.mock import patch
14 +
15 +import pytest
16 +import torch
17 +
18 +from vllm import LLM
19 +from vllm.config import WeightTransferConfig
20 +from vllm.distributed.weight_transfer.base import (
21 + WeightTransferEngine,
22 + WeightTransferInitInfo,
23 + WeightTransferInitRequest,
24 + WeightTransferUpdateInfo,
25 + WeightTransferUpdateRequest,
26 +)
27 +
28 +from ...utils import create_new_process_for_each_test
29 +
30 +# Use a tiny model for fast testing
31 +MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
32 +
33 +
34 +# --- Mock Weight Transfer Engine ---
35 +
36 +
37 +@dataclass
38 +class MockInitInfo(WeightTransferInitInfo):
39 + """Mock initialization info."""
40 +
41 + test_param: str = "test"
42 +
43 +
44 +@dataclass
45 +class MockUpdateInfo(WeightTransferUpdateInfo):
46 + """Mock update info."""
47 +
48 + names: list[str] | None = None
49 + dtype_names: list[str] | None = None
50 + shapes: list[list[int]] | None = None
51 +
52 +
53 +class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
54 + """Mock weight transfer engine that tracks method calls."""
55 +
56 + init_info_cls = MockInitInfo
57 + update_info_cls = MockUpdateInfo
58 +
59 + # Class-level tracking for verification across processes
60 + init_transfer_engine_called: bool = False
61 + receive_weights_called: bool = False
62 + shutdown_called: bool = False
63 + last_init_info: MockInitInfo | None = None
64 + last_update_info: MockUpdateInfo | None = None
65 +
66 + def __init__(self, config, parallel_config):
67 + super().__init__(config, parallel_config)
68 + # Reset tracking on init
69 + MockWeightTransferEngine.init_transfer_engine_called = False
70 + MockWeightTransferEngine.receive_weights_called = False
71 + MockWeightTransferEngine.shutdown_called = False
72 + MockWeightTransferEngine.last_init_info = None
73 + MockWeightTransferEngine.last_update_info = None
74 +
75 + def init_transfer_engine(self, init_info: MockInitInfo) -> None:
76 + MockWeightTransferEngine.init_transfer_engine_called = True
77 + MockWeightTransferEngine.last_init_info = init_info
78 +
79 + def receive_weights(
80 + self,
81 + update_info: MockUpdateInfo,
82 + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
83 + ) -> None:
84 + MockWeightTransferEngine.receive_weights_called = True
85 + MockWeightTransferEngine.last_update_info = update_info
86 + # Simulate loading weights by calling load_weights with empty list
87 + # (In real implementation, this would receive and load actual weights)
88 + load_weights([])
89 +
90 + def shutdown(self) -> None:
91 + MockWeightTransferEngine.shutdown_called = True
92 +
93 +
94 +def mock_create_engine(config, parallel_config):
95 + """Mock factory function that returns our mock engine."""
96 + return MockWeightTransferEngine(config, parallel_config)
97 +
98 +
99 +# --- Tests ---
100 +
101 +
102 +@create_new_process_for_each_test()
103 +def test_get_world_size_tp1():
104 + """Test world_size is correctly configured for TP=1."""
105 + if torch.cuda.device_count() < 1:
106 + pytest.skip("Need at least 1 GPU for this test")
107 +
108 + llm = LLM(
109 + model=MODEL_NAME,
110 + enforce_eager=True,
111 + load_format="dummy",
112 + tensor_parallel_size=1,
113 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
114 + )
115 +
116 + world_size = llm.llm_engine.vllm_config.parallel_config.world_size
117 + assert world_size == 1
118 +
119 +
120 +@create_new_process_for_each_test()
121 +def test_init_weight_transfer_engine_calls_engine():
122 + """Test that init_weight_transfer_engine calls the engine's
123 + init_transfer_engine method."""
124 + if torch.cuda.device_count() < 1:
125 + pytest.skip("Need at least 1 GPU for this test")
126 +
127 + # Enable insecure serialization to allow pickling functions for collective_rpc
128 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
129 +
130 + with patch(
131 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
132 + mock_create_engine,
133 + ):
134 + llm = LLM(
135 + model=MODEL_NAME,
136 + enforce_eager=True,
137 + load_format="dummy",
138 + tensor_parallel_size=1,
139 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
140 + )
141 +
142 + # Verify engine was created
143 + def check_engine_exists(self):
144 + return self.weight_transfer_engine is not None
145 +
146 + results = llm.collective_rpc(check_engine_exists)
147 + assert all(results), "Weight transfer engine should be initialized"
148 +
149 + # Call init_weight_transfer_engine
150 + llm.init_weight_transfer_engine(
151 + WeightTransferInitRequest(init_info={"test_param": "hello"})
152 + )
153 +
154 + # Verify init_transfer_engine was called on the engine
155 + def check_init_called(self):
156 + engine = self.weight_transfer_engine
157 + return (
158 + engine.init_transfer_engine_called,
159 + engine.last_init_info.test_param if engine.last_init_info else None,
160 + )
161 +
162 + results = llm.collective_rpc(check_init_called)
163 + for called, param in results:
164 + assert called, "init_transfer_engine should have been called"
165 + assert param == "hello", f"Expected 'hello', got {param}"
166 +
167 +
168 +@create_new_process_for_each_test()
169 +def test_update_weights_calls_engine():
170 + """Test that update_weights calls the engine's receive_weights method."""
171 + if torch.cuda.device_count() < 1:
172 + pytest.skip("Need at least 1 GPU for this test")
173 +
174 + # Enable insecure serialization to allow pickling functions for collective_rpc
175 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
176 +
177 + with patch(
178 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
179 + mock_create_engine,
180 + ):
181 + llm = LLM(
182 + model=MODEL_NAME,
183 + enforce_eager=True,
184 + load_format="dummy",
185 + tensor_parallel_size=1,
186 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
187 + )
188 +
189 + # First init the weight transfer
190 + llm.init_weight_transfer_engine(
191 + WeightTransferInitRequest(init_info={"test_param": "init"})
192 + )
193 +
194 + # Call update_weights
195 + test_names = ["layer.weight", "layer.bias"]
196 + test_dtypes = ["float32", "float32"]
197 + test_shapes = [[10, 10], [10]]
198 +
199 + llm.update_weights(
200 + WeightTransferUpdateRequest(
201 + update_info={
202 + "names": test_names,
203 + "dtype_names": test_dtypes,
204 + "shapes": test_shapes,
205 + }
206 + )
207 + )
208 +
209 + # Verify receive_weights was called with correct info
210 + def check_update_called(self):
211 + engine = self.weight_transfer_engine
212 + if not engine.receive_weights_called:
213 + return False, None, None, None
214 + info = engine.last_update_info
215 + return (True, info.names, info.dtype_names, info.shapes)
216 +
217 + results = llm.collective_rpc(check_update_called)
218 + for called, names, dtypes, shapes in results:
219 + assert called, "receive_weights should have been called"
220 + assert names == test_names
221 + assert dtypes == test_dtypes
222 + assert shapes == test_shapes
223 +
224 +
225 +@create_new_process_for_each_test()
226 +def test_full_weight_transfer_flow():
227 + """Test the complete weight transfer flow: init -> update."""
228 + if torch.cuda.device_count() < 1:
229 + pytest.skip("Need at least 1 GPU for this test")
230 +
231 + # Enable insecure serialization to allow pickling functions for collective_rpc
232 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
233 +
234 + with patch(
235 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
236 + mock_create_engine,
237 + ):
238 + llm = LLM(
239 + model=MODEL_NAME,
240 + enforce_eager=True,
241 + load_format="dummy",
242 + tensor_parallel_size=1,
243 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
244 + )
245 +
246 + # Step 1: Initialize
247 + llm.init_weight_transfer_engine(
248 + WeightTransferInitRequest(init_info={"test_param": "flow_test"})
249 + )
250 +
251 + # Step 2: Update weights
252 + llm.update_weights(
253 + WeightTransferUpdateRequest(
254 + update_info={
255 + "names": ["test.weight"],
256 + "dtype_names": ["bfloat16"],
257 + "shapes": [[100, 100]],
258 + }
259 + )
260 + )
261 +
262 + # Verify the full flow completed
263 + def check_flow(self):
264 + engine = self.weight_transfer_engine
265 + return {
266 + "init_called": engine.init_transfer_engine_called,
267 + "update_called": engine.receive_weights_called,
268 + "init_param": (
269 + engine.last_init_info.test_param if engine.last_init_info else None
270 + ),
271 + "update_names": (
272 + engine.last_update_info.names if engine.last_update_info else None
273 + ),
274 + }
275 +
276 + results = llm.collective_rpc(check_flow)
277 + for result in results:
278 + assert result["init_called"], "init_transfer_engine should be called"
279 + assert result["update_called"], "receive_weights should be called"
280 + assert result["init_param"] == "flow_test"
281 + assert result["update_names"] == ["test.weight"]
282 +
283 +
284 +@create_new_process_for_each_test()
285 +def test_weight_transfer_config_backend():
286 + """Test that WeightTransferConfig backend is properly configured."""
287 + if torch.cuda.device_count() < 1:
288 + pytest.skip("Need at least 1 GPU for this test")
289 +
290 + # Test with nccl backend
291 + llm = LLM(
292 + model=MODEL_NAME,
293 + enforce_eager=True,
294 + load_format="dummy",
295 + tensor_parallel_size=1,
296 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
297 + )
298 +
299 + config = llm.llm_engine.vllm_config.weight_transfer_config
300 + assert config.backend == "nccl"
vllm/config/__init__.py
+2
@@ -47,6 +47,7 @@ from vllm.config.vllm import (
47 get_layers_from_vllm_config,
48 set_current_vllm_config,
49 )
50 +from vllm.config.weight_transfer import WeightTransferConfig
51
52 # __all__ should only contain classes and functions.
53 # Types and globals should be imported from their respective modules.
@@ -111,4 +112,5 @@ __all__ = [
112 "get_current_vllm_config_or_none",
113 "set_current_vllm_config",
114 "get_layers_from_vllm_config",
115 + "WeightTransferConfig",
116 ]
vllm/config/vllm.py
+4
@@ -42,6 +42,7 @@ from .scheduler import SchedulerConfig
42 from .speculative import EagleModelTypes, SpeculativeConfig
43 from .structured_outputs import StructuredOutputsConfig
44 from .utils import SupportsHash, config, replace
45 +from .weight_transfer import WeightTransferConfig
46
47 if TYPE_CHECKING:
48 from transformers import PretrainedConfig
@@ -255,6 +256,9 @@ class VllmConfig:
256 performance. -02 is used by defult. See OptimizationLevel for full
257 description."""
258
259 + weight_transfer_config: WeightTransferConfig | None = None
260 + """The configurations for weight transfer during RL training."""
261 +
262 def compute_hash(self) -> str:
263 """
264 WARNING: Whenever a new field is added to this config,
vllm/config/weight_transfer.py new
+15
@@ -0,0 +1,15 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass
4 +from typing import Literal
5 +
6 +from vllm.config.utils import config
7 +
8 +
9 +@config
10 +@dataclass
11 +class WeightTransferConfig:
12 + """Configuration for weight transfer during RL training."""
13 +
14 + backend: Literal["nccl"] = "nccl"
15 + """The backend to use for weight transfer."""
vllm/distributed/weight_transfer/__init__.py new
+12
@@ -0,0 +1,12 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Weight transfer engines for syncing model weights from trainers
5 +to inference workers.
6 +"""
7 +
8 +from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
9 +
10 +__all__ = [
11 + "WeightTransferEngineFactory",
12 +]
vllm/distributed/weight_transfer/base.py new
+158
@@ -0,0 +1,158 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Base class for weight transfer engines."""
4 +
5 +from abc import ABC, abstractmethod
6 +from collections.abc import Callable
7 +from dataclasses import KW_ONLY, dataclass, field
8 +from typing import Any, Generic, TypeVar
9 +
10 +import torch
11 +
12 +from vllm.config.parallel import ParallelConfig
13 +from vllm.config.weight_transfer import WeightTransferConfig
14 +
15 +TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
16 +TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
17 +
18 +
19 +# Base protocols for backend-specific dataclasses
20 +@dataclass
21 +class WeightTransferInitInfo(ABC): # noqa: B024
22 + """Base class for backend-specific initialization info."""
23 +
24 + pass
25 +
26 +
27 +@dataclass
28 +class WeightTransferUpdateInfo(ABC): # noqa: B024
29 + """Base class for backend-specific weight update info."""
30 +
31 + _: KW_ONLY
32 + is_checkpoint_format: bool = True
33 + """Set to True if weights are in checkpoint/original model format and need
34 + layerwise processing. Set to False if weights have already been processed
35 + into kernel format (repacking, renaming, etc.)."""
36 +
37 +
38 +# API-level request classes (accept dicts for backend-agnostic serialization)
39 +@dataclass
40 +class WeightTransferInitRequest:
41 + """API-level weight transfer initialization request."""
42 +
43 + init_info: dict[str, Any] = field(default_factory=dict)
44 +
45 +
46 +@dataclass
47 +class WeightTransferUpdateRequest:
48 + """API-level weight update request."""
49 +
50 + update_info: dict[str, Any] = field(default_factory=dict)
51 +
52 +
53 +class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
54 + """
55 + Base class for weight transfer engines that handle transport of model weights
56 + from a trainer to inference workers.
57 +
58 + This abstraction separates weight transfer transport logic from the worker
59 + implementation, allowing different backends (NCCL, CUDA IPC[TODO], RDMA[TODO]) to be
60 + plugged in.
61 +
62 + Subclasses should define:
63 + init_info_cls: Type of backend-specific initialization info
64 + update_info_cls: Type of backend-specific update info
65 + """
66 +
67 + # Subclasses should override these class attributes
68 + init_info_cls: type[TInitInfo]
69 + update_info_cls: type[TUpdateInfo]
70 +
71 + def __init__(
72 + self, config: WeightTransferConfig, parallel_config: ParallelConfig
73 + ) -> None:
74 + """
75 + Initialize the weight transfer engine.
76 +
77 + Args:
78 + config: The configuration for the weight transfer engine
79 + parallel_config: The configuration for the parallel setup
80 + """
81 + self.config = config
82 + self.parallel_config = parallel_config
83 +
84 + def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
85 + """
86 + Construct typed init info from dict with validation.
87 +
88 + Args:
89 + init_dict: Dictionary containing backend-specific initialization parameters
90 +
91 + Returns:
92 + Typed backend-specific init info dataclass
93 +
94 + Raises:
95 + ValueError: If init_dict is invalid for this backend
96 + """
97 + try:
98 + return self.init_info_cls(**init_dict)
99 + except TypeError as e:
100 + raise ValueError(
101 + f"Invalid init_info for {self.__class__.__name__}: {e}"
102 + ) from e
103 +
104 + def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
105 + """
106 + Construct typed update info from dict with validation.
107 +
108 + Args:
109 + update_dict: Dictionary containing backend-specific update parameters
110 +
111 + Returns:
112 + Typed backend-specific update info dataclass
113 +
114 + Raises:
115 + ValueError: If update_dict is invalid for this backend
116 + """
117 + try:
118 + return self.update_info_cls(**update_dict)
119 + except TypeError as e:
120 + raise ValueError(
121 + f"Invalid update_info for {self.__class__.__name__}: {e}"
122 + ) from e
123 +
124 + @abstractmethod
125 + def init_transfer_engine(self, init_info: TInitInfo) -> None:
126 + """
127 + Initialize the weight transfer mechanism.
128 + This is called once at the beginning of training.
129 +
130 + Args:
131 + init_info: Backend-specific initialization info
132 + """
133 + raise NotImplementedError
134 +
135 + @abstractmethod
136 + def receive_weights(
137 + self,
138 + update_info: TUpdateInfo,
139 + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
140 + ) -> None:
141 + """
142 + Receive weights from the trainer and load them incrementally.
143 +
144 + Args:
145 + update_info: Backend-specific update info containing parameter metadata
146 + and any backend-specific data
147 + load_weights: Callable that loads weights into the model. Called
148 + incrementally for each weight to avoid OOM.
149 + """
150 + raise NotImplementedError
151 +
152 + @abstractmethod
153 + def shutdown(self) -> None:
154 + """
155 + Shutdown the weight transfer engine.
156 + This should be called when the worker is shutting down.
157 + """
158 + raise NotImplementedError
vllm/distributed/weight_transfer/factory.py new
+116
@@ -0,0 +1,116 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Factory for weight transfer engines with lazy loading."""
4 +
5 +import importlib
6 +from collections.abc import Callable
7 +from typing import TYPE_CHECKING
8 +
9 +from vllm.distributed.weight_transfer.base import WeightTransferEngine
10 +from vllm.logger import init_logger
11 +
12 +if TYPE_CHECKING:
13 + from vllm.config.parallel import ParallelConfig
14 + from vllm.config.weight_transfer import WeightTransferConfig
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +class WeightTransferEngineFactory:
20 + """Factory for creating weight transfer engines with lazy loading.
21 +
22 + This factory implements a registry pattern that supports:
23 + - Lazy loading: Engine modules are only imported when actually needed
24 + - Extensibility: Custom engines can be registered at runtime
25 + - Centralized registration: All built-in engines registered in one place
26 + """
27 +
28 + _registry: dict[str, Callable[[], type[WeightTransferEngine]]] = {}
29 +
30 + @classmethod
31 + def register_engine(
32 + cls,
33 + name: str,
34 + module_path_or_cls: str | type[WeightTransferEngine],
35 + class_name: str | None = None,
36 + ) -> None:
37 + """Register an engine with lazy-loading or direct class reference.
38 +
39 + Supports two calling conventions:
40 + 1. Lazy loading: register_engine(name, module_path, class_name)
41 + 2. Direct class: register_engine(name, engine_cls)
42 +
43 + Args:
44 + name: The name to register the engine under (e.g., "nccl")
45 + module_path_or_cls: Either a module path string for lazy loading,
46 + or the engine class directly
47 + class_name: Name of the engine class (required if module_path is string)
48 +
49 + Raises:
50 + ValueError: If an engine with the same name is already registered
51 + """
52 + if name in cls._registry:
53 + raise ValueError(f"Weight transfer engine '{name}' is already registered.")
54 +
55 + if isinstance(module_path_or_cls, str):
56 + # Lazy loading path
57 + module_path = module_path_or_cls
58 + if class_name is None:
59 + raise ValueError(
60 + "class_name is required when registering with module path"
61 + )
62 +
63 + def loader() -> type[WeightTransferEngine]:
64 + module = importlib.import_module(module_path)
65 + return getattr(module, class_name)
66 +
67 + cls._registry[name] = loader
68 + else:
69 + # Direct class registration
70 + engine_cls = module_path_or_cls
71 + cls._registry[name] = lambda: engine_cls
72 +
73 + @classmethod
74 + def create_engine(
75 + cls,
76 + config: "WeightTransferConfig",
77 + parallel_config: "ParallelConfig",
78 + ) -> WeightTransferEngine:
79 + """Create a weight transfer engine instance.
80 +
81 + Args:
82 + config: Weight transfer configuration containing the backend name
83 + parallel_config: Parallel configuration for the engine
84 +
85 + Returns:
86 + An initialized weight transfer engine instance
87 +
88 + Raises:
89 + ValueError: If the backend is not registered
90 + """
91 + backend = config.backend
92 + if backend not in cls._registry:
93 + available = list(cls._registry.keys())
94 + raise ValueError(
95 + f"Invalid weight transfer backend: {backend}. "
96 + f"Available engines: {available}"
97 + )
98 + engine_cls = cls._registry[backend]()
99 +
100 + logger.info(
101 + "Creating weight transfer engine: %s",
102 + engine_cls.__name__,
103 + )
104 +
105 + return engine_cls(config, parallel_config)
106 +
107 +
108 +# Register built-in weight transfer engines here.
109 +# Registration should be centralized to ensure lazy loading -
110 +# engine modules are only imported when actually used.
111 +
112 +WeightTransferEngineFactory.register_engine(
113 + "nccl",
114 + "vllm.distributed.weight_transfer.nccl_engine",
115 + "NCCLWeightTransferEngine",
116 +)
vllm/distributed/weight_transfer/nccl_engine.py new
+315
@@ -0,0 +1,315 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NCCL-based weight transfer engine."""
4 +
5 +from collections.abc import Callable, Iterator
6 +from dataclasses import dataclass
7 +from typing import TYPE_CHECKING, Any
8 +
9 +import torch
10 +
11 +if TYPE_CHECKING:
12 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
13 +
14 +from vllm.config.parallel import ParallelConfig
15 +from vllm.config.weight_transfer import WeightTransferConfig
16 +from vllm.distributed.weight_transfer.base import (
17 + WeightTransferEngine,
18 + WeightTransferInitInfo,
19 + WeightTransferUpdateInfo,
20 +)
21 +from vllm.distributed.weight_transfer.packed_tensor import (
22 + DEFAULT_PACKED_BUFFER_SIZE_BYTES,
23 + DEFAULT_PACKED_NUM_BUFFERS,
24 + packed_broadcast_consumer,
25 +)
26 +
27 +
28 +@dataclass
29 +class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
30 + """Initialization info for NCCL weight transfer backend."""
31 +
32 + master_address: str
33 + master_port: int
34 + rank_offset: int
35 + world_size: int
36 +
37 +
38 +@dataclass
39 +class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
40 + """Update info for NCCL weight transfer backend."""
41 +
42 + names: list[str]
43 + dtype_names: list[str]
44 + shapes: list[list[int]]
45 + packed: bool = False
46 + """Whether to use packed tensor broadcasting for efficiency.
47 + When True, multiple tensors are batched together before broadcasting
48 + to reduce NCCL communication overhead."""
49 + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
50 + """Size in bytes for each packed tensor buffer. Default is 1GB.
51 + Both producer and consumer must use the same value."""
52 + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
53 + """Number of buffers for double/triple buffering during packed transfer.
54 + Both producer and consumer must use the same value."""
55 +
56 + def __post_init__(self):
57 + """Validate that all lists have the same length."""
58 + num_params = len(self.names)
59 + if len(self.dtype_names) != num_params:
60 + raise ValueError(
61 + f"`dtype_names` should be of the same size as `names`: "
62 + f"got {len(self.dtype_names)} and {len(self.names)}"
63 + )
64 + if len(self.shapes) != num_params:
65 + raise ValueError(
66 + f"`shapes` should be of the same size as `names`: "
67 + f"got {len(self.shapes)} and {len(self.names)}"
68 + )
69 +
70 +
71 +class NCCLWeightTransferEngine(
72 + WeightTransferEngine[NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo]
73 +):
74 + """
75 + Weight transfer engine using NCCL for communication between trainer and workers.
76 +
77 + This implementation uses NCCL broadcast operations to transfer weights from
78 + the trainer (rank 0) to all inference workers in a process group.
79 + """
80 +
81 + # Define backend-specific dataclass types
82 + init_info_cls = NCCLWeightTransferInitInfo
83 + update_info_cls = NCCLWeightTransferUpdateInfo
84 +
85 + def __init__(
86 + self, config: WeightTransferConfig, parallel_config: ParallelConfig
87 + ) -> None:
88 + """
89 + Initialize the NCCL weight transfer engine.
90 +
91 + Args:
92 + config: The configuration for the weight transfer engine
93 + parallel_config: The configuration for the parallel setup
94 + """
95 + super().__init__(config, parallel_config)
96 + self.model_update_group: PyNcclCommunicator | None = None
97 +
98 + def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
99 + """
100 + Initialize NCCL process group with the trainer.
101 +
102 + Args:
103 + init_info: NCCL initialization info containing master address, port,
104 + rank offset, and world size
105 + """
106 +
107 + # Calculate the global rank in the trainer-worker process group
108 + # Must account for data parallel to get unique ranks across all workers
109 + dp_rank = self.parallel_config.data_parallel_rank
110 + world_size_per_dp = self.parallel_config.world_size # TP * PP
111 + rank_within_dp = self.parallel_config.rank
112 +
113 + # Unique rank across all DP groups
114 + worker_rank = dp_rank * world_size_per_dp + rank_within_dp
115 + rank = worker_rank + init_info.rank_offset
116 + # Create stateless process group
117 + self.model_update_group = (
118 + NCCLWeightTransferEngine._stateless_init_process_group(
119 + init_info.master_address,
120 + init_info.master_port,
121 + rank,
122 + init_info.world_size,
123 + torch.cuda.current_device(),
124 + )
125 + )
126 +
127 + def receive_weights(
128 + self,
129 + update_info: NCCLWeightTransferUpdateInfo,
130 + load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
131 + ) -> None:
132 + """
133 + Receive weights from trainer via NCCL broadcast and load them incrementally.
134 +
135 + If update_info.packed is True, uses packed tensor broadcasting for
136 + efficient transfer of multiple weights in batches. Otherwise, uses simple
137 + one-by-one broadcasting.
138 +
139 + Args:
140 + update_info: NCCL update info containing parameter names, dtypes, shapes,
141 + and packed flag
142 + load_weights: Callable that loads weights into the model. Called
143 + incrementally for each batch of weights to avoid OOM.
144 + """
145 + if self.model_update_group is None:
146 + raise RuntimeError(
147 + "NCCL weight transfer not initialized. "
148 + "Call init_transfer_engine() first."
149 + )
150 +
151 + if update_info.packed:
152 + # Build iterator of (name, (shape, dtype)) from update_info
153 + def state_dict_info_iterator():
154 + for name, dtype_name, shape in zip(
155 + update_info.names, update_info.dtype_names, update_info.shapes
156 + ):
157 + dtype = getattr(torch, dtype_name)
158 + yield (name, (shape, dtype))
159 +
160 + packed_broadcast_consumer(
161 + iterator=state_dict_info_iterator(),
162 + group=self.model_update_group,
163 + src=0,
164 + post_unpack_func=load_weights,
165 + buffer_size_bytes=update_info.packed_buffer_size_bytes,
166 + num_buffers=update_info.packed_num_buffers,
167 + )
168 + else:
169 + # Use simple one-by-one broadcasting
170 + for name, dtype_name, shape in zip(
171 + update_info.names, update_info.dtype_names, update_info.shapes
172 + ):
173 + dtype = getattr(torch, dtype_name)
174 + weight = torch.empty(shape, dtype=dtype, device="cuda")
175 + self.model_update_group.broadcast(
176 + weight, src=0, stream=torch.cuda.current_stream()
177 + )
178 + load_weights([(name, weight)])
179 + del weight
180 +
181 + def shutdown(self) -> None:
182 + if self.model_update_group is not None:
183 + # Clean up the communicator by removing the reference
184 + self.model_update_group = None
185 +
186 + @staticmethod
187 + def trainer_send_weights(
188 + 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,
197 + ) -> None:
198 + """Broadcast weights from trainer to vLLM workers.
199 +
200 + Args:
201 + 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.
215 +
216 + Example:
217 + >>> from vllm.distributed.weight_transfer.nccl_engine import (
218 + ... NCCLWeightTransferEngine,
219 + ... )
220 + >>> param_iter = ((n, p) for n, p in model.named_parameters())
221 + >>> NCCLWeightTransferEngine.trainer_send_weights(
222 + ... param_iter, group, packed=True
223 + ... )
224 + """
225 + if post_iter_func is None:
226 + # Default: extract just the tensor from (name, tensor) tuple
227 + post_iter_func = lambda x: x[1]
228 +
229 + if packed:
230 + # Use packed tensor broadcasting for efficiency
231 + from vllm.distributed.weight_transfer.packed_tensor import (
232 + packed_broadcast_producer,
233 + )
234 +
235 + packed_broadcast_producer(
236 + iterator=iterator,
237 + group=group,
238 + src=src,
239 + post_iter_func=post_iter_func,
240 + buffer_size_bytes=packed_buffer_size_bytes,
241 + num_buffers=packed_num_buffers,
242 + )
243 + else:
244 + # Use simple one-by-one broadcasting
245 + for item in iterator:
246 + tensor = post_iter_func(item)
247 + group.broadcast(
248 + tensor, src=src, stream=stream or torch.cuda.current_stream()
249 + )
250 +
251 + @staticmethod
252 + def trainer_init(
253 + init_info: NCCLWeightTransferInitInfo | dict,
254 + ) -> "PyNcclCommunicator":
255 + """
256 + Initialize NCCL process group for trainer-side weight transfer.
257 +
258 + The trainer is always rank 0 in the process group. Uses the current
259 + CUDA device (torch.cuda.current_device()).
260 +
261 + Args:
262 + init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys:
263 + - master_address: str
264 + - master_port: int
265 + - world_size: int
266 +
267 + Returns:
268 + PyNcclCommunicator for weight transfer.
269 +
270 + Example:
271 + >>> from vllm.distributed.weight_transfer.nccl_engine import (
272 + ... NCCLWeightTransferEngine,
273 + ... )
274 + >>> group = NCCLWeightTransferEngine.trainer_init(
275 + ... dict(
276 + ... master_address=master_address,
277 + ... master_port=master_port,
278 + ... world_size=world_size,
279 + ... ),
280 + ... )
281 + """
282 + if isinstance(init_info, dict):
283 + master_address = init_info["master_address"]
284 + master_port = init_info["master_port"]
285 + world_size = init_info["world_size"]
286 + else:
287 + # NCCLWeightTransferInitInfo object
288 + master_address = init_info.master_address
289 + master_port = init_info.master_port
290 + world_size = init_info.world_size
291 +
292 + # Trainer is always rank 0
293 + return NCCLWeightTransferEngine._stateless_init_process_group(
294 + master_address, master_port, 0, world_size, torch.cuda.current_device()
295 + )
296 +
297 + @staticmethod
298 + def _stateless_init_process_group(
299 + master_address, master_port, rank, world_size, device
300 + ):
301 + """
302 + vLLM provides `StatelessProcessGroup` to create a process group
303 + without considering the global process group in torch.distributed.
304 + It is recommended to create `StatelessProcessGroup`, and then initialize
305 + the data-plane communication (NCCL) between external (train processes)
306 + and vLLM workers.
307 + """
308 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
309 + from vllm.distributed.utils import StatelessProcessGroup
310 +
311 + pg = StatelessProcessGroup.create(
312 + host=master_address, port=master_port, rank=rank, world_size=world_size
313 + )
314 + pynccl = PyNcclCommunicator(pg, device=device)
315 + return pynccl
vllm/distributed/weight_transfer/packed_tensor.py new
+216
@@ -0,0 +1,216 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Packed tensor utilities for efficient weight transfer."""
4 +
5 +import math
6 +from collections.abc import Callable, Iterator
7 +from typing import Any
8 +
9 +import torch
10 +
11 +# Default values for packed tensor configuration.
12 +# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights.
13 +DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB
14 +DEFAULT_PACKED_NUM_BUFFERS = 2
15 +
16 +
17 +def packed_broadcast_producer(
18 + iterator: Iterator[tuple[str, torch.Tensor]],
19 + group: Any,
20 + src: int,
21 + post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor],
22 + buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
23 + num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
24 +) -> None:
25 + """Broadcast tensors in a packed manner from trainer to workers.
26 +
27 + Args:
28 + iterator: Iterator of model parameters. Returns a tuple of (name, tensor)
29 + group: Process group (PyNcclCommunicator)
30 + src: Source rank (0 in current implementation)
31 + post_iter_func: Function to apply to each (name, tensor) pair before
32 + packing, should return a tensor
33 + buffer_size_bytes: Size in bytes for each packed tensor buffer.
34 + Both producer and consumer must use the same value.
35 + num_buffers: Number of buffers for double/triple buffering.
36 + Both producer and consumer must use the same value.
37 +
38 + """
39 + target_packed_tensor_size = buffer_size_bytes
40 +
41 + streams = [torch.cuda.Stream() for _ in range(num_buffers)]
42 + buffer_idx = 0
43 +
44 + packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)]
45 + packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
46 + packed_tensors: list[torch.Tensor] = [
47 + torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers)
48 + ]
49 +
50 + while True:
51 + # Synchronize the current stream
52 + streams[buffer_idx].synchronize()
53 + # Start tasks for the new buffer in a new stream
54 + with torch.cuda.stream(streams[buffer_idx]):
55 + try:
56 + # Initialize the packing tensor list and sizes
57 + packing_tensor_list[buffer_idx] = []
58 + packing_tensor_sizes[buffer_idx] = 0
59 + # Pack the tensors
60 + while True:
61 + # Apply post processing and convert to linearized uint8 tensor
62 + tensor = (
63 + post_iter_func(next(iterator))
64 + .contiguous()
65 + .view(torch.uint8)
66 + .view(-1)
67 + )
68 + packing_tensor_list[buffer_idx].append(tensor)
69 + packing_tensor_sizes[buffer_idx] += tensor.numel()
70 + if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
71 + break
72 + # Pack the tensors and call broadcast collective
73 + packed_tensors[buffer_idx] = torch.cat(
74 + packing_tensor_list[buffer_idx], dim=0
75 + )
76 + group.broadcast(packed_tensors[buffer_idx], src=src)
77 + # Move to the next buffer
78 + buffer_idx = (buffer_idx + 1) % num_buffers
79 + except StopIteration:
80 + # Do the last broadcast if there are remaining tensors
81 + if len(packing_tensor_list[buffer_idx]) > 0:
82 + packed_tensors[buffer_idx] = torch.cat(
83 + packing_tensor_list[buffer_idx], dim=0
84 + )
85 + group.broadcast(packed_tensors[buffer_idx], src=src)
86 + break
87 +
88 +
89 +def packed_broadcast_consumer(
90 + iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]],
91 + group: Any,
92 + src: int,
93 + post_unpack_func: Callable[[list[tuple[str, torch.Tensor]]], None],
94 + buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
95 + num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
96 +) -> None:
97 + """Consume packed tensors and unpack them into a list of tensors.
98 +
99 + Args:
100 + iterator: Iterator of parameter metadata. Returns (name, (shape, dtype))
101 + group: Process group (PyNcclCommunicator)
102 + src: Source rank (0 in current implementation)
103 + post_unpack_func: Function to apply to each list of (name, tensor) after
104 + unpacking
105 + buffer_size_bytes: Size in bytes for each packed tensor buffer.
106 + Both producer and consumer must use the same value.
107 + num_buffers: Number of buffers for double/triple buffering.
108 + Both producer and consumer must use the same value.
109 +
110 + """
111 +
112 + def unpack_tensor(
113 + packed_tensor: torch.Tensor,
114 + names: list[str],
115 + shapes: list[list[int]],
116 + dtypes: list[torch.dtype],
117 + tensor_sizes: list[int],
118 + ) -> list[tuple[str, torch.Tensor]]:
119 + """Unpack a single tensor into a list of tensors.
120 +
121 + Args:
122 + packed_tensor: The packed torch.uint8 tensor to unpack
123 + names: List of tensor names
124 + shapes: List of tensor shapes
125 + dtypes: List of tensor dtypes
126 + tensor_sizes: List of tensor sizes in bytes
127 +
128 + Returns:
129 + unpacked List[(name, tensor)]
130 + """
131 + unpacked_tensors = packed_tensor.split(tensor_sizes)
132 +
133 + unpacked_list = [
134 + (name, tensor.contiguous().view(dtype).view(*shape))
135 + for name, shape, dtype, tensor in zip(
136 + names, shapes, dtypes, unpacked_tensors
137 + )
138 + ]
139 +
140 + return unpacked_list
141 +
142 + target_packed_tensor_size = buffer_size_bytes
143 +
144 + streams = [torch.cuda.Stream() for _ in range(num_buffers)]
145 + buffer_idx = 0
146 +
147 + packing_tensor_meta_data: list[list[tuple[str, list[int], torch.dtype, int]]] = [
148 + [] for _ in range(num_buffers)
149 + ]
150 + packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
151 + packed_tensors: list[torch.Tensor] = [
152 + torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers)
153 + ]
154 +
155 + while True:
156 + # Synchronize the current stream
157 + streams[buffer_idx].synchronize()
158 + with torch.cuda.stream(streams[buffer_idx]):
159 + # Initialize the packing tensor meta data
160 + packing_tensor_meta_data[buffer_idx] = []
161 + packing_tensor_sizes[buffer_idx] = 0
162 + try:
163 + # Form a packed tensor
164 + while True:
165 + name, (shape, dtype) = next(iterator)
166 + tensor_size = math.prod(shape) * dtype.itemsize
167 + packing_tensor_meta_data[buffer_idx].append(
168 + (name, shape, dtype, tensor_size)
169 + )
170 + packing_tensor_sizes[buffer_idx] += tensor_size
171 + if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
172 + break
173 + # Create a packed tensor and broadcast it
174 + packed_tensors[buffer_idx] = torch.empty(
175 + packing_tensor_sizes[buffer_idx], dtype=torch.uint8, device="cuda"
176 + )
177 + group.broadcast(packed_tensors[buffer_idx], src=src)
178 + # Load the packed tensor into the model
179 + names, shapes, dtypes, tensor_sizes = zip(
180 + *packing_tensor_meta_data[buffer_idx]
181 + )
182 + post_unpack_func(
183 + unpack_tensor(
184 + packed_tensors[buffer_idx],
185 + list(names),
186 + list(shapes),
187 + list(dtypes),
188 + list(tensor_sizes),
189 + )
190 + )
191 + # Move to the next buffer
192 + buffer_idx = (buffer_idx + 1) % num_buffers
193 + except StopIteration:
194 + # Do the last broadcast if there are remaining tensors
195 + if len(packing_tensor_meta_data[buffer_idx]) > 0:
196 + # Create a packed tensor and broadcast it
197 + packed_tensors[buffer_idx] = torch.empty(
198 + packing_tensor_sizes[buffer_idx],
199 + dtype=torch.uint8,
200 + device="cuda",
201 + )
202 + group.broadcast(packed_tensors[buffer_idx], src=src)
203 + # Load the packed tensor into the model
204 + names, shapes, dtypes, tensor_sizes = zip(
205 + *packing_tensor_meta_data[buffer_idx]
206 + )
207 + post_unpack_func(
208 + unpack_tensor(
209 + packed_tensors[buffer_idx],
210 + list(names),
211 + list(shapes),
212 + list(dtypes),
213 + list(tensor_sizes),
214 + )
215 + )
216 + break
vllm/engine/arg_utils.py
+14
@@ -54,6 +54,7 @@ from vllm.config import (
54 SpeculativeConfig,
55 StructuredOutputsConfig,
56 VllmConfig,
57 + WeightTransferConfig,
58 get_attr_docs,
59 )
60 from vllm.config.cache import (
@@ -581,6 +582,11 @@ class EngineArgs:
582 kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend
583 tokens_only: bool = False
584
585 + weight_transfer_config: WeightTransferConfig | None = None
586 + """Configuration for weight transfer during RL training.
587 + Accepts a JSON string or dict with backend-specific options.
588 + Example: '{"backend": "nccl"}'"""
589 +
590 def __post_init__(self):
591 # support `EngineArgs(compilation_config={...})`
592 # without having to manually construct a
@@ -591,6 +597,10 @@ class EngineArgs:
597 self.attention_config = AttentionConfig(**self.attention_config)
598 if isinstance(self.eplb_config, dict):
599 self.eplb_config = EPLBConfig(**self.eplb_config)
600 + if isinstance(self.weight_transfer_config, dict):
601 + self.weight_transfer_config = WeightTransferConfig(
602 + **self.weight_transfer_config
603 + )
604 # Setup plugins
605 from vllm.plugins import load_general_plugins
606
@@ -1189,6 +1199,9 @@ class EngineArgs:
1199 vllm_group.add_argument(
1200 "--optimization-level", **vllm_kwargs["optimization_level"]
1201 )
1202 + vllm_group.add_argument(
1203 + "--weight-transfer-config", **vllm_kwargs["weight_transfer_config"]
1204 + )
1205
1206 # Other arguments
1207 parser.add_argument(
@@ -1765,6 +1778,7 @@ class EngineArgs:
1778 profiler_config=self.profiler_config,
1779 additional_config=self.additional_config,
1780 optimization_level=self.optimization_level,
1781 + weight_transfer_config=self.weight_transfer_config,
1782 )
1783
1784 return config
vllm/engine/protocol.py
+14
@@ -6,6 +6,10 @@ from collections.abc import AsyncGenerator, Iterable, Mapping
6 from typing import Any
7
8 from vllm.config import ModelConfig, VllmConfig
9 +from vllm.distributed.weight_transfer.base import (
10 + WeightTransferInitRequest,
11 + WeightTransferUpdateRequest,
12 +)
13 from vllm.inputs.data import PromptType, StreamingInput
14 from vllm.lora.request import LoRARequest
15 from vllm.outputs import PoolingRequestOutput, RequestOutput
@@ -191,3 +195,13 @@ class EngineClient(ABC):
195 async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
196 """Get supported tasks"""
197 raise NotImplementedError
198 +
199 + async def init_weight_transfer_engine(
200 + self, init_request: WeightTransferInitRequest
201 + ) -> None:
202 + """Initialize weight transfer for RL training."""
203 + raise NotImplementedError
204 +
205 + async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
206 + """Batched weight update for RL training."""
207 + raise NotImplementedError
vllm/entrypoints/llm.py
+53
@@ -34,6 +34,10 @@ from vllm.config.model import (
34 RunnerOption,
35 TokenizerMode,
36 )
37 +from vllm.distributed.weight_transfer.base import (
38 + WeightTransferInitRequest,
39 + WeightTransferUpdateRequest,
40 +)
41 from vllm.engine.arg_utils import EngineArgs
42 from vllm.entrypoints.chat_utils import (
43 ChatCompletionMessageParam,
@@ -360,6 +364,23 @@ class LLM:
364 def get_tokenizer(self) -> TokenizerLike:
365 return self.llm_engine.get_tokenizer()
366
367 + def get_world_size(self, include_dp: bool = True) -> int:
368 + """Get the world size from the parallel config.
369 +
370 + Args:
371 + include_dp: If True (default), returns the world size including
372 + data parallelism (TP * PP * DP). If False, returns the world
373 + size without data parallelism (TP * PP).
374 +
375 + Returns:
376 + The world size (tensor_parallel_size * pipeline_parallel_size),
377 + optionally multiplied by data_parallel_size if include_dp is True.
378 + """
379 + parallel_config = self.llm_engine.vllm_config.parallel_config
380 + if include_dp:
381 + return parallel_config.world_size_across_dp
382 + return parallel_config.world_size
383 +
384 def reset_mm_cache(self) -> None:
385 self.input_processor.clear_mm_cache()
386 self.llm_engine.reset_mm_cache()
@@ -1903,6 +1924,38 @@ class LLM:
1924 # its previous requests.
1925 return sorted(outputs, key=lambda x: int(x.request_id))
1926
1927 + def init_weight_transfer_engine(
1928 + self, request: WeightTransferInitRequest | dict
1929 + ) -> None:
1930 + """
1931 + Initialize weight transfer for RL training.
1932 +
1933 + Args:
1934 + request: Weight transfer initialization request with backend-specific info
1935 + """
1936 + init_info_dict = (
1937 + request["init_info"] if isinstance(request, dict) else request.init_info
1938 + )
1939 +
1940 + self.llm_engine.collective_rpc(
1941 + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
1942 + )
1943 +
1944 + def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None:
1945 + """
1946 + Update the weights of the model.
1947 +
1948 + Args:
1949 + request: Weight update request with backend-specific update info
1950 + """
1951 + update_info_dict = (
1952 + request["update_info"] if isinstance(request, dict) else request.update_info
1953 + )
1954 +
1955 + self.llm_engine.collective_rpc(
1956 + "update_weights", kwargs={"update_info": update_info_dict}
1957 + )
1958 +
1959 def __repr__(self) -> str:
1960 """Return a transformers-style hierarchical view of the model."""
1961 # Cache the result to avoid repeated collective_rpc calls
vllm/entrypoints/serve/rlhf/api_router.py
+65 -2
@@ -1,12 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -
4 +import json
5 from http import HTTPStatus
6
7 -from fastapi import APIRouter, FastAPI, Query, Request
7 +from fastapi import APIRouter, FastAPI, HTTPException, Query, Request
8 from fastapi.responses import JSONResponse
9
10 +import vllm.envs as envs
11 +from vllm.distributed.weight_transfer.base import (
12 + WeightTransferInitRequest,
13 + WeightTransferUpdateRequest,
14 +)
15 from vllm.engine.protocol import EngineClient
16 from vllm.logger import init_logger
17
@@ -98,5 +103,63 @@ async def is_paused(raw_request: Request) -> JSONResponse:
103 return JSONResponse(content={"is_paused": paused})
104
105
106 +@router.post("/init_weight_transfer_engine")
107 +async def init_weight_transfer_engine(raw_request: Request):
108 + try:
109 + body = await raw_request.json()
110 + except json.JSONDecodeError as e:
111 + raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904
112 + init_info = body.get("init_info")
113 + if init_info is None:
114 + raise HTTPException(
115 + status_code=HTTPStatus.BAD_REQUEST.value,
116 + detail="Missing 'init_info' in request body",
117 + )
118 + await engine_client(raw_request).init_weight_transfer_engine(
119 + WeightTransferInitRequest(init_info=init_info)
120 + )
121 + return JSONResponse(content={"message": "Weight transfer initialized"})
122 +
123 +
124 +@router.post("/update_weights")
125 +async def update_weights(raw_request: Request):
126 + try:
127 + body = await raw_request.json()
128 + except json.JSONDecodeError as e:
129 + raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904
130 + update_info = body.get("update_info")
131 + if update_info is None:
132 + raise HTTPException(
133 + status_code=HTTPStatus.BAD_REQUEST.value,
134 + detail="Missing 'update_info' in request body",
135 + )
136 + await engine_client(raw_request).update_weights(
137 + request=WeightTransferUpdateRequest(update_info=update_info)
138 + )
139 + return JSONResponse(content={"message": "Weights updated"})
140 +
141 +
142 +@router.get("/get_world_size")
143 +async def get_world_size(
144 + raw_request: Request,
145 + include_dp: bool = Query(True),
146 +):
147 + """Get the world size from the parallel config.
148 +
149 + Args:
150 + include_dp: If True (default), returns the world size including
151 + data parallelism (TP * PP * DP). If False, returns the world
152 + size without data parallelism (TP * PP).
153 + """
154 + parallel_config = engine_client(raw_request).vllm_config.parallel_config
155 + if include_dp:
156 + world_size = parallel_config.world_size_across_dp
157 + else:
158 + world_size = parallel_config.world_size
159 + return JSONResponse(content={"world_size": world_size})
160 +
161 +
162 def attach_router(app: FastAPI):
163 + if not envs.VLLM_SERVER_DEV_MODE:
164 + return
165 app.include_router(router)
vllm/model_executor/layers/quantization/fp8.py
+9
@@ -649,6 +649,9 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
649 )
650 # Activations not quantized for marlin.
651
652 + # Prevent duplicate processing (e.g., during weight reload)
653 + layer._already_called_process_weights_after_loading = True
654 +
655
656 class Fp8MoEMethod(FusedMoEMethodBase):
657 """MoE method for FP8.
@@ -908,6 +911,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
911 layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
912 )
913
914 + # Prevent duplicate processing (e.g., during weight reload)
915 + layer._already_called_process_weights_after_loading = True
916 +
917 def maybe_make_prepare_finalize(
918 self,
919 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
@@ -1241,6 +1247,9 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1247 layer.w2_input_scale,
1248 )
1249
1250 + # Prevent duplicate processing (e.g., during weight reload)
1251 + layer._already_called_process_weights_after_loading = True
1252 +
1253
1254 class Fp8KVCacheMethod(BaseKVCacheMethod):
1255 """
vllm/model_executor/model_loader/reload/layerwise.py
+5
@@ -216,6 +216,11 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
216 # Materialize layer tensors onto device
217 materialize_layer(layer)
218
219 + # Reset FP8 online quantization flag so process_weights_after_loading
220 + # will run again during reload
221 + if hasattr(layer, "_already_called_process_weights_after_loading"):
222 + delattr(layer, "_already_called_process_weights_after_loading")
223 +
224 # Unwrap layerwise loading wrappers
225 for param in get_layer_tensors(layer).values():
226 param.weight_loader = _get_original_loader(param)
vllm/v1/engine/async_llm.py
+45
@@ -14,6 +14,10 @@ import torch
14 import vllm.envs as envs
15 from vllm import TokensPrompt
16 from vllm.config import VllmConfig
17 +from vllm.distributed.weight_transfer.base import (
18 + WeightTransferInitRequest,
19 + WeightTransferUpdateRequest,
20 +)
21 from vllm.engine.arg_utils import AsyncEngineArgs
22 from vllm.engine.protocol import EngineClient
23 from vllm.inputs import PromptType, StreamingInput
@@ -1011,3 +1015,44 @@ class AsyncLLM(EngineClient):
1015 @property
1016 def dead_error(self) -> BaseException:
1017 return EngineDeadError()
1018 +
1019 + async def init_weight_transfer_engine(
1020 + self, request: WeightTransferInitRequest
1021 + ) -> None:
1022 + """
1023 + Initialize weight transfer for RL training.
1024 +
1025 + Args:
1026 + request: Weight transfer initialization request with backend-specific info
1027 + """
1028 + from vllm.distributed.weight_transfer.base import (
1029 + WeightTransferInitRequest,
1030 + )
1031 +
1032 + if isinstance(request, WeightTransferInitRequest):
1033 + init_info_dict = request.init_info
1034 + else:
1035 + raise TypeError(f"Expected WeightTransferInitRequest, got {type(request)}")
1036 +
1037 + await self.collective_rpc(
1038 + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
1039 + )
1040 +
1041 + async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
1042 + """
1043 + Batched weight update for RL training.
1044 +
1045 + Args:
1046 + request: Weight update request with backend-specific update info
1047 + """
1048 +
1049 + if isinstance(request, WeightTransferUpdateRequest):
1050 + update_info_dict = request.update_info
1051 + else:
1052 + raise TypeError(
1053 + f"Expected WeightTransferUpdateRequest, got {type(request)}"
1054 + )
1055 +
1056 + await self.collective_rpc(
1057 + "update_weights", kwargs={"update_info": update_info_dict}
1058 + )
vllm/v1/worker/gpu_worker.py
+77
@@ -33,6 +33,7 @@ from vllm.distributed.parallel_state import (
33 get_pp_group,
34 get_tp_group,
35 )
36 +from vllm.distributed.weight_transfer import WeightTransferEngineFactory
37 from vllm.logger import init_logger
38 from vllm.lora.request import LoRARequest
39 from vllm.model_executor.models.interfaces import is_mixture_of_experts
@@ -89,6 +90,16 @@ class Worker(WorkerBase):
90 # Buffers saved before sleep
91 self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
92
93 + # Weight transfer engine (initialized on-demand)
94 + self.weight_transfer_engine = (
95 + WeightTransferEngineFactory.create_engine(
96 + self.vllm_config.weight_transfer_config,
97 + self.vllm_config.parallel_config,
98 + )
99 + if self.vllm_config.weight_transfer_config is not None
100 + else None
101 + )
102 +
103 # Torch/CUDA profiler. Enabled and configured through profiler_config.
104 self.profiler: Any | None = None
105 profiler_config = vllm_config.profiler_config
@@ -932,6 +943,69 @@ class Worker(WorkerBase):
943 tensorizer_config=tensorizer_config,
944 )
945
946 + def init_weight_transfer_engine(self, init_info: dict) -> None:
947 + """
948 + Initialize weight transfer mechanism.
949 + For NCCL backend, this creates a process group with the trainer.
950 +
951 + Args:
952 + init_info: Dictionary containing backend-specific initialization info
953 + """
954 + if self.weight_transfer_engine is None:
955 + raise RuntimeError(
956 + "Weight transfer not configured. "
957 + "Please set weight_transfer_config to enable weight transfer."
958 + )
959 + # Parse dict into backend-specific typed dataclass
960 + typed_init_info = self.weight_transfer_engine.parse_init_info(init_info)
961 + self.weight_transfer_engine.init_transfer_engine(typed_init_info)
962 +
963 + def update_weights(self, update_info: dict) -> None:
964 + """
965 + Batched weight update from the trainer.
966 +
967 + Args:
968 + update_info: Dictionary containing backend-specific update info
969 + """
970 + if self.weight_transfer_engine is None:
971 + raise RuntimeError(
972 + "Weight transfer not configured. "
973 + "Please set weight_transfer_config to enable weight transfer."
974 + )
975 +
976 + # Parse dict into backend-specific typed dataclass
977 + typed_update_info = self.weight_transfer_engine.parse_update_info(update_info)
978 +
979 + model = self.model_runner.model
980 +
981 + if typed_update_info.is_checkpoint_format:
982 + from vllm.model_executor.model_loader.reload import (
983 + finalize_layerwise_reload,
984 + initialize_layerwise_reload,
985 + )
986 +
987 + # Use layerwise reload pattern for checkpoint format weights
988 + with torch.device(self.device):
989 + initialize_layerwise_reload(model)
990 + self.weight_transfer_engine.receive_weights(
991 + typed_update_info,
992 + load_weights=model.load_weights,
993 + )
994 + finalize_layerwise_reload(model, self.model_config)
995 + else:
996 + # Weights are already in kernel format, copy directly
997 + def load_weights_direct(
998 + weights: list[tuple[str, torch.Tensor]],
999 + ) -> None:
1000 + for name, weight in weights:
1001 + param = model.get_parameter(name)
1002 + param.copy_(weight)
1003 +
1004 + self.weight_transfer_engine.receive_weights(
1005 + typed_update_info,
1006 + load_weights=load_weights_direct,
1007 + )
1008 +
1009 def shutdown(self) -> None:
1010 # has_kv_transfer_group can be None during interpreter shutdown.
1011 if ensure_kv_transfer_shutdown is not None:
@@ -939,6 +1013,9 @@ class Worker(WorkerBase):
1013 if self.profiler is not None:
1014 self.profiler.shutdown()
1015
1016 + if weight_transfer_engine := getattr(self, "weight_transfer_engine", None):
1017 + weight_transfer_engine.shutdown()
1018 +
1019
1020 def init_worker_distributed_environment(
1021 vllm_config: VllmConfig,