TPU v3

本文档介绍了 Cloud TPU v3。

系统架构

每个 v3 TPU 芯片包含两个 TensorCore。每个 TensorCore 都有两个矩阵乘法单元 (MXU),一个 矢量单位和标量单位。下表显示了关键规范 及其值。

密钥规范 v3 Pod 值
每个芯片的计算峰值 每秒 123 万亿次浮点运算 (bf16)
HBM2 容量和带宽 32 GiB,900 GBps
测量的最小/平均/最大功率 123/220/262 瓦特
TPU Pod 大小 1024 条状标签
互连拓扑 二维环面
每个 Pod 的计算峰值 每秒 126 千万亿次浮点运算 (bf16)
每个 Pod 完全减少带宽 340 TB/秒
每个 Pod 的对分带宽 6.4 TB/秒

下图展示了 TPU v3 芯片。

TPU v3 芯片示意图

TPU v3 的架构详情和性能特性,请参阅 用于训练深度神经网络的领域特定超级计算机

TPU v3 相对于 v2 的性能优势

在 TPU v3 配置中提高了每个 TensorCore 的 FLOPS 和内存容量 可以通过以下方式提高模型的性能:

  • TPU v3 配置可显著提升 TensorCore,用于计算受限的模型。TPU v2 上的内存受限模型 如果实际运行这些配置 在 TPU v3 配置上也受内存限制。

  • 在采用 TPU v2 配置时,如果内存无法容纳数据,TPU v3 可以提高性能, 中间值(再具体化)。

  • TPU v3 配置可以运行批次大小不适合 TPU v2 配置的新模型。例如,TPU v3 可能允许使用更深的 ResNet 模型, 处理大型图片。

因训练步骤等待输入而在 TPU v2 上几乎成为受限于输入(“馈入”)的模型,在 Cloud TPU v3 中也可能会受限于输入。通过 流水线性能指南可帮助您解决信息流问题。

配置

TPU v3 Pod 由 1024 个芯片组成,这些芯片通过高速链路互连。接收者 创建 TPU v3 设备或 Pod 切片,使用 --accelerator-type 标志(在 TPU 创建命令 (gcloud compute tpus tpu-vm) 中)。您可以通过指定 TPU 版本和 TPU 核心数。例如,对于单个 v3 TPU,请使用 --accelerator-type=v3-8.对于具有 128 个 TensorCore 的 v3 Pod 切片,请使用 --accelerator-type=v3-128.

以下命令展示了如何创建具有 128 个 TensorCore 的 v3 TPU Pod 切片:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v3-128 \
    --version=tpu-vm-tf-2.17.0-pjrt

下表列出了支持的 v3 TPU 类型:

TPU 版本 支持结束
v3-8 (结束日期尚未确定)
v3-32 (结束日期尚未确定)
v3-128 (结束日期尚未确定)
v3-256 (结束日期尚未确定)
v3-512 (结束日期尚未确定)
v3-1024 (结束日期尚未确定)
v3-2048 (结束日期尚未确定)

如需详细了解如何管理 TPU,请参阅管理 TPU。 如需详细了解 Cloud TPU 的系统架构,请参阅 系统架构