[Bugfix][ModelLoader] Make enable_weights_track check quantized layers (#57108)

Signed-off-by: siliangchen-amd <SiLiang.Chen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Shanshan Shen <467638484@qq.com>

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"}

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():

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