| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only MuseGlimmer multimodal model for vLLM. |
| 4 | |
| 5 | Native port of the MuseGlimmer text decoder (``MuseGlimmerForCausalLM``). The text |
| 6 | stack is a Gemma2 derivative with the |
| 7 | following MuseGlimmer-specific deltas, each of which is handled explicitly here: |
| 8 | |
| 9 | * SiLU-gated MLP (``hidden_activation="silu"``), not Gemma's gelu-tanh. |
| 10 | * Scaleless RMSNorm on the token embeddings (no sqrt(hidden) scaling). |
| 11 | * Per-layer sandwich RMSNorms with a baked ``+1`` weight offset |
| 12 | (``x * (1 + w)``), matching Gemma, but with distinct eps for the |
| 13 | pre/post norms (``rms_norm_eps`` vs ``post_norm_eps``). |
| 14 | * QK-norm (weightless, fp32) applied *before* RoPE, followed by a query |
| 15 | pre-scale of ``qk_scale_factor / sqrt(head_dim)``. |
| 16 | * A per-head sigmoid attention output gate. |
| 17 | * iRoPE layout: NoPE layers use full attention, RoPE layers use sliding |
| 18 | window attention. RoPE is applied NEOX-style (``is_neox_style=True``): |
| 19 | the HF converter (``convert_muse_glimmer_weights_to_hf.py``, 20260806+) permutes |
| 20 | q/k into the half-split (NEOX) layout via ``_permute_for_rope`` so they |
| 21 | pair with ``rotate_half`` — matching the reference's interleaved rotation |
| 22 | on the *native* (unpermuted) weights. Serving the permuted HF weights with |
| 23 | ``is_neox_style=False`` scrambles q/k and causes token-repetition collapse. |
| 24 | * Final logits are pre-scaled by ``output_multiplier`` and then tanh |
| 25 | soft-capped at ``final_logit_softcapping``. |
| 26 | * Untied lm_head. |
| 27 | |
| 28 | The vision path supports variable-resolution images and temporally patched |
| 29 | videos. It mirrors the checkpoint's native vision encoder, including sparse |
| 30 | block attention, 2-D RoPE, pixel-shuffle downsampling, and the two-layer |
| 31 | adapter/projection stack. |
| 32 | """ |
| 33 | |
| 34 | import math |
| 35 | from collections.abc import Iterable, Mapping, Sequence |
| 36 | from itertools import islice |
| 37 | from typing import Annotated, Literal |
| 38 | |
| 39 | import numpy as np |
| 40 | import torch |
| 41 | import torch.nn.functional as F |
| 42 | from PIL import Image |
| 43 | from torch import nn |
| 44 | from transformers import BatchFeature |
| 45 | |
| 46 | from vllm.compilation.decorators import support_torch_compile |
| 47 | from vllm.config import CacheConfig, VllmConfig |
| 48 | from vllm.config.multimodal import MultiModalDummyOptions |
| 49 | from vllm.distributed import ( |
| 50 | divide, |
| 51 | get_pp_group, |
| 52 | get_tensor_model_parallel_world_size, |
| 53 | ) |
| 54 | from vllm.inputs import MultiModalDataDict |
| 55 | from vllm.logger import init_logger |
| 56 | from vllm.model_executor.layers.activation import SiluAndMul |
| 57 | from vllm.model_executor.layers.attention import Attention, MMEncoderAttention |
| 58 | from vllm.model_executor.layers.linear import ( |
| 59 | ColumnParallelLinear, |
| 60 | MergedColumnParallelLinear, |
| 61 | QKVParallelLinear, |
| 62 | RowParallelLinear, |
| 63 | ) |
| 64 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 65 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 66 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 67 | from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb |
| 68 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 69 | ParallelLMHead, |
| 70 | VocabParallelEmbedding, |
| 71 | ) |
| 72 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 73 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 74 | from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems |
| 75 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems |
| 76 | from vllm.multimodal.processing import ( |
| 77 | BaseDummyInputsBuilder, |
| 78 | BaseMultiModalProcessor, |
| 79 | BaseProcessingInfo, |
| 80 | PromptReplacement, |
| 81 | PromptUpdate, |
| 82 | PromptUpdateDetails, |
| 83 | ) |
| 84 | from vllm.renderers import TokenizeParams |
| 85 | from vllm.sequence import IntermediateTensors |
| 86 | from vllm.transformers_utils.processors.muse_glimmer import MuseGlimmerProcessor |
| 87 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 88 | |
| 89 | from .interfaces import ( |
| 90 | EagleModelMixin, |
| 91 | MultiModalEmbeddings, |
| 92 | SupportsEagle3, |
| 93 | SupportsLoRA, |
| 94 | SupportsMultiModal, |
| 95 | SupportsPP, |
| 96 | ) |
| 97 | from .utils import ( |
| 98 | AutoWeightsLoader, |
| 99 | WeightsMapper, |
| 100 | extract_layer_index, |
| 101 | make_empty_intermediate_tensors_factory, |
| 102 | make_layers, |
| 103 | maybe_prefix, |
| 104 | ) |
| 105 | from .vision import get_fp8_padded_hidden_size, is_vit_use_data_parallel |
| 106 | |
| 107 | logger = init_logger(__name__) |
| 108 | |
| 109 | |
| 110 | def _text_config(config): |
| 111 | """MuseGlimmer checkpoints may nest the text config under ``text_config`` |
| 112 | (multimodal ``MuseGlimmerConfig``) or expose it directly |
| 113 | (``MuseGlimmerTextConfig``).""" |
| 114 | return getattr(config, "text_config", config) |
| 115 | |
| 116 | |
| 117 | def _vision_config(config): |
| 118 | return getattr(config, "vision_config", config) |
| 119 | |
| 120 | |
| 121 | def _muse_glimmer_has_vision(config) -> bool: |
| 122 | has_vision = getattr(config, "has_vision", None) |
| 123 | if has_vision is not None: |
| 124 | return bool(has_vision) |
| 125 | return hasattr(config, "vision_config") |
| 126 | |
| 127 | |
| 128 | IMAGE_TOKEN = "<|patch|>" |
| 129 | IMAGE_PROCESSOR_TOKEN = "<|image|>" |
| 130 | VIDEO_TOKEN = "<|video|>" |
| 131 | |
| 132 | |
| 133 | class MuseGlimmerImagePixelInputs(TensorSchema): |
| 134 | """Batched variable-resolution image inputs.""" |
| 135 | |
| 136 | type: Literal["image_pixels"] |
| 137 | pixel_values: Annotated[ |
| 138 | torch.Tensor | list[torch.Tensor], |
| 139 | TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}), |
| 140 | ] |
| 141 | feature_sizes: Annotated[torch.Tensor, TensorShape("bn")] |
| 142 | |
| 143 | |
| 144 | class MuseGlimmerVideoPixelInputs(TensorSchema): |
| 145 | """Batched variable-length, variable-resolution video inputs.""" |
| 146 | |
| 147 | type: Literal["video_pixels"] |
| 148 | pixel_values: Annotated[ |
| 149 | torch.Tensor | list[torch.Tensor], |
| 150 | TensorShape( |
| 151 | "bn", |
| 152 | "ng", |
| 153 | "c", |
| 154 | "h", |
| 155 | "w", |
| 156 | dynamic_dims={"ng", "h", "w"}, |
| 157 | ), |
| 158 | ] |
| 159 | feature_sizes: Annotated[torch.Tensor, TensorShape("bn")] |
| 160 | |
| 161 | |
| 162 | class MuseGlimmerProcessingInfo(BaseProcessingInfo): |
| 163 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 164 | if not _muse_glimmer_has_vision(self.get_hf_config()): |
| 165 | return {} |
| 166 | return {"image": None, "video": None} |
| 167 | |
| 168 | def get_default_tok_params(self) -> TokenizeParams: |
| 169 | return super().get_default_tok_params().with_kwargs(add_special_tokens=False) |
| 170 | |
| 171 | def get_hf_processor(self, **kwargs: object) -> MuseGlimmerProcessor: |
| 172 | return self.ctx.init_processor( |
| 173 | MuseGlimmerProcessor, |
| 174 | tokenizer=self.get_tokenizer(), |
| 175 | **kwargs, |
| 176 | ) |
| 177 | |
| 178 | def get_mm_max_tokens_per_item( |
| 179 | self, |
| 180 | seq_len: int, |
| 181 | mm_counts: Mapping[str, int], |
| 182 | ) -> Mapping[str, int]: |
| 183 | processor = self.get_hf_processor() |
| 184 | image_processor = processor.image_processor |
| 185 | video_processor = processor.video_processor |
| 186 | num_frames = self.get_num_frames_with_most_features(seq_len, mm_counts) |
| 187 | return { |
| 188 | "image": int(image_processor.max_image_tokens), |
| 189 | "video": ( |
| 190 | num_frames |
| 191 | // int(video_processor.patch_temporal) |
| 192 | * int(video_processor.max_video_frame_tokens) |
| 193 | ), |
| 194 | } |
| 195 | |
| 196 | def get_image_size_with_most_features(self) -> ImageSize: |
| 197 | image_processor = self.get_hf_processor().image_processor |
| 198 | grid_size = math.isqrt(int(image_processor.max_image_tokens)) |
| 199 | side = ( |
| 200 | int(image_processor.patch_size) |
| 201 | * int(image_processor.downsample_factor) |
| 202 | * grid_size |
| 203 | ) |
| 204 | return ImageSize(width=side, height=side) |
| 205 | |
| 206 | def get_num_frames_with_most_features( |
| 207 | self, seq_len: int, mm_counts: Mapping[str, int] |
| 208 | ) -> int: |
| 209 | video_processor = self.get_hf_processor().video_processor |
| 210 | num_videos = max(mm_counts.get("video", 0), 1) |
| 211 | groups = max( |
| 212 | 1, |
| 213 | seq_len // num_videos // int(video_processor.max_video_frame_tokens), |
| 214 | ) |
| 215 | patch_temporal = int(video_processor.patch_temporal) |
| 216 | max_groups = max( |
| 217 | 1, |
| 218 | int(video_processor.video_num_frames) // patch_temporal, |
| 219 | ) |
| 220 | return min(groups, max_groups) * patch_temporal |
| 221 | |
| 222 | |
| 223 | class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessingInfo]): |
| 224 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 225 | return IMAGE_TOKEN * mm_counts.get("image", 0) + VIDEO_TOKEN * mm_counts.get( |
| 226 | "video", 0 |
| 227 | ) |
| 228 | |
| 229 | def get_dummy_mm_data( |
| 230 | self, |
| 231 | seq_len: int, |
| 232 | mm_counts: Mapping[str, int], |
| 233 | mm_options: MultiModalDummyOptions, |
| 234 | ) -> MultiModalDataDict: |
| 235 | processor = self.info.get_hf_processor() |
| 236 | video_processor = processor.video_processor |
| 237 | image_width, image_height = self.info.get_image_size_with_most_features() |
| 238 | video_grid = math.isqrt(int(video_processor.max_video_frame_tokens)) |
| 239 | video_size = ( |
| 240 | int(video_processor.patch_size) |
| 241 | * int(video_processor.downsample_factor) |
| 242 | * video_grid |
| 243 | ) |
| 244 | return { |
| 245 | "image": self._get_dummy_images( |
| 246 | width=image_width, |
| 247 | height=image_height, |
| 248 | num_images=mm_counts.get("image", 0), |
| 249 | overrides=mm_options.get("image"), |
| 250 | ), |
| 251 | "video": self._get_dummy_videos( |
| 252 | width=video_size, |
| 253 | height=video_size, |
| 254 | num_frames=self.info.get_num_frames_with_most_features( |
| 255 | seq_len, mm_counts |
| 256 | ), |
| 257 | num_videos=mm_counts.get("video", 0), |
| 258 | overrides=mm_options.get("video"), |
| 259 | ), |
| 260 | } |
| 261 | |
| 262 | |
| 263 | class MuseGlimmerMultiModalProcessor( |
| 264 | BaseMultiModalProcessor[MuseGlimmerProcessingInfo] |
| 265 | ): |
| 266 | def _apply_hf_processor_main( |
| 267 | self, |
| 268 | mm_items: MultiModalDataItems, |
| 269 | hf_kwargs: Mapping[str, object], |
| 270 | ) -> BatchFeature: |
| 271 | mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( |
| 272 | mm_items, hf_kwargs |
| 273 | ) |
| 274 | |
| 275 | if not mm_data: |
| 276 | return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data) |
| 277 | |
| 278 | prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) |
| 279 | |
| 280 | processor = self.info.get_hf_processor(**hf_kwargs) |
| 281 | tokenizer = processor.tokenizer |
| 282 | config = self.info.get_hf_config() |
| 283 | images = mm_data.get("images", ()) |
| 284 | videos = mm_data.get("videos", ()) |
| 285 | if not isinstance(images, Sequence) or not isinstance(videos, Sequence): |
| 286 | raise TypeError("MuseGlimmer multi-modal data must be a sequence") |
| 287 | |
| 288 | prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False) |
| 289 | image_sentinel_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN) |
| 290 | video_sentinel_id = tokenizer.convert_tokens_to_ids(VIDEO_TOKEN) |
| 291 | if prompt_ids.count(image_sentinel_id) != len(images): |
| 292 | raise ValueError("The number of image placeholders does not match images") |
| 293 | if prompt_ids.count(video_sentinel_id) != len(videos): |
| 294 | raise ValueError("The number of video placeholders does not match videos") |
| 295 | |
| 296 | image_pixels: list[torch.Tensor] = [] |
| 297 | image_sizes: list[int] = [] |
| 298 | image_prompt_blocks: list[list[int]] = [] |
| 299 | for image in images: |
| 300 | output = processor( |
| 301 | text=IMAGE_PROCESSOR_TOKEN, |
| 302 | images=[image], |
| 303 | return_tensors=None, |
| 304 | ) |
| 305 | block_ids = output["input_ids"][0] |
| 306 | if isinstance(block_ids, torch.Tensor): |
| 307 | block_ids = block_ids.tolist() |
| 308 | pixels = output["pixel_values"] |
| 309 | if isinstance(pixels, torch.Tensor): |
| 310 | pixels = [pixels] if pixels.ndim == 3 else list(pixels) |
| 311 | if len(pixels) != 1: |
| 312 | raise ValueError( |
| 313 | "MuseGlimmer HF processor must return one tensor per image" |
| 314 | ) |
| 315 | image_pixels.append(pixels[0]) |
| 316 | image_sizes.append(block_ids.count(config.image_token_id)) |
| 317 | image_prompt_blocks.append(block_ids) |
| 318 | |
| 319 | video_pixels: list[torch.Tensor] = [] |
| 320 | video_sizes: list[int] = [] |
| 321 | video_prompt_blocks: list[list[int]] = [] |
| 322 | for video in videos: |
| 323 | if not isinstance(video, np.ndarray): |
| 324 | raise TypeError("MuseGlimmer video input must be a NumPy array") |
| 325 | frames = [Image.fromarray(frame).convert("RGB") for frame in video] |
| 326 | output = processor( |
| 327 | text=VIDEO_TOKEN, |
| 328 | videos=[frames], |
| 329 | return_tensors=None, |
| 330 | ) |
| 331 | block_ids = output["input_ids"][0] |
| 332 | if isinstance(block_ids, torch.Tensor): |
| 333 | block_ids = block_ids.tolist() |
| 334 | pixels = output["pixel_values"] |
| 335 | if isinstance(pixels, torch.Tensor): |
| 336 | pixels = list(pixels) |
| 337 | video_pixels.append(torch.stack(pixels)) |
| 338 | video_sizes.append(block_ids.count(config.video_token_id)) |
| 339 | video_prompt_blocks.append(block_ids) |
| 340 | |
| 341 | combined_prompt_ids: list[int] = [] |
| 342 | image_idx = video_idx = 0 |
| 343 | for token_id in prompt_ids: |
| 344 | if token_id == image_sentinel_id: |
| 345 | combined_prompt_ids.extend(image_prompt_blocks[image_idx]) |
| 346 | image_idx += 1 |
| 347 | elif token_id == video_sentinel_id: |
| 348 | combined_prompt_ids.extend(video_prompt_blocks[video_idx]) |
| 349 | video_idx += 1 |
| 350 | else: |
| 351 | combined_prompt_ids.append(token_id) |
| 352 | |
| 353 | data: dict[str, object] = { |
| 354 | "input_ids": [combined_prompt_ids], |
| 355 | } |
| 356 | if image_pixels: |
| 357 | data.update( |
| 358 | image_pixel_values=image_pixels, |
| 359 | image_feature_sizes=torch.tensor(image_sizes), |
| 360 | ) |
| 361 | if video_pixels: |
| 362 | data.update( |
| 363 | video_pixel_values=video_pixels, |
| 364 | video_feature_sizes=torch.tensor(video_sizes), |
| 365 | ) |
| 366 | processed_data = BatchFeature(data=data, tensor_type=None) |
| 367 | return self._finalize_hf_mm_data( |
| 368 | mm_data, hf_kwargs, passthrough_data, processed_data |
| 369 | ) |
| 370 | |
| 371 | def _get_mm_fields_config( |
| 372 | self, |
| 373 | hf_inputs: BatchFeature, |
| 374 | hf_processor_mm_kwargs: Mapping[str, object], |
| 375 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 376 | fields: dict[str, MultiModalFieldConfig] = {} |
| 377 | if "image_pixel_values" in hf_inputs: |
| 378 | fields.update( |
| 379 | image_pixel_values=MultiModalFieldConfig.batched("image"), |
| 380 | image_feature_sizes=MultiModalFieldConfig.batched( |
| 381 | "image", keep_on_cpu=True |
| 382 | ), |
| 383 | ) |
| 384 | if "video_pixel_values" in hf_inputs: |
| 385 | fields.update( |
| 386 | video_pixel_values=MultiModalFieldConfig.batched("video"), |
| 387 | video_feature_sizes=MultiModalFieldConfig.batched( |
| 388 | "video", keep_on_cpu=True |
| 389 | ), |
| 390 | ) |
| 391 | return fields |
| 392 | |
| 393 | def _get_prompt_updates( |
| 394 | self, |
| 395 | mm_items: MultiModalDataItems, |
| 396 | hf_processor_mm_kwargs: Mapping[str, object], |
| 397 | out_mm_kwargs: MultiModalKwargsItems, |
| 398 | ) -> Sequence[PromptUpdate]: |
| 399 | processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 400 | config = self.info.get_hf_config() |
| 401 | tokenizer = self.info.get_tokenizer() |
| 402 | vocab = tokenizer.get_vocab() |
| 403 | |
| 404 | placeholder = { |
| 405 | "image": vocab[IMAGE_TOKEN], |
| 406 | "video": vocab[VIDEO_TOKEN], |
| 407 | } |
| 408 | image_start_id = vocab["<|image_start|>"] |
| 409 | image_end_id = vocab["<|image_end|>"] |
| 410 | video_start_id = vocab["<|vid_start|>"] |
| 411 | video_end_id = vocab["<|vid_end|>"] |
| 412 | video_separator_id = vocab["<|vid_frame_separator|>"] |
| 413 | |
| 414 | def image_replacement(item_idx: int) -> PromptUpdateDetails: |
| 415 | out_item = out_mm_kwargs["image"][item_idx] |
| 416 | feature_size = out_item["image_feature_sizes"].data |
| 417 | assert isinstance(feature_size, torch.Tensor) |
| 418 | num_tokens = int(feature_size.item()) |
| 419 | replacement = ( |
| 420 | [image_start_id] + [config.image_token_id] * num_tokens + [image_end_id] |
| 421 | ) |
| 422 | return PromptUpdateDetails.select_token_id( |
| 423 | replacement, config.image_token_id |
| 424 | ) |
| 425 | |
| 426 | def video_replacement(item_idx: int) -> PromptUpdateDetails: |
| 427 | out_item = out_mm_kwargs["video"][item_idx] |
| 428 | pixel_values = out_item["video_pixel_values"].data |
| 429 | num_groups = len(pixel_values) |
| 430 | if num_groups == 0: |
| 431 | raise ValueError( |
| 432 | "MuseGlimmer video must contain at least one frame group" |
| 433 | ) |
| 434 | |
| 435 | feature_size = out_item["video_feature_sizes"].data |
| 436 | assert isinstance(feature_size, torch.Tensor) |
| 437 | num_tokens = int(feature_size.item()) |
| 438 | tokens_per_group, remainder = divmod(num_tokens, num_groups) |
| 439 | if remainder: |
| 440 | raise ValueError( |
| 441 | "MuseGlimmer video feature size must be divisible by frame groups" |
| 442 | ) |
| 443 | |
| 444 | video_processor = processor.video_processor |
| 445 | patch_temporal = int(video_processor.patch_temporal) |
| 446 | sampling_fps = float(video_processor.video_sampling_fps) |
| 447 | if sampling_fps <= 0: |
| 448 | raise ValueError("MuseGlimmer video_sampling_fps must be positive") |
| 449 | |
| 450 | replacement = [video_start_id] |
| 451 | for group_idx in range(num_groups): |
| 452 | timestamp = group_idx * patch_temporal / sampling_fps |
| 453 | replacement.extend( |
| 454 | tokenizer.encode( |
| 455 | f"Time: {timestamp:.1f}s", add_special_tokens=False |
| 456 | ) |
| 457 | ) |
| 458 | replacement.extend([config.video_token_id] * tokens_per_group) |
| 459 | replacement.append( |
| 460 | video_separator_id if group_idx < num_groups - 1 else video_end_id |
| 461 | ) |
| 462 | return PromptUpdateDetails.select_token_id( |
| 463 | replacement, config.video_token_id |
| 464 | ) |
| 465 | |
| 466 | return [ |
| 467 | PromptReplacement( |
| 468 | modality="image", |
| 469 | target=[placeholder["image"]], |
| 470 | replacement=image_replacement, |
| 471 | ), |
| 472 | PromptReplacement( |
| 473 | modality="video", |
| 474 | target=[placeholder["video"]], |
| 475 | replacement=video_replacement, |
| 476 | ), |
| 477 | ] |
| 478 | |
| 479 | |
| 480 | def _muse_glimmer_use_qk_norm(config) -> bool: |
| 481 | """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF |
| 482 | ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``). |
| 483 | Treat a missing/None flag as True — only an explicit ``False`` disables it.""" |
| 484 | val = getattr(config, "use_qk_norm", None) |
| 485 | return True if val is None else bool(val) |
| 486 | |
| 487 | |
| 488 | def _muse_glimmer_use_attn_output_gate(config) -> bool: |
| 489 | """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS |
| 490 | applies it; the modular HF ``text_config`` omits ``use_attn_output_gate`` |
| 491 | (reads as ``None``). Missing/None -> True; only explicit ``False`` disables.""" |
| 492 | val = getattr(config, "use_attn_output_gate", None) |
| 493 | return True if val is None else bool(val) |
| 494 | |
| 495 | |
| 496 | def _muse_glimmer_query_prescale(config) -> float: |
| 497 | """Post-QK-norm query pre-scale (``scale_query_by``), normalized across the |
| 498 | two config schemas so the net query scaling matches the native reference. |
| 499 | |
| 500 | HF native modeling computes ``scale_query_by = qk_scale_factor / sqrt(head_dim)`` |
| 501 | where the NATIVE ``qk_scale_factor`` is the raw ``params.json`` value |
| 502 | (~43.784). The modular HF ``text_config`` PRE-FOLDS the ``1/sqrt(head_dim)`` |
| 503 | factor and ships ``qk_scale_factor = 43.784 / sqrt(128) = 3.87`` already, |
| 504 | expecting it applied directly. Both must yield the SAME ``scale_query_by`` |
| 505 | (~3.87), then softmax uses ``scaling = head_dim**-0.5``. |
| 506 | |
| 507 | Precedence: |
| 508 | 1. explicit ``scale_query_by`` (already the final factor) -> use as-is. |
| 509 | 2. else derive from ``qk_scale_factor``: |
| 510 | - if it is already the folded value (``~= qk_scale_factor/sqrt(hd)`` is |
| 511 | NOT what we want; detect the native form and divide) — we decide by |
| 512 | magnitude: the native raw value is ``folded * sqrt(head_dim)``. If |
| 513 | ``qk_scale_factor`` is close to ``folded_expected * sqrt(hd)`` we treat |
| 514 | it as native and divide; otherwise it is already folded, use directly. |
| 515 | """ |
| 516 | head_dim = config.head_dim |
| 517 | sqrt_hd = head_dim**0.5 |
| 518 | |
| 519 | explicit = getattr(config, "scale_query_by", None) |
| 520 | if explicit is not None: |
| 521 | return float(explicit) |
| 522 | |
| 523 | qk_scale = getattr(config, "qk_scale_factor", None) |
| 524 | if qk_scale is None: |
| 525 | # No scale info at all: fall back to the plain 1/sqrt(head_dim) identity |
| 526 | # (net query scaling then just the softmax scaling); should not happen |
| 527 | # for real MuseGlimmer checkpoints, which always carry qk_scale_factor. |
| 528 | return 1.0 |
| 529 | |
| 530 | qk_scale = float(qk_scale) |
| 531 | # Disambiguate native (raw, ~43.78) vs modular (folded, ~3.87). The native |
| 532 | # form, when divided by sqrt(head_dim), yields the folded target; the folded |
| 533 | # form is already the target. Native values are ~sqrt(head_dim)x larger than |
| 534 | # folded. Use a threshold at sqrt(head_dim) (with margin): if qk_scale is |
| 535 | # comparable to or larger than sqrt(head_dim), it is the native raw value and |
| 536 | # must be divided; otherwise it is already folded and used directly. |
| 537 | # head_dim=128 -> sqrt=11.31; native 43.78 > 11.31 (divide -> 3.87), |
| 538 | # folded 3.87 < 11.31 (use as-is). |
| 539 | if qk_scale >= sqrt_hd: |
| 540 | return qk_scale / sqrt_hd |
| 541 | return qk_scale |
| 542 | |
| 543 | |
| 544 | class MuseGlimmerRMSNorm(nn.Module): |
| 545 | """RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end). |
| 546 | |
| 547 | ``normed = _norm(x.float()) * (w.float() + weight_offset)`` cast back to the |
| 548 | input dtype. When ``with_scale`` is False the layer is weightless (used for |
| 549 | QK-norm and the token-embedding norm). |
| 550 | """ |
| 551 | |
| 552 | def __init__( |
| 553 | self, |
| 554 | dim: int | None = None, |
| 555 | eps: float = 1e-6, |
| 556 | with_scale: bool = True, |
| 557 | weight_offset: int = 0, |
| 558 | ) -> None: |
| 559 | super().__init__() |
| 560 | self.eps = eps |
| 561 | self.with_scale = with_scale |
| 562 | self.weight_offset = weight_offset |
| 563 | if with_scale: |
| 564 | assert dim is not None |
| 565 | self.weight = nn.Parameter(torch.zeros(dim)) |
| 566 | else: |
| 567 | self.register_parameter("weight", None) |
| 568 | |
| 569 | def _norm(self, x: torch.Tensor) -> torch.Tensor: |
| 570 | return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| 571 | |
| 572 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 573 | out = self._norm(hidden_states.float()) |
| 574 | if self.with_scale: |
| 575 | out = out * (self.weight.float() + self.weight_offset) |
| 576 | return out.type_as(hidden_states) |
| 577 | |
| 578 | |
| 579 | class MuseGlimmerVisionAttention(nn.Module): |
| 580 | def __init__( |
| 581 | self, |
| 582 | hidden_size: int, |
| 583 | num_heads: int, |
| 584 | prefix: str = "", |
| 585 | ) -> None: |
| 586 | super().__init__() |
| 587 | if hidden_size % num_heads: |
| 588 | raise ValueError("MuseGlimmer vision hidden size must divide num heads") |
| 589 | |
| 590 | self.total_num_heads = num_heads |
| 591 | self.head_dim = hidden_size // num_heads |
| 592 | self.use_data_parallel = is_vit_use_data_parallel(num_heads) |
| 593 | self.tp_size = ( |
| 594 | 1 if self.use_data_parallel else get_tensor_model_parallel_world_size() |
| 595 | ) |
| 596 | self.num_heads = divide(num_heads, self.tp_size) |
| 597 | |
| 598 | self.qkv_proj = QKVParallelLinear( |
| 599 | hidden_size=hidden_size, |
| 600 | head_size=self.head_dim, |
| 601 | total_num_heads=num_heads, |
| 602 | total_num_kv_heads=num_heads, |
| 603 | bias=True, |
| 604 | prefix=f"{prefix}.qkv_proj", |
| 605 | disable_tp=self.use_data_parallel, |
| 606 | ) |
| 607 | self.o_proj = RowParallelLinear( |
| 608 | input_size=hidden_size, |
| 609 | output_size=hidden_size, |
| 610 | bias=True, |
| 611 | prefix=f"{prefix}.o_proj", |
| 612 | disable_tp=self.use_data_parallel, |
| 613 | ) |
| 614 | self.attn = MMEncoderAttention( |
| 615 | num_heads=self.num_heads, |
| 616 | head_size=self.head_dim, |
| 617 | scale=self.head_dim**-0.5, |
| 618 | prefix=f"{prefix}.attn", |
| 619 | ) |
| 620 | self.apply_rotary_emb = ApplyRotaryEmb( |
| 621 | enforce_enable=True, |
| 622 | is_neox_style=True, |
| 623 | enable_fp32_compute=True, |
| 624 | ) |
| 625 | |
| 626 | def forward( |
| 627 | self, |
| 628 | hidden_states: torch.Tensor, |
| 629 | rotary_pos_emb_cos: torch.Tensor, |
| 630 | rotary_pos_emb_sin: torch.Tensor, |
| 631 | cu_seqlens: torch.Tensor, |
| 632 | max_seqlen: torch.Tensor, |
| 633 | sequence_lengths: torch.Tensor | None, |
| 634 | ) -> torch.Tensor: |
| 635 | qkv, _ = self.qkv_proj(hidden_states) |
| 636 | qkv = qkv.view( |
| 637 | hidden_states.shape[0], |
| 638 | 3, |
| 639 | self.num_heads, |
| 640 | self.head_dim, |
| 641 | ) |
| 642 | query, key, value = qkv.unbind(1) |
| 643 | query, key = self.apply_rotary_emb( |
| 644 | torch.stack([query, key]).contiguous(), |
| 645 | rotary_pos_emb_cos, |
| 646 | rotary_pos_emb_sin, |
| 647 | ).unbind(0) |
| 648 | query = query.unsqueeze(0) |
| 649 | key = key.unsqueeze(0) |
| 650 | value = value.unsqueeze(0) |
| 651 | |
| 652 | output = self.attn( |
| 653 | query=query, |
| 654 | key=key, |
| 655 | value=value, |
| 656 | cu_seqlens=cu_seqlens, |
| 657 | max_seqlen=max_seqlen, |
| 658 | sequence_lengths=sequence_lengths, |
| 659 | ) |
| 660 | output = output.reshape(hidden_states.shape[0], -1) |
| 661 | output, _ = self.o_proj(output) |
| 662 | return output |
| 663 | |
| 664 | |
| 665 | class MuseGlimmerVisionMLP(nn.Module): |
| 666 | def __init__(self, hidden_size: int, intermediate_size: int) -> None: |
| 667 | super().__init__() |
| 668 | self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True) |
| 669 | self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True) |
| 670 | |
| 671 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 672 | return self.c_proj(F.gelu(self.c_fc(hidden_states))) |
| 673 | |
| 674 | |
| 675 | class MuseGlimmerVisionBlock(nn.Module): |
| 676 | def __init__( |
| 677 | self, |
| 678 | hidden_size: int, |
| 679 | num_heads: int, |
| 680 | intermediate_size: int, |
| 681 | prefix: str = "", |
| 682 | ) -> None: |
| 683 | super().__init__() |
| 684 | self.ln_1 = nn.LayerNorm(hidden_size) |
| 685 | self.attn = MuseGlimmerVisionAttention( |
| 686 | hidden_size, |
| 687 | num_heads, |
| 688 | prefix=f"{prefix}.attn", |
| 689 | ) |
| 690 | self.ln_2 = nn.LayerNorm(hidden_size) |
| 691 | self.mlp = MuseGlimmerVisionMLP(hidden_size, intermediate_size) |
| 692 | |
| 693 | def forward( |
| 694 | self, |
| 695 | hidden_states: torch.Tensor, |
| 696 | rotary_pos_emb_cos: torch.Tensor, |
| 697 | rotary_pos_emb_sin: torch.Tensor, |
| 698 | cu_seqlens: torch.Tensor, |
| 699 | max_seqlen: torch.Tensor, |
| 700 | sequence_lengths: torch.Tensor | None, |
| 701 | ) -> torch.Tensor: |
| 702 | batch_size, seq_len, hidden_size = hidden_states.shape |
| 703 | flattened = hidden_states.view(batch_size * seq_len, hidden_size) |
| 704 | flattened = flattened + self.attn( |
| 705 | self.ln_1(flattened), |
| 706 | rotary_pos_emb_cos, |
| 707 | rotary_pos_emb_sin, |
| 708 | cu_seqlens, |
| 709 | max_seqlen, |
| 710 | sequence_lengths, |
| 711 | ) |
| 712 | flattened = flattened + self.mlp(self.ln_2(flattened)) |
| 713 | return flattened.view(batch_size, seq_len, hidden_size) |
| 714 | |
| 715 | |
| 716 | class MuseGlimmerVisionEncoder(nn.Module): |
| 717 | def __init__(self, config, prefix: str = "") -> None: |
| 718 | super().__init__() |
| 719 | hidden_size = config.hidden_size |
| 720 | patch_dim = config.patch_temporal * 3 * config.patch_size**2 |
| 721 | self.hidden_size = hidden_size |
| 722 | self.patch_size = config.patch_size |
| 723 | self.patch_temporal = config.patch_temporal |
| 724 | self.merge_kernel_size = config.merge_kernel_size |
| 725 | self.pos_emb_height = config.pos_emb_height |
| 726 | self.pos_emb_width = config.pos_emb_width |
| 727 | self.head_dim = hidden_size // config.num_attention_heads |
| 728 | self.layer_types = list(config.layer_types) |
| 729 | if len(self.layer_types) != config.num_hidden_layers: |
| 730 | raise ValueError( |
| 731 | "MuseGlimmer vision layer_types must match num_hidden_layers" |
| 732 | ) |
| 733 | if self.head_dim % 4: |
| 734 | raise ValueError("MuseGlimmer vision head dimension must be divisible by 4") |
| 735 | |
| 736 | self.conv1_linear = nn.Linear(patch_dim, hidden_size, bias=False) |
| 737 | self.positional_embedding_vlm = nn.Parameter( |
| 738 | torch.zeros(config.pos_emb_height * config.pos_emb_width, hidden_size) |
| 739 | ) |
| 740 | self.ln_pre = nn.LayerNorm(hidden_size) |
| 741 | self.transformer = nn.ModuleList( |
| 742 | [ |
| 743 | MuseGlimmerVisionBlock( |
| 744 | hidden_size, |
| 745 | config.num_attention_heads, |
| 746 | config.intermediate_size, |
| 747 | prefix=f"{prefix}.transformer.{layer_idx}", |
| 748 | ) |
| 749 | for layer_idx in range(config.num_hidden_layers) |
| 750 | ] |
| 751 | ) |
| 752 | vision_attention = self.transformer[0].attn |
| 753 | self.tp_size = vision_attention.tp_size |
| 754 | self.attn_backend = vision_attention.attn.attn_backend |
| 755 | self.fp8_padded_hidden_size = get_fp8_padded_hidden_size( |
| 756 | config.num_attention_heads, self.head_dim |
| 757 | ) |
| 758 | self.ln_post = nn.LayerNorm(hidden_size) |
| 759 | |
| 760 | expected_output = hidden_size * self.merge_kernel_size**2 |
| 761 | if config.output_dim != expected_output: |
| 762 | raise ValueError( |
| 763 | f"MuseGlimmer vision output_dim={config.output_dim} does not match " |
| 764 | f"pixel-shuffle output {expected_output}" |
| 765 | ) |
| 766 | |
| 767 | def _make_2d_rope( |
| 768 | self, grid_height: int, grid_width: int, device: torch.device |
| 769 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 770 | spatial_dim = self.head_dim // 2 |
| 771 | inv_freq = 1.0 / ( |
| 772 | 10000.0 |
| 773 | ** ( |
| 774 | torch.arange(0, spatial_dim, 2, dtype=torch.float32, device=device) |
| 775 | / spatial_dim |
| 776 | ) |
| 777 | ) |
| 778 | height = torch.arange(1, grid_height + 1, dtype=torch.float32, device=device) |
| 779 | width = torch.arange(1, grid_width + 1, dtype=torch.float32, device=device) |
| 780 | height = height.unsqueeze(1).expand(-1, grid_width).reshape(-1) |
| 781 | width = width.unsqueeze(0).expand(grid_height, -1).reshape(-1) |
| 782 | freq_w = torch.outer(width, inv_freq) |
| 783 | freq_h = torch.outer(height, inv_freq) |
| 784 | freqs = torch.cat([freq_w, freq_h], dim=-1) |
| 785 | return torch.cos(freqs), torch.sin(freqs) |
| 786 | |
| 787 | def _get_pos_emb( |
| 788 | self, grid_height: int, grid_width: int, device: torch.device |
| 789 | ) -> torch.Tensor: |
| 790 | h_grid = ( |
| 791 | torch.arange(grid_height, device=device, dtype=torch.float32) + 0.5 |
| 792 | ) * (self.pos_emb_height / grid_height) - 0.5 |
| 793 | w_grid = ( |
| 794 | torch.arange(grid_width, device=device, dtype=torch.float32) + 0.5 |
| 795 | ) * (self.pos_emb_width / grid_width) - 0.5 |
| 796 | h_floor = torch.floor(h_grid).long() |
| 797 | w_floor = torch.floor(w_grid).long() |
| 798 | h_ceil = h_floor + 1 |
| 799 | w_ceil = w_floor + 1 |
| 800 | h_frac = h_grid - h_floor.float() |
| 801 | w_frac = w_grid - w_floor.float() |
| 802 | |
| 803 | h_floor_valid = (h_floor >= 0) & (h_floor < self.pos_emb_height) |
| 804 | h_ceil_valid = (h_ceil >= 0) & (h_ceil < self.pos_emb_height) |
| 805 | w_floor_valid = (w_floor >= 0) & (w_floor < self.pos_emb_width) |
| 806 | w_ceil_valid = (w_ceil >= 0) & (w_ceil < self.pos_emb_width) |
| 807 | h_floor = h_floor.clamp(0, self.pos_emb_height - 1) |
| 808 | h_ceil = h_ceil.clamp(0, self.pos_emb_height - 1) |
| 809 | w_floor = w_floor.clamp(0, self.pos_emb_width - 1) |
| 810 | w_ceil = w_ceil.clamp(0, self.pos_emb_width - 1) |
| 811 | |
| 812 | h_floor_offset = h_floor * self.pos_emb_width |
| 813 | h_ceil_offset = h_ceil * self.pos_emb_width |
| 814 | indices = torch.stack( |
| 815 | [ |
| 816 | (h_floor_offset[:, None] + w_floor[None, :]).flatten(), |
| 817 | (h_floor_offset[:, None] + w_ceil[None, :]).flatten(), |
| 818 | (h_ceil_offset[:, None] + w_floor[None, :]).flatten(), |
| 819 | (h_ceil_offset[:, None] + w_ceil[None, :]).flatten(), |
| 820 | ] |
| 821 | ) |
| 822 | weights = torch.stack( |
| 823 | [ |
| 824 | ( |
| 825 | (1 - h_frac)[:, None] |
| 826 | * (1 - w_frac)[None, :] |
| 827 | * (h_floor_valid[:, None] & w_floor_valid[None, :]) |
| 828 | ).flatten(), |
| 829 | ( |
| 830 | (1 - h_frac)[:, None] |
| 831 | * w_frac[None, :] |
| 832 | * (h_floor_valid[:, None] & w_ceil_valid[None, :]) |
| 833 | ).flatten(), |
| 834 | ( |
| 835 | h_frac[:, None] |
| 836 | * (1 - w_frac)[None, :] |
| 837 | * (h_ceil_valid[:, None] & w_floor_valid[None, :]) |
| 838 | ).flatten(), |
| 839 | ( |
| 840 | h_frac[:, None] |
| 841 | * w_frac[None, :] |
| 842 | * (h_ceil_valid[:, None] & w_ceil_valid[None, :]) |
| 843 | ).flatten(), |
| 844 | ] |
| 845 | ) |
| 846 | return (self.positional_embedding_vlm[indices] * weights[..., None]).sum(0) |
| 847 | |
| 848 | def _pixel_shuffle_downsample( |
| 849 | self, hidden_states: torch.Tensor, grid_height: int, grid_width: int |
| 850 | ) -> torch.Tensor: |
| 851 | factor = self.merge_kernel_size |
| 852 | output_tokens = (grid_height // factor) * (grid_width // factor) |
| 853 | permutation = torch.arange( |
| 854 | grid_height * grid_width, device=hidden_states.device |
| 855 | ) |
| 856 | permutation = permutation.view( |
| 857 | grid_height // factor, factor, grid_width // factor, factor |
| 858 | ) |
| 859 | permutation = permutation.permute(0, 2, 1, 3).reshape(-1) |
| 860 | hidden_states = hidden_states.squeeze(0)[permutation] |
| 861 | hidden_size = hidden_states.shape[-1] |
| 862 | hidden_states = ( |
| 863 | hidden_states.view(output_tokens, factor * factor, hidden_size) |
| 864 | .permute(0, 2, 1) |
| 865 | .contiguous() |
| 866 | .view(output_tokens, hidden_size * factor * factor) |
| 867 | ) |
| 868 | return hidden_states.unsqueeze(0) |
| 869 | |
| 870 | def _get_sparse_permutation( |
| 871 | self, grid_height: int, grid_width: int, device: torch.device |
| 872 | ) -> tuple[torch.Tensor, list[int]]: |
| 873 | block_height = self.pos_emb_height |
| 874 | block_width = self.pos_emb_width |
| 875 | padded_height = math.ceil(grid_height / block_height) * block_height |
| 876 | padded_width = math.ceil(grid_width / block_width) * block_width |
| 877 | indices = torch.arange(grid_height * grid_width, device=device).view( |
| 878 | grid_height, grid_width |
| 879 | ) |
| 880 | indices = F.pad( |
| 881 | indices, |
| 882 | (0, padded_width - grid_width, 0, padded_height - grid_height), |
| 883 | value=-1, |
| 884 | ).flatten() |
| 885 | indices = indices.view( |
| 886 | padded_height // block_height, |
| 887 | block_height, |
| 888 | padded_width // block_width, |
| 889 | block_width, |
| 890 | ) |
| 891 | indices = indices.permute(0, 2, 1, 3).reshape(-1) |
| 892 | valid = (indices != -1).view(-1, block_height * block_width) |
| 893 | return indices[indices != -1], valid.sum(dim=1).tolist() |
| 894 | |
| 895 | def _get_attention_metadata( |
| 896 | self, |
| 897 | seq_lens: Sequence[int], |
| 898 | device: torch.device, |
| 899 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 900 | cu_seqlens_np = np.concatenate( |
| 901 | [ |
| 902 | np.zeros(1, dtype=np.int32), |
| 903 | np.asarray(seq_lens, dtype=np.int32).cumsum(dtype=np.int32), |
| 904 | ] |
| 905 | ) |
| 906 | sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( |
| 907 | self.attn_backend, |
| 908 | cu_seqlens_np, |
| 909 | device, |
| 910 | ) |
| 911 | max_seqlen = torch.tensor( |
| 912 | MMEncoderAttention.compute_max_seqlen( |
| 913 | self.attn_backend, |
| 914 | cu_seqlens_np, |
| 915 | ), |
| 916 | dtype=torch.int32, |
| 917 | ) |
| 918 | cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( |
| 919 | self.attn_backend, |
| 920 | cu_seqlens_np, |
| 921 | self.hidden_size, |
| 922 | self.tp_size, |
| 923 | device, |
| 924 | fp8_padded_hidden_size=self.fp8_padded_hidden_size, |
| 925 | ) |
| 926 | return cu_seqlens, max_seqlen, sequence_lengths |
| 927 | |
| 928 | def _patchify(self, pixels: torch.Tensor) -> torch.Tensor: |
| 929 | patch_size = self.patch_size |
| 930 | _, channels, height, width = pixels.shape |
| 931 | grid_height = height // patch_size |
| 932 | grid_width = width // patch_size |
| 933 | if channels == 3: |
| 934 | patches = pixels.unfold(2, patch_size, patch_size).unfold( |
| 935 | 3, patch_size, patch_size |
| 936 | ) |
| 937 | patches = patches.contiguous().view( |
| 938 | 1, channels, grid_height, grid_width, patch_size, patch_size |
| 939 | ) |
| 940 | patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous() |
| 941 | patches = patches.unsqueeze(3).expand( |
| 942 | -1, -1, -1, self.patch_temporal, -1, -1, -1 |
| 943 | ) |
| 944 | elif channels == self.patch_temporal * 3: |
| 945 | frame_patches = [] |
| 946 | for frame_idx in range(self.patch_temporal): |
| 947 | frame = pixels[:, frame_idx * 3 : (frame_idx + 1) * 3] |
| 948 | frame = frame.unfold(2, patch_size, patch_size).unfold( |
| 949 | 3, patch_size, patch_size |
| 950 | ) |
| 951 | frame = frame.contiguous().view( |
| 952 | 1, 3, grid_height, grid_width, patch_size, patch_size |
| 953 | ) |
| 954 | frame_patches.append(frame.permute(0, 2, 3, 1, 4, 5).contiguous()) |
| 955 | patches = torch.stack(frame_patches, dim=3) |
| 956 | else: |
| 957 | raise ValueError( |
| 958 | f"MuseGlimmer vision input has {channels} channels; expected 3 or " |
| 959 | f"{self.patch_temporal * 3}" |
| 960 | ) |
| 961 | return patches.reshape(1, grid_height * grid_width, -1) |
| 962 | |
| 963 | def forward(self, pixel_values: Sequence[torch.Tensor]) -> torch.Tensor: |
| 964 | device = self.conv1_linear.weight.device |
| 965 | dtype = self.conv1_linear.weight.dtype |
| 966 | has_sparse_layers = any( |
| 967 | layer_type != "full_attention" for layer_type in self.layer_types |
| 968 | ) |
| 969 | |
| 970 | all_hidden_states = [] |
| 971 | all_rotary_pos_emb_cos = [] |
| 972 | all_rotary_pos_emb_sin = [] |
| 973 | sparse_seq_lens: list[int] = [] |
| 974 | global_seq_lens: list[int] = [] |
| 975 | metadata = [] |
| 976 | for pixels in pixel_values: |
| 977 | pixels = pixels.to(device=device, dtype=dtype) |
| 978 | if pixels.ndim == 3: |
| 979 | pixels = pixels.unsqueeze(0) |
| 980 | if pixels.ndim != 4 or pixels.shape[0] != 1: |
| 981 | raise ValueError( |
| 982 | "Each MuseGlimmer vision input must have shape [C, H, W]" |
| 983 | ) |
| 984 | spatial_stride = self.patch_size * self.merge_kernel_size |
| 985 | if pixels.shape[-2] % spatial_stride or pixels.shape[-1] % spatial_stride: |
| 986 | raise ValueError( |
| 987 | "MuseGlimmer vision input dimensions must divide the " |
| 988 | "merged patch stride" |
| 989 | ) |
| 990 | grid_height = pixels.shape[-2] // self.patch_size |
| 991 | grid_width = pixels.shape[-1] // self.patch_size |
| 992 | num_tokens = grid_height * grid_width |
| 993 | hidden_states = self.conv1_linear(self._patchify(pixels)) |
| 994 | hidden_states = hidden_states + self._get_pos_emb( |
| 995 | grid_height, grid_width, device |
| 996 | ).unsqueeze(0).to(dtype) |
| 997 | hidden_states = self.ln_pre(hidden_states.view(-1, self.hidden_size)).view( |
| 998 | 1, -1, self.hidden_size |
| 999 | ) |
| 1000 | rotary_pos_emb_cos, rotary_pos_emb_sin = self._make_2d_rope( |
| 1001 | grid_height, grid_width, device |
| 1002 | ) |
| 1003 | |
| 1004 | permutation = None |
| 1005 | if has_sparse_layers: |
| 1006 | permutation, seq_lens = self._get_sparse_permutation( |
| 1007 | grid_height, grid_width, device |
| 1008 | ) |
| 1009 | hidden_states = hidden_states[:, permutation] |
| 1010 | rotary_pos_emb_cos = rotary_pos_emb_cos[permutation] |
| 1011 | rotary_pos_emb_sin = rotary_pos_emb_sin[permutation] |
| 1012 | sparse_seq_lens.extend(seq_lens) |
| 1013 | |
| 1014 | all_hidden_states.append(hidden_states.squeeze(0)) |
| 1015 | all_rotary_pos_emb_cos.append(rotary_pos_emb_cos) |
| 1016 | all_rotary_pos_emb_sin.append(rotary_pos_emb_sin) |
| 1017 | global_seq_lens.append(num_tokens) |
| 1018 | metadata.append((grid_height, grid_width, num_tokens, permutation)) |
| 1019 | |
| 1020 | hidden_states = torch.cat(all_hidden_states).unsqueeze(0) |
| 1021 | rotary_pos_emb_cos = torch.cat(all_rotary_pos_emb_cos) |
| 1022 | rotary_pos_emb_sin = torch.cat(all_rotary_pos_emb_sin) |
| 1023 | global_attention_metadata = self._get_attention_metadata( |
| 1024 | global_seq_lens, |
| 1025 | device, |
| 1026 | ) |
| 1027 | sparse_attention_metadata = ( |
| 1028 | self._get_attention_metadata(sparse_seq_lens, device) |
| 1029 | if sparse_seq_lens |
| 1030 | else None |
| 1031 | ) |
| 1032 | for layer_type, block in zip(self.layer_types, self.transformer): |
| 1033 | attention_metadata = ( |
| 1034 | global_attention_metadata |
| 1035 | if layer_type == "full_attention" |
| 1036 | else sparse_attention_metadata |
| 1037 | ) |
| 1038 | if attention_metadata is None: |
| 1039 | raise ValueError("MuseGlimmer sparse attention metadata is missing") |
| 1040 | hidden_states = block( |
| 1041 | hidden_states, |
| 1042 | rotary_pos_emb_cos, |
| 1043 | rotary_pos_emb_sin, |
| 1044 | *attention_metadata, |
| 1045 | ) |
| 1046 | |
| 1047 | features = [] |
| 1048 | offset = 0 |
| 1049 | for grid_height, grid_width, num_tokens, permutation in metadata: |
| 1050 | item = hidden_states[:, offset : offset + num_tokens] |
| 1051 | offset += num_tokens |
| 1052 | if permutation is not None: |
| 1053 | inverse = torch.empty_like(permutation) |
| 1054 | inverse[permutation] = torch.arange(len(permutation), device=device) |
| 1055 | item = item[:, inverse] |
| 1056 | item = self.ln_post(item.view(-1, self.hidden_size)).view( |
| 1057 | 1, -1, self.hidden_size |
| 1058 | ) |
| 1059 | features.append( |
| 1060 | self._pixel_shuffle_downsample(item, grid_height, grid_width).squeeze(0) |
| 1061 | ) |
| 1062 | return torch.cat(features) |
| 1063 | |
| 1064 | |
| 1065 | class MuseGlimmerVisionAdapter(nn.Module): |
| 1066 | def __init__(self, config) -> None: |
| 1067 | super().__init__() |
| 1068 | self.c_fc = nn.Linear(config.output_dim, config.adapter_dim, bias=False) |
| 1069 | self.c_proj = nn.Linear(config.adapter_dim, config.adapter_dim, bias=False) |
| 1070 | |
| 1071 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 1072 | return F.gelu(self.c_proj(F.gelu(self.c_fc(hidden_states)))) |
| 1073 | |
| 1074 | |
| 1075 | class MuseGlimmerMLP(nn.Module): |
| 1076 | def __init__( |
| 1077 | self, |
| 1078 | hidden_size: int, |
| 1079 | intermediate_size: int, |
| 1080 | hidden_activation: str, |
| 1081 | quant_config: QuantizationConfig | None = None, |
| 1082 | prefix: str = "", |
| 1083 | ) -> None: |
| 1084 | super().__init__() |
| 1085 | self.gate_up_proj = MergedColumnParallelLinear( |
| 1086 | hidden_size, |
| 1087 | [intermediate_size] * 2, |
| 1088 | bias=False, |
| 1089 | quant_config=quant_config, |
| 1090 | prefix=f"{prefix}.gate_up_proj", |
| 1091 | ) |
| 1092 | self.down_proj = RowParallelLinear( |
| 1093 | intermediate_size, |
| 1094 | hidden_size, |
| 1095 | bias=False, |
| 1096 | quant_config=quant_config, |
| 1097 | prefix=f"{prefix}.down_proj", |
| 1098 | ) |
| 1099 | if hidden_activation != "silu": |
| 1100 | raise ValueError( |
| 1101 | f"MuseGlimmer uses `silu` as the hidden activation; " |
| 1102 | f"got `{hidden_activation}`." |
| 1103 | ) |
| 1104 | self.act_fn = SiluAndMul() |
| 1105 | |
| 1106 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 1107 | gate_up, _ = self.gate_up_proj(x) |
| 1108 | x = self.act_fn(gate_up) |
| 1109 | x, _ = self.down_proj(x) |
| 1110 | return x |
| 1111 | |
| 1112 | |
| 1113 | class MuseGlimmerAttention(nn.Module): |
| 1114 | def __init__( |
| 1115 | self, |
| 1116 | config, |
| 1117 | cache_config: CacheConfig | None = None, |
| 1118 | quant_config: QuantizationConfig | None = None, |
| 1119 | prefix: str = "", |
| 1120 | ) -> None: |
| 1121 | super().__init__() |
| 1122 | self.config = config |
| 1123 | self.layer_idx = extract_layer_index(prefix) |
| 1124 | |
| 1125 | tp_size = get_tensor_model_parallel_world_size() |
| 1126 | self.hidden_size = config.hidden_size |
| 1127 | self.total_num_heads = config.num_attention_heads |
| 1128 | assert self.total_num_heads % tp_size == 0 |
| 1129 | self.num_heads = self.total_num_heads // tp_size |
| 1130 | self.total_num_kv_heads = config.num_key_value_heads |
| 1131 | if self.total_num_kv_heads >= tp_size: |
| 1132 | assert self.total_num_kv_heads % tp_size == 0 |
| 1133 | else: |
| 1134 | assert tp_size % self.total_num_kv_heads == 0 |
| 1135 | self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) |
| 1136 | self.head_dim = config.head_dim |
| 1137 | self.q_size = self.num_heads * self.head_dim |
| 1138 | self.kv_size = self.num_kv_heads * self.head_dim |
| 1139 | # MuseGlimmer overrides Gemma2's query_pre_attn_scalar scaling with the |
| 1140 | # standard 1/sqrt(head_dim). The query is *additionally* pre-scaled by |
| 1141 | # scale_query_by after QK-norm (see forward). |
| 1142 | self.scaling = self.head_dim**-0.5 |
| 1143 | |
| 1144 | # iRoPE: NoPE layers (no_rope_layers[i] == 0) run full attention; RoPE |
| 1145 | # layers run sliding-window attention. |
| 1146 | self.use_rope = config.no_rope_layers[self.layer_idx] == 1 |
| 1147 | |
| 1148 | self.use_qk_norm = _muse_glimmer_use_qk_norm(config) |
| 1149 | if self.use_qk_norm: |
| 1150 | # Weightless, computed in fp32, applied per head over head_dim. |
| 1151 | self.qk_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False) |
| 1152 | # Post-QK-norm query pre-scale, normalized across the native (raw |
| 1153 | # ~43.78) and modular (pre-folded ~3.87) config schemas. |
| 1154 | self.scale_query_by = _muse_glimmer_query_prescale(config) |
| 1155 | |
| 1156 | self.qkv_proj = QKVParallelLinear( |
| 1157 | self.hidden_size, |
| 1158 | self.head_dim, |
| 1159 | self.total_num_heads, |
| 1160 | self.total_num_kv_heads, |
| 1161 | bias=False, |
| 1162 | quant_config=quant_config, |
| 1163 | prefix=f"{prefix}.qkv_proj", |
| 1164 | ) |
| 1165 | self.o_proj = RowParallelLinear( |
| 1166 | self.total_num_heads * self.head_dim, |
| 1167 | self.hidden_size, |
| 1168 | bias=False, |
| 1169 | quant_config=quant_config, |
| 1170 | prefix=f"{prefix}.o_proj", |
| 1171 | ) |
| 1172 | |
| 1173 | self.use_output_gate = _muse_glimmer_use_attn_output_gate(config) |
| 1174 | if self.use_output_gate: |
| 1175 | self.output_gate_proj = ColumnParallelLinear( |
| 1176 | self.hidden_size, |
| 1177 | self.total_num_heads * self.head_dim, |
| 1178 | bias=False, |
| 1179 | gather_output=False, |
| 1180 | quant_config=quant_config, |
| 1181 | prefix=f"{prefix}.output_gate_proj", |
| 1182 | ) |
| 1183 | |
| 1184 | self.rotary_emb = ( |
| 1185 | get_rope( |
| 1186 | self.head_dim, |
| 1187 | max_position=config.max_position_embeddings, |
| 1188 | rope_parameters=config.rope_parameters, |
| 1189 | # HF converter permutes q/k to NEOX (half-split) |
| 1190 | # layout via _permute_for_rope |
| 1191 | is_neox_style=True, |
| 1192 | ) |
| 1193 | if self.use_rope |
| 1194 | else None |
| 1195 | ) |
| 1196 | |
| 1197 | # Full attention on NoPE layers, sliding window otherwise. |
| 1198 | sliding_window = None if not self.use_rope else config.sliding_window |
| 1199 | self.attn = Attention( |
| 1200 | self.num_heads, |
| 1201 | self.head_dim, |
| 1202 | self.scaling, |
| 1203 | num_kv_heads=self.num_kv_heads, |
| 1204 | cache_config=cache_config, |
| 1205 | quant_config=quant_config, |
| 1206 | logits_soft_cap=None, # MuseGlimmer sets attn_logit_softcapping = None |
| 1207 | per_layer_sliding_window=sliding_window, |
| 1208 | prefix=f"{prefix}.attn", |
| 1209 | ) |
| 1210 | |
| 1211 | def forward( |
| 1212 | self, |
| 1213 | positions: torch.Tensor, |
| 1214 | hidden_states: torch.Tensor, |
| 1215 | ) -> torch.Tensor: |
| 1216 | qkv, _ = self.qkv_proj(hidden_states) |
| 1217 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 1218 | |
| 1219 | if self.use_qk_norm: |
| 1220 | # QK-norm over head_dim, fp32, applied BEFORE RoPE; then pre-scale q. |
| 1221 | q = q.reshape(-1, self.head_dim) |
| 1222 | q = self.qk_norm(q).reshape(-1, self.q_size) * self.scale_query_by |
| 1223 | k = k.reshape(-1, self.head_dim) |
| 1224 | k = self.qk_norm(k).reshape(-1, self.kv_size) |
| 1225 | q = q.to(v.dtype) |
| 1226 | k = k.to(v.dtype) |
| 1227 | |
| 1228 | if self.rotary_emb is not None: |
| 1229 | q, k = self.rotary_emb(positions, q, k) |
| 1230 | |
| 1231 | attn_output = self.attn(q, k, v) |
| 1232 | |
| 1233 | if self.use_output_gate: |
| 1234 | # Gate reads the layer input hidden states (not the attn output). |
| 1235 | gate, _ = self.output_gate_proj(hidden_states) |
| 1236 | attn_output = torch.sigmoid(gate) * attn_output |
| 1237 | |
| 1238 | output, _ = self.o_proj(attn_output) |
| 1239 | return output |
| 1240 | |
| 1241 | |
| 1242 | class MuseGlimmerDecoderLayer(nn.Module): |
| 1243 | def __init__( |
| 1244 | self, |
| 1245 | config, |
| 1246 | cache_config: CacheConfig | None = None, |
| 1247 | quant_config: QuantizationConfig | None = None, |
| 1248 | prefix: str = "", |
| 1249 | ) -> None: |
| 1250 | super().__init__() |
| 1251 | self.self_attn = MuseGlimmerAttention( |
| 1252 | config=config, |
| 1253 | cache_config=cache_config, |
| 1254 | quant_config=quant_config, |
| 1255 | prefix=f"{prefix}.self_attn", |
| 1256 | ) |
| 1257 | self.mlp = MuseGlimmerMLP( |
| 1258 | hidden_size=config.hidden_size, |
| 1259 | intermediate_size=config.intermediate_size, |
| 1260 | hidden_activation=config.hidden_activation, |
| 1261 | quant_config=quant_config, |
| 1262 | prefix=f"{prefix}.mlp", |
| 1263 | ) |
| 1264 | # Sandwich norms with baked +1 offset. Pre-norms use rms_norm_eps; the |
| 1265 | # post-norms use the (typically smaller) post_norm_eps. |
| 1266 | self.input_layernorm = MuseGlimmerRMSNorm( |
| 1267 | config.hidden_size, eps=config.rms_norm_eps, weight_offset=1 |
| 1268 | ) |
| 1269 | self.post_attention_layernorm = MuseGlimmerRMSNorm( |
| 1270 | config.hidden_size, eps=config.post_norm_eps, weight_offset=1 |
| 1271 | ) |
| 1272 | self.pre_feedforward_layernorm = MuseGlimmerRMSNorm( |
| 1273 | config.hidden_size, eps=config.rms_norm_eps, weight_offset=1 |
| 1274 | ) |
| 1275 | self.post_feedforward_layernorm = MuseGlimmerRMSNorm( |
| 1276 | config.hidden_size, eps=config.post_norm_eps, weight_offset=1 |
| 1277 | ) |
| 1278 | |
| 1279 | def forward( |
| 1280 | self, |
| 1281 | positions: torch.Tensor, |
| 1282 | hidden_states: torch.Tensor, |
| 1283 | residual: torch.Tensor | None, |
| 1284 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1285 | # Gemma2-style sandwich, replicated explicitly (matches HF MuseGlimmer). |
| 1286 | residual = hidden_states |
| 1287 | hidden_states = self.input_layernorm(hidden_states) |
| 1288 | hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) |
| 1289 | hidden_states = self.post_attention_layernorm(hidden_states) |
| 1290 | hidden_states = residual + hidden_states |
| 1291 | |
| 1292 | residual = hidden_states |
| 1293 | hidden_states = self.pre_feedforward_layernorm(hidden_states) |
| 1294 | hidden_states = self.mlp(hidden_states) |
| 1295 | hidden_states = self.post_feedforward_layernorm(hidden_states) |
| 1296 | hidden_states = residual + hidden_states |
| 1297 | return hidden_states, residual |
| 1298 | |
| 1299 | |
| 1300 | @support_torch_compile |
| 1301 | class MuseGlimmerModel(nn.Module, EagleModelMixin): |
| 1302 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1303 | super().__init__() |
| 1304 | config = _text_config(vllm_config.model_config.hf_config) |
| 1305 | cache_config = vllm_config.cache_config |
| 1306 | quant_config = vllm_config.quant_config |
| 1307 | self.config = config |
| 1308 | self.quant_config = quant_config |
| 1309 | |
| 1310 | self.embed_tokens = VocabParallelEmbedding( |
| 1311 | config.vocab_size, |
| 1312 | config.hidden_size, |
| 1313 | ) |
| 1314 | # MuseGlimmer normalizes token embeddings with a weightless RMSNorm instead of |
| 1315 | # Gemma's sqrt(hidden_size) multiplier. |
| 1316 | self.embed_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False) |
| 1317 | |
| 1318 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 1319 | config.num_hidden_layers, |
| 1320 | lambda prefix: MuseGlimmerDecoderLayer( |
| 1321 | config, cache_config, quant_config, prefix=prefix |
| 1322 | ), |
| 1323 | prefix=f"{prefix}.layers", |
| 1324 | ) |
| 1325 | # Final norm: weight-as-scale, no offset. |
| 1326 | self.norm = MuseGlimmerRMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 1327 | self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( |
| 1328 | ["hidden_states", "residual"], config.hidden_size |
| 1329 | ) |
| 1330 | |
| 1331 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 1332 | return self.embed_norm(self.embed_tokens(input_ids)) |
| 1333 | |
| 1334 | def forward( |
| 1335 | self, |
| 1336 | input_ids: torch.Tensor | None, |
| 1337 | positions: torch.Tensor, |
| 1338 | intermediate_tensors: IntermediateTensors | None, |
| 1339 | inputs_embeds: torch.Tensor | None = None, |
| 1340 | ) -> torch.Tensor | IntermediateTensors: |
| 1341 | if get_pp_group().is_first_rank: |
| 1342 | if inputs_embeds is not None: |
| 1343 | hidden_states = inputs_embeds |
| 1344 | else: |
| 1345 | hidden_states = self.embed_input_ids(input_ids) |
| 1346 | residual = None |
| 1347 | else: |
| 1348 | assert intermediate_tensors is not None |
| 1349 | hidden_states = intermediate_tensors["hidden_states"] |
| 1350 | residual = intermediate_tensors["residual"] |
| 1351 | aux_hidden_states = self._maybe_add_hidden_state( |
| 1352 | [], self.start_layer, hidden_states, None |
| 1353 | ) |
| 1354 | for layer_idx, layer in enumerate( |
| 1355 | islice(self.layers, self.start_layer, self.end_layer), |
| 1356 | start=self.start_layer, |
| 1357 | ): |
| 1358 | hidden_states, residual = layer(positions, hidden_states, residual) |
| 1359 | self._maybe_add_hidden_state( |
| 1360 | aux_hidden_states, layer_idx + 1, hidden_states, None |
| 1361 | ) |
| 1362 | if not get_pp_group().is_last_rank: |
| 1363 | return IntermediateTensors( |
| 1364 | {"hidden_states": hidden_states, "residual": residual} |
| 1365 | ) |
| 1366 | hidden_states = self.norm(hidden_states) |
| 1367 | if len(aux_hidden_states) > 0: |
| 1368 | return hidden_states, aux_hidden_states |
| 1369 | return hidden_states |
| 1370 | |
| 1371 | |
| 1372 | @MULTIMODAL_REGISTRY.register_processor( |
| 1373 | MuseGlimmerMultiModalProcessor, |
| 1374 | info=MuseGlimmerProcessingInfo, |
| 1375 | dummy_inputs=MuseGlimmerDummyInputsBuilder, |
| 1376 | ) |
| 1377 | class MuseGlimmerForCausalLM( |
| 1378 | nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3 |
| 1379 | ): |
| 1380 | # Weight-name normalization. Two checkpoint conventions are supported: |
| 1381 | # |
| 1382 | # * HF MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``): the |
| 1383 | # multimodal ``MuseGlimmerConfig`` prefixes the language model with |
| 1384 | # ``model.language_model.``; the per-layer sandwich norms are already |
| 1385 | # named ``input_layernorm`` / ``post_attention_layernorm`` / |
| 1386 | # ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm``. |
| 1387 | # |
| 1388 | # * Legacy HF export (an earlier checkpoint convention): uses ``model.`` |
| 1389 | # and a different sandwich-norm naming where ``post_attn_norm`` is the |
| 1390 | # true post-attention norm and ``post_attention_layernorm`` is actually |
| 1391 | # the pre-feedforward norm. We remap those to MuseGlimmer's names. |
| 1392 | # |
| 1393 | # CONVENTION DISAMBIGUATION (critical): the two checkpoint families use |
| 1394 | # DIFFERENT sandwich-norm names, and they must not be conflated: |
| 1395 | # |
| 1396 | # * Canonical MuseGlimmer export (current |
| 1397 | # ``convert_muse_glimmer_weights_to_hf.py`` — what partners ship): |
| 1398 | # keys are ``model.language_model.layers.N.*`` and |
| 1399 | # the norms are ALREADY named ``input_layernorm`` / |
| 1400 | # ``post_attention_layernorm`` / ``pre_feedforward_layernorm`` / |
| 1401 | # ``post_feedforward_layernorm``. No norm rename needed — pass through. |
| 1402 | # |
| 1403 | # * Legacy HF export (an earlier checkpoint convention): keys are |
| 1404 | # ``model.layers.N.*`` and the sandwich norms are |
| 1405 | # named ``input_layernorm`` / ``post_attention_layernorm`` (this one is |
| 1406 | # actually the PRE-feedforward norm) / ``post_attn_norm`` (the true |
| 1407 | # post-attention norm) / ``post_ffn_norm``. These must be remapped. |
| 1408 | # |
| 1409 | # The unambiguous discriminator is the PREFIX: legacy keys start with |
| 1410 | # ``model.layers.`` while canonical keys start with |
| 1411 | # ``model.language_model.layers.``. ``orig_to_new_regex`` runs BEFORE the |
| 1412 | # prefix strip (see WeightsMapper._map_name_with_shard), so we anchor the |
| 1413 | # legacy renames on ``^model\.layers\.`` — they fire ONLY on legacy keys and |
| 1414 | # leave canonical/partner checkpoints untouched. Rule order within the regex |
| 1415 | # dict matters: the ``post_attention_layernorm`` -> ``pre_feedforward_...`` |
| 1416 | # rule must precede the ``post_attn_norm`` -> ``post_attention_layernorm`` |
| 1417 | # rule so the latter's output is not re-captured by the former (regex rules |
| 1418 | # apply as a single forward pass). |
| 1419 | hf_to_vllm_mapper = WeightsMapper( |
| 1420 | orig_to_new_substr={ |
| 1421 | "model.vision_tower.patch_embedder.position_embedding_table.weight": ( |
| 1422 | "model.vision_tower.positional_embedding_vlm" |
| 1423 | ), |
| 1424 | "model.vision_tower.layers.": "model.vision_tower.transformer.", |
| 1425 | ".norm1.": ".ln_1.", |
| 1426 | ".norm2.": ".ln_2.", |
| 1427 | ".attn.proj.": ".attn.o_proj.", |
| 1428 | ".mlp.fc1.": ".mlp.c_fc.", |
| 1429 | ".mlp.fc2.": ".mlp.c_proj.", |
| 1430 | ".self_attn.gate_proj": ".self_attn.output_gate_proj", |
| 1431 | }, |
| 1432 | orig_to_new_prefix={ |
| 1433 | "model.rotary_emb.": None, |
| 1434 | "model.language_model.": "model.", |
| 1435 | "language_model.": "model.", |
| 1436 | "model.vision_tower.patch_embedder.patch_embedding.": ( |
| 1437 | "model.vision_tower.conv1_linear." |
| 1438 | ), |
| 1439 | "model.vision_tower.": "vision_encoder.", |
| 1440 | "vision_tower.": "vision_encoder.", |
| 1441 | "model.vision_encoder.": "vision_encoder.", |
| 1442 | "model.vision_adapter.fc1.": "model.vision_adapter.c_fc.", |
| 1443 | "model.vision_adapter.fc2.": "model.vision_adapter.c_proj.", |
| 1444 | "model.vision_adapter.": "vision_adapter.", |
| 1445 | "model.vision_projection.": "vision_projection.", |
| 1446 | "model.perception_emb_norm.": "perception_emb_norm.", |
| 1447 | }, |
| 1448 | orig_to_new_stacked={ |
| 1449 | ".q_proj": (".qkv_proj", "q"), |
| 1450 | ".k_proj": (".qkv_proj", "k"), |
| 1451 | ".v_proj": (".qkv_proj", "v"), |
| 1452 | ".gate_proj": (".gate_up_proj", 0), |
| 1453 | ".up_proj": (".gate_up_proj", 1), |
| 1454 | }, |
| 1455 | ) |
| 1456 | |
| 1457 | packed_modules_mapping = { |
| 1458 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 1459 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 1460 | } |
| 1461 | |
| 1462 | @classmethod |
| 1463 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1464 | if modality.startswith("image"): |
| 1465 | return IMAGE_TOKEN |
| 1466 | if modality.startswith("video"): |
| 1467 | return VIDEO_TOKEN |
| 1468 | raise ValueError(f"Unsupported MuseGlimmer modality: {modality}") |
| 1469 | |
| 1470 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1471 | super().__init__() |
| 1472 | config = vllm_config.model_config.hf_config |
| 1473 | text_config = _text_config(config) |
| 1474 | vision_config = _vision_config(config) |
| 1475 | quant_config = vllm_config.quant_config |
| 1476 | self.config = config |
| 1477 | self.text_config = text_config |
| 1478 | self.quant_config = quant_config |
| 1479 | self.has_vision = _muse_glimmer_has_vision(config) |
| 1480 | |
| 1481 | with self._mark_language_model(vllm_config): |
| 1482 | self.model = MuseGlimmerModel( |
| 1483 | vllm_config=vllm_config, |
| 1484 | prefix=maybe_prefix(prefix, "model"), |
| 1485 | ) |
| 1486 | |
| 1487 | self.vision_encoder: MuseGlimmerVisionEncoder | None |
| 1488 | self.vision_adapter: MuseGlimmerVisionAdapter | None |
| 1489 | if self.has_vision: |
| 1490 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 1491 | self.vision_encoder = MuseGlimmerVisionEncoder( |
| 1492 | vision_config, |
| 1493 | prefix=maybe_prefix(prefix, "vision_encoder"), |
| 1494 | ) |
| 1495 | self.vision_adapter = MuseGlimmerVisionAdapter(vision_config) |
| 1496 | self.vision_projection = nn.Linear( |
| 1497 | vision_config.adapter_dim, |
| 1498 | text_config.hidden_size, |
| 1499 | bias=False, |
| 1500 | ) |
| 1501 | self.perception_emb_norm = ( |
| 1502 | MuseGlimmerRMSNorm(eps=text_config.rms_norm_eps, with_scale=False) |
| 1503 | if text_config.normalize_tok_embeddings |
| 1504 | else nn.Identity() |
| 1505 | ) |
| 1506 | else: |
| 1507 | self.vision_encoder = None |
| 1508 | self.vision_adapter = None |
| 1509 | self.vision_projection = None |
| 1510 | self.perception_emb_norm = None |
| 1511 | |
| 1512 | self.lm_head = ParallelLMHead( |
| 1513 | text_config.vocab_size, |
| 1514 | text_config.hidden_size, |
| 1515 | quant_config=quant_config, |
| 1516 | prefix=maybe_prefix(prefix, "lm_head"), |
| 1517 | ) |
| 1518 | self.output_multiplier = text_config.output_multiplier |
| 1519 | self.final_logit_softcapping = text_config.final_logit_softcapping |
| 1520 | self.logits_processor = LogitsProcessor(text_config.vocab_size) |
| 1521 | self.make_empty_intermediate_tensors = ( |
| 1522 | self.model.make_empty_intermediate_tensors |
| 1523 | ) |
| 1524 | image_token_id = getattr( |
| 1525 | config, "image_token_id", getattr(config, "patch_token_id", 200092) |
| 1526 | ) |
| 1527 | video_token_id = getattr(config, "video_token_id", 200091) |
| 1528 | self.configure_mm_token_handling( |
| 1529 | text_config.vocab_size, [image_token_id, video_token_id] |
| 1530 | ) |
| 1531 | |
| 1532 | def _parse_and_validate_image_input( |
| 1533 | self, |
| 1534 | **kwargs: object, |
| 1535 | ) -> MuseGlimmerImagePixelInputs | None: |
| 1536 | pixel_values = kwargs.pop("image_pixel_values", None) |
| 1537 | feature_sizes = kwargs.pop("image_feature_sizes", None) |
| 1538 | if pixel_values is None and feature_sizes is None: |
| 1539 | return None |
| 1540 | if pixel_values is None or feature_sizes is None: |
| 1541 | raise ValueError( |
| 1542 | "MuseGlimmer image_pixel_values and image_feature_sizes " |
| 1543 | "must be provided together" |
| 1544 | ) |
| 1545 | if not isinstance(feature_sizes, torch.Tensor): |
| 1546 | raise ValueError("MuseGlimmer image_feature_sizes must be a tensor") |
| 1547 | if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 3: |
| 1548 | pixel_values = pixel_values.unsqueeze(0) |
| 1549 | return MuseGlimmerImagePixelInputs( |
| 1550 | type="image_pixels", |
| 1551 | pixel_values=pixel_values, |
| 1552 | feature_sizes=feature_sizes.reshape(-1), |
| 1553 | ) |
| 1554 | |
| 1555 | def _parse_and_validate_video_input( |
| 1556 | self, |
| 1557 | **kwargs: object, |
| 1558 | ) -> MuseGlimmerVideoPixelInputs | None: |
| 1559 | pixel_values = kwargs.pop("video_pixel_values", None) |
| 1560 | feature_sizes = kwargs.pop("video_feature_sizes", None) |
| 1561 | if pixel_values is None and feature_sizes is None: |
| 1562 | return None |
| 1563 | if pixel_values is None or feature_sizes is None: |
| 1564 | raise ValueError( |
| 1565 | "MuseGlimmer video_pixel_values and video_feature_sizes " |
| 1566 | "must be provided together" |
| 1567 | ) |
| 1568 | if not isinstance(feature_sizes, torch.Tensor): |
| 1569 | raise ValueError("MuseGlimmer video_feature_sizes must be a tensor") |
| 1570 | if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 4: |
| 1571 | pixel_values = pixel_values.unsqueeze(0) |
| 1572 | patch_temporal = int(_vision_config(self.config).patch_temporal) |
| 1573 | return MuseGlimmerVideoPixelInputs( |
| 1574 | type="video_pixels", |
| 1575 | pixel_values=pixel_values, |
| 1576 | feature_sizes=feature_sizes.reshape(-1), |
| 1577 | resolve_bindings={"c": patch_temporal * 3}, |
| 1578 | ) |
| 1579 | |
| 1580 | def _encode_pixel_groups( |
| 1581 | self, |
| 1582 | pixel_groups: Sequence[torch.Tensor], |
| 1583 | feature_sizes: Sequence[int], |
| 1584 | ) -> tuple[torch.Tensor, ...]: |
| 1585 | if ( |
| 1586 | self.vision_encoder is None |
| 1587 | or self.vision_adapter is None |
| 1588 | or self.vision_projection is None |
| 1589 | or self.perception_emb_norm is None |
| 1590 | ): |
| 1591 | raise ValueError("This MuseGlimmer checkpoint has no vision tower") |
| 1592 | features = self.vision_encoder(pixel_groups) |
| 1593 | features = self.vision_adapter(features) |
| 1594 | features = self.vision_projection(features) |
| 1595 | features = self.perception_emb_norm(features) |
| 1596 | if features.shape[0] != sum(feature_sizes): |
| 1597 | raise ValueError( |
| 1598 | f"MuseGlimmer produced {features.shape[0]} vision features for " |
| 1599 | f"{sum(feature_sizes)} placeholder tokens" |
| 1600 | ) |
| 1601 | return features.split(list(feature_sizes)) |
| 1602 | |
| 1603 | def _process_image_input( |
| 1604 | self, |
| 1605 | image_input: MuseGlimmerImagePixelInputs, |
| 1606 | ) -> tuple[torch.Tensor, ...]: |
| 1607 | images = list(image_input["pixel_values"]) |
| 1608 | sizes = [int(size) for size in image_input["feature_sizes"].tolist()] |
| 1609 | if len(images) != len(sizes): |
| 1610 | raise ValueError("MuseGlimmer image batch metadata does not match pixels") |
| 1611 | return self._encode_pixel_groups(images, sizes) |
| 1612 | |
| 1613 | def _process_video_input( |
| 1614 | self, |
| 1615 | video_input: MuseGlimmerVideoPixelInputs, |
| 1616 | ) -> tuple[torch.Tensor, ...]: |
| 1617 | videos = list(video_input["pixel_values"]) |
| 1618 | sizes = [int(size) for size in video_input["feature_sizes"].tolist()] |
| 1619 | if len(videos) != len(sizes): |
| 1620 | raise ValueError("MuseGlimmer video batch metadata does not match pixels") |
| 1621 | groups = [group for video in videos for group in video] |
| 1622 | return self._encode_pixel_groups(groups, sizes) |
| 1623 | |
| 1624 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1625 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 1626 | video_input = self._parse_and_validate_video_input(**kwargs) |
| 1627 | |
| 1628 | embeddings: list[torch.Tensor] = [] |
| 1629 | for key in kwargs: |
| 1630 | if key == "image_pixel_values" and image_input is not None: |
| 1631 | embeddings.extend(self._process_image_input(image_input)) |
| 1632 | elif key == "video_pixel_values" and video_input is not None: |
| 1633 | embeddings.extend(self._process_video_input(video_input)) |
| 1634 | return embeddings |
| 1635 | |
| 1636 | def forward( |
| 1637 | self, |
| 1638 | input_ids: torch.Tensor | None, |
| 1639 | positions: torch.Tensor, |
| 1640 | intermediate_tensors: IntermediateTensors | None = None, |
| 1641 | inputs_embeds: torch.Tensor | None = None, |
| 1642 | **kwargs: object, |
| 1643 | ) -> torch.Tensor | IntermediateTensors: |
| 1644 | return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) |
| 1645 | |
| 1646 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 1647 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 1648 | if logits is None: |
| 1649 | return None |
| 1650 | logits = logits * self.output_multiplier |
| 1651 | if self.final_logit_softcapping is not None: |
| 1652 | cap = self.final_logit_softcapping |
| 1653 | logits = cap * torch.tanh(logits / cap) |
| 1654 | return logits |
| 1655 | |
| 1656 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1657 | loader = AutoWeightsLoader(self) |
| 1658 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1659 | |
| 1660 | def get_mm_mapping(self) -> MultiModelKeys: |
| 1661 | """Get the module prefix in multimodal models""" |
| 1662 | return MultiModelKeys.from_string_field( |
| 1663 | language_model="model", |
| 1664 | connector=["vision_adapter.", "vision_projection."], |
| 1665 | tower_model="vision_encoder.", |
| 1666 | ) |