在 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 大小。