Google announces MaxText long-context training

Google Cloud has announced long-context training support in MaxText, its open-source JAX training framework, utilizing four context-parallel strategies combined with the Tokamax Splash Attention kernel. Benchmarks demonstrate that Llama 3.1 8B achieved 70 percent Model FLOPs Utilization on TPU 5p for sequence lengths up to 4 million tokens, with context extendable to 14 million tokens across TPU 5p and TPU 7. Open article →