Add Tensor Parallel support for ALL models

[ArthurZucker]

Just opening this to add support for all models following #34184

Lets bring support to all model! 🤗

• Llama

It would be great to add the support for more architectures such as

• Qwen2
• QwenVl
• Mistral
• Llava

... and many more

For anyone who wants to contribute just open a PR and link it to this issue, and ping me for a review!! 🤗

[farrosalferro]

If it's okay, I want to take the Gemma and Gemma2 @ArthurZucker

[kmehant]

@ArthurZucker I will add support for Granite. Thanks

[VladOS95-cyber]

Hey @ArthurZucker! I am going to work on Mistral.

[VladOS95-cyber]

Hey @ArthurZucker! I am going to start working on Qwen2

[akshayg08]

For tensor parallel, when changing the dimension while applying the heads, why is num_heads made to -1. This essentially means that the head computations are being distributed, which is fine. But, an alternate approach could be to keep the num_heads same per device, and let head_dim be dynamic (-1). Is there a problem with the latter?

[akshayg08]

For tensor parallel, when changing the dimension while applying the heads, why is num_heads made to -1. This essentially means that the head computations are being distributed, which is fine. But, an alternate approach could be to keep the num_heads same per device, and let head_dim be dynamic (-1). Is there a problem with the latter?

I think it might affect the rotatory embeddings so it's better to split the heads