分布式 GPU 训练指南 (SDK v1)

适用于：Python SDK azureml v1

详细了解如何在 Azure 机器学习 (ML) 中使用分布式 GPU 训练代码。 本文不讲解分布式训练， 而是帮助用户在 Azure 机器学习上运行现有的分布式训练代码。 本文针对每种框架提供了提示和示例以供遵照：

• 消息传递接口 (MPI)
  • Horovod
  • DeepSpeed
  • Open MPI 中的环境变量
• PyTorch
  • 进程组初始化
  • 启动选项
  • DistributedDataParallel（每进程启动）
  • 使用 torch.distributed.launch（每节点启动）
  • PyTorch Lightning
  • Hugging Face Transformers
• TensorFlow
  • TensorFlow 的环境变量 (TF_CONFIG)
• 使用 InfiniBand 加速 GPU 训练

必备条件

回顾这些分布式 GPU 训练的基本概念，例如数据并行性、分布式数据并行性和模型并行性。

MPI

Azure 机器学习提供了一个 MPI 作业，用于在每个节点中启动给定数量的进程。 可以采用此方法，使用每进程启动器或每节点启动器运行分布式训练，具体取决于是对于每节点启动器 process_count_per_node 设置为 1（默认值），还是对于每进程启动器该值等于设备/GPU 的数量。 Azure 机器学习在后台构造完整的 MPI 启动命令 (mpirun)。 用户无法提供自己的完整头节点启动器命令，如 mpirun 或 DeepSpeed launcher。

若要使用 MPI 运行分布式训练，请执行以下步骤：

1. 将 Azure 机器学习环境与首选深度学习框架和 MPI 一起使用。 Azure 机器学习针对热门框架提供了特选环境。
2. 使用 process_count_per_node 和 node_count 定义 MpiConfiguration。 如果用户脚本将负责每节点启动进程，则对于每进程启动 process_count_per_node 应等于每节点的 GPU 数，或者对于每节点启动则设置为 1（默认值）。
3. 将 MpiConfiguration 对象传递到 ScriptRunConfig 的 distributed_job_config 参数。

from azureml.core import Workspace, ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import MpiConfiguration

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = MpiConfiguration(process_count_per_node=4, node_count=2)

