[Bugfix][Model Loader] Defer post-load attention weight processing (#49519)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io>
aoshen02 committed
Aug 11, 2026 at 10:16 UTC
99e62b802c2a2fcf8393e0ca87d8b7d8f00297d3
4 files changed
+105
-24
tests/model_executor/model_loader/test_reload.py
+77
@@ -9,6 +9,9 @@ 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.config import ModelConfig
13
+from vllm.model_executor.layers.attention import MMEncoderAttention
14
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
15
from vllm.model_executor.layers.linear import QKVParallelLinear
16
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
17
from vllm.model_executor.model_loader.reload.layerwise import (
@@ -101,6 +104,39 @@ class _NonPersistentBufferLayer(torch.nn.Module):
104
self.register_buffer("scale", torch.tensor(0.25), persistent=False)
105
106
107
+class _ReloadableMMEncoderAttention(MMEncoderAttention):
108
+ """Minimal stand-in to test reload lifecycle without encoder initialization."""
109
+
110
+ def __init__(self):
111
+ torch.nn.Module.__init__(self)
112
+ self.weight = torch.nn.Parameter(torch.ones(2, 2))
113
+ self.weight.weight_loader = default_weight_loader
114
+ self.post_load_called = False
115
+
116
+ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
117
+ self.post_load_called = True
118
+
119
+
120
+class _ReloadableAttentionLayer(
121
+ torch.nn.Module,
122
+ AttentionLayerBase,
123
+):
124
+ def __init__(self):
125
+ super().__init__()
126
+ self.weight = torch.nn.Parameter(torch.ones(2, 2))
127
+ self.weight.weight_loader = default_weight_loader
128
+ self.post_load_called = False
129
+
130
+ def get_attn_backend(self):
131
+ raise NotImplementedError
132
+
133
+ def get_kv_cache_spec(self, vllm_config):
134
+ return None
135
+
136
+ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
137
+ self.post_load_called = True
138
+
139
+
140
def test_move_metatensors():
141
tensor = torch.empty((1, 2, 3))
142
meta_tensor = to_meta_tensor(tensor)
@@ -115,6 +151,47 @@ def test_move_metatensors():
151
assert tensor.__dict__ == meta_tensor.__dict__ == materialized_tensor.__dict__
152
153
154
+@pytest.mark.parametrize(
155
+ "layer_cls",
156
+ [_ReloadableMMEncoderAttention, _ReloadableAttentionLayer],
157
+)
158
+def test_attention_reload_defers_post_load(default_vllm_config, layer_cls):
159
+ default_vllm_config.model_config = ModelConfig()
160
+ layer = layer_cls()
161
+ model = torch.nn.Sequential(layer)
162
+ loaded_weight = torch.full_like(layer.weight, 7.0)
163
+
164
+ record_metadata_for_reloading(model)
165
+ initialize_layerwise_reload(model)
166
+ layer.weight.weight_loader(layer.weight, loaded_weight)
167
+
168
+ assert not layer.post_load_called
169
+
170
+ finalize_layerwise_reload(model, default_vllm_config.model_config)
171
+
172
+ assert layer.post_load_called
173
+ assert torch.equal(layer.weight, loaded_weight)
174
+
175
+
176
+@pytest.mark.parametrize(
177
+ "layer_cls",
178
+ [_ReloadableMMEncoderAttention, _ReloadableAttentionLayer],
179
+)
180
+def test_attention_first_load_processes_weights(default_vllm_config, layer_cls):
181
+ default_vllm_config.model_config = ModelConfig()
182
+ layer = layer_cls()
183
+ model = torch.nn.Sequential(layer)
184
+ loaded_weight = torch.full_like(layer.weight, 7.0)
185
+
186
+ initialize_online_processing(layer)
187
+ layer.weight.weight_loader(layer.weight, loaded_weight)
188
+
189
+ finalize_layerwise_reload(model, default_vllm_config.model_config)
190
+
191
+ assert layer.post_load_called
192
+ assert torch.equal(layer.weight, loaded_weight)
193
+
194
+
195
def test_reload_lifecycle():
196
layer = torch.nn.Linear(2, 3)
197
info = LayerReloadingInfo(
vllm/model_executor/layers/attention/__init__.py
+12
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import torch
5
+
6
from vllm.model_executor.layers.attention.attention import Attention
7
from vllm.model_executor.layers.attention.chunked_local_attention import (
8
ChunkedLocalAttention,
@@ -18,6 +20,15 @@ from vllm.model_executor.layers.attention.rswa_attention import RSWAAttention
20
from vllm.model_executor.layers.attention.static_sink_attention import (
21
StaticSinkAttention,
22
)
23
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
24
+
25
+
26
+def is_deferred_attention_layer(layer: torch.nn.Module) -> bool:
27
+ """Whether an attention-like layer requires deferred post-load processing."""
28
+ return isinstance(layer, (AttentionLayerBase, MMEncoderAttention)) and callable(
29
+ getattr(layer, "process_weights_after_loading", None)
30
+ )
31
+
32
33
__all__ = [
34
"Attention",
@@ -29,4 +40,5 @@ __all__ = [
40
"PrefillPrefixLMAttention",
41
"RSWAAttention",
42
"StaticSinkAttention",
43
+ "is_deferred_attention_layer",
44
]
vllm/model_executor/model_loader/reload/layerwise.py
+14
-17
@@ -9,7 +9,7 @@ import torch
9
10
from vllm.config import ModelConfig
11
from vllm.logger import init_logger
12
-from vllm.model_executor.layers.attention import Attention, MLAAttention
12
+from vllm.model_executor.layers.attention import is_deferred_attention_layer
13
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
14
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
@@ -196,7 +196,7 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla
196
)
197
198
# Do not online process attention layers, must wait until finalize
199
- if isinstance(layer, (Attention, MLAAttention)):
199
+ if is_deferred_attention_layer(layer):
200
return ret
201
202
# Log warnings allocating excessive buffers on device
@@ -249,8 +249,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
249
info.reset()
250
continue
251
252
- # Attention/MLA layers are processed after all other layers
253
- if isinstance(layer, (Attention, MLAAttention)):
252
+ # Deferred attention-like layers are processed after all other layers
253
+ if is_deferred_attention_layer(layer):
254
deferred_attn.append((layer, info))
255
continue
256
@@ -293,15 +293,13 @@ def finalize_layerwise_reload(*args, **kwargs):
293
def _finalize_attention_layer(
294
layer: torch.nn.Module, info: LayerReloadingInfo, model_config: ModelConfig
295
) -> None:
296
- if info.load_numel > 0 and info.kernel_tensors is not None:
296
+ if info.kernel_tensors is None:
297
+ if info.load_numel > 0:
298
+ _layerwise_process(layer, info)
299
+ elif info.load_numel > 0:
300
# Reload with new scale weights from checkpoint
301
_place_kernel_tensors(layer, info)
302
_reload_attention_scales(layer, info)
300
- elif info.load_numel > 0 or info.kernel_tensors is None:
301
- raise ValueError(
302
- "Layerwise loading of attention layers is not supported. "
303
- "Attention must always process after linears."
304
- )
303
else:
304
_place_kernel_tensors(layer, info)
305
layer.process_weights_after_loading(model_config.dtype)
@@ -315,19 +313,18 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -
313
processing, since we use .data.copy_() to preserve kernel tensor
314
references."""
315
quant_method = getattr(layer, "quant_method", None)
318
- if quant_method is None:
319
- return
320
-
321
- # Re-create scale Parameters with sentinel values so unloaded scales
322
- # are correctly detected by process_weights_after_loading
323
- quant_method.create_weights(layer)
316
+ if quant_method is not None:
317
+ # Re-create scale Parameters with sentinel values so unloaded scales
318
+ # are correctly detected by process_weights_after_loading
319
+ quant_method.create_weights(layer)
320
321
for name, args in info.loaded_weights:
322
param = getattr(layer, name)
323
args.arguments["param"] = param
324
_get_weight_loader(param)(*args.args, **args.kwargs)
325
330
- quant_method.process_weights_after_loading(layer)
326
+ if quant_method is not None:
327
+ quant_method.process_weights_after_loading(layer)
328
329
_copy_and_restore_kernel_tensors(layer, info)
330
vllm/model_executor/model_loader/utils.py
+2
-7
@@ -14,10 +14,7 @@ from typing_extensions import assert_never
14
import vllm.envs as envs
15
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
16
from vllm.logger import init_logger
17
-from vllm.model_executor.layers.attention import (
18
- MMEncoderAttention,
19
-)
20
-from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17
+from vllm.model_executor.layers.attention import is_deferred_attention_layer
18
from vllm.model_executor.layers.hpc import HpcModule
19
from vllm.model_executor.layers.quantization.base_config import (
20
QuantizationConfig,
@@ -124,9 +121,7 @@ def process_weights_after_loading(
121
# encoder. NOTE: Happens after other modules so we can easily decompress
122
# weights.
123
for _, module in model.named_modules():
127
- if isinstance(module, (AttentionLayerBase, MMEncoderAttention)) and hasattr(
128
- module, "process_weights_after_loading"
129
- ):
124
+ if is_deferred_attention_layer(module):
125
# TODO(lucas): see if there is a way to unify the signatures
126
# of process_weights_after_loading
127
with device_loading_context(module, target_device):