[Native] Export Muse Glimmer through the native backend

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()

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"],