次の方法で共有


分散 GPU トレーニング ガイド (SDK v2)

適用対象: Python SDK azure-ai-ml v2 (現行)

この記事では、Azure Machine Learning で分散 GPU トレーニング コードを使用する方法について説明します。 PyTorch、DeepSpeed、TensorFlow のヒントと例を使用して、既存のコードを実行する方法について説明します。 InfiniBand を使用した分散 GPU トレーニングの高速化についても説明します。

ヒント

90% を超える時間では、 分散データ並列処理を分散並列処理 の種類として使用する必要があります。

前提条件

  • データの並列処理、分散データの並列処理、モデルの並列処理など、分散 GPU トレーニングの基本的な概念を理解します。
  • バージョン 1.5.0 以降のパッケージがインストールされている azure-ai-ml

PyTorch

Azure Machine Learning では、PyTorch のネイティブ分散トレーニング機能 ( torch.distributed) を使用した分散ジョブの実行がサポートされています。

データの並列処理の場合、 公式の PyTorch ガイダンス では、単一ノードとマルチノードの両方の分散トレーニングに DataParallel 経由で DistributedDataParallel (DDP) を使用することをお勧めします。 PyTorch では、マルチプロセッシング パッケージより DistributedDataParallel を優先して使用することも推奨されています。 そのため、Azure Machine Learning のドキュメントと例では、DistributedDataParallel トレーニングに重点を置いています。

プロセス グループの初期化

分散トレーニングのバックボーンは、バックエンドを使用して相互に認識して通信できるプロセスのグループです。 PyTorch の場合は、すべての分散プロセスで torch.distributed.init_process_group を呼び出してプロセス グループをまとめて形成することで、プロセス グループを作成します。

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

最も一般的な通信バックエンドは、 mpinccl、および glooです。 GPU ベースのトレーニングでは、最適なパフォーマンスを得るための nccl を使用します。

Azure Machine Learning で分散 PyTorch を実行するには、トレーニング コードで init_method パラメーターを使用します。 このパラメーターは、各プロセスが他のプロセスを検出する方法と、通信バックエンドを使用してプロセス グループを初期化および検証する方法を指定します。 既定では、 init_methodを指定しない場合、PyTorch は環境変数の初期化メソッド env://を使用します。

PyTorch によって、初期化のために次の環境変数が検索されます。

  • MASTER_ADDR: ランク 0でプロセスをホストするマシンの IP アドレス。
  • MASTER_PORT: ランク 0でプロセスをホストするマシン上の空きポート。
  • WORLD_SIZE: プロセスの合計数。 分散トレーニングに使用される GPU デバイスの合計数に等しい必要があります。
  • RANK: 現在のプロセスのグローバル ランク。 指定できる値は、0(<world size> - 1)

プロセス グループの初期化の詳細については、PyTorch のドキュメントを参照してください。

環境変数

多くのアプリケーションでは次の環境変数も必要になります:

  • LOCAL_RANK: ノード内のプロセスのローカル相対ランク。 指定できる値は、0(<# of processes on the node> - 1)。 この情報は、データ準備などの多くの操作をノードごとに 1 回だけ実行する必要があるため、通常は local_rank = 0 で役立ちます。
  • NODE_RANK: マルチノード トレーニングのノードのランク。 指定できる値は、0(<total # of nodes> - 1)

分散 PyTorch ジョブを実行する

分散 PyTorch ジョブを実行するために、 torch.distributed.launch などのランチャー ユーティリティを使用する必要はありません。 次のようにすることができます。

  1. トレーニング スクリプトと引数を指定します。

  2. commandを作成し、typePyTorch として指定し、process_count_per_instance パラメーターにdistributionを指定します。

    process_count_per_instanceは、ジョブに対して実行するプロセスの合計数に対応し、通常はノードあたりの GPU の数と等しい必要があります。 process_count_per_instanceを指定しない場合、Azure Machine Learning は既定でノードごとに 1 つのプロセスを起動します。

Azure Machine Learning は、各ノードに MASTER_ADDRMASTER_PORTWORLD_SIZE、および NODE_RANK 環境変数を設定し、プロセス レベルの RANKLOCAL_RANK 環境変数を設定します。

PyTorch の例

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.microsoft.com/azure/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",
}

from azure.ai.ml.entities import ResourceConfiguration

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-acpt-pytorch-2.8-cuda12.6@latest",
    instance_count=2,  # In this, only 2 node cluster was created.
    distribution={
        "type": "PyTorch",
        # set process count to the number of gpus per node
        # NC6s_v3 has only 1 GPU
        "process_count_per_instance": 1,
    },
)
job.resources = ResourceConfiguration(
    instance_type="STANDARD_NC4AS_T4_V3", instance_count=2
)  # Serverless compute resources

PyTorch の例を実行するための完全なノートブックについては、「azureml-examples: CIFAR-10 での PyTorch を使用した分散トレーニング」を参照してください。

DeepSpeed

