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

[アーティクル]
01/29/2024

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

Azure Machine Learning での分散 GPU トレーニングコードの使用の詳細について説明します。この記事は、既存の分散トレーニングコードを実行する際に役立ち、以下の各フレームワークについて従うべきヒントと例を提供します。

メッセージパッシングインターフェイス (MPI)
- Horovod
- Open MPI からの環境変数
PyTorch
TensorFlow
InfiniBand による GPU トレーニングの高速化

前提条件

"データ並列処理"、"分散データ並列処理"、"モデル並列処理" など、分散 GPU トレーニングの基本的な概念を確認します。

ヒント

使用する並列処理の種類がわからない場合は、90% 超の時間は分散データ並列処理を使用するはずです。

MPI

Azure Machine Learning には、各ノードで特定の数のプロセスを起動する MPI ジョブが用意されています。 Azure Machine Learning により、バックグラウンドで完全な MPI 起動コマンド (mpirun) が作成されます。 mpirun や DeepSpeed launcher のように、独自の完全なヘッドノードランチャーコマンドを提供することはできません。

ヒント

Azure Machine Learning MPI ジョブによって使用される基本 Docker イメージに、MPI ライブラリがインストールされている必要があります。 Open MPI は、すべての Azure Machine Learning GPU 基本イメージに含まれています。カスタム Docker イメージを使用する場合は、イメージに MPI ライブラリが含まれていることを自分で確認する必要があります。 Open MPI を使用することをお勧めしますが、Intel MPI などの別の MPI 実装を使用することもできます。 Azure Machine Learning では、人気のあるフレームワーク用のキュレーションされた環境も提供されています。

MPI を使用して分散トレーニングを実行するには、次の手順のようにします。

好みのディープラーニングフレームワークと MPI で Azure Machine Learning 環境を使用します。 Azure Machine Learning では、人気のあるフレームワーク用のキュレーションされた環境が提供されています。または、好みのディープラーニングフレームワークと MPI でカスタム環境を作成します。
instance_count で command を定義します。 instance_count は、プロセス単位の起動の場合はノードあたりの GPU の数と等しい値に設定し、ユーザースクリプトでノードごとにプロセスを起動するノード単位の起動の場合は 1 (既定値) に設定する必要があります。
command の distribution パラメーターを使用して、MpiDistribution の設定を指定します。

from azure.ai.ml import command, MpiDistribution

job = command(
    code="./src",  # local path where the code is stored
    command="python train.py --epochs ${{inputs.epochs}}",
    inputs={"epochs": 1},
    environment="AzureML-tensorflow-2.7-ubuntu20.04-py38-cuda11-gpu@latest",
    compute="gpu-cluster",
    instance_count=2,
    distribution=MpiDistribution(process_count_per_instance=2),
    display_name="tensorflow-mnist-distributed-horovod-example"
    # experiment_name: tensorflow-mnist-distributed-horovod-example
    # description: Train a basic neural network with TensorFlow on the MNIST dataset, distributed via Horovod.
)

Horovod

ディープラーニングフレームワークによる分散トレーニングに Horovod を使用する場合は、MPI ジョブ構成を使用します。

コードが次のヒントに従っていることを確認します。

Azure Machine Learning のパーツを追加する前に、トレーニングコードが Horovod で正しくインストルメント化されています。
Azure Machine Learning 環境に Horovod と MPI が含まれています。 PyTorch および TensorFlow のキュレーションされた GPU 環境に、Horovod とその依存関係が事前に構成されています。
目的の分散で command を作成します。

Horovod の例

Horovod の例を実行するための完全なノートブックについては、「azureml-examples: Horovod を使用して MNIST データセットで分散 MPI を使用して基本的なニューラルネットワークをトレーニングする」を参照してください

Open MPI からの環境変数

Open MPI イメージで MPI ジョブを実行するときは、起動されるプロセスごとに次の環境変数を使用できます。

OMPI_COMM_WORLD_RANK: プロセスのランク
OMPI_COMM_WORLD_SIZE: ワールドサイズ
AZ_BATCH_MASTER_NODE: プライマリアドレスとポート、MASTER_ADDR:MASTER_PORT
OMPI_COMM_WORLD_LOCAL_RANK: ノード上のプロセスのローカルランク
OMPI_COMM_WORLD_LOCAL_SIZE: ノード上のプロセスの数

ヒント

環境変数 OMPI_COMM_WORLD_NODE_RANK は、その名前にもかかわらず、NODE_RANK に対応していません。ノード単位ランチャーを使用するには、process_count_per_node=1 を設定し、NODE_RANK として OMPI_COMM_WORLD_RANK を使用します。

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://', ...)

使用される最も一般的な通信バックエンドは、mpi、nccl、gloo です。 GPU ベースのトレーニングの場合は、最適なパフォーマンスのために nccl が推奨され、可能な限りそれを使用する必要があります。

