diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index 260d0f3eda4..92879323cda 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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( 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 d796a35ac66..35e8ef6e715 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/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.