From 90930784a87617aea9791b3015f62a0956543140 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 14:19:19 -0800 Subject: [PATCH] Export Muse Glimmer through the native backend export_hf_llm loads models with AutoModelForCausalLM, which has no Muse Glimmer: Transformers builds it only as MuseGlimmerForConditionalGeneration, with a vision tower and a scan of input_ids for image and video tokens that the ggml engine cannot run. Build it with AutoModelForImageTextToText and export MuseGlimmerTextLM instead: its language model and LM head under their Hugging Face parameter names, with the output multiplier and logit soft cap. --layers also truncates its per-layer layer_rope_theta. It exports only with --offgraph_cache, which carries its sliding windows. With --gguf, a GGUF checkpoint loads through the muse-glimmer mapping. The first four layers of Muse-Glimmer-30B-KQuant-17GB-Q4_K_M.gguf export with every weight packed, and their golden samples pass on the ggml engine on CPU and Metal. The test needs a transformers with Muse Glimmer and skips otherwise. --- .../native/examples/hf_llm/export_hf_llm.py | 63 ++++++++++++++-- .../examples/hf_llm/test_export_hf_llm.py | 74 +++++++++++++++++++ 2 files changed, 129 insertions(+), 8 deletions(-) diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index a1776fb5202..7335a1951cd 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/backends/native/examples/hf_llm/export_hf_llm.py @@ -55,6 +55,7 @@ AttentionInterface, AutoConfig, AutoModelForCausalLM, + AutoModelForImageTextToText, AutoTokenizer, PreTrainedModel, ) @@ -65,6 +66,7 @@ ALL_MASK_ATTENTION_FUNCTIONS, AttentionMaskInterface, ) +from transformers.modeling_outputs import CausalLMOutput _OFFGRAPH_ATTENTION = "offgraph" @@ -92,37 +94,82 @@ def register_on_meta(module, name, param): nn.Module.register_parameter = register +class MuseGlimmerTextLM(nn.Module): + """Muse Glimmer's language model and LM head, under their Hugging Face + parameter names, without the vision tower. Text-only, so it skips the + multimodal model's scan of input_ids for image and video tokens.""" + + def __init__(self, model: PreTrainedModel) -> None: + super().__init__() + self.config = model.config + self.generation_config = model.generation_config + self.model = nn.ModuleDict({"language_model": model.model.language_model}) + self.lm_head = model.lm_head + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None = None, + logits_to_keep: int = 0, + **_, + ) -> CausalLMOutput: + text_config = self.config.text_config + hidden = self.model["language_model"]( + input_ids=input_ids, position_ids=position_ids, use_cache=False + ).last_hidden_state + if logits_to_keep: + hidden = hidden[:, -logits_to_keep:] + logits = self.lm_head(hidden) * text_config.output_multiplier + cap = text_config.final_logit_softcapping + return CausalLMOutput(logits=torch.tanh(logits / cap) * cap) + + +# Models Transformers builds only with their vision tower: the class that +# builds one, and the text-only module to export instead. +_TEXT_ONLY = {"muse_glimmer": (AutoModelForImageTextToText, MuseGlimmerTextLM)} +# Per-layer config lists, kept for the layers --layers keeps. +_PER_LAYER = ("layer_types", "layer_rope_theta") + + def _load( path: str, layers: int, random_weights: bool, empty_weights: bool, offgraph_cache: bool = False, -) -> PreTrainedModel: +) -> nn.Module: config = AutoConfig.from_pretrained(path) text_config = config.get_text_config() if layers: text_config.num_hidden_layers = layers - if getattr(text_config, "layer_types", None): - text_config.layer_types = text_config.layer_types[:layers] + for name in _PER_LAYER: + if getattr(text_config, name, None): + setattr(text_config, name, getattr(text_config, name)[:layers]) + auto_class, text_only = _TEXT_ONLY.get( + config.model_type, (AutoModelForCausalLM, None) + ) + if text_only and not offgraph_cache: + raise ValueError(f"{config.model_type} exports only with --offgraph_cache") if offgraph_cache: register_offgraph_attention(_OFFGRAPH_ATTENTION) attn = _OFFGRAPH_ATTENTION if offgraph_cache else "sdpa" kwargs = {"dtype": torch.float32, "attn_implementation": attn} if empty_weights: with _parameters_on_meta(): - model = AutoModelForCausalLM.from_config(config, **kwargs) + model = auto_class.from_config(config, **kwargs) elif random_weights: - model = AutoModelForCausalLM.from_config(config, **kwargs) + model = auto_class.from_config(config, **kwargs) else: - model = AutoModelForCausalLM.from_pretrained(path, config=config, **kwargs) + model = auto_class.from_pretrained(path, config=config, **kwargs) + if text_only: + model = text_only(model) model.generation_config.use_cache = True model.generation_config.cache_implementation = "static" return model.eval() def load_gguf( - model: PreTrainedModel, + model: nn.Module, path: str, architecture: str | None = None, raw_iter: RawIter | None = None, @@ -277,7 +324,7 @@ class OffGraphLastTokenLM(nn.Module): ``capacity`` positions the active one. """ - def __init__(self, model: PreTrainedModel, capacity: int) -> None: + def __init__(self, model: nn.Module, capacity: int) -> None: super().__init__() self.model = model config = model.config.get_text_config() diff --git a/backends/native/examples/hf_llm/test_export_hf_llm.py b/backends/native/examples/hf_llm/test_export_hf_llm.py index 9c20b97ec6c..0daeb0a9bff 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/backends/native/examples/hf_llm/test_export_hf_llm.py @@ -125,6 +125,41 @@ def _ops_except_casts(ep: torch.export.ExportedProgram) -> collections.Counter: ) +def _tiny_muse_glimmer_dir(path: str) -> PreTrainedModel | None: + """Saves a 4-layer Muse Glimmer config, window 3 on the first three layers, + to `path` and returns the model, or None where transformers has none.""" + try: + from transformers import MuseGlimmerConfig, MuseGlimmerForConditionalGeneration + except ImportError: + return None + config = MuseGlimmerConfig( + text_config={ + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "sliding_window": 3, + "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, + ) + config.save_pretrained(path) + model = MuseGlimmerForConditionalGeneration(config).eval().requires_grad_(False) + model.config._attn_implementation = "sdpa" + return model + + class ExportHfLlmTest(unittest.TestCase): def test_gguf_weights_reproduce_the_model(self) -> None: ids = torch.tensor([[3, 4, 5, 6, 7, 8]]) @@ -334,6 +369,45 @@ def test_offgraph_cache_matches_the_static_cache(self) -> None: self.assertNotIn("aten.index_copy_.default", ops) self.assertFalse(ep.graph_signature.buffers_to_mutate) + def test_muse_glimmer_exports_text_only_with_an_offgraph_cache(self) -> None: + with tempfile.TemporaryDirectory() as path: + torch.manual_seed(7) + reference = _tiny_muse_glimmer_dir(path) + if reference is None: + self.skipTest("transformers has no Muse Glimmer") + with self.assertRaisesRegex(ValueError, "only with --offgraph_cache"): + _load(path, 0, True, False) + model = _load(path, 0, True, False, offgraph_cache=True) + missing, unexpected = model.load_state_dict( + reference.state_dict(), strict=False + ) + self.assertFalse(missing) + self.assertTrue(all(".language_model." not in k for k in unexpected)) + offgraph = OffGraphLastTokenLM(model, 32).eval() + prompts = [[3, 4, 5, 6, 7, 8, 9], [9, 10]] + + samples = golden_samples(offgraph, prompts, 3) + + for sample in samples.samples: + sequence = [] + for call in sample.calls: + sequence += call.inputs[0][0].tolist() + with torch.no_grad(): + want = reference(torch.tensor([sequence])).logits[:, -1:] + torch.testing.assert_close(call.outputs[0], want, atol=1e-5, rtol=1e-5) + self.assertEqual(offgraph.cache_metadata()["get_windows"], [3, 3, 3, 0]) + seq = torch.export.Dim("seq", min=1, max=31) + with torch.no_grad(): + ep = torch.export.export( + offgraph, + (torch.tensor([[3, 4, 5]]), torch.arange(3)), + dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}}, + strict=False, + ) + ops = _ops_except_casts(ep) + self.assertEqual(ops["kvcache.update_and_attend.default"], 4) + self.assertNotIn("aten.eq.Scalar", ops) + def test_offgraph_cache_excludes_in_graph_cache_flags(self) -> None: for flags in ( ["--max_context_length", "64"],