llama code break with torch compile
Author: samsja
Pull request: no
Labels: name: bug

name: Compilation

### System Info


I am having some graph break when using transformers llama
```bash
[rank0]:W0301 00:23:46.119000 785666 torch/_dynamo/convert_frame.py:844] [9/8]    function: 'forward' (/home/ubuntu/prime-rl/.venv/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:318)
[rank0]:W0301 00:23:46.119000 785666 torch/_dynamo/convert_frame.py:844] [9/8]    last reason: 9/0: L['self']._modules['self_attn'].layer_idx == 0
[rank0]:W0301 00:23:46.119000 785666 torch/_dynamo/convert_frame.py:844] [9/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
[rank0]:W0301 00:23:46.119000 785666 torch/_dynamo/convert_frame.py:844] [9/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/torch.compiler_troubleshooting.html.
```

### Who can help?

_No response_

### 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

use torch.compile with llama

### Expected behavior

Would expect no graph break or at least not there

Comments:
Could you share a small reproducer! 🤗 

(Doc might be missing but you need to pass `cache_implementatin="static"` when you initi the model to make sure you are using static cache

will do

> (Doc might be missing but you need to pass cache_implementatin="static" when you initi the model to make sure you are using static cache

hmm not sure to understand where should I use this feature and what does it do ?


@samsja You can set it in generate_config like `model.generation_config.cache_implementation = "static"` or If you are using generate then pass it in generate function directly `generate(**inputs, cache_implementation = "static"). With static cache the cache max len is fixed so it should help with compilation unlike dynamic cache

> [@samsja](https://github.com/samsja) You can set it in generate_config like `model.generation_config.cache_implementation = "static"` or If you are using generate then pass it in generate function directly `generate(**inputs, cache_implementation = "static"). With static cache the cache max len is fixed so it should help with compilation unlike dynamic cache

this did not fixed my issue. I also cannot find any documentation regarding this features. Thanks for the help tho 😃 

@samsja It would be super helpful for me or the maintainers if you could be provide ur code or minimal reproducer. 

I will try to come up with a minimal example

the full traceback can also help us!

hey, I tried to come up with a minimal reproducible example

```python
import os
import torch

from transformers import (
    AutoTokenizer,
    Qwen2Config,
    Qwen2ForCausalLM,
)

import torch.nn.functional as F
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed import destroy_process_group, init_process_group

torch.set_float32_matmul_precision("high")

init_process_group(backend="nccl")

torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 
tokenizer = AutoTokenizer.from_pretrained(model_name)
config_model = Qwen2Config.from_pretrained(model_name, attn_implementation="sdpa")

model = Qwen2ForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name, config=config_model)

mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=None)

for layer_id, transformer_block in enumerate(model.model.layers):
    reshard_after_forward = layer_id < len(model.model.layers) - 1
    fully_shard(transformer_block, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward)
fully_shard(model, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward)



model = torch.compile(model)

BS = 1
SEQ_LEN = 64

print("start")
for i in range(10):
    
    
    input_ids = torch.randint(0, tokenizer.vocab_size, (BS, SEQ_LEN)).to("cuda")
    
    
    output = model(input_ids).logits
    

    loss = F.cross_entropy(input=output.view(-1, output.shape[-1]), target=input_ids.view(-1))
    
    loss.backward()
    
    print(f"{i} loss: {loss.item()}")
    
    
destroy_process_group()
```

to run with 

```bash
 torchrun --nproc_per_node=2 script.py
```

there is no issue when not using fsdp and torch compile. 

error

```bash
[rank1]:W0305 22:54:17.105000 3067357 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8] torch._dynamo hit config.cache_size_limit (8)
[rank1]:W0305 22:54:17.105000 3067357 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8]    function: 'forward' (/root/prime-rl/.venv/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py:242)
[rank1]:W0305 22:54:17.105000 3067357 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8]    last reason: 9/0: L['self']._modules['self_attn'].layer_idx == 0              
[rank1]:W0305 22:54:17.105000 3067357 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
[rank1]:W0305 22:54:17.105000 3067357 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/torch.compiler_troubleshooting.html.
[rank0]:W0305 22:54:17.226000 3067356 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8] torch._dynamo hit config.cache_size_limit (8)
[rank0]:W0305 22:54:17.226000 3067356 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8]    function: 'forward' (/root/prime-rl/.venv/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py:242)
[rank0]:W0305 22:54:17.226000 3067356 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8]    last reason: 9/0: L['self']._modules['self_attn'].layer_idx == 0              
[rank0]:W0305 22:54:17.226000 3067356 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
[rank0]:W0305 22:54:17.226000 3067356 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:844] [9/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/torch.compiler_troubleshooting.html.
```


torch version: 2.5.1
transformer version: 4.48.3


Thanks in advance for the help 🙏 


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.

cc @vasqu if you have time to have a look 👀 

Sry don't have the bandwith for multi-gpu 😢

@samsja Could you run your script with `TORCH_LOGS="recompiles"`, e.g. `TORCH_LOGS="recompiles" torchrun ...`? This should give more details on what causes recompilation. Otherwise, for fsdp/ddp, you should be able to compile the transformer blocks each first and then after compiling all blocks wrap with fsdp (just an idea not completely sure about that).


> Sry don't have the bandwith for multi-gpu 😢
> 
> [@samsja](https://github.com/samsja) Could you run your script with `TORCH_LOGS="recompiles"`, e.g. `TORCH_LOGS="recompiles" torchrun ...`? This should give more details on what causes recompilation. Otherwise, for fsdp/ddp, you should be able to compile the transformer blocks each first and then after compiling all blocks wrap with fsdp (just an idea not completely sure about that).

I will try to do what you suggest. That being said I think we should see the same issue one only one GPU with FSDP

I think the cache as pointed out by Arthur might be the issue, especially due to the modifications done on the cache (which compile definitely doesn't like). I've modified your script and adjusted it to be done on one single GPU:
```python
import os
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"

import torch
torch.set_float32_matmul_precision("high")

from transformers import (
    AutoTokenizer,
    Qwen2Config,
    Qwen2ForCausalLM,
)

import torch.nn.functional as F
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy


model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

tokenizer = AutoTokenizer.from_pretrained(model_name)
config_model = Qwen2Config.from_pretrained(model_name, attn_implementation="sdpa")

model = Qwen2ForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name, config=config_model)

mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=None)
for layer_id, transformer_block in enumerate(model.model.layers):
    reshard_after_forward = layer_id < len(model.model.layers) - 1
    fully_shard(transformer_block, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward)
fully_shard(model, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward)

model = torch.compile(model)

BS = 1
SEQ_LEN = 64

print("start")
for i in range(10):
    input_ids = torch.randint(0, tokenizer.vocab_size, (BS, SEQ_LEN)).to("cuda")
    # this here is the importatn part! You don't even need a cache in training so why not disable it altogether :)
    output = model(input_ids, use_cache=False).logits

    loss = F.cross_entropy(input=output.view(-1, output.shape[-1]), target=input_ids.view(-1))
    loss.backward()

    print(f"{i} loss: {loss.item()}")
```

Let me know if you still encounter the compile issue with this modification. 

Edit: Just fyi, without the `use_cache=False` I also encountered your initial issue (too many recompilations).

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.