tracker: generate compatibility with torch.compile

[gante]

generate 🤜 🤛 torch.compile

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

• 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

(models tagged as "important models" in our CI + popular models)

Decoder-only:

• GPT-J is compatible
• OPT is compatible
• 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)
• Starcoder2 is compatible
• Phi is compatible (Phi: static cache & compile compatibility #30688)
• Phi3 is compatible (Phi: static cache & compile compatibility #30688)
• BLOOM is compatible (note: this one might be tricky due to cache format)
• Mamba is compatible (requested here) (PR: Add torch.compile Support For Mamba #31247)

Encoder-decoder:

• BART is compatible
• T5 is compatible
• Whisper is compatible (Whisper Static Cache  #30760)

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]

👀