main
py 848 lines 33.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import final
5
6 import torch
7 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
8 from transformers import PreTrainedConfig
9
10 from vllm import envs
11 from vllm.config.model_arch import (
12 ModelArchitectureConfig,
13 )
14 from vllm.config.utils import getattr_iter
15 from vllm.logger import init_logger
16 from vllm.transformers_utils.config import (
17 ConfigFormat,
18 get_safetensors_params_metadata,
19 )
20 from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
21 from vllm.utils.torch_utils import common_broadcastable_dtype
22
23 logger = init_logger(__name__)
24
25
26 class ModelArchConfigConvertorBase:
27 def __init__(
28 self,
29 hf_config: PreTrainedConfig,
30 hf_text_config: PreTrainedConfig,
31 revision: str | None = None,
32 ):
33 self.hf_config = hf_config
34 self.hf_text_config = hf_text_config
35 self.revision = revision
36
37 def get_per_layer_hf_configs(
38 self,
39 ) -> list[tuple[PreTrainedConfig, PreTrainedConfig]] | None:
40 """`(hf_config, hf_text_config)` per layer, or `None` if homogeneous.
41
42 This is the only place that decides whether a checkpoint is heterogeneous
43 and how many layers it has, so a convertor can declare variation that
44 Transformers does not express (see `Gemma4ModelArchConfigConvertor`).
45
46 Each pair is converted independently and the results are diffed, so a
47 convertor that overrides this must not also collapse a varying field in
48 its getters: the collapse would then be applied per layer and every layer
49 would report the same value.
50 """
51 config_het = getattr(self.hf_config, "is_heterogeneous", False)
52 text_het = getattr(self.hf_text_config, "is_heterogeneous", False)
53 if not (config_het or text_het):
54 return None
55 # The text config decides the layer count when it is the heterogeneous one:
56 # its stack is what vLLM builds attention layers for.
57 source = self.hf_text_config if text_het else self.hf_config
58 return [
59 (
60 self.hf_config.per_layer_config[i] if config_het else self.hf_config,
61 self.hf_text_config.per_layer_config[i]
62 if text_het
63 else self.hf_text_config,
64 )
65 for i in range(len(source.per_layer_config))
66 ]
67
68 def get_architectures(self) -> list[str]:
69 # Sometimes we get here from `vllm_config.with_hf_config(text_config)` where
70 # `text_config` is a sub-config from a multi-modal model. If this is the case,
71 # the sub-config will not have `architectures` and it will explicitly be `None`
72 return getattr(self.hf_config, "architectures", None) or []
73
74 def get_num_hidden_layers(self) -> int:
75 return getattr(self.hf_text_config, "num_hidden_layers", 0)
76
77 def get_total_num_attention_heads(self) -> int:
78 return getattr(self.hf_text_config, "num_attention_heads", 0)
79
80 def get_vocab_size(self) -> int:
81 return getattr(self.hf_text_config, "vocab_size", 0)
82
83 def get_hidden_size(self) -> int:
84 return getattr(self.hf_text_config, "hidden_size", 0)
85
86 def get_head_size(self) -> int:
87 if self.is_deepseek_mla():
88 # special case for deepseek_v4
89 if hasattr(self.hf_text_config, "compress_ratios"):
90 return self.hf_text_config.head_dim
91 qk_rope_head_dim = self._get_qk_rope_head_dim()
92 if not envs.VLLM_MLA_DISABLE:
93 return self.hf_text_config.kv_lora_rank + qk_rope_head_dim
94 else:
95 qk_nope_head_dim = getattr(self.hf_text_config, "qk_nope_head_dim", 0)
96 if qk_rope_head_dim and qk_nope_head_dim:
97 return qk_rope_head_dim + qk_nope_head_dim
98
99 # NOTE: Some config classes may set head_dim=None or materialize a missing
100 # head_dim as 0 (for example, DeepseekVLV2TextConfig).
101 if (
102 head_dim := getattr(self.hf_text_config, "head_dim", None)
103 ) is not None and head_dim > 0:
104 return head_dim
105
106 # NOTE: Some models (such as PLaMo2.1) use `hidden_size_per_head`
107 if getattr(self.hf_text_config, "hidden_size_per_head", None) is not None:
108 return self.hf_text_config.hidden_size_per_head
109
110 if (total_num_attention_heads := self.get_total_num_attention_heads()) == 0:
111 return 0
112 # FIXME(woosuk): This may not be true for all models.
113 return self.get_hidden_size() // total_num_attention_heads
114
115 def _get_qk_rope_head_dim(self) -> int:
116 """Get qk_rope_head_dim, fixing the transformers v5.4+ attribute_map bug."""
117 cfg = self.hf_text_config
118 qk_rope_head_dim = getattr(cfg, "qk_rope_head_dim", 0)
119 qk_nope_head_dim = getattr(cfg, "qk_nope_head_dim", 0)
120
121 # In valid MLA configs, qk_rope_head_dim != qk_nope_head_dim.
122 if qk_rope_head_dim == 0 or qk_rope_head_dim != qk_nope_head_dim:
123 return qk_rope_head_dim # not corrupted
124
125 # Read the correct value from raw config.json.
126 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
127
128 model_path = self.hf_config.name_or_path
129 if not model_path:
130 return qk_rope_head_dim
131 raw = get_hf_file_to_dict("config.json", model_path, self.revision)
132 if raw and "qk_rope_head_dim" in raw:
133 correct = raw["qk_rope_head_dim"]
134 if correct != qk_rope_head_dim:
135 logger.info(
136 "Fixing qk_rope_head_dim: %d -> %d "
137 "(transformers v5.4+ attribute_map bug)",
138 qk_rope_head_dim,
139 correct,
140 )
141 # Patch the config so downstream model layers also get
142 # the correct value.
143 cfg.qk_rope_head_dim = correct
144 return correct
145 return qk_rope_head_dim
146
147 def get_total_num_kv_heads(self) -> int:
148 attributes = [
149 # For Falcon:
150 "n_head_kv",
151 "num_kv_heads",
152 # For LLaMA-2:
153 "num_key_value_heads",
154 # For ChatGLM:
155 "multi_query_group_num",
156 # For Step3p5:
157 "num_attention_groups",
158 ]
159 # For non-grouped-query attention models, the number of KV heads is
160 # equal to the number of attention heads.
161 default_factory = self.get_total_num_attention_heads
162 return getattr_iter(
163 self.hf_text_config, attributes, default_factory=default_factory
164 )
165
166 def get_num_experts_from_block_configs(self) -> int:
167 """Check block_configs for heterogeneous models (e.g., NemotronH).
168
169 For heterogeneous models with varying expert counts per layer,
170 returns the MAX to ensure all expert weights can be loaded.
171 """
172 max_experts = 0
173 block_configs = getattr(self.hf_text_config, "block_configs", None)
174 if block_configs:
175 for block in block_configs:
176 if isinstance(block, dict):
177 if block.get("block_type", "") == "moe":
178 max_experts = max(max_experts, block.get("n_routed_experts", 0))
179 else:
180 if getattr(block, "block_type", "") == "moe":
181 max_experts = max(
182 max_experts, getattr(block, "n_routed_experts", 0)
183 )
184 return max_experts
185
186 def get_num_experts(self) -> int:
187 """Returns the number of experts in the model."""
188 num_expert_names = [
189 "num_experts", # Jamba
190 "moe_num_experts", # Dbrx
191 "n_routed_experts", # DeepSeek
192 "num_local_experts", # Mixtral
193 ]
194
195 num_experts = getattr_iter(self.hf_text_config, num_expert_names, 0)
196 if isinstance(num_experts, list):
197 # Ernie VL's remote code uses list[int]...
198 # The values are always the same so we just take the first one.
199 return num_experts[0]
200
201 if not num_experts:
202 num_experts = self.get_num_experts_from_block_configs()
203 return num_experts
204
205 def get_num_experts_per_token(self) -> int:
206 names = [
207 "num_experts_per_tok",
208 "num_experts_per_token",
209 "top_k_experts",
210 "moe_topk",
211 "moe_top_k",
212 ]
213 num_experts_per_token = getattr_iter(self.hf_text_config, names, 0)
214 if isinstance(num_experts_per_token, list):
215 return max(num_experts_per_token, default=0)
216 return num_experts_per_token or 0
217
218 @final
219 @classmethod
220 def get_torch_dtype(
221 cls,
222 hf_config: PreTrainedConfig,
223 model_id: str,
224 revision: str | None,
225 config_format: str | ConfigFormat,
226 ):
227 # NOTE: getattr(config, "dtype", torch.float32) is not correct
228 # because config.dtype can be None.
229 config_dtype = getattr(hf_config, "dtype", None)
230
231 # Fallbacks for multi-modal models if the root config
232 # does not define dtype
233 if config_dtype is None:
234 config_dtype = getattr(hf_config.get_text_config(), "dtype", None)
235 if config_dtype is None and hasattr(hf_config, "vision_config"):
236 config_dtype = getattr(hf_config.vision_config, "dtype", None)
237 if config_dtype is None and hasattr(hf_config, "encoder_config"):
238 config_dtype = getattr(hf_config.encoder_config, "dtype", None)
239
240 # Try to read the dtype of the weights if they are in safetensors format
241 if config_dtype is None:
242 param_mt = get_safetensors_params_metadata(model_id, revision=revision)
243
244 if param_mt:
245 param_dtypes: set[torch.dtype] = {
246 _SAFETENSORS_TO_TORCH_DTYPE[dtype]
247 for info in param_mt.values()
248 if (dtype := info.get("dtype", None))
249 and dtype in _SAFETENSORS_TO_TORCH_DTYPE
250 }
251
252 if param_dtypes:
253 return common_broadcastable_dtype(param_dtypes)
254
255 if config_dtype is None:
256 config_dtype = torch.float32
257
258 return config_dtype
259
260 def _normalize_quantization_config(self, config: PreTrainedConfig):
261 quant_cfg = getattr(config, "quantization_config", None)
262 if quant_cfg is None:
263 # compressed-tensors uses a "compression_config" key
264 quant_cfg = getattr(config, "compression_config", None)
265
266 else:
267 # Set quant_method for ModelOpt models.
268 producer_name = quant_cfg.get("producer", {}).get("name")
269 modelopt_quant_cfg = quant_cfg.get("quantization", {})
270 is_legacy_modelopt = (
271 isinstance(modelopt_quant_cfg, dict)
272 and "modelopt_quant_config" in modelopt_quant_cfg
273 )
274 if producer_name == "modelopt" or is_legacy_modelopt:
275 quant_algo = modelopt_quant_cfg.get("quant_algo")
276 if quant_algo is not None:
277 quant_algo_upper = str(quant_algo).upper()
278 if quant_algo_upper in {
279 "FP8",
280 "FP8_PER_CHANNEL_PER_TOKEN",
281 "FP8_PB_WO",
282 }:
283 quant_cfg["quant_method"] = "modelopt"
284 elif quant_algo_upper == "NVFP4":
285 quant_cfg["quant_method"] = "modelopt_fp4"
286 else:
287 raise ValueError(f"Unknown ModelOpt quant algo: {quant_algo}")
288
289 if quant_cfg is not None:
290 # Use the community standard 'quant_method'
291 quant_method = quant_cfg.get("quant_method", "").lower()
292
293 # Normalize library names
294 quant_method = quant_method.replace(
295 "compressed_tensors", "compressed-tensors"
296 )
297
298 quant_cfg["quant_method"] = quant_method
299
300 return quant_cfg
301
302 def get_quantization_config(self):
303 quant_cfg = self._normalize_quantization_config(self.hf_config)
304 if quant_cfg is None and (
305 text_config := getattr(self.hf_config, "text_config", None)
306 ):
307 # Check the text config as well for multi-modal models.
308 quant_cfg = self._normalize_quantization_config(text_config)
309 return quant_cfg
310
311 def is_deepseek_mla(self) -> bool:
312 if not hasattr(self.hf_text_config, "model_type"):
313 return False
314 elif self.hf_text_config.model_type in (
315 "axk1",
316 "deepseek_v2",
317 "deepseek_v3",
318 "deepseek_v32",
319 "deepseek_v4",
320 "deepseek_v41",
321 "dots3_note",
322 "deepseek_mtp",
323 "k3_dspark",
324 "glm5_next",
325 "glm5_next_text",
326 "glm_moe_dsa",
327 "glm4_moe_lite",
328 "glm4_moe_lite_mtp",
329 "kimi_k2",
330 "kimi_linear",
331 "longcat_flash",
332 "longcat_flash_ngram",
333 "pangu_ultra_moe",
334 "pangu_ultra_moe_mtp",
335 "bailing_hybrid",
336 "bailing_hybrid_mtp",
337 "bailing_hybrid_v3_mtp",
338 "hy_v4",
339 "hy_v4_mtp",
340 ):
341 # check is deepseek_v4 model
342 if hasattr(self.hf_text_config, "compress_ratios"):
343 return getattr(self.hf_text_config, "head_dim", None) is not None
344 else:
345 return getattr(self.hf_text_config, "kv_lora_rank", None) is not None
346 elif self.hf_text_config.model_type == "eagle":
347 # if the model is an EAGLE module, check for the
348 # underlying architecture
349 return (
350 self.hf_text_config.model.model_type
351 in (
352 "axk1",
353 "deepseek_v2",
354 "deepseek_v3",
355 "deepseek_v32",
356 "deepseek_mtp",
357 )
358 and getattr(self.hf_text_config, "kv_lora_rank", None) is not None
359 )
360 return False
361
362 def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
363 """Whether to use bidirectional attention for mm positions.
364
365 ``supports_multimodal`` is False when the deployment is configuration-
366 disabled for multimodal inputs (text-only serving). In that case
367 mm_prefix is unnecessary and must stay off so attention backends
368 without ``supports_mm_prefix()`` remain eligible.
369 """
370 if not supports_multimodal:
371 return False
372 if hasattr(self.hf_config, "is_mm_prefix_lm"):
373 return bool(self.hf_config.is_mm_prefix_lm)
374 # fallback to list of known models
375 MM_PREFIX_LM_MODELS = (
376 "bagel",
377 "gemma3",
378 "molmo2",
379 "moondream3",
380 "paligemma",
381 )
382 if not hasattr(self.hf_config, "model_type"):
383 return False
384 return self.hf_config.model_type in MM_PREFIX_LM_MODELS
385
386 def rswa_window(self) -> int | None:
387 value = getattr(self.hf_config, "rswa_window", None)
388 if value is None:
389 return None
390 return int(value)
391
392 def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
393 derived_max_model_len = float("inf")
394 possible_keys = [
395 # OPT
396 "max_position_embeddings",
397 # GPT-2
398 "n_positions",
399 # MPT
400 "max_seq_len",
401 # ChatGLM2
402 "seq_length",
403 # Command-R
404 "model_max_length",
405 # Whisper
406 "max_target_positions",
407 # Others
408 "max_sequence_length",
409 "max_seq_length",
410 "seq_len",
411 ]
412 # Choose the smallest "max_length" from the possible keys
413 max_len_key = None
414 for key in possible_keys:
415 max_len = getattr(self.hf_text_config, key, None)
416 if max_len is not None:
417 if max_len < derived_max_model_len:
418 max_len_key = key
419 derived_max_model_len = min(derived_max_model_len, max_len)
420
421 # For Command-R / Cohere, Cohere2 models
422 if tmp_max_len := getattr(self.hf_text_config, "model_max_length", None):
423 max_len_key = "model_max_length"
424 derived_max_model_len = tmp_max_len
425 return derived_max_model_len, max_len_key
426
427 def convert(self, supports_multimodal: bool = True) -> ModelArchitectureConfig:
428 if (per_layer := self.get_per_layer_hf_configs()) is None:
429 return self.convert_layer(supports_multimodal)
430
431 if self.is_deepseek_mla():
432 raise NotImplementedError(
433 "Heterogeneous MLA models are not supported: `get_head_size` "
434 "patches the config in place for them, which a per-layer "
435 "conversion would apply to a throwaway copy."
436 )
437 # Convert each layer and let the result work out which fields differ.
438 # Reading a varying attribute off the global config would raise, so the
439 # whole-model config has to be built up from the layers.
440 return ModelArchitectureConfig.from_layers(
441 [
442 type(self)(*configs).convert_layer(supports_multimodal)
443 for configs in per_layer
444 ]
445 )
446
447 def convert_layer(
448 self, supports_multimodal: bool = True
449 ) -> ModelArchitectureConfig:
450 """Convert one homogeneous config, without resolving per-layer values.
451
452 `convert` calls this once per layer, so it must not recurse back into
453 per-layer resolution: the configs it receives are already layer-specific,
454 and they can still carry the attributes that made the model look
455 heterogeneous in the first place.
456 """
457 model_arch_config = ModelArchitectureConfig(
458 architectures=self.get_architectures(),
459 model_type=self.hf_config.model_type,
460 text_model_type=getattr(self.hf_text_config, "model_type", None),
461 hidden_size=self.get_hidden_size(),
462 total_num_hidden_layers=self.get_num_hidden_layers(),
463 total_num_attention_heads=self.get_total_num_attention_heads(),
464 head_size=self.get_head_size(),
465 vocab_size=self.get_vocab_size(),
466 total_num_kv_heads=self.get_total_num_kv_heads(),
467 num_experts=self.get_num_experts(),
468 num_experts_per_token=self.get_num_experts_per_token(),
469 quantization_config=self.get_quantization_config(),
470 is_deepseek_mla=self.is_deepseek_mla(),
471 is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal),
472 rswa_window=self.rswa_window(),
473 derived_max_model_len_and_key=self.derive_max_model_len_and_key(),
474 )
475
476 return model_arch_config
477
478
479 class CohereAsrModelArchConfigConvertor(ModelArchConfigConvertorBase):
480 def get_total_num_attention_heads(self) -> int:
481 return self.hf_text_config.transf_decoder["config_dict"]["num_attention_heads"]
482
483 def get_head_size(self) -> int:
484 hidden_size = self.hf_text_config.transf_decoder["config_dict"]["hidden_size"]
485 num_attention_heads = self.hf_text_config.transf_decoder["config_dict"][
486 "num_attention_heads"
487 ]
488 return hidden_size // num_attention_heads
489
490 def get_total_num_kv_heads(self) -> int:
491 enc_num_kv_heads = self.hf_text_config.encoder["n_heads"]
492 dec_num_kv_heads = self.hf_text_config.transf_decoder["config_dict"][
493 "num_attention_heads"
494 ]
495 assert enc_num_kv_heads == dec_num_kv_heads, (
496 "Encoder and decoder must have the same number of kv heads"
497 )
498 return enc_num_kv_heads
499
500 def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
501 return False
502
503
504 class MambaModelArchConfigConvertor(ModelArchConfigConvertorBase):
505 def get_head_size(self) -> int:
506 return 0
507
508 def get_total_num_kv_heads(self) -> int:
509 return 0
510
511
512 class TerratorchModelArchConfigConvertor(ModelArchConfigConvertorBase):
513 def get_head_size(self) -> int:
514 return 0
515
516 def get_total_num_kv_heads(self) -> int:
517 return 0
518
519
520 class MedusaModelArchConfigConvertor(ModelArchConfigConvertorBase):
521 def get_head_size(self) -> int:
522 return 0
523
524 def get_total_num_kv_heads(self) -> int:
525 return 0
526
527
528 class Zamba2ModelArchConfigConvertor(ModelArchConfigConvertorBase):
529 def get_head_size(self) -> int:
530 return getattr(self.hf_text_config, "attention_head_dim", 0)
531
532
533 class FalconModelArchConfigConvertor(ModelArchConfigConvertorBase):
534 def get_total_num_kv_heads(self) -> int:
535 # NOTE: for falcon, when new_decoder_architecture is True, the
536 # multi_query flag is ignored and we use n_head_kv for the number of
537 # KV heads.
538 new_decoder_arch_falcon = getattr(
539 self.hf_text_config, "new_decoder_architecture", False
540 )
541
542 if not new_decoder_arch_falcon and getattr(
543 self.hf_text_config, "multi_query", False
544 ):
545 # Multi-query attention, only one KV head.
546 return 1
547
548 # Use the base implementation which checks n_head_kv, num_kv_heads, etc.
549 return super().get_total_num_kv_heads()
550
551
552 class DbrxModelArchConfigConvertor(ModelArchConfigConvertorBase):
553 def get_total_num_kv_heads(self) -> int:
554 return getattr(
555 self.hf_text_config.attn_config,
556 "kv_n_heads",
557 self.hf_text_config.num_attention_heads,
558 )
559
560
561 class NemotronNasModelArchConfigConvertor(ModelArchConfigConvertorBase):
562 def get_total_num_kv_heads(self) -> int:
563 for block in self.hf_text_config.block_configs:
564 if not block.attention.no_op:
565 return (
566 self.hf_text_config.num_attention_heads
567 // block.attention.n_heads_in_group
568 )
569 raise RuntimeError(
570 "Could not determine the number of key-value attention heads "
571 "from model configuration. "
572 f"Architecture: {self.get_architectures()}. "
573 "This usually indicates an unsupported model architecture or "
574 "missing configuration. "
575 "Please check if your model is supported at: "
576 "https://docs.vllm.ai/en/latest/models/supported_models.html"
577 )
578
579
580 class DeepSeekMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
581 def get_num_hidden_layers(self) -> int:
582 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
583
584
585 class DeepseekV4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
586 def __init__(
587 self,
588 hf_config: PreTrainedConfig,
589 hf_text_config: PreTrainedConfig,
590 revision: str | None = None,
591 ):
592 # DeepSeek-V4-Flash-Vision-Exp ships the same architectures/model_type
593 # as the text-only DeepSeek-V4-Flash; route to the VL wrapper class
594 # when the config carries a vision tower. Mutate (not just override
595 # get_architectures) because model-class resolution reads the raw
596 # hf_config.architectures (get_model_architecture).
597 # Only rewrite the stock text architecture: speculative-draft configs
598 # (DSparkDraftModel / DeepSeekV4MTPModel) keep their own classes, and
599 # the VL wrapper marks the config copy it hands to the inner text
600 # backbone with _dsv4_vl_inner so this rewrite does not recurse.
601 if (
602 getattr(hf_config, "vision_n_layers", 0) > 0
603 and getattr(hf_config, "architectures", None) == ["DeepseekV4ForCausalLM"]
604 and not getattr(hf_config, "_dsv4_vl_inner", False)
605 ):
606 hf_config.architectures = ["DeepseekV4ForConditionalGeneration"]
607 super().__init__(hf_config, hf_text_config, revision)
608
609 def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
610 # The vision variant needs the mm-prefix plumbing: it makes the
611 # scheduler prefill image spans atomically (disable_chunked_mm_input)
612 # and routes per-request image ranges to the sparse-SWA metadata
613 # builder, which widens the sliding window bidirectionally in-kernel.
614 if not supports_multimodal:
615 return False
616 return getattr(self.hf_config, "vision_n_layers", 0) > 0
617
618
619 class MimoMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
620 def get_num_hidden_layers(self) -> int:
621 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
622
623
624 def _strip_mimo_v2_attention_chunk_size(
625 hf_config: PreTrainedConfig, hf_text_config: PreTrainedConfig
626 ) -> None:
627 # MiMo-V2-Flash's config.json sets `attention_chunk_size=128` but the
628 # architecture does not actually use chunked local attention. Leaving it
629 # set makes vLLM disable the hybrid KV cache manager
630 for cfg in (hf_text_config, hf_config):
631 if cfg is not None and hasattr(cfg, "attention_chunk_size"):
632 delattr(cfg, "attention_chunk_size")
633
634
635 class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase):
636 def __init__(
637 self,
638 hf_config: PreTrainedConfig,
639 hf_text_config: PreTrainedConfig,
640 revision: str | None = None,
641 ):
642 if getattr(hf_config, "vision_config", None):
643 hf_config.architectures = ["MiMoV2OmniForCausalLM"]
644 super().__init__(hf_config, hf_text_config, revision)
645 _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)
646
647
648 class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
649 def __init__(
650 self,
651 hf_config: PreTrainedConfig,
652 hf_text_config: PreTrainedConfig,
653 revision: str | None = None,
654 ):
655 super().__init__(hf_config, hf_text_config, revision)
656 _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)
657
658 def get_num_hidden_layers(self) -> int:
659 n = getattr(self.hf_text_config, "num_nextn_predict_layers", None)
660 if n is not None:
661 return n
662 # Fall back to n_predict set by hf_config_override
663 n = getattr(self.hf_text_config, "n_predict", None)
664 return n if n is not None else 0
665
666
667 class GLM4MoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
668 def get_num_hidden_layers(self) -> int:
669 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
670
671
672 class ErnieMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
673 def get_num_hidden_layers(self) -> int:
674 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
675
676
677 class Qwen3NextMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
678 def get_num_hidden_layers(self) -> int:
679 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
680
681
682 class Qwen4ExpMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
683 def get_num_hidden_layers(self) -> int:
684 return getattr(
685 self.hf_text_config,
686 "mtp_num_hidden_layers",
687 getattr(self.hf_text_config, "num_nextn_predict_layers", 0),
688 )
689
690
691 class BailingHybridMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
692 def get_num_hidden_layers(self) -> int:
693 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
694
695
696 class BailingHybridV3MTPModelArchConfigConvertor(
697 BailingHybridMTPModelArchConfigConvertor
698 ):
699 pass
700
701
702 class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
703 def get_num_hidden_layers(self) -> int:
704 return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0)
705
706
707 class Step3p5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
708 def get_num_hidden_layers(self) -> int:
709 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
710
711
712 class PanguUltraMoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
713 def get_num_hidden_layers(self) -> int:
714 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
715
716
717 class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
718 def get_num_hidden_layers(self) -> int:
719 return getattr(self.hf_text_config, "num_nextn_predict_layers", 1)
720
721
722 class Gemma4MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
723 def get_hidden_size(self) -> int:
724 # The speculator buffer must match the backbone (target) model's
725 # hidden dimension, not the draft model's smaller dimension.
726 return getattr(
727 self.hf_config, "backbone_hidden_size", super().get_hidden_size()
728 )
729
730 def get_num_hidden_layers(self) -> int:
731 return getattr(self.hf_text_config, "num_hidden_layers", 0)
732
733
734 class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
735 def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
736 if not supports_multimodal:
737 return False
738 return (
739 getattr(self.hf_text_config, "use_bidirectional_attention", None)
740 == "vision"
741 )
742
743 def get_per_layer_hf_configs(
744 self,
745 ) -> list[tuple[PreTrainedConfig, PreTrainedConfig]] | None:
746 # Gemma4 uses a larger head dimension, and sometimes more KV heads, on its
747 # full attention layers than on its sliding ones. Transformers >= 5.15.0
748 # says so in the config; before that the values are flat attributes picked
749 # apart by `layer_types`, so build the per-layer configs here instead. Both
750 # then reach the base convertor as ordinary homogeneous configs.
751 if getattr(self.hf_text_config, "is_heterogeneous", False):
752 return super().get_per_layer_hf_configs()
753 text_config = self.hf_text_config
754 if not getattr(text_config, "layer_types", None):
755 return None
756 # `num_hidden_layers` is what the stack is built from, and it is what
757 # Transformers sizes `per_layer_config` by. `layer_types` can be longer,
758 # e.g. when a test truncates the layer count to build a small model.
759 return [
760 (self.hf_config, gemma4_layer_config(text_config, layer_idx))
761 for layer_idx in range(self.get_num_hidden_layers())
762 ]
763
764
765 class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):
766 def _language_config(self) -> PreTrainedConfig:
767 return self.hf_config.language_config
768
769 def get_num_hidden_layers(self) -> int:
770 return getattr(self._language_config(), "num_hidden_layers", 0)
771
772 def get_total_num_attention_heads(self) -> int:
773 return getattr(self._language_config(), "num_attention_heads", 0)
774
775 def get_vocab_size(self) -> int:
776 return getattr(self._language_config(), "vocab_size", 0)
777
778 def get_hidden_size(self) -> int:
779 return getattr(self._language_config(), "hidden_size", 0)
780
781 def get_head_size(self) -> int:
782 head_dim = getattr(self._language_config(), "head_dim", None)
783 if head_dim is not None:
784 return head_dim
785 total_num_attention_heads = self.get_total_num_attention_heads()
786 if total_num_attention_heads == 0:
787 return 0
788 return self.get_hidden_size() // total_num_attention_heads
789
790 def get_total_num_kv_heads(self) -> int:
791 return getattr(
792 self._language_config(),
793 "num_key_value_heads",
794 self.get_total_num_attention_heads(),
795 )
796
797 def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
798 language_config = self._language_config()
799 max_position_embeddings = getattr(
800 language_config,
801 "max_position_embeddings",
802 None,
803 )
804 if max_position_embeddings is None:
805 return super().derive_max_model_len_and_key()
806 return max_position_embeddings, "language_config.max_position_embeddings"
807
808
809 # hf_config.model_type -> convertor class
810 MODEL_ARCH_CONFIG_CONVERTORS = {
811 "bailing_hybrid_mtp": BailingHybridMTPModelArchConfigConvertor,
812 "cohere_asr": CohereAsrModelArchConfigConvertor,
813 "dbrx": DbrxModelArchConfigConvertor,
814 "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
815 "deepseek_v4": DeepseekV4ModelArchConfigConvertor,
816 "diffusion_gemma": Gemma4ModelArchConfigConvertor,
817 "diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
818 "ernie_mtp": ErnieMTPModelArchConfigConvertor,
819 "falcon": FalconModelArchConfigConvertor,
820 "falcon_mamba": MambaModelArchConfigConvertor,
821 "gemma4": Gemma4ModelArchConfigConvertor,
822 "gemma4_mtp": Gemma4MTPModelArchConfigConvertor,
823 "gemma4_text": Gemma4ModelArchConfigConvertor,
824 "gemma4_unified": Gemma4ModelArchConfigConvertor,
825 "gemma4_unified_text": Gemma4ModelArchConfigConvertor,
826 "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor,
827 "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor,
828 "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor,
829 "mamba": MambaModelArchConfigConvertor,
830 "medusa": MedusaModelArchConfigConvertor,
831 "mimo_mtp": MimoMTPModelArchConfigConvertor,
832 "mimo_v2": MimoV2ModelArchConfigConvertor,
833 "mimo_v2_flash": MimoV2ModelArchConfigConvertor,
834 "mimo_v2_mtp": MimoV2MTPModelArchConfigConvertor,
835 "mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor,
836 "moss_audio": MossAudioModelArchConfigConvertor,
837 "nemotron-nas": NemotronNasModelArchConfigConvertor,
838 "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor,
839 "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,
840 "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor,
841 "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,
842 "qwen4_exp_mtp": Qwen4ExpMTPModelArchConfigConvertor,
843 "RefinedWeb": FalconModelArchConfigConvertor,
844 "RefinedWebModel": FalconModelArchConfigConvertor,
845 "step3p5_mtp": Step3p5MTPModelArchConfigConvertor,
846 "timm_wrapper": TerratorchModelArchConfigConvertor,
847 "zamba2": Zamba2ModelArchConfigConvertor,
848 }