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

[アーティクル]
06/03/2023

Azure Machine Learning (ML) で分散 GPU トレーニングコードを使用する方法の詳細について説明します。この記事では、分散トレーニングについては説明しません。これは、既存の分散トレーニングコードを Azure Machine Learning で実行するのに役立ちます。各フレームワークを使用するためのヒントと例が提供されています。

メッセージパッシングインターフェイス (MPI)
- Horovod
- DeepSpeed
- Open MPI からの環境変数
PyTorch
- プロセスグループの初期化
- 起動のオプション
- DistributedDataParallel (プロセス単位の起動)
- torch.distributed.launch の使用 (ノード単位の起動)
- PyTorch Lightning
- Hugging Face Transformers
TensorFlow
- TensorFlow の環境変数 (TF_CONFIG)
InfiniBand による GPU トレーニングの高速化

前提条件

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

ヒント

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

MPI

Azure Machine Learning には、各ノードで特定の数のプロセスを起動する MPI ジョブが用意されています。このアプローチを採用すると、process_count_per_node の設定 (ノード単位ランチャーの場合は 1 (既定値)、プロセス単位ランチャーの場合はデバイスまたは GPU の数と等しい値) に応じて、プロセス単位ランチャーまたはノード単位ランチャーを使用して分散トレーニングを実行できます。 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 では、人気のあるフレームワーク用のキュレーションされた環境が提供されています。
process_count_per_node と node_count を使用して MpiConfiguration を定義します。 process_count_per_node は、プロセス単位の起動の場合はノードあたりの GPU の数と等しい値に設定し、ユーザースクリプトでノードごとにプロセスを起動するノード単位の起動の場合は 1 (既定値) に設定する必要があります。
MpiConfiguration オブジェクトを ScriptRunConfig の distributed_job_config パラメーターに渡します。

from azureml.core import Workspace, ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import MpiConfiguration

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = MpiConfiguration(process_count_per_node=4, node_count=2)

run_config = ScriptRunConfig(
  source_directory= './src',
  script='train.py',
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

# submit the run configuration to start the job
run = Experiment(ws, "experiment_name").submit(run_config)

Horovod

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

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

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

Horovod の例

azureml-examples: Horovod を使用する TensorFlow 分散トレーニング

DeepSpeed

Azure Machine Learning 上の DeepSpeed ライブラリを使用して分散トレーニングを実行するには、DeepSpeed のカスタムランチャーを使用しないでください。代わりに、MPI を使用してトレーニングジョブを起動するように MPI ジョブを構成します。

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

Azure Machine Learning 環境に、DeepSpeed とその依存関係、Open MPI、mpi4py が含まれています。
目的の分散で MpiConfiguration を作成します。

DeepSpeed の例

azureml-examples: CIFAR-10 で DeepSpeed を使用する分散トレーニング

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) です。この情報は、データの準備などの多くの操作はノードごとに 1 回だけ実行する必要があるため便利です (通常は、local_rank = 0 で)。
NODE_RANK - マルチノードトレーニングのノードのランク。指定できる値は 0 から (ノードの合計数 - 1) です。

PyTorch の起動オプション

Azure Machine Learning の PyTorch ジョブでは、分散トレーニングを開始するために 2 種類のオプションがサポートされています。

プロセス単位ランチャー: プロセスグループを設定するためのすべての関連情報 (環境変数など) を使用して、システムによりすべての分散プロセスが自動的に起動されます。
ノード単位ランチャー: 各ノードで実行されるユーティリティランチャーを Azure Machine Learning に提供します。ユーティリティランチャーによって、特定のノード上での各プロセスの起動を処理します。各ノード内でローカルに、RANK と LOCAL_RANK がランチャーによって設定されます。 torch.distributed.launch ユーティリティと PyTorch Lightning はどちらも、このカテゴリに属します。

