[Native] Map Muse Glimmer GGUF tensors onto its HF model

backends/native/examples/hf_llm/export_hf_llm.py

@@ -27,6 +27,7 @@
 import torch
 import torch.nn as nn
 from executorch.backends.native.examples.hf_llm.gguf_hf import (
+    check_config_tensors,
     gguf_architecture,
     gguf_to_hf,
 )
@@ -39,6 +40,7 @@
     SequenceReferenceCache,
 )
 from executorch.extension.llm.cache.update_and_attend import REGISTRY
+from executorch.extension.llm.export.gguf import iter_gguf
 from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter
 from executorch.extension.llm.export.model_metadata import write_cache_geometry
 from torch.export import Dim, export
@@ -130,9 +132,9 @@ def load_gguf(
     ``ExportableGGUFTensor`` so the PTN carries the GGUF blocks. ``architecture``
     defaults to the file's ``general.architecture``. ``truncated`` loads a model
     built with fewer layers from the file's first blocks."""
-    key_map, convert = gguf_to_hf(
-        model, architecture or gguf_architecture(path), truncated
-    )
+    architecture = architecture or gguf_architecture(path)
+    key_map, convert = gguf_to_hf(model, architecture, truncated)
+    read = raw_iter or iter_gguf
     shapes = {n: p.shape for n, p in model.named_parameters(remove_duplicate=False)}
     expected = set(shapes)
     loaded = set()
@@ -142,7 +144,7 @@ def load_gguf(
         convert=convert,
         tie_map={"model.embed_tokens.weight": ("lm_head.weight", False)},
         dtype=torch.float32,
-        raw_iter=raw_iter,
+        raw_iter=lambda p: check_config_tensors(model, architecture, read(p)),
     ):
         if fqn not in expected:
             raise ValueError(f"GGUF tensor maps to unknown parameter {fqn!r}")

backends/native/examples/hf_llm/gguf_hf.py

@@ -9,14 +9,16 @@
 Use with ``iter_checkpoint`` / ``load_checkpoint``: :func:`gguf_to_hf` returns the
 ``key_map`` and ``convert`` that load a GGUF file of a supported architecture into
 a Hugging Face model. Names come from gguf-py's standard tensor map, the one
-llama.cpp's ``convert_hf_to_gguf.py`` writes with, and ``convert`` undoes the
+llama.cpp's ``convert_hf_to_gguf.py`` writes with, or a table here for an
+architecture gguf-py does not know; ``convert`` undoes the
 weight changes that converter makes for the architecture. Those changes move whole
 rows or touch only unquantized norms, so quantized GGUF blocks stay byte-exact.
 """
 
 from __future__ import annotations
 
 import re
+from collections.abc import Iterable, Iterator
 from typing import Callable
 
 import gguf
@@ -81,12 +83,12 @@ def convert(fqn: str, value: torch.Tensor) -> torch.Tensor:
     return convert
 
 
-def _remove_norm_offset(config) -> Convert:
-    """Subtract the 1 llama.cpp adds to Gemma norm weights; HF Gemma adds it in
-    its forward instead."""
+def _norm_offset_remover(suffix: str) -> Convert:
+    """Subtract the 1 llama.cpp adds to the norm weights ending in ``suffix``,
+    for HF norms that add it in their forward instead."""
 
     def convert(fqn: str, value: torch.Tensor) -> torch.Tensor:
-        if not fqn.endswith("norm.weight"):
+        if not fqn.endswith(suffix):
             return value
         if isinstance(value, ExportableGGUFTensor):
             raise ValueError(f"{fqn}: expected an unquantized norm weight")
@@ -95,6 +97,61 @@ def convert(fqn: str, value: torch.Tensor) -> torch.Tensor:
     return convert
 
 
+def _remove_norm_offset(config) -> Convert:
+    """Gemma adds the 1 in every norm."""
+    return _norm_offset_remover("norm.weight")
+
+
+def _muse_glimmer_convert(config) -> Convert:
+    """Muse Glimmer rotates like Llama, and adds the 1 in its decoder layer
+    norms but not its final norm."""
+    restore_qk = _restore_qk_rows(config)
+    remove_offset = _norm_offset_remover("layernorm.weight")
+    return lambda fqn, value: remove_offset(fqn, restore_qk(fqn, value))
+
+
+_MG_TEXT = "model.language_model"
+# GGUF name -> HF module, for the blocks and the rest; None skips a tensor.
+_MUSE_GLIMMER_BLOCK = {
+    "attn_q": "self_attn.q_proj",
+    "attn_k": "self_attn.k_proj",
+    "attn_v": "self_attn.v_proj",
+    "attn_output": "self_attn.o_proj",
+    "attn_gate": "self_attn.gate_proj",
+    "attn_norm": "input_layernorm",
+    "post_attention_norm": "post_attention_layernorm",
+    "ffn_norm": "pre_feedforward_layernorm",
+    "post_ffw_norm": "post_feedforward_layernorm",
+    "ffn_gate": "mlp.gate_proj",
+    "ffn_up": "mlp.up_proj",
+    "ffn_down": "mlp.down_proj",
+    # The QK norm has no weight in HF; see _muse_glimmer_config_tensors.
+    "attn_q_norm": None,
+    "attn_k_norm": None,
+}
+_MUSE_GLIMMER_TOP = {
+    "token_embd": f"{_MG_TEXT}.embed_tokens",
+    "output_norm": f"{_MG_TEXT}.norm",
+    "output": "lm_head",
+}
+
+
+def _muse_glimmer_names(layers: int) -> dict[str, str | None]:
+    names = {f"{g}.weight": f"{h}.weight" for g, h in _MUSE_GLIMMER_TOP.items()}
+    for i in range(layers):
+        for g, h in _MUSE_GLIMMER_BLOCK.items():
+            names[f"blk.{i}.{g}.weight"] = (
+                f"{_MG_TEXT}.layers.{i}.{h}.weight" if h else None
+            )
+    return names
+
+
+def _muse_glimmer_config_tensors(config) -> dict[str, float]:
+    """llama.cpp's converter stores the query scale HF reads from the config as
+    the Q norm weight, and ones as the K norm weight."""
+    return {"attn_q_norm.weight": config.qk_scale_factor, "attn_k_norm.weight": 1.0}
+
+
 # GGUF architecture -> (the Hugging Face model types it holds, what undoes
 # llama.cpp's weight changes). llama.cpp stores Mistral as "llama".
 _ARCHITECTURES: dict[str, tuple[frozenset[str], Callable[..., Convert] | None]] = {
@@ -105,8 +162,44 @@ def convert(fqn: str, value: torch.Tensor) -> torch.Tensor:
     "gemma": (frozenset({"gemma"}), _remove_norm_offset),
     "gemma2": (frozenset({"gemma2"}), _remove_norm_offset),
     "gemma3": (frozenset({"gemma3_text"}), _remove_norm_offset),
+    "muse-glimmer": (frozenset({"muse_glimmer_text"}), _muse_glimmer_convert),
+}
+
+# Architectures gguf-py has no tensor map for: layer count -> GGUF name -> HF
+# parameter, or None to skip.
+_NAMES: dict[str, Callable[[int], dict[str, str | None]]] = {
+    "muse-glimmer": _muse_glimmer_names,
 }
 
+# Skipped GGUF block tensors that hold a value Hugging Face reads from its
+# config: config -> tensor -> that value.
+_CONFIG_TENSORS: dict[str, Callable[..., dict[str, float]]] = {
+    "muse-glimmer": _muse_glimmer_config_tensors,
+}
+
+
+def check_config_tensors(
+    model: torch.nn.Module,
+    architecture: str,
+    tensors: Iterable[tuple[str, torch.Tensor]],
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Pass a GGUF file's ``(name, tensor)`` pairs through, raising for a
+    skipped block tensor that disagrees with the value ``model``'s config gives
+    it, which loading would otherwise ignore."""
+    make = _CONFIG_TENSORS.get(architecture)
+    want = make(model.config.get_text_config()) if make else {}
+    for name, value in tensors:
+        block = re.match(r"blk\.\d+\.(.+)", name)
+        if block and block[1] in want:
+            expected = torch.tensor(want[block[1]])
+            if isinstance(value, ExportableGGUFTensor) or not torch.allclose(
+                value.float(), expected, rtol=1e-3, atol=0
+            ):
+                raise ValueError(
+                    f"{name} must hold {want[block[1]]}, as the config gives"
+                )
+        yield name, value
+
 
 def gguf_to_hf(
     model: torch.nn.Module, architecture: str, truncated: bool = False
@@ -129,13 +222,16 @@ def gguf_to_hf(
         raise ValueError(
             f"a {architecture!r} GGUF cannot load a {config.model_type!r} model"
         )
-    arch = next(a for a, n in gguf.MODEL_ARCH_NAMES.items() if n == architecture)
-    tensor_map = gguf.get_tensor_name_map(arch, config.num_hidden_layers)
-    names = {}
-    for fqn in model.state_dict():
-        name = tensor_map.get_name(fqn, try_suffixes=(".weight", ".bias"))
-        if name is not None:
-            names[name] = fqn
+    if architecture in _NAMES:
+        names = _NAMES[architecture](config.num_hidden_layers)
+    else:
+        arch = next(a for a, n in gguf.MODEL_ARCH_NAMES.items() if n == architecture)
+        tensor_map = gguf.get_tensor_name_map(arch, config.num_hidden_layers)
+        names = {}
+        for fqn in model.state_dict():
+            name = tensor_map.get_name(fqn, try_suffixes=(".weight", ".bias"))
+            if name is not None:
+                names[name] = fqn
 
     def key_map(name: str) -> str | None:
         if name in names:

backends/native/examples/hf_llm/tests/test_gguf_hf.py

@@ -12,6 +12,7 @@
 import numpy as np
 import torch
 from executorch.backends.native.examples.hf_llm.gguf_hf import (
+    check_config_tensors,
     gguf_architecture,
     gguf_to_hf,
 )
@@ -41,6 +42,38 @@ def _tiny(config_cls, heads: int = 4, kv_heads: int = 2, head_dim: int = 16, **k
     return AutoModelForCausalLM.from_config(config)
 
 
+def _tiny_muse_glimmer():
+    """A 4-layer Muse Glimmer with its vision tower, or None where transformers
+    has no Muse Glimmer."""
+    try:
+        from transformers import MuseGlimmerConfig, MuseGlimmerForConditionalGeneration
+    except ImportError:
+        return None
+    config = MuseGlimmerConfig(
+        text_config={
+            "vocab_size": 32,
+            "hidden_size": 64,
+            "intermediate_size": 32,
+            "num_hidden_layers": 4,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 16,
+            "sliding_window": 4,
+            "layer_types": ["sliding_attention"] * 3 + ["full_attention"],
+            "layer_rope_theta": [10000.0] * 3 + [0],
+        },
+        vision_config={
+            "hidden_size": 32,
+            "intermediate_size": 32,
+            "num_hidden_layers": 1,
+            "num_attention_heads": 2,
+        },
+        projector_hidden_size=32,
+        out_hidden_size=32,
+    )
+    return MuseGlimmerForConditionalGeneration(config)
+
+
 def _to_gguf_rows(weight: torch.Tensor, heads: int) -> torch.Tensor:
     """llama.cpp's convert_hf_to_gguf.py Q/K permutation."""
     return (
@@ -193,6 +226,77 @@ def test_removes_the_gemma_norm_offset(self) -> None:
                 ),
             )
 
+    def test_maps_muse_glimmer_names(self) -> None:
+        model = _tiny_muse_glimmer()
+        if model is None:
+            self.skipTest("transformers has no Muse Glimmer")
+        key_map, _ = gguf_to_hf(model, "muse-glimmer")
+        text = "model.language_model"
+        names = {
+            "token_embd.weight": f"{text}.embed_tokens.weight",
+            "output_norm.weight": f"{text}.norm.weight",
+            "output.weight": "lm_head.weight",
+            "blk.3.attn_gate.weight": f"{text}.layers.3.self_attn.gate_proj.weight",
+            "blk.0.ffn_norm.weight": f"{text}.layers.0.pre_feedforward_layernorm.weight",
+            "blk.1.post_ffw_norm.weight": (
+                f"{text}.layers.1.post_feedforward_layernorm.weight"
+            ),
+            "blk.2.attn_q_norm.weight": None,
+            "blk.2.attn_k_norm.weight": None,
+        }
+
+        self.assertEqual({n: key_map(n) for n in names}, names)
+        params = dict(model.named_parameters())
+        self.assertTrue(all(fqn in params for fqn in names.values() if fqn))
+        with self.assertRaisesRegex(ValueError, "maps to no muse_glimmer_text"):
+            key_map("blk.4.attn_q.weight")
+
+    def test_restores_muse_glimmer_weights(self) -> None:
+        model = _tiny_muse_glimmer()
+        if model is None:
+            self.skipTest("transformers has no Muse Glimmer")
+        _, convert = gguf_to_hf(model, "muse-glimmer")
+        text = "model.language_model"
+        q = torch.randn(64, 64)
+        decoder_norm = torch.randn(64)
+        final_norm = torch.randn(64)
+
+        self.assertTrue(
+            torch.equal(
+                convert(
+                    f"{text}.layers.0.self_attn.q_proj.weight", _to_gguf_rows(q, 4)
+                ),
+                q,
+            )
+        )
+        torch.testing.assert_close(
+            convert(f"{text}.layers.0.input_layernorm.weight", decoder_norm + 1),
+            decoder_norm,
+        )
+        self.assertIs(convert(f"{text}.norm.weight", final_norm), final_norm)
+
+    def test_checks_muse_glimmer_qk_norms_against_the_config(self) -> None:
+        model = _tiny_muse_glimmer()
+        if model is None:
+            self.skipTest("transformers has no Muse Glimmer")
+        scale = torch.full((16,), model.config.get_text_config().qk_scale_factor)
+        ones = torch.ones(16)
+        stored = [
+            ("blk.0.attn_q_norm.weight", scale),
+            ("blk.0.attn_k_norm.weight", ones),
+            ("blk.0.attn_q.weight", torch.zeros(64, 64)),
+        ]
+
+        checked = list(check_config_tensors(model, "muse-glimmer", stored))
+
+        self.assertEqual([n for n, _ in checked], [n for n, _ in stored])
+        for name, value in (
+            ("blk.1.attn_q_norm.weight", ones),
+            ("blk.1.attn_k_norm.weight", scale),
+        ):
+            with self.assertRaisesRegex(ValueError, name):
+                list(check_config_tensors(model, "muse-glimmer", [(name, value)]))
+
     def test_keeps_qwen_weights_as_stored(self) -> None:
         _, convert = gguf_to_hf(_tiny(Qwen2Config), "qwen2")
         weight = torch.randn(64, 64)