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 exaflops (bf16 或 int8)
每个 Pod 的 All-reduce 带宽	1.1 PB/s
每个 Pod 的对分带宽	24 TB/s

下图展示了 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 节点 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 工作负载（例如图片分类、推荐工作负载），无论使用哪种框架。
• 如果您使用的 TPU 运行时版本不带 -pod 后缀（例如 tpu-vm-tf-2.10.0-v4）。

其他内存系统差异：

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

TensorCores

• MXU 数量翻了一番，时钟速率更高，可提供 275 TFLOPS 的最大性能。
• 转置和排列带宽翻倍。
• 通用内存 (Cmem) 的加载-存储内存访问模型。
• 更快的 MXU 权重加载带宽和 8 位模式支持，可缩减批处理大小并缩短推理延迟时间。

芯片间互连

每个芯片有 6 个互连链接，可实现网络直径较小的网络拓扑。

其他

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

配置

TPU v4 Pod 由 4096 个芯片组成，这些芯片通过可重配置的高速链路相互连接。借助 TPU v4 的灵活网络功能，您可以通过多种方式连接同一大小的 Pod 切片中的芯片。创建 TPU Podslice 时，您可以指定所需的 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.18.0-pod-pjrt

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

使用 AcceleratorConfig

如果您想自定义 TPU 切片的物理拓扑，请使用 AcceleratorConfig。通常，如果 Pod 分片大于 256 个芯片，则需要执行此操作才能进行性能调整。

如需使用 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 slice，其中 128 个 TPU 芯片以 4x4x8 阵列的形式排列：

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.18.0-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 模型上进行的一些实验表明，使用扭曲拓扑结构可将 MFU 提高 1-2 个。

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

• 改进了负载均衡
• 更高的二分带宽
• 数据包路线更短

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

TPU 软件支持在切片上使用扭曲的环面，其中每个维度的大小等于或为最小维度的两倍。例如，4x4x8、4×8×8 或 12x12x24。

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

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

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

与上一个示意图相比，此示意图中唯一的变化是 Y 环绕边缘。它们已改为连接到坐标为 X+2 mod 4 的 TPU，而不是连接到具有相同 X 坐标的其他 TPU。

同样的想法可以推广到不同的维度大小和不同的维度数量。只要每个维度等于或为最小维度的两倍，生成的网络就是对称的。

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

下表显示了支持的扭绞型拓扑以及与非扭绞型拓扑相比，使用这些拓扑时分片带宽的理论增幅。

扭曲拓扑	与非扭曲环面相比，分割带宽的理论增幅
4×4×8_twisted	约 70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

注意： 除非使用 --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 PodSlice，TPU 拓扑维度为 8、16 和 16。使用 8 路或 16 路模型并行处理（映射到物理 TPU 拓扑维度之一）可以提高性能。在这种拓扑中，4 路模型并行性不太理想，因为它与任何 TPU 拓扑维度都不一致，但在相同数量的芯片上，4x16x32 拓扑是最佳选择。

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

小型 v4 拓扑

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

名称（基于 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
…	…	…