Make load_weights completely optional

tests/models/test_utils.py

@@ -89,82 +89,6 @@ def weight_generator():
     assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
 
 
-@pytest.mark.cpu_test
-def test_module_skip_prefix():
-    """Ensure the auto weight loader can skip prefix."""
-    mod = ModuleWithNestedBatchNorm()
-    # Run some data through the module with batchnorm
-    mod(torch.Tensor([[1, 2], [3, 4]]))
-
-    # Try to load the weights to a new instance
-    def weight_generator():
-        # weights needed to be filtered out
-        redundant_weights = {
-            "prefix.bn.weight": torch.Tensor([1, 2]),
-            "prefix.bn.bias": torch.Tensor([3, 4]),
-        }
-        yield from (mod.state_dict() | redundant_weights).items()
-
-    new_mod = ModuleWithNestedBatchNorm()
-
-    assert not torch.all(
-        new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
-    )
-    assert not torch.all(
-        new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var
-    )
-    assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0
-
-    loader = AutoWeightsLoader(new_mod, skip_prefixes=["prefix."])
-    loader.load_weights(weight_generator())
-
-    # Ensure the stats are updated
-    assert torch.all(
-        new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
-    )
-    assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)
-    assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
-
-
-@pytest.mark.cpu_test
-def test_module_skip_substr():
-    """Ensure the auto weight loader can skip prefix."""
-    mod = ModuleWithNestedBatchNorm()
-    # Run some data through the module with batchnorm
-    mod(torch.Tensor([[1, 2], [3, 4]]))
-
-    # Try to load the weights to a new instance
-    def weight_generator():
-        # weights needed to be filtered out
-        redundant_weights = {
-            "nested_mod.0.substr.weight": torch.Tensor([1, 2]),
-            "nested_mod.0.substr.bias": torch.Tensor([3, 4]),
-            "nested_mod.substr.weight": torch.Tensor([1, 2]),
-            "nested_mod.substr.bias": torch.Tensor([3, 4]),
-        }
-        yield from (mod.state_dict() | redundant_weights).items()
-
-    new_mod = ModuleWithNestedBatchNorm()
-
-    assert not torch.all(
-        new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
-    )
-    assert not torch.all(
-        new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var
-    )
-    assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0
-
-    loader = AutoWeightsLoader(new_mod, skip_substrs=["substr."])
-    loader.load_weights(weight_generator())
-
-    # Ensure the stats are updated
-    assert torch.all(
-        new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
-    )
-    assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)
-    assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
-
-
 class raise_if_cuda_sync:
     def __enter__(self):
         self.previous_debug_mode = torch.cuda.get_sync_debug_mode()

tests/v1/shutdown/test_startup_error.py

@@ -61,7 +61,7 @@ def test_async_llm_startup_error(
         pytest.skip(reason="Not enough CUDA devices")
 
     # Monkeypatch an error in the model.
-    monkeypatch.setattr(LlamaForCausalLM, failing_method, evil_method)
+    monkeypatch.setattr(LlamaForCausalLM, failing_method, evil_method, raising=False)
 
     engine_args = AsyncEngineArgs(
         model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size

vllm/distributed/weight_transfer/ipc_engine.py

@@ -20,6 +20,7 @@
     WeightTransferInitInfo,
     WeightTransferUpdateInfo,
 )
+from vllm.model_executor.model_loader.utils import autoload_weights
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -274,7 +275,7 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None:
                 weight = rebuild_cuda_tensor(*list_args)
                 weights.append((name, weight))
 
-        self.model.load_weights(weights)
+        autoload_weights(self.model, weights)
 
     def shutdown(self) -> None:
         pass

vllm/distributed/weight_transfer/nccl_engine.py

@@ -27,6 +27,7 @@
     DEFAULT_PACKED_NUM_BUFFERS,
     packed_nccl_broadcast_consumer,
 )
