Announcing Long-Context Training on TPUs with MaxText Context Parallelism
Summary
- Google Cloud announced long-context training support in MaxText v0.2.4, its open-source JAX training framework for Cloud TPUs.
- The release introduces four context-parallel strategies: All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention), and the A2A + P2P hybrid (USP), powered by the Tokamax Splash Attention kernel.
- Benchmarks on a 64-chip TPU 5p slice running Llama 3.1 8B achieved 70% MFU up to 4 million tokens, with context length scaling up to 14 million tokens across TPU 5p and TPU 7 chips.
Why it matters
Efficient long-context training enables large language models to process multi-million token sequences natively without exceeding HBM capacity, overcoming major memory and compute bottlenecks on hardware accelerators.