Remove unnecessary load_weights methods (vllm-project#44589)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

Author: hmellor

tests/model_executor/test_weight_utils.py

@@ -160,5 +160,126 @@ def test_missing_target_returns_none(self):
         assert result is None
 
 
+class TestKvCacheScaleMapper:
+    """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the
+    per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of
+    checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent
+    so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`."""
+
+    def _mapper(self):
+        # `get_cache_scale_mapper` does not use `self`; call it on the base
+        # class to get the default (non-config-specific) mapper.
+        from vllm.model_executor.layers.quantization.base_config import (
+            QuantizationConfig,
+        )
+
+        return QuantizationConfig.get_cache_scale_mapper()
+
+    def _map(self, name: str) -> str | None:
+        return self._mapper()._map_name(name)
+
+    @pytest.mark.parametrize(
+        "name,expected",
+        [
+            # Qwen3-MoE / llm-compressor fused qkv_proj
+            (
+                "model.layers.0.self_attn.qkv_proj.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            (
+                "model.layers.0.self_attn.qkv_proj.v_scale",
+                "model.layers.0.self_attn.attn.v_scale",
+            ),
+            # ModelOpt / NVFP4 k_proj/v_proj
+            (
+                "model.layers.0.self_attn.k_proj.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            (
+                "model.layers.0.self_attn.v_proj.v_scale",
+                "model.layers.0.self_attn.attn.v_scale",
+            ),
+            # deprecated fused kv_scale and bare scales
+            (
+                "model.layers.0.self_attn.kv_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            (
+                "model.layers.0.self_attn.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            # NemotronH mixer
+            (
+                "model.layers.0.mixer.k_proj.k_scale",
+                "model.layers.0.mixer.attn.k_scale",
+            ),
+            # already in vLLM form -> unchanged (idempotent)
+            (
+                "model.layers.0.self_attn.attn.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            # non-kv scales must not be touched
+            (
+                "model.layers.0.self_attn.k_proj.weight_scale",
+                "model.layers.0.self_attn.k_proj.weight_scale",
+            ),
+            (
+                "model.layers.0.self_attn.k_proj.input_scale",
+                "model.layers.0.self_attn.k_proj.input_scale",
+            ),
+            # regular weights untouched
+            (
+                "model.layers.0.self_attn.q_proj.weight",
+                "model.layers.0.self_attn.q_proj.weight",
+            ),
+        ],
+    )
+    def test_remap(self, name, expected):
+        assert self._map(name) == expected
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "model.layers.0.self_attn.k_scale",
+            "model.layers.0.self_attn.k_proj.k_scale",
+            "model.layers.0.self_attn.qkv_proj.v_scale",
+            "model.layers.0.mixer.k_proj.k_scale",
+        ],
+    )
+    def test_idempotent(self, name):
+        once = self._map(name)
+        assert once is not None
+        assert self._map(once) == once
+
+    def test_composes_with_qkv_mapper(self):
+        """Applied together with a model's qkv/gate_up mapper, the regex scale
+        rules run before the substr rename, so scales are normalized to `.attn.`
+        and regular projections are still fused correctly."""
+        from vllm.model_executor.models.utils import WeightsMapper
+
+        model_mapper = WeightsMapper(
+            orig_to_new_substr={
+                ".q_proj": ".qkv_proj.q",
+                ".k_proj": ".qkv_proj.k",
+                ".v_proj": ".qkv_proj.v",
+            }
+        )
+        # AutoWeightsLoader does `mapper |= cache_scale_mapper`
+        combined = model_mapper | self._mapper()
+
+        assert (
+            combined._map_name("model.layers.0.self_attn.q_proj.weight")
+            == "model.layers.0.self_attn.qkv_proj.q.weight"
+        )
+        assert (
+            combined._map_name("model.layers.0.self_attn.k_proj.k_scale")
+            == "model.layers.0.self_attn.attn.k_scale"
+        )
+        assert (
+            combined._map_name("model.layers.0.self_attn.k_scale")
+            == "model.layers.0.self_attn.attn.k_scale"
+        )
+
+
 if __name__ == "__main__":
     test_download_weights_from_hf()

vllm/lora/worker_manager.py

@@ -128,9 +128,13 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel:
             peft_helper.validate_legal(self.lora_config)
 
             # For some models like Qwen2VL, we need to use hf_to_vllm_mapper
-            # to ensure correct loading of lora weights.
+            # to ensure correct loading of lora weights. Drop the QKV/MLP fusion
+            # substr maps so constituent names (e.g. `q_proj`) survive for the
+            # LoRA manager to pack, while keeping genuine renames/prefixes.
             model = self._adapter_manager.model
             hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None)
+            if hf_to_vllm_mapper is not None:
+                hf_to_vllm_mapper = hf_to_vllm_mapper.get_unstacked_mapper()
 
             # Get model-defined prefixes to skip during LoRA loading.
             lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None)

vllm/model_executor/layers/linear.py

@@ -3,9 +3,12 @@
 
 import itertools
 from abc import abstractmethod
+from collections.abc import Iterable
+from typing import Any
 
 import torch
 from torch.nn.parameter import Parameter
