tracker: `generate` compatibility with `torch.compile`
Author: gante
Pull request: no
Labels: name: WIP

# `generate` 🤜 🤛  `torch.compile`

Part of the [PyTorch 2024 H2 roadmap](https://drive.google.com/file/d/1Ucm17fyUeF0PWSd2g7jonM144XMx2t2r/view).

This issue is a tracker of the compatibility between `.generate` and `torch.compile` ([intro docs by pytorch](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html)). The goal is to enable `fullgraph=True` compilation on the main `generate` use cases.

⚠️ Is *your* `generate` use case not covered by this tracker? Check if it was requested below and upvote it if it was. Otherwise, add a comment. We will consider expanding the selection below on widely requested use cases 🤗 

### Decoding Strategies (end-to-end compilation)
- [x] `greedy_search` / `sample` are compatible (https://github.com/huggingface/transformers/pull/30788)
- [ ] `beam_search` / `beam_sample` are compatible, depends on the step above
- [ ] `assisted_decoding` (aka speculative decoding) is compatible, depends on the steps above

### Generate Flags and Options
- [ ] all `LogitsProcessor` classes were checked for compatibility (and the appropriate exceptions are raised when not compatible)
- [ ] all `StoppingCriteria` classes were checked for compatibility (and the appropriate exceptions are raised when not compatible)

### Models

Notes:
1. models tagged as "important models" in our CI + popular models
2. language models released starting from v4.42 should ALL support compile

Decoder-only:
- [x] GPT-J is compatible (#31421)
- [ ] GPT2 is compatible
- [x] Llama is compatible (#27931)
- [x] Gemma is compatible (#29167)
- [x] Llava is compatible (https://github.com/huggingface/transformers/pull/34502 removed dynamic control flow)
- [ ] Llava-Next is compatible (#29891)
- [x] Mistral is compatible (https://github.com/huggingface/transformers/pull/30642)
- [ ] Mixtral is compatible (https://github.com/huggingface/transformers/pull/30793/)
- [X] Phi is compatible (https://github.com/huggingface/transformers/pull/32617)
- [ ] Phi3 is compatible (https://github.com/huggingface/transformers/pull/30688)
- [X] BLOOM is compatible (note: this one might be tricky due to cache format) https://github.com/huggingface/transformers/pull/32617
- [x] Mamba is compatible (requested [here](https://github.com/huggingface/transformers/issues/29699#issuecomment-2072233223)) (PR: https://github.com/huggingface/transformers/pull/31247)
- [X] Persimmon is compatible (https://github.com/huggingface/transformers/pull/32617)
- [x] Qwen2 https://github.com/huggingface/transformers/pull/32617
- [x] Qwen2-VL https://github.com/huggingface/transformers/pull/32617
- [x] Falcon https://github.com/huggingface/transformers/pull/32617
- [x] GPTNeoX https://github.com/huggingface/transformers/pull/32617
- [x] Starcoder2 https://github.com/huggingface/transformers/pull/32617
- [x] StableLM https://github.com/huggingface/transformers/pull/32617

Encoder-decoder:
- [ ] BART is compatible (https://github.com/huggingface/transformers/pull/35314)
- [x] T5 is compatible (https://github.com/huggingface/transformers/pull/34089)
- [x] Whisper is compatible (https://github.com/huggingface/transformers/pull/31166)


### Quantization
- [ ] BNB support
- [ ] GPTQ support
- [ ] AWQ support

### Others
- [x] We have a benchmark script to quickly compare the impact of PRs
- [x] Add section to existing docs on the topic
- [x] Confirm that pipelines work after compiling generate


Comments:
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.

👀 

👀

👀

Hi @gante, I'm planning on looking at these issues from the torch.compile side - two questions: (1) is there someone from HF that's committed to working on this in the event I identify some model changes that would help run `torch.compile` more smoothly? and (2) are these the right models we should be focusing on? are there others that are also in need of torch.compile support? 

Let me know when you can, also happy to chat on slack (I'm on the HF slack already)

Hi, is there any updates? I am especially interested in the things related to https://github.com/huggingface/transformers/issues/30647. Thanks!

Hi! Do you plan on adding DONUT too? It would be highly appreciated :) As its decoder is a BART one, do you plan on doing the two alltogether for instance? 
Thanks!

@vbonnivardprobayes T5 compatibility is close to being done (https://github.com/huggingface/transformers/pull/34089), BART and related models will get the changes next :) (cc @zucchini-nlp )

@fzyzcjy Not on beam search. I'll be separating prefil into a separate function as my next task, then I'll probably work on vectorized beam search :)

(see other tracker: https://github.com/huggingface/transformers/issues/30810)

Are there any plans to add cuda graph support for models that are partitioned over multiple GPUs?

@tsengalb99 if that's possible then yes :D (multi-device is not my speciality, cc @SunMarc )

Hey, @tsengalb99 , we are integrating TP with transformers and it is also compatible with torch.compile . Could you confirm that it is comptible with cuda graph @kwen2501 ? Otherwise, PP should work with torch.compile cuda graph also with the latest pytorch 2.5

@SunMarc so this applies to models scattered across multiple GPUs with DeepSpeed via Accelerate too?

@gante Can you update the description to reflect any new model compile support?

I updated the models which have compile compatibilty already and added a few more links for open PRs. This list is not comprehensive and shows commonly used models. At this moment all decoder-only models should be compile friendly, and the encoder-decoder ones will come in the next batch

This tracker is mostly complete, and we can better track support by looking for `cache_position` in the signature of `model.forward` -- if the model accepts `past_key_values` but DOESN'T accept `cache_position`, then it doesn't have support for compilation 🤗 

End-to-end compilation will not be pursued, at least not until `torch` has solid support for data-dependent conditional branching. Instead, we will be working towards enabling the compilation of each decoding iteration, which contains the bulk of the computations.

Hi @gante . Is there any plan to support static cache for opt model?

@jiqing-feng it will be added very soon with https://github.com/huggingface/transformers/pull/35724 :)

> [@jiqing-feng](https://github.com/jiqing-feng) it will be added very soon with [#35724](https://github.com/huggingface/transformers/pull/35724) :)

Could you please help  by sharing inputs on a possible fix for this specific error :  [AttributeError: 'StaticCache' object has no attribute 'max_batch_size'. Did you mean: 'batch_size'?](https://github.com/huggingface/parler-tts/issues/214)

Thanks in advance

`max_batch_size` was renamed into `batch_size` -- updating the name of the attribute on your end should fix it :)

> This tracker is mostly complete, and we can better track support by looking for `cache_position` in the signature of `model.forward` -- if the model accepts `past_key_values` but DOESN'T accept `cache_position`, then it doesn't have support for compilation 🤗

I've tried to fullgraph-compile a recent `PleIAs/Modal` (for training time, not inference), and it succeeds with SDPA backend, but probably not impossible to make it work with FA/varlen too:
- https://github.com/huggingface/transformers/issues/42950

@gante Do you have any ideas of feasibility of fixing un/pad for fullgraph-ability? IMO in the worst case it would entail adding an optional model arg for user to pass pre-computed un/pad metainfo, but otherwise should be feasible

---

Also, are there CI-automated tests for checking popular models under various torch.compile options?

Hey @vadimkantorov , FA2 is known to be hard to compile in fullgraph mode. We don't support it yet and compile in non-fullgraph mode for FA2 models. If you have any proposal to make it fullgraph-able, feel free to raise a PR

> Also, are there CI-automated tests for checking popular models under various torch.compile options?

Yes, we have tests for checking compile and checking that there are no graph breaks with all model architectures. Generation tests are in `tests/utils/test_utils.py`, e.g. `test_generate_with_static_cache` or `test_generate_compile_model_forward_fullgraph`

> Yes, we have tests for checking compile and checking that there are no graph breaks with all model architectures.

Oh, very nice! Is the resulting nightly/release results table published anywhere?

> If you have any proposal to make it fullgraph-able

Are there any other specific issues with it besides pad/unpad utils in HF? (as this meta-info could be lifter to model args in the worst case similar to `cache_positions` you mentioned in https://github.com/huggingface/transformers/issues/43011#issuecomment-3685184713 - or at least these methods could be maybe monkey-patched to return pre-calculated user-provided values?)

---

Also, is there a non-inplace version of Cache - i.e. so that the new cache is returned as output as torch.cat'ing of the old and new cached values? thus supporting `torch.func.grad` (it currently doesn't support any inplace):
- https://github.com/pytorch/pytorch/issues/171158

Or does DynamicCache do this?

> FA2 is known to be hard to compile in fullgraph mode.

There is also some very recent merged PR:
- https://github.com/Dao-AILab/flash-attention/pull/1769

But not sure if it solves fullgraph (and if it also concerns FAv2)

> Are there any other specific issues with it besides pad/unpad utils in HF?

I don't really remember. Looking at the above linked PR, could be that `unpad` is the only issue left now, though it needs validation

> Also, is there a non-inplace version of Cache?

DynamicCache will do the cat'ting, yes. Though the cache class is always updated in-place, i.e. `cache.update(keys, values)` does not return a new instance of `Cache`

> Hey @vadimkantorov , FA2 is known to be hard to compile in fullgraph mode. We don't support it yet and compile in non-fullgraph mode for FA2 models. If you have any proposal to make it fullgraph-able, feel free to raise a PR

Do you have a reproducer for this case? I've added support for torch.compile in FA3 and I'm happy to take a look at this. This would also ensure that we can use `fullgraph=True` in recent versions of flash attention.