@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 159 lines 5.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import copy
5 import json
6 from pathlib import Path
7 from typing import Any
8
9 import pytest
10 from transformers import AutoConfig
11
12 from vllm.transformers_utils.config import get_config, uses_mrope
13 from vllm.transformers_utils.configs.bailing_moe_v3_vl import (
14 BailingMoeV3TextConfig,
15 BailingMoeV3VisionConfig,
16 BailingMoeV3VLConfig,
17 )
18
19
20 def _bailing_v3_vl_config() -> dict[str, Any]:
21 return {
22 "architectures": ["BailingMoeV3VLForConditionalGeneration"],
23 "auto_map": {
24 "AutoConfig": ("configuration_bailing_moe_v3_vl.BailingMoeV3VLConfig")
25 },
26 "model_type": "bailing_moe_v3_vl",
27 "mrope_section": [8, 12, 12],
28 "text_config": {
29 "num_hidden_layers": 42,
30 "layer_group_size": 6,
31 "rope_theta": 6_000_000,
32 },
33 "vision_config": {"disable_merger_proj": True},
34 }
35
36
37 def _write_config(path: Path, config: dict[str, Any]) -> None:
38 path.mkdir()
39 (path / "config.json").write_text(json.dumps(config), encoding="utf-8")
40
41
42 def test_bailing_v3_vl_config_loads_without_remote_code(tmp_path: Path):
43 model_path = tmp_path / "model"
44 _write_config(model_path, _bailing_v3_vl_config())
45
46 config = get_config(model_path, trust_remote_code=False)
47
48 assert isinstance(config, BailingMoeV3VLConfig)
49 assert isinstance(config.text_config, BailingMoeV3TextConfig)
50 assert isinstance(config.vision_config, BailingMoeV3VisionConfig)
51 assert isinstance(
52 AutoConfig.from_pretrained(model_path, trust_remote_code=False),
53 BailingMoeV3VLConfig,
54 )
55 assert config.mrope_section == [8, 12, 12]
56 assert uses_mrope(config)
57
58 text_config = config.text_config
59 assert text_config.model_type == "bailing_hybrid"
60 assert text_config.architectures == ["BailingMoeV3ForCausalLM"]
61 assert text_config.rope_parameters["mrope_section"] == [8, 12, 12]
62 assert text_config.rope_parameters["rope_theta"] == 6_000_000
63 assert text_config.layer_types.count("linear_attention") == 35
64 assert [
65 layer_idx
66 for layer_idx, layer_type in enumerate(text_config.layer_types)
67 if layer_type == "full_attention"
68 ] == [5, 11, 17, 23, 29, 35, 41]
69 assert config.vision_config.disable_merger_proj is True
70
71 normalized = config.to_dict()
72 restored = BailingMoeV3VLConfig.from_dict(normalized)
73 assert isinstance(restored.text_config, BailingMoeV3TextConfig)
74 assert isinstance(restored.vision_config, BailingMoeV3VisionConfig)
75 assert restored.to_dict() == normalized
76
77
78 def test_bailing_v3_vl_config_preserves_explicit_nested_values():
79 config_dict = _bailing_v3_vl_config()
80 config_dict["architectures"] = ["CustomBailingVLForConditionalGeneration"]
81 text_config = config_dict["text_config"]
82 text_config["architectures"] = ["CustomBailingForCausalLM"]
83 text_config["layer_types"] = (["linear_attention"] * 5 + ["full_attention"]) * 7
84 text_config["rope_theta"] = 1_000_000
85 text_config["rope_parameters"] = {
86 "rope_type": "default",
87 "rope_theta": 1_000_000,
88 "mrope_section": [8, 12, 12],
89 }
90 original = copy.deepcopy(config_dict)
91
92 config = BailingMoeV3VLConfig(**config_dict)
93
94 assert config_dict == original
95 assert config.architectures == ["CustomBailingVLForConditionalGeneration"]
96 assert config.text_config.architectures == ["CustomBailingForCausalLM"]
97 assert config.text_config.layer_types == text_config["layer_types"]
98 assert config.text_config.rope_parameters["rope_theta"] == 1_000_000
99 assert config.text_config.rope_parameters["mrope_section"] == [8, 12, 12]
100
101
102 def test_bailing_v3_vl_config_normalizes_config_instance():
103 text_config = BailingMoeV3TextConfig()
104
105 config = BailingMoeV3VLConfig(
106 text_config=text_config,
107 mrope_section=[8, 12, 12],
108 )
109
110 assert config.text_config is text_config
111 assert text_config.rope_parameters["mrope_section"] == [8, 12, 12]
112
113
114 def test_bailing_v3_text_config_trailing_and_legacy_layers():
115 expected = ["linear_attention", "linear_attention", "full_attention"] * 2 + [
116 "full_attention"
117 ] * 2
118 text_config = BailingMoeV3TextConfig(
119 num_hidden_layers=8,
120 layer_group_size=3,
121 )
122 assert text_config.layer_types == expected
123
124 text_config = BailingMoeV3TextConfig(
125 num_hidden_layers=8,
126 layer_group_size=3,
127 layers_block_type=["mamba", "mamba", "attention"] * 2 + ["attention"] * 2,
128 )
129 assert text_config.layer_types == expected
130
131
132 def test_bailing_v3_text_config_accepts_legacy_rope_type():
133 text_config = BailingMoeV3TextConfig(rope_scaling={"type": "linear", "factor": 2.0})
134
135 assert text_config.rope_parameters["type"] == "linear"
136 assert text_config.rope_parameters["rope_type"] == "linear"
137
138
139 def test_bailing_v3_vl_config_rejects_invalid_values():
140 with pytest.raises(ValueError, match="layer_group_size must be positive"):
141 BailingMoeV3TextConfig(layer_group_size=0)
142
143 with pytest.raises(ValueError, match="must match the layer_group_size schedule"):
144 BailingMoeV3TextConfig(
145 num_hidden_layers=8,
146 layer_group_size=3,
147 layer_types=["full_attention"] * 8,
148 )
149
150 with pytest.raises(ValueError, match="conflicting M-RoPE sections"):
151 BailingMoeV3VLConfig(
152 text_config={
153 "rope_parameters": {
154 "rope_type": "default",
155 "mrope_section": [4, 14, 14],
156 }
157 },
158 mrope_section=[8, 12, 12],
159 )