diff --git a/docs/models.md b/docs/models.md index 2a1600aaf..d2c811ba0 100644 --- a/docs/models.md +++ b/docs/models.md @@ -30,6 +30,7 @@ These families accept image input by default; pass `--text-model-only` to skip t | Gemma-4 26B-A4B, 31B (`gemma4`: ViT tower, streamed under `--mm-encoder-weights host`) | one of the soft-token budgets 70 / 140 / 280 / 560 / 1120, every image scaled to its budget as far as the aspect ratio allows | the maximum picks the largest budget within it, below 70 is refused at start-up; the minimum has no effect | `{"max_soft_tokens": 1120}` | | Gemma-4 12B (`gemma4_unified`: linear patch embedder, resident under either placement) | same budgets, one 48x48 super-patch per soft token | same as the tower releases | same | | GLM-5.3-Flash (`glm5_next`: ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution on a canvas zero-padded to a 28-multiple | token counts, passed through as the processor's `min_image_tokens` / `max_image_tokens`; checkpoint defaults 16 to 8000 tokens | `{"max_image_tokens": 2048}` | +| Muse-Glimmer-30B (windowed ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, aspect ratio kept under a token cap; checkpoint default 4096 tokens | the maximum is the cap (`max_image_tokens`); the minimum has no effect | `{"max_image_tokens": 1024}` | ## MoE strategies diff --git a/python/freetoken/mm/processors/muse_glimmer.py b/python/freetoken/mm/processors/muse_glimmer.py new file mode 100644 index 000000000..1e14387bc --- /dev/null +++ b/python/freetoken/mm/processors/muse_glimmer.py @@ -0,0 +1,79 @@ +"""Muse-Glimmer image processor: patch-grid ViT with a 2x2 pixel shuffle, a token-maximum budget and 1-D rope; the replacement adds the start/end wrapper the checkpoint processor renders.""" + +from __future__ import annotations + +import struct +from typing import Any + +import torch + +from freetoken.message import MMItem +from freetoken.mm import mm_pad_value +from freetoken.mm.config import MultimodalConfig +from freetoken.mm.processor import MMProcessor, PromptReplacement, content_hash + +# the tokenizer's <|image_start|> / <|image_end|>; the config only names the <|patch|> id +_IMAGE_START_ID = 200080 +_IMAGE_END_ID = 200081 + + +class MuseGlimmerMMProcessor(MMProcessor): + def __init__(self, hf_config: Any, model_path: str, mm: MultimodalConfig) -> None: + super().__init__(model_path, mm) + vc = hf_config.vision_config + self.image_token_id = hf_config.image_token_id + # the template renders one <|patch|> per image + self.placeholder = [self.image_token_id] + self.merge = vc.merge_size + self.patch_dim = vc.patch_temporal * 3 * vc.patch_size**2 + + def get_mm_processor_kwargs(self, mm: MultimodalConfig) -> dict[str, Any]: + kwargs: dict[str, Any] = {"return_tensors": "pt"} + # the image processor's only budget is a token maximum; a minimum has no knob + if mm.image_max_tokens is not None: + kwargs["max_image_tokens"] = mm.image_max_tokens + return {**kwargs, **mm.processor_kwargs} + + def process(self, images: list[Any]) -> list[MMItem]: + processor = self._image_processor() + kwargs = self.get_mm_processor_kwargs(self.mm) + items: list[MMItem] = [] + for pil in images: + out = processor(images=pil, **kwargs) + pixel_values = out["pixel_values"].to(torch.bfloat16) + grid = out["image_grid_thw"][0].tolist() + t, h, w = int(grid[0]), int(grid[1]), int(grid[2]) + if t != 1: + raise ValueError("video/temporal input is not supported") + # hash the bf16 wire buffer: half the bytes, and the exact tensor the cache and radix key agree on + item_hash = content_hash(pixel_values, struct.pack("<3i", t, h, w)) + items.append( + MMItem( + modality="image", + hash=item_hash, + pad_value=mm_pad_value(item_hash), + offsets=[], + feature=pixel_values, + model_specific_data={"grid_thw": [t, h, w]}, + ) + ) + return items + + def prompt_replacement(self, item: MMItem) -> PromptReplacement: + t, h, w = item.grid_thw + pads = [self.image_token_id] * ((t * h * w) // (self.merge * self.merge)) + return PromptReplacement.select_token_id([_IMAGE_START_ID, *pads, _IMAGE_END_ID], self.image_token_id) + + def dummy_items(self, dtype: torch.dtype, device: torch.device) -> list[MMItem]: + merge = self.merge + return [MMItem( + modality="image", + hash=0, + pad_value=0, + offsets=[[0, 1]], + feature=torch.zeros(merge * merge, self.patch_dim, dtype=dtype, device=device), + model_specific_data={"grid_thw": [1, merge, merge]}, + )] + + +__all__ = ["MuseGlimmerMMProcessor"] diff --git a/python/freetoken/models/muse_glimmer/__init__.py b/python/freetoken/models/muse_glimmer/__init__.py index 307d28660..6f08348e3 100644 --- a/python/freetoken/models/muse_glimmer/__init__.py +++ b/python/freetoken/models/muse_glimmer/__init__.py @@ -1,11 +1,16 @@ from .attention import MuseGlimmerAttention -from .config import parse_config -from .model import MuseGlimmerForCausalLM +from .config import VisionConfig, parse_config, parse_vision_config +from .model import MuseGlimmerForCausalLM, MuseGlimmerForConditionalGeneration +from .vision import MuseGlimmerVisionModel from .weight import iter_weights __all__ = [ "MuseGlimmerAttention", "MuseGlimmerForCausalLM", + "MuseGlimmerForConditionalGeneration", + "MuseGlimmerVisionModel", + "VisionConfig", "parse_config", + "parse_vision_config", "iter_weights", ] diff --git a/python/freetoken/models/muse_glimmer/config.py b/python/freetoken/models/muse_glimmer/config.py index 4cc756cb6..6916b0025 100644 --- a/python/freetoken/models/muse_glimmer/config.py +++ b/python/freetoken/models/muse_glimmer/config.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Any from freetoken.models.config import ( @@ -12,13 +13,76 @@ _SWA_TYPE = "sliding_attention" _FULL_TYPE = "full_attention" +_VISION_LAYER_TYPES = {"window_attention", "full_attention"} + + +@dataclass(frozen=True) +class VisionConfig: + hidden_size: int + intermediate_size: int + num_layers: int + num_heads: int + layer_types: tuple[str, ...] + patch_size: int + temporal_patch_size: int + merge_size: int + pos_emb_side: int + layer_norm_eps: float + rope_theta: float + projector_hidden_size: int + text_hidden_size: int + text_rms_norm_eps: float + + @property + def patch_dim(self) -> int: + return self.temporal_patch_size * 3 * self.patch_size**2 + + @property + def out_hidden_size(self) -> int: + return self.hidden_size * self.merge_size**2 + + +def parse_vision_config(hf_config: Any) -> VisionConfig | None: + """None when the config carries no vision section, which is how a text-only engine asks for no tower.""" + vc = getattr(hf_config, "vision_config", None) + if vc is None: + return None + if vc.hidden_act != "gelu" or hf_config.projector_hidden_act != "gelu": + raise NotImplementedError( + f"muse_glimmer vision activations {vc.hidden_act!r} / {hf_config.projector_hidden_act!r}; only gelu is implemented" + ) + if vc.pos_emb_height != vc.pos_emb_width: + raise NotImplementedError("muse_glimmer vision tower: the position table must be square") + layer_types = tuple(vc.layer_types) + assert set(layer_types) <= _VISION_LAYER_TYPES, f"unknown vision layer types in {set(layer_types)}" + rope_params = getattr(vc, "rope_parameters", None) or {} + text = _text_config(hf_config) + config = VisionConfig( + hidden_size=vc.hidden_size, + intermediate_size=vc.intermediate_size, + num_layers=vc.num_hidden_layers, + num_heads=vc.num_attention_heads, + layer_types=layer_types, + patch_size=vc.patch_size, + temporal_patch_size=vc.patch_temporal, + merge_size=vc.merge_size, + pos_emb_side=vc.pos_emb_height, + layer_norm_eps=vc.layer_norm_eps, + rope_theta=float(rope_params.get("rope_theta", 10000.0)), + projector_hidden_size=hf_config.projector_hidden_size, + text_hidden_size=text.hidden_size, + text_rms_norm_eps=text.rms_norm_eps, + ) + assert hf_config.out_hidden_size == config.out_hidden_size, ( + f"adapter input {hf_config.out_hidden_size} != pixel-shuffled tower width {config.out_hidden_size}" + ) + return config def _text_config(hf_config: Any) -> Any: """Muse Glimmer ships as a multimodal wrapper (MuseGlimmerForConditionalGeneration): the text tower lives in ``text_config`` and the weights carry a ``language_model.`` - prefix. Served text-only -- the ~1.8B ViT perception encoder is never built, so - ``ModelConfig.vision_config`` stays None and the loader drops the vision tensors.""" + prefix.""" text = getattr(hf_config, "text_config", None) return text if text is not None else hf_config @@ -108,7 +172,7 @@ def parse_config(hf_config: Any) -> ModelConfig: attn_sm_scale=attn_sm_scale, final_logit_softcapping=getattr(text, "final_logit_softcapping", None), output_multiplier=getattr(text, "output_multiplier", None), - vision_config=None, # served text-only + vision_config=parse_vision_config(hf_config), image_token_id=getattr(hf_config, "image_token_id", None), attn_quant=quant, dense_quant=quant, @@ -133,4 +197,4 @@ def parse_config(hf_config: Any) -> ModelConfig: ) -__all__ = ["parse_config"] +__all__ = ["VisionConfig", "parse_config", "parse_vision_config"] diff --git a/python/freetoken/models/muse_glimmer/model.py b/python/freetoken/models/muse_glimmer/model.py index 1e1072ece..50bd45067 100644 --- a/python/freetoken/models/muse_glimmer/model.py +++ b/python/freetoken/models/muse_glimmer/model.py @@ -17,12 +17,14 @@ VocabParallelEmbedding, silu_and_mul, ) -from freetoken.models.blocks import BaseLLMModel +from freetoken.models.blocks import BaseLLMModel, embed_input_ids from freetoken.utils import nvtx_annotate from .attention import MuseGlimmerAttention +from .vision import MuseGlimmerVisionModel if TYPE_CHECKING: + from freetoken.message import MMItem from freetoken.models.config import ModelConfig @@ -83,6 +85,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return residual + h +class _NormedEmbedding: + """embed_tokens followed by the weightless norm as one lookup for embed_input_ids, so image rows (already normed by the tower) scatter in after it like in the reference.""" + + def __init__(self, embed_tokens: VocabParallelEmbedding, embed_norm: GemmaRMSNorm): + self.embed_tokens = embed_tokens + self.embed_norm = embed_norm + + @property + def num_embeddings(self) -> int: + return self.embed_tokens.num_embeddings + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_norm.forward(self.embed_tokens.forward(input_ids)) + + class MuseGlimmerModel(BaseOP): def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( @@ -95,6 +112,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_norm = GemmaRMSNorm( config.hidden_size, eps=config.rms_norm_eps, with_scale=False ) + self._normed_embedding = _NormedEmbedding(self.embed_tokens, self.embed_norm) self.layers = OPList( [ MuseGlimmerDecoderLayer(config, layer_id, prefix=f"{prefix}.layers.{layer_id}") @@ -106,7 +124,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward(self, input_ids: torch.Tensor) -> torch.Tensor: - x = self.embed_norm.forward(self.embed_tokens.forward(input_ids)) + x = embed_input_ids(self._normed_embedding, input_ids, get_global_ctx().batch) for layer in self.layers.op_list: x = layer.forward(x) return self.norm.forward(x) @@ -139,4 +157,17 @@ def forward(self) -> torch.Tensor: return logits -__all__ = ["MuseGlimmerForCausalLM"] +class MuseGlimmerForConditionalGeneration(MuseGlimmerForCausalLM): + def __init__(self, config: ModelConfig): + super().__init__(config) + if config.is_multimodal: + self.vision_tower = MuseGlimmerVisionModel(config.vision_config) + + def place_encoder_weights(self, mode: str) -> None: + self.vision_tower.place_weights(mode) + + def encode(self, item: MMItem) -> torch.Tensor: + return self.vision_tower.forward(item.feature, [item.grid_thw]) + + +__all__ = ["MuseGlimmerForCausalLM", "MuseGlimmerForConditionalGeneration"] diff --git a/python/freetoken/models/muse_glimmer/vision.py b/python/freetoken/models/muse_glimmer/vision.py new file mode 100644 index 000000000..782d26f75 --- /dev/null +++ b/python/freetoken/models/muse_glimmer/vision.py @@ -0,0 +1,185 @@ +"""Muse-Glimmer vision tower: linear patch embedding plus a resampled learned position table, LayerNorm blocks with 2-D rope that attend inside 32x32-patch windows (full attention every fourth block), a 2x2 pixel shuffle, the GELU adapter, the projection into the text width and the weightless perception norm.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List + +import torch +import torch.nn.functional as F +from freetoken.layers import BaseOP, GemmaRMSNorm, LayerNorm, LinearReplicated, OPList +from freetoken.models.weight_stream import BlockWeightStreamer + +if TYPE_CHECKING: + from .config import VisionConfig + + +class _Table(BaseOP): + """A bare weight under the key of the checkpoint's nn.Embedding.""" + + def __init__(self, rows: int, cols: int): + self.weight = torch.empty(rows, cols) + + forward = None + + +class MuseGlimmerVisionPatchEmbedder(BaseOP): + def __init__(self, vc: VisionConfig): + self.patch_embedding = LinearReplicated(vc.patch_dim, vc.hidden_size, has_bias=False) + self.position_embedding_table = _Table(vc.pos_emb_side**2, vc.hidden_size) + self._side = vc.pos_emb_side + + def forward(self, pixel_values: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + from transformers.vision_utils import get_vision_interpolation_indices_and_weights + + x = self.patch_embedding.forward(pixel_values.to(self.patch_embedding.weight.dtype)) + # the reference resamples the table with grid_sample(align_corners=False, padding_mode="zeros") + taps, weights = ( + t.to(x.device) + for t in get_vision_interpolation_indices_and_weights( + grid, self._side, mode="bilinear", align_corners=False, spatial_merge_size=1, padding="zeros" + ) + ) + # one tap at a time: fp32 like the reference's weighted sum, without the [S, 4, hidden] gather + table = self.position_embedding_table.weight + pos = table[taps[:, 0]] * weights[:, 0, None] + for t in range(1, taps.shape[1]): + pos = pos + table[taps[:, t]] * weights[:, t, None] + return x + pos.to(x.dtype) + + +class MuseGlimmerVisionAttention(BaseOP): + def __init__(self, vc: VisionConfig): + self.num_heads = vc.num_heads + self.head_dim = vc.hidden_size // vc.num_heads + self.qkv = LinearReplicated(vc.hidden_size, 3 * vc.hidden_size, has_bias=True) + self.proj = LinearReplicated(vc.hidden_size, vc.hidden_size, has_bias=True) + + def forward(self, x: torch.Tensor, cache: torch.Tensor, positions: torch.Tensor, lengths: List[int]) -> torch.Tensor: + from freetoken.kernel.triton.rope import apply_rope_with_cos_sin_cache_inplace + + S, H, D = x.shape[0], self.num_heads, self.head_dim + q, k, v = self.qkv.forward(x).view(S, 3, H * D).unbind(1) + apply_rope_with_cos_sin_cache_inplace(positions, q, k, D, cache, is_neox=True) + q, k, v = (t.view(S, H, D).transpose(0, 1).unsqueeze(0) for t in (q, k, v)) + # bidirectional attention within each segment: a window or a whole image + outs = [ + F.scaled_dot_product_attention(qs, ks, vs) + for qs, ks, vs in zip(torch.split(q, lengths, dim=2), torch.split(k, lengths, dim=2), torch.split(v, lengths, dim=2)) + ] + o = outs[0] if len(outs) == 1 else torch.cat(outs, dim=2) + return self.proj.forward(o[0].transpose(0, 1).reshape(S, H * D)) + + +class MuseGlimmerVisionMLP(BaseOP): + def __init__(self, vc: VisionConfig): + self.fc1 = LinearReplicated(vc.hidden_size, vc.intermediate_size, has_bias=True) + self.fc2 = LinearReplicated(vc.intermediate_size, vc.hidden_size, has_bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2.forward(F.gelu(self.fc1.forward(x))) + + +class MuseGlimmerVisionBlock(BaseOP): + def __init__(self, vc: VisionConfig): + # the reference builds the block norms with the torch default eps, not the config's + self.norm1 = LayerNorm(vc.hidden_size, eps=1e-5) + self.norm2 = LayerNorm(vc.hidden_size, eps=1e-5) + self.attn = MuseGlimmerVisionAttention(vc) + self.mlp = MuseGlimmerVisionMLP(vc) + + def forward(self, x: torch.Tensor, cache: torch.Tensor, positions: torch.Tensor, lengths: List[int]) -> torch.Tensor: + x = x + self.attn.forward(self.norm1.forward(x), cache, positions, lengths) + return x + self.mlp.forward(self.norm2.forward(x)) + + +class MuseGlimmerVisionAdapter(BaseOP): + def __init__(self, vc: VisionConfig): + self.fc1 = LinearReplicated(vc.out_hidden_size, vc.projector_hidden_size, has_bias=False) + self.fc2 = LinearReplicated(vc.projector_hidden_size, vc.projector_hidden_size, has_bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(self.fc2.forward(F.gelu(self.fc1.forward(x)))) + + +def _shuffle_index(grid: torch.Tensor, merge: int) -> torch.Tensor: + """Row order that puts every merge x merge patch block of each image in merge**2 consecutive rows.""" + parts, offset = [], 0 + for t, h, w in grid.tolist(): + assert t == 1, "images only" + block = torch.arange(h * w, device=grid.device).view(h // merge, merge, w // merge, merge) + parts.append(block.permute(0, 2, 1, 3).reshape(-1) + offset) + offset += h * w + return torch.cat(parts) + + +class MuseGlimmerVisionModel(BaseOP): + """Pixels -> [sum(h*w) / merge^2, text_hidden] soft tokens (the reference's get_image_features), one grid_thw row per image; adapter, projection and perception norm ride along so the whole image path is one module.""" + + def __init__(self, vc: VisionConfig): + self.patch_embedder = MuseGlimmerVisionPatchEmbedder(vc) + self.ln_pre = LayerNorm(vc.hidden_size, eps=vc.layer_norm_eps) + self.layers = OPList([MuseGlimmerVisionBlock(vc) for _ in range(vc.num_layers)]) + self.ln_post = LayerNorm(vc.hidden_size, eps=vc.layer_norm_eps) + self.adapter = MuseGlimmerVisionAdapter(vc) + self.projection = LinearReplicated(vc.projector_hidden_size, vc.text_hidden_size, has_bias=False) + self.perception_emb_norm = GemmaRMSNorm(vc.text_hidden_size, eps=vc.text_rms_norm_eps, with_scale=False) + self._vc = vc + self._inv_freq: torch.Tensor | None = None + self._streamer: BlockWeightStreamer | None = None + + def place_weights(self, mode: str) -> None: + """gpu: every tensor resident; host: block tensors in pinned banks streamed two blocks at a time (embedder, adapter and projection stay resident).""" + if mode == "host" and self._streamer is None: + self._streamer = BlockWeightStreamer(self.layers.op_list, self.patch_embedder.patch_embedding.weight.device) + elif mode == "gpu" and self._streamer is not None: + self._streamer.unstream() + self._streamer = None + elif mode not in ("gpu", "host"): + raise ValueError(f"unknown vision weight placement {mode!r}") + + def _layers(self): + if self._streamer is None: + return enumerate(self.layers.op_list) + return self._streamer.blocks(self.layers.op_list) + + def _rope_table(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + vc = self._vc + device = position_ids.device + if self._inv_freq is None or self._inv_freq.device != device: + # the reference splits the head in two and gives each axis the frequency ladder of a head_dim/2 rope + half = (vc.hidden_size // vc.num_heads) // 2 + self._inv_freq = 1.0 / (vc.rope_theta ** (torch.arange(0, half, 2, dtype=torch.float32, device=device) / half)) + freqs = (position_ids.float().unsqueeze(-1) * self._inv_freq).flatten(1) + # per-token rope rows [cos | sin] consumed by the NeoX kernel via positions=arange + cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1).contiguous() + return cache, torch.arange(position_ids.shape[0], dtype=torch.int32, device=device) + + @torch.inference_mode() + def forward(self, pixel_values: torch.Tensor, grid_thw: List[List[int]]) -> torch.Tensor: + from transformers.vision_utils import get_vision_cu_seqlens, get_vision_position_ids, get_vision_window_index + + vc = self._vc + device = self.patch_embedder.patch_embedding.weight.device + # the grid stays on the host: the index helpers loop in Python, so a device grid would sync per call + grid = torch.tensor(grid_thw, dtype=torch.long) + # window layers attend inside pos_emb_side x pos_emb_side patch tiles; tokens are reordered tile by tile for the whole stack + window_index, cu_window = get_vision_window_index( + grid, spatial_merge_size=1, window_size=vc.pos_emb_side * vc.patch_size, patch_size=vc.patch_size + ) + x = self.ln_pre.forward(self.patch_embedder.forward(pixel_values.to(device), grid))[window_index.to(device)] + # the reference feeds rope (w + 1, h + 1) + position_ids = (get_vision_position_ids(grid, 1).flip(-1) + 1)[window_index].to(device) + cache, positions = self._rope_table(position_ids) + lengths = { + "full_attention": torch.diff(get_vision_cu_seqlens(grid)).tolist(), + "window_attention": torch.diff(cu_window).tolist(), + } + for i, layer in self._layers(): + x = layer.forward(x, cache, positions, lengths[vc.layer_types[i]]) + x = self.ln_post.forward(x[torch.argsort(window_index).to(device)]) + m = vc.merge_size + x = x[_shuffle_index(grid, m).to(device)].view(-1, m * m, vc.hidden_size).permute(0, 2, 1).reshape(-1, vc.out_hidden_size) + return self.perception_emb_norm.forward(self.projection.forward(self.adapter.forward(x))) + + +__all__ = ["MuseGlimmerVisionModel"] diff --git a/python/freetoken/models/muse_glimmer/weight.py b/python/freetoken/models/muse_glimmer/weight.py index bf82b3ab8..b9ce56f31 100644 --- a/python/freetoken/models/muse_glimmer/weight.py +++ b/python/freetoken/models/muse_glimmer/weight.py @@ -17,15 +17,13 @@ from freetoken.utils import cached_load_hf_config from tqdm import tqdm -# Vision stack of the multimodal wrapper -- served text-only, always dropped. -_VISION_PREFIXES = ( - "model.vision_tower.", - "model.vision_adapter.", - "model.vision_projection.", - "vision_tower.", - "vision_adapter.", - "vision_projection.", +# The image path loads under the wrapper's vision_tower, adapter and projection included, so one prefix covers it. +_VISION_RENAMES = ( + ("model.vision_tower.", "vision_tower."), + ("model.vision_adapter.", "vision_tower.adapter."), + ("model.vision_projection.", "vision_tower.projection."), ) +_VISION_QKV_PARTS = (".attn.q_proj", ".attn.k_proj", ".attn.v_proj") # Fused projections, concatenated on the output dim in this exact order to match the # model's merged-linear splits. The attention gate rides the q/k/v fusion (it is computed @@ -39,10 +37,8 @@ } -def _rename(raw_name: str) -> str | None: - """HF key -> FreeToken state-dict key, or None to skip (vision stack).""" - if raw_name.startswith(_VISION_PREFIXES): - return None +def _rename(raw_name: str) -> str: + """HF text key -> FreeToken state-dict key.""" if raw_name.startswith("model.language_model."): return "model." + raw_name[len("model.language_model.") :] if raw_name.startswith("language_model."): @@ -50,12 +46,36 @@ def _rename(raw_name: str) -> str | None: return raw_name # lm_head.weight +def _vision_name(raw_name: str) -> str | None: + """FreeToken state-dict key of an image-path tensor, None for text keys.""" + for hf_prefix, prefix in _VISION_RENAMES: + if raw_name.startswith(hf_prefix): + return prefix + raw_name[len(hf_prefix) :] + return None + + +def _vision_tensors(name: str, tensor: torch.Tensor, buf: dict) -> list[tuple[str, torch.Tensor]]: + """What one image-path tensor yields: q/k/v weights and biases wait in buf and leave as one fused attn.qkv.""" + stem, _, leaf = name.rpartition(".") + for idx, part in enumerate(_VISION_QKV_PARTS): + if stem.endswith(part): + key = f"{stem[: -len(part)]}.attn.qkv.{leaf}" + slots = buf.setdefault(key, {}) + slots[idx] = tensor + if len(slots) < len(_VISION_QKV_PARTS): + return [] + del buf[key] + return [(key, torch.cat([slots[i] for i in range(len(_VISION_QKV_PARTS))], dim=0))] + return [(name, tensor)] + + def iter_weights( model_path: str, device: torch.device, *, include_moe_experts: bool, include_non_moe: bool, + include_vision: bool = True, ) -> Iterator[tuple[str, torch.Tensor]]: if not include_non_moe: return # dense model: there is no experts-only pass @@ -63,10 +83,11 @@ def iter_weights( raise NotImplementedError("muse_glimmer weight loading currently supports TP=1 only") if detect_compressed_tensors_nvfp4(cached_load_hf_config(model_path)): - yield from _iter_weights_compressed_tensors(model_path, device) + yield from _iter_weights_compressed_tensors(model_path, device, include_vision) return fuse_buf: dict = {} + vision_buf: dict = {} for file in tqdm( iter_weight_files(model_path), desc="Loading weights", @@ -74,9 +95,12 @@ def iter_weights( ): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: for raw_name in f.keys(): - name = _rename(raw_name) - if name is None: + vision = _vision_name(raw_name) + if vision is not None: + if include_vision: + yield from _vision_tensors(vision, f.get_tensor(raw_name), vision_buf) continue + name = _rename(raw_name) tensor = f.get_tensor(raw_name) if name.endswith(".weight"): emit = ct_bf16_fuse(name[: -len(".weight")], tensor, fuse_buf, _FUSIONS) @@ -88,10 +112,11 @@ def iter_weights( yield name, tensor assert not fuse_buf, f"Incomplete projection fusions: {list(fuse_buf.keys())}" + assert not vision_buf, f"Incomplete vision qkv fusions: {list(vision_buf.keys())}" def _iter_weights_compressed_tensors( - model_path: str, device: torch.device + model_path: str, device: torch.device, include_vision: bool ) -> Iterator[tuple[str, torch.Tensor]]: """Dense pass for the compressed-tensors NVFP4 checkpoint. @@ -99,10 +124,11 @@ def _iter_weights_compressed_tensors( (W4A16): ``.weight`` (uint8 packed) + ``.weight_scale`` (fp8 block) + ``.weight_global`` (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v/gate fuse into ``qkvg_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping - its own scales, so the fused FP4 weights are exact. Embeddings, norms and lm_head are - bf16 (the checkpoint's ignore list). Scale lookups go through the shard-map reader: - they can land in a different shard than their weight_packed.""" + its own scales, so the fused FP4 weights are exact. Embeddings, norms, lm_head and the + vision tower are bf16 (the checkpoint's ignore list). Scale lookups go through the + shard-map reader: they can land in a different shard than their weight_packed.""" nvfp4_buf: dict = {} + vision_buf: dict = {} reader = ShardReader(model_path, device) try: for file in tqdm( @@ -113,9 +139,12 @@ def _iter_weights_compressed_tensors( for raw_name in reader.names_in(file): if raw_name.endswith(CT_SCALE_SUFFIXES): continue # consumed with their weight_packed - name = _rename(raw_name) - if name is None: + vision = _vision_name(raw_name) + if vision is not None: + if include_vision: + yield from _vision_tensors(vision, reader.get_tensor(raw_name), vision_buf) continue + name = _rename(raw_name) if raw_name.endswith(".weight_packed"): base = name[: -len(".weight_packed")] parts = nvfp4_parts_ct(reader, raw_name[: -len(".weight_packed")]) @@ -135,6 +164,7 @@ def _iter_weights_compressed_tensors( reader.close() assert not nvfp4_buf, f"Incomplete NVFP4 fusions: {list(nvfp4_buf.keys())}" + assert not vision_buf, f"Incomplete vision qkv fusions: {list(vision_buf.keys())}" __all__ = ["iter_weights"] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index d84d2ca66..1d7ac3e5d 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -77,6 +77,8 @@ class ModelSpec: _GEMMA4_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4MMProcessor" _GEMMA4_UNIFIED_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4UnifiedMMProcessor" _GEMMA4_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),) +_MUSE_GLIMMER_PROCESSOR = "freetoken.mm.processors.muse_glimmer:MuseGlimmerMMProcessor" +_MUSE_GLIMMER_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),) _MINIMAX_M3_PACKED = _DENSE_PACKED + ( ("index_qk_proj", ("index_q_proj", "index_k_proj")), ) + _EXPERTS_W123_PACKED @@ -196,18 +198,21 @@ class ModelSpec: unquantized_modules=_QWEN3_5_UNQUANTIZED, ), # Muse-Glimmer-30B (model_type muse_glimmer): multimodal wrapper config (text tower in - # text_config, weights under model.language_model.); served text-only. Dense gated GQA - # with a [SWA x3, full] pattern -- full layers are NoPE -- weightless qk norms, centered - # (1+w) sandwich norms and softcapped logits; the NVFP4 release is compressed-tensors - # W4A16 on every text Linear. + # text_config, weights under model.language_model.); the windowed ViT under + # model.vision_tower. serves image input. Dense gated GQA with a [SWA x3, full] + # pattern -- full layers are NoPE -- weightless qk norms, centered (1+w) sandwich norms + # and softcapped logits; the NVFP4 release is compressed-tensors W4A16 on every text + # Linear. "MuseGlimmerForConditionalGeneration": ModelSpec( "freetoken.models.muse_glimmer", - "MuseGlimmerForCausalLM", + "MuseGlimmerForConditionalGeneration", checkpoint_roots=_LANGUAGE_MODEL_ROOT, packed_modules_mapping=( ("qkvg_proj", ("q_proj", "k_proj", "v_proj", "gate_proj")), ("gate_up_proj", ("gate_proj", "up_proj")), ), + mm_processor=_MUSE_GLIMMER_PROCESSOR, + encoders=_MUSE_GLIMMER_ENCODERS, ), "MistralForCausalLM": ModelSpec( "freetoken.models.mistral", diff --git a/tests/README.md b/tests/README.md index 9b1d7fdc7..1f90f6a38 100644 --- a/tests/README.md +++ b/tests/README.md @@ -53,6 +53,7 @@ checkpoint is set: | `FREETOKEN_GEMMA4_MODEL` | `models/test_gemma4_vision.py`, `tokenizer/test_mm_tokenize.py` — local Gemma-4 `gemma4` checkpoint (tower parity, boi/eoi expansion; the tokenizer test also takes the 12B `gemma4_unified` checkpoint) | | `FREETOKEN_GEMMA4_UNIFIED_MODEL` | `models/test_gemma4_vision.py` — local gemma-4-12B `gemma4_unified` checkpoint (embedder parity) | | `FREETOKEN_GLM53_MODEL` | `tokenizer/test_mm_tokenize.py` — local GLM-5.3-Flash checkpoint (image-token expansion) | +| `FREETOKEN_MUSE_MODEL` | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py` — local Muse-Glimmer-30B checkpoint (tower parity, image_start/image_end expansion) | | `FREETOKEN_TEST_MOE_CACHE_SIZE` | `e2e/test_aime.py` — >0 switches to the offload MoE backend with this cache size | | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) | | `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) | diff --git a/tests/mm/test_processor.py b/tests/mm/test_processor.py index 86668972c..c3049256a 100644 --- a/tests/mm/test_processor.py +++ b/tests/mm/test_processor.py @@ -14,6 +14,7 @@ from freetoken.mm.processor import MMProcessor, PromptReplacement, image_positions from freetoken.mm.processors.gemma4 import Gemma4MMProcessor, Gemma4UnifiedMMProcessor from freetoken.mm.processors.glm5_next import Glm5NextMMProcessor +from freetoken.mm.processors.muse_glimmer import MuseGlimmerMMProcessor from freetoken.mm.processors.qwen_vl import QwenVLMMProcessor PLACEHOLDER = 7 @@ -312,6 +313,43 @@ def __call__(self, images, **kwargs): assert dummy.feature.shape == (4, 1176) and dummy.num_tokens == 1 +def _muse_config(): + return SimpleNamespace( + architectures=["MuseGlimmerForConditionalGeneration"], + image_token_id=200092, + vision_config=SimpleNamespace(merge_size=2, patch_temporal=2, patch_size=14), + text_config=SimpleNamespace(vocab_size=202048), + ) + + +def test_muse_wraps_the_pads_in_image_start_end_and_takes_only_a_token_maximum(): + calls = [] + + class _FakeImageProcessor: + def __call__(self, images, **kwargs): + calls.append(kwargs) + return {"pixel_values": torch.zeros(16, 1176), "image_grid_thw": torch.tensor([[1, 4, 4]])} + + mm = MultimodalConfig(image_min_tokens=32, image_max_tokens=1024, processor_kwargs={"do_convert_rgb": False}) + proc = MuseGlimmerMMProcessor(_muse_config(), "/nonexistent", mm) + proc._image_processor = lambda: _FakeImageProcessor() + (item,) = proc.process([object()]) + # the image processor has no minimum knob: only the maximum reaches it, the extra kwarg rides along + assert calls == [{"return_tensors": "pt", "max_image_tokens": 1024, "do_convert_rgb": False}] + assert item.grid_thw == [1, 4, 4] and item.feature.dtype == torch.bfloat16 + # the template renders one <|patch|>; the replacement adds the start/end wrapper the checkpoint processor renders + repl = proc.prompt_replacement(item) + assert proc.placeholder == [200092] and repl.full == [200080, *[200092] * 4, 200081] + assert repl.embed_spans() == [[1, 5]] and proc.positions(10, [item]) is None + proc = MuseGlimmerMMProcessor(_muse_config(), "/nonexistent", MultimodalConfig()) + proc._image_processor = lambda: _FakeImageProcessor() + proc.process([object()]) + assert calls[-1] == {"return_tensors": "pt"} + (dummy,) = proc.dummy_items(torch.bfloat16, torch.device("cpu")) + dummy.validate() + assert dummy.feature.shape == (4, 1176) and dummy.num_tokens == 1 + + def test_registry_resolves_by_architecture(monkeypatch): import freetoken.utils from freetoken.mm.processor import get_mm_processor @@ -321,6 +359,7 @@ def test_registry_resolves_by_architecture(monkeypatch): "/gemma": _gemma_config(), "/gemma-unified": _gemma_unified_config(), "/glm": _glm_config(), + "/muse": _muse_config(), "/text-only": _hf_config(vision=False), "/unknown-vlm": _hf_config(arch="SomeOtherForConditionalGeneration"), } @@ -330,6 +369,7 @@ def test_registry_resolves_by_architecture(monkeypatch): assert isinstance(get_mm_processor("/gemma"), Gemma4MMProcessor) assert isinstance(get_mm_processor("/gemma-unified"), Gemma4UnifiedMMProcessor) assert isinstance(get_mm_processor("/glm"), Glm5NextMMProcessor) + assert isinstance(get_mm_processor("/muse"), MuseGlimmerMMProcessor) assert get_mm_processor("/text-only") is None assert get_mm_processor("/unknown-vlm") is None assert get_mm_processor("/missing") is None # a config that fails to load means no vision diff --git a/tests/models/test_muse_glimmer.py b/tests/models/test_muse_glimmer.py index f89330cdf..d35ff503a 100644 --- a/tests/models/test_muse_glimmer.py +++ b/tests/models/test_muse_glimmer.py @@ -21,10 +21,31 @@ class _Cfg: def __init__(self, data: dict): for k, v in data.items(): - setattr(self, k, _Cfg(v) if isinstance(v, dict) and k == "text_config" else v) + setattr(self, k, _Cfg(v) if isinstance(v, dict) and k in ("text_config", "vision_config") else v) + + +def _vision_section() -> dict: + """A tiny vision_config in the checkpoint's shape (the real tower is 50 x 1536).""" + return { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 4, + "num_attention_heads": 2, + "hidden_act": "gelu", + "layer_types": ["window_attention"] * 3 + ["full_attention"], + "patch_size": 2, + "patch_temporal": 2, + "merge_size": 2, + "pos_emb_height": 4, + "pos_emb_width": 4, + "layer_norm_eps": 1e-5, + "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, + "max_position_embeddings": 16, + "model_type": "muse_glimmer_vision", + } -def _hf_config(num_layers: int = 52, quantized: bool = False) -> _Cfg: +def _hf_config(num_layers: int = 52, quantized: bool = False, vision: bool = False) -> _Cfg: pattern = ["sliding_attention"] * 3 + ["full_attention"] thetas = [500000.0] * 3 + [0.0] reps = (num_layers + 3) // 4 @@ -55,6 +76,13 @@ def _hf_config(num_layers: int = 52, quantized: bool = False) -> _Cfg: "model_type": "muse_glimmer_text", }, } + if vision: + data |= { + "vision_config": _vision_section(), + "out_hidden_size": 16 * 4, + "projector_hidden_size": 24, + "projector_hidden_act": "gelu", + } if quantized: data["quantization_config"] = { "quant_method": "compressed-tensors", @@ -91,7 +119,7 @@ def test_parse_config_full_model(): assert cfg.final_logit_softcapping == 20.0 assert cfg.output_multiplier == pytest.approx(0.19611613513818404) assert cfg.embedding_scale is None # NormedEmbedding, not Gemma's sqrt(hidden) scale - assert cfg.vision_config is None # served text-only + assert cfg.vision_config is None # no vision_config section: the engine asked for no tower groups = {g.name: g for g in cfg.attention_groups} assert groups["swa"].layer_ids == tuple(i for i in range(52) if (i + 1) % 4 != 0) @@ -112,6 +140,27 @@ def test_parse_config_full_model(): assert cfg.lm_head_quant == "none" +def test_parse_vision_config(): + from freetoken.models.muse_glimmer.config import parse_vision_config + + vc = parse_config(_hf_config(vision=True)).vision_config + assert (vc.hidden_size, vc.intermediate_size, vc.num_layers, vc.num_heads) == (16, 32, 4, 2) + assert vc.layer_types == ("window_attention",) * 3 + ("full_attention",) + assert (vc.patch_size, vc.temporal_patch_size, vc.merge_size, vc.pos_emb_side) == (2, 2, 2, 4) + assert vc.patch_dim == 2 * 3 * 2 * 2 and vc.out_hidden_size == 64 + assert (vc.projector_hidden_size, vc.text_hidden_size, vc.text_rms_norm_eps) == (24, 6656, 1e-5) + assert vc.layer_norm_eps == 1e-5 and vc.rope_theta == 10000.0 + + hf = _hf_config(vision=True) + hf.vision_config.hidden_act = "silu" + with pytest.raises(NotImplementedError, match="only gelu"): + parse_vision_config(hf) + hf = _hf_config(vision=True) + hf.out_hidden_size = 16 # the adapter must take the pixel-shuffled width + with pytest.raises(AssertionError, match="pixel-shuffled"): + parse_vision_config(hf) + + def test_parse_config_nvfp4_checkpoint(): cfg = parse_config(_hf_config(quantized=True)) # compressed-tensors NVFP4 quantizes every text Linear (attention incl. the gate, @@ -162,7 +211,9 @@ def test_registry_resolves_architecture(): spec = get_model_spec("MuseGlimmerForConditionalGeneration") assert spec.module == "freetoken.models.muse_glimmer" - assert spec.model_cls == "MuseGlimmerForCausalLM" + assert spec.model_cls == "MuseGlimmerForConditionalGeneration" + assert spec.mm_processor == "freetoken.mm.processors.muse_glimmer:MuseGlimmerMMProcessor" + assert [(e.kind, e.config_key, e.modalities) for e in spec.encoders] == [("vision", "vision_config", ("image",))] def test_aot_table_covers_the_checkpoints(): @@ -179,17 +230,27 @@ def test_weight_rename_and_fusion(): import torch from freetoken.models.loader import ct_bf16_fuse - from freetoken.models.muse_glimmer.weight import _FUSIONS, _rename + from freetoken.models.muse_glimmer.weight import _FUSIONS, _rename, _vision_name, _vision_tensors - # Text tower renamed, vision dropped, lm_head untouched. + # Text tower renamed, lm_head untouched; the image path lands under the wrapper's vision_tower. assert _rename("model.language_model.layers.0.self_attn.q_proj.weight") == ( "model.layers.0.self_attn.q_proj.weight" ) assert _rename("model.language_model.embed_tokens.weight") == "model.embed_tokens.weight" assert _rename("lm_head.weight") == "lm_head.weight" - assert _rename("model.vision_tower.layers.0.attn.q_proj.weight") is None - assert _rename("model.vision_adapter.fc1.weight") is None - assert _rename("model.vision_projection.weight") is None + assert _vision_name("model.language_model.norm.weight") is None + assert _vision_name("model.vision_tower.layers.0.attn.q_proj.weight") == "vision_tower.layers.0.attn.q_proj.weight" + assert _vision_name("model.vision_adapter.fc1.weight") == "vision_tower.adapter.fc1.weight" + assert _vision_name("model.vision_projection.weight") == "vision_tower.projection.weight" + + # vision q/k/v fuse into attn.qkv in that order, biases like weights; other tensors pass through + vbuf: dict = {} + assert _vision_tensors("vision_tower.layers.0.attn.q_proj.bias", torch.zeros(2), vbuf) == [] + assert _vision_tensors("vision_tower.layers.0.attn.v_proj.bias", torch.full((2,), 2.0), vbuf) == [] + ((key, bias),) = _vision_tensors("vision_tower.layers.0.attn.k_proj.bias", torch.ones(2), vbuf) + assert key == "vision_tower.layers.0.attn.qkv.bias" and bias.tolist() == [0, 0, 1, 1, 2, 2] and not vbuf + ((key, weight),) = _vision_tensors("vision_tower.layers.0.attn.proj.weight", torch.zeros(2, 2), vbuf) + assert key == "vision_tower.layers.0.attn.proj.weight" and torch.equal(weight, torch.zeros(2, 2)) # q/k/v + the attention gate fuse into qkvg_proj in declaration order. buf: dict = {} @@ -234,6 +295,41 @@ def _write_shards(tmp_path, shards: dict[str, dict]): ) +def _vision_checkpoint_tensors(hf) -> dict: + """The image path in the checkpoint's naming: separate q/k/v with biases, an nn.Embedding position table, bias-free adapter and projection.""" + import torch + + vc = hf.vision_config + H, I, P = vc.hidden_size, vc.intermediate_size, hf.projector_hidden_size + bf16 = torch.bfloat16 + tensors = { + "model.vision_tower.patch_embedder.patch_embedding.weight": torch.randn(H, vc.patch_temporal * 3 * vc.patch_size**2, dtype=bf16), + "model.vision_tower.patch_embedder.position_embedding_table.weight": torch.randn(vc.pos_emb_height * vc.pos_emb_width, H, dtype=bf16), + "model.vision_tower.ln_pre.weight": torch.randn(H, dtype=bf16), + "model.vision_tower.ln_pre.bias": torch.randn(H, dtype=bf16), + "model.vision_tower.ln_post.weight": torch.randn(H, dtype=bf16), + "model.vision_tower.ln_post.bias": torch.randn(H, dtype=bf16), + "model.vision_adapter.fc1.weight": torch.randn(P, hf.out_hidden_size, dtype=bf16), + "model.vision_adapter.fc2.weight": torch.randn(P, P, dtype=bf16), + "model.vision_projection.weight": torch.randn(hf.text_config.hidden_size, P, dtype=bf16), + } + for i in range(vc.num_hidden_layers): + lp = f"model.vision_tower.layers.{i}." + for name in ("attn.q_proj", "attn.k_proj", "attn.v_proj", "attn.proj"): + tensors[lp + name + ".weight"] = torch.randn(H, H, dtype=bf16) + tensors[lp + name + ".bias"] = torch.randn(H, dtype=bf16) + tensors |= { + lp + "mlp.fc1.weight": torch.randn(I, H, dtype=bf16), + lp + "mlp.fc1.bias": torch.randn(I, dtype=bf16), + lp + "mlp.fc2.weight": torch.randn(H, I, dtype=bf16), + lp + "mlp.fc2.bias": torch.randn(H, dtype=bf16), + } + for name in ("norm1", "norm2"): + tensors[lp + name + ".weight"] = torch.randn(H, dtype=bf16) + tensors[lp + name + ".bias"] = torch.randn(H, dtype=bf16) + return tensors + + def _bf16_checkpoint_tensors(hf) -> dict: import torch @@ -246,11 +342,7 @@ def _bf16_checkpoint_tensors(hf) -> dict: p + "embed_tokens.weight": torch.randn(text.vocab_size, H, dtype=torch.bfloat16), p + "norm.weight": torch.randn(H, dtype=torch.bfloat16), "lm_head.weight": torch.randn(text.vocab_size, H, dtype=torch.bfloat16), - # vision tensors must be dropped, not loaded - "model.vision_tower.ln_pre.weight": torch.randn(4, dtype=torch.bfloat16), - "model.vision_adapter.fc1.weight": torch.randn(4, 4, dtype=torch.bfloat16), - "model.vision_projection.weight": torch.randn(4, 4, dtype=torch.bfloat16), - } + } | _vision_checkpoint_tensors(hf) for i in range(text.num_hidden_layers): lp = f"{p}layers.{i}." tensors |= { @@ -270,33 +362,43 @@ def _bf16_checkpoint_tensors(hf) -> dict: return tensors -def test_iter_weights_bf16_matches_model_state_dict(tmp_path, monkeypatch): +@pytest.mark.parametrize("include_vision", [False, True]) +def test_iter_weights_bf16_matches_model_state_dict(tmp_path, monkeypatch, include_vision): """The BF16 loader must produce exactly the model's state-dict keys with the - right shapes (rename + qkvg / gate_up fusion, vision dropped, norms raw).""" + right shapes (rename + qkvg / gate_up fusion, norms raw); the image path only + when an encoder is built, then with q/k/v fused into attn.qkv.""" import torch from freetoken.distributed import set_tp_info, try_get_tp_info if try_get_tp_info() is None: set_tp_info(rank=0, size=1) - from freetoken.models.muse_glimmer.model import MuseGlimmerForCausalLM + from freetoken.models.muse_glimmer.model import MuseGlimmerForConditionalGeneration from freetoken.models.muse_glimmer.weight import iter_weights - hf = _hf_config(num_layers=4) + hf = _hf_config(num_layers=4, vision=True) tensors = _bf16_checkpoint_tensors(hf) _write_shards(tmp_path, {"model-00001-of-00001.safetensors": tensors}) import freetoken.models.muse_glimmer.weight as w monkeypatch.setattr(w, "cached_load_hf_config", lambda _p: hf) - loaded = dict( - iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True) - ) - model = MuseGlimmerForCausalLM(parse_config(hf)) + loaded = dict(iter_weights( + str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True, include_vision=include_vision, + )) + if not include_vision: + hf.vision_config = None # what the engine's text-only model_config looks like + model = MuseGlimmerForConditionalGeneration(parse_config(hf)) expected = model.state_dict() assert set(loaded) == set(expected) + assert any(k.startswith("vision_tower.") for k in expected) is include_vision for k in expected: assert loaded[k].shape == expected[k].shape, k + if include_vision: + v = "model.vision_tower.layers.1.attn." + assert torch.equal(loaded["vision_tower.layers.1.attn.qkv.weight"], torch.cat([tensors[v + f"{n}_proj.weight"] for n in "qkv"])) + assert torch.equal(loaded["vision_tower.layers.1.attn.qkv.bias"], torch.cat([tensors[v + f"{n}_proj.bias"] for n in "qkv"])) + assert torch.equal(loaded["vision_tower.projection.weight"], tensors["model.vision_projection.weight"]) # fusion order [q, k, v, gate] against the raw parts fused = loaded["model.layers.0.self_attn.qkvg_proj.weight"] p = "model.language_model.layers.0.self_attn." @@ -311,12 +413,18 @@ def test_iter_weights_nvfp4_cross_shard_scales(tmp_path, monkeypatch): """compressed-tensors loader: native FP4 parts fused with per-part scales, the reciprocal global, and sibling scales resolved through the index even when they land in a different shard than their weight_packed (the real checkpoint splits - layer 49's down_proj across the shard boundary).""" + layer 49's down_proj across the shard boundary); the bf16 image path rides along + with its q/k/v fused across shards.""" import torch + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + from freetoken.models.muse_glimmer import MuseGlimmerVisionModel, parse_vision_config from freetoken.models.muse_glimmer.weight import iter_weights - hf = _hf_config(num_layers=1, quantized=True) + hf = _hf_config(num_layers=1, quantized=True, vision=True) text = hf.text_config text.hidden_size, text.intermediate_size = 64, 96 text.num_attention_heads, text.num_key_value_heads, text.head_dim = 4, 2, 16 @@ -357,6 +465,11 @@ def nvfp4(base, out_f, in_f, global_scale, shard_for_scales=None): "pre_feedforward_layernorm", "post_feedforward_layernorm", ): shard1[p + name + ".weight"] = torch.randn(H, dtype=torch.bfloat16) + # the image path stays bf16; layer 0's q_proj lands in shard 1 and its k/v in shard 2 + vision = _vision_checkpoint_tensors(hf) + q_keys = {k for k in vision if ".layers.0.attn.q_proj." in k} + shard1 |= {k: vision[k] for k in q_keys} + shard2 |= {k: v for k, v in vision.items() if k not in q_keys} _write_shards(tmp_path, { "model-00001-of-00002.safetensors": shard1, "model-00002-of-00002.safetensors": shard2, @@ -381,6 +494,16 @@ def nvfp4(base, out_f, in_f, global_scale, shard_for_scales=None): assert loaded[dp + ".input_scale"].item() == pytest.approx(1.0) assert loaded[qkvg + ".input_scale"].shape == () assert not any(k.endswith(".input_global_scale") for k in loaded) + expected_vision = {"vision_tower." + k for k in MuseGlimmerVisionModel(parse_vision_config(hf)).state_dict()} + assert {k for k in loaded if k.startswith("vision_tower.")} == expected_vision + v = "model.vision_tower.layers.0.attn." + assert torch.equal(loaded["vision_tower.layers.0.attn.qkv.weight"], torch.cat([vision[v + f"{n}_proj.weight"] for n in "qkv"])) + assert torch.equal(loaded["vision_tower.layers.0.attn.qkv.bias"], torch.cat([vision[v + f"{n}_proj.bias"] for n in "qkv"])) + assert loaded["vision_tower.projection.weight"].dtype == torch.bfloat16 + text_only = dict(iter_weights( + str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True, include_vision=False, + )) + assert not any(k.startswith("vision_tower.") for k in text_only) and qkvg + ".weight" in text_only def test_raw_config_shim_serves_unknown_model_type(tmp_path): diff --git a/tests/models/test_muse_glimmer_vision.py b/tests/models/test_muse_glimmer_vision.py new file mode 100644 index 000000000..6e55d34b1 --- /dev/null +++ b/tests/models/test_muse_glimmer_vision.py @@ -0,0 +1,193 @@ +"""Muse-Glimmer vision tower: layout against a tiny reference tower and host streaming on random weights, parity with the reference on a real checkpoint.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from freetoken.models.muse_glimmer import MuseGlimmerVisionModel, VisionConfig + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +CHECKPOINT = os.environ.get("FREETOKEN_MUSE_MODEL", "") +needs_checkpoint = pytest.mark.skipif(not os.path.exists(os.path.join(CHECKPOINT, "config.json")), reason="FREETOKEN_MUSE_MODEL not set") + + +@pytest.fixture(autouse=True) +def _single_rank(): + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +def _tiny_vc(): + # pos_emb_side 4 makes the window 4 patches wide, so a 6x6 grid has ragged 4/2 windows + return VisionConfig( + hidden_size=64, intermediate_size=128, num_layers=4, num_heads=4, + layer_types=("window_attention", "window_attention", "window_attention", "full_attention"), + patch_size=4, temporal_patch_size=2, merge_size=2, pos_emb_side=4, layer_norm_eps=1e-5, rope_theta=1e4, + projector_hidden_size=48, text_hidden_size=40, text_rms_norm_eps=1e-5, + ) + + +def _build(vc): + # bf16 like the engine: the weightless norm kernel takes no fp32 + torch.set_default_dtype(torch.bfloat16) + try: + with torch.device("cuda"): + tower = MuseGlimmerVisionModel(vc) + finally: + torch.set_default_dtype(torch.float32) + for p in tower.state_dict().values(): + p.normal_(0, 0.02) + return tower + + +def _reference(vc, state, dtype): + """The reference's get_image_features on our tower's weights (attn.qkv split back into q/k/v).""" + from transformers.models.muse_glimmer.configuration_muse_glimmer import MuseGlimmerVisionConfig + from transformers.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerRMSNorm, MuseGlimmerVisionAdapter + from transformers.models.muse_glimmer.modeling_muse_glimmer import MuseGlimmerVisionModel as HFVision + + config = MuseGlimmerVisionConfig( + hidden_size=vc.hidden_size, intermediate_size=vc.intermediate_size, num_hidden_layers=vc.num_layers, + num_attention_heads=vc.num_heads, layer_types=list(vc.layer_types), patch_size=vc.patch_size, + patch_temporal=vc.temporal_patch_size, merge_size=vc.merge_size, pos_emb_height=vc.pos_emb_side, + pos_emb_width=vc.pos_emb_side, layer_norm_eps=vc.layer_norm_eps, max_position_embeddings=vc.pos_emb_side**2, + rope_parameters={"rope_theta": vc.rope_theta, "rope_type": "default"}, + ) + config._attn_implementation = "sdpa" + wrapper = SimpleNamespace(out_hidden_size=vc.out_hidden_size, projector_hidden_size=vc.projector_hidden_size, projector_hidden_act="gelu") + # the default dtype sets the parameters; the rotary table is built fp32 regardless, as from_pretrained keeps it + torch.set_default_dtype(dtype) + try: + tower = HFVision(config).to("cuda").eval() + adapter = MuseGlimmerVisionAdapter(wrapper).to("cuda") + projection = torch.nn.Linear(vc.projector_hidden_size, vc.text_hidden_size, bias=False).to("cuda") + finally: + torch.set_default_dtype(torch.float32) + norm = MuseGlimmerRMSNorm(eps=vc.text_rms_norm_eps, with_scale=False) + hf_state = {} + for key, value in state.items(): + if ".attn.qkv." in key: + for name, part in zip(("q_proj", "k_proj", "v_proj"), value.chunk(3, dim=0)): + hf_state[key.replace("attn.qkv", f"attn.{name}")] = part.to(dtype) + else: + hf_state[key] = value.to(dtype) + tower.load_state_dict({k: v for k, v in hf_state.items() if not k.startswith(("adapter.", "projection."))}, strict=True) + adapter.load_state_dict({k[len("adapter.") :]: v for k, v in hf_state.items() if k.startswith("adapter.")}, strict=True) + projection.load_state_dict({"weight": hf_state["projection.weight"]}, strict=True) + return lambda pixels, grid: norm(projection(adapter(tower(pixels.to(dtype), grid_thw=grid).last_hidden_state))) + + +def test_tiny_tower_matches_the_reference_layout(): + """Rope layout, window permutation, pixel shuffle, adapter and norm against the reference on random weights.""" + vc = _tiny_vc() + tower = _build(vc) + reference = _reference(vc, tower.state_dict(), torch.bfloat16) + pixels = torch.randn(36 + 16, vc.patch_dim, device="cuda", dtype=torch.bfloat16) + grid = [[1, 6, 6], [1, 4, 4]] + with torch.inference_mode(): + ref = reference(pixels, torch.tensor(grid, device="cuda")).float() + ours = tower.forward(pixels, grid).float() + assert ours.shape == ref.shape == (9 + 4, 40) + cos = F.cosine_similarity(ref, ours, dim=-1) + rel = ((ref - ours).norm() / ref.norm()).item() + print(f"muse_glimmer tiny tower vs reference: min cos {cos.min():.5f}, relative error {rel:.4f}") + assert cos.min() > 0.99 and rel < 5e-2 + keys = tower.state_dict() + assert keys["layers.0.attn.qkv.bias"].shape == (192,) and keys["patch_embedder.position_embedding_table.weight"].shape == (16, 64) + assert keys["adapter.fc1.weight"].shape == (48, 256) and keys["projection.weight"].shape == (40, 48) + assert "adapter.fc1.bias" not in keys and "projection.bias" not in keys and "perception_emb_norm.weight" not in keys + + +def test_packed_images_match_the_single_image_calls(): + vc = _tiny_vc() + tower = _build(vc) + big = torch.randn(36, vc.patch_dim, device="cuda", dtype=torch.bfloat16) + small = torch.randn(16, vc.patch_dim, device="cuda", dtype=torch.bfloat16) + packed = tower.forward(torch.cat([big, small]), [[1, 6, 6], [1, 4, 4]]).float() + singles = torch.cat([tower.forward(big, [[1, 6, 6]]), tower.forward(small, [[1, 4, 4]])]).float() + # windows and positions are per image; only the GEMM tiling differs between the two calls + assert ((packed - singles).norm() / singles.norm()).item() < 1e-2 + + +def test_host_streamed_weights_compute_the_same_output(): + vc = _tiny_vc() + tower = _build(vc) + feature = torch.randn(64, vc.patch_dim, device="cuda", dtype=torch.bfloat16) + resident = tower.forward(feature, [[1, 8, 8]]) + keys = tower.state_dict() + tower.place_weights("host") + assert tower.state_dict().keys() == keys.keys() + assert tower.state_dict()["layers.0.attn.qkv.weight"].device.type == "cpu" + assert tower.state_dict()["adapter.fc1.weight"].device.type == "cuda" + for _ in range(2): # a second forward reuses the staging buffers behind the release events + assert torch.equal(tower.forward(feature, [[1, 8, 8]]), resident) + tower.place_weights("gpu") + assert tower.state_dict()["layers.0.attn.qkv.weight"].device.type == "cuda" + assert torch.equal(tower.forward(feature, [[1, 8, 8]]), resident) + + +def _textured_image(width, height): + """Gradient, shapes and mild noise: on flat colour fields the bf16 reference itself drifts far from its fp32 result.""" + from PIL import Image, ImageDraw + + xs = torch.linspace(0, 255, width).expand(height, width) + ys = torch.linspace(0, 255, height)[:, None].expand(height, width) + rgb = torch.stack([xs, ys, 127 + (ys - xs) / 2], dim=-1) + rgb = rgb + torch.randn(height, width, 3, generator=torch.Generator().manual_seed(0)) * 12 + img = Image.fromarray(rgb.clamp(0, 255).to(torch.uint8).numpy()) + draw = ImageDraw.Draw(img) + draw.ellipse((width * 0.2, height * 0.15, width * 0.7, height * 0.7), fill=(255, 215, 0)) + draw.rectangle((width * 0.1, height * 0.8, width * 0.9, height * 0.9), fill=(220, 20, 60)) + return img + + +@pytest.mark.needs_weights +@needs_checkpoint +def test_tower_matches_the_reference_on_a_checkpoint(): + """The bf16 tower stays within the reference's own bf16 noise around its fp32 get_image_features.""" + import json + + from safetensors import safe_open + from transformers import AutoConfig, AutoImageProcessor + + from freetoken.mm.config import MultimodalConfig + from freetoken.mm.processors.muse_glimmer import MuseGlimmerMMProcessor + from freetoken.models.muse_glimmer import parse_config + from freetoken.models.muse_glimmer.weight import _vision_name, _vision_tensors + + torch.set_float32_matmul_precision("highest") + hf_config = AutoConfig.from_pretrained(CHECKPOINT) + config = parse_config(hf_config) + index = json.load(open(os.path.join(CHECKPOINT, "model.safetensors.index.json")))["weight_map"] + state, buf = {}, {} + for name, file in index.items(): + if _vision_name(name) is not None: + with safe_open(os.path.join(CHECKPOINT, file), "pt") as f: + for key, value in _vision_tensors(_vision_name(name)[len("vision_tower.") :], f.get_tensor(name).cuda(), buf): + state[key] = value + tower = _build(config.vision_config) + tower.load_state_dict(dict(state)) + + # 896x616 resizes to a 64x44 patch grid: four windows, the bottom row of them 12 patches high + img = _textured_image(896, 616) + out = AutoImageProcessor.from_pretrained(CHECKPOINT)(images=img, return_tensors="pt") + pixels, grid = out["pixel_values"].cuda(), out["image_grid_thw"].cuda() + with torch.inference_mode(): + ref32 = _reference(config.vision_config, state, torch.float32)(pixels, grid) + ref16 = _reference(config.vision_config, state, torch.bfloat16)(pixels, grid).float() + (item,) = MuseGlimmerMMProcessor(hf_config, CHECKPOINT, MultimodalConfig()).process([img]) + ours = tower.forward(item.feature, [item.grid_thw]).float() + assert item.grid_thw == [1, 44, 64] + assert ours.shape == ref32.shape == (44 * 64 // 4, config.hidden_size) + noise, error = (ref32 - ref16).norm(), (ref32 - ours).norm() + print(f"muse_glimmer tower: ours vs fp32 mean cos {F.cosine_similarity(ref32, ours, dim=-1).mean():.5f}, error/noise {error / noise:.3f}") + assert F.cosine_similarity(ref32, ours, dim=-1).mean() > 0.99 + assert error <= 1.2 * noise diff --git a/tests/tokenizer/test_mm_tokenize.py b/tests/tokenizer/test_mm_tokenize.py index b3866f7d4..05f883bdb 100644 --- a/tests/tokenizer/test_mm_tokenize.py +++ b/tests/tokenizer/test_mm_tokenize.py @@ -1,4 +1,4 @@ -"""Image tokenization against real Qwen VL, Gemma-4 and GLM-5.3 checkpoints (skipped when absent).""" +"""Image tokenization against real Qwen VL, Gemma-4, GLM-5.3 and Muse-Glimmer checkpoints (skipped when absent).""" from __future__ import annotations @@ -18,12 +18,13 @@ GEMMA = os.environ.get("FREETOKEN_GEMMA4_MODEL", "") GLM = os.environ.get("FREETOKEN_GLM53_MODEL", "") +MUSE = os.environ.get("FREETOKEN_MUSE_MODEL", "") pytestmark = [ pytest.mark.needs_weights, pytest.mark.skipif( - not MODELS and not any(os.path.exists(os.path.join(p, "config.json")) for p in (GEMMA, GLM)), - reason="no FREETOKEN_QWEN36_MODEL / FREETOKEN_QWEN3VL_MODEL / FREETOKEN_GEMMA4_MODEL / FREETOKEN_GLM53_MODEL checkpoint", + not MODELS and not any(os.path.exists(os.path.join(p, "config.json")) for p in (GEMMA, GLM, MUSE)), + reason="no FREETOKEN_QWEN36_MODEL / FREETOKEN_QWEN3VL_MODEL / FREETOKEN_GEMMA4_MODEL / FREETOKEN_GLM53_MODEL / FREETOKEN_MUSE_MODEL checkpoint", ), ] @@ -163,3 +164,28 @@ def test_glm_expands_the_image_token_between_the_template_wrappers(): assert bool((ids[start:end] == item.pad_value).all()) and r.mrope_positions is None budgeted = TokenizeManager(tokenizer, get_mm_processor(GLM, MultimodalConfig(image_max_tokens=100))).tokenize([_msg([_png(640, 400)])])[0] assert budgeted.mm_items[0].num_tokens <= 100 + + +@pytest.mark.skipif(not os.path.exists(os.path.join(MUSE, "config.json")), reason="FREETOKEN_MUSE_MODEL not set") +def test_muse_wraps_the_pad_span_in_image_start_end_and_clamps_to_a_token_maximum(): + from freetoken.mm.config import MultimodalConfig + from freetoken.mm.processor import get_mm_processor + from freetoken.tokenizer.tokenize import TokenizeManager + from freetoken.utils.hf import load_tokenizer + + tokenizer = load_tokenizer(MUSE) + r = TokenizeManager(tokenizer, get_mm_processor(MUSE)).tokenize([_msg([_png(640, 400)])])[0] + item = r.mm_items[0] + ((start, end),) = item.offsets + ids = r.input_ids + assert ids[start - 1].item() == tokenizer.convert_tokens_to_ids("<|image_start|>") + assert ids[end].item() == tokenizer.convert_tokens_to_ids("<|image_end|>") + # 640x400 resizes to a 44x28 patch grid; the single <|patch|> the template renders is gone + assert item.grid_thw == [1, 28, 44] and end - start == item.num_tokens == 308 and r.mrope_positions is None + assert bool((ids[start:end] == item.pad_value).all()) and not bool((ids == 200092).any()) + budgeted = TokenizeManager(tokenizer, get_mm_processor(MUSE, MultimodalConfig(image_max_tokens=1024))) + assert budgeted.tokenize([_msg([_png(3840, 2160)])])[0].mm_items[0].num_tokens <= 1024 < item.num_tokens * 4 + # an explicit processor kwarg replaces the checkpoint's 4096 cap: 640x400 under 256 tokens is a 24x40 grid + override = TokenizeManager(tokenizer, get_mm_processor(MUSE, MultimodalConfig(processor_kwargs={"max_image_tokens": 256}))) + small = override.tokenize([_msg([_png(640, 400)])])[0].mm_items[0] + assert small.grid_thw == [1, 24, 40] and small.num_tokens == 240