分布式训练

本页面介绍如何在 Vertex AI 上运行分布式训练作业。

代码要求

使用支持分布式训练的机器学习框架。在训练代码中，您可以使用 CLUSTER_SPEC 或 TF_CONFIG 环境变量来引用训练集群的特定部分。

训练集群的结构

如果您使用 Vertex AI 运行分布式训练作业，则需要在训练集群中指定多个机器（节点）。训练服务为您指定的机器类型分配资源。给定节点上正在运行的作业称为副本。一组具有相同配置的副本被称为一个工作器池。

训练集群中的每个副本在分布式训练中被分配单个角色或任务。例如：

  • 主要副本：您只能将一个副本指定为主要副本。这项任务可管理其他任务，并报告作业的整体状态。

  • 工作器：可以将一个或多个副本指定为工作器。这些副本执行您在作业配置中指定的部分工作。

  • 参数服务器：如果您的机器学习框架支持，您可以将一个或多个副本指定为参数服务器。这些副本会存储模型参数，并协调工作器之间的共享模型状态。

  • 评估程序：如果您的机器学习框架支持，系统可能会将一个或多个副本指定为评估器。这些副本可用于评估您的模型。如果您使用的是 TensorFlow，请注意 TensorFlow 通常要求您只能使用一个评估器。

配置分布式训练作业

您可以通过定义多个工作器池，将任何自定义训练作业配置为分布式训练作业。您还可以在训练流水线或超参数调节作业中运行分布式训练。

如需配置分布式训练作业，请定义您的工作器池 (workerPoolSpecs[]) 列表，并为每种类型设计一种 WorkerPoolSpec：

workerPoolSpecs[] 中的位置 在集群中执行任务
第一 (workerPoolSpecs[0]) 主要、首要、调度程序或“主实例”
第二 (workerPoolSpecs[1]) 辅助、副本、工作器
第三 (workerPoolSpecs[2]) 参数服务器、Reduction Server
第四 (workerPoolSpecs[3]) 评估程序

您必须指定主要副本，该副本协调了所有其他副本完成的工作。仅对第一个副本使用第一个工作器池规范，并将其 replicaCount 设置为 1：

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, primary replica, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {},
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

指定其他工作器池

根据您的机器学习框架，您可以指定其他工作器池以用于其他目的。例如，如果您使用的是 TensorFlow，则可以指定工作器池来配置工作器副本、参数服务器副本和评估程序副本。

您在 workerPoolSpecs[] 列表中指定的工作器池的顺序决定了工作器池的类型。设置您不想使用的工作器池的空值，以便在 workerPoolSpecs[] 列表中跳过这些值，以指定要使用的工作器池。例如：

如要指定仅具有主副本和参数服务器工作器池的作业，则必须为工作器池 1 设置空值：

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

使用 Reduction Server 缩短训练时间

当您使用多个节点训练大型 ML 模型时，节点之间的通信梯度可能会导致显著的延迟。Redution Server 是一种全缩减算法，可以提高分布式训练的吞吐量并减少延迟。Vertex AI 使 Reduction Server 在 Docker 容器映像中可用，您可以在分布式训练期间将其用于您的一个工作池。

如需了解 Reduulation Server 的工作原理，请参阅在 Vertex AI 上使用 Reduulation Server 加快分布式 GPU 训练。

前提条件

如果您满足以下要求，则可以使用 Reduction Server：

  • 您正在使用 GPU 工作器执行分布式训练。

  • 您的训练代码使用 TensorFlow 或 PyTorch，并配置为使用 NCCL 全缩减通过 GPU 进行多主机数据并行训练。（您可能还可以使用其他采用 NCCL 的机器学习框架。）

  • 在主节点 (workerPoolSpecs[0]) 和工作器 (workerPoolSpecs[1]) 上运行的容器支持 Reduction Server。具体而言，每个容器是以下项之一：

    • 预构建的 TensorFlow 训练容器 2.3 版或更高版本。

    • 预构建的 Pytorch 训练容器 1.4 版或更高版本。

    • 安装了 NCCL 2.7 或更高版本以及 google-reduction-server 软件包的自定义容器。您可以通过将以下行添加到 Dockerfile 中，在自定义容器映像上安装此软件包：

      RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list && \
    curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
    apt update && apt install -y google-reduction-server