Azure Machine Learning では、モデル サイズの増加と GPU 数の増加に対して、ほぼ直線的なスケーラビリティを備えた分散ジョブを実行するための最上位機能として DeepSpeed がサポートされています。

PyTorch ディストリビューションまたはメッセージ パッシング インターフェイス (MPI) を使用して、分散トレーニングを実行するために DeepSpeed を有効にすることができます。 Azure Machine Learning では、分散トレーニングを開始するための DeepSpeed ランチャーと、最適な ds 構成を得るための自動チューニングがサポートされています。

DeepSpeed トレーニング ジョブには、DeepSpeed、ONNX (Open Neural Network Exchange) Runtime (ORT)、Microsoft Collective Communication Library (MSSCCL)、PyTorch など、最新のテクノロジを使用して キュレーションされた環境 を使用できます。

DeepSpeed の例

DeepSpeed のトレーニングと自動チューニングの例については、 https://github.com/Azure/azureml-examples/cli/jobs/deepspeedを参照してください。

TensorFlow

TensorFlow 2.x API などのトレーニング コードでtf.distribute.Strategy を使用する場合は、distribution パラメーターまたはTensorFlowDistribution オブジェクトを使用して、Azure Machine Learning を使用して分散ジョブを起動できます。

# 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.16-cuda12@latest",
    compute="cpu-cluster",
    instance_count=2,
    # distribution = {"type": "mpi", "process_count_per_instance": 1},
    # distribution={
    #     "type": "tensorflow",
    #     "parameter_server_count": 1,  # for legacy TensorFlow 1.x
    #     "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(worker_count=2)

従来の TensorFlow 1.x など、トレーニング スクリプトで分散トレーニングにパラメーター サーバー戦略を使用する場合は、ジョブで使用するパラメーター サーバーの数も指定する必要があります。 前の例では、"parameter_server_count" : 1"worker_count": 2 パラメーター内にdistributioncommandを指定します。

TF_CONFIG

TensorFlow の複数のマシンでトレーニングするには、 TF_CONFIG 環境変数を使用します。 TensorFlow ジョブの場合、Azure Machine Learning はトレーニング スクリプトを実行する前に、各ワーカーの 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 の例を実行するための完全なノートブックについては、 tensorflow-mnist-distributed-example を参照してください。

InfiniBand

モデルをトレーニングする仮想マシン (VM) の数を増やすと、そのモデルのトレーニングに必要な時間は、トレーニング VM の数に比例して減少します。 たとえば、1 つの仮想マシン (VM) でモデルをトレーニングするのに 100 秒かかる場合、2 台の VM で同じモデルをトレーニングするには 50 秒かかります。4 台の VM でのモデルのトレーニングには 25 秒かかる必要があります。

InfiniBand は、クラスター内のノード間で低待機時間の GPU 間通信を可能にすることで、この線形スケーリングを実現するのに役立ちます。 InfiniBand には、Azure VM NC、ND、H シリーズなどの特殊なハードウェアが必要です。 これらの VM には、シングル ルート I/O 仮想化 (SR-IOV) と InfiniBand サポートを備えたリモート ダイレクト メモリ アクセス (RDMA) 対応の VM があります。

これらの VM は、イーサネット ベースの接続よりも優れたパフォーマンスを発揮する、低待機時間および高帯域幅の InfiniBand ネットワーク経由で通信します。 InfiniBand 用の SR-IOV は、NVIDIA Collective Communications Library (NCCL) などの多くの分散トレーニング フレームワークやツールで使用される、MPI ライブラリのベアメタル パフォーマンスに近いパフォーマンスを実現します。

これらのストック キーピング ユニット (SKU) は、計算負荷の高い GPU で高速化された機械学習ワークロードのニーズを満たすことを目的としています。 詳細については、 による Azure Machine Learning での分散トレーニングの高速化に関する記事を参照してください。

通常、RDMA を参照する名前に r が含まれる VM SKU にのみ、必要な InfiniBand ハードウェアが含まれます。 たとえば、VM SKU Standard_NC24rs_v3 は InfiniBand 対応ですが、 Standard_NC24s_v3 は有効ではありません。 これら 2 つの SKU の仕様は、InfiniBand 機能を除いてほとんど同じです。 どちらの SKU にも、24 コア、448 GB の RAM、同じ SKU の 4 つの GPU があります。 RDMA および InfiniBand 対応マシン SKU の詳細については、「 ハイ パフォーマンス コンピューティング」を参照してください。

旧世代のマシン SKU Standard_NC24r は RDMA 対応ですが、InfiniBand に必要な SR-IOV ハードウェアは含まれません。

これらの RDMA 対応の InfiniBand 対応サイズのいずれかを使用して AmlCompute クラスターを作成する場合、INFiniBand がプレインストールおよび事前構成済みで有効にするために必要な Mellanox OpenFabrics Enterprise Distribution (OFED) ドライバーが OS イメージに付属しています。