Google Cloud MaxText 新增上下文并行,支持在 TPU 上进行长上下文训练
Google Cloud 在其开源 JAX 训练框架 MaxText v0.2.4 中发布了长上下文训练支持,在 Cloud TPU 上新增四种上下文并行策略:All-gather、A2A(DeepSpeed Ulysses)、P2P(Ring Attention)以及 A2A + P2P 混合(USP)。Llama 3.1 8B 的基准测试在 TPU 5p 上、序列长度最高达 400 万 token 时达到 70% MFU,上下文在 TPU 5p 和 TPU 7 上均扩展至 1400 万 token,其中 TPU 7 在低于 600 万 token 时保持 50% MFU。该版本已在 PyPI、Docker 镜像和 GitHub main 上提供,并附有 TPU 7 和 TPU 5p 的示例配方。
推荐理由:MaxText 在 Cloud TPU 上的四种上下文并行策略,以及 Llama 3.1 8B 的基准 MFU 数据,为团队选择如何在数百万 token 序列长度下进行训练提供了具体参考。