TPU7x (Ironwood) 性能优化

本指南介绍了通过高效管理 TPU7x (Ironwood) 多层级内存系统之间的数据移动来优化其性能的几种方法。这包括低精度训练、分片、通信优化、激活重新具体化、作用域向量内存 (VMEM) 调优和自定义加速器内核等技术。

如需使用 TPU7x 优化性能，您必须先熟悉 Ironwood 架构，尤其是内存层次结构和互连拓扑。如需了解详情，请参阅 TPU7x (Ironwood)。

使用 FP8 进行低精度训练

FP8（8 位浮点）是一种高效的数值数据格式，主要用于加速模型训练和推理。通过使用 8 位（而非标准的 16 位格式 [FP16 或 BF16] 和 32 位 [FP32]）表示数字，TPU 可以显著加快数据处理速度并减少内存用量。

TPU7x 支持针对 FP8 数据类型的内置硬件加速，每个芯片的理论峰值性能为 4614 TFLOPS。此功能可以显著缩短端到端训练时间。对于兼容的操作，尤其是 AI 工作负载中常见的密集矩阵乘法，使用 FP8 可以使性能比标准 BF16 训练提高 1.3 倍。与 BF16 相比，FP8 可将峰值 FLOP 提高一倍，并将权重和激活的内存占用减半。对于计算受限的工作负载以及受内存容量或带宽限制的场景，FP8 应作为主要的调优手段。

使用 FP8 可带来以下性能优势：

• 降低高带宽内存 (HBM) 压力：内存占用空间较小，可让较大的模型或推理期间具有较大 KV 缓存的模型完全容纳在 192 GiB 的 HBM 中。这样可以避免将数据卸载到速度较慢的主机内存中，从而节省成本。
• 有效批次大小增加：通过减少激活所需的内存，FP8 允许使用更大的批次大小。这可以提高数据并行性，并提高吞吐量和计算单元的利用率。
• 降低内存带宽要求：每次操作移动的数据量减半，从而降低了对 HBM 到 MXU 数据路径的需求。在数据移动是常见瓶颈的系统上，这有助于使 MXU 保持工作饱和状态。

若要使用 FP8 且性能不下降或仅略微下降，需要谨慎选择量化技术。以下是一些可供您参考的 FP8 训练最佳实践：

• 缩放粒度：首先以每个张量伸缩作为基准。如果存在质量或性能问题，请切换到按轴伸缩。可能不需要子渠道伸缩。
• 缩放模式：动态缩放会在运行时计算伸缩比例，是保持质量的理想默认设置。虽然静态伸缩可以通过消除计算来显著提升性能，但它需要仔细分析以确定正确的伸缩比例，并且可能不适用于所有用例，尤其是在模型配置发生变化时。相反，一些稳健的模型和配置可以将缩放比例固定在 FP8 权重或激活值的上限，从而在保持准确率和提升性能的同时，降低量化开销。
• FP8 格式（E4M3 和 E5M2）：一种常见且有效的方法是混合使用 FP8 格式。例如，在前向传递中，对权重和激活使用 E4M3，以利用 E4M3 的更高精度；在反向传递中，对梯度使用 E5M2，以适应梯度的更宽动态范围。
• 舍入：对梯度使用“舍入到最接近的偶数”(RNE) 而不是随机舍入，可以在保持质量的同时提供更好的性能和可再现性。
• 在 MaxText 中启用 FP8：MaxText 通过 QWIX 量化库支持 FP8 训练。如需启用量化，请在配置中设置以下标志：use_qwix_quantization=true。

分片和并行处理

分片是指将大型模型或其训练数据切分成较小的片段，并将其分布在多个 TPU 芯片或核心上的过程。选择正确的分片策略对于在 TPU7x 上实现高性能至关重要。

如果只是单纯地最大化并行度，往往会因为通信受限而导致性能不佳。最佳方法通常是选择满足内存限制的最简单分片策略，因为这样可以最大限度地减少通信开销，并使计算单元得到高效利用。

在选择分片策略之前，任何性能调优工作的第一步都应该是算术强度分析。此分析可确定给定计算是否受计算、内存带宽或互连带宽的限制。其计算方式为浮点运算与必须移动的数据字节的比率。

