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 瓦
TPU Pod 大小	4096 条状标签
互连拓扑	3D 网格
每个 Pod 的计算峰值	1.1 exaflops（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 网格配置的效果要好。如需更多信息 请参阅扭曲 Tori 拓扑。

与 v3 相比，TPU v4 的性能优势

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

内存系统

非统一内存访问 (NUMA) 是一种机器的计算机内存架构 并支持多个 CPU 的计算方法每个 CPU 都可以直接访问高速块 内存。CPU 及其内存称为 NUMA 节点。NUMA 个节点已连接到 NUMA 个节点直接相邻。来自一个 NUMA 节点的 CPU 可以访问另一个 NUMA 节点中的内存，但是这种访问比访问 NUMA 节点内的内存空间。

在多 CPU 计算机上运行的软件可以将 CPU 所需的数据 NUMA 节点，从而增加内存吞吐量。有关 NUMA 的更多信息， 请参阅非统一内存访问 。

您可以通过绑定训练脚本来利用 NUMA 地区的优势 更改为 NUMA 节点 0。

如需启用 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 性能配置文件 5120 亿个粒度。

TensorCores

• MXU 数量加倍，时钟频率更高，TFLOPS 最高为 275。
• 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。配置 v4 TPU 请使用 AcceleratorType，在创建--accelerator-type TPU Pod 切片。将 --accelerator-type 设置为包含 TPU 的字符串 版本和您要使用的 TensorCore 数量。例如，要创建 具有 32 个 TensorCore 的 v4 Pod 切片，您将使用 --accelerator-type=v4-32。

以下命令使用 --accelerator-type 标志：

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

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

使用 AcceleratorConfig

如果您想自定义物理拓扑，请使用 AcceleratorConfig 使用 TPU 切片使用 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 这三个维度的条状标签数量都相同。例如，创建一个 v4 Pod 对于具有 16 个条状标签的 Slice，您需要设置 --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.17.0-pod-pjrt

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

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

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

扭曲 Tori 拓扑

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

对于将数据并行用作唯一并行策略的工作负载，扭曲了 但拓扑的性能可能略好一些。对于 LLM，使用扭曲的 拓扑可能因并行类型（DP、MP 等）而异。最佳 使用和不使用扭曲拓扑来训练 LLM 这能为模型提供最佳性能。我们针对 FSDP MaxText 模型 使用扭曲拓扑改进了 1-2 个 MFU。

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

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

这些优势最终会转化为许多全球 通信模式。

TPU 软件支持在切片上扭曲的 Tori，每个切片的大小 尺寸等于或小于最小尺寸的两倍。对于 例如 4x4x8、4×8×8 或 12x12x24。

以这个 4×2 环形拓扑为例，其中 TPU 标有 (X,Y) 坐标：

为清晰起见，此拓扑图中的边显示为无向边。在 每个边缘都是 TPU 之间的双向连接。我们指的是 将网格的一侧和另一侧之间的边设为包围边， 如图所示。

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

上图与上图之间唯一不同的是，Y 环绕边缘不连接到具有相同 X 的另一个 TPU 坐标，它们已改为使用 Coordinate 连接到 TPU X+2 mod 4。

相同的概念可以泛化为不同的维度大小和数量不同的 维度。只要每个维度都是 等于或大于最小尺寸的两倍。

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

下表显示了支持的扭曲拓扑和一个理论 与非扭曲拓扑相比，采用双切法的拓扑可增加 2 倍带宽。

扭曲拓扑	理论上对分率的升幅 带宽与非扭曲环面
4×4×8（扭曲型）	约 70%
8x8x16_扭曲
12 × 12 × 24 扭曲
4×8×8（扭曲型）	约 40%
8×16×16（扭曲型）

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

TPU v4 拓扑变体

一些包含相同芯片数量的拓扑可以分布在不同的 方法。例如，具有 512 个芯片 (1024 TensorCore) 的 TPU Pod 切片可以 使用以下拓扑配置：4x4x32、4x8x16 或 8x8x8。TPU Pod 切片 配备 2048 个芯片 (4096 TensorCore)，可提供更多拓扑选项：4x4x128、 4x8x64、4x16x32 和 8x16x16。

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

具有多个并行维度的模型在 并行处理维度映射到 TPU 拓扑维度。这些通常是 和模型并行大语言模型 (LLM)。例如，对于 TPU v4 拓扑为 8x16x16 的 Pod 切片，TPU 拓扑维度为 8、16 和 16。时间是 使用 8 向或 16 向模型并行处理（映射到 物理 TPU 拓扑维度）。四向模型并行处理将不是最佳的 因为这个拓扑与任何 TPU 拓扑维度都不一致 但在相同数量的芯片上采用 4x16x32 拓扑会是最佳选择。

TPU v4 配置由两组组成，其中一组的拓扑小于 64 个芯片（小型拓扑）和拓扑超过 64 个芯片的那些 （大型拓扑）。

小型 v4 拓扑

Cloud TPU 支持以下小于 64 个芯片的 TPU v4 Pod 切片： 一个 4x4x4 立方体。您可以使用 基于 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
…	…	…