Google Cloud MaxText 新增上下文平行處理，支援 TPU 上的長上下文訓練

Google Cloud 在 MaxText v0.2.4 中發布長上下文訓練支援，這是其開源的 JAX 訓練框架，在 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 的範例配方。