Symbolic trance with past_key_values input is not supported yet for the qwen2.
Author: My-captain
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.46.2
- Platform: Linux-6.1.0-17-amd64-x86_64-with-glibc2.35
- Python version: 3.10.0
- Huggingface_hub version: 0.29.3
- Safetensors version: 0.5.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.1.2+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA RTX A6000

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python3
transformers.utils.fx.symbolic_trace(Qwen2ForCausalLM, input_names=["attention_mask", "position_ids", "input_embeds", "past_key_values"])
```

### Expected behavior

1. I am using Qualcomm AIMET to quantize Qwen2.5 for deployment on the hexagon NPU.
2. AIMET relies on `fx.symbolic_trace` to convert the model into a static computation graph. 
3. However, when I attempted to use `transformers.utils.fx.symbolic_trace(Qwen2ForCausalLM, input_names=["attention_mask", "position_ids", "input_embeds", "past_key_values"])` for graph tracing, it failed.

Hope for your response.

Comments:
some related info:
https://github.com/huggingface/transformers/pull/30259

https://github.com/huggingface/transformers/blob/70b07d97cf2c5f61fff55700b65528a1b6845cd2/src/transformers/utils/fx.py#L1052-L1068

Hi @My-captain, as the error suggests, we don't support FX tracing with `past_key_values` input for `Qwen2` right now! We'll need someone to make a PR to add support before your code will work.

Sure, I still have a few questions I'd like to ask:
1. After removing the `past_key_values` parameter, does the resulting `fx.GraphModule` support `kv_cache`?
2. What efforts are required to make Qwen2 support `symbolic_trace` in a way that uses `past_key_values`?

any response is helpful.

Hi, @michaelbenayoun .

1. I'm trying to use `transformers.utils.fx.symbolic_trace` to symbolically trace `Qwen2ForCausalLM`.  
2. Qwen2 uses `DynamicCache` as the type for the `past_key_values` parameter, and during the forward process, it checks whether it is `None`. If it is, the cache is created internally.  
3. I traced the source code of `transformers.utils.fx` and found that although `transformers.cache_utils.Dynamic` is dynamically replaced with `ProxyableDynamicCache`, when `create_proxy` is called, the `past_key_values` parameter is created as an `HFProxy` instead of a `ProxyableDynamicCache`.  
4. After analysis, I discovered that the `proxy_factory_fn` in `transformers.utils.fx.HFTracer.create_proxy` does not correctly point to the `proxy_factory_fn` of `ProxyableDynamicCache`.  

My questions for you are:  
1. What was your specific design approach when creating the `ProxyableDynamicCache` class? During the `symbolic_trace` process, at what point is the execution flow intercepted and directed to `ProxyableDynamicCache`?  

>  [https://github.com/huggingface/transformers/commit/fbabd6746f7f129adc66f0290bb45ecbfe94236c#diff-06392bad3b9e97be9ade60d4ac46f73b6809388f4d507c2ba1384ab872711c51](url)

2. Could you provide a code example that correctly creates an instance of the `ProxyableDynamicCache` class for reference?
3. Looking forward to your reply.



This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.