TPU 类型和拓扑
TPU 配置
创建 TPU 时,您需要指定 TPU 配置。您可以基于 TensorCore 或 TPU 芯片指定 TPU 配置。只有在训练大型模型并调整其性能的高级场景中,才需要使用芯片指定 TPU 配置。
基于 TensorCore 的配置
您可以将基于 TensorCore 的配置与 v2 - v4 TPU 搭配使用。请使用 TPU 版本和您希望要使用的 TensorCore 数量来描述基于 TensorCore 的配置。例如,v4-128
指定了使用 v4 TPU 和 128 TensorCore 的配置。您可以使用 gcloud compute tpus tpu-vm create 命令的 --accelerator-type
标志指定基于 TensorCore 的配置:
$ gcloud compute tpus tpu-vm create tpu-name --zone=zone --accelerator-type=v3-128 --version=tpu-vm-tf-2.12.0
基于芯片的 TPU 配置
您可以将基于 TPU 芯片的配置与 v4 TPU 和更高 TPU 搭配使用。您使用 TPU 版本和拓扑(物理芯片的 TPU 排列方式)描述基于芯片的配置。您可以使用类似于“v4
”的字符串来指定 TPU 版本。您可以使用 3 个元组指定拓扑,描述芯片的三维排列方式。您可以使用 gcloud compute tpus tpu-vm create 命令的 --version
和 --topology
标志指定基于芯片的配置。例如,以下命令会将 TPU v4 指定为 TPU 版本,并将 2x2x4
(16 个芯片,每个芯片 2 个 TensorCore)指定为拓扑:
$ gcloud compute tpus tpu-vm create tpu-name --zone=zone --type=v4 --topology=2x2x4 --version=tpu-vm-tf-2.12.0
如需详细了解不同的 Cloud TPU 版本和拓扑,请参阅系统架构。
v4 拓扑的表示法为 AxBxC
,其中 A<=B<=C,且 A、B、C 均为 <= 4 或 4 的所有整数倍数。值 A、B 和 C 是这三个维度中的每个条状标签的数量。2A=B=C 或 2A=2B=C 的拓扑还包含已针对全包通信进行优化的拓扑变体,例如 4×4×8、8×8×16 和 12×12×24。
如需详细了解 v4 拓扑,请参阅 TPU v4 配置。
拓扑变体
与给定芯片组关联的默认拓扑是与立方体最相似的拓扑(请参阅拓扑形状)。此形状可能是数据并行机器学习训练的最佳选择。其他拓扑适用于存在多种并行性(例如,模型和数据并行处理或模拟空间分区)的工作负载。如果拓扑与所用的并行性相匹配,则这些工作负载的效果最佳。例如,将 4 向模型并行处理放置在 X 维度和 256 路数据并行处理中,Y 维度和 Z 维度均与 4x16x16 拓扑相匹配。
具有多个并行维度的模型将并行性维度映射到 TPU 拓扑维度时效果最佳。这些模型通常是数据 + 模型并行大型语言模型 (LLM)。例如,对于拓扑为 8x16x16 的 TPU v4 pod 切片,TPU 拓扑尺寸为 8、16 和 16。使用 8 行或 16 型模型并行处理(映射到其中一个物理 TPU 拓扑维度)性能更高。四向模型并行处理在这种拓扑中并不理想,因为它与任何 TPU 拓扑尺寸都不一致,但对于具有相同数量的芯片的 4x16x32 拓扑,效果最佳。
小型 v4 拓扑
Cloud TPU 支持以下小于 64 个芯片和一个 4x4x4 立方体的 TPU v4 切片。 您可以使用基于 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 切片以 64 个芯片为增量提供,形状是所有三个维度的 4 的倍数。维度必须按升序排列。下表显示了一些示例。其中一些拓扑是只能使用拓扑 API 启动的“自定义”拓扑,因为它们与芯片数量相同,其芯片所用名称更常用。
名称(基于 TensorCore 计数) | 条状标签 | 拓扑 |
v4-128 | 64 | 4x4x4 |
v4-256 | 128 | 4x4x8 |
v4-512 | 256 | 4x8x8 |
请参阅 Topology API。 | 256 | 4x4x16 |
v4-1024 | 512 | 8x8x8 |
v4-1536 | 768 | 8x8x12 |
v4-2048 | 1,024 | 8x8x16 |
请参阅 Topology API。 | 1,024 | 4x16x16 |
v4-4096 | 2048 年 | 8x16x16 |
... | ... | ... |
Topology API
为了使用自定义拓扑创建 Cloud TPU Pod 切片,gcloud
TPU API 可以按如下方式使用:
$ gcloud compute tpus tpu-vm create tpu-name \ --zone=us-central2-b \ --subnetwork=tpusubnet \ --type=v4 \ --topology=4x4x16 \ --version=runtime-version
拓扑形状
下图显示了三种可能的拓扑形状及其关联的 topology
标志设置。
TPU 类型兼容性
您可以将 TPU 类型更改为具有相同 TensorCore 或条状标签数量的其他 TPU 类型(例如 v3-128
和 v4-128
),并运行训练脚本而无需更改代码。但是,如果您更改为具有更多或更少数量的芯片或 TensorCore 的 TPU 类型,则需要执行重要的调整和优化。如需了解详情,请参阅在 TPU Pod 上运行训练。