Architecture improvements

[ylacombe]

Supersedes #50 and #55:

• Add Grouped-Query Attention
• Clean attention
• Add and fix FA2 and SDPA from Flash Attention Support #55
• Add prompts in cross attention + RoPE from [Modelling] ROPE and Prompt Cross-Attention #50
• Fix generation from [Modelling] ROPE and Prompt Cross-Attention #50
• Make RoPE 25% faster
• Add training flexibility in using attention implementaton
• Makes smoother multi-nodes training
• Fix some torch compile and padding issues

[Artyom17]

Any ETA on this to be completed (at least partially) and committed?

[sanchit-gandhi]

Should we merge this one when ready @ylacombe to make development of new features more straightforward? After merging this PR, we can open follow-ups to iterate on the design further. Should help with the integration of static cache (#89) among others! cc @eustlb

[ylacombe]

Let's merge it ASAP, would you like to make a quick review first? It's been quite thoroughly "tested" since I've trained new checkpoints with the current architecture and the new one, as well as testing generation when evaluating the model

Feel free to merge when you read the message if you don't make the review!

[sang-nguyen-ts]

@ylacombe is there any update on this?

[ylacombe]

@sang-nguyen-ts it's just been merged !