[Bugfix] Wait for the linear bias before layerwise online processing (#49805)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Harry Mellor committed
Jul 25, 2026 at 18:56 UTC
9321aff536c1a73f0a7fefc68aef1f2630a904cc
4 files changed
+67
-10
tests/model_executor/model_loader/test_reload.py
+54
@@ -10,9 +10,11 @@ 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.layers.quantization.base_config import QuantizeMethodBase
14
from vllm.model_executor.model_loader.reload.layerwise import (
15
finalize_layerwise_reload,
16
initialize_layerwise_reload,
17
+ initialize_online_processing,
18
record_metadata_for_reloading,
19
)
20
from vllm.model_executor.model_loader.reload.meta import (
@@ -278,6 +280,58 @@ def test_layerwise_reload_composed_loader_does_not_drop_params(monkeypatch):
280
assert torch.equal(layer.D, loaded["D"])
281
282
283
+class _RecordingQuantMethod(QuantizeMethodBase):
284
+ """Records the layer's bias at the moment processing runs."""
285
+
286
+ uses_meta_device = True
287
+
288
+ def __init__(self):
289
+ self.bias_at_process = None
290
+
291
+ def create_weights(self, layer, *weight_args, **extra_weight_attrs):
292
+ pass
293
+
294
+ def apply(self, layer, *args, **kwargs):
295
+ raise NotImplementedError
296
+
297
+ def process_weights_after_loading(self, layer):
298
+ self.bias_at_process = layer.bias.detach().clone()
299
+
300
+
301
+class _LateBiasLayer(torch.nn.Module):
302
+ """Mimics an online-quantized linear: `weight` is created on meta by
303
+ `create_weights()`, which wraps the loaders, and the linear base registers
304
+ `bias` afterwards."""
305
+
306
+ def __init__(self, quant_method):
307
+ super().__init__()
308
+ self.quant_method = quant_method
309
+ weight = torch.nn.Parameter(torch.empty(4, 2, device="meta"))
310
+ weight.weight_loader = default_weight_loader
311
+ self.register_parameter("weight", weight)
312
+ initialize_online_processing(self)
313
+ bias = torch.nn.Parameter(torch.zeros(4))
314
+ bias.weight_loader = default_weight_loader
315
+ self.register_parameter("bias", bias)
316
+
317
+
318
+def test_online_processing_waits_for_late_registered_bias():
319
+ # Regression test: `bias` is skipped by the meta device paths, but it is
320
+ # still loaded by a weight loader. Excluding it from the processing trigger
321
+ # finalized the layer one load early, so the trailing bias was written into
322
+ # an already-processed layer (e.g. over FP8 Marlin's permuted bias).
323
+ quant_method = _RecordingQuantMethod()
324
+ layer = _LateBiasLayer(quant_method)
325
+ loaded_bias = torch.full((4,), 3.0)
326
+
327
+ layer.weight.weight_loader(layer.weight, torch.full((4, 2), 2.0))
328
+ assert quant_method.bias_at_process is None
329
+
330
+ layer.bias.weight_loader(layer.bias, loaded_bias)
331
+ assert quant_method.bias_at_process is not None
332
+ assert torch.equal(quant_method.bias_at_process, loaded_bias)
333
+
334
+
335
def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch):
336
layer = _AliasedBufferLayer()
337
model = torch.nn.Sequential(layer)
vllm/model_executor/model_loader/reload/layerwise.py
+2
-2
@@ -14,7 +14,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBa
14
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
16
from .meta import (
17
- SKIP_TENSORS,
17
+ SKIP_LOAD_TENSORS,
18
capture_layer_to_meta,
19
get_numel_loaded,
20
materialize_layer,
@@ -140,7 +140,7 @@ def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None:
140
"""Wrap each parameter's weight loader."""
141
# Note that nested wrapping will occur for shared tensors
142
for name, tensor in get_layer_tensors(layer).items():
143
- if name in SKIP_TENSORS:
143
+ if name in SKIP_LOAD_TENSORS:
144
continue
145
if _get_weight_loader(tensor).__name__ != "online_process_loader":
146
tensor.weight_loader = make_online_process_loader(layer, name)
vllm/model_executor/model_loader/reload/meta.py
+7
-4
@@ -20,20 +20,23 @@ __all__ = [
20
"get_numel_loaded",
21
]
22
23
+# Modules whose tensors are never moved to, or materialized from, the meta device.
24
SKIP_MODULES: set[str] = {"HadamardTransform"}
25
25
-SKIP_TENSORS: set[str] = {
26
+# Tensors never loaded by a weight loader, so the layerwise trigger ignores them.
27
+SKIP_LOAD_TENSORS: set[str] = {
28
"_expert_map",
29
"expert_mask",
30
"expert_global_to_physical",
31
"expert_physical_to_global",
32
"expert_local_to_global",
33
"e_score_correction_bias",
32
- # Built after create_weights(), so it is not tracked by the layerwise-reload
33
- # trigger and would be re-materialized into uninitialized memory. Skip it.
34
- "bias",
34
}
35
36
+# Tensors which are never moved to, or materialized from, the meta device.
37
+# `bias` is built after create_weights(), so it is never on meta to begin with.
38
+SKIP_TENSORS: set[str] = SKIP_LOAD_TENSORS | {"bias"}
39
+
40
41
def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor:
42
"""Convert a tensor to a meta tensor while preserving class and attributes."""
vllm/model_executor/model_loader/reload/utils.py
+4
-4
@@ -33,15 +33,15 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
33
def get_layer_size(layer: torch.nn.Module) -> int:
34
"""Calculate total number of elements across loadable tensors in a layer.
35
36
- Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta
37
- device and never loaded via weight_loader during layerwise reload.
36
+ Excludes SKIP_LOAD_TENSORS (e.g. _expert_map) which are never loaded via
37
+ weight_loader during layerwise reload.
38
"""
39
- from .meta import SKIP_TENSORS
39
+ from .meta import SKIP_LOAD_TENSORS
40
41
return sum(
42
tensor.numel()
43
for name, tensor in get_layer_tensors(layer).items()
44
- if name not in SKIP_TENSORS
44
+ if name not in SKIP_LOAD_TENSORS
45
)
46
47