| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MuseGlimmer model configuration for vLLM. |
| 4 | |
| 5 | Native vLLM copy of the MuseGlimmer HuggingFace configs |
| 6 | (``configuration_muse_glimmer.py``). MuseGlimmer's |
| 7 | ``model_type`` is not yet registered in released transformers, so vLLM ships |
| 8 | this config so partners can serve MuseGlimmer checkpoints *without* trust_remote_code. |
| 9 | |
| 10 | The field set and defaults are kept byte-for-byte in sync with the HF reference |
| 11 | so a checkpoint's ``config.json`` deserializes identically here. Both the text |
| 12 | and vision configs are consumed by the native multimodal serving path. |
| 13 | """ |
| 14 | |
| 15 | from __future__ import annotations |
| 16 | |
| 17 | from transformers import Qwen3Config |
| 18 | from transformers.configuration_utils import PreTrainedConfig |
| 19 | |
| 20 | |
| 21 | def _default_no_rope_layers(num_hidden_layers: int) -> list[int]: |
| 22 | # iRoPE mask: NoPE every 4 layers, counted backward from the last layer. |
| 23 | stride = 4 |
| 24 | return [ |
| 25 | 0 if (num_hidden_layers - 1 - i) % stride == 0 else 1 |
| 26 | for i in range(num_hidden_layers) |
| 27 | ] |
| 28 | |
| 29 | |
| 30 | class MuseGlimmerTextConfig(PreTrainedConfig): |
| 31 | model_type = "muse_glimmer_text" |
| 32 | keys_to_ignore_at_inference = ["past_key_values"] |
| 33 | |
| 34 | def __init__( |
| 35 | self, |
| 36 | vocab_size: int = 202_048, |
| 37 | hidden_size: int = 6656, |
| 38 | intermediate_size: int = 19968, |
| 39 | num_hidden_layers: int = 52, |
| 40 | num_attention_heads: int = 32, |
| 41 | num_key_value_heads: int = 2, |
| 42 | head_dim: int = 128, |
| 43 | hidden_activation: str = "silu", |
| 44 | max_position_embeddings: int = 16_384, |
| 45 | initializer_range: float = 0.02, |
| 46 | rms_norm_eps: float = 1e-5, |
| 47 | use_cache: bool = True, |
| 48 | pad_token_id: int | None = None, |
| 49 | eos_token_id: int | list[int] | None = 200_001, |
| 50 | bos_token_id: int | None = 200_000, |
| 51 | tie_word_embeddings: bool = False, |
| 52 | rope_parameters: dict | None = None, |
| 53 | rope_theta: float | None = None, |
| 54 | attention_bias: bool = False, |
| 55 | attention_dropout: float = 0.0, |
| 56 | query_pre_attn_scalar: int = 256, |
| 57 | sliding_window: int | None = 2048, |
| 58 | layer_types: list[str] | None = None, |
| 59 | final_logit_softcapping: float | None = 20.0, |
| 60 | attn_logit_softcapping: float | None = None, |
| 61 | use_bidirectional_attention: bool | None = None, |
| 62 | # MuseGlimmer-specific |
| 63 | qk_scale_factor: float = 43.7840518911, |
| 64 | use_qk_norm: bool = True, |
| 65 | use_attn_output_gate: bool = True, |
| 66 | output_multiplier: float = 0.19611613513818404, |
| 67 | normalize_tok_embeddings: bool = True, |
| 68 | post_norm_eps: float = 1e-8, |
| 69 | no_rope_layers: list[int] | None = None, |
| 70 | **kwargs, |
| 71 | ) -> None: |
| 72 | self.vocab_size = vocab_size |
| 73 | self.hidden_size = hidden_size |
| 74 | self.intermediate_size = intermediate_size |
| 75 | self.num_hidden_layers = num_hidden_layers |
| 76 | self.num_attention_heads = num_attention_heads |
| 77 | self.num_key_value_heads = num_key_value_heads |
| 78 | self.head_dim = head_dim |
| 79 | self.hidden_activation = hidden_activation |
| 80 | self.max_position_embeddings = max_position_embeddings |
| 81 | self.initializer_range = initializer_range |
| 82 | self.rms_norm_eps = rms_norm_eps |
| 83 | self.use_cache = use_cache |
| 84 | self.attention_bias = attention_bias |
| 85 | self.attention_dropout = attention_dropout |
| 86 | self.query_pre_attn_scalar = query_pre_attn_scalar |
| 87 | self.sliding_window = sliding_window |
| 88 | self.final_logit_softcapping = final_logit_softcapping |
| 89 | self.attn_logit_softcapping = attn_logit_softcapping |
| 90 | self.use_bidirectional_attention = use_bidirectional_attention |
| 91 | |
| 92 | # RoPE: accept either an explicit rope_parameters dict (HF 5.x) or a |
| 93 | # bare rope_theta; normalize to rope_parameters for vLLM's get_rope. |
| 94 | if rope_parameters is None: |
| 95 | theta = rope_theta if rope_theta is not None else 500_000.0 |
| 96 | rope_parameters = {"rope_type": "default", "rope_theta": theta} |
| 97 | self.rope_parameters = rope_parameters |
| 98 | # vLLM reads rope_theta off the config in some codepaths. |
| 99 | self.rope_theta = rope_parameters.get("rope_theta", 500_000.0) |
| 100 | |
| 101 | # MuseGlimmer-specific fields |
| 102 | self.qk_scale_factor = qk_scale_factor |
| 103 | self.use_qk_norm = use_qk_norm |
| 104 | self.use_attn_output_gate = use_attn_output_gate |
| 105 | self.output_multiplier = output_multiplier |
| 106 | self.normalize_tok_embeddings = normalize_tok_embeddings |
| 107 | self.post_norm_eps = post_norm_eps |
| 108 | |
| 109 | self.no_rope_layers = ( |
| 110 | no_rope_layers |
| 111 | if no_rope_layers is not None |
| 112 | else _default_no_rope_layers(num_hidden_layers) |
| 113 | ) |
| 114 | if layer_types is None: |
| 115 | layer_types = [ |
| 116 | "full_attention" if self.no_rope_layers[i] == 0 else "sliding_attention" |
| 117 | for i in range(num_hidden_layers) |
| 118 | ] |
| 119 | self.layer_types = layer_types |
| 120 | |
| 121 | super().__init__( |
| 122 | pad_token_id=pad_token_id, |
| 123 | bos_token_id=bos_token_id, |
| 124 | eos_token_id=eos_token_id, |
| 125 | tie_word_embeddings=tie_word_embeddings, |
| 126 | **kwargs, |
| 127 | ) |
| 128 | |
| 129 | |
| 130 | class MuseGlimmerVisionConfig(PreTrainedConfig): |
| 131 | model_type = "muse_glimmer_vision" |
| 132 | |
| 133 | def __init__( |
| 134 | self, |
| 135 | patch_size: int = 14, |
| 136 | pos_emb_height: int = 32, |
| 137 | pos_emb_width: int = 32, |
| 138 | num_attention_heads: int = 16, |
| 139 | num_hidden_layers: int = 50, |
| 140 | hidden_size: int = 1536, |
| 141 | intermediate_size: int = 8960, |
| 142 | hidden_act: str = "gelu", |
| 143 | merge_kernel_size: int = 2, |
| 144 | rope_parameters: dict | None = None, |
| 145 | max_position_embeddings: int = 32 * 32, |
| 146 | output_dim: int = 6144, |
| 147 | patch_temporal: int = 2, |
| 148 | adapter_dim: int = 4096, |
| 149 | layer_norm_eps: float = 1e-5, |
| 150 | layer_types: list[str] | None = None, |
| 151 | **kwargs, |
| 152 | ) -> None: |
| 153 | super().__init__(**kwargs) |
| 154 | self.patch_size = patch_size |
| 155 | self.pos_emb_height = pos_emb_height |
| 156 | self.pos_emb_width = pos_emb_width |
| 157 | self.num_attention_heads = num_attention_heads |
| 158 | self.num_hidden_layers = num_hidden_layers |
| 159 | self.hidden_size = hidden_size |
| 160 | self.intermediate_size = intermediate_size |
| 161 | self.hidden_act = hidden_act |
| 162 | self.merge_kernel_size = merge_kernel_size |
| 163 | self.rope_parameters = rope_parameters |
| 164 | self.max_position_embeddings = max_position_embeddings |
| 165 | self.output_dim = output_dim |
| 166 | self.patch_temporal = patch_temporal |
| 167 | self.adapter_dim = adapter_dim |
| 168 | self.layer_norm_eps = layer_norm_eps |
| 169 | if layer_types is None: |
| 170 | stride = 4 |
| 171 | layer_types = [ |
| 172 | "full_attention" |
| 173 | if (i + 1) % stride == 0 or i == num_hidden_layers - 1 |
| 174 | else "sliding_attention" |
| 175 | for i in range(num_hidden_layers) |
| 176 | ] |
| 177 | self.layer_types = layer_types |
| 178 | |
| 179 | |
| 180 | class MuseGlimmerConfig(PreTrainedConfig): |
| 181 | model_type = "muse_glimmer" |
| 182 | sub_configs = { |
| 183 | "text_config": MuseGlimmerTextConfig, |
| 184 | "vision_config": MuseGlimmerVisionConfig, |
| 185 | } |
| 186 | |
| 187 | # --- Flat (legacy-converter) -> canonical normalization ------------------- |
| 188 | # MuseGlimmer checkpoints exist in two config layouts in the wild: |
| 189 | # |
| 190 | # * CANONICAL (current HF converter, transformers 5.15): nested |
| 191 | # ``text_config`` / ``vision_config`` sub-dicts with canonical field |
| 192 | # names (``hidden_activation``, ``final_logit_softcapping``, |
| 193 | # ``rope_parameters``, ``vision_config.hidden_size`` ...). |
| 194 | # |
| 195 | # * FLAT (older converter, e.g. Ruan's ``rl_v1/hf``, transformers 5.9): |
| 196 | # every field is a top-level key, with different names |
| 197 | # (``hidden_act``, ``output_soft_cap_temp``, ``rope_theta``, |
| 198 | # ``vision_latent_dim`` ...) and NO ``text_config`` nesting. |
| 199 | # |
| 200 | # Without normalization a flat config silently deserializes to an |
| 201 | # ALL-DEFAULT text config (every checkpoint value ignored) — a dangerous |
| 202 | # correctness bug: a checkpoint whose arch differs from the defaults would |
| 203 | # load into a wrong-shaped model with no error. So when we detect a flat |
| 204 | # config we hoist the flat fields into ``text_config`` / ``vision_config`` |
| 205 | # and rename them to canonical names. |
| 206 | |
| 207 | # flat text field name -> canonical MuseGlimmerTextConfig field name |
| 208 | _FLAT_TEXT_RENAMES = { |
| 209 | "hidden_act": "hidden_activation", |
| 210 | "output_soft_cap_temp": "final_logit_softcapping", |
| 211 | } |
| 212 | # canonical MuseGlimmerTextConfig constructor params that may appear flat |
| 213 | _FLAT_TEXT_KEYS = frozenset( |
| 214 | { |
| 215 | "vocab_size", |
| 216 | "hidden_size", |
| 217 | "intermediate_size", |
| 218 | "num_hidden_layers", |
| 219 | "num_attention_heads", |
| 220 | "num_key_value_heads", |
| 221 | "head_dim", |
| 222 | "hidden_activation", |
| 223 | "hidden_act", |
| 224 | "max_position_embeddings", |
| 225 | "initializer_range", |
| 226 | "rms_norm_eps", |
| 227 | "use_cache", |
| 228 | "tie_word_embeddings", |
| 229 | "rope_parameters", |
| 230 | "rope_theta", |
| 231 | "attention_bias", |
| 232 | "attention_dropout", |
| 233 | "query_pre_attn_scalar", |
| 234 | "sliding_window", |
| 235 | "layer_types", |
| 236 | "final_logit_softcapping", |
| 237 | "output_soft_cap_temp", |
| 238 | "attn_logit_softcapping", |
| 239 | "use_bidirectional_attention", |
| 240 | "qk_scale_factor", |
| 241 | "use_qk_norm", |
| 242 | "use_attn_output_gate", |
| 243 | "output_multiplier", |
| 244 | "normalize_tok_embeddings", |
| 245 | "post_norm_eps", |
| 246 | "no_rope_layers", |
| 247 | } |
| 248 | ) |
| 249 | # flat vision field name -> canonical MuseGlimmerVisionConfig field name |
| 250 | _FLAT_VISION_RENAMES = { |
| 251 | "vision_latent_dim": "hidden_size", |
| 252 | "vision_heads": "num_attention_heads", |
| 253 | "vision_layers": "num_hidden_layers", |
| 254 | "vision_output_dim": "output_dim", |
| 255 | "vision_patch_size": "patch_size", |
| 256 | "vision_patch_temporal": "patch_temporal", |
| 257 | "vision_adapter_dim": "adapter_dim", |
| 258 | "vision_pos_emb_grid_h": "pos_emb_height", |
| 259 | "vision_pos_emb_grid_w": "pos_emb_width", |
| 260 | "vision_downsample_factor": "merge_kernel_size", |
| 261 | } |
| 262 | |
| 263 | @classmethod |
| 264 | def _looks_flat(cls, kwargs: dict) -> bool: |
| 265 | # Flat if there is no explicit text_config but there ARE text-level |
| 266 | # fields at the top level (e.g. hidden_size / num_hidden_layers). |
| 267 | if kwargs.get("text_config") is not None: |
| 268 | return False |
| 269 | return any(k in kwargs for k in ("hidden_size", "num_hidden_layers")) |
| 270 | |
| 271 | @classmethod |
| 272 | def _normalize_flat(cls, kwargs: dict) -> dict: |
| 273 | kwargs = dict(kwargs) |
| 274 | text: dict = {} |
| 275 | for key in list(kwargs.keys()): |
| 276 | if key in cls._FLAT_TEXT_KEYS: |
| 277 | canon = cls._FLAT_TEXT_RENAMES.get(key, key) |
| 278 | text.setdefault(canon, kwargs.pop(key)) |
| 279 | |
| 280 | vision_mlp_ratio = kwargs.pop("vision_mlp_ratio", None) |
| 281 | sparse_factor = kwargs.pop("vision_sparse_attention_factor", None) |
| 282 | vision: dict = {} |
| 283 | for key in list(kwargs.keys()): |
| 284 | if key in cls._FLAT_VISION_RENAMES: |
| 285 | vision.setdefault(cls._FLAT_VISION_RENAMES[key], kwargs.pop(key)) |
| 286 | |
| 287 | if vision_mlp_ratio is not None: |
| 288 | hidden_size = vision.get("hidden_size", 1536) |
| 289 | vision["intermediate_size"] = int(vision_mlp_ratio * hidden_size) |
| 290 | if sparse_factor is not None: |
| 291 | sparse_factor = int(sparse_factor) |
| 292 | if sparse_factor <= 0: |
| 293 | raise ValueError("vision_sparse_attention_factor must be positive") |
| 294 | num_layers = vision.get("num_hidden_layers", 50) |
| 295 | vision["layer_types"] = [ |
| 296 | "full_attention" |
| 297 | if (layer_idx + 1) % sparse_factor == 0 or layer_idx == num_layers - 1 |
| 298 | else "sliding_attention" |
| 299 | for layer_idx in range(num_layers) |
| 300 | ] |
| 301 | |
| 302 | if text: |
| 303 | kwargs["text_config"] = text |
| 304 | if vision: |
| 305 | kwargs["vision_config"] = vision |
| 306 | return kwargs |
| 307 | |
| 308 | def __init__( |
| 309 | self, |
| 310 | text_config: dict | MuseGlimmerTextConfig | None = None, |
| 311 | vision_config: dict | MuseGlimmerVisionConfig | None = None, |
| 312 | image_token_id: int = 200092, |
| 313 | video_token_id: int = 200091, |
| 314 | **kwargs, |
| 315 | ) -> None: |
| 316 | # Detect + fold a flat config into nested text/vision before building |
| 317 | # the sub-configs. (image/video token ids are handled below; a flat |
| 318 | # ``patch_token_id`` alias is mapped too.) |
| 319 | if text_config is None and self._looks_flat(kwargs): |
| 320 | kwargs = self._normalize_flat(kwargs) |
| 321 | text_config = kwargs.pop("text_config", None) |
| 322 | vision_config = kwargs.pop("vision_config", vision_config) |
| 323 | # flat image-token alias |
| 324 | if "patch_token_id" in kwargs and "image_token_id" not in kwargs: |
| 325 | image_token_id = kwargs.pop("patch_token_id") |
| 326 | |
| 327 | if text_config is None: |
| 328 | self.text_config = MuseGlimmerTextConfig() |
| 329 | elif isinstance(text_config, dict): |
| 330 | self.text_config = MuseGlimmerTextConfig(**text_config) |
| 331 | else: |
| 332 | self.text_config = text_config |
| 333 | |
| 334 | if vision_config is None: |
| 335 | self.vision_config = MuseGlimmerVisionConfig() |
| 336 | elif isinstance(vision_config, dict): |
| 337 | self.vision_config = MuseGlimmerVisionConfig(**vision_config) |
| 338 | else: |
| 339 | self.vision_config = vision_config |
| 340 | |
| 341 | self.image_token_id = image_token_id |
| 342 | self.video_token_id = video_token_id |
| 343 | self.patch_token_id = image_token_id |
| 344 | super().__init__(**kwargs) |
| 345 | |
| 346 | |
| 347 | class MuseGlimmerAssistantConfig(Qwen3Config): |
| 348 | """Config for the Muse Glimmer DFlash draft head. |
| 349 | |
| 350 | The head is Qwen3-shaped and runs on vLLM's generic ``qwen3_dflash`` |
| 351 | implementation, so this derives from ``Qwen3Config``. It cannot BE |
| 352 | ``Qwen3Config``, because two of the checkpoint's values do not survive that |
| 353 | class: |
| 354 | |
| 355 | * ``sliding_window`` is gated behind ``use_sliding_window``, which defaults |
| 356 | to False -- ``Qwen3Config(sliding_window=2048).sliding_window`` is None. |
| 357 | The checkpoint declares ``sliding_window: 2048`` and five |
| 358 | ``sliding_attention`` layers, so the window silently disappears and the |
| 359 | DFlash path then raises "sliding attention requires a window size". |
| 360 | * ``vocab_size`` is absent from the checkpoint, so Qwen3's default of |
| 361 | 151936 applies instead of Muse Glimmer's 202048. That one is *silent*: it |
| 362 | builds an all-zero ``draft_id_to_target_id`` remap, and also puts |
| 363 | pad/bos/eos/mask token ids out of range. |
| 364 | |
| 365 | Both defaults are set here so an unmodified checkpoint loads correctly. |
| 366 | """ |
| 367 | |
| 368 | model_type = "muse_glimmer_assistant" |
| 369 | |
| 370 | def __init__( |
| 371 | self, |
| 372 | vocab_size: int = 202048, |
| 373 | use_sliding_window: bool = True, |
| 374 | dflash_config: dict | None = None, |
| 375 | **kwargs, |
| 376 | ) -> None: |
| 377 | super().__init__( |
| 378 | vocab_size=vocab_size, use_sliding_window=use_sliding_window, **kwargs |
| 379 | ) |
| 380 | # muse_glimmer_assistant uses non-causal attention. |
| 381 | self.dflash_config = {"causal": False} | (dflash_config or {}) |
| 382 | |
| 383 | |
| 384 | __all__ = [ |
| 385 | "MuseGlimmerTextConfig", |
| 386 | "MuseGlimmerVisionConfig", |
| 387 | "MuseGlimmerConfig", |
| 388 | "MuseGlimmerAssistantConfig", |
| 389 | ] |