@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 266b9d9c64
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 2, 07:15 UTC

[Frontend][Core] Add sparse NCCL weight transfer support for in-place updates (#40096)

Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Siddharth Bedekar committed Jun 1, 2026 at 12:37 UTC 266b9d9c64ddb64c719d422280d4522382a229d3
12 files changed +1425 -77
docs/training/weight_transfer/nccl.md
+18 -3
@@ -84,7 +84,10 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT
84
85 ## Receiving Weights (Inference Side)
86
87 -The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
87 +The inference side triggers weight reception using the four-phase protocol:
88 +`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
89 +`finish_weight_update`. The init phase is shown [above](#initialization). The
90 +remaining three steps are:
91
92 ```python
93 from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
@@ -108,12 +111,24 @@ llm.update_weights(
111 llm.finish_weight_update()
112 ```
113
111 -The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
114 +The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
115 +must match the order in which the trainer iterates over its parameters.
116
113 -`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
117 +`start_weight_update` must be called before `update_weights`, and
118 +`finish_weight_update` must be called after all weight chunks have been
119 +transferred. The `is_checkpoint_format` flag controls whether layerwise reload
120 +processing is applied (`True` for checkpoint-format weights, `False` for
121 +pre-processed kernel-format weights).
122 +
123 +Sparse NCCL patches still use `update_kind="sparse_flat"` inside
124 +`update_info`, but they should be wrapped in
125 +`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
126 +directly to runtime/kernel-format parameters. The current sparse MVP requires
127 +`TP=1` and `PP=1`.
128
129 ## Examples
130
131 - [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
132 +- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
133 - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
134 - [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
examples/rl/rlhf_sparse_nccl.py new
+526
@@ -0,0 +1,526 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
5 +
6 +This example mirrors the validation story used for the sparse NCCL MVP:
7 +both the dense update path and the sparse patch path start from the same real
8 +checkpoint and apply the same deterministic trainer-side patch. The script then
9 +checks that greedy 1-token outputs match between the dense and sparse vLLM
10 +engines after the update.
11 +
12 +The example performs the following steps:
13 +* Load a training model on one GPU via a Ray actor.
14 +* Launch a vLLM engine with the same real model on a second GPU.
15 +* Verify trainer vs vLLM baseline agreement before any update.
16 +* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
17 +* Run a dense NCCL update into a fresh vLLM engine and collect post-update
18 + outputs.
19 +* Reset the trainer back to the baseline checkpoint.
20 +* Apply the same deterministic patch again.
21 +* Run a sparse NCCL update into another fresh vLLM engine and collect
22 + post-update outputs.
23 +* Compare dense vs sparse baseline outputs, dense vs sparse post-update
24 + outputs, estimated payload sizes, and trainer-side send times.
25 +
26 +Current sparse weight transfer MVP limitations:
27 +* ``TP=1`` and ``PP=1`` only
28 +* sparse updates use runtime/kernel-format parameter names
29 +* sparse updates are not composable with checkpoint-format or packed updates
30 +
31 +This example assumes a single-node cluster with two GPUs.
32 +"""
33 +
34 +import hashlib
35 +import os
36 +import time
37 +from collections.abc import Sequence
38 +
39 +import ray
40 +import torch
41 +from ray.util.placement_group import placement_group
42 +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
43 +from transformers import AutoModelForCausalLM, AutoTokenizer
44 +
45 +from vllm import LLM, SamplingParams
46 +from vllm.config import WeightTransferConfig
47 +from vllm.distributed.weight_transfer.base import SparseWeightPatch
48 +from vllm.distributed.weight_transfer.nccl_engine import (
49 + NCCLTrainerSendWeightsArgs,
50 + NCCLWeightTransferEngine,
51 +)
52 +from vllm.utils.network_utils import get_ip, get_open_port
53 +
54 +MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
55 +PATCHED_PARAM_NAME = "model.embed_tokens.weight"
56 +MAX_PATCH_ROWS = 32
57 +PROMPTS = [
58 + "Hello, my name is",
59 + "The president of the United States is",
60 + "The capital of France is",
61 + "The future of AI is",
62 +]
63 +SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
64 +
65 +
66 +class MyLLM(LLM):
67 + """Configure the vLLM worker for Ray placement group execution."""
68 +
69 + def __init__(self, *args, **kwargs):
70 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
71 + super().__init__(*args, **kwargs)
72 +
73 +
74 +@ray.remote(num_gpus=1)
75 +class TrainModel:
76 + """Ray actor that owns the trainer-side model and deterministic patch state."""
77 +
78 + def __init__(self, model_name: str):
79 + self.model_name = model_name
80 + self.tokenizer = AutoTokenizer.from_pretrained(model_name)
81 + if self.tokenizer.pad_token_id is None:
82 + self.tokenizer.pad_token = self.tokenizer.eos_token
83 +
84 + self.model = None
85 + self.patched_param = None
86 + self.pending_sparse_patches: list[SparseWeightPatch] | None = None
87 + self.model_update_group = None
88 + self.master_address = get_ip()
89 + self.port = get_open_port()
90 + self.reset_model()
91 +
92 + def reset_model(self) -> None:
93 + self.model = AutoModelForCausalLM.from_pretrained(
94 + self.model_name,
95 + torch_dtype=torch.bfloat16,
96 + ).to("cuda:0")
97 + self.model.eval()
98 +
99 + try:
100 + self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
101 + except AttributeError as exc:
102 + raise RuntimeError(
103 + f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
104 + ) from exc
105 +
106 + self.pending_sparse_patches = None
107 +
108 + def create_rendezvous(self) -> tuple[str, int]:
109 + self.port = get_open_port()
110 + return self.master_address, self.port
111 +
112 + def init_weight_transfer_group(self, world_size: int) -> None:
113 + self.model_update_group = NCCLWeightTransferEngine.trainer_init(
114 + dict(
115 + master_address=self.master_address,
116 + master_port=self.port,
117 + world_size=world_size,
118 + )
119 + )
120 +
121 + def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
122 + names = []
123 + dtype_names = []
124 + shapes = []
125 + payload_bytes = 0
126 + for name, param in self.model.named_parameters():
127 + names.append(name)
128 + dtype_names.append(str(param.dtype).split(".")[-1])
129 + shapes.append(list(param.shape))
130 + payload_bytes += param.numel() * param.element_size()
131 +
132 + return (
133 + dict(
134 + names=names,
135 + dtype_names=dtype_names,
136 + shapes=shapes,
137 + packed=packed,
138 + ),
139 + payload_bytes,
140 + )
141 +
142 + @torch.inference_mode()
143 + def generate(
144 + self,
145 + prompts: Sequence[str],
146 + max_new_tokens: int = 1,
147 + ) -> list[dict[str, object]]:
148 + generations = []
149 + for prompt in prompts:
150 + model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
151 + output = self.model.generate(
152 + **model_inputs,
153 + max_new_tokens=max_new_tokens,
154 + do_sample=False,
155 + pad_token_id=self.tokenizer.pad_token_id,
156 + )
157 + new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
158 + generations.append(
159 + {
160 + "token_ids": new_token_ids,
161 + "text": self.tokenizer.decode(
162 + new_token_ids,
163 + skip_special_tokens=False,
164 + ),
165 + }
166 + )
167 + return generations
168 +
169 + def prepare_sparse_patch(
170 + self,
171 + prompts: Sequence[str],
172 + max_patch_rows: int = MAX_PATCH_ROWS,
173 + ) -> tuple[dict[str, object], list[int], str, int]:
174 + selected_token_ids: list[int] = []
175 + special_ids = set(self.tokenizer.all_special_ids)
176 + for prompt in prompts:
177 + token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
178 + for token_id in token_ids:
179 + if token_id in special_ids or token_id in selected_token_ids:
180 + continue
181 + selected_token_ids.append(token_id)
182 + if len(selected_token_ids) == max_patch_rows:
183 + break
184 + if len(selected_token_ids) == max_patch_rows:
185 + break
186 +
187 + if not selected_token_ids:
188 + raise ValueError("Could not derive any non-special token IDs to patch")
189 +
190 + vocab_size = self.patched_param.shape[0]
191 + next_token_id = selected_token_ids[-1]
192 + while len(selected_token_ids) < max_patch_rows:
193 + next_token_id = (next_token_id + 1) % vocab_size
194 + if next_token_id in special_ids or next_token_id in selected_token_ids:
195 + continue
196 + selected_token_ids.append(next_token_id)
197 +
198 + row_ids = torch.tensor(
199 + selected_token_ids,
200 + device=self.patched_param.device,
201 + dtype=torch.long,
202 + )
203 + hidden_size = self.patched_param.shape[1]
204 + column_offsets = torch.arange(
205 + hidden_size,
206 + device=self.patched_param.device,
207 + dtype=torch.long,
208 + )
209 +
210 + with torch.no_grad():
211 + # Rotate the selected embedding rows instead of zeroing them so the
212 + # patch remains deterministic while avoiding a degenerate collapse
213 + # to the same special token after the update.
214 + replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
215 + self.patched_param[row_ids] = replacement_rows
216 +
217 + flat_indices = (
218 + row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
219 + )
220 + flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
221 + self.pending_sparse_patches = [
222 + SparseWeightPatch(
223 + name=PATCHED_PARAM_NAME,
224 + indices=flat_indices.to(torch.int32),
225 + values=flat_values,
226 + )
227 + ]
228 + patch_digest = hashlib.sha256(
229 + self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
230 + + self.pending_sparse_patches[0]
231 + .values.detach()
232 + .float()
233 + .cpu()
234 + .numpy()
235 + .tobytes()
236 + ).hexdigest()
237 +
238 + sparse_payload_bytes = (
239 + flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
240 + + flat_values.numel() * flat_values.element_size()
241 + )
242 + update_info = dict(
243 + names=[PATCHED_PARAM_NAME],
244 + dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
245 + shapes=[list(self.patched_param.shape)],
246 + num_updates_list=[flat_indices.numel()],
247 + update_kind="sparse_flat",
248 + )
249 + return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
250 +
251 + def broadcast_weights(self, packed: bool = False) -> float:
252 + if self.model_update_group is None:
253 + raise RuntimeError("Weight transfer group is not initialized")
254 +
255 + trainer_args = NCCLTrainerSendWeightsArgs(
256 + group=self.model_update_group,
257 + packed=packed,
258 + )
259 + start = time.perf_counter()
260 + NCCLWeightTransferEngine.trainer_send_weights(
261 + iterator=self.model.named_parameters(),
262 + trainer_args=trainer_args,
263 + )
264 + torch.accelerator.synchronize()
265 + return (time.perf_counter() - start) * 1000.0
266 +
267 + def broadcast_pending_sparse_patch(self) -> float:
268 + if self.model_update_group is None:
269 + raise RuntimeError("Weight transfer group is not initialized")
270 + if self.pending_sparse_patches is None:
271 + raise RuntimeError("Sparse patch has not been prepared")
272 +
273 + start = time.perf_counter()
274 + NCCLWeightTransferEngine.trainer_send_sparse_weights(
275 + iter(self.pending_sparse_patches),
276 + NCCLTrainerSendWeightsArgs(group=self.model_update_group),
277 + )
278 + torch.accelerator.synchronize()
279 + self.pending_sparse_patches = None
280 + return (time.perf_counter() - start) * 1000.0
281 +
282 +
283 +def launch_llm(
284 + scheduling_inference: PlacementGroupSchedulingStrategy,
285 +):
286 + return ray.remote(
287 + num_cpus=0,
288 + num_gpus=0,
289 + scheduling_strategy=scheduling_inference,
290 + )(MyLLM).remote(
291 + model=MODEL_NAME,
292 + enforce_eager=True,
293 + tensor_parallel_size=1,
294 + distributed_executor_backend="ray",
295 + gpu_memory_utilization=0.7,
296 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
297 + )
298 +
299 +
300 +def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
301 + outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
302 + generations = []
303 + for output in outputs:
304 + generations.append(
305 + {
306 + "token_ids": output.outputs[0].token_ids,
307 + "text": output.outputs[0].text,
308 + }
309 + )
310 + return generations
311 +
312 +
313 +def token_sequences_match(
314 + left: Sequence[dict[str, object]],
315 + right: Sequence[dict[str, object]],
316 +) -> bool:
317 + return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
318 +
319 +
320 +def print_generations(label: str, prompts: Sequence[str], generations) -> None:
321 + print(f"\n{label}")
322 + print("-" * 50)
323 + for prompt, generation in zip(prompts, generations):
324 + print(f"Prompt: {prompt!r}")
325 + print(f"Token IDs: {generation['token_ids']}")
326 + print(f"Text: {generation['text']!r}")
327 + print("-" * 50)
328 +
329 +
330 +def run_dense_phase(
331 + train_model,
332 + scheduling_inference: PlacementGroupSchedulingStrategy,
333 +) -> dict[str, object]:
334 + ray.get(train_model.reset_model.remote())
335 + llm = launch_llm(scheduling_inference)
336 + try:
337 + dense_before = collect_vllm_generations(llm)
338 +
339 + ray.get(llm.sleep.remote(level=0))
340 + master_address, master_port = ray.get(train_model.create_rendezvous.remote())
341 + world_size = ray.get(llm.get_world_size.remote()) + 1
342 + inference_init = llm.init_weight_transfer_engine.remote(
343 + dict(
344 + init_info=dict(
345 + master_address=master_address,
346 + master_port=master_port,
347 + rank_offset=1,
348 + world_size=world_size,
349 + )
350 + )
351 + )
352 + trainer_init = train_model.init_weight_transfer_group.remote(world_size)
353 + ray.get([trainer_init, inference_init])
354 + ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
355 +
356 + dense_update_info, dense_payload_bytes = ray.get(
357 + train_model.get_dense_update_info.remote()
358 + )
359 + _, selected_token_ids, patch_digest, _ = ray.get(
360 + train_model.prepare_sparse_patch.remote(PROMPTS)
361 + )
362 +
363 + inference_update = llm.update_weights.remote(
364 + dict(update_info=dense_update_info)
365 + )
366 + dense_send_ms, _ = ray.get(
367 + [
368 + train_model.broadcast_weights.remote(packed=False),
369 + inference_update,
370 + ]
371 + )
372 + ray.get(llm.finish_weight_update.remote())
373 + ray.get(llm.wake_up.remote(tags=["scheduling"]))
374 +
375 + dense_after = collect_vllm_generations(llm)
376 +
377 + return {
378 + "dense_before": dense_before,
379 + "dense_after": dense_after,
380 + "selected_token_ids": selected_token_ids,
381 + "patch_digest": patch_digest,
382 + "dense_payload_bytes": dense_payload_bytes,
383 + "dense_send_ms": dense_send_ms,
384 + }
385 + finally:
386 + ray.kill(llm)
387 +
388 +
389 +def run_sparse_phase(
390 + train_model,
391 + scheduling_inference: PlacementGroupSchedulingStrategy,
392 +) -> dict[str, object]:
393 + ray.get(train_model.reset_model.remote())
394 + llm = launch_llm(scheduling_inference)
395 + try:
396 + sparse_before = collect_vllm_generations(llm)
397 +
398 + ray.get(llm.sleep.remote(level=0))
399 + master_address, master_port = ray.get(train_model.create_rendezvous.remote())
400 + world_size = ray.get(llm.get_world_size.remote()) + 1
401 + inference_init = llm.init_weight_transfer_engine.remote(
402 + dict(
403 + init_info=dict(
404 + master_address=master_address,
405 + master_port=master_port,
406 + rank_offset=1,
407 + world_size=world_size,
408 + )
409 + )
410 + )
411 + trainer_init = train_model.init_weight_transfer_group.remote(world_size)
412 + ray.get([trainer_init, inference_init])
413 + ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
414 +
415 + sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
416 + ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
417 + )
418 +
419 + inference_update = llm.update_weights.remote(
420 + dict(update_info=sparse_update_info)
421 + )
422 + sparse_send_ms, _ = ray.get(
423 + [
424 + train_model.broadcast_pending_sparse_patch.remote(),
425 + inference_update,
426 + ]
427 + )
428 + ray.get(llm.finish_weight_update.remote())
429 + ray.get(llm.wake_up.remote(tags=["scheduling"]))
430 +
431 + sparse_after = collect_vllm_generations(llm)
432 +
433 + return {
434 + "sparse_before": sparse_before,
435 + "sparse_after": sparse_after,
436 + "selected_token_ids": selected_token_ids,
437 + "patch_digest": patch_digest,
438 + "sparse_payload_bytes": sparse_payload_bytes,
439 + "sparse_send_ms": sparse_send_ms,
440 + }
441 + finally:
442 + ray.kill(llm)
443 +
444 +
445 +ray.init()
446 +
447 +try:
448 + train_model = TrainModel.remote(MODEL_NAME)
449 +
450 + pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
451 + ray.get(pg_inference.ready())
452 + scheduling_inference = PlacementGroupSchedulingStrategy(
453 + placement_group=pg_inference,
454 + placement_group_capture_child_tasks=True,
455 + placement_group_bundle_index=0,
456 + )
457 +
458 + dense_results = run_dense_phase(train_model, scheduling_inference)
459 + sparse_results = run_sparse_phase(train_model, scheduling_inference)
460 +
461 + baseline_equal = token_sequences_match(
462 + dense_results["dense_before"],
463 + sparse_results["sparse_before"],
464 + )
465 + patch_selection_equal = (
466 + dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
467 + )
468 + patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
469 + after_equal = token_sequences_match(
470 + dense_results["dense_after"],
471 + sparse_results["sparse_after"],
472 + )
473 + any_output_changed = any(
474 + before["token_ids"] != after["token_ids"]
475 + for before, after in zip(
476 + dense_results["dense_before"],
477 + dense_results["dense_after"],
478 + )
479 + )
480 + dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
481 + sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
482 +
483 + print_generations(
484 + "Dense baseline outputs",
485 + PROMPTS,
486 + dense_results["dense_before"],
487 + )
488 + print_generations(
489 + "Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
490 + )
491 + print_generations(
492 + "Dense outputs after update", PROMPTS, dense_results["dense_after"]
493 + )
494 + print_generations(
495 + "Sparse outputs after update",
496 + PROMPTS,
497 + sparse_results["sparse_after"],
498 + )
499 +
500 + print(f"patched_token_ids = {dense_results['selected_token_ids']}")
501 + print(f"patch_selection_equal = {patch_selection_equal}")
502 + print(f"dense_patch_digest = {dense_results['patch_digest']}")
503 + print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
504 + print(f"patch_digest_equal = {patch_digest_equal}")
505 + print(f"baseline_equal = {baseline_equal}")
506 + print(f"after_equal = {after_equal}")
507 + print(f"any_output_changed = {any_output_changed}")
508 + print(f"dense_payload_mb = {dense_payload_mb:.2f}")
509 + print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
510 + print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
511 + print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
512 +
513 + if not baseline_equal:
514 + raise RuntimeError(
515 + "Dense and sparse phases did not start from the same baseline"
516 + )
517 + if not patch_selection_equal:
518 + raise RuntimeError("Dense and sparse phases used different sparse patches")
519 + if not patch_digest_equal:
520 + raise RuntimeError("Dense and sparse phases produced different patch values")
521 + if not after_equal:
522 + raise RuntimeError("Dense and sparse updates produced different outputs")
523 + if not any_output_changed:
524 + raise RuntimeError("Patch did not change the observed outputs")
525 +finally:
526 + ray.shutdown()
tests/distributed/test_weight_transfer.py
+330 -3
@@ -18,6 +18,7 @@ from torch.multiprocessing.reductions import reduce_tensor
18 from vllm.config.parallel import ParallelConfig
19 from vllm.config.weight_transfer import WeightTransferConfig
20 from vllm.distributed.weight_transfer import WeightTransferEngineFactory
21 +from vllm.distributed.weight_transfer.base import SparseWeightPatch
22 from vllm.distributed.weight_transfer.ipc_engine import (
23 IPCWeightTransferEngine,
24 IPCWeightTransferInitInfo,
@@ -89,6 +90,67 @@ class TestNCCLWeightTransferUpdateInfoValidation:
90 )
91 assert len(info.names) == 0
92
93 + def test_valid_sparse_update_info(self):
94 + """Test creating valid sparse NCCL update info."""
95 + info = NCCLWeightTransferUpdateInfo(
96 + names=["layer.weight", "layer.bias"],
97 + dtype_names=["float32", "bfloat16"],
98 + shapes=[[10, 10], [10]],
99 + num_updates_list=[4, 2],
100 + update_kind="sparse_flat",
101 + )
102 + assert info.update_kind == "sparse_flat"
103 + assert info.num_updates_list == [4, 2]
104 +
105 + def test_sparse_update_requires_num_updates_list(self):
106 + with pytest.raises(ValueError, match="`num_updates_list` is required"):
107 + NCCLWeightTransferUpdateInfo(
108 + names=["layer.weight"],
109 + dtype_names=["float32"],
110 + shapes=[[10, 10]],
111 + update_kind="sparse_flat",
112 + )
113 +
114 + def test_sparse_update_rejects_empty_num_updates_list(self):
115 + with pytest.raises(ValueError, match="cannot be empty"):
116 + NCCLWeightTransferUpdateInfo(
117 + names=[],
118 + dtype_names=[],
119 + shapes=[],
120 + num_updates_list=[],
121 + update_kind="sparse_flat",
122 + )
123 +
124 + def test_sparse_update_rejects_packed(self):
125 + with pytest.raises(ValueError, match="cannot be combined with `packed=True`"):
126 + NCCLWeightTransferUpdateInfo(
127 + names=["layer.weight"],
128 + dtype_names=["float32"],
129 + shapes=[[10, 10]],
130 + num_updates_list=[3],
131 + update_kind="sparse_flat",
132 + packed=True,
133 + )
134 +
135 + def test_sparse_update_rejects_mismatched_num_updates(self):
136 + with pytest.raises(ValueError, match="`num_updates_list`"):
137 + NCCLWeightTransferUpdateInfo(
138 + names=["layer.weight", "layer.bias"],
139 + dtype_names=["float32", "float32"],
140 + shapes=[[10, 10], [10]],
141 + num_updates_list=[3],
142 + update_kind="sparse_flat",
143 + )
144 +
145 + def test_dense_update_rejects_sparse_metadata(self):
146 + with pytest.raises(ValueError, match="Sparse metadata"):
147 + NCCLWeightTransferUpdateInfo(
148 + names=["layer.weight"],
149 + dtype_names=["float32"],
150 + shapes=[[10, 10]],
151 + num_updates_list=[3],
152 + )
153 +
154
155 # --- Unit Tests: Engine Parsing ---
156
@@ -222,6 +284,27 @@ def test_nccl_receive_weights_without_init_raises():
284 engine.receive_weights(update_info, lambda x: None)
285
286
287 +def test_nccl_receive_sparse_weights_without_init_raises():
288 + """Test that sparse receive raises if init_transfer_engine wasn't called."""
289 + if torch.accelerator.device_count() < 1:
290 + pytest.skip("Need at least 1 GPU for this test")
291 +
292 + config = WeightTransferConfig(backend="nccl")
293 + parallel_config = create_mock_parallel_config()
294 + engine = NCCLWeightTransferEngine(config, parallel_config)
295 +
296 + update_info = NCCLWeightTransferUpdateInfo(
297 + names=["w"],
298 + dtype_names=["float32"],
299 + shapes=[[10]],
300 + num_updates_list=[2],
301 + update_kind="sparse_flat",
302 + )
303 +
304 + with pytest.raises(RuntimeError, match="not initialized"):
305 + engine.receive_sparse_weights(update_info, lambda x: None)
306 +
307 +
308 # --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
309
310
@@ -379,6 +462,136 @@ def test_nccl_weight_transfer_between_processes():
462 )
463
464
465 +@ray.remote(num_gpus=1)
466 +def trainer_broadcast_sparse_tensor(
467 + master_address: str,
468 + master_port: int,
469 + world_size: int,
470 +) -> bool:
471 + """Trainer task that broadcasts sparse patches via NCCL."""
472 + import torch
473 +
474 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
475 + from vllm.distributed.utils import StatelessProcessGroup
476 + from vllm.distributed.weight_transfer.base import SparseWeightPatch
477 + from vllm.distributed.weight_transfer.nccl_engine import (
478 + NCCLTrainerSendWeightsArgs,
479 + NCCLWeightTransferEngine,
480 + )
481 +
482 + pg = StatelessProcessGroup.create(
483 + host=master_address,
484 + port=master_port,
485 + rank=0,
486 + world_size=world_size,
487 + )
488 + comm = PyNcclCommunicator(pg, device=0)
489 +
490 + patch = SparseWeightPatch(
491 + name="test.weight",
492 + indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"),
493 + values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"),
494 + )
495 + NCCLWeightTransferEngine.trainer_send_sparse_weights(
496 + iter([patch]),
497 + NCCLTrainerSendWeightsArgs(group=comm),
498 + )
499 + torch.accelerator.synchronize()
500 + return True
501 +
502 +
503 +@ray.remote(num_gpus=1)
504 +def inference_receive_sparse_tensor(
505 + master_address: str,
506 + master_port: int,
507 + world_size: int,
508 +) -> dict:
509 + """Inference task that receives sparse patches via NCCLWeightTransferEngine."""
510 + from unittest.mock import MagicMock
511 +
512 + import torch
513 +
514 + from vllm.config.parallel import ParallelConfig
515 + from vllm.config.weight_transfer import WeightTransferConfig
516 + from vllm.distributed.weight_transfer.nccl_engine import (
517 + NCCLWeightTransferEngine,
518 + NCCLWeightTransferInitInfo,
519 + NCCLWeightTransferUpdateInfo,
520 + )
521 +
522 + config = WeightTransferConfig(backend="nccl")
523 + parallel_config = MagicMock(spec=ParallelConfig)
524 + parallel_config.rank = 0
525 + parallel_config.world_size = 1
526 + parallel_config.data_parallel_rank = 0
527 + parallel_config.data_parallel_index = 0
528 +
529 + engine = NCCLWeightTransferEngine(config, parallel_config)
530 + engine.init_transfer_engine(
531 + NCCLWeightTransferInitInfo(
532 + master_address=master_address,
533 + master_port=master_port,
534 + rank_offset=1,
535 + world_size=world_size,
536 + )
537 + )
538 +
539 + target = torch.zeros(30, dtype=torch.float32, device="cuda")
540 +
541 + def apply_sparse_patches(patches: list[SparseWeightPatch]):
542 + for patch in patches:
543 + target.index_copy_(0, patch.indices.to(torch.long), patch.values)
544 +
545 + update_info = NCCLWeightTransferUpdateInfo(
546 + names=["test.weight"],
547 + dtype_names=["float32"],
548 + shapes=[[30]],
549 + num_updates_list=[3],
550 + update_kind="sparse_flat",
551 + )
552 + engine.receive_sparse_weights(update_info, apply_sparse_patches)
553 + torch.accelerator.synchronize()
554 +
555 + expected = torch.zeros(30, dtype=torch.float32, device="cuda")
556 + expected[[1, 7, 25]] = torch.tensor(
557 + [10.0, 20.0, 30.0], dtype=torch.float32, device="cuda"
558 + )
559 + success = torch.equal(target, expected)
560 + engine.shutdown()
561 + return {
562 + "success": success,
563 + "selected_values": target[[1, 7, 25]].cpu().tolist(),
564 + }
565 +
566 +
567 +@pytest.mark.skipif(
568 + torch.accelerator.device_count() < 2,
569 + reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
570 +)
571 +def test_nccl_sparse_weight_transfer_between_processes():
572 + """Test NCCL sparse weight transfer from trainer to inference process."""
573 + ray.init(ignore_reinit_error=True)
574 +
575 + master_address = "127.0.0.1"
576 + master_port = get_open_port()
577 + world_size = 2
578 +
579 + inference_future = inference_receive_sparse_tensor.remote(
580 + master_address, master_port, world_size
581 + )
582 + trainer_future = trainer_broadcast_sparse_tensor.remote(
583 + master_address, master_port, world_size
584 + )
585 +
586 + trainer_result, result = ray.get([trainer_future, inference_future])
587 +
588 + assert trainer_result, "Trainer should complete successfully"
589 + assert result["success"], (
590 + "Sparse weight transfer failed. "
591 + f"Received selected values: {result['selected_values']}"
592 + )
593 +
594 +
595 # --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
596
597
@@ -461,9 +674,101 @@ class TestIPCWeightTransferUpdateInfoValidation:
674 ipc_handles=ipc_handles,
675 )
676
464 - def test_missing_ipc_handles_raises(self):
465 - """Test that omitting ipc_handles raises TypeError."""
466 - with pytest.raises(TypeError):
677 + def test_sparse_update_kind_rejected(self):
678 + """Test that IPC backend rejects sparse update metadata."""
679 + if torch.accelerator.device_count() < 1:
680 + pytest.skip("Need at least 1 GPU for this test")
681 +
682 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
683 + ipc_handle = reduce_tensor(dummy_tensor)
684 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
685 + ipc_handles = [{gpu_uuid: ipc_handle}]
686 +
687 + with pytest.raises(NotImplementedError, match="dense updates"):
688 + IPCWeightTransferUpdateInfo(
689 + names=["layer.weight"],
690 + dtype_names=["float32"],
691 + shapes=[[10, 10]],
692 + num_updates_list=[1],
693 + ipc_handles=ipc_handles,
694 + update_kind="sparse_flat",
695 + )
696 +
697 + def test_sparse_methods_not_supported(self):
698 + """Test that IPC engine inherits sparse rejection from the base class."""
699 + config = WeightTransferConfig(backend="ipc")
700 + parallel_config = create_mock_parallel_config()
701 + engine = IPCWeightTransferEngine(
702 + config, parallel_config, MagicMock(spec=torch.nn.Module)
703 + )
704 +
705 + with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
706 + engine.receive_sparse_weights(MagicMock(), lambda _: None)
707 + with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
708 + engine.trainer_send_sparse_weights(
709 + iter([]),
710 + {"mode": "http", "url": "http://localhost:8000"},
711 + )
712 +
713 + def test_valid_update_info_from_pickled(self, monkeypatch):
714 + """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
715 + if torch.accelerator.device_count() < 1:
716 + pytest.skip("Need at least 1 GPU for this test")
717 +
718 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
719 +
720 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
721 + ipc_handle = reduce_tensor(dummy_tensor)
722 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
723 + ipc_handles = [{gpu_uuid: ipc_handle}]
724 +
725 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
726 +
727 + info = IPCWeightTransferUpdateInfo(
728 + names=["layer.weight"],
729 + dtype_names=["float32"],
730 + shapes=[[10, 10]],
731 + ipc_handles_pickled=pickled,
732 + )
733 + assert info.ipc_handles == ipc_handles
734 + assert info.ipc_handles_pickled is None
735 +
736 + def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
737 + """Test that pickled handles are rejected unless env flag is enabled."""
738 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
739 +
740 + with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
741 + IPCWeightTransferUpdateInfo(
742 + names=[],
743 + dtype_names=[],
744 + shapes=[],
745 + ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
746 + )
747 +
748 + def test_both_handles_and_pickled_raises(self):
749 + """Test that providing both ipc_handles and ipc_handles_pickled raises."""
750 + if torch.accelerator.device_count() < 1:
751 + pytest.skip("Need at least 1 GPU for this test")
752 +
753 + dummy_tensor = torch.ones(10, 10, device="cuda:0")
754 + ipc_handle = reduce_tensor(dummy_tensor)
755 + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
756 + ipc_handles = [{gpu_uuid: ipc_handle}]
757 +
758 + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
759 +
760 + with pytest.raises(ValueError, match="Cannot specify both"):
761 + IPCWeightTransferUpdateInfo(
762 + names=["layer.weight"],
763 + dtype_names=["float32"],
764 + shapes=[[10, 10]],
765 + ipc_handles=ipc_handles,
766 + ipc_handles_pickled=pickled,
767 + )
768 +
769 + def test_neither_handles_nor_pickled_raises(self):
770 + """Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
771 + with pytest.raises(ValueError, match="must be provided"):
772 IPCWeightTransferUpdateInfo(
773 names=["layer.weight"],
774 dtype_names=["float32"],
@@ -558,6 +863,28 @@ class TestIPCEngineParsing:
863 assert gpu_uuid in update_info.ipc_handles[0]
864 assert gpu_uuid in update_info.ipc_handles[1]
865
866 + def test_parse_update_info_ignores_none_pickled_handles(self):
867 + """Test Ray/asdict payloads with a null pickled field use ipc_handles."""
868 + config = WeightTransferConfig(backend="ipc")
869 + parallel_config = create_mock_parallel_config()
870 + engine = IPCWeightTransferEngine(
871 + config, parallel_config, MagicMock(spec=torch.nn.Module)
872 + )
873 + ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
874 +
875 + update_info = engine.parse_update_info(
876 + {
877 + "names": ["w1"],
878 + "dtype_names": ["float32"],
879 + "shapes": [[1]],
880 + "ipc_handles": ipc_handles,
881 + "ipc_handles_pickled": None,
882 + }
883 + )
884 +
885 + assert isinstance(update_info, IPCWeightTransferUpdateInfo)
886 + assert update_info.ipc_handles == ipc_handles
887 +
888 def test_parse_update_info_both_handles_and_pickled_raises(self):
889 """Test that providing both ipc_handles and ipc_handles_pickled raises."""
890 if torch.accelerator.device_count() < 1:
tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
+66 -5
@@ -48,6 +48,7 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
48 names: list[str] | None = None
49 dtype_names: list[str] | None = None
50 shapes: list[list[int]] | None = None
51 + num_updates_list: list[int] | None = None
52
53
54 class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
@@ -87,6 +88,15 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
88 # (In real implementation, this would receive and load actual weights)
89 load_weights([])
90
91 + def receive_sparse_weights(
92 + self,
93 + update_info: MockUpdateInfo,
94 + apply_patches: Callable[[list], None],
95 + ) -> None:
96 + MockWeightTransferEngine.receive_weights_called = True
97 + MockWeightTransferEngine.last_update_info = update_info
98 + apply_patches([])
99 +
100 def shutdown(self) -> None:
101 MockWeightTransferEngine.shutdown_called = True
102
@@ -198,8 +208,6 @@ def test_update_weights_calls_engine():
208 llm.init_weight_transfer_engine(
209 WeightTransferInitRequest(init_info={"test_param": "init"})
210 )
201 -
202 - # Start weight update (required before update_weights)
211 llm.start_weight_update(is_checkpoint_format=True)
212
213 # Call update_weights
@@ -232,14 +240,67 @@ def test_update_weights_calls_engine():
240 assert dtypes == test_dtypes
241 assert shapes == test_shapes
242
235 - # Finish weight update
243 + llm.finish_weight_update()
244 +
245 +
246 +@create_new_process_for_each_test()
247 +def test_update_weights_passes_sparse_metadata():
248 + """Test sparse update metadata is forwarded unchanged to the engine."""
249 + if torch.accelerator.device_count() < 1:
250 + pytest.skip("Need at least 1 GPU for this test")
251 +
252 + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
253 + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
254 +
255 + with patch(
256 + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
257 + mock_create_engine,
258 + ):
259 + llm = LLM(
260 + model=MODEL_NAME,
261 + enforce_eager=True,
262 + load_format="dummy",
263 + tensor_parallel_size=1,
264 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
265 + )
266 +
267 + llm.init_weight_transfer_engine(
268 + WeightTransferInitRequest(init_info={"test_param": "init"})
269 + )
270 + llm.start_weight_update(is_checkpoint_format=False)
271 +
272 + llm.update_weights(
273 + WeightTransferUpdateRequest(
274 + update_info={
275 + "names": ["layer.weight"],
276 + "dtype_names": ["bfloat16"],
277 + "shapes": [[100]],
278 + "num_updates_list": [3],
279 + "update_kind": "sparse_flat",
280 + }
281 + )
282 + )
283 +
284 + def check_sparse_update_called(self):
285 + engine = self.weight_transfer_engine
286 + if not engine.receive_weights_called:
287 + return None
288 + info = engine.last_update_info
289 + return (
290 + info.update_kind,
291 + info.num_updates_list,
292 + )
293 +
294 + results = llm.collective_rpc(check_sparse_update_called)
295 + for result in results:
296 + assert result == ("sparse_flat", [3])
297 +
298 llm.finish_weight_update()
299
300
301 @create_new_process_for_each_test()
302 def test_full_weight_transfer_flow():
241 - """Test the complete weight transfer flow:
242 - init -> start -> update -> finish."""
303 + """Test the complete weight transfer flow: init -> start -> update -> finish."""
304 if torch.accelerator.device_count() < 1:
305 pytest.skip("Need at least 1 GPU for this test")
306
tests/v1/worker/test_gpu_model_runner.py
+69
@@ -7,6 +7,7 @@ from unittest.mock import Mock
7 import numpy as np
8 import pytest
9 import torch
10 +import torch.nn as nn
11
12 import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module
13 from vllm.config import (
@@ -22,6 +23,7 @@ from vllm.distributed.parallel_state import (
23 init_distributed_environment,
24 initialize_model_parallel,
25 )
26 +from vllm.distributed.weight_transfer.base import SparseWeightPatch
27 from vllm.model_executor.layers.attention import Attention
28 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
29 from vllm.platforms import current_platform
@@ -784,6 +786,73 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler():
786 assert torch.equal(passed_draft_probs, expected_draft_probs)
787
788
789 +def test_apply_sparse_weight_patches_updates_only_selected_entries():
790 + class DummyModel(nn.Module):
791 + def __init__(self):
792 + super().__init__()
793 + self.weight = nn.Parameter(torch.zeros(6, dtype=torch.float32))
794 +
795 + runner = object.__new__(GPUModelRunner)
796 + runner.model = DummyModel()
797 +
798 + runner.apply_sparse_weight_patches(
799 + [
800 + SparseWeightPatch(
801 + name="weight",
802 + indices=torch.tensor([1, 4], dtype=torch.int32),
803 + values=torch.tensor([3.5, -2.0], dtype=torch.float32),
804 + )
805 + ]
806 + )
807 +
808 + expected = torch.tensor([0.0, 3.5, 0.0, 0.0, -2.0, 0.0], dtype=torch.float32)
809 + assert torch.equal(runner.get_model().weight.data, expected)
810 +
811 +
812 +def test_apply_sparse_weight_patches_rejects_mismatched_lengths():
813 + class DummyModel(nn.Module):
814 + def __init__(self):
815 + super().__init__()
816 + self.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
817 +
818 + runner = object.__new__(GPUModelRunner)
819 + runner.model = DummyModel()
820 +
821 + with pytest.raises(ValueError, match="matching lengths"):
822 + runner.apply_sparse_weight_patches(
823 + [
824 + SparseWeightPatch(
825 + name="weight",
826 + indices=torch.tensor([1, 2], dtype=torch.int32),
827 + values=torch.tensor([1.0], dtype=torch.float32),
828 + )
829 + ]
830 + )
831 +
832 +
833 +def test_apply_sparse_weight_patches_rejects_non_contiguous_param():
834 + class DummyModel(nn.Module):
835 + def __init__(self):
836 + super().__init__()
837 + self.weight = nn.Parameter(
838 + torch.arange(12, dtype=torch.float32).view(3, 4).t()
839 + )
840 +
841 + runner = object.__new__(GPUModelRunner)
842 + runner.model = DummyModel()
843 +
844 + with pytest.raises(NotImplementedError, match="contiguous params"):
845 + runner.apply_sparse_weight_patches(
846 + [
847 + SparseWeightPatch(
848 + name="weight",
849 + indices=torch.tensor([1], dtype=torch.int32),
850 + values=torch.tensor([1.0], dtype=torch.float32),
851 + )
852 + ]
853 + )
854 +
855 +
856 def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):
857 torch.set_default_dtype(torch.float16)
858 layer_0 = "model.layers.0.self_attn.attn"
tests/v1/worker/test_gpu_worker_weight_transfer.py new
+155
@@ -0,0 +1,155 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from unittest.mock import MagicMock
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.config.parallel import ParallelConfig
11 +from vllm.config.weight_transfer import WeightTransferConfig
12 +from vllm.distributed.weight_transfer.base import SparseWeightPatch
13 +from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine
14 +from vllm.v1.worker.gpu_worker import Worker
15 +
16 +
17 +def _make_nccl_engine() -> NCCLWeightTransferEngine:
18 + parallel_config = MagicMock(spec=ParallelConfig)
19 + parallel_config.rank = 0
20 + parallel_config.world_size = 1
21 + parallel_config.data_parallel_rank = 0
22 + parallel_config.data_parallel_index = 0
23 + return NCCLWeightTransferEngine(
24 + WeightTransferConfig(backend="nccl"),
25 + parallel_config,
26 + MagicMock(spec=torch.nn.Module),
27 + )
28 +
29 +
30 +def test_update_weights_sparse_dispatches_to_sparse_receive(monkeypatch):
31 + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
32 +
33 + worker = object.__new__(Worker)
34 + worker.device = "cpu"
35 + worker.parallel_config = SimpleNamespace(world_size=1)
36 + worker.weight_transfer_engine = _make_nccl_engine()
37 + worker._weight_update_active = True
38 + worker._is_checkpoint_format = False
39 +
40 + applied_patches = []
41 +
42 + def apply_sparse_weight_patches(patches):
43 + applied_patches.extend(patches)
44 +
45 + worker.model_runner = SimpleNamespace(
46 + apply_sparse_weight_patches=apply_sparse_weight_patches,
47 + )
48 +
49 + received_kinds = []
50 +
51 + def receive_sparse_weights(update_info, apply_patches):
52 + received_kinds.append(update_info.update_kind)
53 + apply_patches(
54 + [
55 + SparseWeightPatch(
56 + name="layer.weight",
57 + indices=torch.tensor([1], dtype=torch.int32),
58 + values=torch.tensor([2.0], dtype=torch.float32),
59 + )
60 + ]
61 + )
62 +
63 + worker.weight_transfer_engine.receive_sparse_weights = receive_sparse_weights
64 +
65 + Worker.update_weights(
66 + worker,
67 + {
68 + "names": ["layer.weight"],
69 + "dtype_names": ["float32"],
70 + "shapes": [[4]],
71 + "num_updates_list": [1],
72 + "update_kind": "sparse_flat",
73 + },
74 + )
75 +
76 + assert received_kinds == ["sparse_flat"]
77 + assert len(applied_patches) == 1
78 + assert torch.equal(applied_patches[0].indices, torch.tensor([1], dtype=torch.int32))
79 +
80 +
81 +def test_update_weights_sparse_rejects_tp_or_pp(monkeypatch):
82 + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
83 +
84 + worker = object.__new__(Worker)
85 + worker.device = "cpu"
86 + worker.parallel_config = SimpleNamespace(world_size=2)
87 + worker.weight_transfer_engine = _make_nccl_engine()
88 + worker._weight_update_active = True
89 + worker._is_checkpoint_format = False
90 + worker.model_runner = SimpleNamespace(apply_sparse_weight_patches=lambda _: None)
91 +
92 + with pytest.raises(NotImplementedError, match="TP=1 and PP=1"):
93 + Worker.update_weights(
94 + worker,
95 + {
96 + "names": ["layer.weight"],
97 + "dtype_names": ["float32"],
98 + "shapes": [[4]],
99 + "num_updates_list": [1],
100 + "update_kind": "sparse_flat",
101 + },
102 + )
103 + assert worker._weight_update_active is False
104 + assert worker._is_checkpoint_format is True
105 +
106 +
107 +def test_update_weights_sparse_rejects_checkpoint_format(monkeypatch):
108 + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
109 +
110 + worker = object.__new__(Worker)
111 + worker.device = "cpu"
112 + worker.parallel_config = SimpleNamespace(world_size=1)
113 + worker.weight_transfer_engine = _make_nccl_engine()
114 + worker._weight_update_active = True
115 + worker._is_checkpoint_format = True
116 + worker.model_runner = SimpleNamespace(model=MagicMock())
117 +
118 + with pytest.raises(ValueError, match="start_weight_update"):
119 + Worker.update_weights(
120 + worker,
121 + {
122 + "names": ["layer.weight"],
123 + "dtype_names": ["float32"],
124 + "shapes": [[4]],
125 + "num_updates_list": [1],
126 + "update_kind": "sparse_flat",
127 + },
128 + )
129 + assert worker._weight_update_active is False
130 + assert worker._is_checkpoint_format is True
131 +
132 +
133 +def test_update_weights_resets_state_when_update_info_is_invalid(monkeypatch):
134 + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
135 +
136 + worker = object.__new__(Worker)
137 + worker.device = "cpu"
138 + worker.parallel_config = SimpleNamespace(world_size=1)
139 + worker.weight_transfer_engine = _make_nccl_engine()
140 + worker._weight_update_active = True
141 + worker._is_checkpoint_format = False
142 +
143 + with pytest.raises(ValueError, match="cannot be empty"):
144 + Worker.update_weights(
145 + worker,
146 + {
147 + "names": [],
148 + "dtype_names": [],
149 + "shapes": [],
150 + "num_updates_list": [],
151 + "update_kind": "sparse_flat",
152 + },
153 + )
154 + assert worker._weight_update_active is False
155 + assert worker._is_checkpoint_format is True
vllm/distributed/weight_transfer/base.py
+58 -3
@@ -4,8 +4,8 @@
4
5 from abc import ABC, abstractmethod
6 from collections.abc import Callable, Iterator
7 -from dataclasses import dataclass, field
8 -from typing import Any, Generic, TypeVar
7 +from dataclasses import KW_ONLY, dataclass, field
8 +from typing import Any, Generic, Literal, TypeVar
9
10 import torch
11
@@ -28,7 +28,44 @@ class WeightTransferInitInfo(ABC): # noqa: B024
28 class WeightTransferUpdateInfo(ABC): # noqa: B024
29 """Base class for backend-specific weight update info."""
30
31 - pass
31 + _: KW_ONLY
32 + update_kind: Literal["dense", "sparse_flat"] = "dense"
33 + """Weight update format."""
34 + num_updates_list: list[int] | None = None
35 + """Number of sparse entries to receive for each parameter in ``names``."""
36 +
37 + def __post_init__(self) -> None:
38 + if self.update_kind not in ("dense", "sparse_flat"):
39 + raise ValueError(f"Unsupported update_kind: {self.update_kind}")
40 + if self.update_kind == "dense":
41 + if self.num_updates_list is not None:
42 + raise ValueError(
43 + "Sparse metadata is only supported for `update_kind='sparse_flat'`"
44 + )
45 + return
46 +
47 + if self.num_updates_list is None:
48 + raise ValueError("`num_updates_list` is required for sparse updates")
49 + if len(self.num_updates_list) == 0:
50 + raise ValueError("`num_updates_list` cannot be empty for sparse updates")
51 + if any(num_updates < 0 for num_updates in self.num_updates_list):
52 + raise ValueError("Sparse `num_updates_list` entries must be non-negative")
53 +
54 + names = getattr(self, "names", None)
55 + if names is not None and len(self.num_updates_list) != len(names):
56 + raise ValueError(
57 + f"`num_updates_list` should be of the same size as `names`: "
58 + f"got {len(self.num_updates_list)} and {len(names)}"
59 + )
60 +
61 +
62 +@dataclass
63 +class SparseWeightPatch:
64 + """A sparse in-place patch for one existing parameter."""
65 +
66 + name: str
67 + indices: torch.Tensor
68 + values: torch.Tensor
69
70
71 # API-level request classes (accept dicts for backend-agnostic serialization)
@@ -150,6 +187,16 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
187 """
188 raise NotImplementedError
189
190 + def receive_sparse_weights(
191 + self,
192 + update_info: TUpdateInfo,
193 + apply_patches: Callable[[list[SparseWeightPatch]], None],
194 + ) -> None:
195 + """Receive sparse weight patches from the trainer."""
196 + raise NotImplementedError(
197 + f"{self.__class__.__name__} does not support sparse weight updates"
198 + )
199 +
200 @abstractmethod
201 def shutdown(self) -> None:
202 """
@@ -184,3 +231,11 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
231 >>> engine.trainer_send_weights(param_iter, trainer_args)
232 """
233 raise NotImplementedError
234 +
235 + @staticmethod
236 + def trainer_send_sparse_weights(
237 + _iterator: Iterator[SparseWeightPatch],
238 + _trainer_args: dict[str, Any] | Any,
239 + ) -> None:
240 + """Send sparse weight patches from trainer to inference workers."""
241 + raise NotImplementedError("Sparse weight updates are not supported")
vllm/distributed/weight_transfer/ipc_engine.py
+29 -4
@@ -74,10 +74,12 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
74 names: list[str]
75 dtype_names: list[str]
76 shapes: list[list[int]]
77 - ipc_handles: list[dict[str, tuple]] | dict[str, tuple]
77 + ipc_handles: list[dict[str, tuple]] | dict[str, tuple] | None = None
78 """IPC handles mapping physical GPU UUID to rebuild_cuda_tensor args.
79 For non-packed mode: list of per-parameter handle dicts.
80 For packed mode: single handle dict for the packed buffer."""
81 + ipc_handles_pickled: str | None = None
82 + """Base64-encoded pickled IPC handles, used for HTTP transport."""
83 tensor_sizes: list[int] | None = None
84 """Per-parameter sizes in bytes within the packed buffer.
85 Required when packed=True, unused otherwise."""
@@ -85,6 +87,29 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
87 """Whether this update uses packed tensor format."""
88
89 def __post_init__(self):
90 + super().__post_init__()
91 + if self.update_kind != "dense":
92 + raise NotImplementedError("IPC weight transfer only supports dense updates")
93 +
94 + if self.ipc_handles_pickled is not None:
95 + if self.ipc_handles is not None:
96 + raise ValueError(
97 + "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
98 + )
99 +
100 + if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
101 + raise ValueError(
102 + "Refusing to deserialize `ipc_handles_pickled` without "
103 + "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
104 + )
105 +
106 + self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled))
107 + self.ipc_handles_pickled = None
108 +
109 + if self.ipc_handles is None:
110 + raise ValueError(
111 + "Either `ipc_handles` or `ipc_handles_pickled` must be provided"
112 + )
113 num_params = len(self.names)
114 if len(self.dtype_names) != num_params:
115 raise ValueError(
@@ -153,8 +178,9 @@ class IPCWeightTransferEngine(
178 Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
179 payload is deserialized via ``pickle.loads``.
180 """
156 - if "ipc_handles_pickled" in update_dict:
157 - if "ipc_handles" in update_dict:
181 + pickled = update_dict.pop("ipc_handles_pickled", None)
182 + if pickled is not None:
183 + if update_dict.get("ipc_handles") is not None:
184 raise ValueError(
185 "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
186 )
@@ -165,7 +191,6 @@ class IPCWeightTransferEngine(
191 "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
192 )
193
168 - pickled = update_dict.pop("ipc_handles_pickled")
194 update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
195
196 return super().parse_update_info(update_dict)
vllm/distributed/weight_transfer/nccl_engine.py
+71
@@ -14,6 +14,7 @@ if TYPE_CHECKING:
14 from vllm.config.parallel import ParallelConfig
15 from vllm.config.weight_transfer import WeightTransferConfig
16 from vllm.distributed.weight_transfer.base import (
17 + SparseWeightPatch,
18 WeightTransferEngine,
19 WeightTransferInitInfo,
20 WeightTransferUpdateInfo,
@@ -81,6 +82,7 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
82
83 def __post_init__(self):
84 """Validate that all lists have the same length."""
85 + super().__post_init__()
86 num_params = len(self.names)
87 if len(self.dtype_names) != num_params:
88 raise ValueError(
@@ -92,6 +94,13 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
94 f"`shapes` should be of the same size as `names`: "
95 f"got {len(self.shapes)} and {len(self.names)}"
96 )
97 + if self.update_kind == "dense":
98 + return
99 +
100 + if self.packed:
101 + raise ValueError(
102 + "`update_kind='sparse_flat'` cannot be combined with `packed=True`"
103 + )
104
105
106 class NCCLWeightTransferEngine(
@@ -178,6 +187,11 @@ class NCCLWeightTransferEngine(
187 "NCCL weight transfer not initialized. "
188 "Call init_transfer_engine() first."
189 )
190 + if update_info.update_kind != "dense":
191 + raise ValueError(
192 + "Sparse updates must use `receive_sparse_weights`, not "
193 + "`receive_weights`"
194 + )
195
196 if update_info.packed:
197 # Build iterator of (name, (shape, dtype)) from update_info
@@ -209,6 +223,42 @@ class NCCLWeightTransferEngine(
223 load_weights([(name, weight)])
224 del weight
225
226 + def receive_sparse_weights(
227 + self,
228 + update_info: NCCLWeightTransferUpdateInfo,
229 + apply_patches: Callable[[list[SparseWeightPatch]], None],
230 + ) -> None:
231 + """Receive sparse flat-index patches from trainer via NCCL."""
232 + if self.model_update_group is None:
233 + raise RuntimeError(
234 + "NCCL weight transfer not initialized. "
235 + "Call init_transfer_engine() first."
236 + )
237 + if update_info.update_kind != "sparse_flat":
238 + raise ValueError("Sparse receive path requires `update_kind='sparse_flat'`")
239 + assert update_info.num_updates_list is not None
240 +
241 + for name, dtype_name, num_updates in zip(
242 + update_info.names,
243 + update_info.dtype_names,
244 + update_info.num_updates_list,
245 + ):
246 + dtype = getattr(torch, dtype_name)
247 + device = torch.accelerator.current_device_index()
248 + indices = torch.empty(num_updates, dtype=torch.int32, device=device)
249 + values = torch.empty(num_updates, dtype=dtype, device=device)
250 + self.model_update_group.broadcast(
251 + indices, src=0, stream=torch.cuda.current_stream()
252 + )
253 + self.model_update_group.broadcast(
254 + values, src=0, stream=torch.cuda.current_stream()
255 + )
256 + apply_patches(
257 + [SparseWeightPatch(name=name, indices=indices, values=values)]
258 + )
259 + del indices
260 + del values
261 +
262 def shutdown(self) -> None:
263 if self.model_update_group is not None:
264 # Clean up the communicator by removing the reference
@@ -272,6 +322,27 @@ class NCCLWeightTransferEngine(
322 stream=args.stream or torch.cuda.current_stream(),
323 )
324
325 + @staticmethod
326 + def trainer_send_sparse_weights(
327 + iterator: Iterator[SparseWeightPatch],
328 + trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
329 + ) -> None:
330 + """Broadcast sparse flat-index patches from trainer to vLLM workers."""
331 + if isinstance(trainer_args, dict):
332 + args = NCCLTrainerSendWeightsArgs(**trainer_args)
333 + else:
334 + args = trainer_args
335 +
336 + if args.packed:
337 + raise ValueError(
338 + "Sparse NCCL updates cannot be combined with `packed=True`"
339 + )
340 +
341 + stream = args.stream or torch.cuda.current_stream()
342 + for patch in iterator:
343 + args.group.broadcast(patch.indices, src=args.src, stream=stream)
344 + args.group.broadcast(patch.values, src=args.src, stream=stream)
345 +
346 @staticmethod
347 def trainer_init(
348 init_info: NCCLWeightTransferInitInfo | dict,
vllm/entrypoints/llm.py
+2 -11
@@ -873,14 +873,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
873 )
874
875 def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
876 - """
877 - Start a new weight update.
878 -
879 - Args:
880 - is_checkpoint_format: Whether incoming weights are in checkpoint
881 - format (need layerwise processing) or kernel format (direct
882 - copy).
883 - """
876 + """Start a new weight update."""
877 self.llm_engine.collective_rpc(
878 "start_weight_update",
879 kwargs={"is_checkpoint_format": is_checkpoint_format},
@@ -902,9 +895,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
895 )
896
897 def finish_weight_update(self) -> None:
905 - """
906 - Finish the current weight update.
907 - """
898 + """Finish the current weight update."""
899 self.llm_engine.collective_rpc("finish_weight_update")
900
901 def __repr__(self) -> str:
vllm/v1/worker/gpu_model_runner.py
+39
@@ -49,6 +49,7 @@ from vllm.distributed.parallel_state import (
49 is_global_first_rank,
50 prepare_communication_buffer_for_model,
51 )
52 +from vllm.distributed.weight_transfer.base import SparseWeightPatch
53 from vllm.forward_context import (
54 BatchDescriptor,
55 set_forward_context,
@@ -3189,6 +3190,44 @@ class GPUModelRunner(
3190 return self.model.unwrap()
3191 return self.model
3192
3193 + def apply_sparse_weight_patches(self, patches: Iterable[SparseWeightPatch]) -> None:
3194 + """Apply sparse flat-index patches directly to existing model params."""
3195 + model = self.get_model()
3196 + for patch in patches:
3197 + param = model.get_parameter(patch.name)
3198 + if not param.data.is_contiguous():
3199 + raise NotImplementedError(
3200 + "Sparse weight updates currently require contiguous params: "
3201 + f"{patch.name}"
3202 + )
3203 +
3204 + if patch.indices.dtype != torch.int32:
3205 + raise ValueError(
3206 + "Sparse weight updates currently require int32 indices: "
3207 + f"{patch.name}"
3208 + )
3209 + if patch.indices.ndim != 1 or patch.values.ndim != 1:
3210 + raise ValueError(
3211 + f"Sparse weight patches must be 1D flattened updates: {patch.name}"
3212 + )
3213 + if patch.indices.numel() != patch.values.numel():
3214 + raise ValueError(
3215 + "`indices` and `values` must have matching lengths for "
3216 + f"{patch.name}"
3217 + )
3218 + if patch.values.dtype != param.dtype:
3219 + raise ValueError(
3220 + f"Sparse values dtype {patch.values.dtype} does not match "
3221 + f"parameter dtype {param.dtype} for {patch.name}"
3222 + )
3223 +
3224 + flat_param = param.data.view(-1)
3225 + flat_param.index_copy_(
3226 + 0,
3227 + patch.indices.to(device=flat_param.device, dtype=torch.long),
3228 + patch.values.to(device=flat_param.device),
3229 + )
3230 +
3231 def get_supported_generation_tasks(self) -> list[GenerationTask]:
3232 model = self.get_model()
3233 supported_tasks = list[GenerationTask]()
vllm/v1/worker/gpu_worker.py
+62 -48
@@ -992,23 +992,19 @@ class Worker(WorkerBase):
992
993 def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
994 """
995 - Start a new weight update.
996 -
997 - Prepares the model for receiving weights. For checkpoint format,
998 - this initializes state for layerwise processing. For kernel format, this is
999 - a no-op but must still be called for consistency.
995 + Start a new weight update session.
996
997 Args:
998 is_checkpoint_format: Whether incoming weights are in checkpoint
999 format (need layerwise processing) or kernel format (direct
1004 - copy). Stored as state for finish_weight_update.
1000 + copy / sparse patch application).
1001 """
1002 self._check_weight_transfer_engine()
1003
1004 if self._weight_update_active:
1005 raise RuntimeError(
1010 - "start_weight_update called while a weight update is "
1011 - "already active. Call finish_weight_update first."
1006 + "start_weight_update called while a weight update is already "
1007 + "active. Call finish_weight_update first."
1008 )
1009
1010 if is_checkpoint_format:
@@ -1020,16 +1016,15 @@ class Worker(WorkerBase):
1016 with torch.device(self.device):
1017 initialize_layerwise_reload(model)
1018
1023 - # Store state so update_weights/finish_weight_update can check
1019 self._is_checkpoint_format = is_checkpoint_format
1020 self._weight_update_active = True
1021
1022 def update_weights(self, update_info: dict) -> None:
1023 """
1029 - Receive weights from the trainer (one or more chunks).
1024 + Receive one weight update chunk from the trainer.
1025
1026 start_weight_update must be called before update_weights and
1032 - finish_weight_update must be called after.
1027 + finish_weight_update must be called after all chunks have been sent.
1028
1029 Args:
1030 update_info: Dictionary containing backend-specific update info
@@ -1042,52 +1037,72 @@ class Worker(WorkerBase):
1037 "start_weight_update must be called before update_weights."
1038 )
1039
1045 - # Parse dict into backend-specific typed dataclass
1046 - typed_update_info = self.weight_transfer_engine.parse_update_info(update_info)
1047 -
1048 - model = self.model_runner.model
1040 + update_succeeded = False
1041 + try:
1042 + # Parse dict into backend-specific typed dataclass
1043 + typed_update_info = self.weight_transfer_engine.parse_update_info(
1044 + update_info
1045 + )
1046
1050 - with torch.device(self.device):
1051 - if self._is_checkpoint_format:
1052 - self.weight_transfer_engine.receive_weights(
1053 - typed_update_info,
1054 - load_weights=model.load_weights,
1055 - )
1056 - else:
1057 - # Weights are already in kernel format, copy directly
1058 - def load_weights_direct(
1059 - weights: list[tuple[str, torch.Tensor]],
1060 - ) -> None:
1061 - for name, weight in weights:
1062 - param = model.get_parameter(name)
1063 - param.copy_(weight)
1064 -
1065 - self.weight_transfer_engine.receive_weights(
1066 - typed_update_info,
1067 - load_weights=load_weights_direct,
1068 - )
1047 + with torch.device(self.device):
1048 + if self._is_checkpoint_format:
1049 + if typed_update_info.update_kind != "dense":
1050 + raise ValueError(
1051 + "Sparse weight updates require "
1052 + "`start_weight_update(is_checkpoint_format=False)`."
1053 + )
1054 +
1055 + model = self.model_runner.model
1056 +
1057 + # Use layerwise reload pattern for checkpoint format weights
1058 + self.weight_transfer_engine.receive_weights(
1059 + typed_update_info,
1060 + load_weights=model.load_weights,
1061 + )
1062 + elif typed_update_info.update_kind == "sparse_flat":
1063 + if self.parallel_config.world_size != 1:
1064 + raise NotImplementedError(
1065 + "Sparse weight updates currently require TP=1 and PP=1"
1066 + )
1067 + self.weight_transfer_engine.receive_sparse_weights(
1068 + typed_update_info,
1069 + apply_patches=self.model_runner.apply_sparse_weight_patches,
1070 + )
1071 + else:
1072 + model = self.model_runner.model
1073 +
1074 + # Weights are already in kernel format, copy directly.
1075 + def load_weights_direct(
1076 + weights: list[tuple[str, torch.Tensor]],
1077 + ) -> None:
1078 + for name, weight in weights:
1079 + param = model.get_parameter(name)
1080 + param.copy_(weight)
1081 +
1082 + self.weight_transfer_engine.receive_weights(
1083 + typed_update_info,
1084 + load_weights=load_weights_direct,
1085 + )
1086
1070 - # NCCL broadcast/packed path are asynchronous.
1071 - # Sync here so the next step uses the new weights.
1072 - torch.accelerator.synchronize()
1087 + # NCCL broadcast/packed path are asynchronous.
1088 + # Sync here so the next step uses the new weights.
1089 + torch.accelerator.synchronize()
1090 + update_succeeded = True
1091 + finally:
1092 + if not update_succeeded:
1093 + self._weight_update_active = False
1094 + self._is_checkpoint_format = True
1095
1096 def finish_weight_update(self) -> None:
1075 - """
1076 - Finish the current weight update.
1077 -
1078 - For checkpoint format, this runs layerwise postprocessing.
1079 - Uses the is_checkpoint_format state stored by start_weight_update.
1080 - """
1097 + """Finish the current weight update session."""
1098 self._check_weight_transfer_engine()
1099
1100 if not self._weight_update_active:
1101 raise RuntimeError(
1085 - "start_weight_update must be called before finish_weight_update."
1102 + "finish_weight_update called without a matching start_weight_update."
1103 )
1104
1088 - is_checkpoint_format = self._is_checkpoint_format
1089 -
1090 - if is_checkpoint_format:
1105 + if self._is_checkpoint_format:
1106 from vllm.model_executor.model_loader.reload import (
1107 finalize_layerwise_reload,
1108 )
@@ -1096,7 +1111,6 @@ class Worker(WorkerBase):
1111 with torch.device(self.device):
1112 finalize_layerwise_reload(model, self.model_config)
1113
1099 - # Reset state
1114 self._weight_update_active = False
1115 self._is_checkpoint_format = True
1116