Google adds long-context training to MaxText on Cloud TPUs
Overview
Google Cloud announces long-context training support in MaxText, its open-source JAX training framework, with four context-parallel strategies on Cloud TPU. The feature is available in MaxText v0.2.4 on PyPI, in the corresponding Docker images, and on GitHub main.
Google says Llama 3.1 8B reaches 70% Model FLOPs Utilization on TPU 5p at sequence lengths up to 4 million tokens. It says context can extend to 14 million tokens on TPU 7, with MFU staying at 50% below 6 million tokens.
Written by AI from the articles below · updated Nov 9, 3:44 PM ET
Check the sources:
Article timeline
The articles in this story. Times are ET.
- Google Developers BlogOfficialGoogle adds long-context training support to MaxText on Cloud TPUs
AIGoogle Cloud announces long-context training support in MaxText, its open-source JAX training framework, with four context-parallel strategies on Cloud TPU. Google says Llama 3.1 8B reaches 70% Model FLOPs Utilization on TPU 5p at sequence lengths up to 4 million tokens, and context can extend to 14 million tokens with TPU 7 maintaining 50% MFU below 6 million tokens. The feature is available in MaxText v0.2.4.
Heat Trend
- Current comparable heat
- 5
- Peak comparable heat
- 6Nov 9, 4:00 PM ET
- Change in past 24 hours
- –
The trend compares only continuously observed participants, so its scope may be smaller than the current heat total. Hover or tap for hourly heat; use arrow keys with a keyboard.