Ring Attention & Friends

How Gemini 1.5 Scaled To 10 Million Tokens of Context

​

2nd March - Yannic’s Discord

Contents

1. Background
1. Some context on context
2. The technical work in a 4 part journey
• Vanilla Transformer Attention
• Flash Attention
• Ring Attention
• Striped Attention
3. What long context enables:
• World Models
• Gemini 1.5
4. Questions and After Hours

🔑 Key Points Will Be…

​

• Reducing Memory Footprint is unreasonably effective for a lot of problems
• To that end, tiling is the trick that seems to work
• Play to the hardware - Always think about GPUs
• Avoid premature optimisation - look at the actual bottleneck
• Sometimes small ideas make the largest difference

​

Chapters

• Background
• Some context on context
• The technical work in a 4 part journey
• Vanilla Transformer
• Flash Attention
• Ring Attention
• Striped Attention
• What long context enables:
• World Models
• Gemini 1.5
• Questions and After Hours

Some Context on Context

Some Context on Context

Chapters

• ✅ Background
• ✅ Some context on context
• The technical work in a 4 part journey
• 🤖 Vanilla Transformer Attention
• Flash Attention
• Ring Attention
• Striped Attention
• What long context enables:
• World Models
• Gemini 1.5
• Questions and After Hours

🤖 Vanilla Transformer

🤖 Vanilla Transformer

🤖 Vanilla Transformer

Row-wise softmax normalisation constant

🤖 Vanilla Transformer

O(N²) compute and memory requirements

🤖 Vanilla Transformer

What’s the bottleneck?

Chapters

• ✅ Background
• ✅ Some context on context
• The technical work in a 4 part journey
• ✅ Vanilla Transformer
• ✨ Flash Attention
• Ring Attention
• Striped Attention
• What long context enables:
• World Models
• Gemini 1.5
• Questions and After Hours

✨ Flash Attention - Naive Attention

We can split Q into Q_i tiles pretty easily but we don’t seem to be able to split K,V because they’re coupled along the sequence dimension by the normalisation constant 🤔

✨ Flash Attention - 1st Attempt at Tiling

Red before Blue

✨ Flash Attention - Tiling with the Softmax Rescaling Trick

✨ Flash Attention - Tiling Video

✨ Flash Attention - Upshot

Enables 4x larger context window

Also Train GPT-3 >2x faster than previous methods

Reduces memory footprint of attention algorithm from O(n²) to O(logn)

Chapters

• ✅ Background
• ✅ Some context on context
• The technical work in a 4 part journey
• ✅ Vanilla Transformer
• ✅ Flash Attention
• 💍 Ring Attention
• Striped Attention
• What long context enables:
• World Models
• Gemini 1.5
• Questions and After Hours

💍 Ring Attention

Suppose we want to increase the sequence length to 100 million tokens.

• What’s the new bottleneck?

💍 Ring Attention

💍 Ring Attention

Chapters

• ✅ Background
• ✅ Some context on context
• The technical work in a 4 part journey
• ✅ Vanilla Transformer
• ✅ Flash Attention
• ✅ Ring Attention
• 🦓 Striped Attention
• What long context enables:
• World Models
• Gemini 1.5
• Questions and After Hours

🦓 Striped Attention

Can we make Ring Attention even faster?

• If we’re doing autoregressive generation, then apparently yes!

Chapters

• ✅ Background
• ✅ Some context on context
• ✅ The technical work in a 4 part journey
• ✅ Vanilla Transformer
• ✅ Flash Attention
• ✅ Ring Attention
• ✅ Striped Attention
• What long context enables:
• World Models
• Gemini 1.5
• Questions and After Hours

🌍 Large World Models

Chapters

• ✅ Background
• ✅ Some context on context
• ✅ The technical work in a 4 part journey
• ✅ Vanilla Transformer
• ✅ Flash Attention
• ✅ Ring Attention
• ✅ Striped Attention
• What long context enables:
• ✅ World Models
• Gemini 1.5
• Questions and After Hours

🔑 Recap of Key Points

​

• Reducing Memory Footprint is unreasonably effective for a lot of problems
• To that end, tiling is the trick that seems to work
• Play to the hardware - Always think about GPUs
• Avoid premature optimisation - look at the actual bottleneck
• Sometimes small ideas make the largest difference

​

Chapters

• ✅ Background
• ✅ Some context on context
• ✅ The technical work in a 4 part journey
• ✅ Vanilla Transformer
• ✅ Flash Attention
• ✅ Ring Attention
• ✅ Striped Attention
• ✅ What long context enables:
• ✅ World Models
• ✅ Gemini 1.5
• Questions, References and After Hours

References & Marginalia

Core Papers:

• 🌟 Flash Attention
• 🌟 Ring Attention
• Striped Attention
• Large World Models
• 🌟 Gemini 1.5 Technical Report

Other resources:

• Self Attention Does Not Need O(n²) memory
• Blockwise Parallel Transformer for Large Context Models
• 🌟 Making Deep Learning Go Brrrr From First Principles
• 🌟 Tri Dao on Flash Attention (video)
• 🌟 Mamba Explained
• 🌟 Ring Attention PyTorch Implementation
• Attention Tiling (short video)
• The Illustrated Transformer
• Descriptive Matrix Ops with Einsums