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

项目
09/13/2024

适用范围：Python SDK azure-ai-ml v2（最新版）

详细了解如何在 Azure 机器学习中使用分布式 GPU 训练代码。本文可帮助你运行现有的分布式训练代码，并提供每个框架要遵循的提示和示例：

PyTorch
TensorFlow
使用 InfiniBand 加速 GPU 训练

先决条件

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

提示

如果不知道要使用哪种类型的并行性，那么 90% 以上的情况应该使用分布式数据并行性。

PyTorch

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

提示

对于数据并行性，PyTorch 官方指南使用 DistributedDataParallel (DDP) over DataParallel 进行单节点和多节点分布式训练。 PyTorch 还建议对多处理包使用 DistributedDataParallel。因此，Azure 机器学习文档和示例重点介绍 DistributedDataParallel 训练。

进程组初始化

任何分布式训练的主干都基于一组彼此知晓且可以使用后端相互通信的进程。对于 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）。

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

指定训练脚本和参数。
创建 command，将类型指定为 PyTorch，并在 distribution 参数中指定 process_count_per_instance。 process_count_per_instance 对应于要为作业运行的进程总数。 process_count_per_instance 通常应等于 # of GPUs per node。如果未指定 process_count_per_instance，Azure 机器学习将默认每节点启动一个进程。

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

from azure.ai.ml import command
from azure.ai.ml.entities import Data
from azure.ai.ml import Input
from azure.ai.ml import Output
from azure.ai.ml.constants import AssetTypes

# === Note on path ===
# can be can be a local path or a cloud path. AzureML supports https://`, `abfss://`, `wasbs://` and `azureml://` URIs.
# Local paths are automatically uploaded to the default datastore in the cloud.
# More details on supported paths: https://docs.azure.cn/machine-learning/how-to-read-write-data-v2#supported-paths

inputs = {
    "cifar": Input(
        type=AssetTypes.URI_FOLDER, path=returned_job.outputs.cifar.path
    ),  # path="azureml:azureml_stoic_cartoon_wgb3lgvgky_output_data_cifar:1"), #path="azureml://datastores/workspaceblobstore/paths/azureml/stoic_cartoon_wgb3lgvgky/cifar/"),
    "epoch": 10,
    "batchsize": 64,
    "workers": 2,
    "lr": 0.01,
    "momen": 0.9,
    "prtfreq": 200,
    "output": "./outputs",
}

job = command(
    code="./src",  # local path where the code is stored
    command="python train.py --data-dir ${{inputs.cifar}} --epochs ${{inputs.epoch}} --batch-size ${{inputs.batchsize}} --workers ${{inputs.workers}} --learning-rate ${{inputs.lr}} --momentum ${{inputs.momen}} --print-freq ${{inputs.prtfreq}} --model-dir ${{inputs.output}}",
    inputs=inputs,
    environment="azureml:AzureML-pytorch-1.9-ubuntu18.04-py37-cuda11-gpu:6",
    compute="gpu-cluster",  # Change the name to the gpu cluster of your workspace.
    instance_count=2,  # In this, only 2 node cluster was created.
    distribution={
        "type": "PyTorch",
        # set process count to the number of gpus per node
        # NV6 has only 1 GPU
        "process_count_per_instance": 1,
    },
)

Pytorch 示例

有关运行 Pytorch 示例的完整笔记本，请参阅 azureml-examples：CIFAR-10 上使用 PyTorch 的分布式训练。

DeepSpeed

Azure 机器学习支持将 DeepSpeed 作为一等公民，可在以下方面以近乎线性的可缩放性运行分布式作业：

模型大小增加
GPU 数量增加

可以使用 Pytorch 发行版或 MPI 启用 DeepSpeed 来运行分布式训练。 Azure 机器学习支持使用 DeepSpeed 启动器来启动分布式训练和自动优化，以获得最佳 ds 配置。

可以将特选环境用于开箱即用的环境，其中包含用于 DeepSpeed 训练作业的最新先进技术，例如 DeepSpeed、ORT、MSSCCL 和 Pytorch。

DeepSpeed 示例

有关 DeepSpeed 训练和自动优化示例，请参阅这些文件夹。

TensorFlow

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

# create the command
job = command(
    code="./src",  # local path where the code is stored
    command="python main.py --epochs ${{inputs.epochs}} --model-dir ${{inputs.model_dir}}",
    inputs={"epochs": 1, "model_dir": "outputs/keras-model"},
    environment="AzureML-tensorflow-2.4-ubuntu18.04-py37-cuda11-gpu@latest",
    compute="cpu-cluster",
    instance_count=2,
    # distribution = {"type": "mpi", "process_count_per_instance": 1},
    distribution={
        "type": "tensorflow",
        "parameter_server_count": 1,
        "worker_count": 2,
        "added_property": 7,
    },
    # distribution = {
    #        "type": "pytorch",
    #        "process_count_per_instance": 4,
    #        "additional_prop": {"nested_prop": 3},
    #    },
    display_name="tensorflow-mnist-distributed-example"
    # experiment_name: tensorflow-mnist-distributed-example
    # description: Train a basic neural network with TensorFlow on the MNIST dataset, distributed via TensorFlow.
)

# can also set the distribution in a separate step and using the typed objects instead of a dict
job.distribution = TensorFlowDistribution(parameter_server_count=1, worker_count=2)

如果训练脚本使用参数服务器策略进行分布式训练（例如，对于传统的 TensorFlow 1.x），则还需要在 command 的 distribution 参数内指定要在作业中使用的参数服务器数量。例如，在上面 "parameter_server_count" : 1 和 "worker_count": 2。

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 示例

有关运行 TensorFlow 示例的完整笔记本，请参阅 azureml-examples：结合使用 Tensorflow 和 Horovod 在 MNIST 数据集上通过分布式 MPI 训练基本神经网络。

使用 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 的引用，它代表远程直接内存访问。）例如，VM SKU Standard_NC24rs_v3 已启用 InfiniBand，但 SKU Standard_NC24s_v3 未启用。除了 InfiniBand 功能外，这两个 SKU 的规格大致相同。对于同一 SKU，两者都有 24 个核心、448 GB RAM、4 个 GPU，等等。详细了解启用了 RDMA 和 InfiniBand 的计算机 SKU。

警告

较旧代系的计算机 SKU Standard_NC24r 启用了 RDMA，但它不包含 InfiniBand 所需的 SR-IOV 硬件。

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

通过