| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for ModelArchitectureConfig and its integration with ModelConfig.""" |
| 4 | |
| 5 | import json |
| 6 | from copy import copy |
| 7 | from pathlib import Path |
| 8 | from types import SimpleNamespace |
| 9 | from typing import cast |
| 10 | from unittest.mock import Mock |
| 11 | |
| 12 | import pytest |
| 13 | from transformers import PreTrainedConfig |
| 14 | from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig |
| 15 | |
| 16 | from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig |
| 17 | from vllm.config.model_arch import ModelArchitectureConfig |
| 18 | from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig |
| 19 | from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config |
| 20 | from vllm.transformers_utils.model_arch_config_convertor import ( |
| 21 | MODEL_ARCH_CONFIG_CONVERTORS, |
| 22 | Gemma4ModelArchConfigConvertor, |
| 23 | ModelArchConfigConvertorBase, |
| 24 | ) |
| 25 | |
| 26 | BASE_TRUST_REMOTE_CODE_MODELS = { |
| 27 | "nvidia/Llama-3_3-Nemotron-Super-49B-v1", |
| 28 | "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", |
| 29 | "XiaomiMiMo/MiMo-7B-RL", |
| 30 | "stepfun-ai/Step-3.5-Flash", |
| 31 | # Excluded: Not available online right now |
| 32 | # "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1", |
| 33 | "meituan-longcat/LongCat-Flash-Chat", |
| 34 | } |
| 35 | |
| 36 | BASE_MODELS_TO_TEST = [ |
| 37 | "state-spaces/mamba-130m-hf", |
| 38 | "mistralai/Mamba-Codestral-7B-v0.1", |
| 39 | # Excluded: terratorch/torchgeo version mismatch in CPU CI environment |
| 40 | # (NonGeoDataset import error). Tested in model initialization tests. |
| 41 | # "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", |
| 42 | "Zyphra/Zamba2-7B-instruct", |
| 43 | # FIXME: databricks/dbrx-instruct has been deleted |
| 44 | # "databricks/dbrx-instruct", |
| 45 | "tiiuae/falcon-7b", |
| 46 | "tiiuae/falcon-40b", |
| 47 | "luccafong/deepseek_mtp_main_random", |
| 48 | "Qwen/Qwen3-Next-80B-A3B-Instruct", |
| 49 | "tiny-random/qwen3-next-moe", |
| 50 | "zai-org/GLM-4.5", |
| 51 | "baidu/ERNIE-4.5-21B-A3B-PT", |
| 52 | # Models using base convertor |
| 53 | "lmsys/gpt-oss-20b-bf16", |
| 54 | "deepseek-ai/DeepSeek-V3.2-Exp", |
| 55 | "meta-llama/Llama-4-Scout-17B-16E-Instruct", |
| 56 | ] + list(BASE_TRUST_REMOTE_CODE_MODELS) |
| 57 | |
| 58 | # (target_model, draft_model, trust_remote_code) |
| 59 | SPECULATIVE_MODELS = [ |
| 60 | ("JackFram/llama-68m", "abhigoyal/vllm-medusa-llama-68m-random", False), |
| 61 | ("luccafong/deepseek_mtp_main_random", "luccafong/deepseek_mtp_draft_random", True), |
| 62 | ("eagle618/deepseek-v3-random", "eagle618/eagle-deepseek-v3-random", True), |
| 63 | ("meta-llama/Meta-Llama-3-8B-Instruct", "yuhuili/EAGLE-LLaMA3-Instruct-8B", True), |
| 64 | ("meta-llama/Llama-3.1-8B-Instruct", "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", True), |
| 65 | ] |
| 66 | |
| 67 | |
| 68 | def _load_groundtruth(filename: str) -> dict: |
| 69 | """Load groundtruth JSON from the test directory.""" |
| 70 | groundtruth_path = Path(__file__).parent / filename |
| 71 | with open(groundtruth_path) as f: |
| 72 | return json.load(f) |
| 73 | |
| 74 | |
| 75 | def _assert_model_arch_config( |
| 76 | model_config, expected: dict, check_head_size: bool = True |
| 77 | ): |
| 78 | """Assert model_arch_config matches expected values.""" |
| 79 | model_arch_config = model_config.model_arch_config |
| 80 | assert model_arch_config.architectures == expected["architectures"] |
| 81 | assert model_arch_config.model_type == expected["model_type"] |
| 82 | assert model_arch_config.text_model_type == expected["text_model_type"] |
| 83 | assert model_arch_config.hidden_size == expected["hidden_size"] |
| 84 | assert ( |
| 85 | model_arch_config.total_num_hidden_layers == expected["total_num_hidden_layers"] |
| 86 | ) |
| 87 | assert ( |
| 88 | model_arch_config.total_num_attention_heads |
| 89 | == expected["total_num_attention_heads"] |
| 90 | ) |
| 91 | assert model_arch_config.vocab_size == expected["vocab_size"] |
| 92 | assert model_arch_config.total_num_kv_heads == expected["total_num_kv_heads"] |
| 93 | assert model_arch_config.num_experts == expected["num_experts"] |
| 94 | assert model_arch_config.is_deepseek_mla == expected["is_deepseek_mla"] |
| 95 | |
| 96 | torch_dtype = ModelArchConfigConvertorBase.get_torch_dtype( |
| 97 | model_config.hf_config, |
| 98 | model_config.model, |
| 99 | revision=model_config.revision, |
| 100 | config_format="hf", |
| 101 | ) |
| 102 | assert str(torch_dtype) == expected["dtype"] |
| 103 | |
| 104 | if check_head_size: |
| 105 | assert model_arch_config.head_size == expected["head_size"] |
| 106 | |
| 107 | |
| 108 | def _assert_model_config_methods( |
| 109 | model_config, expected: dict, check_head_size: bool = True |
| 110 | ): |
| 111 | """Assert model_config methods return expected values.""" |
| 112 | assert model_config.architectures == expected["architectures"] |
| 113 | assert model_config.get_vocab_size() == expected["vocab_size"] |
| 114 | assert model_config.get_hidden_size() == expected["hidden_size"] |
| 115 | assert model_config.get_total_num_kv_heads() == expected["total_num_kv_heads"] |
| 116 | assert model_config.get_num_experts() == expected["num_experts"] |
| 117 | assert ( |
| 118 | model_config.get_total_num_hidden_layers() |
| 119 | == expected["total_num_hidden_layers"] |
| 120 | ) |
| 121 | |
| 122 | if check_head_size: |
| 123 | assert model_config.get_head_size() == expected["head_size"] |
| 124 | |
| 125 | |
| 126 | def test_head_size_falls_back_when_head_dim_is_zero(): |
| 127 | """Regression test for configs that materialize missing head_dim as 0.""" |
| 128 | hf_config = PreTrainedConfig( |
| 129 | model_type="deepseek_vl_v2", |
| 130 | hidden_size=1280, |
| 131 | num_attention_heads=10, |
| 132 | num_key_value_heads=10, |
| 133 | head_dim=0, |
| 134 | kv_lora_rank=None, |
| 135 | ) |
| 136 | |
| 137 | convertor = ModelArchConfigConvertorBase(hf_config, hf_config) |
| 138 | |
| 139 | assert convertor.get_head_size() == 128 |
| 140 | |
| 141 | |
| 142 | def test_qk_rope_head_dim_recovery_uses_model_revision(monkeypatch: pytest.MonkeyPatch): |
| 143 | hf_config = PreTrainedConfig( |
| 144 | model_type="deepseek_v2", |
| 145 | qk_rope_head_dim=128, |
| 146 | qk_nope_head_dim=128, |
| 147 | ) |
| 148 | hf_config.name_or_path = "org/model" |
| 149 | get_hf_file_to_dict = Mock(return_value={"qk_rope_head_dim": 64}) |
| 150 | |
| 151 | monkeypatch.setattr( |
| 152 | "vllm.transformers_utils.repo_utils.get_hf_file_to_dict", |
| 153 | get_hf_file_to_dict, |
| 154 | ) |
| 155 | |
| 156 | convertor = ModelArchConfigConvertorBase( |
| 157 | hf_config, |
| 158 | hf_config, |
| 159 | revision="pinned-revision", |
| 160 | ) |
| 161 | |
| 162 | assert convertor._get_qk_rope_head_dim() == 64 |
| 163 | get_hf_file_to_dict.assert_called_once_with( |
| 164 | "config.json", "org/model", "pinned-revision" |
| 165 | ) |
| 166 | |
| 167 | |
| 168 | @pytest.mark.parametrize( |
| 169 | "attribute", |
| 170 | [ |
| 171 | "num_experts_per_tok", |
| 172 | "num_experts_per_token", |
| 173 | "top_k_experts", |
| 174 | "moe_topk", |
| 175 | "moe_top_k", |
| 176 | ], |
| 177 | ) |
| 178 | def test_num_experts_per_tok_aliases(attribute: str): |
| 179 | hf_config = PreTrainedConfig(**{attribute: 4}) |
| 180 | convertor = ModelArchConfigConvertorBase(hf_config, hf_config) |
| 181 | model_config = cast( |
| 182 | ModelConfig, |
| 183 | SimpleNamespace( |
| 184 | model_arch_config=SimpleNamespace( |
| 185 | num_experts_per_token=convertor.get_num_experts_per_token() |
| 186 | ) |
| 187 | ), |
| 188 | ) |
| 189 | |
| 190 | assert ModelConfig.get_num_experts_per_tok(model_config) == 4 |
| 191 | |
| 192 | |
| 193 | def test_num_experts_per_tok_none_is_normalized(): |
| 194 | hf_config = PreTrainedConfig(top_k_experts=None) |
| 195 | convertor = ModelArchConfigConvertorBase(hf_config, hf_config) |
| 196 | |
| 197 | assert convertor.get_num_experts_per_token() == 0 |
| 198 | |
| 199 | |
| 200 | def test_legacy_modelopt_config_without_producer_is_normalized(): |
| 201 | quantization_config = { |
| 202 | "quantization": { |
| 203 | "quant_algo": "NVFP4", |
| 204 | "group_size": 16, |
| 205 | "kv_cache_quant_algo": None, |
| 206 | "exclude_modules": [], |
| 207 | "modelopt_quant_config": {"quant_cfg": {}}, |
| 208 | } |
| 209 | } |
| 210 | hf_config = PreTrainedConfig(quantization_config=quantization_config) |
| 211 | |
| 212 | convertor = ModelArchConfigConvertorBase(hf_config, hf_config) |
| 213 | |
| 214 | assert convertor.get_quantization_config()["quant_method"] == "modelopt_fp4" |
| 215 | |
| 216 | |
| 217 | def _layer(**overrides) -> ModelArchitectureConfig: |
| 218 | fields = dict( |
| 219 | architectures=["X"], |
| 220 | model_type="x", |
| 221 | text_model_type=None, |
| 222 | hidden_size=64, |
| 223 | total_num_hidden_layers=3, |
| 224 | total_num_attention_heads=8, |
| 225 | head_size=16, |
| 226 | vocab_size=32, |
| 227 | total_num_kv_heads=4, |
| 228 | num_experts=0, |
| 229 | num_experts_per_token=0, |
| 230 | quantization_config=None, |
| 231 | is_deepseek_mla=False, |
| 232 | is_mm_prefix_lm=False, |
| 233 | rswa_window=None, |
| 234 | derived_max_model_len_and_key=(1.0, None), |
| 235 | ) |
| 236 | return ModelArchitectureConfig(**(fields | overrides)) |
| 237 | |
| 238 | |
| 239 | def test_from_layers_collapses_to_the_largest_layer(): |
| 240 | """The whole-model config must size buffers for every layer.""" |
| 241 | arch = ModelArchitectureConfig.from_layers( |
| 242 | [_layer(), _layer(), _layer(head_size=32, total_num_kv_heads=8)] |
| 243 | ) |
| 244 | |
| 245 | assert (arch.head_size, arch.total_num_kv_heads) == (32, 8) |
| 246 | assert [arch[i].head_size for i in range(3)] == [16, 16, 32] |
| 247 | assert [arch[i].total_num_kv_heads for i in range(3)] == [4, 4, 8] |
| 248 | # A layer view is itself homogeneous, so reading from it cannot recurse. |
| 249 | assert arch[0].per_layer_overrides is None |
| 250 | |
| 251 | |
| 252 | def test_layer_view_follows_later_edits_to_the_whole_model_config(): |
| 253 | """`model_arch_config` is edited after construction in a few places.""" |
| 254 | arch = ModelArchitectureConfig.from_layers( |
| 255 | [_layer(), _layer(), _layer(head_size=32)] |
| 256 | ) |
| 257 | |
| 258 | arch.is_mm_prefix_lm = True |
| 259 | |
| 260 | assert all(arch[i].is_mm_prefix_lm for i in range(3)) |
| 261 | assert [arch[i].head_size for i in range(3)] == [16, 16, 32] |
| 262 | |
| 263 | |
| 264 | def test_uniform_layers_stay_homogeneous(): |
| 265 | """A checkpoint can vary attributes vLLM never reads.""" |
| 266 | arch = ModelArchitectureConfig.from_layers([_layer()] * 3) |
| 267 | |
| 268 | assert arch.per_layer_overrides is None |
| 269 | assert arch[1] is arch |
| 270 | |
| 271 | |
| 272 | def test_getitem_rejects_negative_indices(): |
| 273 | """Wrapping would only misbehave on heterogeneous models, so reject both.""" |
| 274 | homogeneous = _layer() |
| 275 | heterogeneous = ModelArchitectureConfig.from_layers( |
| 276 | [_layer(), _layer(), _layer(head_size=32)] |
| 277 | ) |
| 278 | |
| 279 | for arch in (homogeneous, heterogeneous): |
| 280 | with pytest.raises(IndexError): |
| 281 | arch[-1] |
| 282 | |
| 283 | |
| 284 | @pytest.mark.parametrize( |
| 285 | "varying", |
| 286 | [ |
| 287 | # `bool` is an `int`; collapsing this one with `max` would make `use_mla` |
| 288 | # true model wide and silently discard every per-layer KV head count. |
| 289 | {"is_deepseek_mla": True}, |
| 290 | {"quantization_config": {"quant_method": "fp8"}}, |
| 291 | {"rswa_window": 512}, |
| 292 | ], |
| 293 | ) |
| 294 | def test_from_layers_rejects_fields_with_no_whole_model_value(varying: dict): |
| 295 | with pytest.raises(ValueError, match="varies across layers"): |
| 296 | ModelArchitectureConfig.from_layers([_layer(), _layer(), _layer(**varying)]) |
| 297 | |
| 298 | |
| 299 | def test_from_layers_rejects_a_layer_count_mismatch(): |
| 300 | """Draft configs can inherit a target's per-layer spec.""" |
| 301 | with pytest.raises(ValueError, match="per-layer configs"): |
| 302 | ModelArchitectureConfig.from_layers([_layer(), _layer(head_size=32)]) |
| 303 | |
| 304 | |
| 305 | def _gemma4_text_config(**overrides) -> Gemma4TextConfig: |
| 306 | """A six layer Gemma4 whose last layer is wider than the rest.""" |
| 307 | fields = dict( |
| 308 | num_hidden_layers=6, |
| 309 | hidden_size=64, |
| 310 | num_attention_heads=8, |
| 311 | num_key_value_heads=4, |
| 312 | head_dim=16, |
| 313 | global_head_dim=32, |
| 314 | # Transformers forces the last layer to full attention. |
| 315 | layer_types=["sliding_attention"] * 5 + ["full_attention"], |
| 316 | ) |
| 317 | return Gemma4TextConfig(**(fields | overrides)) |
| 318 | |
| 319 | |
| 320 | def test_gemma4_head_dims_vary_by_layer_type(): |
| 321 | """Gemma4's full attention layers are wider than its sliding ones. |
| 322 | |
| 323 | Transformers >= 5.15.0 says so in the config; this exercises the convertor |
| 324 | building the same per-layer view from the flat attributes used before that. |
| 325 | """ |
| 326 | text_config = _gemma4_text_config( |
| 327 | num_global_key_value_heads=8, attention_k_eq_v=True |
| 328 | ) |
| 329 | |
| 330 | arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert() |
| 331 | |
| 332 | assert (arch.head_size, arch.total_num_kv_heads) == (32, 8) |
| 333 | assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32] |
| 334 | assert [arch[i].total_num_kv_heads for i in range(6)] == [4] * 5 + [8] |
| 335 | # The model files resolve each layer through the same helper, so the KV cache |
| 336 | # vLLM allocates and the projections the model builds cannot disagree. |
| 337 | assert [gemma4_layer_config(text_config, i).head_dim for i in range(6)] == [ |
| 338 | arch[i].head_size for i in range(6) |
| 339 | ] |
| 340 | |
| 341 | |
| 342 | def test_gemma4_without_global_kv_heads(): |
| 343 | """`num_global_key_value_heads` defaults to `None`, not to a head count.""" |
| 344 | text_config = _gemma4_text_config() |
| 345 | |
| 346 | arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert() |
| 347 | |
| 348 | assert arch.total_num_kv_heads == 4 |
| 349 | assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32] |
| 350 | |
| 351 | |
| 352 | def test_gemma4_layer_count_comes_from_num_hidden_layers(): |
| 353 | """`dummy_hf_overrides` shrinks the stack but leaves `layer_types` long.""" |
| 354 | text_config = _gemma4_text_config() |
| 355 | text_config.num_hidden_layers = 3 |
| 356 | |
| 357 | arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert() |
| 358 | |
| 359 | assert arch.total_num_hidden_layers == 3 |
| 360 | # Only sliding layers survive the truncation, so nothing varies. |
| 361 | assert arch.per_layer_overrides is None |
| 362 | |
| 363 | |
| 364 | def test_gemma4_uniform_head_dims_are_homogeneous(): |
| 365 | text_config = _gemma4_text_config(global_head_dim=16) |
| 366 | |
| 367 | arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert() |
| 368 | |
| 369 | assert arch.per_layer_overrides is None |
| 370 | assert arch[3] is arch |
| 371 | |
| 372 | |
| 373 | class _HeterogeneousConfig(PreTrainedConfig): |
| 374 | """A heterogeneous config with no convertor of its own. |
| 375 | |
| 376 | Mirrors the parts vLLM uses: per-layer configs are shallow copies with the |
| 377 | varying attributes applied and heterogeneity stripped, so they do not recurse.""" |
| 378 | |
| 379 | is_heterogeneous = True |
| 380 | |
| 381 | def __init__(self, per_layer: dict[str, list], **kwargs): |
| 382 | # Set before `super().__init__` because it validates `per_layer_config` |
| 383 | self._per_layer = per_layer |
| 384 | super().__init__(**kwargs) |
| 385 | |
| 386 | @property |
| 387 | def per_layer_config(self) -> list[PreTrainedConfig]: |
| 388 | layers = [] |
| 389 | for i in range(self.num_hidden_layers): |
| 390 | layer = copy(self) |
| 391 | layer.is_heterogeneous = False |
| 392 | for name, values in self._per_layer.items(): |
| 393 | setattr(layer, name, values[i]) |
| 394 | layers.append(layer) |
| 395 | return layers |
| 396 | |
| 397 | |
| 398 | def test_transformers_heterogeneous_config_is_resolved_per_layer(): |
| 399 | hf_config = _HeterogeneousConfig( |
| 400 | per_layer={"head_dim": [16, 16, 32], "num_key_value_heads": [4, 4, 8]}, |
| 401 | num_hidden_layers=3, |
| 402 | hidden_size=64, |
| 403 | num_attention_heads=8, |
| 404 | ) |
| 405 | |
| 406 | arch = ModelArchConfigConvertorBase(hf_config, hf_config).convert() |
| 407 | |
| 408 | assert (arch.head_size, arch.total_num_kv_heads) == (32, 8) |
| 409 | assert [arch[i].head_size for i in range(3)] == [16, 16, 32] |
| 410 | assert [arch[i].total_num_kv_heads for i in range(3)] == [4, 4, 8] |
| 411 | |
| 412 | |
| 413 | def test_heterogeneous_config_varying_nothing_vllm_reads(): |
| 414 | """Transformers prunes nothing for us: the diff has to notice.""" |
| 415 | hf_config = _HeterogeneousConfig( |
| 416 | per_layer={"bos_token_id": [1, 2, 3]}, |
| 417 | num_hidden_layers=3, |
| 418 | hidden_size=64, |
| 419 | num_attention_heads=8, |
| 420 | head_dim=16, |
| 421 | ) |
| 422 | |
| 423 | arch = ModelArchConfigConvertorBase(hf_config, hf_config).convert() |
| 424 | |
| 425 | assert arch.per_layer_overrides is None |
| 426 | assert arch[2] is arch |
| 427 | |
| 428 | |
| 429 | @pytest.mark.parametrize("model", BASE_MODELS_TO_TEST) |
| 430 | def test_base_model_arch_config(model: str): |
| 431 | """Test model architecture config for base models.""" |
| 432 | groundtruth = _load_groundtruth("base_model_arch_groundtruth.json") |
| 433 | expected = groundtruth[model] |
| 434 | |
| 435 | model_config = ModelConfig( |
| 436 | model, trust_remote_code=model in BASE_TRUST_REMOTE_CODE_MODELS |
| 437 | ) |
| 438 | |
| 439 | _assert_model_arch_config(model_config, expected) |
| 440 | _assert_model_config_methods(model_config, expected) |
| 441 | |
| 442 | |
| 443 | @pytest.mark.parametrize( |
| 444 | "target_model,draft_model,trust_remote_code", SPECULATIVE_MODELS |
| 445 | ) |
| 446 | def test_draft_model_arch_config( |
| 447 | target_model: str, draft_model: str, trust_remote_code: bool |
| 448 | ): |
| 449 | """Test model architecture config for draft/speculative models.""" |
| 450 | groundtruth = _load_groundtruth("draft_model_arch_groundtruth.json") |
| 451 | expected = groundtruth[draft_model] |
| 452 | |
| 453 | target_model_config = ModelConfig(target_model, trust_remote_code=trust_remote_code) |
| 454 | speculative_config = SpeculativeConfig( |
| 455 | model=draft_model, |
| 456 | num_speculative_tokens=1, |
| 457 | target_model_config=target_model_config, |
| 458 | target_parallel_config=ParallelConfig(), |
| 459 | ) |
| 460 | model_config = speculative_config.draft_model_config |
| 461 | |
| 462 | # For medusa models, head_size may cause division by zero before |
| 463 | # model_arch_config was introduced, so we conditionally check it |
| 464 | check_head_size = isinstance(expected["head_size"], int) |
| 465 | |
| 466 | _assert_model_arch_config(model_config, expected, check_head_size=check_head_size) |
| 467 | _assert_model_config_methods( |
| 468 | model_config, expected, check_head_size=check_head_size |
| 469 | ) |
| 470 | |
| 471 | |
| 472 | def test_deepseek_v4_convertor_splits_vision_architecture(): |
| 473 | """The vision checkpoint shares model_type/architectures with the |
| 474 | text-only DeepSeek-V4; the convertor routes it to the VL wrapper class |
| 475 | by rewriting hf_config.architectures (model-class resolution reads the |
| 476 | raw attribute).""" |
| 477 | from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config |
| 478 | from vllm.transformers_utils.model_arch_config_convertor import ( |
| 479 | DeepseekV4ModelArchConfigConvertor, |
| 480 | ) |
| 481 | |
| 482 | text_cfg = DeepseekV4Config(architectures=["DeepseekV4ForCausalLM"]) |
| 483 | conv = DeepseekV4ModelArchConfigConvertor(text_cfg, text_cfg) |
| 484 | assert conv.get_architectures() == ["DeepseekV4ForCausalLM"] |
| 485 | assert text_cfg.architectures == ["DeepseekV4ForCausalLM"] |
| 486 | |
| 487 | vl_cfg = DeepseekV4Config( |
| 488 | architectures=["DeepseekV4ForCausalLM"], vision_n_layers=32 |
| 489 | ) |
| 490 | conv = DeepseekV4ModelArchConfigConvertor(vl_cfg, vl_cfg) |
| 491 | assert conv.get_architectures() == ["DeepseekV4ForConditionalGeneration"] |
| 492 | assert vl_cfg.architectures == ["DeepseekV4ForConditionalGeneration"] |
| 493 | |
| 494 | # Speculative-draft configs (set by SpeculativeConfig.hf_config_override) |
| 495 | # keep their own architecture even with a vision tower in the config. |
| 496 | for draft_arch in ("DSparkDraftModel", "DeepSeekV4MTPModel"): |
| 497 | draft_cfg = DeepseekV4Config(architectures=[draft_arch], vision_n_layers=32) |
| 498 | conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg) |
| 499 | assert conv.get_architectures() == [draft_arch] |
| 500 | assert draft_cfg.architectures == [draft_arch] |
| 501 | |
| 502 | |
| 503 | def test_diffusion_gemma_head_dims_vary_by_layer_type(): |
| 504 | """Checkpoints carry model_type `diffusion_gemma` at the top level. Under |
| 505 | the base convertor every layer reports the sliding head dim, Gemma4Config |
| 506 | skips its head-dim handling, and the full-attention KV-cache group falls to |
| 507 | FlashInfer, which cannot take a per-request causal tensor.""" |
| 508 | text_config = dict( |
| 509 | num_hidden_layers=6, |
| 510 | hidden_size=64, |
| 511 | num_attention_heads=8, |
| 512 | num_key_value_heads=4, |
| 513 | head_dim=16, |
| 514 | global_head_dim=32, |
| 515 | layer_types=["sliding_attention"] * 5 + ["full_attention"], |
| 516 | ) |
| 517 | hf_config = DiffusionGemmaConfig(text_config=text_config) |
| 518 | |
| 519 | convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS[hf_config.model_type] |
| 520 | assert convertor_cls is Gemma4ModelArchConfigConvertor |
| 521 | arch = convertor_cls(hf_config, hf_config.text_config).convert() |
| 522 | |
| 523 | assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32] |