Announcing Long-Context Training on TPUs with MaxText Context Parallelism

Google Cloud announced long-context training support in MaxText, its open-source JAX training framework, via four context-parallel strategies on Cloud TPU: All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention), and the A2A + P2P hybrid (USP), all using the Tokamax Splash Attention kernel. Benchmarks report Llama 3.1 8B reaching 70% MFU on TPU 5p for sequences up to 4 million tokens, with context extendable to 14 million tokens on TPU 5p and TPU 7, where 14M-token P2P training on 64 TPU 7 chips reached 48.7% MFU versus 37.1% on TPU 5p. The feature is available in MaxText v0.2.4 on PyPI, Docker images and GitHub main, with sample recipes for TPU 7 and TPU 5p.