+from typing_extensions import TypeIs
 
 import vllm.envs as envs
 from vllm.distributed import (
@@ -632,31 +635,31 @@ def __init__(
             disable_tp=disable_tp,
         )
 
-    def validate_shard_id(self, loaded_shard_id: int | tuple[int, ...] | None):
-        if loaded_shard_id is None:
-            return
-        if isinstance(loaded_shard_id, tuple):
-            for idx in loaded_shard_id:
+    def validate_shard_id(self, shard_id: Any) -> TypeIs[int | tuple[int, ...] | None]:
+        if isinstance(shard_id, int):
+            if shard_id < 0 or shard_id >= len(self.output_sizes):
+                raise ValueError(
+                    f"Shard id should be between 0 and {len(self.output_sizes) - 1}. "
+                    f"Got shard id {shard_id}."
+                )
+            return True
+        if shard_id is None:
+            return True
+        if isinstance(shard_id, tuple):
+            for idx in shard_id:
                 if not (0 <= idx < len(self.output_sizes)):
                     raise ValueError(
                         f"Shard id index {idx} should be between 0 and "
-                        f"{len(self.output_sizes) - 1}. Got shard id {loaded_shard_id}."
+                        f"{len(self.output_sizes) - 1}. Got shard id {shard_id}."
                     )
-            if len(loaded_shard_id) > 1 and any(
-                b - a != 1 for a, b in zip(loaded_shard_id[:-1], loaded_shard_id[1:])
+            if len(shard_id) > 1 and any(
+                b - a != 1 for a, b in zip(shard_id[:-1], shard_id[1:])
             ):
                 raise ValueError(
                     "Shard id with multiple indices should be consecutive. "
-                    f"Got shard id {loaded_shard_id}."
+                    f"Got shard id {shard_id}."
                 )
-            return
-        elif isinstance(loaded_shard_id, int):
-            if loaded_shard_id < 0 or loaded_shard_id >= len(self.output_sizes):
-                raise ValueError(
-                    f"Shard id should be between 0 and {len(self.output_sizes) - 1}. "
-                    f"Got shard id {loaded_shard_id}."
-                )
-            return
+            return True
         raise ValueError("This line should not be reached")
 
     def weight_loader(
@@ -910,6 +913,31 @@ def weight_loader_v2(
             tp_rank=self.tp_rank,
         )
 
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[str]:
+        for name, loaded_weight in weights:
+            shard_id = getattr(loaded_weight, "shard_id", None)
+            self.validate_shard_id(shard_id)
+            # Load into self if name is not an attr of self or its submodules
+            param: Parameter
+            if "." in name:
+                submodule, _, attr = name.rpartition(".")
+                param = getattr(self.get_submodule(submodule), attr, self)
+            else:
+                param = getattr(self, name, self)
+            if param is None and name == "bias":
+                continue
+            param.weight_loader(param, loaded_weight, shard_id)
+            logger.debug(
+                "Loaded shard %s with shape %s into %s.%s",
+                shard_id,
+                loaded_weight.shape,
+                self.prefix,
+                name,
+            )
+            yield name
+
 
 class QKVParallelLinear(ColumnParallelLinear):
     """Linear layers for the attention's QKV transformation.
@@ -996,17 +1024,13 @@ def __init__(
             disable_tp=disable_tp,
         )
 
-    def validate_shard_id(self, loaded_shard_id: str | None):
-        if loaded_shard_id is None:
-            return
-        if isinstance(loaded_shard_id, str):
-            if loaded_shard_id not in ["q", "k", "v"]:
-                raise ValueError(
-                    "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', "
-                    f"got shard id {loaded_shard_id}."
-                )
-            return
-        raise ValueError("This line should not be reached")
+    def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]:
+        if shard_id in {"q", "k", "v"} or shard_id is None:
+            return True
+        raise ValueError(
+            "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', "
+            f"got shard id {shard_id}."
+        )
 
     def _get_shard_offset_mapping(self, loaded_shard_id: str):
         shard_offset_mapping = {
@@ -1302,6 +1326,31 @@ def weight_loader(
         assert param_data.shape == loaded_weight.shape
         param_data.copy_(loaded_weight)
 
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[str]:
+        for name, loaded_weight in weights:
+            shard_id = getattr(loaded_weight, "shard_id", None)
+            self.validate_shard_id(shard_id)
+            # Load into self if name is not an attr of self or its submodules
+            param: Parameter
+            if "." in name:
+                submodule, _, attr = name.rpartition(".")
+                param = getattr(self.get_submodule(submodule), attr, self)
+            else:
+                param = getattr(self, name, self)
+            if param is None and name == "bias":
+                continue
+            param.weight_loader(param, loaded_weight, shard_id)
+            logger.debug(
+                "Loaded shard %s with shape %s into %s.%s",
+                shard_id,
+                loaded_weight.shape,
+                self.prefix,
+                name,
+            )
+            yield name
+
 
 class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):
     """QKV projection fused with a lightning-indexer's index_q/index_k.
@@ -1387,15 +1436,14 @@ def __init__(
             prefix=prefix,
         )
 
-    def validate_shard_id(self, loaded_shard_id: str | None) -> None:
-        if loaded_shard_id is None:
-            return
-        if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"):
-            raise ValueError(
-                "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "
-                "'q', 'k', 'v', 'index_q', 'index_k'; got "
-                f"{loaded_shard_id}."
-            )
+    def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]:
+        if shard_id in {"q", "k", "v", "index_q", "index_k"} or shard_id is None:
+            return True
+        raise ValueError(
+            "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "
+            "'q', 'k', 'v', 'index_q', 'index_k'; got "
+            f"{shard_id}."
+        )
 
     def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:
         h = self.head_size