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

[アーティクル]
09/03/2024

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

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

PyTorch
TensorFlow
InfiniBand による GPU トレーニングの高速化

前提条件

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

ヒント

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

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:AzureML-acpt-pytorch-2.2-cuda12.1@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_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.12-cuda11@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)

前提条件

PyTorch

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

Pytorch の例

DeepSpeed

DeepSpeed の例

TensorFlow

TF_CONFIG

TensorFlow の例

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

次のステップ

フィードバック

その他のリソース

次の方法で共有

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

前提条件

PyTorch

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

Pytorch の例

DeepSpeed

DeepSpeed の例

TensorFlow

TF_CONFIG

TensorFlow の例

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

次のステップ

フィードバック

その他のリソース

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

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