在 TPU Pod 上训练

概览

TPU 的设计使其可扩容为 TPU Pod。TPU Pod 是一组通过专用高速网络接口连接的 TPU 设备。借助 TPU Pod，您可以跨多个 TPU 分配处理负载。每个 TPU 板都连接到一个基于 CPU 的高性能主机，用于执行加载和预处理数据等任务。要充分利用大量的 TPU，您必须调整多个训练任务参数。

使用 TPU Pod 进行训练的设置因框架而异。请使用以下链接详细了解如何在 Pod 上使用各个框架进行训练：

• TensorFlow

• PyTorch

• JAX

以下部分介绍了一些常见问题、您需要在模型中做出的更改以及减少或避免 Pod 故障的最佳实践。

扩缩批量大小和训练步数

要在更大的 TPU 类型上实现线性扩缩，请保持每核心批量大小相同。

例如，如果您在 v2-8 上使用 1024 作为批量大小，则在 v2-32 上使用 4096 (4 * 1024) 作为批量大小。这充分利用了 TPU 硬件。您可以使用较小的批次大小，但训练不会线性扩缩。

某些模型包含 train_steps 标志，其中一步对应处理单个批次的数据。增加批次大小时，请缩减训练步数，以使训练样本总数保持不变。

例如，如果对于 100 步，您的批次大小为 1000，则在训练期间会处理 10 万个样本。如果您现在有 4 个工作器且有效批次大小为 4000，则必须将步数调整为 25，才能处理这 10 万个样本。如果模型使用 epochs 标志，则无需扩缩步数。

较大的批次大小可以改变模型的收敛行为，因此您可能还需要调整一些超参数，例如学习速率。

使用 TPU Pod 所在区域中的区域 Google Cloud Storage 存储分区

一般而言，TPU 训练的最佳做法是始终使用同一区域中的资源。使用 TPU Pod 时，资源区域特别重要，因为当您的 Google Cloud Storage 存储桶和 TPU 位于同一区域时，数据转移的速率会更高。

确保您使用 TPU 所在区域中的区域 Google Cloud Storage 存储分区来训练数据集和检查点。

在 TPU Pod 上进行开发的工作流最佳做法

在开发新的 TPU 工作负载时，通常最好在最小的 TPU 上开始开发，并逐渐迭代为较大的 TPU 大小。首先使用较小的 TPU 版本（例如 v2-8 或 v3-8）。

• 测试工作负载的功能
• 使用性能工具测试和验证性能

您的工作负载正常运行并达到性能目标后，纵向扩容到更大的 TPU 类型，例如 v2-32 或 v3-32。在验证可伸缩性（功能和性能）的同时，逐渐迭代地增加 TPU 大小，直到达到所需的 TPU 大小。