run_config = ScriptRunConfig(
  source_directory= './src',
  script='train.py',
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

# submit the run configuration to start the job
run = Experiment(ws, "experiment_name").submit(run_config)

Horovod

当使用 Horovod 通过深度学习框架进行分布式训练时，请使用 MPI 作业配置。

确保代码遵照以下提示：

• 在添加 Azure 机器学习部件之前，已使用 Horovod 正确检测了训练代码
• Azure 机器学习环境包含 Horovod 和 MPI。 PyTorch 和 TensorFlow 特选 GPU 环境预配置了 Horovod 及其依赖项。
• 使用所需的分布创建 MpiConfiguration。

Horovod 示例

• azureml-examples：使用 Horovod 的 TensorFlow 分布式训练

DeepSpeed

不要使用 DeepSpeed 的自定义启动器在 Azure 机器学习上使用 DeepSpeed 库运行分布式训练。 相反，请配置 MPI 作业以使用 MPI 启动训练作业。

确保代码遵照以下提示：

• Azure 机器学习环境包含 DeepSpeed 及其依赖项、Open MPI 和 mpi4py。
• 使用分布创建 MpiConfiguration。

DeepSpeed 示例

• azureml-examples：CIFAR-10 上使用 DeepSpeed 的分布式训练

Open MPI 中的环境变量

使用 Open MPI 映像运行 MPI 作业时，启动每个进程的以下环境变量：

1. OMPI_COMM_WORLD_RANK - 进程的排名
2. OMPI_COMM_WORLD_SIZE - 世界大小
3. AZ_BATCH_MASTER_NODE - 主地址，含端口，MASTER_ADDR:MASTER_PORT
4. OMPI_COMM_WORLD_LOCAL_RANK - 节点上该进程的本地排名
5. OMPI_COMM_WORLD_LOCAL_SIZE - 节点上的进程数

PyTorch

Azure 机器学习支持使用 PyTorch 的本机分布式训练功能来运行分布式作业（torch.distributed）。

进程组初始化

任何分布式训练的主干都基于一组彼此知晓且可以使用后端相互通信的进程。 对于 PyTorch，通过在所有分布式进程（共同构成一个进程组）中调用 torch.distributed.init_process_group 来创建进程组。

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

最常用的通信后端是 mpi、nccl 和 gloo。 对于基于 GPU 的训练，建议使用 nccl 以获得最佳性能，并且应尽可能使用。

init_method 告知每个进程如何发现彼此，如何使用通信后端初始化和验证进程组。 默认情况下，如果未指定 init_method，PyTorch 将使用环境变量初始化方法 (env://)。 init_method 是建议在训练代码中使用的用于在 Azure 机器学习上运行分布式 PyTorch 的初始化方法。 PyTorch 将查找以下用于初始化的环境变量：

• MASTER_ADDR：将承载排名为 0 的进程的计算机的 IP 地址。
• MASTER_PORT：将承载排名为 0 的进程的计算机的空闲端口。
• WORLD_SIZE：进程总数。 应该等于用于分布式训练的设备 (GPU) 总数。
• RANK：当前进程的（全局）排名。 可能的值为 0 到（世界大小 - 1）。

有关进程组初始化详细信息，请参阅 PyTorch 文档。

除此之外，许多应用程序还需要以下环境变量：

• LOCAL_RANK：节点内进程的本地（相对）排名。 可能的值为 0 到（节点上的进程数 - 1）。 此信息很有用，因为许多操作（例如数据准备）在每个节点上只能执行一次，通常是在 local_rank = 0 上。
• NODE_RANK：用于多节点训练的节点的排名。 可能的值为 0 到（节点总数 - 1）。

PyTorch 启动选项

Azure 机器学习 PyTorch 作业支持用两种类型的选项来启动分布式训练：

• 每进程启动器：系统将启动所有分布式进程，并包含用于设置进程组的所有相关信息（例如环境变量）。
• 每节点启动器：为 Azure 机器学习提供一个将在每个节点上运行的实用工具启动器。 该实用工具启动器将负责在给定节点上启动每个进程。 在每个节点内，RANK 和 LOCAL_RANK 由启动器以本地方式设置。 torch.distributed.launch 实用工具和 PyTorch Lightning 都属于这一类别。

这些启动选项之间没有根本的区别。 选择哪种启动选项主要取决于用户偏好或基于原版 PyTorch 构建的框架/库的约定（如 Lightning 或 Hugging Face）。

以下部分详细介绍了如何为每个启动选项配置 Azure 机器学习 PyTorch 作业。

DistributedDataParallel（每进程启动）

无需使用类似 torch.distributed.launch 的启动器实用工具。 要运行分布式 PyTorch 作业：

1. 指定训练脚本和参数
2. 创建 PyTorchConfiguration 并指定 process_count 和 node_count。 process_count 对应于要为作业运行的进程总数。 process_count 通常应等于 # GPUs per node x # nodes。 如果未指定 process_count，Azure 机器学习将默认每节点启动一个进程。

Azure 机器学习将在每个节点上设置 MASTER_ADDR、MASTER_PORT、WORLD_SIZE 和 NODE_RANK 环境变量，并设置进程级的 RANK 和 LOCAL_RANK 环境变量。

要将此选项用于每节点多进程的训练，请使用 Azure 机器学习 Python SDK >= 1.22.0。 Process_count 是在 1.22.0 中引入。

from azureml.core import ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import PyTorchConfiguration

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = PyTorchConfiguration(process_count=8, node_count=2)

run_config = ScriptRunConfig(
  source_directory='./src',
  script='train.py',
  arguments=['--epochs', 50],
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

run = Experiment(ws, 'experiment_name').submit(run_config)

arguments=['--epochs', 50, '--local_rank', $LOCAL_RANK]

Pytorch 每进程启动示例

• azureml-examples：CIFAR-10 上使用 PyTorch 的分布式训练

使用 torch.distributed.launch（每节点启动）

PyTorch 在 torch.distributed.launch 中提供了一个启动实用工具，可以使用它在每节点上启动多个进程。 torch.distributed.launch 模块将在每个节点上生成多个训练进程。

以下步骤演示了如何在 Azure 机器学习上使用每节点启动器配置 PyTorch 作业。 作业实现的效果等同于运行以下命令：

python -m torch.distributed.launch --nproc_per_node <num processes per node> \
  --nnodes <num nodes> --node_rank $NODE_RANK --master_addr $MASTER_ADDR \
  --master_port $MASTER_PORT --use_env \
  <your training script> <your script arguments>

1. 向 ScriptRunConfig 构造函数的 command 参数提供 torch.distributed.launch 命令。 Azure 机器学习在训练群集的每个节点上运行此命令。 --nproc_per_node 应小于或等于每个节点上可用的 GPU 数。 MASTER_ADDR、MASTER_PORT 和 NODE_RANK 都是由 Azure 机器学习设置，因此可以在命令中引用环境变量。 Azure 机器学习将 MASTER_PORT 设置为 6105，但如果需要，可以将不同的值传递给 torch.distributed.launch 命令的 --master_port 参数。 （启动实用工具将重置环境变量。）
2. 创建 PyTorchConfiguration 并指定 node_count。

from azureml.core import ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import PyTorchConfiguration

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = PyTorchConfiguration(node_count=2)
launch_cmd = "python -m torch.distributed.launch --nproc_per_node 4 --nnodes 2 --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT --use_env train.py --epochs 50".split()

run_config = ScriptRunConfig(
  source_directory='./src',
  command=launch_cmd,
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

run = Experiment(ws, 'experiment_name').submit(run_config)

launch_cmd = "python -m torch.distributed.launch --nproc_per_node 4 --use_env train.py --epochs 50".split()

run_config = ScriptRunConfig(
 source_directory='./src',
 command=launch_cmd,
 compute_target=compute_target,
 environment=pytorch_env,
)

PyTorch 每节点启动示例

• azureml-examples：CIFAR-10 上使用 PyTorch 的分布式训练

PyTorch Lightning

PyTorch Lightning 是一个轻量级的开源库，为 PyTorch 提供高级接口。 Lightning 抽象出了许多原版 PyTorch 所需的较低级别的分布式训练配置。 Lightning 支持在单个 GPU、单节点多 GPU 和多节点多 GPU 设置中运行训练脚本。 Lightning 在后台启动类似于 torch.distributed.launch 的多个进程。

对于单节点训练（包括单节点多 GPU），可以在 Azure 机器学习上运行代码，无需指定 distributed_job_config。 若要使用多个节点和多个 GPU 运行试验，有两个选择：

• 使用 PyTorch 配置（建议）：定义 PyTorchConfiguration 并指定 communication_backend="Nccl"、node_count 和 process_count（请注意，这是进程总数，即 num_nodes * process_count_per_node）。 在“Lightning 训练器”模块中，指定 num_nodes 和 gpus 与 PyTorchConfiguration 保持一致。 例如，num_nodes = node_count 和 gpus = process_count_per_node。

• 使用 MPI 配置：

  • 定义 MpiConfiguration 并指定 node_count 和 process_count_per_node。 在“Lightning 训练器”中，指定 num_nodes 和 gpus 分别与 MpiConfiguration 中的 node_count 和 process_count_per_node 相同。

  • 对于使用 MPI 的多节点训练，Lightning 要求在训练群集的每个节点上设置以下环境变量：

    • MASTER_ADDR
    • MASTER_PORT
    • NODE_RANK
    • LOCAL_RANK

    在主训练脚本中手动设置 Lightning 所需的这些环境变量：

  import os
  from argparse import ArgumentParser
  
  def set_environment_variables_for_mpi(num_nodes, gpus_per_node, master_port=54965):
       if num_nodes > 1:
           os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"] = os.environ["AZ_BATCH_MASTER_NODE"].split(":")
       else:
           os.environ["MASTER_ADDR"] = os.environ["AZ_BATCHAI_MPI_MASTER_NODE"]
           os.environ["MASTER_PORT"] = str(master_port)
  
       try:
           os.environ["NODE_RANK"] = str(int(os.environ.get("OMPI_COMM_WORLD_RANK")) // gpus_per_node)
           # additional variables
           os.environ["MASTER_ADDRESS"] = os.environ["MASTER_ADDR"]
           os.environ["LOCAL_RANK"] = os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]
           os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"]
       except:
           # fails when used with pytorch configuration instead of mpi
           pass
  
  if __name__ == "__main__":
       parser = ArgumentParser()
       parser.add_argument("--num_nodes", type=int, required=True)
       parser.add_argument("--gpus_per_node", type=int, required=True)
       args = parser.parse_args()
       set_environment_variables_for_mpi(args.num_nodes, args.gpus_per_node)
  
       trainer = Trainer(
        num_nodes=args.num_nodes,
        gpus=args.gpus_per_node
    )
  

  Lightning 从训练器标志 --gpus 和 --num_nodes 计算世界大小。

  from azureml.core import ScriptRunConfig, Experiment
  from azureml.core.runconfig import MpiConfiguration
  
  nnodes = 2
  gpus_per_node = 4
  args = ['--max_epochs', 50, '--gpus_per_node', gpus_per_node, '--accelerator', 'ddp', '--num_nodes', nnodes]
  distr_config = MpiConfiguration(node_count=nnodes, process_count_per_node=gpus_per_node)
  
  run_config = ScriptRunConfig(
     source_directory='./src',
     script='train.py',
     arguments=args,
     compute_target=compute_target,
     environment=pytorch_env,
     distributed_job_config=distr_config,
  )
  
  run = Experiment(ws, 'experiment_name').submit(run_config)
  

Hugging Face Transformers

Hugging Face 针对将其 Transformers 库与 torch.distributed.launch 结合使用来运行分布式训练，提供了很多示例。 要使用 Transformers Trainer API 运行这些示例和自己的自定义训练脚本，请遵照使用 torch.distributed.launch 一节。

作业配置代码示例，可在一个含 8 个 GPU 的节点上使用 run_glue.py 脚本针对文本分类 MNLI 任务微调 BERT 大模型：

from azureml.core import ScriptRunConfig
from azureml.core.runconfig import PyTorchConfiguration

distr_config = PyTorchConfiguration() # node_count defaults to 1
launch_cmd = "python -m torch.distributed.launch --nproc_per_node 8 text-classification/run_glue.py --model_name_or_path bert-large-uncased-whole-word-masking --task_name mnli --do_train --do_eval --max_seq_length 128 --per_device_train_batch_size 8 --learning_rate 2e-5 --num_train_epochs 3.0 --output_dir /tmp/mnli_output".split()

run_config = ScriptRunConfig(
  source_directory='./src',
  command=launch_cmd,
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

还可以使用每进程启动选项在不使用 torch.distributed.launch 的情况下运行分布式训练。 如果使用此方法，需要注意的一点是，转换器 TrainingArguments 预期将本地排名作为参数 (--local_rank) 来传递。 为 --use_env=False 时，则由 torch.distributed.launch 处理这种情况，但如果在使用“每进程启动”，则需要将本地排名作为参数显式传递到训练脚本 --local_rank=$LOCAL_RANK，因为 Azure 机器学习仅仅设置 LOCAL_RANK 环境变量。

TensorFlow

如果在训练代码中使用本机分布式 TensorFlow（比如 TensorFlow 2.x 的 tf.distribute.Strategy API），则可以使用 TensorflowConfiguration 通过 Azure 机器学习来启动分布式作业。

为此，请为 ScriptRunConfig 构造函数的 distributed_job_config 参数指定 TensorflowConfiguration 对象。 如果在使用 tf.distribute.experimental.MultiWorkerMirroredStrategy，请在与训练作业的节点数相对应的 TensorflowConfiguration 中指定 worker_count。

from azureml.core import ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import TensorflowConfiguration

curated_env_name = 'AzureML-TensorFlow-2.3-GPU'
tf_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = TensorflowConfiguration(worker_count=2, parameter_server_count=0)

run_config = ScriptRunConfig(
  source_directory='./src',
  script='train.py',
  compute_target=compute_target,
  environment=tf_env,
  distributed_job_config=distr_config,
)

# submit the run configuration to start the job
run = Experiment(ws, "experiment_name").submit(run_config)

如果训练脚本使用参数服务器策略进行分布式训练（例如，对于传统的 TensorFlow 1.x），则还需要指定要在作业中使用的参数服务器数量，例如 tf_config = TensorflowConfiguration(worker_count=2, parameter_server_count=1)。

TF_CONFIG

在 TensorFlow 中，在多台计算机上训练需要 TF_CONFIG 环境变量。 对于 TensorFlow 作业，在执行训练脚本之前，Azure 机器学习会相应地为每个工作器配置和设置 TF_CONFIG 变量。

如果需要，可以从训练脚本访问 TF_CONFIG：os.environ['TF_CONFIG']。

在主要工作器节点上设置的 TF_CONFIG 示例：

TF_CONFIG='{
    "cluster": {
        "worker": ["host0:2222", "host1:2222"]
    },
    "task": {"type": "worker", "index": 0},
    "environment": "cloud"
}'

TensorFlow 示例

• azureml-examples：使用 MultiWorkerMirroredStrategy 的分布式 TensorFlow 训练

使用 InfiniBand 加速分步式 GPU 训练

随着训练某个模型的 VM 数量的增加，训练该模型所需的时间应当会减少。 理想情况下，时间的减少应该与训练 VM 的数量成线性比例。 例如，如果在一个 VM 上训练某个模型需要 100 秒，那么理想情况下，在两个 VM 上训练同一模型应当需要 50 秒。 在四个 VM 上训练模型应当需要 25 秒，依此类推。

InfiniBand 可能是实现这种线性缩放的一个重要因素。 InfiniBand 可以在群集中节点之间实现低延迟的 GPU 到 GPU 通信。 InfiniBand 需要专用硬件来运行。 某些 Azure VM 系列（具体而言是 NC、ND 和 H 系列）现在具有即支持 RDMA 功能又支持 SR-IOV 和 InfiniBand 的 VM。 这些 VM 在低延迟、高带宽的 InfiniBand 网络上进行通信，这比基于以太网的连接的性能更高。 适用于 InfiniBand 的 SR-IOV 可为任何 MPI 库提供接近裸机的性能（MPI 被许多分布式训练框架和工具采用，包括 NVIDIA 的 NCCL 软件）。这些 SKU 旨在满足计算密集型、GPU 加速的机器学习工作负荷的需求。 有关详细信息，请参阅在 Azure 机器学习中采用 SR-IOV 加速分布式训练。

通常，名称中包含“r”的 VM SKU 包含所需的 InfiniBand 硬件，而不带“r”的 VM SKU 通常不包含。 （“r”指的是 RDMA，表示“remote direct memory access”。）例如，VM SKU Standard_NC24rs_v3 启用了 InfiniBand，但 SKU Standard_NC24s_v3 没有。 除了 InfiniBand 功能外，这两个 SKU 的规格大致相同 - 对于同一 SKU，两者都有 24 个核心、448 GB RAM、4 个 GPU，等等。 详细了解启用了 RDMA 和 InfiniBand 的计算机 SKU。

如果创建了其中一个即支持 RDMA 功能又支持 InfiniBand 大小的 AmlCompute 群集，则操作系统映像将附带启用预安装和预配置 InfiniBand 所需的 Mellanox OFED 驱动程序。

后续步骤

• 将机器学习模型部署到 Azure
• Azure 中分布式深度学习训练的参考体系结构