TPU7x (Ironwood)

本页面介绍了 TPU7x 的架构和可用配置。TPU7x 是目前 Google Cloud上提供的最新一代 TPU。 作为 Ironwood 系列的首款力作，TPU7x 是 Google Cloud推出的第七代 TPU。Ironwood 系列专为大规模 AI 训练和推理而设计。

TPU7x 的每个 Pod 拥有 9,216 个芯片，在诸多方面与 TPU v5p 一脉相承。TPU7x 可为大规模密集模型和 MoE 模型、预训练、采样以及解码密集型推理提供卓越的性能表现。

您可以结合 Google Kubernetes Engine (GKE) 或 Compute Engine 来使用 TPU7x。如需详细了解如何在 GKE 中使用 TPU，请参阅关于 GKE 中的 TPU。

您还可以通过 TPU Cluster Director 协同使用 TPU7x 与 GKE。TPU Cluster Director 通过“全容量”模式预留提供。该模式赋予您对所有预留容量的完全访问权限（无任何保留），并让您能够全方位洞察 TPU 硬件拓扑、利用率状态以及健康状况。如需了解详情，请参阅全容量模式概览。

系统架构

每个 TPU7x 芯片包含两个 TensorCore 和四个 SparseCore。下表列出了 TPU7x 的核心规格及相应参数值，并将其与前几代产品进行了对比。

规格 v5p v6e (Trillium) TPU7x (Ironwood)
每个 Pod 的芯片数 8960 256 9216
单芯片峰值算力 (BF16) (TFLOPs) 459 918 2307
单芯片峰值算力 (FP8) (TFLOPs) 459 918 4614
单芯片 HBM 容量 (GiB) 95 32 192
单芯片 HBM 带宽 (GBps) 2765 1638 7380
vCPU 数量（4 芯片虚拟机） 208 180 224
RAM (GB)（4 芯片虚拟机） 448 720 960
单芯片 TensorCore 数量 2 1 2
单芯片 SparseCore 数量 4 2 4
单芯片双向芯片间互连 (ICI) 带宽 (GBps) 1200 800 1200
单芯片数据中心网络 (DCN) 带宽 (Gbps) 50 100 100

下图展示了 Ironwood 的架构：

Ironwood 架构图

内存层次结构

TPU7x 采用多层级内存系统，管理这些层级之间的数据移动对于性能至关重要：

  • 高带宽内存 (HBM)：每个芯片都配备了 192 GB 的 HBM，带宽约为 7.37 TB/s。充足的 HBM 容量可支持更大的批次大小，从而有效提升吞吐量。然而，尽管 HBM 容量充足，但在面对受内存限制的向量运算或低效的数据访问模式时，它仍可能成为性能瓶颈。
  • 向量内存 (VMEM)：VMEM 是一种体积更小、集成在芯片上的 SRAM（静态随机存取存储器），其连接至矩阵乘法单元 (MXU) 的带宽远高于 HBM。此内存可充当自定义内核的高速暂存器。此缓冲区的大小是一个可调参数。优化缓冲区空间对于调优自定义 Pallas 内核至关重要，因为这些内核的块大小通常受限于可用的 VMEM。
  • 主机内存和 PCIe：每四个 TPU 芯片为一组，通过 PCIe 网络连接到 CPU 主机。虽然此连接的带宽远低于 HBM，但您可以利用主机的主内存来分流激活或优化器状态，从而腾出 HBM 空间。在处理大型模型时，这项技术对于缓解内存压力尤为奏效。

如需详细了解如何高效管理 TPU7x 内存层次结构各层级之间的数据移动，请参阅 Ironwood 性能优化。

双芯粒架构

借助 Ironwood 编程模型，您可以调用两个 TPU 芯粒，而不再受限于前几代产品（TPU v4 和 v5p）所采用的单逻辑核心（也称为 MegaCore）架构。这一改进提升了芯片制造的成本效益与生产效率。虽然这标志着架构上的重大飞跃，但全新的设计方案确保了您只需进行微调，即可沿用现有的软件模型。

Ironwood TPU 由两个独立的芯粒组成，每个芯粒都拥有专属的内存空间。这与 MegaCore 架构所采用的统一内存空间设计有所不同。

  • 芯粒构成：每个芯粒都是一个自包含单元，集成了 1 个 TensorCore、2 个 SparseCore 以及 96 GB 的高带宽内存 (HBM)。

  • 高速互连：这两个芯粒通过裸片间 (D2D) 接口实现互连，其通信速度比一维芯片间互连 (ICI) 链路快 6 倍。芯粒间通信通过集体操作进行管理。

编程模型与框架呈现

Ironwood 的编程模型与 v4 之前的 TPU 系列产品（例如 TPU v3）类似。这种全新的架构通过以下方式呈现：

  • 单芯片双设备：JAX 等框架会将每个 Ironwood 芯片呈现为两个独立的“设备”，分别对应其中的一个芯粒。

  • 芯粒指定：您可以明确指定将哪个芯粒用于计算。JAX 在拓扑规范中引入了第四个维度，以便精准区分不同的芯粒。这种设计方案让您只需进行微调，即可复用现有的软件模型。

如需详细了解如何充分发挥双芯粒架构的极致性能，请参阅针对 Ironwood 双芯粒架构的性能建议

支持的配置

TPU7x 芯片在三个维度上均与最邻近的芯片直接相连，从而构建起一个 3D 网状网络连接体系。当芯片数量超过 64 个时，切片将由一个或多个 4x4x4 的芯片“立方体”构建而成。

TPU7x 芯片采用 3D 环面互连拓扑结构。借助这种拓扑结构，切片可纵向扩容至多达 9216 个芯片。该拓扑结构可为 Pod 内的芯片间通信提供每轴 200 GBps 的双向带宽。

下表列出了 TPU7x 支持的常见 3D 切片形状：

拓扑 TPU 芯片数 主机数 虚拟机数 立方体数 范围
2x2x1 4 1 1 1/16 单主机
2x2x2 8 2 2 1/8 多主机
2x2x4 16 4 4 1/4 多主机
2x4x4 32 8 8 1/2 多主机
4x4x4 64 16 16 1 多主机
4x4x8 128 32 32 2 多主机
4x8x8 256 64 64 4 多主机
8x8x8 512 128 128 8 多主机
8x8x16 1024 256 256 16 多主机
8x16x16 2048 512 512 32 多主机

TPU7x 虚拟机

每台 TPU7x 虚拟机包含 4 个芯片。每台虚拟机均可访问两个 NUMA 节点。如需深入了解 NUMA 节点，请参阅维基百科上的非统一内存访问条目。

所有 TPU7x 切片均采用整机配置，由 4 芯片虚拟机组成。TPU7x 虚拟机的技术规格如下：

  • 单台虚拟机的 vCPU 数量：224
  • 单台虚拟机的 RAM：960 GB
  • 单台虚拟机的 NUMA 节点数：2

Hyperdisk

默认情况下，TPU7x 的虚拟机启动磁盘为 Hyperdisk Balanced。您还可以为 TPU 虚拟机挂接更多磁盘，以获取额外的存储空间。TPU7x 支持以下磁盘类型：

  • Hyperdisk Balanced
  • Hyperdisk ML

如需详细了解 Hyperdisk，请参阅 Hyperdisk 概览。如需详细了解 Cloud TPU 的存储方案，请参阅 Cloud TPU 数据的存储选项。

后续步骤