| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import get_args |
| 5 | |
| 6 | from transformers import PreTrainedConfig |
| 7 | |
| 8 | from vllm.config.speculative import MTPModelTypes, SpeculativeConfig |
| 9 | from vllm.transformers_utils.model_arch_config_convertor import ( |
| 10 | BailingHybridMTPModelArchConfigConvertor, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | def _bailing_config() -> PreTrainedConfig: |
| 15 | config = PreTrainedConfig( |
| 16 | architectures=["BailingMoeV2_5ForCausalLM"], |
| 17 | hidden_size=4096, |
| 18 | kv_lora_rank=512, |
| 19 | num_attention_heads=32, |
| 20 | num_experts=256, |
| 21 | num_hidden_layers=32, |
| 22 | num_key_value_heads=32, |
| 23 | num_nextn_predict_layers=1, |
| 24 | qk_rope_head_dim=64, |
| 25 | vocab_size=157184, |
| 26 | ) |
| 27 | config.model_type = "bailing_hybrid" |
| 28 | return config |
| 29 | |
| 30 | |
| 31 | def test_bailing_hybrid_mtp_hf_config_override(): |
| 32 | config = _bailing_config() |
| 33 | |
| 34 | overridden = SpeculativeConfig.hf_config_override(config) |
| 35 | |
| 36 | assert overridden.model_type == "bailing_hybrid_mtp" |
| 37 | assert overridden.architectures == ["BailingMoeV25MTPModel"] |
| 38 | assert overridden.n_predict == 1 |
| 39 | assert "bailing_hybrid_mtp" in get_args(MTPModelTypes) |
| 40 | |
| 41 | |
| 42 | def test_bailing_hybrid_mtp_model_arch_config(): |
| 43 | config = _bailing_config() |
| 44 | config.model_type = "bailing_hybrid_mtp" |
| 45 | config.architectures = ["BailingMoeV25MTPModel"] |
| 46 | |
| 47 | model_arch_config = BailingHybridMTPModelArchConfigConvertor( |
| 48 | config, config |
| 49 | ).convert() |
| 50 | |
| 51 | assert model_arch_config.model_type == "bailing_hybrid_mtp" |
| 52 | assert model_arch_config.architectures == ["BailingMoeV25MTPModel"] |
| 53 | assert model_arch_config.total_num_hidden_layers == 1 |
| 54 | assert model_arch_config.is_deepseek_mla |