@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 174 lines 5.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """MuseGlimmer config normalization.
5
6 Two schema pairs have to converge on the same model, and each has produced a
7 silent-wrong-output bug:
8
9 1. flat (legacy converter) vs nested (canonical) ``config.json``. A FLAT
10 config once deserialized to an all-default text config, silently dropping
11 every checkpoint value.
12 2. native vs modular attention config. The modular HF text_config OMITS
13 use_qk_norm / use_attn_output_gate (read as None) and ships a PRE-FOLDED
14 qk_scale_factor (43.784/sqrt(128)=3.87). Missing flags must read as True
15 (MuseGlimmer always applies QK-norm + output gate) and the query pre-scale
16 must normalize so native and modular land on the same number.
17 """
18
19 import math
20 from typing import Any
21
22 from vllm.model_executor.models.muse_glimmer import (
23 _muse_glimmer_query_prescale,
24 _muse_glimmer_use_attn_output_gate,
25 _muse_glimmer_use_qk_norm,
26 )
27 from vllm.transformers_utils.configs.muse_glimmer import MuseGlimmerConfig
28
29 # A representative FLAT config (Ruan rl_v1/hf shape), trimmed.
30 FLAT: dict[str, Any] = {
31 "architectures": ["MuseGlimmerForCausalLM"],
32 "model_type": "muse_glimmer",
33 "has_vision": True,
34 "bos_token_id": 200000,
35 "eos_token_id": 200001,
36 "vocab_size": 202048,
37 "hidden_size": 6656,
38 "intermediate_size": 19968,
39 "num_hidden_layers": 52,
40 "num_attention_heads": 32,
41 "num_key_value_heads": 2,
42 "head_dim": 128,
43 "hidden_act": "silu",
44 "max_position_embeddings": 16384,
45 "rms_norm_eps": 1e-5,
46 "post_norm_eps": 1e-8,
47 "qk_scale_factor": 43.7840518911,
48 "use_qk_norm": True,
49 "use_attn_output_gate": True,
50 "output_multiplier": 0.19611613513818404,
51 "output_soft_cap_temp": 20.0,
52 "normalize_tok_embeddings": True,
53 "rope_theta": 500000.0,
54 "sliding_window": 2048,
55 "patch_token_id": 200092,
56 "vision_latent_dim": 1536,
57 "vision_heads": 16,
58 "vision_layers": 50,
59 "vision_output_dim": 6144,
60 "vision_patch_size": 14,
61 "vision_patch_temporal": 2,
62 "vision_adapter_dim": 4096,
63 "vision_pos_emb_grid_h": 32,
64 "vision_pos_emb_grid_w": 32,
65 }
66
67 NESTED: dict[str, Any] = {
68 "architectures": ["MuseGlimmerForCausalLM"],
69 "model_type": "muse_glimmer",
70 "image_token_id": 200092,
71 "text_config": {
72 "model_type": "muse_glimmer_text",
73 "vocab_size": 202048,
74 "hidden_size": 6656,
75 "num_hidden_layers": 52,
76 "hidden_activation": "silu",
77 "final_logit_softcapping": 20.0,
78 "qk_scale_factor": 43.7840518911,
79 "rope_parameters": {"rope_type": "default", "rope_theta": 500000.0},
80 },
81 "vision_config": {
82 "model_type": "muse_glimmer_vision",
83 "hidden_size": 1536,
84 "num_hidden_layers": 50,
85 },
86 }
87
88 HEAD_DIM = 128
89 SQRT_HD = math.sqrt(HEAD_DIM)
90 NATIVE = 43.7840518911
91 FOLDED = NATIVE / SQRT_HD # 3.8700...
92
93
94 class Cfg:
95 """Attention-config stand-in carrying only the fields under test."""
96
97 def __init__(self, **kw):
98 self.head_dim = HEAD_DIM
99 for k, v in kw.items():
100 setattr(self, k, v)
101
102
103 # ------------------------------------------------------------ flat vs nested
104
105
106 def test_flat_config_values_respected():
107 c = MuseGlimmerConfig(**FLAT)
108 t = c.text_config
109 assert t.hidden_size == 6656
110 assert t.num_hidden_layers == 52
111 assert t.vocab_size == 202048
112 assert t.head_dim == 128
113 assert t.hidden_activation == "silu" # renamed from hidden_act
114 assert t.final_logit_softcapping == 20.0 # renamed from output_soft_cap_temp
115 assert abs(t.output_multiplier - 0.19611613513818404) < 1e-12
116 assert abs(t.qk_scale_factor - 43.7840518911) < 1e-9
117 assert t.rope_parameters["rope_theta"] == 500000.0 # from flat rope_theta
118 # vision hoisted + renamed
119 assert c.vision_config.hidden_size == 1536
120 assert c.vision_config.num_hidden_layers == 50
121 assert c.vision_config.output_dim == 6144
122 # flat patch_token_id -> image_token_id
123 assert c.image_token_id == 200092
124
125
126 def test_flat_config_no_silent_default():
127 # The regression: a non-default value MUST be honored, not silently dropped.
128 flat = dict(FLAT)
129 flat["hidden_size"] = 4096
130 flat["num_hidden_layers"] = 40
131 c = MuseGlimmerConfig(**flat)
132 assert c.text_config.hidden_size == 4096, "flat hidden_size silently ignored!"
133 assert c.text_config.num_hidden_layers == 40, "flat num_hidden_layers ignored!"
134
135
136 def test_nested_config_unchanged():
137 c = MuseGlimmerConfig(**NESTED)
138 assert c.text_config.hidden_size == 6656
139 assert c.text_config.hidden_activation == "silu"
140 assert c.text_config.final_logit_softcapping == 20.0
141 assert c.vision_config.hidden_size == 1536
142 assert c.vision_config.num_hidden_layers == 50
143 assert c.image_token_id == 200092
144
145
146 # -------------------------------------------------------- native vs modular
147
148
149 def test_qk_norm_missing_defaults_true():
150 assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=None)) is True # modular
151 assert _muse_glimmer_use_qk_norm(Cfg()) is True # absent
152 assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=True)) is True # native
153 assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=False)) is False # explicit off
154
155
156 def test_output_gate_missing_defaults_true():
157 assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=None)) is True
158 assert _muse_glimmer_use_attn_output_gate(Cfg()) is True
159 assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=False)) is False
160
161
162 def test_query_prescale_native_and_modular_converge():
163 # Both schemas must yield the SAME final scale_query_by (~3.87).
164 assert (
165 abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=NATIVE)) - FOLDED) < 1e-9
166 )
167 assert (
168 abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=FOLDED)) - FOLDED) < 1e-9
169 )
170
171
172 def test_query_prescale_explicit_wins():
173 c = Cfg(scale_query_by=FOLDED, qk_scale_factor=NATIVE)
174 assert abs(_muse_glimmer_query_prescale(c) - FOLDED) < 1e-9