为 PyTorch 配置分布式训练

本文档介绍如何创建分布式 PyTorch 训练作业。当您创建分布式训练作业时，AI Platform Training 会在具有支持分布式 PyTorch 训练的环境变量的虚拟机实例（也称为节点）集群上运行代码。这有助于您的训练作业扩容以处理大量数据。

本指南假设您要使用预构建的 PyTorch 容器进行训练，如 PyTorch 使用入门中所述。调整 PyTorch 代码以用于分布式训练只需最少量的更改。

指定训练集群结构

对于分布式 PyTorch 训练，请将您的作业配置为使用一个主工作器节点以及一个或多个工作器节点。这些角色具有以下行为：

  • 主工作器：rank 为 0 的虚拟机。此节点在集群中的节点之间建立连接。
  • 工作器：集群中的其余节点。根据训练应用代码的指定，每个节点进行一部分训练。

如需了解如何为训练集群指定主工作器节点和工作器节点，请参阅指定机器类型或容量层级。

指定容器映像

创建训练作业时，请在 trainingInput.masterConfig.imageUri 字段中为主工作器指定要使用的 Docker 容器映像，并在 trainingInput.workerConfig.imageUri 字段中为每个工作器指定 Docker 容器映像。请参阅预构建的 PyTorch 容器列表。

如果您使用 gcloud ai-platform jobs submit training 命令创建训练作业，则可以使用 --master-image-uri 和 --worker-image-uri 标志指定这些字段。

但是，如果您未指定 trainingInput.workerConfig.imageUri 字段，则其值默认为 trainingInput.masterConfig.imageUri 的值。在每个节点上使用相同的预构建 PyTorch 容器通常是一个合理的选择。

更新训练代码

在训练应用中添加以下代码以初始化训练集群：

import torch

torch.distributed.init_process_group(
    backend='BACKEND',
    init_method='env://'
)

将 BACKEND 替换为下一部分中介绍的受支持的分布式训练后端。init_method='env://' 关键字参数指示 PyTorch 使用环境变量来初始化集群中的通信。如需了解详情，请参阅本指南的环境变量部分。

此外，请更新训练代码以使用 torch.nn.parallel.DistributedDataParallel 类。例如，如果您已在代码中创建了名为 model 的 PyTorch 模块，请添加以下行：

model = torch.nn.parallel.DistributedDataParallel(model)

如需详细了解如何配置分布式训练，请参阅 PyTorch 文档的分布式训练指南。

分布式训练后端

对于分布式 PyTorch 训练，AI Platform Training 支持以下后端：

  • gloo：建议用于 CPU 训练作业
  • nccl：建议用于 GPU 训练作业

了解后端之间的差异。

环境变量

当您创建分布式 PyTorch 训练作业时，AI Platform Training 会在每个节点上设置以下环境变量：

  • WORLD_SIZE：集群中的总节点数。此变量在每个节点上的值都相同。
  • RANK：每个节点的唯一标识符。在主工作器上，此变量设置为 0。它在每个工作器上的值都不同，范围从 1 到 WORLD_SIZE - 1。
  • MASTER_ADDR：主工作器节点的主机名。此变量在每个节点上的值都相同。
  • MASTER_PORT：主工作器节点进行通信的端口。此变量在每个节点上的值都相同。

PyTorch 使用这些环境变量来初始化集群。

后续步骤