TPU v4

本文档介绍了 Cloud TPU v4 的架构和支持的配置。

系统架构

每个 TPU v4 芯片包含两个 TensorCore。每个 TensorCore 均有四个矩阵乘法单元 (MXU)、一个矢量单元和一个标量单元。下表显示了 v4 TPU Pod 的密钥规范。

关键规范	v4 Pod 值
每个芯片的峰值计算	每秒 275 万亿次浮点运算（bf16 或 int8）
HBM2 容量和带宽	32 GiB、1200 GBps
测量的最小/平均/最大功率	90/170/192 W
TPU Pod 大小	4096 芯片
互连拓扑	3D 网格
每个 Pod 的计算峰值	1.1 exaflop（bf16 或 int8）
每个 Pod 的全减少带宽	1.1 PB/秒
每个 Pod 的对分带宽	24 TB/秒

下图展示了 TPU v4 芯片。

3D 网格和 3D 环面

v4 TPU 在 3 维空间内直接连接到最近的邻近芯片，从而形成网络连接的 3D 网格。连接可以在切片上配置为 3D 环面，其中拓扑 AxBxC 为 2A=B=C 或 2A=2B=C，其中每个维度都是 4 的倍数。例如，4x4x8、4x8x8 或 12x12x24。一般来说，3D 环面配置的性能要优于 3D 网格配置。如需了解详情，请参阅扭曲拓扑。

TPU v4 与 v3 相比的性能优势

本部分介绍 TPU v4 的性能优势

内存系统

非统一内存访问 (NUMA) 是一种适用于具有多个 CPU 的机器的计算机内存架构。每个 CPU 都可以直接访问一块高速内存。CPU 及其内存称为 NUMA 节点。NUMA 节点连接到直接相邻的其他 NUMA 节点。一个 NUMA 节点中的 CPU 可以访问另一个 NUMA 节点中的内存，但这种访问比访问 NUMA 节点中的内存要慢。

在多 CPU 机器上运行的软件可以将 CPU 所需的数据存储在其 NUMA 节点中，从而提高内存吞吐量。如需详细了解 NUMA，请参阅维基百科上的非统一内存访问。

您可以通过将训练脚本绑定到 NUMA Node 0 来利用 NUMA 位置优势。

要启用 NUMA 节点绑定，请执行以下操作：

1. 安装 numactl 命令行工具。

    $ sudo apt-get update
    $ sudo apt-get install numactl
   

2. 将您的脚本代码绑定到 NUMA 节点 0。

    $ numactl --cpunodebind=0 python3 your-training-script
   

在以下情况下，请启用 NUMA 节点绑定：

• 如果您的工作负载严重依赖 CPU 工作负载（例如，映像分类、推荐工作负载），而不考虑框架。
• 如果您使用的是不带 -pod 后缀的 TPU 运行时版本（例如 tpu-vm-tf-2.10.0-v4），

其他内存系统差异：

• v4 TPU 芯片在整个芯片上拥有统一的 32 GiB HBM 内存空间，有助于在两个芯片上 TensorCore 之间更好地协调。
• 使用最新的内存标准和速度改进了 HBM 性能。
• 改进了 DMA 性能配置文件，并且内置对 512B 粒度的高性能步进支持。

TensorCores

• MXU 的数量翻倍，时钟频率提高，最多可提供 275 TFLOPS。
• 2 倍转置和排列带宽。
• 通用内存 (Cmem) 的加载和存储内存访问模型。
• 更快的 MXU 权重加载带宽和 8 位模式支持，可以降低批次大小并缩短推断延迟时间。

芯片间互连

每个芯片有六个互连链路，可实现网络直径更小的网络拓扑。

其他

• x16 PCIE gen3 接口到主机（直接连接）。
• 改进了安全模型。
• 提高了能源效率。

配置

