@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6c4aa8bec3

[RL] Support checkpoint-coordinate sparse NCCL weight updates (#53751)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>

Shuolei Wang committed Aug 27, 2026 at 18:19 UTC 6c4aa8bec3029ec03d32194dc7aa0d1252253eaf
5 files changed +436 -798
docs/training/weight_transfer/README.md
+1 -1
@@ -33,7 +33,7 @@ engine drives on your behalf:
33 | ------- | --------- | -------- |
34 | [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
35 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
36 -| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
36 +| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches |
37 | [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) |
38
39 ## Quickstart
docs/training/weight_transfer/nccl.md
+28 -22
@@ -132,17 +132,19 @@ pass — this is the invariant to test first.
132
133 ## Sparse NCCL
134
135 -Sparse, flat-index weight patches use a separate backend,
136 -`WeightTransferConfig(backend="sparse_nccl")`. It shares only NCCL process-group
137 -initialization with the dense engine; patches are applied directly in place to
138 -existing parameters, with no layerwise reload. The current sparse MVP requires
139 -`TP=1` and `PP=1`, and assumes a single-rank trainer.
140 -
141 -Sparse is a **delta** backend: each round's patches are a fresh set of deltas
142 -from the latest optimizer step, not a stable stream of the model's parameters. So
143 -the engine takes **no `WeightSource`** — passing one raises — and each round's
144 -patches go straight to `send_weights(patches)`. A round with no patches is a
145 -no-op.
135 +Sparse, flat-index weight patches use
136 +`WeightTransferConfig(backend="sparse_nccl")`. Names, full shapes, and flat
137 +indices are interpreted in **checkpoint/Hugging Face coordinates**. Every
138 +inference rank receives the same checkpoint-global patch, then the model's native
139 +`load_weights()` maps it to rank-local TP/EP and packed runtime parameters.
140 +
141 +Tensor parallelism is supported and does not have to match the trainer's layout.
142 +Pipeline-parallel behavior is not covered by the current sparse NCCL GPU test.
143 +
144 +Sparse is a **delta** backend: each call supplies replacement patches rather
145 +than a stable stream of the model's parameters. The engine therefore takes no
146 +`WeightSource`; patches go straight to `send_weights(patches)`, and an empty
147 +patch list is a no-op.
148
149 ```python
150 from vllm.distributed.weight_transfer import (
@@ -154,6 +156,7 @@ from vllm.distributed.weight_transfer.sparse_nccl_engine import (
156 SparseWeightPatch,
157 )
158
159 +client = RayVLLMWeightSyncClient(llm)
160 engine = WeightTransferTrainerFactory.trainer_init(
161 init_info=SparseNCCLTrainerInitInfo(
162 master_address=master_address,
@@ -161,7 +164,7 @@ engine = WeightTransferTrainerFactory.trainer_init(
164 world_size=world_size,
165 rank=0,
166 ),
164 - client=RayVLLMWeightSyncClient(llm),
167 + client=client,
168 )
169
170 patches = [
@@ -175,21 +178,24 @@ patches = [
178 engine.send_weights(patches)
179 ```
180
178 -`SparseNCCLTrainerInitInfo` takes the same rendezvous fields as the dense
179 -backend and no packed params — sparse transfers are never packed.
181 +Each `send_weights(patches)` call owns a complete one-shot `start` / `update` / `finish` lifecycle.
182 +
183 +RL infrastructure that owns the generic worker lifecycle can keep one logical update open across bounded chunks without adding trainer-side session state:
184 +
185 +```python
186 +client.start_weight_update()
187 +for patches in patch_chunks:
188 + engine.send_weight_chunk(patches)
189 +client.finish_weight_update()
190 +```
191
181 -Patches are validated on the trainer *before* `start_weight_update`: int32
182 -indices, 1-D, and matching `indices`/`values` lengths, with `full_shape` set. The
183 -worker checks the same invariants when applying, but by then the broadcasts are
184 -under way, where a size mismatch wedges both sides instead of raising.
192 +Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still uses `O(N)` staging through the native loader, and the caller owns export/diff state and restart or reseed after a partial failure.
193
186 -A trainer can hold a dense engine and a sparse engine at once — a full resync
187 -through the dense engine, cheap sparse deltas in between — as
188 -[`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates.
194 +[`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates per-expert checkpoint patches with a Qwen3 MoE TP2/EP2 engine.
195
196 ## Examples
197
198 - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
199 - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
194 -- [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 `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
200 +- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs
201 - [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; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient`
examples/rl/rlhf_sparse_nccl.py
+158 -390
@@ -1,64 +1,45 @@
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.
3 +"""Demonstrate checkpoint-coordinate sparse NCCL updates with expert parallelism.
4 +
5 +The trainer and vLLM start from the same Qwen3 MoE checkpoint. The trainer
6 +modifies rows from two global experts, converts its fused expert storage back to
7 +per-expert checkpoint coordinates, and sends both patches through one
8 +``send_weights()`` lifecycle. Every vLLM rank receives both patches; the native
9 +loader applies its local expert and skips the foreign expert.
10 +
11 +This example uses three GPUs on one node: one for the full BF16 Hugging Face
12 +trainer model and two for a TP2/EP2 vLLM inference engine. Unspecified checkpoint
13 +elements keep their initialized values, so sparse updates require a known shared
14 +baseline. ``SPARSE_NCCL_MODEL`` may point to a compatible local Qwen3 MoE
15 +checkpoint.
16 """
17
34 -import hashlib
18 import os
36 -import time
37 -from collections.abc import Sequence
19 +from contextlib import suppress
20
21 import ray
22 import torch
23 from ray.util.placement_group import placement_group
24 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
43 -from transformers import AutoModelForCausalLM, AutoTokenizer
25 +from transformers import AutoModelForCausalLM
26
27 from vllm import LLM, SamplingParams
28 from vllm.config import WeightTransferConfig
29 from vllm.distributed.weight_transfer import (
48 - ModuleSource,
30 RayVLLMWeightSyncClient,
31 WeightTransferTrainerFactory,
32 )
52 -from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
33 from vllm.distributed.weight_transfer.sparse_nccl_engine import (
34 SparseNCCLTrainerInitInfo,
35 SparseWeightPatch,
36 )
37 from vllm.utils.network_utils import get_ip, get_open_port
38
59 -MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
60 -PATCHED_PARAM_NAME = "model.embed_tokens.weight"
61 -MAX_PATCH_ROWS = 32
39 +MODEL_NAME = os.environ.get("SPARSE_NCCL_MODEL", "Qwen/Qwen3-30B-A3B")
40 +INFERENCE_TP_SIZE = 2
41 +PATCHED_LAYER = 0
42 +PATCHED_ROWS = 2
43 PROMPTS = [
44 "Hello, my name is",
45 "The president of the United States is",
@@ -72,414 +53,201 @@ class MyLLM(LLM):
53 """Configure the vLLM worker for Ray placement group execution."""
54
55 def __init__(self, *args, **kwargs):
75 - os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
56 + os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(
57 + str(index) for index in range(INFERENCE_TP_SIZE)
58 + )
59 super().__init__(*args, **kwargs)
60
61
62 @ray.remote(num_gpus=1)
63 class TrainModel:
81 - """Ray actor that owns the trainer-side model and deterministic patch state."""
64 + """Own the trainer model and sparse NCCL sender on one GPU."""
65
66 def __init__(self, model_name: str):
84 - self.model_name = model_name
85 - self.tokenizer = AutoTokenizer.from_pretrained(model_name)
86 - if self.tokenizer.pad_token_id is None:
87 - self.tokenizer.pad_token = self.tokenizer.eos_token
88 -
89 - self.model = None
90 - self.patched_param = None
91 - self.pending_sparse_patches: list[SparseWeightPatch] | None = None
92 - self.dense_engine = None
93 - self.sparse_engine = None
94 - self.master_address = get_ip()
95 - self.port = get_open_port()
96 - self.reset_model()
97 -
98 - def reset_model(self) -> None:
67 self.model = AutoModelForCausalLM.from_pretrained(
100 - self.model_name,
68 + model_name,
69 torch_dtype=torch.bfloat16,
70 ).to("cuda:0")
71 self.model.eval()
72
73 + config = self.model.config
74 + if config.model_type != "qwen3_moe" or config.num_experts % INFERENCE_TP_SIZE:
75 + raise RuntimeError(
76 + "This recipe requires a Qwen3 MoE model whose experts divide "
77 + f"evenly across TP{INFERENCE_TP_SIZE}"
78 + )
79 + self.expert_ids = (0, config.num_experts // 2)
80 + self.expert_intermediate_size = config.moe_intermediate_size
81 + fused_name = f"model.layers.{PATCHED_LAYER}.mlp.experts.gate_up_proj"
82 try:
106 - self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
83 + self.gate_up_proj = self.model.get_parameter(fused_name)
84 except AttributeError as exc:
85 raise RuntimeError(
109 - f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
86 + f"Expected trainer model to expose `{fused_name}`"
87 ) from exc
88
112 - self.pending_sparse_patches = None
89 + expected_shape = (
90 + config.num_experts,
91 + 2 * self.expert_intermediate_size,
92 + config.hidden_size,
93 + )
94 + if self.gate_up_proj.shape != expected_shape:
95 + raise RuntimeError(
96 + f"Unexpected fused expert shape: {self.gate_up_proj.shape} "
97 + f"!= {expected_shape}"
98 + )
99
114 - def create_rendezvous(self) -> tuple[str, int]:
115 - self.port = get_open_port()
116 - return self.master_address, self.port
100 + self.master_address = get_ip()
101 + self.master_port = get_open_port()
102 + self.engine = None
103
118 - def init_dense_engine(self, world_size: int, llm_handle) -> None:
119 - """Build the dense NCCL trainer engine"""
120 - self.dense_engine = WeightTransferTrainerFactory.trainer_init(
121 - init_info=NCCLTrainerInitInfo(
104 + def init_sparse_engine(self, world_size: int, llm_handle) -> None:
105 + self.engine = WeightTransferTrainerFactory.trainer_init(
106 + init_info=SparseNCCLTrainerInitInfo(
107 master_address=self.master_address,
123 - master_port=self.port,
108 + master_port=self.master_port,
109 world_size=world_size,
125 - rank=0, # single-GPU trainer is the sole (sender) rank
126 - packed=True,
110 + rank=0,
111 ),
112 client=RayVLLMWeightSyncClient(llm_handle),
129 - source=ModuleSource(self.model),
130 - )
131 -
132 - def shutdown_engines(self) -> None:
133 - """Drop the NCCL communicators before the vLLM engine they point at
134 - goes away, so the next phase rendezvouses from a clean slate."""
135 - for attr in ("dense_engine", "sparse_engine"):
136 - engine = getattr(self, attr)
137 - if engine is not None:
138 - engine.shutdown()
139 - setattr(self, attr, None)
140 -
141 - def dense_payload_bytes(self) -> int:
142 - return sum(
143 - param.numel() * param.element_size() for param in self.model.parameters()
113 )
114
146 - @torch.inference_mode()
147 - def generate(
148 - self,
149 - prompts: Sequence[str],
150 - max_new_tokens: int = 1,
151 - ) -> list[dict[str, object]]:
152 - generations = []
153 - for prompt in prompts:
154 - model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
155 - output = self.model.generate(
156 - **model_inputs,
157 - max_new_tokens=max_new_tokens,
158 - do_sample=False,
159 - pad_token_id=self.tokenizer.pad_token_id,
115 + @torch.no_grad()
116 + def patch_and_send(self) -> tuple[list[str], int]:
117 + if self.engine is None:
118 + raise RuntimeError("Sparse NCCL engine is not initialized")
119 +
120 + patches = []
121 + for expert_id in self.expert_ids:
122 + checkpoint_weight = self.gate_up_proj[
123 + expert_id, : self.expert_intermediate_size
124 + ]
125 + original_rows = checkpoint_weight[:PATCHED_ROWS].clone()
126 + replacement_rows = original_rows.flip(0)
127 + if torch.equal(original_rows, replacement_rows):
128 + raise RuntimeError(f"Expert {expert_id} patch would be a no-op")
129 + checkpoint_weight[:PATCHED_ROWS] = replacement_rows
130 +
131 + hidden_size = checkpoint_weight.shape[1]
132 + flat_indices = torch.arange(
133 + PATCHED_ROWS * hidden_size,
134 + device=checkpoint_weight.device,
135 + dtype=torch.int32,
136 )
161 - new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
162 - generations.append(
163 - {
164 - "token_ids": new_token_ids,
165 - "text": self.tokenizer.decode(
166 - new_token_ids,
167 - skip_special_tokens=False,
137 + patches.append(
138 + SparseWeightPatch(
139 + name=(
140 + f"model.layers.{PATCHED_LAYER}.mlp.experts."
141 + f"{expert_id}.gate_proj.weight"
142 ),
169 - }
143 + full_shape=tuple(checkpoint_weight.shape),
144 + indices=flat_indices,
145 + values=replacement_rows.reshape(-1).contiguous(),
146 + )
147 )
171 - return generations
172 -
173 - def prepare_sparse_patch(
174 - self,
175 - prompts: Sequence[str],
176 - max_patch_rows: int = MAX_PATCH_ROWS,
177 - ) -> tuple[list[int], str, int]:
178 - selected_token_ids: list[int] = []
179 - special_ids = set(self.tokenizer.all_special_ids)
180 - for prompt in prompts:
181 - token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
182 - for token_id in token_ids:
183 - if token_id in special_ids or token_id in selected_token_ids:
184 - continue
185 - selected_token_ids.append(token_id)
186 - if len(selected_token_ids) == max_patch_rows:
187 - break
188 - if len(selected_token_ids) == max_patch_rows:
189 - break
190 -
191 - if not selected_token_ids:
192 - raise ValueError("Could not derive any non-special token IDs to patch")
193 -
194 - vocab_size = self.patched_param.shape[0]
195 - next_token_id = selected_token_ids[-1]
196 - while len(selected_token_ids) < max_patch_rows:
197 - next_token_id = (next_token_id + 1) % vocab_size
198 - if next_token_id in special_ids or next_token_id in selected_token_ids:
199 - continue
200 - selected_token_ids.append(next_token_id)
201 -
202 - row_ids = torch.tensor(
203 - selected_token_ids,
204 - device=self.patched_param.device,
205 - dtype=torch.long,
206 - )
207 - hidden_size = self.patched_param.shape[1]
208 - column_offsets = torch.arange(
209 - hidden_size,
210 - device=self.patched_param.device,
211 - dtype=torch.long,
212 - )
148
214 - with torch.no_grad():
215 - # Rotate the selected embedding rows instead of zeroing them so the
216 - # patch remains deterministic while avoiding a degenerate collapse
217 - # to the same special token after the update.
218 - replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
219 - self.patched_param[row_ids] = replacement_rows
220 -
221 - flat_indices = (
222 - row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
149 + self.engine.send_weights(patches)
150 + return [patch.name for patch in patches], sum(
151 + patch.indices.numel() for patch in patches
152 )
224 - flat_values = self.patched_param[row_ids].detach().reshape(-1).contiguous()
225 - self.pending_sparse_patches = [
226 - SparseWeightPatch(
227 - name=PATCHED_PARAM_NAME,
228 - indices=flat_indices.to(torch.int32),
229 - values=flat_values,
230 - full_shape=tuple(self.patched_param.shape),
231 - )
232 - ]
233 - patch_digest = hashlib.sha256(
234 - self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
235 - + self.pending_sparse_patches[0].values.float().cpu().numpy().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 - return selected_token_ids, patch_digest, sparse_payload_bytes
243 -
244 - def send_dense_weights(self) -> float:
245 - """One call drives start/update/finish on the inference side
246 - concurrently with the NCCL broadcast."""
247 - if self.dense_engine is None:
248 - raise RuntimeError("Dense engine is not initialized")
153
250 - start = time.perf_counter()
251 - self.dense_engine.send_weights()
252 - torch.accelerator.synchronize()
253 - return (time.perf_counter() - start) * 1000.0
154 + def shutdown_engine(self) -> None:
155 + if self.engine is not None:
156 + self.engine.shutdown()
157 + self.engine = None
158
255 - def init_sparse_engine(self, world_size: int, llm_handle) -> None:
256 - """Build the sparse trainer engine"""
257 - self.sparse_engine = WeightTransferTrainerFactory.trainer_init(
258 - init_info=SparseNCCLTrainerInitInfo(
259 - master_address=self.master_address,
260 - master_port=self.port,
261 - world_size=world_size,
262 - rank=0, # single-GPU trainer is the sole (sender) rank
263 - ),
264 - client=RayVLLMWeightSyncClient(llm_handle),
265 - )
159
267 - def send_pending_sparse_patch(self) -> float:
268 - """One call drives start/update/finish on the inference side
269 - concurrently with the patch broadcasts. Sparse deltas differ each round,
270 - so the patches are passed to `send_weights` rather than fixed at init."""
271 - if self.sparse_engine is None:
272 - raise RuntimeError("Sparse engine is not initialized")
273 - if self.pending_sparse_patches is None:
274 - raise RuntimeError("Sparse patch has not been prepared")
275 -
276 - start = time.perf_counter()
277 - self.sparse_engine.send_weights(self.pending_sparse_patches)
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 - backend: str = "nccl",
286 -):
160 +def launch_llm(scheduling_strategy: PlacementGroupSchedulingStrategy):
161 return ray.remote(
162 num_cpus=0,
163 num_gpus=0,
290 - scheduling_strategy=scheduling_inference,
164 + scheduling_strategy=scheduling_strategy,
165 )(MyLLM).remote(
166 model=MODEL_NAME,
167 enforce_eager=True,
294 - tensor_parallel_size=1,
168 + tensor_parallel_size=INFERENCE_TP_SIZE,
169 + enable_expert_parallel=True,
170 + expert_placement_strategy="linear",
171 + moe_backend="triton",
172 distributed_executor_backend="ray",
173 gpu_memory_utilization=0.7,
297 - weight_transfer_config=WeightTransferConfig(backend=backend),
174 + enable_prefix_caching=False,
175 + weight_transfer_config=WeightTransferConfig(backend="sparse_nccl"),
176 )
177
178
301 -def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
179 +def generate(llm_handle) -> list[dict[str, object]]:
180 outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
303 - generations = []
304 - for output in outputs:
305 - generations.append(
306 - {
307 - "token_ids": output.outputs[0].token_ids,
308 - "text": output.outputs[0].text,
309 - }
310 - )
311 - return generations
312 -
313 -
314 -def token_sequences_match(
315 - left: Sequence[dict[str, object]],
316 - right: Sequence[dict[str, object]],
317 -) -> bool:
318 - return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
181 + return [
182 + {
183 + "token_ids": output.outputs[0].token_ids,
184 + "text": output.outputs[0].text,
185 + }
186 + for output in outputs
187 + ]
188
189
321 -def print_generations(label: str, prompts: Sequence[str], generations) -> None:
190 +def print_generations(label: str, generations: list[dict[str, object]]) -> None:
191 print(f"\n{label}")
323 - print("-" * 50)
324 - for prompt, generation in zip(prompts, generations):
325 - print(f"Prompt: {prompt!r}")
326 - print(f"Token IDs: {generation['token_ids']}")
327 - print(f"Text: {generation['text']!r}")
328 - print("-" * 50)
329 -
330 -
331 -def run_dense_phase(
332 - train_model,
333 - scheduling_inference: PlacementGroupSchedulingStrategy,
334 -) -> dict[str, object]:
335 - ray.get(train_model.reset_model.remote())
336 - llm = launch_llm(scheduling_inference, backend="nccl")
337 - try:
338 - dense_before = collect_vllm_generations(llm)
339 -
340 - ray.get(llm.sleep.remote(level=0))
341 - ray.get(train_model.create_rendezvous.remote())
342 - world_size = ray.get(llm.get_world_size.remote()) + 1
343 - ray.get(train_model.init_dense_engine.remote(world_size, llm))
344 -
345 - dense_payload_bytes = ray.get(train_model.dense_payload_bytes.remote())
346 - selected_token_ids, patch_digest, _ = ray.get(
347 - train_model.prepare_sparse_patch.remote(PROMPTS)
192 + for prompt, generation in zip(PROMPTS, generations):
193 + print(
194 + f" {prompt!r} -> {generation['text']!r} "
195 + f"(token_ids={generation['token_ids']})"
196 )
197
350 - # One call drives start/update/finish + the NCCL broadcast.
351 - dense_send_ms = ray.get(train_model.send_dense_weights.remote())
352 - ray.get(llm.wake_up.remote(tags=["scheduling"]))
353 -
354 - dense_after = collect_vllm_generations(llm)
198
356 - return {
357 - "dense_before": dense_before,
358 - "dense_after": dense_after,
359 - "selected_token_ids": selected_token_ids,
360 - "patch_digest": patch_digest,
361 - "dense_payload_bytes": dense_payload_bytes,
362 - "dense_send_ms": dense_send_ms,
363 - }
364 - finally:
365 - ray.get(train_model.shutdown_engines.remote())
366 - ray.kill(llm)
367 -
368 -
369 -def run_sparse_phase(
370 - train_model,
371 - scheduling_inference: PlacementGroupSchedulingStrategy,
372 -) -> dict[str, object]:
373 - ray.get(train_model.reset_model.remote())
374 - llm = launch_llm(scheduling_inference, backend="sparse_nccl")
199 +def main() -> None:
200 + ray.init()
201 + train_model = None
202 + pg_inference = None
203 + llm = None
204 try:
376 - sparse_before = collect_vllm_generations(llm)
205 + train_model = TrainModel.remote(MODEL_NAME)
206 + pg_inference = placement_group(
207 + [{"GPU": 1, "CPU": 0}] * INFERENCE_TP_SIZE,
208 + strategy="STRICT_PACK",
209 + )
210 + ray.get(pg_inference.ready())
211 + scheduling_strategy = PlacementGroupSchedulingStrategy(
212 + placement_group=pg_inference,
213 + placement_group_capture_child_tasks=True,
214 + placement_group_bundle_index=0,
215 + )
216 + llm = launch_llm(scheduling_strategy)
217 +
218 + before = generate(llm)
219 + print_generations("BEFORE sparse update", before)
220
221 ray.get(llm.sleep.remote(level=0))
379 - ray.get(train_model.create_rendezvous.remote())
222 world_size = ray.get(llm.get_world_size.remote()) + 1
223 ray.get(train_model.init_sparse_engine.remote(world_size, llm))
382 -
383 - selected_token_ids, patch_digest, sparse_payload_bytes = ray.get(
384 - train_model.prepare_sparse_patch.remote(PROMPTS)
385 - )
386 - sparse_send_ms = ray.get(train_model.send_pending_sparse_patch.remote())
224 + patched_names, num_updates = ray.get(train_model.patch_and_send.remote())
225 ray.get(llm.wake_up.remote(tags=["scheduling"]))
226
389 - sparse_after = collect_vllm_generations(llm)
390 -
391 - return {
392 - "sparse_before": sparse_before,
393 - "sparse_after": sparse_after,
394 - "selected_token_ids": selected_token_ids,
395 - "patch_digest": patch_digest,
396 - "sparse_payload_bytes": sparse_payload_bytes,
397 - "sparse_send_ms": sparse_send_ms,
398 - }
399 - finally:
400 - ray.get(train_model.shutdown_engines.remote())
401 - ray.kill(llm)
402 -
403 -
404 -ray.init()
405 -
406 -try:
407 - train_model = TrainModel.remote(MODEL_NAME)
408 -
409 - pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
410 - ray.get(pg_inference.ready())
411 - scheduling_inference = PlacementGroupSchedulingStrategy(
412 - placement_group=pg_inference,
413 - placement_group_capture_child_tasks=True,
414 - placement_group_bundle_index=0,
415 - )
416 -
417 - dense_results = run_dense_phase(train_model, scheduling_inference)
418 - sparse_results = run_sparse_phase(train_model, scheduling_inference)
419 -
420 - baseline_equal = token_sequences_match(
421 - dense_results["dense_before"],
422 - sparse_results["sparse_before"],
423 - )
424 - patch_selection_equal = (
425 - dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
426 - )
427 - patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
428 - after_equal = token_sequences_match(
429 - dense_results["dense_after"],
430 - sparse_results["sparse_after"],
431 - )
432 - any_output_changed = any(
433 - before["token_ids"] != after["token_ids"]
434 - for before, after in zip(
435 - dense_results["dense_before"],
436 - dense_results["dense_after"],
227 + after = generate(llm)
228 + print_generations("AFTER sparse update", after)
229 + print(f"patched_checkpoint_names={patched_names}")
230 + print(f"num_sparse_values={num_updates}")
231 + outputs_changed = any(
232 + old["token_ids"] != new["token_ids"]
233 + for old, new in zip(before, after, strict=True)
234 )
438 - )
439 - dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
440 - sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
441 -
442 - print_generations(
443 - "Dense baseline outputs",
444 - PROMPTS,
445 - dense_results["dense_before"],
446 - )
447 - print_generations(
448 - "Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
449 - )
450 - print_generations(
451 - "Dense outputs after update", PROMPTS, dense_results["dense_after"]
452 - )
453 - print_generations(
454 - "Sparse outputs after update",
455 - PROMPTS,
456 - sparse_results["sparse_after"],
457 - )
458 -
459 - print(f"patched_token_ids = {dense_results['selected_token_ids']}")
460 - print(f"patch_selection_equal = {patch_selection_equal}")
461 - print(f"dense_patch_digest = {dense_results['patch_digest']}")
462 - print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
463 - print(f"patch_digest_equal = {patch_digest_equal}")
464 - print(f"baseline_equal = {baseline_equal}")
465 - print(f"after_equal = {after_equal}")
466 - print(f"any_output_changed = {any_output_changed}")
467 - print(f"dense_payload_mb = {dense_payload_mb:.2f}")
468 - print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
469 - print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
470 - print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
471 -
472 - if not baseline_equal:
473 - raise RuntimeError(
474 - "Dense and sparse phases did not start from the same baseline"
475 - )
476 - if not patch_selection_equal:
477 - raise RuntimeError("Dense and sparse phases used different sparse patches")
478 - if not patch_digest_equal:
479 - raise RuntimeError("Dense and sparse phases produced different patch values")
480 - if not after_equal:
481 - raise RuntimeError("Dense and sparse updates produced different outputs")
482 - if not any_output_changed:
483 - raise RuntimeError("Patch did not change the observed outputs")
484 -finally:
485 - ray.shutdown()
235 + print(f"outputs_changed={outputs_changed}")
236 + finally:
237 + if train_model is not None:
238 + with suppress(Exception):
239 + ray.get(train_model.shutdown_engine.remote())
240 + if llm is not None:
241 + with suppress(Exception):
242 + ray.kill(llm)
243 + if train_model is not None:
244 + with suppress(Exception):
245 + ray.kill(train_model)
246 + if pg_inference is not None:
247 + with suppress(Exception):
248 + ray.util.remove_placement_group(pg_inference)
249 + ray.shutdown()
250 +
251 +
252 +if __name__ == "__main__":
253 + main()
tests/distributed/test_weight_transfer.py
+122 -239
@@ -9,6 +9,7 @@ Integration tests for NCCL and IPC weight transfer between processes using Ray.
9 import pickle
10 import threading
11 import time
12 +from contextlib import nullcontext
13 from unittest.mock import MagicMock
14
15 import pybase64 as base64
@@ -313,103 +314,6 @@ class TestEngineRegistry:
314 )
315
316
316 -# --- Unit Tests: Sparse patch application (CPU) ---
317 -
318 -
319 -class TestSparseNCCLPatchApplication:
320 - """Test SparseNCCLWeightTransferEngine._apply_patch on a real param."""
321 -
322 - def _make_engine(self, model):
323 - config = WeightTransferConfig(backend="sparse_nccl")
324 - return SparseNCCLWeightTransferEngine(
325 - config, create_mock_vllm_config(), torch.device("cpu"), model
326 - )
327 -
328 - def _make_model(self, numel: int = 8):
329 - model = torch.nn.Module()
330 - model.register_parameter(
331 - "w", torch.nn.Parameter(torch.zeros(numel), requires_grad=False)
332 - )
333 -
334 - def get_parameter(name):
335 - assert name == "w"
336 - return model.w
337 -
338 - model.get_parameter = get_parameter
339 - return model
340 -
341 - def test_apply_patch_updates_only_selected_entries(self):
342 - model = self._make_model(8)
343 - engine = self._make_engine(model)
344 - engine._apply_patch(
345 - SparseWeightPatch(
346 - name="w",
347 - indices=torch.tensor([1, 3], dtype=torch.int32),
348 - values=torch.tensor([5.0, 7.0], dtype=torch.float32),
349 - )
350 - )
351 - expected = torch.zeros(8)
352 - expected[1] = 5.0
353 - expected[3] = 7.0
354 - assert torch.equal(model.w.data, expected)
355 -
356 - def test_apply_patch_rejects_mismatched_lengths(self):
357 - model = self._make_model(8)
358 - engine = self._make_engine(model)
359 - with pytest.raises(ValueError, match="matching lengths"):
360 - engine._apply_patch(
361 - SparseWeightPatch(
362 - name="w",
363 - indices=torch.tensor([1, 3], dtype=torch.int32),
364 - values=torch.tensor([5.0], dtype=torch.float32),
365 - )
366 - )
367 -
368 - def test_apply_patch_rejects_non_int32_indices(self):
369 - model = self._make_model(8)
370 - engine = self._make_engine(model)
371 - with pytest.raises(ValueError, match="int32 indices"):
372 - engine._apply_patch(
373 - SparseWeightPatch(
374 - name="w",
375 - indices=torch.tensor([1], dtype=torch.int64),
376 - values=torch.tensor([5.0], dtype=torch.float32),
377 - )
378 - )
379 -
380 - def test_apply_patch_rejects_dtype_mismatch(self):
381 - model = self._make_model(8)
382 - engine = self._make_engine(model)
383 - with pytest.raises(ValueError, match="does not match"):
384 - engine._apply_patch(
385 - SparseWeightPatch(
386 - name="w",
387 - indices=torch.tensor([1], dtype=torch.int32),
388 - values=torch.tensor([5.0], dtype=torch.bfloat16),
389 - )
390 - )
391 -
392 - def test_apply_patch_rejects_non_contiguous_param(self):
393 - model = torch.nn.Module()
394 - model.register_parameter(
395 - "w",
396 - torch.nn.Parameter(
397 - torch.arange(12, dtype=torch.float32).view(3, 4).t(),
398 - requires_grad=False,
399 - ),
400 - )
401 - model.get_parameter = lambda name: model.w
402 - engine = self._make_engine(model)
403 - with pytest.raises(NotImplementedError, match="contiguous params"):
404 - engine._apply_patch(
405 - SparseWeightPatch(
406 - name="w",
407 - indices=torch.tensor([1], dtype=torch.int32),
408 - values=torch.tensor([1.0], dtype=torch.float32),
409 - )
410 - )
411 -
412 -
317 # --- Test receive_weights without init raises ---
318
319
@@ -628,162 +532,135 @@ def test_nccl_weight_transfer_between_processes():
532 )
533
534
631 -@ray.remote(num_gpus=1)
632 -def trainer_broadcast_sparse_tensor(
633 - master_address: str,
634 - master_port: int,
635 - world_size: int,
636 -) -> bool:
637 - """Trainer task that broadcasts sparse patches via the trainer engine.
638 -
639 - The worker task drives its own init/receive directly (it is not an RPC
640 - endpoint), so the engine gets a no-op control-plane client; the NCCL
641 - rendezvous and the patch broadcasts are the real thing.
642 - """
643 - import torch
535 +def test_sparse_nccl_checkpoint_chunks_to_ep_local_experts_cpu(monkeypatch):
536 + """Replay global expert patches through two EP-local loaders on CPU."""
537
645 - device = _set_ray_assigned_device()
646 -
647 - from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
648 - from vllm.distributed.weight_transfer.sparse_nccl_engine import (
649 - SparseNCCLTrainerInitInfo,
650 - SparseWeightPatch,
651 - )
652 -
653 - class NoopClient:
654 - def init_weight_transfer_engine(self, init_info):
655 - pass
538 + class CaptureClient:
539 + def __init__(self):
540 + self.order: list[str] = []
541 + self.update_infos: list[dict] = []
542
543 def start_weight_update(self):
658 - pass
544 + self.order.append("start")
545
546 def update_weights(self, update_info):
661 - pass
547 + self.order.append("update")
548 + self.update_infos.append(update_info)
549
550 def finish_weight_update(self):
664 - pass
665 -
666 - patch = SparseWeightPatch(
667 - name="test.weight",
668 - indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device),
669 - values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device),
670 - full_shape=(10, 10),
671 - )
672 - engine = WeightTransferTrainerFactory.trainer_init(
673 - init_info=SparseNCCLTrainerInitInfo(
674 - master_address=master_address,
675 - master_port=master_port,
676 - world_size=world_size,
677 - rank=0,
678 - ),
679 - client=NoopClient(),
551 + self.order.append("finish")
552 +
553 + monkeypatch.setattr(torch.cuda, "current_stream", lambda *_, **__: None)
554 + monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
555 + monkeypatch.setattr(torch.accelerator, "device_index", lambda *_: nullcontext())
556 +
557 + client = CaptureClient()
558 + sender = SparseNCCLTrainerWeightTransferEngine(client=client)
559 + sender.model_update_group = MagicMock()
560 + sender.model_update_group.device = torch.device("cpu")
561 + wire_payloads = []
562 + sender.model_update_group.broadcast.side_effect = lambda tensor, **_: (
563 + wire_payloads.append(tensor.clone())
564 )
681 - engine.send_weights([patch])
682 - torch.accelerator.synchronize()
683 - engine.shutdown()
684 - return True
685 -
686 -
687 -@ray.remote(num_gpus=1)
688 -def inference_receive_sparse_tensor(
689 - master_address: str,
690 - master_port: int,
691 - world_size: int,
692 -) -> dict:
693 - """Inference task that receives sparse patches via the sparse engine."""
694 - from unittest.mock import MagicMock
695 -
696 - import torch
565
698 - device = _set_ray_assigned_device()
699 -
700 - from vllm.config.parallel import ParallelConfig
701 - from vllm.config.weight_transfer import WeightTransferConfig
702 - from vllm.distributed.weight_transfer.sparse_nccl_engine import (
703 - SparseNCCLWeightTransferEngine,
704 - SparseNCCLWeightTransferUpdateInfo,
705 - )
706 -
707 - config = WeightTransferConfig(backend="sparse_nccl")
708 - vllm_config = MagicMock()
709 - parallel_config = MagicMock(spec=ParallelConfig)
710 - parallel_config.rank = 0
711 - parallel_config.world_size = 1
712 - parallel_config.data_parallel_rank = 0
713 - parallel_config.data_parallel_index = 0
714 - vllm_config.parallel_config = parallel_config
715 - vllm_config.model_config = MagicMock()
716 -
717 - # Real module holding the target parameter the patch will modify.
718 - model = torch.nn.Module()
719 - model.register_parameter(
720 - "w", torch.nn.Parameter(torch.zeros(30, device="cuda"), requires_grad=False)
721 - )
722 - model.get_parameter = lambda name: model.w
566 + expert_names = [
567 + "model.layers.0.mlp.experts.0.gate_proj.weight",
568 + "model.layers.0.mlp.experts.1.gate_proj.weight",
569 + ]
570 + client.start_weight_update()
571 + for name, index, value in zip(
572 + expert_names,
573 + (0, 3),
574 + (5.0, 7.0),
575 + strict=True,
576 + ):
577 + sender.send_weight_chunk(
578 + [
579 + SparseWeightPatch(
580 + name=name,
581 + indices=torch.tensor([index], dtype=torch.int32),
582 + values=torch.tensor([value]),
583 + full_shape=(2, 2),
584 + )
585 + ]
586 + )
587 + client.finish_weight_update()
588
724 - update_info = SparseNCCLWeightTransferUpdateInfo(
725 - names=["w"],
726 - dtype_names=["float32"],
727 - shapes=[[30]],
728 - num_updates_list=[3],
729 - )
589 + assert client.order == ["start", "update", "update", "finish"]
590 + assert [info["names"] for info in client.update_infos] == [
591 + [expert_names[0]],
592 + [expert_names[1]],
593 + ]
594 + assert [info["shapes"] for info in client.update_infos] == [
595 + [[2, 2]],
596 + [[2, 2]],
597 + ]
598 + assert [info["num_updates_list"] for info in client.update_infos] == [
599 + [1],
600 + [1],
601 + ]
602
731 - engine = SparseNCCLWeightTransferEngine(
732 - config, vllm_config, torch.device("cuda"), model
733 - )
734 - from vllm.distributed.weight_transfer.nccl_common import (
735 - NCCLWeightTransferInitInfo,
603 + expected_gates = (
604 + [[5.0, -1.0], [-1.0, -1.0]],
605 + [[-1.0, -1.0], [-1.0, 7.0]],
606 )
737 -
738 - engine.init_transfer_engine(
739 - NCCLWeightTransferInitInfo(
740 - master_address=master_address,
741 - master_port=master_port,
742 - rank_offset=1,
743 - world_size=world_size,
607 + for ep_rank, expected_gate in enumerate(expected_gates):
608 + model = torch.nn.Module()
609 + model.register_parameter(
610 + "w13_weight",
611 + torch.nn.Parameter(torch.full((1, 4, 2), -1.0), requires_grad=False),
612 + )
613 + local_name = expert_names[ep_rank]
614 + seen_names: list[str] = []
615 + loaded_names: set[str] = set()
616 + copy_count = [0]
617 +
618 + def load_weights(
619 + weights,
620 + local_name=local_name,
621 + target=model,
622 + seen=seen_names,
623 + loaded=loaded_names,
624 + count=copy_count,
625 + ):
626 + for name, checkpoint_weight in weights:
627 + seen.append(name)
628 + if name != local_name:
629 + continue
630 + target.w13_weight.data[0, :2].copy_(checkpoint_weight)
631 + loaded.add(name)
632 + count[0] += 1
633 + return loaded
634 +
635 + model.load_weights = load_weights
636 + receiver = SparseNCCLWeightTransferEngine(
637 + WeightTransferConfig(backend="sparse_nccl"),
638 + create_mock_vllm_config(rank=ep_rank, world_size=2),
639 + torch.device("cpu"),
640 + model,
641 + )
642 + receiver.model_update_group = MagicMock()
643 + receiver.model_update_group.device = torch.device("cpu")
644 + payloads = iter(wire_payloads)
645 + receiver.model_update_group.broadcast.side_effect = (
646 + lambda tensor, payloads=payloads, **_: tensor.copy_(next(payloads))
647 )
745 - )
746 - engine.receive_weights(update_info)
747 - torch.accelerator.synchronize()
748 -
749 - expected = torch.zeros(30, dtype=torch.float32, device=device)
750 - expected[[1, 7, 25]] = torch.tensor(
751 - [10.0, 20.0, 30.0], dtype=torch.float32, device=device
752 - )
753 - success = torch.equal(model.w.data, expected)
754 - engine.shutdown()
755 - return {
756 - "success": success,
757 - "selected_values": model.w.data[[1, 7, 25]].cpu().tolist(),
758 - }
759 -
760 -
761 -@pytest.mark.skipif(
762 - torch.accelerator.device_count() < 2,
763 - reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
764 -)
765 -def test_nccl_sparse_weight_transfer_between_processes():
766 - """Test NCCL sparse weight transfer from trainer to inference process."""
767 - _init_ray_for_weight_transfer()
768 -
769 - master_address = "127.0.0.1"
770 - master_port = get_open_port()
771 - world_size = 2
648
773 - inference_future = inference_receive_sparse_tensor.remote(
774 - master_address, master_port, world_size
775 - )
776 - trainer_future = trainer_broadcast_sparse_tensor.remote(
777 - master_address, master_port, world_size
778 - )
649 + receiver.start_weight_update()
650 + for update_info in client.update_infos:
651 + receiver.receive_weights(SparseNCCLWeightTransferUpdateInfo(**update_info))
652 + receiver.finish_weight_update()
653
780 - trainer_result, result = ray.get([trainer_future, inference_future])
654 + assert seen_names == expert_names
655 + assert loaded_names == {local_name}
656 + assert copy_count == [1]
657 + assert model.w13_weight[0, :2].tolist() == expected_gate
658 + assert model.w13_weight[0, 2:].eq(-1).all()
659 + foreign_value = 7.0 if ep_rank == 0 else 5.0
660 + assert not model.w13_weight.eq(foreign_value).any()
661 + receiver.shutdown()
662
782 - assert trainer_result, "Trainer should complete successfully"
783 - assert result["success"], (
784 - "Sparse weight transfer failed. "
785 - f"Received selected values: {result['selected_values']}"
786 - )
663 + sender.shutdown()
664
665
666 # --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
@@ -1995,11 +1872,11 @@ def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch):
1872
1873
1874 def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
1998 - """send_weights takes the round's patches and ships per-patch metadata
1999 - (names / shapes / num_updates_list) + broadcasts indices + values each."""
1875 + """One-shot send drives start, update, and finish in order."""
1876 client = RecordingClient()
1877 engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1878 engine.model_update_group = MagicMock()
1879 + engine.model_update_group.device = torch.device("cpu")
1880 # The group is a mock, so the stream is just a handle it is handed (and a
1881 # handle _post_send_sync can synchronize).
1882 monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
@@ -2013,6 +1890,7 @@ def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
1890 assert client.last_update_info["num_updates_list"] == [2]
1891 # One broadcast for indices + one for values per patch.
1892 assert engine.model_update_group.broadcast.call_count == 2
1893 + engine.shutdown()
1894
1895
1896 def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
@@ -2024,6 +1902,8 @@ def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
1902
1903 engine.send_weights([])
1904 engine.send_weights() # no argument is also a no-op round
1905 + engine.send_weight_chunk([])
1906 + engine.send_weight_chunk()
1907
1908 assert client.order == []
1909
@@ -2033,6 +1913,7 @@ def test_sparse_nccl_trainer_send_weights_requires_full_shape():
1913 patch.full_shape = None
1914 engine = SparseNCCLTrainerWeightTransferEngine(client=RecordingClient())
1915 engine.model_update_group = MagicMock()
1916 + engine.model_update_group.device = torch.device("cpu")
1917
1918 with pytest.raises(ValueError, match="full_shape"):
1919 engine.send_weights([patch])
@@ -2054,6 +1935,7 @@ def test_sparse_nccl_trainer_validates_patch_before_any_rpc():
1935 client = RecordingClient()
1936 engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1937 engine.model_update_group = MagicMock()
1938 + engine.model_update_group.device = torch.device("cpu")
1939
1940 mismatched = SparseWeightPatch(
1941 name="w",
@@ -2092,4 +1974,5 @@ def test_sparse_nccl_trainer_non_sender_skips_client():
1974 assert engine.model_update_group is None
1975 assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine)
1976 engine.send_weights([_sparse_patch()])
1977 + engine.send_weight_chunk([_sparse_patch()])
1978 assert client.order == []
vllm/distributed/weight_transfer/sparse_nccl_engine.py
+127 -146
@@ -2,20 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Sparse NCCL weight transfer engine.
4
5 -A standalone engine (not a subclass of `NCCLWeightTransferEngine`) for applying
6 -sparse, flat-index weight patches in place. It shares only NCCL process-group
7 -initialization with the dense engine (via `nccl_common`); the update path
8 -applies index/value patches directly to existing model parameters and never runs
9 -layerwise reload.
10 -
11 -MVP limitations:
12 -* TP=1 and PP=1 only
13 -* uses runtime/kernel-format parameter names
14 -* not composable with checkpoint-format or packed updates
5 +Sparse patches use checkpoint names, shapes, and flat indices. The model's native
6 +weight loader maps them to rank-local runtime parameters, including TP shards and
7 +packed parameters.
8 """
9
10 from collections.abc import Iterable
11 from concurrent.futures import ThreadPoolExecutor
12 +from contextlib import nullcontext
13 from dataclasses import asdict, dataclass
14 from typing import TYPE_CHECKING, ClassVar
15
@@ -42,6 +36,10 @@ from vllm.distributed.weight_transfer.nccl_common import (
36 from vllm.distributed.weight_transfer.nccl_common import (
37 trainer_init as open_trainer_endpoint,
38 )
39 +from vllm.model_executor.model_loader.checkpoint_weight_patch import (
40 + CheckpointWeightPatch,
41 + load_checkpoint_weight_patches,
42 +)
43
44 __all__ = [
45 "SparseWeightPatch",
@@ -54,15 +52,13 @@ __all__ = [
52
53 @dataclass
54 class SparseWeightPatch:
57 - """A sparse in-place patch for one existing parameter."""
55 + """A sparse patch in checkpoint coordinates."""
56
57 name: str
58 indices: torch.Tensor
59 values: torch.Tensor
62 - full_shape: tuple[int, ...] | None = None
63 - """Full shape of the patched parameter. Required when the patch is sent
64 - via `SparseNCCLTrainerWeightTransferEngine` (it ships in the per-round
65 - update info); the worker-side apply path does not read it."""
60 + full_shape: tuple[int, ...]
61 + """Full checkpoint shape."""
62
63
64 @dataclass
@@ -119,10 +115,9 @@ class SparseNCCLWeightTransferEngine(
115 """
116 Sparse weight transfer engine using NCCL.
117
122 - Receives flat-index (indices, values) patches broadcast from the trainer
123 - (rank 0) and applies them in place to existing model parameters. Weights are
124 - applied directly without layerwise reload, so `start_weight_update` and
125 - `finish_weight_update` are no-ops.
118 + Receives checkpoint-coordinate patches broadcast from the trainer and applies
119 + them through the model's native weight loader. Sparse updates modify initialized
120 + model tensors in place, so the layerwise reload lifecycle is not used.
121 """
122
123 init_info_cls = NCCLWeightTransferInitInfo
@@ -147,10 +142,7 @@ class SparseNCCLWeightTransferEngine(
142
143 def start_weight_update(self) -> None:
144 """No-op: sparse patches are applied in place, no layerwise reload."""
150 - if self.parallel_config.world_size != 1:
151 - raise NotImplementedError(
152 - "Sparse weight updates currently require TP=1 and PP=1"
153 - )
145 + pass
146
147 def finish_weight_update(self) -> None:
148 """No-op: sparse patches are applied in place, no layerwise reload."""
@@ -164,64 +156,34 @@ class SparseNCCLWeightTransferEngine(
156 "Call init_transfer_engine() first."
157 )
158
167 - # Use the worker's assigned device rather than the ambient current
168 - # device: the receive path is no longer wrapped in
169 - # `with torch.device(self.device)` by the caller, so the current device
170 - # is not guaranteed to match self.device. The recv buffers must live on
171 - # the same device as the NCCL communicator (created on self.device).
172 - device = self.device
173 - for name, dtype_name, num_updates in zip(
174 - update_info.names,
175 - update_info.dtype_names,
176 - update_info.num_updates_list,
177 - ):
178 - dtype = getattr(torch, dtype_name)
179 - indices = torch.empty(num_updates, dtype=torch.int32, device=device)
180 - values = torch.empty(num_updates, dtype=dtype, device=device)
181 - self.model_update_group.broadcast(
182 - indices, src=0, stream=torch.cuda.current_stream()
183 - )
184 - self.model_update_group.broadcast(
185 - values, src=0, stream=torch.cuda.current_stream()
186 - )
187 - self._apply_patch(
188 - SparseWeightPatch(name=name, indices=indices, values=values)
189 - )
190 - del indices
191 - del values
192 -
193 - def _apply_patch(self, patch: SparseWeightPatch) -> None:
194 - """Apply a single sparse flat-index patch to an existing model param."""
195 - param = self.model.get_parameter(patch.name)
196 - if not param.data.is_contiguous():
197 - raise NotImplementedError(
198 - "Sparse weight updates currently require contiguous params: "
199 - f"{patch.name}"
200 - )
201 - if patch.indices.dtype != torch.int32:
202 - raise ValueError(
203 - f"Sparse weight updates currently require int32 indices: {patch.name}"
204 - )
205 - if patch.indices.ndim != 1 or patch.values.ndim != 1:
206 - raise ValueError(
207 - f"Sparse weight patches must be 1D flattened updates: {patch.name}"
208 - )
209 - if patch.indices.numel() != patch.values.numel():
210 - raise ValueError(
211 - f"`indices` and `values` must have matching lengths for {patch.name}"
212 - )
213 - if patch.values.dtype != param.dtype:
214 - raise ValueError(
215 - f"Sparse values dtype {patch.values.dtype} does not match "
216 - f"parameter dtype {param.dtype} for {patch.name}"
217 - )
218 -
219 - flat_param = param.data.view(-1)
220 - flat_param.index_copy_(
221 - 0,
222 - patch.indices.to(device=flat_param.device, dtype=torch.long),
223 - patch.values.to(device=flat_param.device),
224 - )
159 + # Receive on the communicator's device rather than relying on the
160 + # ambient current device in the RPC thread.
161 + device = self.model_update_group.device
162 + with torch.accelerator.device_index(device.index):
163 + stream = torch.cuda.current_stream(device=device)
164 + patches = []
165 + for name, dtype_name, shape, num_updates in zip(
166 + update_info.names,
167 + update_info.dtype_names,
168 + update_info.shapes,
169 + update_info.num_updates_list,
170 + strict=True,
171 + ):
172 + dtype = getattr(torch, dtype_name)
173 + indices = torch.empty(num_updates, dtype=torch.int32, device=device)
174 + values = torch.empty(num_updates, dtype=dtype, device=device)
175 + self.model_update_group.broadcast(indices, src=0, stream=stream)
176 + self.model_update_group.broadcast(values, src=0, stream=stream)
177 + patches.append(
178 + CheckpointWeightPatch(
179 + name=name,
180 + shape=tuple(shape),
181 + dtype=dtype,
182 + values=values,
183 + indices=indices,
184 + )
185 + )
186 + load_checkpoint_weight_patches(self.model, patches)
187
188 def shutdown(self) -> None:
189 if self.model_update_group is not None:
@@ -235,15 +197,16 @@ class SparseNCCLTrainerWeightTransferEngine(
197
198 Broadcasts flat-index (indices, values) patches from NCCL rank 0 while the
199 inference-side `update_weights` runs concurrently on a side thread (the
238 - worker's recvs rendezvous inside the same NCCL broadcasts), then finishes
239 - the update. Unlike the full-resync backends, sparse patches differ every
240 - round (a fresh set of deltas from each optimizer step), so they are not a
241 - stable `WeightSource`: the engine takes no `source`, and each round's
242 - patches are passed straight to `send_weights(patches)`. A round with no
243 - patches is a no-op.
244 -
245 - The sparse backend assumes a single-rank trainer (matching its TP=1 / PP=1
246 - MVP scope), so non-sender ranks skip `send_weights` entirely.
200 + worker's recvs rendezvous inside the same NCCL broadcasts). `send_weights`
201 + owns a complete one-shot lifecycle. RL infrastructure that owns the generic
202 + client lifecycle can call `send_weight_chunk` between one start and finish.
203 +
204 + Sparse patches differ every round, so they are not a stable `WeightSource`:
205 + the engine takes no `source`, and patches are passed directly to the send
206 + methods. An empty patch list is a no-op.
207 +
208 + Only the designated trainer sender joins the transfer group; other trainer
209 + ranks skip sparse sends.
210 """
211
212 init_info_cls = SparseNCCLTrainerInitInfo
@@ -303,70 +266,88 @@ class SparseNCCLTrainerWeightTransferEngine(
266 return engine
267
268 def send_weights(self, patches: Iterable[SparseWeightPatch] | None = None) -> None:
306 - """Broadcast this round's sparse patches. `patches` is the per-round
307 - payload (sparse deltas differ every round), so it is passed here rather
308 - than fixed at init. Every patch must set `full_shape`."""
309 - if not self.is_sender:
269 + """Broadcast one sparse update through a one-shot lifecycle."""
270 + patches = self._prepare_patches(patches)
271 + if not patches:
272 return
273
312 - patches = list(patches) if patches is not None else []
274 + self.client.start_weight_update()
275 + self._broadcast_chunk(patches)
276 + self.client.finish_weight_update()
277 +
278 + def send_weight_chunk(
279 + self, patches: Iterable[SparseWeightPatch] | None = None
280 + ) -> None:
281 + """Broadcast one chunk inside a caller-owned weight update lifecycle."""
282 + patches = self._prepare_patches(patches)
283 if not patches:
284 return
285 + self._broadcast_chunk(patches)
286
316 - shapes = []
317 - for patch in patches:
318 - if patch.full_shape is None:
319 - raise ValueError(
320 - "SparseWeightPatch.full_shape must be set to send via the "
321 - f"trainer engine: {patch.name}"
322 - )
323 - self._validate_patch(patch)
324 - shapes.append(list(patch.full_shape))
287 + def _prepare_patches(
288 + self, patches: Iterable[SparseWeightPatch] | None
289 + ) -> list[SparseWeightPatch]:
290 + if not self.is_sender:
291 + return []
292
326 - update_info = SparseNCCLWeightTransferUpdateInfo(
327 - names=[patch.name for patch in patches],
328 - dtype_names=[str(patch.values.dtype).split(".")[-1] for patch in patches],
329 - shapes=shapes,
330 - num_updates_list=[patch.indices.numel() for patch in patches],
293 + prepared = list(patches) if patches is not None else []
294 + if not prepared:
295 + return []
296 + if self.model_update_group is None:
297 + raise RuntimeError("trainer_init() must be called before sending weights.")
298 + for patch in prepared:
299 + self._validate_patch(patch)
300 + return prepared
301 +
302 + def _broadcast_chunk(self, patches: list[SparseWeightPatch]) -> None:
303 + assert self.model_update_group is not None
304 + device = self.model_update_group.device
305 + device_context = (
306 + torch.accelerator.device_index(device.index)
307 + if device.type != "cpu"
308 + else nullcontext()
309 )
310 + with device_context:
311 + update_info = SparseNCCLWeightTransferUpdateInfo(
312 + names=[patch.name for patch in patches],
313 + dtype_names=[
314 + str(patch.values.dtype).split(".")[-1] for patch in patches
315 + ],
316 + shapes=[list(patch.full_shape) for patch in patches],
317 + num_updates_list=[patch.indices.numel() for patch in patches],
318 + )
319
333 - assert self.model_update_group is not None, (
334 - "trainer_init() must be called before send_weights()."
335 - )
336 - self.client.start_weight_update()
337 - # update_weights (workers receive) must run concurrently with the
338 - # trainer-side broadcasts — both rendezvous inside the same NCCL calls.
339 - exe = ThreadPoolExecutor(max_workers=1)
340 - try:
341 - future = exe.submit(self.client.update_weights, asdict(update_info))
342 - # Cheap best-effort: if update_weights already failed (e.g. a bad
343 - # request rejected before any NCCL call), surface it now instead of
344 - # hanging in broadcast waiting for a peer that will never arrive.
345 - if future.done():
346 - future.result()
347 - stream = torch.cuda.current_stream()
348 - for patch in patches:
349 - self.model_update_group.broadcast(patch.indices, src=0, stream=stream)
350 - self.model_update_group.broadcast(patch.values, src=0, stream=stream)
351 - future.result() # surface inference-side errors
352 - finally:
353 - # Never wait for the RPC thread here: if a broadcast raised, the
354 - # worker is still blocked in the matching NCCL call and will never
355 - # return, so joining would turn the error into a permanent hang.
356 - # See NCCLTrainerWeightTransferEngine.send_weights.
357 - exe.shutdown(wait=False)
358 - self.client.finish_weight_update()
359 - self._post_send_sync()
320 + # update_weights (workers receive) must run concurrently with the
321 + # trainer-side broadcasts — both rendezvous inside the same NCCL calls.
322 + executor = ThreadPoolExecutor(max_workers=1)
323 + try:
324 + future = executor.submit(
325 + self.client.update_weights, asdict(update_info)
326 + )
327 + # Surface an RPC that failed before the worker entered NCCL rather
328 + # than broadcasting to a peer that will never receive.
329 + if future.done():
330 + future.result()
331 + stream = torch.cuda.current_stream()
332 + for patch in patches:
333 + self.model_update_group.broadcast(
334 + patch.indices, src=0, stream=stream
335 + )
336 + self.model_update_group.broadcast(
337 + patch.values, src=0, stream=stream
338 + )
339 + future.result() # surface inference-side errors
340 + finally:
341 + # A failed broadcast can leave the RPC blocked in the matching
342 + # receive. Waiting for that thread would hide the original error.
343 + executor.shutdown(wait=False, cancel_futures=True)
344 + self._post_send_sync()
345
346 @staticmethod
347 def _validate_patch(patch: SparseWeightPatch) -> None:
363 - """Reject a malformed patch before any NCCL call.
364 -
365 - The worker checks the same invariants in `_apply_patch`, but by then the
366 - broadcasts are already under way: a mismatch surfaces as a size mismatch
367 - mid-transfer, which wedges both sides instead of raising. Checking here
368 - keeps the failure on the trainer, before `start_weight_update`.
369 - """
348 + """Reject a malformed patch before starting the NCCL transfer."""
349 + if patch.full_shape is None:
350 + raise ValueError(f"Sparse patch requires full_shape: {patch.name}")
351 if patch.indices.dtype != torch.int32:
352 raise ValueError(
353 f"Sparse weight updates require int32 indices: {patch.name}"
@@ -382,10 +363,10 @@ class SparseNCCLTrainerWeightTransferEngine(
363
364 def _post_send_sync(self) -> None:
365 """Wait for the broadcasts to land before returning, so a caller may
385 - rebuild or free the patch tensors as soon as `send_weights` returns
386 - rather than relying on same-stream ordering. See
366 + rebuild or free the patch tensors as soon as a send method returns rather
367 + than relying on same-stream ordering. See
368 `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no
388 - cross-rank barrier (and sparse is single-rank on the trainer anyway)."""
369 + cross-rank barrier."""
370 if torch.cuda.is_available():
371 torch.cuda.current_stream().synchronize()
372