在 TPU Pod 上进行训练

概览

TPU 的设计使其可扩容为 TPU Pod。TPU Pod 是一组通过专用高速网络接口连接的 TPU 设备。一个 TPU Pod 最多可以有 2048 个 TPU 核心,允许您跨多个 TPU 板分配处理负载。每个 TPU 板都连接到一个基于 CPU 的高性能主机,用于执行加载和预处理数据等任务。要充分利用大量的 TPU,您必须调整多个训练任务参数。本文档介绍了一些常见问题、您需要在模型中做出的更改以及减少或避免 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 存储分区来训练数据集和检查点。

使用 NFS 进行数据存储

创建 TPU Pod 时,系统会为每个 TPU 节点创建一个单独的虚拟机。默认情况下,系统会为每个 TPU 虚拟机分配一个不同的用户 ID (UID)。如果您尝试从多个节点访问同一 NFS 目录,这会导致问题:同一目录将在不同节点上有不同的所有者,并且标准 Linux 权限不适用于节点。例如,一个节点中的进程无法写入另一个节点创建的日志目录。

如需解决此问题,您可以启用 OS Login。您可以为特定虚拟机实例或项目配置 OS Login。如需了解详情,请参阅设置 OS Login

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

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

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

您的工作负载正常运行并达到性能目标后,纵向扩容到 v2-32 或 v3-32。在验证扩缩能力(功能和性能)的同时,逐渐迭代地增加 TPU 大小,直到达到所需的 TPU 大小。