これらの起動オプションの間に基本的な違いはありません。ほとんどの場合、選択は、ユーザーの好みや、vanilla PyTorch (Lightning や Hugging Face など) 上に構築されたフレームワークまたはライブラリの規則に基づいています。

以下のセクションでは、各起動オプションに対する Azure Machine Learning PyTorch ジョブの詳細な構成方法について説明します。

DistributedDataParallel (プロセス単位の起動)

torch.distributed.launch のようなランチャーユーティリティを使用する必要はありません。分散 PyTorch ジョブを実行するには:

トレーニングスクリプトと引数を指定します。
PyTorchConfiguration を作成し、process_count と node_count を指定します。 process_count は、ジョブに対して実行するプロセスの合計数に対応しています。通常、process_count は # GPUs per node x # nodes と等しくする必要があります。 process_count を指定しないと、Azure Machine Learning によって既定でノードごとに 1 つのプロセスが起動されます。

Azure Machine Learning によって、各ノードで環境変数 MASTER_ADDR、MASTER_PORT、WORLD_SIZE、NODE_RANK が設定され、プロセスレベルで RANK および LOCAL_RANK 環境変数が設定されます。

ノードごとにマルチプロセスのトレーニングでこのオプションを使用するには、Azure Machine Learning Python SDK >= 1.22.0 を使用します。 process_count は 1.22.0 で導入されました。

from azureml.core import ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import PyTorchConfiguration

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = PyTorchConfiguration(process_count=8, node_count=2)

