Google Cloud 开源框架 MaxText 支持 TPU 长上下文训练

Google 跟踪代码管理器（noscript）结束 Google 跟踪代码管理器（noscript） HTML 宣布在 TPU 上通过 MaxText 上下文并行实现长上下文训练 2026 年 11 月 9 日 Michael Tran 软件工程实习生 Google Cloud Nuojin Cheng 软件工程师 Google Cloud Deepak Patil 集团产品经理

长上下文训练——包括预训练和后训练——已成为前沿 AI 模型的基础。更长的原生序列长度为大语言模型（LLM）提供了工作记忆，使其能够读取整个代码仓库或摄取数小时的视频，并对多步骤的智能体任务进行推理。三年间，上下文窗口已从 128K–200K token 增长到 1M–2M，而面向 1000 万 token 训练的模型也正在到来。在这些长度上进行训练（而不仅仅是提供服务）正是系统瓶颈显现之处。

具体而言，在数百万 token 的序列长度上进行高效训练会遇到以下瓶颈：

注意力计算： 在长上下文下，注意力 FLOPs 随序列长度呈二次方增长，注意力在训练步骤中占据主导地位。吞吐量取决于注意力内核。
内存与通信： 激活值（包括每个 token 的键和值的聚合量）会随序列长度增长。像朴素的全收集（all-gather）那样在每个芯片上复制它们，会耗尽高带宽内存（HBM）并触发内存不足故障。此外，与内存相关的约束会迫使进行通信和重计算，这可能对模型 FLOPs 利用率（MFU）产生负面影响。这里的解决方案是上下文并行（CP），它流式传输键和值，而不是将它们完整物化。

今天，我们很高兴地宣布在 MaxText中支持长上下文训练，这是 Google Cloud 的开源 JAX 训练框架。MaxText 通过在 Cloud TPU: 上支持四种上下文并行策略来实现这一点：全收集（All-gather）、A2A（DeepSpeed Ulysses）、P2P（Ring Attention），以及 A2A + P2P 混合（统一序列并行，即 USP），并结合负载均衡和序列打包。这四种策略均使用 Tokamax Splash Attention 内核. 结合 3D 芯片间互连（ICI）（在 TPU 5p 或 TPU 7（Ironwood）等 TPU 上支持），这使得单个序列可以跨越整个切片，同时通信与计算重叠，且每个芯片的内存占用保持平稳。在我们的基准测试中，Llama 3.1 8B（25 个 head dim，16 个 query heads）展现出卓越的长上下文效率，在 TPU 5p 上、序列长度最高达 400 万 token 时实现了 70% 的模型 FLOPs 利用率（MFU） 。在 TPU 5p 和 TPU 7 上，上下文均可无缝扩展至 1400 万 token——其中 TPU 7 在低于 600 万 token 的所有序列规模下始终维持 50% MFU 。

此功能现已可在 MaxText v0.2.4(文档).

MaxText 中的上下文并行

上下文并行将长序列拆分到多个芯片上，因此没有任何一个芯片持有整个序列的激活值。上下文并行可以与数据并行和张量并行组合使用，而通常最合适的做法是选择能将你的序列放入 HBM 的最小上下文并行规模。这同样适用于更稀疏的注意力变体，它们减少了注意力计算量，但并不会减少其他每一层的激活值。

注意力是唯一一个全局操作，其中每个 token 都需要所有前序 token 的键和值，因此每种上下文并行策略都必须将这些键和值重新汇集到一起。每种策略在内存（一个芯片持有多少 KV）、通信（每层移动多少数据，以及是否与计算重叠）和模型形状（拆分多少个注意力头）之间做出不同的权衡。

全收集（All-gather）： 每个芯片保留自己那部分序列，并在每一层从其他所有芯片收集键和值。这是最简单的选项。然而，由于因果掩码会使较后的 token 关注更长的前缀，全收集需要在各设备之间平衡计算负载。