+from vllm.model_executor.model_loader.utils import autoload_weights
 
 # Re-exported for backward compatibility; canonical home is nccl_common.
 __all__ = [
@@ -177,11 +178,14 @@ def state_dict_info_iterator():
                     dtype = getattr(torch, dtype_name)
                     yield (name, (shape, dtype))
 
+            def load_weights(w: list[tuple[str, torch.Tensor]]) -> None:
+                autoload_weights(self.model, w)
+
             packed_nccl_broadcast_consumer(
                 iterator=state_dict_info_iterator(),
                 group=self.model_update_group,
                 src=0,
-                post_unpack_func=self.model.load_weights,
+                post_unpack_func=load_weights,
                 buffer_size_bytes=update_info.packed_buffer_size_bytes,
                 num_buffers=update_info.packed_num_buffers,
                 device=self.device,
@@ -196,7 +200,7 @@ def state_dict_info_iterator():
                 self.model_update_group.broadcast(
                     weight, src=0, stream=torch.cuda.current_stream()
                 )
-                self.model.load_weights([(name, weight)])
+                autoload_weights(self.model, [(name, weight)])
                 del weight
 
     def shutdown(self) -> None:

vllm/model_executor/layers/linear.py

@@ -973,6 +973,9 @@ class QKVParallelLinear(ColumnParallelLinear):
                         (e.g. model.layers.0.qkv_proj)
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, weights matrix won't be sharded through tp rank.
+        v_head_size: size of each attention value head.
+            If None, assume v_head_size = head_size.
+        fused_qkv_interleaved: If true, QKV weights are fused and interleaved on disk.
     """
 
     def __init__(
@@ -990,10 +993,12 @@ def __init__(
         return_bias: bool = True,
         disable_tp: bool = False,
         v_head_size: int | None = None,
+        fused_qkv_interleaved: bool = False,
     ):
         self.hidden_size = hidden_size
         self.head_size = head_size
         self.v_head_size = v_head_size if v_head_size is not None else head_size
+        self.fused_qkv_interleaved = fused_qkv_interleaved
         self.total_num_heads = total_num_heads
         if total_num_kv_heads is None:
             total_num_kv_heads = total_num_heads
@@ -1054,6 +1059,27 @@ def _get_shard_size_mapping(self, loaded_shard_id: str):
         }
         return shard_size_mapping.get(loaded_shard_id)
 
+    def _deinterleave_fused_qkv(self, loaded_weight: torch.Tensor) -> torch.Tensor:
+        """De-interleave a per-KV-group fused qkv weight/bias to [Q|K|V].
+
+        The on-disk layout groups each KV head's query heads, key and value
+        together: ``[q_0..q_{g-1}, k, v]`` repeated per KV head, where
+        ``g = total_num_heads // total_num_kv_heads``. This reorders it to the
+        block-contiguous ``[Q_all | K_all | V_all]`` that the fused split below
+        expects. Operates on the output dim (0); trailing dims (hidden, or none
+        for a bias) are preserved. Assumes a uniform head size across q/k/v.
+        """
+        heads = self.total_num_heads
+        kv_heads = self.total_num_kv_heads
+        hs = self.head_size
+        groups = heads // kv_heads
+        rest = loaded_weight.shape[1:]
+        x = loaded_weight.reshape(kv_heads, groups + 2, hs, *rest)
+        q = x[:, :groups].reshape(heads * hs, *rest)
+        k = x[:, groups : groups + 1].reshape(kv_heads * hs, *rest)
+        v = x[:, groups + 1 : groups + 2].reshape(kv_heads * hs, *rest)
+        return torch.cat((q, k, v), dim=0)
+
     def _load_fused_module_from_checkpoint(
         self, param: BasevLLMParameter, loaded_weight: torch.Tensor
     ):
@@ -1066,6 +1092,8 @@ def _load_fused_module_from_checkpoint(
         An example of a model with these fused layers:
         https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
         """
+        if self.fused_qkv_interleaved:
+            loaded_weight = self._deinterleave_fused_qkv(loaded_weight)
         shard_offsets = [
             # (shard_id, shard_offset, shard_size)
             ("q", 0, self.total_num_heads * self.head_size),
@@ -1173,6 +1201,8 @@ def weight_loader(
                 assert param_data.shape == loaded_weight.shape
                 param_data.copy_(loaded_weight)
                 return
+            if self.fused_qkv_interleaved:
+                loaded_weight = self._deinterleave_fused_qkv(loaded_weight)
             shard_offsets = [
                 # (shard_id, shard_offset, shard_size)
                 ("q", 0, self.total_num_heads * self.head_size),

vllm/model_executor/model_loader/bitsandbytes_loader.py

@@ -31,7 +31,7 @@
     RowParallelLinear,
 )
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
-from vllm.model_executor.model_loader.utils import ParamMapping
+from vllm.model_executor.model_loader.utils import ParamMapping, autoload_weights
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
@@ -336,9 +336,8 @@ def _unquantized_generator(
 
         global_tp_size = get_tensor_model_parallel_world_size()
         global_tp_rank = get_tensor_model_parallel_rank()
-        check_match = (
-            lambda weight_name, module_name: weight_name.removesuffix(".weight")
-            == module_name
+        check_match = lambda weight_name, module_name: (
+            weight_name.removesuffix(".weight") == module_name
         )
         for (
             org_weight_name,
@@ -524,12 +523,6 @@ def _verify_model_compatibility(
         """
         Verify that the model is compatible with BitsAndBytes quantization.
         """
-        if not hasattr(model, "load_weights"):
-            raise AttributeError(
-                "The required method 'load_weights' is not defined in class"
-                f" {type(model).__name__}."
-            )
-
         if not hasattr(model, "packed_modules_mapping"):
             raise AttributeError(
                 f"Model {type(model).__name__} does not support BitsAndBytes "
@@ -808,7 +801,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
             model_config.revision,
         )
         weights_to_load = {name for name, _ in model.named_parameters()}
-        loaded_weights = model.load_weights(qweight_iterator)
+        loaded_weights = autoload_weights(model, qweight_iterator)
         # Some models may have weights loading tracker unimplemented.
         if loaded_weights is not None:
             weights_not_loaded = weights_to_load - loaded_weights

vllm/model_executor/model_loader/default_loader.py

@@ -19,6 +19,7 @@
 from vllm.model_executor.model_loader.ep_weight_filter import (
     compute_local_expert_ids,
 )
+from vllm.model_executor.model_loader.utils import autoload_weights
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
@@ -424,7 +425,8 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
 
         self._init_ep_weight_filter(model_config)
 
-        loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
+        weights = self.get_all_weights(model_config, model)
+        loaded_weights = autoload_weights(model, weights)
 
         self.counter_after_loading_weights = time.perf_counter()
         logger.info_once(

vllm/model_executor/model_loader/runai_streamer_loader.py

@@ -10,6 +10,7 @@
 from vllm.config import ModelConfig
 from vllm.config.load import LoadConfig
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+from vllm.model_executor.model_loader.utils import autoload_weights
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
@@ -135,6 +136,5 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         model_weights = model_config.model
         if model_weights_override := model_config.model_weights:
             model_weights = model_weights_override
-        model.load_weights(
-            self._get_weights_iterator(model_weights, model_config.revision)
-        )
+        weights = self._get_weights_iterator(model_weights, model_config.revision)
+        autoload_weights(model, weights)

vllm/model_executor/model_loader/tensorizer_loader.py

@@ -20,6 +20,7 @@
     tensorizer_weights_iterator,
 )
 from vllm.model_executor.model_loader.utils import (
+    autoload_weights,
     get_model_architecture,
     initialize_model,
 )
@@ -83,7 +84,7 @@ def _load_model_serialized_cpu(
             with torch.device(device_config.device):
                 model = initialize_model(vllm_config=vllm_config, prefix=prefix)
 
-            model.load_weights(self._get_weights_iterator())
+            autoload_weights(model, self._get_weights_iterator())
         return model.eval()
 
     def download_model(self, model_config: ModelConfig) -> None:
@@ -110,7 +111,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
             tensorizer_config = self._patch_tensorizer_config(model_config)
             deserialize_tensorizer_model(model, tensorizer_config)
         else:
-            model.load_weights(self._get_weights_iterator())
+            autoload_weights(model, self._get_weights_iterator())
 
     def load_model(
         self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""

vllm/model_executor/model_loader/utils.py

@@ -4,6 +4,7 @@
 
 import inspect
 import warnings
+from collections.abc import Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Any
@@ -310,3 +311,23 @@ def configure_quant_config(
             quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_unstacked_mapper())
         if packed_mapping is not None:
             quant_config.packed_modules_mapping = packed_mapping
+
+
+def autoload_weights(
+    model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]
+) -> set[str]:
+    """Load `weights` into `model` via its `load_weights`, or AutoWeightsLoader.
+
+    Models whose loading is fully handled by `AutoWeightsLoader` (mapper as a
+    class attribute, tied lm_head auto-skipped) need not define a trivial
+    `load_weights`. This is the single entry point every caller should use so
+    such models load correctly whether or not the method exists.
+    """
+    # Imported lazily to avoid a circular import: `AutoWeightsLoader` lives in
+    # `models.utils`, which imports from `model_loader`.
+    from vllm.model_executor.models.utils import AutoWeightsLoader
+
+    model_load_weights = getattr(model, "load_weights", None)
+    if callable(model_load_weights):
+        return model_load_weights(weights)
+    return AutoWeightsLoader(model).load_weights(weights)

vllm/model_executor/models/AXK1.py

@@ -77,7 +77,6 @@
 
 from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP
 from .utils import (
-    AutoWeightsLoader,
     PPMissingLayer,
     get_spec_layer_idx_from_weight_name,
     is_pp_missing_parameter,
@@ -1140,7 +1139,3 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
             num_experts=self.config.n_routed_experts,
             num_redundant_experts=0,
         )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)

vllm/model_executor/models/adapters.py

@@ -219,7 +219,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
 
             def default_load_weights(weights):
                 loader = AutoWeightsLoader(self)
-                return loader.load_weights(weights)
+                mapper = getattr(self, "hf_to_vllm_mapper", None)
+                return loader.load_weights(weights, mapper=mapper)
 
             load_weights = getattr(super(), "load_weights", default_load_weights)
             return load_weights(mapped_weights)

vllm/model_executor/models/afmoe.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Inference-only AfMoE model compatible with HuggingFace weights."""
 
-from collections.abc import Iterable
 from itertools import islice
 
 import torch
@@ -43,7 +42,6 @@
 )
 from vllm.model_executor.models.llama import LlamaMLP as AfmoeMLP
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
     extract_layer_index,
@@ -591,7 +589,3 @@ def forward(
     def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
         logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)