tracker: generate compatibility with torch.compile

[gante]

generate 🤜 🤛 torch.compile

Part of the PyTorch 2024 H2 roadmap.

This issue is a tracker of the compatibility between .generate and torch.compile (intro docs by pytorch). 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)

• greedy_search / sample are compatible (Generate: end-to-end compilation #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:

• GPT-J is compatible (Cache: new Cache format in decoder-only models #31421)
• GPT2 is compatible
• Llama is compatible ([Core generation] Adds support for static KV cache #27931)
• Gemma is compatible ([ gemma] Adds support for Gemma 💎 #29167)
• Llava is compatible (LLaVA torch.compile implementation #29891)
• Mistral is compatible (Add torch.compile for Mistral #30642)
• Mixtral is compatible (Add torch compile for mixtral #30793)
• Phi is compatible (Compile compatibilty for decoder-only models #32617)
• Phi3 is compatible (Phi: static cache & compile compatibility #30688)
• BLOOM is compatible (note: this one might be tricky due to cache format) Compile compatibilty for decoder-only models #32617
• Mamba is compatible (requested here) (PR: Add torch.compile Support For Mamba #31247)
• Persimmon is compatible (Compile compatibilty for decoder-only models #32617)
• Qwen2 Compile compatibilty for decoder-only models #32617
• Qwen2-VL Compile compatibilty for decoder-only models #32617
• Falcon Compile compatibilty for decoder-only models #32617
• GPTNeoX Compile compatibilty for decoder-only models #32617
• Starcoder2 Compile compatibilty for decoder-only models #32617
• StableLM Compile compatibilty for decoder-only models #32617

Encoder-decoder:

• BART is compatible
• T5 is compatible (T5 compile compatibilty #34089)
• Whisper is compatible ([whisper] static kv cache #31166)

Quantization

• BNB support
• GPTQ support
• AWQ support

Others

• We have a benchmark script to quickly compare the impact of PRs
• Add section to existing docs on the topic
• Confirm that pipelines work after compiling generate

[github-actions]

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 are likely to be ignored.

[kadirnar]

👀

[guangy10]

👀

[jiqing-feng]

👀

[mlazos]

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)

[fzyzcjy]

Hi, is there any updates? I am especially interested in the things related to #30647. Thanks!

[vbonnivardprobayes]

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!

[gante]

@vbonnivardprobayes T5 compatibility is close to being done (#34089), BART and related models will get the changes next :) (cc @zucchini-nlp )

[gante]

@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: #30810)

[tsengalb99]

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

[gante]

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

[SunMarc]

Hey, @tsengalb99 , we are integration 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

[joanvelja]

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