TPU v4 Pod 由 4096 个芯片组成，通过可重新配置的高速链路互连。TPU v4 的灵活网络使您能够以多种方式连接相同大小的 Pod 切片中的芯片。创建 TPU Pod 切片时，您可以指定所需的 TPU 版本和 TPU 资源数量。创建 TPU v4 Pod 切片时，您可以通过以下两种方式之一指定其类型和大小：AcceleratorType 和 AccleratorConfig。

注意 ：AcceleratorType 和 AccleratorConfig 是互斥的；在创建 TPU 时，您只能使用其中一个。

使用 AcceleratorType

如果未指定拓扑，请使用 AcceleratorType。如需使用 AcceleratorType 配置 v4 TPU，请在创建 TPU Pod 切片时使用 --accelerator-type 标志。将 --accelerator-type 设置为包含您要使用的 TPU 版本和 TensorCore 数量的字符串。例如，如需创建包含 32 个 TensorCore 的 v4 Pod 切片，您可以使用 --accelerator-type=v4-32。

以下命令使用 --accelerator-type 标志创建一个包含 512 个 TensorCore 的 v4 TPU Pod 切片：

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt

TPU 版本后面的数字 (v4) 指定 TensorCore 的数量。v4 TPU 中有两个 TensorCore，因此 TPU 芯片的数量将为 512/2 = 256。

使用 AcceleratorConfig

如果要自定义 TPU 切片的物理拓扑，请使用 AcceleratorConfig。对于超过 256 个芯片的 Pod 切片的性能调整，通常需要这样做。

如需使用 AcceleratorConfig 配置 v4 TPU，请使用 --version 和 --topology 标志。将 --version 设置为您要使用的 TPU 版本，并将 --topology 设置为 Pod 切片中 TPU 芯片的物理排列方式。

您可以使用 3 元组 AxBxC 指定 TPU 拓扑，其中 A<=B<=C 且 A、B、C 全部 <= 4 或全部为 4 的整数倍。值 A、B 和 C 分别是这三个维度中的条状标签计数。例如，如需创建一个包含 16 个芯片的 v4 Pod 切片，您需要设置 --version=v4 和 --topology=2x2x4。

以下命令将创建一个 v4 TPU Pod 切片，其中有 128 个 TPU 芯片以 4x4x8 数组排列：

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt

2A=B=C 或 2A=2B=C 的拓扑也具有针对全对通信进行了优化的拓扑变体，例如 4×4×8、8×8×16 和 12×12×24。这类拓扑称为“扭曲鸟居”。

下图显示了一些常见的 TPU v4 拓扑。

较大的 Pod 切片可以通过一个或多个 4x4x4“立方体”芯片构建。

扭曲的鸟居拓扑

某些 v4 3D 环面切片形状可以选择使用所谓的扭曲环面。例如，可以将两个 v4 立方体排列为 4x4x8 切片或 4x4x8_twisted。扭曲拓扑提供的平分带宽明显更高。增加对分带宽对于使用全球通信模式的工作负载非常有用。扭曲拓扑可以提高大多数模型的性能，其中大型 TPU 嵌入工作负载受益最大。

对于将数据并行用作唯一并行策略的工作负载，扭曲拓扑的表现可能会略有提升。对于 LLM，使用扭曲拓扑的性能可能因并行类型（DP、MP 等）而异。最佳实践是在使用和不使用扭曲拓扑的情况下训练 LLM，以确定哪种拓扑可为模型提供最佳性能。FSDP MaxText 模型的一些实验使用扭曲的拓扑，实现了 1-2 项 MFU 方面的改进。

扭曲拓扑的主要优势在于，它将非对称环形拓扑（例如 4×4×8）转换为密切相关的对称拓扑。对称拓扑具有诸多优势：

• 改进了负载均衡
• 更高的对分带宽
• 缩短数据包路由

这些优势最终会转化为许多全球通信模式的性能的提升。

TPU 软件支持在切片上扭转 tori，其中每个维度的大小是最小维度的大小，或者是最小维度大小的两倍。例如，4x4x8、4x8x8 或 12x12x24。

例如，请考虑以下 4×2 环面拓扑，其中 TPU 在切片中带有 (X,Y) 坐标标签：