图 1. 全收集，针对一个拆分到四个设备（TPU 芯片）上的序列。每个设备收集所有键和值，然后在本地进行注意力计算。带轮廓的方块表示正在计算的注意力；带轮廓的小方块表示传输中的数据。

全对全（Ulysses）： 收集键和值会让每个芯片都持有完整序列，因此无论增加多少芯片，注意力内存都不会缩小。全对全改变的是每个芯片所拥有的内容。注意力在头之间是独立的，但在 token 之间不是，因此头轴可以自由分片，而 token 轴则不行。注意力之前的全对全会将分片从一种轴轮换到另一种轴：每个芯片放弃自己不负责的头，并接收它所保留的头所缺失的 token，最终得到的是少数几个头的完整序列，而不是所有头的少数几个 token。这也是 Ulysses 不需要 token 负载均衡的原因。随后，注意力以普通的单芯片注意力方式、使用标准内核运行，第二次全对全再将输出轮换回来。每个芯片的 KV 会按上下文并行因子缩小，但你不能拆分到比模型头数更多的芯片上——而张量并行分片的正是同一维度，因此 TP × CP 受头数上限约束（在 GQA 下，受 KV 头上限约束：Llama 3 8B 为八个）。不过，全对全本身并不与计算重叠——注意力无法在部分序列上开始，因此每个芯片都要等待来自其他所有芯片的数据，然后才能开始。

图 2. 全对全（Ulysses）。全对全用 token 块换取注意力头（各行）。每个设备对其头在完整序列上进行注意力计算，输出通过第二次全对全返回。

P2P（环形注意力）： Ulysses 的上限是头数；环形注意力通过始终不离开 token 轴来消除这一上限。键和值永远不会完整物化。KV 块沿着芯片环一次一跳地传递，从邻居到邻居。每个芯片在下一块传输途中时，对其持有的块执行查询的注意力计算，并在过程中合并部分结果。一个芯片持有自己的块和一个来访块，从不持有完整序列，并且通信与计算重叠。然而，将注意力计算划分为更小的块可能会降低本地内核效率，而且与全对全 Ulysses 不同，环形注意力需要显式的负载均衡。

图 3. P2P（环形注意力）。KV 块一次一跳地轮转。每个设备在下一块到达时对其持有的块进行注意力计算，并且一次只持有一个来访块。

A2A + P2P（USP）： A2A + P2P（USP）：两种策略单独使用都无法扩展——全对全受限于头数，而长环在累积跳数的同时会稀释使其奏效的重叠。统一序列并行在芯片网格的两个轴上组合它们，因此上下文并行度成为两者的乘积：沿一个轴做全对全，其规模等于或小于 KV 头数；沿另一个轴做环，以扩展到超出该上限。由于每个全对全组已经汇聚了其 token，现在一个环步携带的块大了全对全因子那么多，环也相应更短——每一跳有更多计算来掩盖传输，累积延迟和因果不均衡的跳数也更少。在 Cloud TPU 上，两者映射到切片芯片网格的两个轴上。

图 4. A2A + P2P（USP）。全对全在每对设备内运行，P2P 环跨对运行，因此每个设备以更少的跳数对更长的块进行注意力计算。

这些策略覆盖了当今主流模型中的注意力变体，从多头注意力和分组查询注意力到 DeepSeek 风格的多头潜在注意力。

构建在快速注意力内核之上

只有当每个芯片上的注意力计算保持高效时，拆分序列才划算。全部四种策略都运行在 Splash Attention 上，这是来自 Tokamax的块稀疏 flash 注意力内核，Tokamax 是 Google 开源的 TPU 内核库。它会跳过完全被掩码的块，并在无填充的情况下处理打包序列。P2P 在每个环步调用它一次，作用于本地查询和刚刚到达的 KV 块，而下一步的传输在其下方运行，前向和反向传播皆如此。每个芯片执行与没有上下文并行时相同的融合块级工作，通信则隐藏在其后。

上下文并行策略与 TPU 性能

