| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | from transformers import PreTrainedConfig, WhisperConfig |
| 6 | |
| 7 | from vllm.logger import init_logger |
| 8 | |
| 9 | logger = init_logger(__name__) |
| 10 | |
| 11 | |
| 12 | def adapt_config_dict( |
| 13 | config_dict: dict[str, Any], |
| 14 | defaults: dict[str, Any], |
| 15 | ) -> PreTrainedConfig: |
| 16 | config_dict = _remap_general_mistral_args(config_dict) |
| 17 | config_dict = _remap_mistral_sliding_window(config_dict) |
| 18 | |
| 19 | if bool(config_dict.get("quantization")): |
| 20 | config_dict = _remap_mistral_quantization_args(config_dict) |
| 21 | |
| 22 | is_mla = bool(config_dict.get("qk_nope_head_dim")) |
| 23 | if is_mla: |
| 24 | config_dict = _remap_mistral_mla_args(config_dict) |
| 25 | |
| 26 | is_moe = bool(config_dict.get("moe")) |
| 27 | is_mistral_large_3 = ( |
| 28 | is_moe and (config_dict["moe"].get("num_shared_experts") or 0) > 0 |
| 29 | ) |
| 30 | if config_dict.get("model_type") == "mamba": |
| 31 | config_dict["architectures"] = ["Mamba2ForCausalLM"] |
| 32 | elif is_moe and is_mistral_large_3: |
| 33 | config_dict = _remap_moe_args(config_dict) |
| 34 | config_dict["model_type"] = "deepseek_v3" |
| 35 | config_dict["architectures"] = ["MistralLarge3ForCausalLM"] |
| 36 | |
| 37 | assert "llama_4_scaling" in config_dict, ( |
| 38 | "MistralLarge3 expect llama4 scaling config." |
| 39 | ) |
| 40 | llama_4_scaling_config_keys = ["original_max_position_embeddings", "beta"] |
| 41 | assert all( |
| 42 | [ |
| 43 | key in config_dict["llama_4_scaling"] |
| 44 | for key in llama_4_scaling_config_keys |
| 45 | ] |
| 46 | ), ( |
| 47 | "llama_4_scaling config should define the keys: " |
| 48 | f"{','.join(llama_4_scaling_config_keys)}" |
| 49 | ) |
| 50 | elif is_moe: |
| 51 | config_dict["architectures"] = ["MixtralForCausalLM"] |
| 52 | else: |
| 53 | config_dict["architectures"] = ["MistralForCausalLM"] |
| 54 | |
| 55 | if bool(config_dict.get("yarn")): |
| 56 | config_dict = _remap_mistral_yarn_args(config_dict) |
| 57 | |
| 58 | if bool(config_dict.get("llama_4_scaling")): |
| 59 | llama_4_scaling_config_keys = ["original_max_position_embeddings", "beta"] |
| 60 | assert all( |
| 61 | [ |
| 62 | key in config_dict["llama_4_scaling"] |
| 63 | for key in llama_4_scaling_config_keys |
| 64 | ] |
| 65 | ), ( |
| 66 | "llama_4_scaling config should define the keys: " |
| 67 | f"{','.join(llama_4_scaling_config_keys)}" |
| 68 | ) |
| 69 | |
| 70 | is_vision = (config_dict.get("multimodal") or {}).get( |
| 71 | "vision_encoder_args" |
| 72 | ) or config_dict.get("vision_encoder") |
| 73 | is_audio = bool( |
| 74 | ((config_dict.get("multimodal") or {}).get("whisper_model_args") or {}).get( |
| 75 | "encoder_args" |
| 76 | ) |
| 77 | ) |
| 78 | |
| 79 | assert not (is_vision and is_audio), "Vision and audio are mutually exclusive" |
| 80 | |
| 81 | if is_vision: |
| 82 | config_dict = _remap_mistral_vision_args(config_dict) |
| 83 | if is_audio: |
| 84 | config_dict = _remap_mistral_audio_args(config_dict) |
| 85 | |
| 86 | for k, v in defaults.items(): |
| 87 | config_dict.setdefault(k, v) |
| 88 | |
| 89 | config = PreTrainedConfig.from_dict(config_dict) |
| 90 | |
| 91 | logger.debug("Initialized config %s", config) |
| 92 | |
| 93 | return config |
| 94 | |
| 95 | |
| 96 | def _remap_mistral_vision_args(config: dict) -> dict: |
| 97 | if config.get("multimodal"): |
| 98 | vision_config = config.pop("multimodal") |
| 99 | else: |
| 100 | vision_config = config.pop("vision_encoder") |
| 101 | |
| 102 | quant_config = config.get("quantization_config") |
| 103 | config = { |
| 104 | "model_type": "pixtral", |
| 105 | "architectures": ["PixtralForConditionalGeneration"], |
| 106 | "text_config": PreTrainedConfig.from_dict(config), |
| 107 | "vision_config": PreTrainedConfig.from_dict(vision_config), |
| 108 | } |
| 109 | if quant_config: |
| 110 | config["quantization_config"] = quant_config |
| 111 | return config |
| 112 | |
| 113 | |
| 114 | def _remap_mistral_yarn_args(config: dict) -> dict: |
| 115 | yarn_config_map = { |
| 116 | "factor": ("factor", float), |
| 117 | "original_max_position_embeddings": ("original_max_position_embeddings", int), |
| 118 | "beta": ("beta_fast", float), |
| 119 | "alpha": ("beta_slow", float), |
| 120 | } |
| 121 | |
| 122 | yarn_config = config.get("yarn") or {} |
| 123 | config["rope_parameters"] = { |
| 124 | "rope_type": "yarn", |
| 125 | "mscale_all_dim": 1, |
| 126 | } |
| 127 | |
| 128 | if rope_theta := config.pop("rope_theta", None): |
| 129 | config["rope_parameters"]["rope_theta"] = rope_theta |
| 130 | |
| 131 | for old_name, (new_name, cast) in yarn_config_map.items(): |
| 132 | if old_name in yarn_config: |
| 133 | # Cast to remove Transformers > v5 type warnings |
| 134 | config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name)) |
| 135 | |
| 136 | # `apply_scale: false` means no magnitude correction, which Transformers |
| 137 | # spells as an explicit attention_factor of 1. |
| 138 | if not yarn_config.pop("apply_scale", True): |
| 139 | config["rope_parameters"]["attention_factor"] = 1.0 |
| 140 | |
| 141 | assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}" |
| 142 | |
| 143 | return config |
| 144 | |
| 145 | |
| 146 | def _remap_general_mistral_args(config: dict) -> dict: |
| 147 | # Mistral key -> HF key |
| 148 | config_mapping = { |
| 149 | "dim": "hidden_size", |
| 150 | "norm_eps": "rms_norm_eps", |
| 151 | "n_kv_heads": "num_key_value_heads", |
| 152 | "n_layers": "num_hidden_layers", |
| 153 | "n_heads": "num_attention_heads", |
| 154 | "hidden_dim": "intermediate_size", |
| 155 | } |
| 156 | # HF key -> (Mistral key, default value) |
| 157 | top_level_mapping_with_default = { |
| 158 | "model_type": ("model_type", "transformer"), |
| 159 | "hidden_act": ("activation", "silu"), |
| 160 | "tie_word_embeddings": ("tied_embeddings", False), |
| 161 | "max_seq_len": ("max_seq_len", config.get("max_position_embeddings", 128_000)), |
| 162 | "max_position_embeddings": ("max_position_embeddings", 128_000), |
| 163 | "dtype": ("dtype", config.get("dtype")), |
| 164 | } |
| 165 | |
| 166 | for key, new_key in config_mapping.items(): |
| 167 | if key in config: |
| 168 | config[new_key] = config.pop(key) |
| 169 | |
| 170 | for new_key, (key, default_value) in top_level_mapping_with_default.items(): |
| 171 | config[new_key] = config.pop(key, default_value) |
| 172 | |
| 173 | return config |
| 174 | |
| 175 | |
| 176 | def _remap_mistral_sliding_window(config: dict) -> dict: |
| 177 | # Remap sliding_window (list) -> layer_types (list) + sliding window (int) |
| 178 | # for HF compatibility |
| 179 | # Mistral configs may define sliding_window as list[int]. Convert it |
| 180 | # to int and add the layer_types list[str] to make it HF compatible |
| 181 | if sliding_window := config.get("sliding_window"): |
| 182 | if isinstance(sliding_window, list): |
| 183 | pattern_repeats = config["num_hidden_layers"] // len(sliding_window) |
| 184 | layer_types = sliding_window * pattern_repeats |
| 185 | config["layer_types"] = [ |
| 186 | "full_attention" if layer_type is None else "sliding_attention" |
| 187 | for layer_type in layer_types |
| 188 | ] |
| 189 | assert len(set(sliding_window) - {None}) <= 1, sliding_window |
| 190 | config["sliding_window"] = next(filter(None, sliding_window), None) |
| 191 | elif isinstance(sliding_window, int) and config.get("layer_types") is None: |
| 192 | config["layer_types"] = ["sliding_attention"] * config["num_hidden_layers"] |
| 193 | else: |
| 194 | raise ValueError(f"Unsupported sliding_window type: {sliding_window}") |
| 195 | |
| 196 | return config |
| 197 | |
| 198 | |
| 199 | def _remap_mistral_quantization_args(config: dict) -> dict: |
| 200 | if config.get("quantization"): |
| 201 | quantization = config.pop("quantization", {}) |
| 202 | if quantization.get("qformat_weight") == "fp8_e4m3": |
| 203 | qscheme_act = quantization.get("qscheme_act") |
| 204 | assert qscheme_act in ("NO_SCALES", "TENSOR", None), ( |
| 205 | "Only NO_SCALES and TENSOR (default) are supported for qscheme_act" |
| 206 | ) |
| 207 | is_dynamic = qscheme_act == "NO_SCALES" |
| 208 | config["quantization_config"] = { |
| 209 | "quant_method": "fp8", |
| 210 | "activation_scheme": "dynamic" if is_dynamic else "static", |
| 211 | } |
| 212 | elif ( |
| 213 | str(quantization.get("quant_method", "")).lower().replace("_", "-") |
| 214 | == "compressed-tensors" |
| 215 | ): |
| 216 | # Pass through compressed-tensors config, while normalizing |
| 217 | # quant_method to the canonical community spelling. |
| 218 | quantization["quant_method"] = "compressed-tensors" |
| 219 | config["quantization_config"] = quantization |
| 220 | else: |
| 221 | raise ValueError(f"Found unknown quantization='{quantization}' in config") |
| 222 | |
| 223 | return config |
| 224 | |
| 225 | |
| 226 | def _remap_mistral_audio_args(config: dict) -> dict: |
| 227 | whisper_args = config["multimodal"].pop("whisper_model_args") |
| 228 | encoder_args = whisper_args["encoder_args"] |
| 229 | downsample_args = whisper_args["downsample_args"] |
| 230 | downsample_factor = downsample_args["downsample_factor"] |
| 231 | |
| 232 | # make sure that k/v blocks can be allocated with |
| 233 | # unified k/v cache class and pool whisper k/v cache blocks |
| 234 | # with downsample_factor:1 ratio |
| 235 | if encoder_args.get("causal"): |
| 236 | block_pool_size = downsample_factor |
| 237 | config["projection_size"] = downsample_factor * encoder_args["dim"] |
| 238 | else: |
| 239 | block_pool_size = 1 |
| 240 | |
| 241 | architecture = ( |
| 242 | "VoxtralRealtimeGeneration" |
| 243 | if encoder_args.get("causal") |
| 244 | else "VoxtralForConditionalGeneration" |
| 245 | ) |
| 246 | |
| 247 | quant_config = config.get("quantization_config") |
| 248 | config = { |
| 249 | "model_type": "voxtral", |
| 250 | "architectures": [architecture], |
| 251 | "text_config": PreTrainedConfig.from_dict(config), |
| 252 | "audio_config": WhisperConfig( |
| 253 | num_mel_bins=encoder_args["audio_encoding_args"]["num_mel_bins"], |
| 254 | window_size=encoder_args["audio_encoding_args"]["window_size"], |
| 255 | sampling_rate=encoder_args["audio_encoding_args"]["sampling_rate"], |
| 256 | hop_length=encoder_args["audio_encoding_args"]["hop_length"], |
| 257 | downsample_factor=downsample_factor, |
| 258 | d_model=encoder_args["dim"], |
| 259 | encoder_layers=encoder_args["n_layers"], |
| 260 | encoder_ffn_dim=encoder_args["hidden_dim"], |
| 261 | encoder_attention_heads=encoder_args["n_heads"], |
| 262 | encoder_head_dim=encoder_args["head_dim"], |
| 263 | vocab_size=encoder_args["vocab_size"], |
| 264 | is_encoder_decoder=False, # Override WhisperConfig default |
| 265 | is_causal=encoder_args.get("causal", False), |
| 266 | sliding_window=encoder_args.get("sliding_window", None), |
| 267 | block_pool_size=block_pool_size, |
| 268 | pos_embed=encoder_args.get("pos_embed", "sinusoidal"), |
| 269 | global_log_mel_max=encoder_args["audio_encoding_args"].get( |
| 270 | "global_log_mel_max" |
| 271 | ), |
| 272 | # only needed for RoPE |
| 273 | max_position_embeddings=block_pool_size * config["max_position_embeddings"], |
| 274 | ), |
| 275 | } |
| 276 | # Sometimes max_source_positions is explicitly set to None in params.json but this |
| 277 | # is not a valid value for WhisperConfig (or downstream code that uses it). |
| 278 | if (max_source_positions := encoder_args.get("max_source_positions")) is not None: |
| 279 | config["audio_config"].max_source_positions = max_source_positions |
| 280 | if quant_config: |
| 281 | config["quantization_config"] = quant_config |
| 282 | return config |
| 283 | |
| 284 | |
| 285 | def _remap_moe_args(config: dict) -> dict: |
| 286 | moe_config_map = { |
| 287 | "route_every_n": "moe_layer_freq", |
| 288 | "first_k_dense_replace": "first_k_dense_replace", |
| 289 | "num_experts_per_tok": "num_experts_per_tok", |
| 290 | "num_experts": "n_routed_experts", |
| 291 | "expert_hidden_dim": "moe_intermediate_size", |
| 292 | "routed_scale": "routed_scaling_factor", |
| 293 | "num_shared_experts": "n_shared_experts", |
| 294 | "num_expert_groups": "n_group", |
| 295 | "num_expert_groups_per_tok": "topk_group", |
| 296 | } |
| 297 | moe_config = config.get("moe", {}) |
| 298 | for old_name, new_name in moe_config_map.items(): |
| 299 | if old_name in moe_config: |
| 300 | value = moe_config.pop(old_name) |
| 301 | config[new_name] = value |
| 302 | |
| 303 | config["topk_method"] = None |
| 304 | config["norm_topk_prob"] = True |
| 305 | config["scoring_func"] = "softmax" |
| 306 | |
| 307 | return config |
| 308 | |
| 309 | |
| 310 | def _remap_mistral_mla_args(config: dict) -> dict: |
| 311 | if not config.get("moe"): |
| 312 | moe = { |
| 313 | "num_experts": 1, |
| 314 | "first_k_dense_replace": config.get("num_hidden_layers"), |
| 315 | "route_every_n": 1, |
| 316 | "num_shared_experts": 1, |
| 317 | "expert_hidden_dim": config.get("intermediate_size"), |
| 318 | "num_experts_per_tok": 1, |
| 319 | "routed_scale": 1.0, |
| 320 | "renorm_strategy": "WEIGHTS", |
| 321 | "use_load_balancing_bias": False, |
| 322 | "num_expert_groups": 1, |
| 323 | "num_expert_groups_per_tok": 1, |
| 324 | } |
| 325 | config["moe"] = moe |
| 326 | return config |