From 02aa10539f15c8c73aa7fd4a5ee0f691cbf88dcb Mon Sep 17 00:00:00 2001 From: siliangchen-amd Date: Sat, 10 Oct 2026 10:44:06 +0800 Subject: [PATCH] [Bugfix][ModelLoader] Make enable_weights_track check quantized layers (#57108) Signed-off-by: siliangchen-amd Co-authored-by: Cursor Co-authored-by: Shanshan Shen <467638484@qq.com> --- .../model_loader/test_registry.py | 82 +++++++++++++++++++ .../layers/fused_moe/runner/moe_runner.py | 5 +- .../model_loader/default_loader.py | 36 +++++++- 3 files changed, 119 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/model_loader/test_registry.py b/tests/model_executor/model_loader/test_registry.py index a795497a11..7e4b6c537a 100644 --- a/tests/model_executor/model_loader/test_registry.py +++ b/tests/model_executor/model_loader/test_registry.py @@ -2,13 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.linear import UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.model_loader import get_model_loader, register_model_loader from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader @register_model_loader("custom_load_format") @@ -88,3 +93,80 @@ def test_default_loader_hf_still_falls_back_to_pt(tmp_path): ) assert use_safetensors is False assert any(f.endswith("model.pt") for f in files) + + +class _QuantMethod: + uses_meta_device = False + + def process_weights_after_loading(self, layer: nn.Module) -> None: + pass + + +def _layer(quant_method=None, **shapes) -> nn.Module: + layer = nn.Module() + for name, shape in shapes.items(): + param = nn.Parameter(torch.zeros(shape), requires_grad=False) + layer.register_parameter(name, param) + layer.quant_method = quant_method + return layer + + +def _track(model: nn.Module, loaded: set[str], quantized: bool) -> str: + loader = DefaultModelLoader(LoadConfig()) + try: + loader.track_weights_loading(model, set(loaded), quantized=quantized) + except ValueError as e: + return str(e) + return "" + + +def test_weights_track_reports_missing_quantized_weights(): + # Quantized checkpoints may only omit KV-cache quantization params, empty + # placeholders, fbgemm_fp8's input_scale_ub and online-quant params. + online = _QuantMethod() + online.uses_meta_device = True + model = nn.Module() + model.proj = _layer(_QuantMethod(), weight=(4, 4), weight_scale_inv=(1, 1)) + model.unquantized = _layer(UnquantizedLinearMethod(), weight=(4, 4), bias=(4,)) + model.gptq = _layer(_QuantMethod(), qweight=(4, 4), qzeros=(0,)) + model.fbgemm = _layer(_QuantMethod(), weight=(4, 4), input_scale_ub=()) + model.attn = _layer(BaseKVCacheMethod(None), k_scale=(), k_zero_point=()) + model.online = _layer(online, weight=(4, 4)) + loaded = {"proj.weight", "unquantized.weight", "gptq.qweight", "fbgemm.weight"} + error = _track(model, loaded, quantized=True) + assert "proj.weight_scale_inv" in error and "unquantized.bias" in error + assert not any(n in error for n in ("gptq", "fbgemm", "attn", "online")) + + +def test_weights_track_keeps_unquantized_models_unchanged(): + # The check is on by default here, so unquantized layers stay exempt. + model = nn.Module() + model.score = _layer(_QuantMethod(), weight=(2, 4)) + model.norm = _layer(weight=(4,)) + assert _track(model, {"norm.weight"}, quantized=False) == "" + assert "norm.weight" in _track(model, set(), quantized=False) + + +def test_moe_runner_reports_registered_expert_names(): + # AutoWeightsLoader prefixes these names, so they must include routed_experts. + class Experts(nn.Module): + def __init__(self): + super().__init__() + self.w13_weight = nn.Parameter(torch.zeros(2, 8, 4), requires_grad=False) + + def load_weights(self, weights): + for _ in weights: + yield "w13_weight" + + class Runner(nn.Module): + load_weights = MoERunner.load_weights + + def __init__(self): + super().__init__() + self.routed_experts = Experts() + + model = nn.Module() + model.experts = Runner() + weights = [("experts.w13_weight", torch.zeros(2, 8, 4))] + loaded = AutoWeightsLoader(model).load_weights(weights) + assert loaded == {"experts.routed_experts.w13_weight"} diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 9cb7cfbb0f..d220f8da36 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -301,7 +301,10 @@ class MoERunner(MoERunnerInterface): def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[str]: - return self.routed_experts.load_weights(weights) + # routed_experts reports names relative to itself, but callers such as + # AutoWeightsLoader qualify what this returns with the runner's prefix. + for name in self.routed_experts.load_weights(weights): + yield f"routed_experts.{name}" def _select_forward(self) -> Callable: if current_platform.is_tpu(): diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 9f1c25dd2e..b162f60b05 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -16,6 +16,7 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.ep_weight_filter import ( @@ -47,6 +48,9 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# fbgemm_fp8 builds input_scale_ub from its config; checkpoints do not carry it. +_CONFIG_BUILT_PARAMS = ("input_scale_ub",) + class DefaultModelLoader(BaseModelLoader): """Model loader that can load different file types from disk.""" @@ -506,10 +510,15 @@ class DefaultModelLoader(BaseModelLoader): else default_enable_weights_track ) if enable_weights_track: - self.track_weights_loading(model, loaded_weights) + self.track_weights_loading( + model, loaded_weights, quantized=model_config.quantization is not None + ) def track_weights_loading( - self, model: nn.Module, loaded_weights: set[str] | None + self, + model: nn.Module, + loaded_weights: set[str] | None, + quantized: bool = False, ) -> None: weights_to_load = {name for name, _ in model.named_parameters()} if loaded_weights is not None: @@ -523,7 +532,28 @@ class DefaultModelLoader(BaseModelLoader): # ignore kv_cache scale and online quant scale, # which can be missing in checkpoints if has_online_quant or has_postprocess_quant: - for param_name, _ in module.named_parameters(): + for param_name, param in module.named_parameters(): + # On a quantized model, every parameter has to come from + # the checkpoint except those a checkpoint cannot carry: + # - online quantization (uses_meta_device) quantizes the + # weights while loading, so its parameters do not map + # to checkpoint tensors; + # - KV-cache quantization parameters (BaseKVCacheMethod) + # fall back to defaults when a checkpoint has none; + # - empty placeholders hold no data, e.g. the qzeros + # moe_wna16 registers for symmetric GPTQ; + # - _CONFIG_BUILT_PARAMS come from the quantization config. + # Unquantized models, where this check runs by default, + # keep the module-wide exemption. + if ( + quantized + and not has_online_quant + and not isinstance(quant_method, BaseKVCacheMethod) + and param.numel() > 0 + and param_name.rsplit(".", 1)[-1] + not in _CONFIG_BUILT_PARAMS + ): + continue full_name = f"{name}.{param_name}" if name else param_name loaded_weights.add(full_name) weights_not_loaded = weights_to_load - loaded_weights