使用缩减服务器进行训练

如需使用 Reduction Server，请在创建自定义训练资源时执行以下操作：

  1. 在第三个工作器池 (workerPoolSpecs[2]) 的 containerSpec.imageUri 字段中指定以下 URI 之一：

    • us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • europe-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • asia-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

    选择最靠近您执行自定义训练的位置的多区域可缩短延迟时间。

  2. 为第三个工作器池选择机器类型和节点数时，请确保第三个工作器池的总网络带宽与第一个和第二个工作器池的总网络带宽匹配，或超出其总网络带宽。

    如需了解第二个工作器池中每个节点的最大可用带宽，请参阅网络带宽和 GPU。

    对 Reduction Server 节点不使用 GPU。如需了解第三个工作器池中每个节点的最大可用带宽，请参阅通用机器系列中的“最大出站带宽 (Gbps)”列。

    例如，如果将第一个和第二个工作器池配置为使用 5 个 n1-highmem-96 节点，每个节点具有 8 个 NVIDIA_TESLA_V100 GPU，则每个节点的最大可用带宽为 100 Gbps，总带宽为 500 Gbps。为了与第三个工作器池中的带宽匹配，您可以使用 16 个 n1-highcpu-16 节点，每个节点的最大带宽为 32 Gbps，总带宽为 512 Gbps。

    我们建议您对 Reduction Serve 节点使用 n1-highcpu-16 机器类型，因为此机器类型可为其资源提供相对高的带宽。

以下命令举例说明了如何创建使用 Reduction Server 的 CustomJob 资源：

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=4,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highcpu-16,replica-count=16,container-image-uri=us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

如需了解详情，请参阅创建 CustomJob 指南。

使用缩减服务器进行训练的最佳做法

机器类型和数量

在缩减服务器训练中，每个工作器都需要连接到所有缩减器主机。为了最大限度地减少工作器主机上的连接数，请对缩减器主机使用具有最高网络带宽的机器类型。

最好选择这样的缩减器主机：至少具有 16 个 vCPU 且提供 32 Gbps 出站流量带宽的通用 N1/N2 虚拟机，例如 n1-highcpu-16 和 n2-highcpu-16。N1/N2 虚拟机的第 1 层虚拟机带宽将出站流量带宽上限增加到 50 Gbps 到 100 Gbps 之间，因此非常适合缩减器虚拟机节点。

工作器和缩减器的总出站带宽应相同。例如，如果您使用 8 个 a2-megagpu-16g 虚拟机作为工作器，则应至少使用 25 个 n1-highcpu-16 虚拟机作为缩减器。

`(8 worker VMs * 100 Gbps) / 32 Gbps egress = 25 reducer VMs`.

批量处理小型消息

如果聚合的消息足够大，缩减服务器效果最佳。大多数机器学习框架已经使用不同术语提供技术，以在执行全缩减之前对小梯度张量进行批处理。

Horovod

Horovod 支持 Tensor Fusion 批量处理小张量以实现全缩减。张量会填充到融合缓冲区中，直到缓冲区完全填满，此时缓冲区上的全缩减操作将会执行。您可以通过设置 HOROVOD_FUSION_THRESHOLD 环境变量来调整融合缓冲区的大小。

HOROVOD_FUSION_THRESHOLD 环境变量的建议值至少为 128 MB。在这种情况下，请将 HOROVOD_FUSION_THRESHOLD 环境变量设置为 134217728 (128 * 1024 * 1024)。

PyTorch

