[Bugfix] Preserve unloaded non-persistent buffers during layerwise reload (#44371)
Signed-off-by: Joan Velja <joan.velja22@gmail.com> Co-authored-by: Dakai An <77474977+andakai@users.noreply.github.com>
Joan committed
Jul 15, 2026 at 01:46 UTC
0bd6b85a1f32cc3b32dc56d24ae8546897518ef3
4 files changed
+143
-5
tests/model_executor/model_loader/test_reload.py
+124
@@ -71,6 +71,18 @@ class _ParentAliasedChildBufferLayer(torch.nn.Module):
71
)
72
73
74
+class _ChildAliasOnlyBufferLayer(torch.nn.Module):
75
+ def __init__(self):
76
+ super().__init__()
77
+ self.conv1d = torch.nn.Linear(3, 2, bias=False)
78
+ self.conv1d.weight.data.copy_(
79
+ torch.arange(6, dtype=torch.float32).reshape(2, 3)
80
+ )
81
+ self.register_buffer(
82
+ "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False
83
+ )
84
+
85
+
86
class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer):
87
def __init__(self):
88
super().__init__()
@@ -80,6 +92,13 @@ class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer):
92
)
93
94
95
+class _NonPersistentBufferLayer(torch.nn.Module):
96
+ def __init__(self):
97
+ super().__init__()
98
+ self.weight = torch.nn.Parameter(torch.ones(2, 2))
99
+ self.register_buffer("scale", torch.tensor(0.25), persistent=False)
100
+
101
+
102
def test_move_metatensors():
103
tensor = torch.empty((1, 2, 3))
104
meta_tensor = to_meta_tensor(tensor)
@@ -289,6 +308,8 @@ def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypat
308
assert layer.weight_view.untyped_storage().data_ptr() == (
309
layer.weight.untyped_storage().data_ptr()
310
)
311
+ assert "weight_view" in layer._non_persistent_buffers_set
312
+ assert "0.weight_view" not in model.state_dict()
313
314
315
def test_capture_layer_to_meta_skips_uninitialized_parameter_storage_ptrs():
@@ -332,6 +353,109 @@ def test_layerwise_reload_skips_child_parameter_alias_buffers(monkeypatch):
353
assert layer.conv_weights.untyped_storage().data_ptr() == (
354
layer.conv1d.weight.untyped_storage().data_ptr()
355
)
356
+ assert "conv_weights" in layer._non_persistent_buffers_set
357
+ assert "0.conv_weights" not in model.state_dict()
358
+
359
+
360
+def test_layerwise_reload_restores_alias_buffer_on_zero_size_layer(monkeypatch):
361
+ layer = _ChildAliasOnlyBufferLayer()
362
+ model = torch.nn.Sequential(layer)
363
+ loaded_conv = torch.full_like(layer.conv1d.weight, 7.0)
364
+
365
+ def materialize_with_sentinel(meta_tensor):
366
+ tensor = torch.empty_strided(
367
+ size=tuple(meta_tensor.size()),
368
+ stride=tuple(meta_tensor.stride()),
369
+ dtype=meta_tensor.dtype,
370
+ requires_grad=False,
371
+ )
372
+ tensor.fill_(-123.0)
373
+ tensor.__class__ = meta_tensor.__class__
374
+ tensor.__dict__ = meta_tensor.__dict__.copy()
375
+ return tensor
376
+
377
+ monkeypatch.setattr(
378
+ reload_meta, "materialize_meta_tensor", materialize_with_sentinel
379
+ )
380
+
381
+ record_metadata_for_reloading(model)
382
+ initialize_layerwise_reload(model)
383
+ layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv)
384
+ finalize_layerwise_reload(model, model_config=None)
385
+
386
+ assert torch.equal(layer.conv_weights, loaded_conv.view(-1))
387
+ assert layer.conv_weights.untyped_storage().data_ptr() == (
388
+ layer.conv1d.weight.untyped_storage().data_ptr()
389
+ )
390
+ assert "conv_weights" in layer._non_persistent_buffers_set
391
+ assert "0.conv_weights" not in model.state_dict()
392
+
393
+
394
+def test_layerwise_reload_preserves_unloaded_non_persistent_buffers(monkeypatch):
395
+ layer = _NonPersistentBufferLayer()
396
+ model = torch.nn.Sequential(layer)
397
+ loaded_weight = torch.full_like(layer.weight, 7.0)
398
+ original_scale = layer.scale.clone()
399
+
400
+ def materialize_with_sentinel(meta_tensor):
401
+ tensor = torch.empty_strided(
402
+ size=tuple(meta_tensor.size()),
403
+ stride=tuple(meta_tensor.stride()),
404
+ dtype=meta_tensor.dtype,
405
+ requires_grad=False,
406
+ )
407
+ tensor.fill_(-123.0)
408
+ tensor.__class__ = meta_tensor.__class__
409
+ tensor.__dict__ = meta_tensor.__dict__.copy()
410
+ return tensor
411
+
412
+ monkeypatch.setattr(
413
+ reload_meta, "materialize_meta_tensor", materialize_with_sentinel
414
+ )
415
+
416
+ record_metadata_for_reloading(model)
417
+ initialize_layerwise_reload(model)
418
+ layer.weight.weight_loader(layer.weight, loaded_weight)
419
+ finalize_layerwise_reload(model, model_config=None)
420
+
421
+ assert torch.equal(layer.weight, loaded_weight)
422
+ assert torch.equal(layer.scale, original_scale)
423
+ assert "scale" in layer._non_persistent_buffers_set
424
+ assert "0.scale" not in model.state_dict()
425
+
426
+
427
+def test_layerwise_reload_updates_loaded_non_persistent_buffers(monkeypatch):
428
+ layer = _NonPersistentBufferLayer()
429
+ model = torch.nn.Sequential(layer)
430
+ loaded_weight = torch.full_like(layer.weight, 7.0)
431
+ loaded_scale = torch.full_like(layer.scale, 0.5)
432
+
433
+ def materialize_with_sentinel(meta_tensor):
434
+ tensor = torch.empty_strided(
435
+ size=tuple(meta_tensor.size()),
436
+ stride=tuple(meta_tensor.stride()),
437
+ dtype=meta_tensor.dtype,
438
+ requires_grad=False,
439
+ )
440
+ tensor.fill_(-123.0)
441
+ tensor.__class__ = meta_tensor.__class__
442
+ tensor.__dict__ = meta_tensor.__dict__.copy()
443
+ return tensor
444
+
445
+ monkeypatch.setattr(
446
+ reload_meta, "materialize_meta_tensor", materialize_with_sentinel
447
+ )
448
+
449
+ record_metadata_for_reloading(model)
450
+ initialize_layerwise_reload(model)
451
+ layer.weight.weight_loader(layer.weight, loaded_weight)
452
+ layer.scale.weight_loader(layer.scale, loaded_scale)
453
+ finalize_layerwise_reload(model, model_config=None)
454
+
455
+ assert torch.equal(layer.weight, loaded_weight)
456
+ assert torch.equal(layer.scale, loaded_scale)
457
+ assert "scale" in layer._non_persistent_buffers_set
458
+ assert "0.scale" not in model.state_dict()
459
460
461
@pytest.mark.parametrize(
vllm/model_executor/model_loader/reload/layerwise.py
+13
-4
@@ -109,6 +109,8 @@ def initialize_layerwise_reload(model: torch.nn.Module):
109
110
# Save current tensors for later copying
111
info.kernel_tensors = get_layer_params_buffers(layer)
112
+ # snapshot now: restore_layer_on_meta drops alias buffers from the live set
113
+ info.kernel_non_persistent_buffers = set(layer._non_persistent_buffers_set)
114
115
# Restore layer parameters/buffers onto meta device
116
restore_layer_on_meta(layer, info)
@@ -259,10 +261,12 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
261
_layerwise_process(layer, info)
262
continue
263
262
- # reloading: place kernel tensors back as a fallback
263
- elif info.load_numel_total > 0: # type: ignore[operator]
264
+ # reloading: place kernel tensors back as a fallback. Always place, even
265
+ # when nothing is loadable (load_numel_total == 0), so parameter-alias
266
+ # buffers on such layers are restored rather than left deleted.
267
+ if info.load_numel_total > 0: # type: ignore[operator]
268
logger.warning("%s: Failed to load weights", layer.__class__.__name__)
265
- _place_kernel_tensors(layer, info)
269
+ _place_kernel_tensors(layer, info)
270
271
# Process non-attention layers which did not load all elements. This can happen
272
# if the created weight has extra padding elements which are not loaded
@@ -388,11 +392,15 @@ def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadin
392
kernel tensor references on the layer. Preserves cudagraph references."""
393
assert info.kernel_tensors is not None
394
parameters, buffers = info.kernel_tensors
395
+ non_persistent = info.kernel_non_persistent_buffers
396
+ loaded_tensor_names = {name for name, _ in info.loaded_weights}
397
for name, param in parameters.items():
398
param.data.copy_(getattr(layer, name))
399
for name, buffer in buffers.items():
400
if name not in layer._buffers:
401
continue
402
+ if name in non_persistent and name not in loaded_tensor_names:
403
+ continue
404
buffer.data.copy_(getattr(layer, name))
405
406
_place_kernel_tensors(layer, info)
@@ -404,7 +412,8 @@ def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
412
413
assert info.kernel_tensors is not None
414
parameters, buffers = info.kernel_tensors
415
+ non_persistent = info.kernel_non_persistent_buffers
416
for name, param in parameters.items():
417
layer.register_parameter(name, param)
418
for name, buffer in buffers.items():
410
- layer.register_buffer(name, buffer)
419
+ layer.register_buffer(name, buffer, persistent=name not in non_persistent)
vllm/model_executor/model_loader/reload/meta.py
+2
-1
@@ -117,6 +117,7 @@ def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):
117
if layer.__class__.__name__ in SKIP_MODULES:
118
return
119
120
+ non_persistent = set(layer._non_persistent_buffers_set)
121
for name in get_layer_tensors(layer):
122
if name not in SKIP_TENSORS:
123
delattr(layer, name)
@@ -130,7 +131,7 @@ def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):
131
for name, buffer in restore_buffers.items():
132
if name not in SKIP_TENSORS:
133
buffer = restore_layer_refs(buffer, layer)
133
- layer.register_buffer(name, buffer)
134
+ layer.register_buffer(name, buffer, persistent=name not in non_persistent)
135
136
137
def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo):
vllm/model_executor/model_loader/reload/types.py
+4
@@ -29,6 +29,10 @@ class LayerReloadingInfo:
29
# kernel formatted tensors, copied into by `_layerwise_process` when reloading
30
kernel_tensors: LayerTensors | None = None
31
32
+ # non-persistent buffer names captured with `kernel_tensors`, so buffer
33
+ # persistence survives `_non_persistent_buffers_set` being mutated during reload
34
+ kernel_non_persistent_buffers: set[str] = field(default_factory=set)
35
+
36
def reset(self):
37
self.__init__( # type: ignore[misc]
38
restore_metadata=self.restore_metadata, restore_device=self.restore_device