[Native] Run the HF LLM export on newer transformers

backends/native/examples/hf_llm/export_hf_llm.py

@@ -20,6 +20,7 @@
 
 import argparse
 import contextlib
+import inspect
 import logging
 from collections.abc import Iterator
 
@@ -56,7 +57,6 @@
     PreTrainedModel,
 )
 from transformers.cache_utils import StaticLayer
-from transformers.integrations.accelerate import init_empty_weights
 from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM
 from transformers.integrations.sdpa_attention import sdpa_attention_forward
 from transformers.masking_utils import (
@@ -68,6 +68,28 @@
 _OFFGRAPH_ATTENTION = "offgraph"
 
 
+@contextlib.contextmanager
+def _parameters_on_meta() -> Iterator[None]:
+    """Create parameters on the meta device and buffers as usual, so computed
+    buffers such as RoPE frequencies keep their values. Weights are assigned
+    later. What accelerate's ``init_empty_weights`` did, which transformers no
+    longer re-exports."""
+    register = nn.Module.register_parameter
+
+    def register_on_meta(module, name, param):
+        register(module, name, param)
+        if param is not None:
+            module._parameters[name] = nn.Parameter(
+                param.to("meta"), requires_grad=param.requires_grad
+            )
+
+    nn.Module.register_parameter = register_on_meta
+    try:
+        yield
+    finally:
+        nn.Module.register_parameter = register
+
+
 def _load(
     path: str,
     layers: int,
@@ -84,9 +106,9 @@ def _load(
     if offgraph_cache:
         register_offgraph_attention(_OFFGRAPH_ATTENTION)
     attn = _OFFGRAPH_ATTENTION if offgraph_cache else "sdpa"
-    kwargs = {"torch_dtype": torch.float32, "attn_implementation": attn}
+    kwargs = {"dtype": torch.float32, "attn_implementation": attn}
     if empty_weights:
-        with init_empty_weights():
+        with _parameters_on_meta():
             model = AutoModelForCausalLM.from_config(config, **kwargs)
     elif random_weights:
         model = AutoModelForCausalLM.from_config(config, **kwargs)
@@ -183,7 +205,11 @@ def _store_kv_as(self, model: PreTrainedModel, dtype: torch.dtype) -> None:
             if type(old) is not StaticLayer:
                 raise ValueError(f"cannot change the KV dtype of {type(old).__name__}")
             layer = _CastingStaticLayer(old.max_cache_len)
-            layer.lazy_initialization(old.keys[:, :, :1].to(dtype))
+            states = [old.keys[:, :, :1].to(dtype), old.values[:, :, :1].to(dtype)]
+            # transformers 5.4+ also sizes the values from value_states.
+            if len(inspect.signature(layer.lazy_initialization).parameters) == 1:
+                states = states[:1]
+            layer.lazy_initialization(*states)
             self.cache.layers[i] = layer
             self.cached.register_buffer(f"key_cache_{i}", layer.keys, persistent=False)
             self.cached.register_buffer(

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -43,6 +43,7 @@
     Qwen2Config,
     Qwen3Config,
 )
+from transformers.cache_utils import StaticLayer
 
 # Architectures gguf_hf maps; the GGUF test covers each.
 _CONFIGS = {
@@ -136,6 +137,9 @@ def test_gguf_weights_reproduce_the_model(self) -> None:
                     weight.uniform_(-0.5, 0.5)
                 torch.manual_seed(2)
                 loaded = _tiny(arch, tie)
+                # Newer transformers' sdpa rejects Gemma 2's attention softcap.
+                for model in (reference, loaded):
+                    model.config._attn_implementation = "eager"
                 pairs = _as_gguf_tensors(reference, arch)
 
                 load_gguf(loaded, "model.gguf", arch, lambda _, p=pairs: iter(p))
@@ -373,8 +377,14 @@ def test_golden_samples_tolerance_follows_weight_quantization(self) -> None:
             self.assertEqual(samples.tol, tol)
 
     def test_fp16_kv_cache_rejects_sliding_window_layers(self) -> None:
+        model = _tiny("gemma3")
+        if all(
+            type(layer) is StaticLayer
+            for layer in LastTokenLM(copy.deepcopy(model), 32).cache.layers
+        ):
+            self.skipTest("this transformers windows static layers with the mask")
         with self.assertRaisesRegex(ValueError, "cannot change the KV dtype"):
-            LastTokenLM(_tiny("gemma3"), 32, torch.float16)
+            LastTokenLM(model, 32, torch.float16)
 
     def test_offgraph_cache_windows_sliding_layers_like_hf(self) -> None:
         # Mistral windows every layer; Gemma3 alternates sliding and full.