🐛 Describe the bug

Huggingface model: https://huggingface.co/CohereForAI/aya-23-8B
Reference code: https://huggingface.co/CohereForAI/aya-23-8B/blob/main/Aya_23_notebook.ipynb

Graph breaks

There are two types of graph breaks.

1. UserWarning: Graph break due to unsupported builtin flash_attn_2_cuda.PyCapsule.fwd.
2. Skipping frame because there is a graph break in a for/while loop

To Reproduce

https://github.com/FindHao/ml_scripts/blob/main/inductor/cases/aya/aya_full.py

You may need to add the following code and your own token for the first run to download weights etc.

from huggingface_hub import login
login(token='your_hf_token')

Versions

nightly

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire