Repository: pytorch
Issue: torch.export support for the latest transformers `DynamicCache` as input
Labels: triaged,oncall: pt2,export-triaged,oncall: export

Hugging Face `transformers` is moving to use the `DynamicCache` class as part of the model inputs for the kv cache values. Currently `torch.export` will complain that it is not a tensor. So all models that take `DynamicCache` as input will not be exportable. This affects torch.onnx and other exporters dependent on torch.export alike.

```
transformers==4.44.2
```

```python
from typing import List

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, DynamicCache


# Get position_ids from attention_mask
def get_position_ids(attention_mask: torch.Tensor, use_past_kv: bool):
    position_ids = attention_mask.long().cumsum(-1) - 1
    position_ids.masked_fill_(attention_mask == 0, 1)
    if use_past_kv:
        # Shape: (batch_size, 1)
        position_ids = position_ids[:, -1].unsqueeze(-1)

    # Shape: (batch_size, sequence_length)
    return position_ids


# Create empty past_key_values
def get_past_kv_inputs(config: AutoConfig, batch_size: int):
    num_heads = config.num_key_value_heads
    head_size = (
        config.head_dim
        if hasattr(config, "head_dim")
        else config.hidden_size // config.num_attention_heads
    )
    past_kv = [
        (
            torch.rand(batch_size, num_heads, 0, head_size, dtype=torch.float32),
            torch.rand(batch_size, num_heads, 0, head_size, dtype=torch.float32),
        )
        for _ in range(config.num_hidden_layers)
    ]
    return past_kv


def get_merged_model_dynamic_axes(input_names: List[str], output_names: List[str]):
    dynamic_axes = {}
    for name in input_names + output_names:
        if name in {"input_ids", "position_ids"}:
            # shape is (batch_size, sequence_length)
            dynamic_axes[name] = {0: "batch_size", 1: "sequence_length"}
        elif name == "attention_mask":
            # shape is (batch_size, past_sequence_length + sequence_length) = (batch_size, total_sequence_length)
            # for prompt generation, past_sequence_length = 0
            # for token generation, sequence_length = 1
            dynamic_axes[name] = {0: "batch_size", 1: "total_sequence_length"}
        elif "past" in name:
            # shape is (batch_size, num_heads, past_sequence_length, head_size)
            dynamic_axes[name] = {0: "batch_size", 2: "past_sequence_length"}
        elif name == "logits":
            # shape is (batch_size, sequence_length, vocab_size)
            dynamic_axes[name] = {0: "batch_size", 1: "sequence_length"}
        elif "present" in name:
            # shape is (batch_size, num_heads, past_sequence_length + sequence_length, head_size) = (batch_size, num_heads, total_sequence_length, head_size)
            # for prompt generation, past_sequence_length = 0
            # for token generation, sequence_length = 1
            dynamic_axes[name] = {0: "batch_size", 2: "total_sequence_length"}
        else:
            raise ValueError("Unknown input or output name found")
    return dynamic_axes


model_name = "google/gemma-2-2b"
cache_dir = "./cache_dir"

config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)

batch_size, sequence_length = 2, 8
inputs = {
    "input_ids": torch.randint(
        low=1, high=config.vocab_size, size=(batch_size, sequence_length), dtype=torch.int64
    ),
    "attention_mask": torch.ones((batch_size, sequence_length), dtype=torch.int64),
    "position_ids": get_position_ids(
        torch.ones((batch_size, sequence_length), dtype=torch.int64), use_past_kv=False
    ),
    "past_key_values": DynamicCache.from_legacy_cache(get_past_kv_inputs(config, batch_size)),
}

program = torch.export.export(
    model,
    args=tuple(inputs.values()),
    strict=False
)
```

Will raise an error

```pytb
Traceback (most recent call last):
  File "/workspace/ONNXConverter/optimum_export.py", line 83, in <module>
    program = torch.export.export(
              ^^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/__init__.py", line 366, in export
    return _export(
           ^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 1017, in wrapper
    raise e
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 990, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/exported_program.py", line 114, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 1880, in _export
    export_artifact = export_func(  # type: ignore[operator]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 1643, in _non_strict_export
    ) = make_fake_inputs(
        ^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/_export/non_strict_utils.py", line 193, in make_fake_inputs
    fake_args, fake_kwargs = tree_map_with_path(
                             ^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/utils/_pytree.py", line 1608, in tree_map_with_path
    return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/utils/_pytree.py", line 803, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/utils/_pytree.py", line 1608, in <genexpr>
    return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves))
                              ^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/_export/non_strict_utils.py", line 194, in <lambda>
    lambda kp, val: fakify(fake_mode, kp, val, t_constraints, sources),
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/_export/non_strict_utils.py", line 95, in fakify
    raise ValueError(f"Unsupported input type {type(t)}")
ValueError: Unsupported input type <class 'transformers.cache_utils.DynamicCache'>
```


cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @kunal-vaishnavi @xadupre @gramalingam