Export Huggingface models with StaticCache

[yuanyao-nv]

🐛 Describe the bug

I'm trying to export a llama2 model with StaticCache https://huggingface.co/docs/transformers/v4.52.3/en/internal/generation_utils#transformers.StaticCache

from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache


model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

inputs = tokenizer(text="My name is Llama", return_tensors="pt")

# Prepare a cache class and pass it to model's forward
# Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate
max_generated_length = inputs.input_ids.shape[1] + 10
past_key_values = StaticCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype)
outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)  # model can run with the given inputs

import torch
ep = torch.export.export(model, (inputs['input_ids'], inputs['attention_mask']), {'past_key_values': past_key_values, 'use_cache': True})  # export fails

I hit this error

Traceback (most recent call last):
  File "/workspace/GenAI/static_cache.py", line 19, in <module>
    ep = torch.export.export(model, (inputs['input_ids'], inputs['attention_mask']), {'past_key_values': past_key_values, 'use_cache': True})
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/__init__.py", line 319, in export
    raise e
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/__init__.py", line 286, in export
    return _export(
           ^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1159, in wrapper
    raise e
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1125, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 2172, in _export
    ep = _export_for_training(
         ^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1159, in wrapper
    raise e
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1125, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 2033, in _export_for_training
    export_artifact = export_func(
                      ^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1933, in _non_strict_export
    ) = make_fake_inputs(
        ^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/_export/non_strict_utils.py", line 347, in make_fake_inputs
    fake_args, fake_kwargs = tree_map_with_path(
                             ^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/utils/_pytree.py", line 2077, in tree_map_with_path
    return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/utils/_pytree.py", line 1197, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/utils/_pytree.py", line 2077, in <genexpr>
    return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves))
                              ^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/_export/non_strict_utils.py", line 348, in <lambda>
    lambda kp, val: fakify(fake_mode, kp, val, t_constraints, sources),
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/_export/non_strict_utils.py", line 124, in fakify
    raise ValueError(
ValueError: Unsupported input type <class 'transformers.cache_utils.StaticCache'>. Export only supports pytree containers of basic types (Tensor, int, float, ...) as input. To register a custom dataclass, use torch.export.register_dataclass. To register a custom container type, use torch.utils._pytree.register_pytree_node. To register a constant input, use torch.utils._pytree.register_constant

Is this usage of StaticCache correct?

I also tried playing with https://huggingface.co/docs/transformers/en/main_classes/executorch#transformers.convert_and_export_with_cache but could not set it up correctly.

Versions

PyTorch version: 2.8.0.dev20250612+cu126
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.1 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version: 18.1.3 (1ubuntu1)
CMake version: version 3.27.0
Libc version: glibc-2.39

Python version: 3.12.3 (main, Jan 17 2025, 18:03:48) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.8.0-52-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 12.8.93
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A5000
Nvidia driver version: 560.28.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

[yuanyao-nv]

cc @xadupre

[zhxchen17]

@yuanyao-nv Hi, as the error message suggests, StaticCache shall be registered as a pytree type from transformers side. Feel free to file an issue to huggingface as torch.export require input types to be registered as pytree type before using.

[yuanyao-nv]

Thanks for the reply.
I figured out how to export using the https://huggingface.co/docs/transformers/en/main_classes/executorch#transformers.convert_and_export_with_cache API.

from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache, GenerationConfig
import transformers
import torch

device = "cpu"
dtype = "bfloat16"
cache_implementation = "static"
attn_implementation = "sdpa"
batch_size = 1
max_cache_len = 1234
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    device_map=device,
    torch_dtype=dtype,
    attn_implementation=attn_implementation,
    generation_config=GenerationConfig(
        use_cache=True,
        cache_implementation=cache_implementation,
        max_length=max_cache_len,
        cache_config={
            "batch_size": batch_size,
            "max_cache_len": max_cache_len,
            "device": device,
        },
    ),
)

exported_program = transformers.convert_and_export_with_cache(model)

inputs = []
for buffer_name, buffer in exported_program.named_buffers():
    print(buffer_name, buffer.shape, buffer.dtype)
    if buffer.dtype == torch.bool:
        inputs.append(torch.rand(*buffer.shape) > 0.5)
    else:
        inputs.append(torch.rand(*buffer.shape, dtype=buffer.dtype))


torch.onnx.export(exported_program, inputs, 'model.onnx')

I can see the named buffer list printed out as follows:

mask torch.Size([1234, 1234]) torch.bool
key_cache_0 torch.Size([1, 1, 1234, 256]) torch.bfloat16
model.model.rotary_emb.inv_freq torch.Size([128]) torch.float32
value_cache_0 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_1 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_1 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_2 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_2 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_3 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_3 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_4 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_4 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_5 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_5 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_6 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_6 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_7 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_7 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_8 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_8 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_9 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_9 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_10 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_10 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_11 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_11 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_12 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_12 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_13 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_13 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_14 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_14 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_15 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_15 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_16 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_16 torch.Size([1, 1, 1234, 256]) torch.bfloat16
key_cache_17 torch.Size([1, 1, 1234, 256]) torch.bfloat16
value_cache_17 torch.Size([1, 1, 1234, 256]) torch.bfloat16

However, probably due to the fact that convert_and_export_with_cache registers the cache tensors as buffers, in the exported ONNX model the cache tensors are not showing up as inputs but as model constants with max_cache_length.

cc: @xadupre @gramalingam @justinchuby @titaiwangms

[justinchuby]

That's unfortunate. Can we change transformers so that they are not registered as buffers?

[yuanyao-nv]

That's unfortunate. Can we change transformers so that they are not registered as buffers?

If I comment out
https://github.com/huggingface/transformers/blob/d5d007a1a0f0c11a726a54c8f00bd71825f84d02/src/transformers/integrations/executorch.py#L284-L286
then I get

[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ❌
Traceback (most recent call last):
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_core.py", line 1405, in export
    decomposed_program = _prepare_exported_program_for_export(
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_core.py", line 969, in _prepare_exported_program_for_export
    exported_program = _fx_passes.decompose_with_registry(exported_program, registry)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_fx_passes.py", line 19, in decompose_with_registry
    return exported_program.run_decompositions(decomp_table)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 1386, in run_decompositions
    return _decompose_exported_program(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 854, in _decompose_exported_program
    ) = _decompose_and_get_gm_with_new_signature_constants(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 496, in _decompose_and_get_gm_with_new_signature_constants
    _override_graph_signature_for_temp_registered_constants(
  File "/home/yuanyao/.local/lib/python3.12/site-packages/torch/_export/utils.py", line 133, in _override_graph_signature_for_temp_registered_constants
    raise RuntimeError(
RuntimeError: Constant self___static_cache_key_cache_0 is mutated in the forward method. Pls register it as buffer

I guess the StaticCache class's update() method also needs to be updated to work with export?

[justinchuby]

Would it be a good idea to keep buffers as inputs then? @titaiwangms @tugsbayasgalan What are some other usages of buffers?

[justinchuby]

Maybe only for non-persistent buffers?

[xadupre]

I looked at the code of convert_and_export_with_cache. The model is wrapped and buffers are added for the cache. The right way to should consists in extending the serialization implemented in https://github.com/huggingface/transformers/blob/main/src/transformers/cache_utils.py#L654 to support StaticCache.

[titaiwangms]

That's unfortunate. Can we change transformers so that they are not registered as buffers?

Do we still get a valid ONNX if we convert from registered buffers included exported program?

[justinchuby]

I hope so?