My architecture/optimizations overhaul proposal to shrink the gap with the consumer hardware

[kabachuha]

Hi, xAI!

From what I saw in its official code, it's highly sub-optimal (it uses the "vanillest" implementation of Attention, not even the Flash one!). I believe it will benefit greatly for pre-existing and recent optimizations from the NLP community. Examples include:

1. XFormers Flash Attention
2. Torch 2 Scaled dot product attention
3. Ring Attention (from LWM - Large World Model), to scale the context to millions of tokens
4. Linear "Rebased" Flash Attention - eliminates the quadratic attention cost and requires less rework than switching to Mamba, more compatible to existing architecture
5. Hybrid attention head sequence parallelism (e.g. FastSeq) including ZeRO (see Accelerated Transformer)
6. Layer offloading
7. Bitsandbytes/AutoGPTQ Quantization

And the best thing is that all these points (except for 1/2 and 3/4) can be combined!

With these optimizations we can hope to shrink the gap between the current requirements and the consumer/low spectrum server requirement

Can't wait to have it in the Transformers library 😬 Track

Edit: add links to resources

Feel free to give feedback!