From ebd18efdfe31472f24b4fb95dcf614577f6ee92c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 14:11:08 -0800 Subject: [PATCH] Map Muse Glimmer GGUF tensors onto its Hugging Face model gguf_hf takes tensor names from gguf-py's tensor map, which has no entry for Muse Glimmer's muse-glimmer architecture, so add a table from its GGUF names to the Hugging Face MuseGlimmerForConditionalGeneration parameters under model.language_model. The all-ones attn_q_norm/attn_k_norm tensors are skipped: the HF QK norm has no weight. The conversion undoes llama.cpp's Q/K row permutation, as for Llama, and the 1 it adds to the four centered decoder norms; the final norm is a plain gain and stays as stored. Checked on the first four layers of Muse-Glimmer-30B-KQuant-17GB-Q4_K_M.gguf against Muse Glimmer's own model and GGUF loader: FP32 logits agree to 5e-5. The tests need a transformers with Muse Glimmer and skip otherwise. --- .../native/examples/hf_llm/export_hf_llm.py | 10 +- backends/native/examples/hf_llm/gguf_hf.py | 120 ++++++++++++++++-- .../examples/hf_llm/tests/test_gguf_hf.py | 104 +++++++++++++++ 3 files changed, 218 insertions(+), 16 deletions(-) diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index bda253bfec3..a1776fb5202 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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}") diff --git a/backends/native/examples/hf_llm/gguf_hf.py b/backends/native/examples/hf_llm/gguf_hf.py index e4bda9d72c7..a94d7c8449d 100644 --- a/backends/native/examples/hf_llm/gguf_hf.py +++ b/backends/native/examples/hf_llm/gguf_hf.py @@ -9,7 +9,8 @@ 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. """ @@ -17,6 +18,7 @@ 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: diff --git a/backends/native/examples/hf_llm/tests/test_gguf_hf.py b/backends/native/examples/hf_llm/tests/test_gguf_hf.py index b836020b782..fa62c67c762 100644 --- a/backends/native/examples/hf_llm/tests/test_gguf_hf.py +++ b/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)