Skipping cudagraphs for unknown reason
Author: hychiang-git
Pull request: no
Labels: name: Good Second Issue

name: Feature request

name: Compilation

name: Cache

### System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.
                                      
- `transformers` version: 4.41.2
- Platform: Linux-5.15.0-112-generic-x86_64-with-glibc2.35
- Python version: 3.10.13                                                    
- Huggingface_hub version: 0.23.4                                            
- Safetensors version: 0.4.2                                                 
- Accelerate version: 0.28.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2 (True) 
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker

### Information

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

### Tasks

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

### Reproduction

I read [issue 30055](https://github.com/huggingface/transformers/issues/30055) and [issue 30351](https://github.com/huggingface/transformers/issues/30351), and Llama works well with  `cache_implementation="static"`. However, I am trying to use `torch.compile` for other models such as `pythia` and `phi-2` where the  `cache_implementation="static"` is not appliable, and it will produce errors like:
```
[2024-06-26 18:18:57,065] [0/0] torch._inductor.fx_passes.split_cat: [WARNING] example value absent for node: cat_3                                        
[2024-06-26 18:18:57,065] [0/0] torch._inductor.fx_passes.split_cat: [WARNING] example value absent for node: cat_2                                        
[2024-06-26 18:18:57,065] [0/0] torch._inductor.fx_passes.split_cat: [WARNING] example value absent for node: cat_1
[2024-06-26 18:18:57,065] [0/0] torch._inductor.fx_passes.split_cat: [WARNING] example value absent for node: cat
skipping cudagraphs for unknown reason
...
...
  File "/tmp/torchinductor_hc29225/bi/cbiig6bqpidhtncuswvfxwqqjwoiiswlwlrnh7eobbwm4wjlvpts.py", line 15465, in call
    extern_kernels.addmm(arg4_1, reinterpret_tensor(buf3, (16, 2560), (2560, 1), 0), reinterpret_tensor(arg3_1, (2560, 2560), (1, 2560), 0), alpha=1, beta=
1, out=buf4)
RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
```


Here is my code for reproducing the errors.
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Llama works well with  cache_implementation="static", but other types of models do not have the configuration.
# model     = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", torch_dtype=torch.float16, attn_implementation="sdpa", token=access_token).cuda().eval()
# tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf", token=access_token)

# model     = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-2.8b" , torch_dtype=torch.float16, trust_remote_code=True).cuda().eval()
# tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-2.8b", trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype=torch.float16, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)

model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

inputs = tokenizer(["
 
 Write an essay about large language models 
"], return_tensors="pt").to(model.device)

max_new_tokens = 64
fn = lambda: model.generate(
      **inputs,
      do_sample=False,
      # cache_implementation="static", # this only works for Llama-2
      max_new_tokens=max_new_tokens,
      pad_token_id=tokenizer.pad_token_id,
      temperature=None,
      top_p=None
)
fn() 
```

### Expected behavior

The models such as `pythia` and `phi-2` can run with `torch.compile` and a clear latency improvement can be observed. 

Comments:
cc @gante too :)

Hey! These models are not supported yet! We could actually open this to the community and get all decoder models to support compile soon-ish. 
Do you want to open a PR? 

Hi @ArthurZucker, I would love to contribute. However, I am not sure what you mean by `get all decoder models to support compile`. Aren't `nn.Module`s compile friendly out of the box?

#31421 will be a good stepping stone! 

It's adding the cache positions api for all models

I don't think this issue is relevant any more, at least for phi-2 model. I was able to successfully run with static cache. I also see the relevant changes here for this model: https://github.com/zucchini-nlp/transformers/blob/08f52e2178a0bada437da02ed7c1395ae54b3309/src/transformers/models/phi/modular_phi.py#L128

@ArthurZucker perhaps it'd be better to list some models who still needs cache support somehow 😊 

Yep I think we can at least close this one! 🤗