算术强度高表示工作负载受限于计算资源。算术强度较低表示工作负载受内存或通信限制，其性能受限于从 HBM 或通过 ICI 网络移动数据的速度。此分析可帮助您确定理想的批次大小和分片策略。例如，受通信限制的工作负载不会从引入更多通信的分片策略（例如高度张量并行处理）中受益。

分片策略决策框架

MaxText 提供了多种分片策略。最佳选择取决于模型架构、序列长度以及平衡计算负载与通信开销的需求。

• 完全分片数据并行处理 (FSDP)：这是数据并行处理的首选默认策略。FSDP 会在数据并行设备之间对模型权重、梯度和优化器状态进行分片。在计算期间，每个设备都会执行 All-Gather 操作，以检索其本地微批次所需的完整权重。只要每个设备的批次大小足够大，能够掩盖这种 All-Gather 通信的延迟，FSDP 就会非常有效。对于混合专家 (MoE) 模型，算术强度计算必须考虑稀疏性。
• 张量并行处理 (TP)：TP 会在设备之间对单个张量进行分片。通常，这些张量是多层感知器 (MLP) 和注意力块中的权重矩阵。硬件的高算术强度 (11.5k) 对模型的维度提出了非常高的要求，以便通过 ICI 实现 TP，而尝试使用 TP 可能会导致系统受通信限制。
• 专家并行 (EP)：这是训练 MoE 模型的标准和必要策略。EP 会在一组设备上对“专家”层进行分片，并使用 All-to-All 通信集合将 token 路由到其指定的专家设备。如果模型的 MLP 维度足够大，可以接近上限，则 EP 可能会很高效。
• 上下文并行 (CP)：CP 是一种专用策略，对于训练序列长度非常长的模型至关重要。其主要功能是管理激活的内存消耗，该消耗会随着序列长度呈二次方增长，并可能超出 HBM 容量。CP 会对激活张量的序列维度进行分片，从而允许使用分数形式的每个设备批次大小。由于 CP 引入的通信比 FSDP 多，因此一般规则是使用满足内存限制所需的最小 CP 度，并确保批处理轴分片保持为整数。

下表将常见工作负载类型映射到最佳分片策略：

工作负载类型	建议的主分片	辅助分片	主要瓶颈	理由
密集模型 - 短序列	FSDP	不适用	重新具体化、FF 矩阵乘法	FSDP 可提供最佳平衡。对于短序列，激活内存可能不是主要问题。关键在于全局批次要足够大，以便隐藏 FSDP 的权重 All-Gather。随着批次大小的增加，激活大小也会增加，因此需要采用合适的重新具体化政策，以确保此配置不会出现内存不足的情况。
密集模型 - 长序列	FSDP	CP	Flash Attention、激活内存	激活内存成为主要限制。必须使用 CP 才能启用分数形式的每设备批次大小，并避免内存不足 (OOM) 问题。Flash Attention 是计算和时间浪费的主要来源。
MoE 模型 - 短序列	FSDP + EP	不适用	All-to-All（专家路由）、重新具体化	MoE 模型需要 EP 来对专家进行分片。用于令牌路由的多对多通信是一个主要瓶颈，必须进行重叠。重新具体化也是造成浪费的一个重要原因。
MoE 模型 - 规模非常大	FSDP + EP + PP	模型并行处理 (MP)	所有前述瓶颈，外加流水线气泡	对于超出单个 Pod 内存的模型，需要使用 PP 在 Pod 之间对层进行分片。这会产生 DCN 通信和流水线气泡开销。这是一种高度复杂的配置，需要仔细调优。

沟通优化

在 TPU7x 上重叠通信和计算的主要机制称为 SparseCore 集合分流。Ironwood 架构包含专用的 SparseCore 单元，这些单元充当独立的控制线程，能够管理 ICI 结构上的数据移动。这样，集体通信操作（如 All-Gather 或 Reduce-Scatter）就可以与 TensorCore 上发生的主要计算并行执行。这是 TPU7x 上异步集合的推荐方法。使用推荐的标志为最常见的集合启用分流。

激活重新具体化

激活重计算（也称为梯度检查点）是减少模型 HBM 占用空间的一项基本技术。它不会将前向传播中的所有中间激活存储在 HBM 中以供反向传播期间使用，而是仅保存一些关键激活（检查点），并在反向传播期间根据需要重新计算其他激活。这会节省大量内存，但会增加计算量（标准 Transformer 块的 FLOP 额外增加约 25-30%）。

