[Core][RL] Support sparse checkpoint updates through native weight loaders (#50723)
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Shuolei Wang committed
Aug 23, 2026 at 01:39 UTC
236f78cc5c39362f80c0c7df1ffe810466525f33
2 files changed
+383
tests/model_executor/model_loader/test_checkpoint_weight_patch.py
new
+110
@@ -0,0 +1,110 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.model_loader.checkpoint_weight_patch import (
8
+ CheckpointWeightPatch,
9
+ load_checkpoint_weight_patches,
10
+)
11
+
12
+pytestmark = pytest.mark.cpu_test
13
+
14
+
15
+class _PackedTPModel(torch.nn.Module):
16
+ """Model that packs the TP rank 1 shards of Q and K into one parameter."""
17
+
18
+ def __init__(self):
19
+ super().__init__()
20
+ self.packed_weight = torch.nn.Parameter(torch.full((4,), -1.0))
21
+ self.load_calls: list[list[str]] = []
22
+
23
+ def load_weights(self, weights):
24
+ weights = list(weights)
25
+ self.load_calls.append([name for name, _ in weights])
26
+ loaded_names = set()
27
+ for name, loaded_weight in weights:
28
+ if name == "q_proj.weight":
29
+ destination = self.packed_weight.data[:2]
30
+ elif name == "k_proj.weight":
31
+ destination = self.packed_weight.data[2:]
32
+ else:
33
+ raise AssertionError(name)
34
+ # Slice each full checkpoint tensor and load its second half for TP rank 1.
35
+ destination.copy_(loaded_weight.narrow(0, 2, 2))
36
+ loaded_names.add(name)
37
+ return loaded_names
38
+
39
+
40
+def _make_patch(
41
+ name: str,
42
+ *,
43
+ values: list[float],
44
+ indices: list[int] | None = None,
45
+) -> CheckpointWeightPatch:
46
+ return CheckpointWeightPatch(
47
+ name=name,
48
+ shape=(4,),
49
+ dtype=torch.float32,
50
+ values=torch.tensor(values),
51
+ indices=None if indices is None else torch.tensor(indices, dtype=torch.int32),
52
+ )
53
+
54
+
55
+def test_dense_and_sparse_patches_follow_packed_tp_loader():
56
+ model = _PackedTPModel()
57
+
58
+ # Seed the TP-local packed weights through the dense path.
59
+ dense_loaded = load_checkpoint_weight_patches(
60
+ model,
61
+ [
62
+ _make_patch("q_proj.weight", values=[0.0, 1.0, 2.0, 3.0]),
63
+ _make_patch("k_proj.weight", values=[10.0, 11.0, 12.0, 13.0]),
64
+ ],
65
+ )
66
+ assert dense_loaded == {"q_proj.weight", "k_proj.weight"}
67
+ assert torch.equal(model.packed_weight, torch.tensor([2.0, 3.0, 12.0, 13.0]))
68
+
69
+ # Sparse indices address the full checkpoint tensor; index 0 is outside this shard.
70
+ model.load_calls.clear()
71
+ original_copy = torch.Tensor.copy_
72
+ sparse_loaded = load_checkpoint_weight_patches(
73
+ model,
74
+ [
75
+ _make_patch(
76
+ "q_proj.weight",
77
+ indices=[0, 3],
78
+ values=[100.0, 30.0],
79
+ ),
80
+ _make_patch("k_proj.weight", indices=[2], values=[20.0]),
81
+ _make_patch("q_proj.weight", indices=[2], values=[22.0]),
82
+ ],
83
+ )
84
+
85
+ assert sparse_loaded == {"q_proj.weight", "k_proj.weight"}
86
+ assert model.load_calls == [
87
+ ["q_proj.weight", "k_proj.weight"],
88
+ ["q_proj.weight"],
89
+ ]
90
+ assert torch.equal(model.packed_weight, torch.tensor([22.0, 30.0, 20.0, 13.0]))
91
+ assert torch.Tensor.copy_ is original_copy
92
+
93
+ # The empty result tells the helper that the loader intentionally made no write.
94
+ class NoLocalWeightModel(torch.nn.Module):
95
+ """Model that deliberately owns none of the supplied weights."""
96
+
97
+ def __init__(self):
98
+ super().__init__()
99
+ self.weight = torch.nn.Parameter(torch.tensor([1.0]))
100
+
101
+ def load_weights(self, _weights):
102
+ return set()
103
+
104
+ no_local_weight_model = NoLocalWeightModel()
105
+ no_local_weight_loaded = load_checkpoint_weight_patches(
106
+ no_local_weight_model,
107
+ [_make_patch("q_proj.weight", indices=[3], values=[30.0])],
108
+ )
109
+ assert no_local_weight_loaded == set()
110
+ assert torch.equal(no_local_weight_model.weight, torch.tensor([1.0]))
vllm/model_executor/model_loader/checkpoint_weight_patch.py
new
+273
@@ -0,0 +1,273 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Apply dense or sparse weight updates in checkpoint coordinates.
4
+
5
+The model's ``load_weights`` method still handles checkpoint-name mapping,
6
+TP slicing, and packed parameters. Sparse patches use NaN to mark unchanged
7
+checkpoint elements. The supported sparse-loader contract permits only final
8
+same-shaped floating-point ``Tensor.copy_`` writes. Composed, multi-stage, and
9
+custom-write loaders are unsupported.
10
+"""
11
+
12
+from __future__ import annotations
13
+
14
+import math
15
+from collections.abc import Iterable
16
+from typing import NamedTuple
17
+
18
+import torch
19
+
20
+_DEFAULT_PATCH_CHUNK_BYTES = 512 << 20
21
+
22
+__all__ = [
23
+ "CheckpointWeightPatch",
24
+ "load_checkpoint_weight_patches",
25
+]
26
+
27
+
28
+class CheckpointWeightPatch(NamedTuple):
29
+ """Describe one dense or sparse checkpoint-coordinate update.
30
+
31
+ Attributes:
32
+ name: Checkpoint weight name passed to the model loader.
33
+ shape: Full checkpoint tensor shape.
34
+ dtype: Checkpoint tensor dtype.
35
+ values: The flattened full tensor for a dense patch, or the values at
36
+ ``indices`` for a sparse patch.
37
+ indices: Flat indices into the full checkpoint tensor described by
38
+ ``shape``. ``None`` makes ``values`` a dense replacement.
39
+ """
40
+
41
+ name: str
42
+ shape: tuple[int, ...]
43
+ dtype: torch.dtype
44
+ values: torch.Tensor
45
+ indices: torch.Tensor | None = None
46
+
47
+
48
+def _load_nan_masked_weights(
49
+ model: torch.nn.Module,
50
+ weights: list[tuple[str, torch.Tensor]],
51
+) -> Iterable[str] | None:
52
+ """Load NaN-masked checkpoint tensors while preserving runtime values at NaNs."""
53
+ # Save the original copy_ so it can be restored after loading.
54
+ original_copy = torch.Tensor.copy_
55
+
56
+ # Enforce the sparse-loader contract and preserve destination values where
57
+ # the corresponding source is NaN.
58
+ def copy_non_nan_(
59
+ destination: torch.Tensor,
60
+ source: torch.Tensor,
61
+ non_blocking: bool = False,
62
+ ) -> torch.Tensor:
63
+ if not (
64
+ destination.is_floating_point()
65
+ and source.is_floating_point()
66
+ and destination.shape == source.shape
67
+ ):
68
+ raise NotImplementedError(
69
+ "Sparse checkpoint patches require a same-shaped "
70
+ "floating-point final copy"
71
+ )
72
+ source = source.to(dtype=destination.dtype, device=destination.device)
73
+ return torch.where(
74
+ torch.isnan(source),
75
+ destination,
76
+ source,
77
+ out=destination,
78
+ )
79
+
80
+ torch.Tensor.copy_ = copy_non_nan_
81
+ try:
82
+ return model.load_weights(weights)
83
+ finally:
84
+ torch.Tensor.copy_ = original_copy
85
+
86
+
87
+def _validate_patch_structure(patch: CheckpointWeightPatch) -> int:
88
+ if not patch.name:
89
+ raise ValueError("Checkpoint weight name must be non-empty")
90
+ if not isinstance(patch.dtype, torch.dtype):
91
+ raise TypeError(f"{patch.name}: dtype must be a torch.dtype")
92
+ if any(dim < 0 for dim in patch.shape):
93
+ raise ValueError(f"{patch.name}: shape dimensions must be non-negative")
94
+ if patch.values.ndim != 1:
95
+ raise ValueError(f"{patch.name}: patch values must be a 1D tensor")
96
+
97
+ numel = math.prod(patch.shape)
98
+ if patch.indices is None:
99
+ if patch.values.numel() != numel:
100
+ raise ValueError(
101
+ f"{patch.name}: dense patch has {patch.values.numel()} values "
102
+ f"for a {numel}-element tensor"
103
+ )
104
+ return numel
105
+
106
+ if not patch.dtype.is_floating_point:
107
+ raise TypeError(
108
+ f"{patch.name}: sparse checkpoint patches require a floating dtype "
109
+ "because NaN is the unchanged-value sentinel"
110
+ )
111
+ if patch.indices.ndim != 1:
112
+ raise ValueError(f"{patch.name}: sparse patch indices must be a 1D tensor")
113
+ if patch.indices.dtype not in (torch.int32, torch.int64):
114
+ raise TypeError(f"{patch.name}: sparse patch indices must be int32 or int64")
115
+ if patch.indices.numel() != patch.values.numel():
116
+ raise ValueError(
117
+ f"{patch.name}: sparse indices and values must have matching lengths"
118
+ )
119
+
120
+ return numel
121
+
122
+
123
+def _load_weight_chunk(
124
+ model: torch.nn.Module,
125
+ weights: list[tuple[str, torch.Tensor]],
126
+ *,
127
+ sparse: bool,
128
+) -> set[str]:
129
+ if not weights:
130
+ return set()
131
+ if sparse:
132
+ loaded_names = _load_nan_masked_weights(model, weights)
133
+ else:
134
+ loaded_names = model.load_weights(weights)
135
+ return set() if loaded_names is None else set(loaded_names)
136
+
137
+
138
+@torch.no_grad()
139
+def load_checkpoint_weight_patches(
140
+ model: torch.nn.Module,
141
+ patches: Iterable[CheckpointWeightPatch],
142
+ *,
143
+ max_chunk_bytes: int = _DEFAULT_PATCH_CHUNK_BYTES,
144
+ validate_unique_indices: bool = True,
145
+) -> set[str]:
146
+ """Load ordered patches through the model's checkpoint loader.
147
+
148
+ A call may contain either dense or sparse patches, but not both. A repeated
149
+ name starts a new loader call, so patches for that weight are applied in input
150
+ order. A later patch may update positions changed by an earlier patch. Each
151
+ sparse patch creates a full checkpoint-shaped tensor whose NaNs mark unchanged
152
+ elements. Each locally applied sparse patch must use one final same-shaped
153
+ floating-point ``Tensor.copy_``. Intermediate or unrelated ``copy_`` calls,
154
+ composed loaders, and custom write paths are unsupported. ``max_chunk_bytes``
155
+ is a batching target; one tensor may exceed it.
156
+
157
+ Online checkpoint-format dense updates require the caller to manage vLLM's
158
+ layerwise reload lifecycle. Sparse online updates modify initialized model
159
+ tensors in place and must not use that lifecycle. Patch shapes, dtypes, value
160
+ lengths, and sparse indices are validated before loading begins.
161
+
162
+ This function does not roll back a partial update. If a loader fails after a
163
+ write, some destinations may already be changed. Do not serve from the affected
164
+ worker until a known baseline has been restored or the worker has been
165
+ restarted.
166
+
167
+ Args:
168
+ model: Model whose native ``load_weights`` method applies the patches.
169
+ patches: Dense replacements or sparse updates in checkpoint coordinates.
170
+ max_chunk_bytes: Target checkpoint tensor bytes per
171
+ ``model.load_weights`` call. A single tensor may exceed this target.
172
+ validate_unique_indices: Whether to reject duplicate sparse indices within
173
+ each patch. Repeated patches may update the same positions. Disable
174
+ only for a trusted producer that already guarantees unique positions
175
+ within each patch.
176
+
177
+ Returns:
178
+ Union of the weight names reported by all ``model.load_weights`` calls.
179
+ """
180
+
181
+ if max_chunk_bytes <= 0:
182
+ raise ValueError("max_chunk_bytes must be positive")
183
+
184
+ patch_list = list(patches)
185
+ patch_numels = [_validate_patch_structure(patch) for patch in patch_list]
186
+ # All patches in one call must use the same dense or sparse representation.
187
+ sparse_flags = {patch.indices is not None for patch in patch_list}
188
+ if len(sparse_flags) > 1:
189
+ raise ValueError("Dense and sparse checkpoint patches cannot be mixed")
190
+ sparse = sparse_flags == {True}
191
+
192
+ # Keep validation results on their devices and read one combined result per
193
+ # device, instead of synchronizing the CPU once for every patch.
194
+ checks_by_device: dict[torch.device, list[torch.Tensor]] = {}
195
+ if sparse:
196
+ for patch, numel in zip(patch_list, patch_numels, strict=True):
197
+ assert patch.indices is not None
198
+
199
+ # NaN is reserved to mean "leave this checkpoint position unchanged,"
200
+ # so it cannot also be supplied as a new sparse value.
201
+ patch_checks = [torch.isnan(patch.values).any()]
202
+ indices = patch.indices.to(device=patch.values.device)
203
+ if indices.numel():
204
+ # Sparse indices address the flattened full checkpoint tensor.
205
+ patch_checks.append(
206
+ torch.logical_or(indices < 0, indices >= numel).any()
207
+ )
208
+ if validate_unique_indices:
209
+ # Equal neighboring indices after sorting indicate duplicates.
210
+ sorted_indices = torch.sort(indices).values
211
+ patch_checks.append(
212
+ (sorted_indices[1:] == sorted_indices[:-1]).any()
213
+ )
214
+
215
+ checks_by_device.setdefault(patch.values.device, []).extend(patch_checks)
216
+
217
+ for device_checks in checks_by_device.values():
218
+ if torch.stack(device_checks).any().item():
219
+ invalid_reasons = (
220
+ "NaN values, out-of-range indices, or duplicate indices"
221
+ if validate_unique_indices
222
+ else "NaN values or out-of-range indices"
223
+ )
224
+ raise ValueError(f"Sparse checkpoint patches contain {invalid_reasons}")
225
+
226
+ loaded_names: set[str] = set()
227
+ # Collect checkpoint tensors for one model.load_weights() call.
228
+ weight_chunk: list[tuple[str, torch.Tensor]] = []
229
+ chunk_names: set[str] = set()
230
+ chunk_bytes = 0
231
+
232
+ for patch, numel in zip(patch_list, patch_numels, strict=True):
233
+ # Count the full checkpoint tensor, including sparse NaN staging.
234
+ checkpoint_tensor_bytes = (
235
+ numel * torch.empty((), dtype=patch.dtype).element_size()
236
+ )
237
+ # Flush before exceeding the batching target. Repeated names also start a
238
+ # new loader call so their patches are applied in input order.
239
+ if weight_chunk and (
240
+ chunk_bytes + checkpoint_tensor_bytes > max_chunk_bytes
241
+ or patch.name in chunk_names
242
+ ):
243
+ loaded_names.update(_load_weight_chunk(model, weight_chunk, sparse=sparse))
244
+ weight_chunk = []
245
+ chunk_names = set()
246
+ chunk_bytes = 0
247
+
248
+ if patch.indices is None:
249
+ checkpoint_tensor = patch.values.to(dtype=patch.dtype).view(patch.shape)
250
+ else:
251
+ # NaNs mark unchanged checkpoint positions; index_copy_ fills only
252
+ # the positions supplied by this sparse patch.
253
+ flat = torch.full(
254
+ (numel,),
255
+ float("nan"),
256
+ dtype=patch.dtype,
257
+ device=patch.values.device,
258
+ )
259
+ if patch.values.numel():
260
+ flat.index_copy_(
261
+ 0,
262
+ patch.indices.to(device=flat.device, dtype=torch.long),
263
+ patch.values.to(dtype=patch.dtype, device=flat.device),
264
+ )
265
+ checkpoint_tensor = flat.view(patch.shape)
266
+
267
+ weight_chunk.append((patch.name, checkpoint_tensor))
268
+ chunk_names.add(patch.name)
269
+ chunk_bytes += checkpoint_tensor_bytes
270
+
271
+ # Load the final partial chunk; the helper also accepts an empty chunk.
272
+ loaded_names.update(_load_weight_chunk(model, weight_chunk, sparse=sparse))
273
+ return loaded_names