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 3, 13:00 UTC
[Bugfix] Fix layerwise reload alias-buffer corruption (#42481)
Signed-off-by: rasdani <73563550+rasdani@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io>
rasdani committed
May 15, 2026 at 15:20 UTC
1ccdf87507407cb02460ec2e7a3e1a4cac9b0a4a
3 files changed
+154
tests/model_executor/model_loader/test_reload.py
+114
@@ -6,8 +6,15 @@ from weakref import WeakKeyDictionary, ref
6
7
import pytest
8
import torch
9
+from torch.nn.parameter import UninitializedParameter
10
11
+import vllm.model_executor.model_loader.reload.meta as reload_meta
12
from vllm.model_executor.layers.linear import QKVParallelLinear
13
+from vllm.model_executor.model_loader.reload.layerwise import (
14
+ finalize_layerwise_reload,
15
+ initialize_layerwise_reload,
16
+ record_metadata_for_reloading,
17
+)
18
from vllm.model_executor.model_loader.reload.meta import (
19
capture_layer_to_meta,
20
get_numel_loaded,
@@ -21,6 +28,38 @@ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
28
from vllm.platforms import current_platform
29
30
31
+class _AliasedBufferLayer(torch.nn.Module):
32
+ def __init__(self):
33
+ super().__init__()
34
+ weight = torch.arange(6, dtype=torch.float32).reshape(2, 3)
35
+ self.weight = torch.nn.Parameter(weight)
36
+ self.register_buffer(
37
+ "weight_view", self.weight.detach().view(-1), persistent=False
38
+ )
39
+
40
+
41
+class _ParentAliasedChildBufferLayer(torch.nn.Module):
42
+ def __init__(self):
43
+ super().__init__()
44
+ self.scale = torch.nn.Parameter(torch.ones(1))
45
+ self.conv1d = torch.nn.Linear(3, 2, bias=False)
46
+ self.conv1d.weight.data.copy_(
47
+ torch.arange(6, dtype=torch.float32).reshape(2, 3)
48
+ )
49
+ self.register_buffer(
50
+ "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False
51
+ )
52
+
53
+
54
+class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer):
55
+ def __init__(self):
56
+ super().__init__()
57
+ self.child = torch.nn.Module()
58
+ self.child.register_parameter(
59
+ "lazy_weight", UninitializedParameter(requires_grad=False)
60
+ )
61
+
62
+
63
def test_move_metatensors():
64
tensor = torch.empty((1, 2, 3))
65
meta_tensor = to_meta_tensor(tensor)
@@ -123,6 +162,81 @@ def test_get_numel_loaded():
162
assert ret == "value"
163
164
165
+def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch):
166
+ layer = _AliasedBufferLayer()
167
+ model = torch.nn.Sequential(layer)
168
+ loaded_weight = torch.full_like(layer.weight, 7.0)
169
+
170
+ def materialize_with_sentinel(meta_tensor):
171
+ tensor = torch.empty_strided(
172
+ size=tuple(meta_tensor.size()),
173
+ stride=tuple(meta_tensor.stride()),
174
+ dtype=meta_tensor.dtype,
175
+ requires_grad=False,
176
+ )
177
+ tensor.fill_(-123.0)
178
+ tensor.__class__ = meta_tensor.__class__
179
+ tensor.__dict__ = meta_tensor.__dict__.copy()
180
+ return tensor
181
+
182
+ monkeypatch.setattr(
183
+ reload_meta, "materialize_meta_tensor", materialize_with_sentinel
184
+ )
185
+
186
+ record_metadata_for_reloading(model)
187
+ initialize_layerwise_reload(model)
188
+ layer.weight.weight_loader(layer.weight, loaded_weight)
189
+ finalize_layerwise_reload(model, model_config=None)
190
+
191
+ assert torch.equal(layer.weight, loaded_weight)
192
+ assert layer.weight_view.untyped_storage().data_ptr() == (
193
+ layer.weight.untyped_storage().data_ptr()
194
+ )
195
+
196
+
197
+def test_capture_layer_to_meta_skips_uninitialized_parameter_storage_ptrs():
198
+ layer = _AliasedBufferWithUninitializedChildLayer()
199
+
200
+ _, buffers = capture_layer_to_meta(layer)
201
+
202
+ assert "weight_view" not in buffers
203
+
204
+
205
+def test_layerwise_reload_skips_child_parameter_alias_buffers(monkeypatch):
206
+ layer = _ParentAliasedChildBufferLayer()
207
+ model = torch.nn.Sequential(layer)
208
+ loaded_conv = torch.full_like(layer.conv1d.weight, 7.0)
209
+ loaded_scale = torch.full_like(layer.scale, 3.0)
210
+
211
+ def materialize_with_sentinel(meta_tensor):
212
+ tensor = torch.empty_strided(
213
+ size=tuple(meta_tensor.size()),
214
+ stride=tuple(meta_tensor.stride()),
215
+ dtype=meta_tensor.dtype,
216
+ requires_grad=False,
217
+ )
218
+ tensor.fill_(-123.0)
219
+ tensor.__class__ = meta_tensor.__class__
220
+ tensor.__dict__ = meta_tensor.__dict__.copy()
221
+ return tensor
222
+
223
+ monkeypatch.setattr(
224
+ reload_meta, "materialize_meta_tensor", materialize_with_sentinel
225
+ )
226
+
227
+ record_metadata_for_reloading(model)
228
+ initialize_layerwise_reload(model)
229
+ layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv)
230
+ layer.scale.weight_loader(layer.scale, loaded_scale)
231
+ finalize_layerwise_reload(model, model_config=None)
232
+
233
+ assert torch.equal(layer.conv1d.weight, loaded_conv)
234
+ assert torch.equal(layer.conv_weights, loaded_conv.view(-1))
235
+ assert layer.conv_weights.untyped_storage().data_ptr() == (
236
+ layer.conv1d.weight.untyped_storage().data_ptr()
237
+ )
238
+
239
+
240
@pytest.mark.parametrize(
241
"tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
242
)
vllm/model_executor/model_loader/reload/layerwise.py
+2
@@ -380,6 +380,8 @@ def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadin
380
for name, param in parameters.items():
381
param.data.copy_(getattr(layer, name))
382
for name, buffer in buffers.items():
383
+ if name not in layer._buffers:
384
+ continue
385
buffer.data.copy_(getattr(layer, name))
386
387
_place_kernel_tensors(layer, info)
vllm/model_executor/model_loader/reload/meta.py
+38
@@ -4,6 +4,7 @@ import inspect
4
from collections.abc import Callable
5
6
import torch
7
+from torch.nn.parameter import UninitializedParameter
8
from torch.utils._python_dispatch import TorchDispatchMode
9
10
from .sanitize import restore_layer_refs, sanitize_layer_refs
@@ -55,11 +56,45 @@ def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor:
56
return tensor
57
58
59
+def _is_non_persistent_parameter_alias_buffer(
60
+ layer: torch.nn.Module,
61
+ name: str,
62
+ buffer: torch.Tensor,
63
+ parameter_storage_ptrs: set[int],
64
+) -> bool:
65
+ if name not in layer._non_persistent_buffers_set:
66
+ return False
67
+
68
+ buffer_storage_ptr = _tensor_storage_ptr(buffer)
69
+ return (
70
+ buffer_storage_ptr is not None and buffer_storage_ptr in parameter_storage_ptrs
71
+ )
72
+
73
+
74
+def _tensor_storage_ptr(tensor: torch.Tensor) -> int | None:
75
+ if isinstance(tensor, UninitializedParameter):
76
+ return None
77
+
78
+ try:
79
+ return tensor.untyped_storage().data_ptr()
80
+ except (RuntimeError, ValueError):
81
+ return None
82
+
83
+
84
+def _parameter_storage_ptrs(layer: torch.nn.Module) -> set[int]:
85
+ return {
86
+ storage_ptr
87
+ for param in layer.parameters(recurse=True)
88
+ if (storage_ptr := _tensor_storage_ptr(param)) is not None
89
+ }
90
+
91
+
92
def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors:
93
if layer.__class__.__name__ in SKIP_MODULES:
94
return ({}, {})
95
96
params, buffers = get_layer_params_buffers(layer)
97
+ parameter_storage_ptrs = _parameter_storage_ptrs(layer)
98
return (
99
{
100
name: sanitize_layer_refs(to_meta_tensor(param), layer)
@@ -70,6 +105,9 @@ def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors:
105
name: sanitize_layer_refs(to_meta_tensor(buffer), layer)
106
for name, buffer in buffers.items()
107
if name not in SKIP_TENSORS
108
+ and not _is_non_persistent_parameter_alias_buffer(
109
+ layer, name, buffer, parameter_storage_ptrs
110
+ )
111
},
112
)
113