PyTorch DistributedDataParallel 支持将批量消息作为“梯度分桶”。设置 DistributedDataParallel 构造函数中的 bucket_cap_mb 参数以控制批量存储桶的大小。默认大小为 25 MB。

最佳做法：bucket_cap_mb 的推荐值为 64 (64 MB)。

集群的环境变量

Vertex AI 会在每个副本上填充环境变量 CLUSTER_SPEC，以描述整个集群的设置方式。与 TensorFlow 的 TF_CONFIG 类似，CLUSTER_SPEC 描述集群中的每个副本，包括其索引和角色（主副本、工作器、参数服务器或评估器）。

当您使用 TensorFlow 运行分布式训练时，系统会解析 TF_CONFIG 以构建 tf.train.ClusterSpec。同样，当您使用其他机器学习框架运行分布式训练时，必须解析 CLUSTER_SPEC 以填充框架所需的任何环境变量或设置。

CLUSTER_SPEC 的格式

CLUSTER_SPEC 环境变量是一个 JSON 字符串，格式如下：

键 说明
"cluster"

自定义容器的集群说明。与 TF_CONFIG 一样，此对象也会采用 TensorFlow 集群规范中的格式，且可以传递给 tf.train.ClusterSpec 的构造函数。

集群说明包含您指定的每个工作器池的副本名称列表。
"workerpool0" 所有分布式训练作业在第一个工作器池中都有一个主要副本。
"workerpool1" 如果您在创建作业时指定了这些工作器副本，则此工作器池包含工作器副本。
"workerpool2" 如果您在创建作业时指定了这些参数，则此工作器池包含参数服务器。
"workerpool3" 如果您在创建作业时指定了评估器，则此工作器池包含评估程序。
"environment" 字符串 cloud。
"task" 描述运行代码的特定节点的任务。您可以使用此信息为分布式作业中的特定工作器编写代码。此条目是一个含有以下键的字典：
"type" 此任务运行的工作器池的类型。例如，"workerpool0" 表示主副本。
"index"

任务的索引（从零开始）。例如，如果训练作业包含两个工作器，则其中一个值设为 0，另一个设为 1。
"trial" 当前运行的超参数调节试验的标识符。为作业配置超参数调节时，可以设置多个要训练的试验。此值为您提供了一种区分代码中正在运行的试验的方法。标识符是一个包含试验编号的字符串值，从 1 开始。
job

您创建当前训练作业时提供的 CustomJobSpec，以字典形式表示。

CLUSTER_SPEC 示例

下面是一个示例值：


{
   "cluster":{
      "workerpool0":[
         "cmle-training-workerpool0-ab-0:2222"
      ],
      "workerpool1":[
         "cmle-training-workerpool1-ab-0:2222",
         "cmle-training-workerpool1-ab-1:2222"
      ],
      "workerpool2":[
         "cmle-training-workerpool2-ab-0:2222",
         "cmle-training-workerpool2-ab-1:2222"
      ],
      "workerpool3":[
         "cmle-training-workerpool3-ab-0:2222",
         "cmle-training-workerpool3-ab-1:2222",
         "cmle-training-workerpool3-ab-2:2222"
      ]
   },
   "environment":"cloud",
   "task":{
      "type":"workerpool0",
      "index":0,
      "trial":"TRIAL_ID"
   },
   "job": {
      ...
   }
}

TF_CONFIG 的格式

除了 CLUSTER_SPEC 之外，Vertex AI 在所有分布式训练作业的每个副本上设置 TF_CONFIG 环境变量。Vertex AI 不会为单副本训练作业设置 TF_CONFIG。

CLUSTER_SPEC 和 TF_CONFIG 具有相同的值，但它们的格式不同。这两个环境变量都包含 TensorFlow 所需的其他字段。

使用自定义容器时，使用 TensorFlow 的分布式训练与使用预构建容器相同。

TF_CONFIG 环境变量是一个 JSON 字符串，格式如下：

TF_CONFIG 个字段
cluster