P2P（Ring Attention）是超长上下文训练的主要主力，而 All-Gather 是大多数其他 KV 能够容纳情况下的标准基线。这些是实践中使用的主要 CP 策略。我们比较了在固定的 64 芯片 TPU 5p 切片上，使用 Llama 3.1 8B（256 头维度，16 个查询头）时，每种策略能将上下文长度推到多远，并测量了它们的 MFU。在理想设置下，即两种策略都能有足够大的块（S/CP）和本地批大小——这是驱动 TPU 5p 上 MFU 的两个主要变量——我们观察到，P2P 在从 64K 到 14M token 的几乎所有上下文长度上都能优于 All-Gather，此时通信在 Ring 循环传输调度下基本被隐藏。

TPU v7x（Ironwood）将这一方案进一步推进。凭借每芯片两倍的 HBM，在 64 个 TPU 7 芯片上进行 1400 万 token 的 P2P 训练达到了 48.7% 的模型 FLOPs 利用率（MFU），而 TPU 5p 上为 37.1%。我们观察到，在 TPU 7 芯片上，ring attention 的表现优于其他上下文并行策略——例如 Ulysses 和 USP。

选择合适的策略

上下文并行策略的选择主要取决于汇聚后的 KV 是否能放入 HBM，以及模型有多少个 KV 头：

当汇聚后的 KV 能放入 HBM 时： 使用 All-gather，并开启负载均衡。其内存随整个序列增长，因此最先耗尽，但当它能放下时，它是最简单的选项，在我们的运行中是最快或接近最快的。对于分组查询模型，汇聚的代价可能很低，因为其 KV 大小只占激活值的一小部分。
在少量芯片上，最多到 KV 头数： 仅使用 A2A。一次全对全输入，一次输出，没有环；它不支持多查询注意力。
当汇聚后的 KV 不再能容纳时： 环形家族的内存随每块芯片的分块增长，而不是随整个序列增长。只要 KV 头数允许 A2A 因子达到 2 或以上，我们就推荐混合方案（USP）。它更长的分块和更少的跳数在许多设置下能比 P2P 更快，并且在 all-gather 还能容纳时与 all-gather 保持同步；一个序列跨越的芯片越多，这一点就越重要。对于 DeepSeek 等 MLA 模型，或当头数无法整除时，P2P 仍是唯一选择。

对于所有 CP 策略，当内存不足时：先卸载，再重算。MaxText 可以逐个张量地将保存的激活值、优化器状态和参数移至主机内存。在长上下文长度下，这些传输被计算隐藏，使卸载实际上不产生开销。先卸载注意力输出，然后是优化器状态、层输入、参数，最后才将重算作为最后手段。在卸载步骤中 MFU 保持平稳，而 all-gather 则下降并最终无法容纳。更多细节请参阅重算调优指南.

下一步

展望未来，我们看到两个自然的扩展方向。第一个是更广泛的模型覆盖。混合滑动窗口层和全局层的模型（如 Gemma）将在窗口层上通过相邻芯片之间的短交换来运行，而全局层则使用环形方案。我们还计划支持线性注意力机制，如 Gated Delta Net（GDN）和 Kimi Delta Attention（KDA）。第二个是性能，进一步调优 Tokamax 内核和 P2P 反向传播，以缩小短上下文下与 all-gather 的剩余差距，并为 TPU 5p 和 7 芯片提供按模型和上下文长度划分的配方。

立即开始

上下文并行现已可用，位于 MaxText v0.2.4 已在 PyPI 上发布，相应的 Docker 镜像以及 GitHub main 分支上也可用。示例配方可在此处获取，适用于 TPU 7 和 TPU 5p选择一种策略，使用 context_parallel_strategy（ring、ulysses 或 usp），设置上下文并行大小，并保持其余运行设置不变。分片指南包含每种策略的标志和约束，以及发布说明以及这项工作背后拉取请求的链接。我们迫不及待想看到你训练出的成果。

发布于：