TPU v4

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

系统架构

每个 TPU v4 芯片包含两个 TensorCore。每个 TensorCore 都有四个矩阵乘法单元 (MXU)、一个矢量单元和一个标量单元。下表显示了关键规范 运行 25 万个 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 的 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 网格配置。如需更多信息 请参阅扭曲 Tori 拓扑。

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 的最大性能。
• 2 倍转置和排列带宽。
• 通用内存的加载存储内存访问模型 (Cmem)。
• 更快的 MXU 权重加载带宽和 8 位模式支持 并缩短了推理延迟时间。

芯片间互连

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

其他

• 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.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 切片或 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 配置。

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

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

注意 ：除非后缀 _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。

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

具有多个并行维度的模型在将其并行维度映射到 TPU 拓扑维度后，性能会达到最佳。这些通常是数据和模型并行的大语言模型 (LLM)。例如，对于拓扑为 8x16x16 的 TPU v4 PodSlice，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
…	…	…