[Bugfix] Avoid repeated layerwise reload warning scans (#49505)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
aoshen02 committed
Aug 12, 2026 at 13:53 UTC
f2906b492e9173ce39920a3927699f202b2cc2e5
2 files changed
+33
-2
tests/model_executor/model_loader/test_reload.py
+31
@@ -2,12 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import gc
4
import inspect
5
+from unittest.mock import Mock
6
from weakref import WeakKeyDictionary, ref
7
8
import pytest
9
import torch
10
from torch.nn.parameter import UninitializedParameter
11
12
+import vllm.model_executor.model_loader.reload.layerwise as reload_layerwise
13
import vllm.model_executor.model_loader.reload.meta as reload_meta
14
from vllm.config import ModelConfig
15
from vllm.model_executor.layers.attention import MMEncoderAttention
@@ -571,6 +573,35 @@ def test_get_numel_loaded_caps_at_param_size():
573
assert num_loaded == 10
574
575
576
+def test_layerwise_loading_warning_only_checks_new_layers(monkeypatch):
577
+ layers = [torch.nn.Linear(16, 1, bias=False) for _ in range(2)]
578
+
579
+ def partial_weight_loader(param, loaded_weight):
580
+ param.view(-1)[: loaded_weight.numel()].copy_(loaded_weight)
581
+
582
+ for layer in layers:
583
+ layer.weight.requires_grad_(False)
584
+ layer.weight.weight_loader = partial_weight_loader
585
+ reload_layerwise.initialize_online_processing(layer)
586
+
587
+ monkeypatch.setattr(reload_layerwise, "has_device_tensors", lambda _: True)
588
+ get_info_size = Mock(return_value=0)
589
+ warning_once = Mock()
590
+ monkeypatch.setattr(reload_layerwise, "get_info_size", get_info_size)
591
+ monkeypatch.setattr(reload_layerwise.logger, "warning_once", warning_once)
592
+
593
+ reload_layerwise.LOADING_LAYERS.clear()
594
+ try:
595
+ for layer in layers:
596
+ for _ in range(3):
597
+ layer.weight.weight_loader(layer.weight, torch.ones(1))
598
+ finally:
599
+ reload_layerwise.LOADING_LAYERS.clear()
600
+
601
+ assert get_info_size.call_count == 2
602
+ warning_once.assert_called_once()
603
+
604
+
605
class _ComposedLoaderLayer(torch.nn.Module):
606
"""Mimics a Mamba2 mixer's equal-numel direct params (A, D, dt_bias).
607
vllm/model_executor/model_loader/reload/layerwise.py
+2
-2
@@ -200,9 +200,9 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla
200
return ret
201
202
# Log warnings allocating excessive buffers on device
203
- if has_device_tensors(bound_args):
203
+ if has_device_tensors(bound_args) and layer not in LOADING_LAYERS:
204
LOADING_LAYERS.add(layer)
205
- if len(LOADING_LAYERS) >= 2:
205
+ if len(LOADING_LAYERS) == 2:
206
names = sorted([layer.__class__.__name__ for layer in LOADING_LAYERS])
207
mem_used = sum(
208
get_info_size(LAYERWISE_INFO[layer]) for layer in LOADING_LAYERS