TensorFlow 集群描述。将一个或多个任务名称（chief、worker、ps 或 master）映射到运行这些任务的网络地址列表的字典。对于给定的训练作业，此字典在每个虚拟机上都是相同的。

这是 tf.train.ClusterSpec 构造函数的有效第一个参数。请注意，此字典不会包含 evaluator 作为键，因为即使您将评估程序用于作业，评估程序也不会被视为训练集群的一部分。
task

设置此环境变量的虚拟机的任务说明。 对于给定的训练作业，此字典在每个虚拟机上都是不同的。您可以使用此信息自定义分布式训练作业中每个虚拟机上运行的代码。您还可以使用它来更改训练代码的行为，使其适应超参数调节作业的不同试验。

此字典包含以下键值对：

task 个字段
type

此虚拟机正在执行的任务类型。此值在工作器中设置为 worker，在参数服务器上设置为 ps，而在评估程序上设置为 evaluator。在作业的主工作器中，该值设置为 chief 或 master。
index

任务的索引（从零开始）。例如，如果训练作业包含两个工作器，则其中一个值设为 0，另一个设为 1。
trial

当前在此虚拟机上运行的超参数调节试验的 ID。仅在当前训练作业是超参数调节作业时才设置此字段。

对于超参数调节作业，Vertex AI 会在许多试验中反复运行您的训练代码，每次使用不同的超参数。此字段包含当前的试验编号，从第一次试验的 1 开始。
cloud

Vertex AI 内部使用的 ID。您可以忽略该字段。
job

您创建当前训练作业时提供的 CustomJobSpec，以字典形式表示。
environment

字符串 cloud。

TF_CONFIG 示例

以下示例代码将 TF_CONFIG 环境变量打印到您的训练日志：

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

在运行时版本 2.1 或更高版本中运行的超参数调节作业使用一个主工作器、两个工作器和一个参数服务器，此代码在该作业中在第一次超参数调节试验期间为其中一个工作器生成以下日志。示例输出为了简洁会隐藏 job 字段，并用一般值替换某些 ID。

{
  "cluster": {
    "chief": [
      "training-workerpool0-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "training-workerpool2-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "training-workerpool1-[ID_STRING_1]-0:2222",
      "training-workerpool1-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

何时使用 TF_CONFIG

仅为分布式训练作业设置 TF_CONFIG。

您可能不需要直接在训练代码中与 TF_CONFIG 环境变量互动。仅在 TensorFlow 的分布策略和 Vertex AI 的标准超参数调优工作流（以下部分将对这两者进行介绍）都不适用于您作业的情况下，才应访问 TF_CONFIG 环境变量。

分布式训练

Vertex AI 设置 TF_CONFIG 环境变量来扩展 TensorFlow 进行分布式训练所需的规范。

如要使用 TensorFlow 执行分布式训练，请使用 tf.distribute.Strategy API。我们尤其建议您将 Keras API 与 MultiWorkerMirroredStrategy 结合使用，或者，如果您为作业指定参数服务器，那么将 Keras API 与 ParameterServerStrategy 结合使用。不过请注意，TensorFlow 仅为这些策略提供实验性支持。

这些分布策略使用 TF_CONFIG 环境变量为训练作业中的每个虚拟机分配角色，并促进虚拟机之间的通信。您无需直接在训练代码中访问 TF_CONFIG 环境变量，因为 TensorFlow 会为您处理。

如果要自定义运行训练作业的不同虚拟机的行为方式，仅需直接解析 TF_CONFIG 环境变量。

超参数调节

当您运行超参数调节作业时，Vertex AI 会为每次试验的训练代码提供不同的参数。您的训练代码不必知道当前正在运行的试验。此外，您可以在 Google Cloud 控制台中监控超参数调优作业的进度。

如果需要，您的代码可以从 TF_CONFIG 环境变量的 task 字段的 trial 字段中读取当前的试验编号。

后续步骤