init_method により、各プロセスが相互を検出する方法と、通信バックエンドを使用してプロセスグループを初期化および検証する方法が指示されます。既定では、init_method が指定されていない場合、PyTorch では環境変数の初期化メソッド (env://) が使用されます。 init_method は、Azure Machine Learning で分散 PyTorch を実行するためにトレーニングコードで使用することが推奨される初期化方法です。 PyTorch によって、初期化のために次の環境変数が検索されます。

MASTER_ADDR: ランクが 0 のプロセスをホストするマシンの IP アドレス
MASTER_PORT: ランクが 0 のプロセスをホストするマシンの空きポート
WORLD_SIZE: プロセスの合計数。分散トレーニングに使用されるデバイス (GPU) の合計数と同じである必要があります
RANK: 現在のプロセスの (グローバル) ランク。指定できる値は 0 から (ワールドサイズ - 1) です

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

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

LOCAL_RANK: ノード内のプロセスのローカル (相対) ランク。指定できる値は 0 から (ノードのプロセス数 - 1) です。データの準備などの多くの操作は、通常は local_rank = 0 でノードごとに 1 回だけ実行する必要があるため、この情報は便利です。
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 Machine Learning によって既定でノードごとに 1 つのプロセスが起動されます。

Azure Machine Learning によって、各ノードで環境変数 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.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-acpt-pytorch-2.0-cuda11.7:26",
    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_NC6s_v3", instance_count=2
)  # Serverless compute resources

Pytorch の例

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

DeepSpeed

Azure Machine Learning では、DeepSpeed を第一級オブジェクトとしてサポートしており、次の点において、分散されたジョブをほぼ直線的な拡張性で実行します。

モデルサイズの増加
GPU 数の増加

DeepSpeed は、分散トレーニングを実行するために Pytorch ディストリビューションまたは MPI を使用して有効にできます。 Azure Machine Learning では、分散型のトレーニングと、最適な ds 構成を得るための自動チューニングを起動する DeepSpeed ランチャーをサポートしています。

DeepSpeed トレーニングジョブには、DeepSpeed、ORT、MSSCCL、Pytorch など、最新式のテクノロジを備えた、面倒な設定の要らない環境として、キュレーションされた環境を利用できます。

DeepSpeed の例

DeepSpeed のトレーニングと自動チューニングの例が必要な場合、こちらのフォルダーをご覧ください。

TensorFlow

TensorFlow 2.x の tf.distribute.Strategy API など、ネイティブの分散 TensorFlow をトレーニングコードで使用している場合は、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.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,  # 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 など、分散トレーニングのためにパラメータサーバー戦略を使用している場合は、command の distribution パラメータ内で、ジョブで使用するパラメータサーバーの数も指定する必要があります。上記では、例えば、"parameter_server_count" : 1 および "worker_count": 2 です。

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 の例を実行するための完全なノートブックについては、「azureml-examples: Tensorflow と Horovod を使用して、MNIST データセットで分散 MPI を使用して基本的なニューラルネットワークをトレーニングする」を参照してください。

InfiniBand による分散 GPU トレーニングの高速化

モデルをトレーニングする VM の数が増えるにつれて、そのモデルのトレーニングに必要な時間が短縮されます。理想的であれば、時間の短縮は、トレーニング VM の数に直線的に比例するはずです。たとえば、1 つの VM でモデルのトレーニングに 100 秒かかる場合、同じモデルを 2 つの VM でトレーニングすると、理想的には 50 秒かかります。 4 つの VM でモデルをトレーニングすると、25 秒かかるはずです。以下同様です。

このような直線的なスケーリングを実現するには、InfiniBand が重要な要素です。 InfiniBand を使用すると、クラスター内のノード間で、低遅延の GPU 間通信が可能になります。 InfiniBand が動作するには、特別なハードウェアが必要です。特定の Azure VM シリーズ (具体的には、NC、ND、H シリーズ) には、SR-IOV と InfiniBand をサポートする RDMA 対応 VM が用意されています。これらの VM は、低遅延で高帯域幅の InfiniBand ネットワークを介して通信し、イーサネットベースの接続よりパフォーマンスがはるかに高くなります。 InfiniBand 用の SR-IOV は、あらゆる MPI ライブラリでベアメタルに近いパフォーマンスを実現します (MPI は NVIDIA の NCCL ソフトウェアなど、多くの分散トレーニングフレームワークとツールによって使われています)。これらの SKU は、計算量が多く、GPU により高速化される機械学習ワークロードのニーズを満たすことを目的としています。詳細については、による Azure Machine Learning での分散トレーニングの高速化に関する記事を参照してください。

通常、名前に "r" が含まれる VM SKU には必要な InfiniBand ハードウェアが含まれ、"r" を含まない VM SKU には通常は含まれません。 ("r" は RDMA への参照であり、"リモートダイレクトメモリアクセス"を表します。)たとえば、VM SKU Standard_NC24rs_v3 は InfiniBand 対応ですが、SKU Standard_NC24s_v3 は対応していません。 InfiniBand の機能を除くと、これら 2 つの SKU の仕様はほぼ同じです。どちらも 24 個のコア、448 GB の RAM、同じ SKU の 4 個の GPU などを備えています。 RDMA 対応および InfiniBand 対応のマシンの SKU の詳細を確認してください。

警告

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

このような RDMA 対応で InfiniBand が有効なサイズのいずれかの AmlCompute クラスターを作成する場合、OS イメージには、InfiniBand を有効にするために必要な Mellanox OFED ドライバーがプレインストールされて事前構成されています。

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

前提条件

MPI

Horovod

Horovod の例

Open MPI からの環境変数

PyTorch

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

Pytorch の例

DeepSpeed

DeepSpeed の例

TensorFlow

TF_CONFIG

TensorFlow の例

InfiniBand による分散 GPU トレーニングの高速化

次のステップ

その他のリソース

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

前提条件

MPI

Horovod

Horovod の例

Open MPI からの環境変数

PyTorch

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

Pytorch の例

DeepSpeed

DeepSpeed の例

TensorFlow

TF_CONFIG

TensorFlow の例

InfiniBand による分散 GPU トレーニングの高速化

次のステップ

その他のリソース

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

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