torch.compile graph break when tuning llama with FA2
Author: SilverSoldier
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.50.3
- Platform: Linux-5.14.0-284.73.1.el9_2.x86_64-x86_64-with-glibc2.31
- Python version: 3.12.9
- Huggingface_hub version: 0.29.3
- Safetensors version: 0.5.3
- Accelerate version: 1.0.1
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (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?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

I am getting 2 graph breaks when tuning llama3-8b model with torch.compile. Both are failing due to Dynamic Control Flow issues with the following error message:
```
Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow.
```

The thing is, both of these lines were actually last touched by PRs to remove graph breaks, so clearly something changed in pytorch compile that is breaking them again.

1. `if attention_mask is not None and (attention_mask == 0.0).any()` in [modeling_llama](https://github.com/huggingface/transformers/blob/800510c67bfc5cedd0bb7635648a07f39719be43/src/transformers/models/llama/modeling_llama.py#L637)
    This was fixed just 4 months ago in [this PR](https://github.com/huggingface/transformers/pull/35187)
3.  `max_length_q is not None or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all())` [in modeling_flash_attention_utils](https://github.com/huggingface/transformers/blame/800510c67bfc5cedd0bb7635648a07f39719be43/src/transformers/modeling_flash_attention_utils.py#L378)
This was fixed as part of [this PR](https://github.com/huggingface/transformers/pull/33932) 6 months ago.

The error message suggests to use `torch.cond` but it is not trivial to use for both these cases. I attempted to fix #1, but the return values of the if and else branch are None and attention_mask respectively while `cond` expects the outputs of the 2 branches to be of same type and shape.
For #2 it is the `torch.diff` clause which is causing the problem. I haven't looked closely enough but the multiple condition checks might make clean separation into `cond` syntax slightly difficult.

### Information

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

### Tasks

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

### Reproduction

```
accelerate launch --num_processes=1  -m tuning.sft_trainer --output_dir ./train_output --max_steps 5 --learning_rate 2e-5 --training_data_path=/data/data/ei_5.jsonl --save_steps=50 --torch_dtype bfloat16 --logging_strategy steps --logging_steps 1 --per_device_train_batch_size 8 --max_seq_length 1024 --include_tokens_per_second true --data_formatter_template "### Input: {input} \n\n### Response: {output}" --response_template "\n### Response:" --torch_compile True --model_name_or_path /data/models/llama3-8b --use_flash_attn True --packing
```

Checked with pytorch 2.5 and 2.6, with and without padding_free with no change in graph breaks

### Expected behavior

Expected no graph breaks.

Curious if anyone knows why this is popping up now and how to fix it, is it only through `cond` or is there some other way (I am willing to fix but not sure how).
Thanks!

Comments:
The original PR authors were @winglian and @Abhishek-TAMU - do you have any idea why the graph breaks would have started to reappear?

For 1. A quick fix is to simply remove `(attention_mask == 0.0).any()`, which avoids the initial issue, but a graph break still occurs in `unpad_input`.

For 2. #37332 addresses the `batch_size > 1` case, but doesn't work for `batch_size == 1`.

Ideally, when using FA2, the input should be packed and the corresponding `cu_seq_lens` and `max_seq_lens` should be prepared during the data collator stage.

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.