[Misc] Clean up Kimi-audio whisper encoder loading

vllm/model_executor/model_loader/default_loader.py

@@ -52,6 +52,9 @@ class Source:
         revision: str | None
         """The optional model revision."""
 
+        subfolder: str | None = None
+        """The subfolder inside the model repo."""
+
         prefix: str = ""
         """A prefix to prepend to all weights."""
 
@@ -81,6 +84,7 @@ def __init__(self, load_config: LoadConfig):
     def _prepare_weights(
         self,
         model_name_or_path: str,
+        subfolder: str | None,
         revision: str | None,
         fall_back_to_pt: bool,
         allow_patterns_overrides: list[str] | None,
@@ -97,6 +101,9 @@ def _prepare_weights(
         load_format = self.load_config.load_format
         use_safetensors = False
         index_file = SAFE_WEIGHTS_INDEX_NAME
+        index_file_with_subfolder = (
+            f"{subfolder}/{index_file}" if subfolder else index_file
+        )
 
         # First check for 'auto' format that mistral files format are present.
         # This is to load mistral models with official format by default.
@@ -148,6 +155,9 @@ def _prepare_weights(
         else:
             hf_folder = model_name_or_path
 
+        if subfolder is not None:
+            hf_folder = os.path.join(hf_folder, subfolder)
+
         hf_weights_files: list[str] = []
         for pattern in allow_patterns:
             hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
@@ -165,7 +175,7 @@ def _prepare_weights(
             if not is_local:
                 download_safetensors_index_file_from_hf(
                     model_name_or_path,
-                    index_file,
+                    index_file_with_subfolder,
                     self.load_config.download_dir,
                     revision,
                 )
@@ -189,6 +199,7 @@ def _get_weights_iterator(
         extra_config = self.load_config.model_loader_extra_config
         hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
             source.model_or_path,
+            source.subfolder,
             source.revision,
             source.fall_back_to_pt,
             source.allow_patterns_overrides,
@@ -269,8 +280,9 @@ def get_all_weights(
 
     def download_model(self, model_config: ModelConfig) -> None:
         self._prepare_weights(
-            model_config.model,
-            model_config.revision,
+            model_name_or_path=model_config.model,
+            subfolder=None,
+            revision=model_config.revision,
             fall_back_to_pt=True,
             allow_patterns_overrides=None,
         )

vllm/model_executor/models/kimi_audio.py

@@ -3,25 +3,21 @@
 
 """Inference-only Kimi-Audio model compatible with HuggingFace weights."""
 
-import os
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, ClassVar, Literal
 
 import numpy as np
 import torch
 import torch.nn as nn
-from huggingface_hub import snapshot_download
-from safetensors import safe_open
 from transformers import BatchFeature
 from transformers import WhisperConfig as HFWhisperConfig
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.inputs.data import PromptType, TokensPrompt
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-)
+from vllm.model_executor.model_loader import DefaultModelLoader
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (
     SupportsMultiModal,
     SupportsPP,
@@ -64,15 +60,6 @@
 KIMIA_WHISPER_SUBFOLDER = "whisper-large-v3"
 
 
-def _get_whisper_local_path(repo_id: str):
-    if os.path.exists(repo_id):
-        repo_local_path = repo_id
-    else:
-        repo_local_path = snapshot_download(repo_id, local_files_only=True)
-
-    return os.path.join(repo_local_path, KIMIA_WHISPER_SUBFOLDER)
-
-
 def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor:
     """Compute output lengths after Whisper feature extraction.
 
@@ -93,7 +80,6 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
     # packed_modules_mapping for Q/K/V fusion during weight loading
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
-        "kv_proj": ["k_proj", "v_proj"],
     }
 
     def __init__(
@@ -104,19 +90,49 @@ def __init__(
         model_path = vllm_config.model_config.model
 
         # Load WhisperConfig from the subfolder
-        whisper_dir = _get_whisper_local_path(model_path)
-        whisper_config = HFWhisperConfig.from_pretrained(whisper_dir)
-
-        # Temporarily replace hf_config for WhisperEncoder.__init__()
-        original_config = vllm_config.model_config.hf_config
-        vllm_config.model_config.hf_config = whisper_config
+        whisper_config = HFWhisperConfig.from_pretrained(
+            model_path,
+            subfolder=KIMIA_WHISPER_SUBFOLDER,
+        )
 
         super().__init__(
-            vllm_config=vllm_config, prefix=prefix, init_in_fp32=init_in_fp32
+            vllm_config=vllm_config.with_hf_config(whisper_config),
+            prefix=prefix,
+            init_in_fp32=init_in_fp32,
         )
 
-        # Restore original config
-        vllm_config.model_config.hf_config = original_config
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
 
 
 # -----------------------------------------------------------------------------
@@ -374,6 +390,8 @@ class KimiAudioForConditionalGeneration(
 
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
+            # audio tower
+            "model.encoder.": "audio_tower.",
             # Audio projector (VQ-Adaptor)
             "model.vq_adaptor.layers.0.": "multi_modal_projector.vq_adaptor_layers_0.",
             "model.vq_adaptor.layers.3.": "multi_modal_projector.vq_adaptor_layers_3.",
@@ -384,7 +402,11 @@ class KimiAudioForConditionalGeneration(
             "model.embed_tokens.": "language_model.model.embed_tokens.",
             "model.norm.": "language_model.model.norm.",
             "lm_head.": "language_model.lm_head.",
-        }
+        },
+        orig_to_new_substr={
+            ".fc1.": ".mlp.fc1.",
+            ".fc2.": ".mlp.fc2.",
+        },
     )
 
     # Audio placeholder token sequence
@@ -401,6 +423,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.multimodal_config = vllm_config.model_config.multimodal_config
         self.model_path = vllm_config.model_config.model
 
+        self.secondary_weights = [
+            DefaultModelLoader.Source(
+                model_or_path=vllm_config.model_config.model,
+                subfolder="whisper-large-v3",
+                revision=None,
+            )
+        ]
+
         self.audio_tower = KimiAudioWhisperEncoder(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "audio_tower"),
@@ -577,99 +607,19 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights, skipping MIMO layers (TTS-only) for ASR."""
         # Filter out MIMO/TTS weights since we only do ASR (speech-to-text)
         skipped_patterns = [
+            # Audio tower
+            "model.",
+            # MIMO/TTS
             "mimo_layers.",
             "mimo_output.",
             "mimo_norm.",
-            "audio_decoder.",
-        ]
-
-        # Filter weights
-        filtered_weights = [
-            (name, param)
-            for name, param in weights
-            if not any(pattern in name for pattern in skipped_patterns)
-        ]
-
-        # Separate main weights (non-Whisper) from Whisper weights
-        main_weights = [
-            (name, param)
-            for name, param in filtered_weights
-            if not name.startswith("audio_tower.")
         ]
 
         # Load main model weights (LLM + projector) with mapper
-        loader = AutoWeightsLoader(self)
-        loaded = loader.load_weights(main_weights, mapper=self.hf_to_vllm_mapper)
-
-        # Load Whisper encoder weights from subfolder
-        whisper_dir = _get_whisper_local_path(self.model_path)
-        whisper_path = os.path.join(whisper_dir, "model.safetensors")
-        if os.path.exists(whisper_path):
-            whisper_loaded = self._load_whisper_weights_from_file(whisper_path)
-            loaded.update(whisper_loaded)
-
+        loader = AutoWeightsLoader(self, skip_prefixes=skipped_patterns)
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         return loaded
 
-    def _load_whisper_weights_from_file(self, whisper_path: str) -> set[str]:
-        """Load Whisper encoder weights from safetensors file with transformations."""
-        if not os.path.exists(whisper_path):
-            return set()
-
-        # Step 1: Load raw weights from safetensors file
-        whisper_weights = []
-        with safe_open(whisper_path, framework="pt") as f:
-            for key in f.keys():  # noqa: SIM118
-                if key.startswith("model.encoder.") and "embed_positions" not in key:
-                    new_key = key.replace("model.encoder.", "")
-                    whisper_weights.append((new_key, f.get_tensor(key)))
-
-        # Step 2: Apply fc → mlp mapping using WeightsMapper
-        fc_mapper = WeightsMapper(
-            orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}
-        )
-        whisper_mapped = list(fc_mapper.apply(whisper_weights))
-
-        # Step 3: Apply Q/K/V fusion manually
-        stacked_params_mapping = [
-            (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
-            (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
-            (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
-        ]
-
-        params_dict = dict(self.audio_tower.named_parameters())
-        whisper_loaded: set[str] = set()
-
-        for name, loaded_weight in whisper_mapped:
-            fused = False
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                fused_name = name.replace(weight_name, param_name)
-                if fused_name not in params_dict:
-                    continue
-
-                param = params_dict[fused_name]
-                param.weight_loader(param, loaded_weight, shard_id)
-                whisper_loaded.add(f"audio_tower.{fused_name}")
-                fused = True
-                break
-
-            if not fused:
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-                whisper_loaded.add(f"audio_tower.{name}")
-
-        # Add embed_positions which is initialized randomly
-        whisper_loaded.add("audio_tower.embed_positions.weight")
-
-        return whisper_loaded
-
     @classmethod
     def get_speech_to_text_config(
         cls, model_config: ModelConfig, task_type: str