如何积极应用重新具体化是一个关键的调优参数，完全取决于主要瓶颈，而瓶颈通常会随序列长度而变化。

对于长序列工作负载（例如 128k）：在这些情况下，激活张量的大小是 HBM 的主要消耗者。工作负载通常受内存限制。因此，应用激进的重新具体化政策非常有益。节省的内存可让训练继续进行，而不会出现内存不足错误，并且还允许更大的批次大小，而重新计算的计算开销是值得的。

对于短序列工作负载（例如 8k）：在这些情况下，激活内存不太重要，工作负载更有可能受计算限制。重新具体化的计算开销可能是导致效率低下的最大单一来源。

在 MaxText 中调整重新具体化政策

MaxText 通过一组预设和自定义政策（使用 remat_policy 标志进行配置）提供对重新具体化的精细控制。

预设政策

MaxText 提供以下内置政策：

• full：最激进的政策，几乎会重新具体化所有内容。这样可以最大限度地减少 HBM 使用量，但会最大限度地增加重新计算开销。非常适合内存极度受限、长序列的场景。
• minimal：最不激进的政策，会存储大多数激活。这样可以最大限度地提高 HBM 使用率，同时最大限度地减少重新计算。最适合短序列、计算密集型工作负载，无需考虑内存。
• 中间策略：save_dot_with_context_except_mlp、save_qkv_proj 和 save_out_proj 等选项通过选择性地对昂贵的点积运算的输出进行检查点设置，同时重新具体化更便宜的逐元素运算，提供了各种权衡。

自定义政策

为了获得更高级别的控制，您可以将 remat_policy 设置为 custom。这样，您就可以为模型解码模块中的各个层指定行为。每个层都可以分配以下三种行为之一：

• device：激活存储在 TPU 设备上的 HBM 中。
• remat：激活会被舍弃，并在反向传播期间重新实例化。
• offload：激活从 HBM 移至 CPU 主机的内存，从而释放 HBM，但会产生 PCIe 传输延迟。

限定范围的 VMEM 调优

内核性能（例如 FlashAttention）取决于内核中选择的图块大小，而图块大小受限于可用的向量内存 (VMEM)。TPU7x 芯片中的两个 TensorCore 各有 64 MiB 的向量内存 (VMEM)。此 VMEM 容量可以在当前范围（限定范围的 VMEM）和未来权重预取之间进行分配。增加限定范围的 VMEM 可以增加内核中的图块大小，从而可能减少内存停滞并提高内核性能。您可以通过设置 xla_tpu_scoped_vmem_limit_kib（在 LIBTPU_INIT_ARGS 中）来更改限定范围的 VMEM 大小，这可用于探索内核性能以及端到端性能限制。优化作用域 VMEM 大小可以间接影响自定义 Pallas 内核性能，因为增加作用域 VMEM 会为内核内图块大小解锁更大的超参数搜索空间。

Tokamax 内核

Tokamax 是一个高性能 JAX 内核库，包含许多高度优化的 TPU 内核，可解决几个常见的硬件特定瓶颈：

• Splash 注意力机制：Splash 注意力机制用作主要注意力机制实现，可消除标准注意力机制的 HBM 瓶颈，并在 TPU 上使用最高效的注意力机制实现。
• Megablox 分组矩阵乘法 (GMM)：对于 MoE 工作负载，Megablox 通过对不规则激活表示进行计算，高效处理分组矩阵乘法。它能高效地映射不规则维度，计算 LHS 中不规则行组与相应专家矩阵之间的矩阵乘法，从而避免需要将批次填充到固定大小。
• 使用 tune-jax 进行经验性调优：tune-jax 库具有用于对最佳块大小执行经验性搜索的实用程序。默认内核大小通常不是最优的；通过调优，可以选择对硬件友好的 VMEM 图块大小，从而最大限度地提高硬件利用率。
• 最大 logits 估计值：通过为 max_logit_const 设置值，可以进一步优化 Tokamax Splash 注意力内核。如果设置，它会在注意力机制的 softmax 运算 (softmax(Q * KT)) 期间替换最大 logit 的缩减计算，从而减少一些计算和同步开销。在 MaxText 中，它由配置 use_max_logits_estimate 实现，可设置为 None（停用）或浮点值。验证特定模型的 logit 范围是否仍与估计值兼容，以防止数值溢出。 如果设置了此值，建议进行收敛测试。