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

适用范围：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，并在 process_count_per_instance 参数中指定 distribution。 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、ORT、MSSCCL 和 PyTorch）用于 DeepSpeed 训练作业。

DeepSpeed 示例

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

TensorFlow

如果在训练代码（例如 TensorFlow 2.x 的 API）中使用tf.distribute.Strategy，则可以使用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），则还需要在 distribution 的 command 参数内指定要在作业中使用的参数服务器数量。在前面的示例中，你指定 "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，通过分布式 MPI 在 MNIST 数据集上训练基本神经网络。

使用 InfiniBand 加速分步式 GPU 训练

增加训练模型的 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 的规格大致相同。两者都有 24 个核心、448 GB RAM、4 个同一 SKU 的 GPU，等等。详细了解启用了 RDMA 和 InfiniBand 的计算机 SKU。

警告

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

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

后续步骤

Last updated on 2025-11-27

通过