run_config = ScriptRunConfig(
  source_directory='./src',
  script='train.py',
  arguments=['--epochs', 50],
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

run = Experiment(ws, 'experiment_name').submit(run_config)

ヒント

トレーニングスクリプトで、スクリプト引数としてローカルランクやランクなどの情報を渡す場合は、引数で環境変数を参照できます。

arguments=['--epochs', 50, '--local_rank', $LOCAL_RANK]

Pytorch のプロセス単位の起動の例

azureml-examples: CIFAR-10 で PyTorch を使用する分散トレーニング

torch.distributed.launch の使用 (ノードごとの起動)

PyTorch には、ノードごとに複数のプロセスを起動するために使用できる起動ユーティリティが torch.distributed.launch に用意されています。 torch.distributed.launch モジュールによって、各ノードで複数のトレーニングプロセスが生成されます。

次の手順では、Azure Machine Learning でノード単位ランチャーを使用して PyTorch ジョブを構成する方法を示します。このジョブにより、次のコマンドを実行するのと同等の機能が実現されます。

python -m torch.distributed.launch --nproc_per_node <num processes per node> \
  --nnodes <num nodes> --node_rank $NODE_RANK --master_addr $MASTER_ADDR \
  --master_port $MASTER_PORT --use_env \
  <your training script> <your script arguments>

ScriptRunConfig コンストラクターの command パラメーターに torch.distributed.launch コマンドを指定します。 Azure Machine Learning によって、トレーニングクラスターの各ノードでこのコマンドが実行されます。 --nproc_per_node は、各ノードで使用できる GPU の数以下である必要があります。 MASTER_ADDR、MASTER_PORT、NODE_RANK はすべて Azure Machine Learning によって設定されるので、コマンドで環境変数を参照するだけで済みます。 Azure Machine Learning によって MASTER_PORT は 6105 に設定されますが、必要に応じて、torch.distributed.launch コマンドの --master_port 引数に別の値を渡すことができます。 (起動ユーティリティによって環境変数がリセットされます。)
PyTorchConfiguration を作成し、node_count を指定します。

from azureml.core import ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import PyTorchConfiguration

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = PyTorchConfiguration(node_count=2)
launch_cmd = "python -m torch.distributed.launch --nproc_per_node 4 --nnodes 2 --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT --use_env train.py --epochs 50".split()

run_config = ScriptRunConfig(
  source_directory='./src',
  command=launch_cmd,
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

run = Experiment(ws, 'experiment_name').submit(run_config)

ヒント

単一ノードマルチ GPU のトレーニング: 起動ユーティリティを使用して、単一ノードマルチ GPU の PyTorch トレーニングを実行する場合は、ScriptRunConfig の distributed_job_config パラメーターを指定する必要はありません。

launch_cmd = "python -m torch.distributed.launch --nproc_per_node 4 --use_env train.py --epochs 50".split()

run_config = ScriptRunConfig(
 source_directory='./src',
 command=launch_cmd,
 compute_target=compute_target,
 environment=pytorch_env,
)

PyTorch のノード単位の起動の例

azureml-examples: CIFAR-10 で PyTorch を使用する分散トレーニング

PyTorch Lightning

PyTorch Lightning は、PyTorch 用の高レベルインターフェイスを提供する軽量のオープンソースライブラリです。 Lightning により、vanilla PyTorch に必要な低レベルの分散トレーニング構成の多くが抽象化されます。 Lightning を使用すると、単一 GPU、単一ノードのマルチ GPU、マルチノードのマルチ GPU の設定で、トレーニングスクリプトを実行できます。バックグラウンドでは、torch.distributed.launch と同様に複数のプロセスが自動的に起動されます。

単一ノードトレーニング (単一ノードのマルチ GPU を含む) の場合は、distributed_job_config を指定することなく、Azure Machine Learning でコードを実行できます。複数の GPU を含む複数のノードを使って実験を実行するには、2 つのオプションがあります。

PyTorch の設定を使用する(推奨): PyTorchConfiguration を定義し、communication_backend="Nccl"、node_count、process_count (これはプロセスの合計数、つまり num_nodes * process_count_per_node です) を指定します。 Lightning Trainer モジュールで、PyTorchConfiguration と一致するように num_nodes と gpus の両方を指定します。たとえば、num_nodes = node_count とgpus = process_count_per_node です。

MPI 構成を使用する:

MpiConfiguration を定義して、node_count と process_count_per_node の両方を指定します。 Lightning Trainer で、num_nodes と gpus の両方を、それぞれ MpiConfiguration の node_count および process_count_per_node と同じになるように指定します。
MPI によるマルチノードトレーニングの場合の Lightning には、トレーニングクラスターの各ノードで次の環境変数を設定する必要があります。
- MASTER_ADDR
- MASTER_PORT
- NODE_RANK
- LOCAL_RANK
メインのトレーニングスクリプトで、Lightning に必要な環境変数を手動で設定します。

import os
from argparse import ArgumentParser

def set_environment_variables_for_mpi(num_nodes, gpus_per_node, master_port=54965):
     if num_nodes > 1:
         os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"] = os.environ["AZ_BATCH_MASTER_NODE"].split(":")
     else:
         os.environ["MASTER_ADDR"] = os.environ["AZ_BATCHAI_MPI_MASTER_NODE"]
         os.environ["MASTER_PORT"] = str(master_port)

     try:
         os.environ["NODE_RANK"] = str(int(os.environ.get("OMPI_COMM_WORLD_RANK")) // gpus_per_node)
         # additional variables
         os.environ["MASTER_ADDRESS"] = os.environ["MASTER_ADDR"]
         os.environ["LOCAL_RANK"] = os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]
         os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"]
     except:
         # fails when used with pytorch configuration instead of mpi
         pass

if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("--num_nodes", type=int, required=True)
     parser.add_argument("--gpus_per_node", type=int, required=True)
     args = parser.parse_args()
     set_environment_variables_for_mpi(args.num_nodes, args.gpus_per_node)

     trainer = Trainer(
      num_nodes=args.num_nodes,
      gpus=args.gpus_per_node
  )

Lightning によって、Trainer フラグ --gpus と --num_nodes からのワールドサイズの計算が処理されます。

from azureml.core import ScriptRunConfig, Experiment
from azureml.core.runconfig import MpiConfiguration

nnodes = 2
gpus_per_node = 4
args = ['--max_epochs', 50, '--gpus_per_node', gpus_per_node, '--accelerator', 'ddp', '--num_nodes', nnodes]
distr_config = MpiConfiguration(node_count=nnodes, process_count_per_node=gpus_per_node)

run_config = ScriptRunConfig(
   source_directory='./src',
   script='train.py',
   arguments=args,
   compute_target=compute_target,
   environment=pytorch_env,
   distributed_job_config=distr_config,
)

run = Experiment(ws, 'experiment_name').submit(run_config)

Hugging Face Transformers

Hugging Face には、その Transformers ライブラリ torch.distributed.launch を使用して分散トレーニングを実行するための、多くの例が用意されています。 Transformers Trainer API を使用してこれらの例と独自のカスタムトレーニングスクリプトを実行するには、「torch.distributed.launch の使用」セクションに従います。

8 GPU の 1 ノードで run_glue.py スクリプトを使用して、テキスト分類 MNLI タスクで BERT ラージモデルを微調整するジョブ構成コードの例:

from azureml.core import ScriptRunConfig
from azureml.core.runconfig import PyTorchConfiguration

distr_config = PyTorchConfiguration() # node_count defaults to 1
launch_cmd = "python -m torch.distributed.launch --nproc_per_node 8 text-classification/run_glue.py --model_name_or_path bert-large-uncased-whole-word-masking --task_name mnli --do_train --do_eval --max_seq_length 128 --per_device_train_batch_size 8 --learning_rate 2e-5 --num_train_epochs 3.0 --output_dir /tmp/mnli_output".split()

run_config = ScriptRunConfig(
  source_directory='./src',
  command=launch_cmd,
  compute_target=compute_target,
  environment=pytorch_env,
  distributed_job_config=distr_config,
)

プロセス単位の起動オプションを使用して、torch.distributed.launch を使用せずに分散トレーニングを実行することもできます。この方法を使用する場合に注意することの 1 つは、Transformers の TrainingArguments には引数としてローカルランクを渡す必要があるということです (--local_rank)。 --use_env=False のときは、これは torch.distributed.launch によって処理されますが、プロセス単位の起動を使用している場合は、Azure Machine Learning によって設定されるのは LOCAL_RANK 環境変数だけなので、ユーザーはトレーニングスクリプト --local_rank=$LOCAL_RANK に引数としてローカルランクを明示的に渡す必要があります。

TensorFlow

TensorFlow 2.x の tf.distribute.Strategy API など、トレーニングコードでネイティブの TensorFlow を使用している場合は、TensorflowConfiguration を使用して Azure Machine Learning から分散ジョブを起動することができます。

そのためには、ScriptRunConfig コンストラクターの distributed_job_config パラメーターに TensorflowConfiguration オブジェクトを指定します。 tf.distribute.experimental.MultiWorkerMirroredStrategy を使用している場合は、TensorflowConfiguration で、トレーニングジョブのノード数に対応する worker_count を指定します。

from azureml.core import ScriptRunConfig, Environment, Experiment
from azureml.core.runconfig import TensorflowConfiguration

curated_env_name = 'AzureML-TensorFlow-2.3-GPU'
tf_env = Environment.get(workspace=ws, name=curated_env_name)
distr_config = TensorflowConfiguration(worker_count=2, parameter_server_count=0)

run_config = ScriptRunConfig(
  source_directory='./src',
  script='train.py',
  compute_target=compute_target,
  environment=tf_env,
  distributed_job_config=distr_config,
)

# submit the run configuration to start the job
run = Experiment(ws, "experiment_name").submit(run_config)

トレーニングスクリプトで、レガシの TensorFlow 1.x など、分散トレーニングのためにパラメーターサーバー戦略を使用している場合は、ジョブで使用するパラメーターサーバーの数も指定する必要があります (例: tf_config = TensorflowConfiguration(worker_count=2, parameter_server_count=1))。

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

azureml-examples: MultiWorkerMirroredStrategy を使用する分散 TensorFlow トレーニング

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 v1)