PyTorch 用の分散トレーニングの構成

このドキュメントでは、分散 PyTorch トレーニング ジョブを作成する方法を説明します。分散トレーニング ジョブを作成すると、AI Platform Training は、分散 PyTorch トレーニングをサポートする環境変数を使用して、仮想マシン（VM）インスタンスのクラスタ（ノードとも呼ばれる）でコードを実行します。これにより、トレーニング ジョブをスケーリングして大量のデータを処理できるようになります。

このガイドでは、PyTorch スタートガイドで説明されているように、トレーニングに事前に構築された PyTorch コンテナを使用することを前提としています。分散トレーニングに PyTorch コードを適応させるには、最小限の変更が必要です。

トレーニング クラスタ構造の指定

分散 PyTorch トレーニングでは、1 つのマスター ワーカーノードと 1 つ以上のワーカーノードを使用するようにジョブを構成します。これらのロールは、次のように動作します。

  • マスター ワーカー: ランクが 0 の VM。このノードはクラスタ内のノード間の接続を設定します。
  • ワーカー: クラスタ内の残りのノード。各ノードは、トレーニング アプリケーション コードで指定されているように、トレーニングの一部を実行します。

トレーニング クラスタのマスター ワーカー ノードとワーカーノードを指定する方法については、マシンタイプまたはスケール階層の指定をご覧ください。

コンテナ イメージの指定

トレーニング ジョブを作成するときに、マスター ワーカーの Docker コンテナのイメージを trainingInput.masterConfig.imageUri フィールドで使用するように指定し、各ワーカーの Docker コンテナのイメージを trainingInput.workerConfig.imageUri フィールドで使用するように指定します。事前に構築された PyTorch コンテナのリストをご覧ください。

gcloud ai-platform jobs submit training コマンドを使用してトレーニング ジョブを作成する場合、--master-image-uri フラグと --worker-image-uri フラグを使用してこれらのフィールドを指定できます。

ただし、trainingInput.workerConfig.imageUri フィールドを指定しない場合、値はデフォルトの trainingInput.masterConfig.imageUri になります。多くの場合、すべてのノードで同じ事前に構築された PyTorch コンテナを使用します。

トレーニング コードの更新

トレーニング アプリケーションで次のコードを追加して、トレーニング クラスタを初期化します。

import torch

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

BACKEND は、サポートされている分散トレーニング バックエンドのいずれかに置き換えます（次のセクションを参照）。init_method='env://' キーワード引数は、環境変数を使用してクラスタ内の通信を初期化するよう PyTorch に指示します。詳しくは、このガイドの環境変数セクションをご覧ください。

また、torch.nn.parallel.DistributedDataParallel クラスを使用するようにトレーニング コードを更新します。たとえば、コード内に model という PyTorch モジュールを作成した場合は、次の行を追加します。

model = torch.nn.parallel.DistributedDataParallel(model)

分散トレーニングの構成の詳細については、PyTorch のドキュメントの分散トレーニング ガイドをご覧ください。

分散トレーニングのバックエンド

AI Platform Training では、分散 PyTorch のトレーニング用に次のバックエンドがサポートされます。

  • gloo: CPU トレーニング ジョブで推奨
  • nccl: GPU トレーニング ジョブで推奨

バックエンドの違いをご覧ください。

環境変数

分散 PyTorch トレーニング ジョブを作成すると、AI Platform Training は各ノードに次の環境変数を設定します。

  • WORLD_SIZE: クラスタ内のノードの合計数。この変数の値は、すべてのノードで同じです。
  • RANK: 各ノードの一意の識別子。マスター ワーカーでは、この識別子は 0 に設定されます。各ワーカーでは、1 から WORLD_SIZE - 1 までの異なる値に設定されます。
  • MASTER_ADDR: マスター ワーカー ノードのホスト名。この変数の値は、すべてのノードで同じです。
  • MASTER_PORT: マスター ワーカー ノードが通信するポート。この変数の値は、すべてのノードで同じです。

PyTorch ではこれらの環境変数を使用してクラスタを初期化します。

次のステップ