为清晰起见，此拓扑图中的边显示为无定向边。实际上，每个边缘都是 TPU 之间的双向连接。我们将此网格的一侧和另一侧之间的边称为环绕边，如图所示。

通过扭转此拓扑，我们最终得到一个完全对称的 4×2 扭曲环面拓扑：

上图与上图之间唯一的变化是 Y 环绕边。这些 TPU 已改为连接到具有 X+2 mod 4 坐标的 TPU，而不是连接到另一个具有相同 X 坐标的 TPU。

这种思路可以泛化到不同的维度大小和不同数量的维度。只要每个维度等于或等于最小维度大小的两倍，生成的网络就是对称的网络。

如需详细了解如何在创建 Cloud TPU 时指定扭曲的 tori 配置，请参阅使用 AcceleratorConfig。

下表显示了支持的扭曲拓扑，以及使用它们与无解拓扑相比在理论上增加了对分带宽。

扭曲拓扑	与非扭环面相比，带宽在理论上增加了
4×4×8_扭转	约 70%
8x8x16 扭曲
12×12×24_扭转
4×8×8_扭转	约 40%
8×16×16 扭曲

注意 ：除非使用 --topology 标志指定后缀 _twisted，否则系统不会自动使用扭曲拓扑。例如，如需指定 4x4x8 扭曲拓扑，可以将 topology 标志指定为：4x4x8_twisted。

TPU v4 拓扑变体

某些包含相同数量芯片的拓扑可以以不同的方式排列。例如，可以使用以下拓扑配置包含 512 个芯片（1024 个 TensorCore）的 TPU Pod 切片：4x4x32、4x8x16 或 8x8x8。包含 2048 个芯片（4096 个 TensorCore）的 TPU Pod 切片可提供更多拓扑选项：4x4x128、4x8x64、4x16x32 和 8x16x16。

与给定芯片数量关联的默认拓扑是最类似于立方体的拓扑。此形状可能是数据并行机器学习训练的最佳选择。其他拓扑对于具有多种并行处理（例如，模型和数据并行，或模拟的空间分区）的工作负载非常有用。如果拓扑与所用的并行性相匹配，则这些工作负载的性能最佳。例如，如果在 X 维度上并行放置 4 路模型，并在 Y 维度和 Z 维度上并行放置 256 路数据并行处理，则会与 4x16x16 拓扑匹配。

具有多个并行维度的模型在并行维度映射到 TPU 拓扑维度时效果最佳。这些模型通常是数据和模型并行大语言模型 (LLM)。例如，对于拓扑为 8x16x16 的 TPU v4 Pod 切片，TPU 拓扑维度为 8、16 和 16。使用 8 路或 16 路模型并行处理（映射到其中一个物理 TPU 拓扑维度）性能更高。采用此拓扑时，四路模型的并行性不是最优，因为它与任何 TPU 拓扑维度均不对齐，但对于相同数量的芯片，采用 4x16x32 拓扑时效果最佳。

TPU v4 配置由两组组成，一组芯片小于 64 个芯片（小型拓扑），另一组拓扑大于 64 个芯片（大型拓扑）。

小型 v4 拓扑

Cloud TPU 支持小于 64 个芯片的以下 TPU v4 Pod 切片（4x4x4 立方体）。您可以使用这些小型 v4 拓扑来创建其基于 TensorCore 的名称（例如 v4-32）或其拓扑（例如 2x2x4）：

名称（基于 TensorCore 数量）	条状标签数量	拓扑
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

大型 v4 拓扑

TPU v4 Pod 切片以 64 个芯片为增量提供，形状在所有三个维度上都是 4 的倍数。维度必须按递增顺序排列。下表显示了几个示例。其中一些拓扑是“自定义”拓扑，只能使用 --type 和 --topology 标志创建，因为排列芯片的方式有多种。

名称（基于 TensorCore 数量）	条状标签数量	拓扑
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
不适用 - 必须使用 --type 和 --topology 标志	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
不适用 - 必须使用 --type 和 --topology 标志	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…