分散トレーニング

このページでは、Vertex AI で分散トレーニング ジョブを実行する方法について説明します。

コードの必要条件

分散トレーニングをサポートする ML フレームワークを使用します。トレーニング コードでは、CLUSTER_SPEC または TF_CONFIG 環境変数を使用して、トレーニング クラスタの特定の部分を参照できます。

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

Vertex AI で分散トレーニング ジョブを実行する場合、トレーニング クラスタで複数のマシン(ノード)を指定します。トレーニング サービスにより、指定したマシンタイプに応じてリソースが割り当てられます。特定のノード上で実行中のジョブは「レプリカ」と呼ばれます。同じ構成のレプリカのグループはワーカープールと呼ばれます。

トレーニング クラスタ内の各レプリカには、分散トレーニングにおけるロールまたはタスクが 1 つ割り当てられます。次に例を示します。

  • プライマリ レプリカ: 1 つのレプリカのみがプライマリ レプリカに指定されます。このタスクは他のレプリカを管理し、ジョブ全体のステータスを報告します。

  • ワーカー: 1 つ以上のレプリカを「ワーカー」として指定できます。これらのレプリカに担当させる作業は、ジョブ構成で指定します。

  • パラメータ サーバー: ML フレームワークでサポートされている場合、1 つ以上のレプリカをパラメータ サーバーとして指定できます。これらのレプリカは、モデル パラメータを格納し、ワーカー間で共有されるモデルの状態を調整します。

  • エバリュエータ: ML フレームワークでサポートされている場合は、1 つ以上のレプリカを「エバリュエータ」に指定します。これらのレプリカを使用して、モデルを評価できます。なお、TensorFlow を使用している場合、TensorFlow では通常、エバリュエータを 1 つだけ使用することを想定します。

分散トレーニング ジョブを構成する

カスタム トレーニング ジョブを分散トレーニング ジョブとして構成するには、複数のワーカープールを定義します。トレーニング パイプラインまたはハイパーパラメータ調整ジョブで分散トレーニングを実行することもできます。

分散トレーニング ジョブを構成するには、ワーカープール リスト(workerPoolSpecs[])を定義し、タスクの種類ごとに 1 つの WorkerPoolSpec を指定します。

workerPoolSpecs[] の位置 クラスタで実行されたタスク
最初(workerPoolSpecs[0] メイン、チーフ、スケジューラ、マスター
2 番目(workerPoolSpecs[1] セカンダリ、レプリカ、ワーカー
3 番目(workerPoolSpecs[2] パラメータ サーバー、Reduction Server
4 番目(workerPoolSpecs[3] エバリュエータ

プライマリ レプリカを指定する必要があります。プライマリ レプリカは、他のすべてのレプリカによって行われる作業を調整します。最初のレプリカの仕様をプライマリ レプリカにのみ使用し、replicaCount1 に設定します。

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, primary replica, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {},
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

追加のワーカープールを指定する

ML フレームワークによっては、他の目的で追加のワーカープールを指定することもできます。たとえば、TensorFlow を使用している場合は、ワーカー レプリカ、パラメータ サーバー レプリカ、エバリュエータ レプリカを構成するようにワーカープールを指定できます。

workerPoolSpecs[] リストに指定するワーカープールの順序によって、ワーカープールのタイプが決まります。使用しないワーカープールに対して空の値を設定します。このようにすると、使用しないワーカープールが workerPoolSpecs[] リストでスキップされるため、使用するワーカープールを指定できます。次に例を示します。

プライマリ レプリカとパラメータ サーバー ワーカープールのみが含まれるジョブを指定する場合は、ワーカープール 1 に空の値を設定する必要があります。

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

Reduction Server を使用してトレーニング時間を短縮する

複数のノードを使用して大規模な ML モデルをトレーニングする場合、ノード間で勾配を伝達すると大きなレイテンシが発生する可能性があります。Reduction Server は、分散トレーニングのスループットを向上させ、レイテンシを短縮できる all-reduce アルゴリズムです。Vertex AI は、分散トレーニング中にワーカープールの 1 つに使用可能な Docker コンテナ イメージで Reduction Server を使用できるようにします。

Reduction Server を使用するには、次の要件を満たす必要があります。

  • GPU ワーカーで分散トレーニングを実行している。

  • トレーニング コードは TensorFlow または PyTorch を使用し、NCCL all-reduce を使用する GPU を備えたマルチホスト データ並列型トレーニングで構成されている(また、NCCL を使用する他の ML フレームワークを使用できる場合もあります)。

  • プライマリ ノード(workerPoolSpecs[0])とワーカー(workerPoolSpecs[1])で実行されているコンテナが Reduction Server をサポートしている。具体的には、各コンテナは次のいずれかになります。

    • ビルド済みの TensorFlow トレーニング コンテナ、バージョン 2.3 以降。

    • ビルド済みの Pytorch トレーニング コンテナ、バージョン 1.4 以降。

    • NCCL 2.7 以降と google-reduction-server パッケージがインストールされたカスタム コンテナ。このパッケージをカスタム コンテナ イメージにインストールするには、Dockerfile に次の行を追加します。

      RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list && \
          curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
          apt update && apt install -y google-reduction-server
      

Reduction Server を使用するには、カスタム トレーニング リソースの作成時に次の操作を行います。

  1. 3 番目のワーカープール(workerPoolSpecs[2])の containerSpec.imageUri フィールドで、次のいずれかの URI を指定します。

    • us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • europe-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • asia-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

    カスタム トレーニングを実行する場所に最も近いマルチリージョンを選択すると、レイテンシが短縮される可能性があります。

  2. 3 番目のワーカープールのマシンタイプとノード数を選択するときは、3 番目のワーカープールのネットワーク帯域幅の合計が、最初と 2 番目のワーカープールのネットワーク帯域幅の合計と一致するか、その値を超えていることを確認してください。

    2 番目のワーカープールの各ノードで利用できる最大帯域幅については、ネットワーク帯域幅と GPU をご覧ください。

    Reduction Server ノードには GPU を使用しません。3 番目のワーカープールの各ノードで使用可能な最大帯域幅については、汎用マシン ファミリーの下り(外向き)帯域幅の上限(Gbps)の列をご覧ください。

    たとえば、5 つの n1-highmem-96 ノードを使用するように最初と 2 番目のワーカープールを構成し、それぞれ 8 個の NVIDIA_TESLA_V100 GPU を持っている場合、各ノードで使用できる最大帯域幅は 100 Gbps、合計帯域幅は 500 Gbps になります。3 番目のワーカープールでこの帯域幅を一致させるには、16 個の n1-highcpu-16 ノードを使用します。各ノードの最大帯域幅は 32 Gbps、合計帯域幅は 512 Gbps になります。

    Reduction Server ノードには、n1-highcpu-16 マシンタイプを使用することをおすすめします。このマシンタイプを使用すると、そのリソースに対して比較的高い帯域幅を提供できます。

次のコマンドでは、Reduction Server を使用する CustomJob リソースを作成しています。

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=4,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highcpu-16,replica-count=16,container-image-uri=us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

詳細については、CustomJob の作成ガイドをご覧ください。

クラスタの環境変数

Vertex AI はすべてのレプリカで、環境変数 CLUSTER_SPEC に値を代入します。これは、クラスタ全体がどのように設定されているかを示す環境変数です。TensorFlow の TF_CONFIG と同様に、CLUSTER_SPEC にはクラスタ内の各レプリカの詳細、たとえばインデックスとロール(プライマリ レプリカ、ワーカー、パラメータ サーバー、または評価者)に関する情報が記述されます。

TensorFlow で分散トレーニングを実行すると、TF_CONFIG が解析されて tf.train.ClusterSpec が作成されます。同様に、他の ML フレームワークで分散トレーニングを実行するときは、CLUSTER_SPEC を解析して、フレームワークに必要なすべての環境変数や設定に値を代入する必要があります。

CLUSTER_SPEC の形式

CLUSTER_SPEC 環境変数は、次の形式の JSON 文字列です。

キー 説明
"cluster"

カスタム コンテナのクラスタの説明。TF_CONFIG と同様に、このオブジェクトは TensorFlow クラスタ仕様として形式設定され、tf.train.ClusterSpec のコンストラクタに渡すことができます。

クラスタの説明には、指定した各ワーカープールのレプリカ名のリストが含まれます。

"workerpool0" すべての分散トレーニング ジョブで、最初のワーカープール内にプライマリ レプリカが 1 つあります。
"workerpool1" ジョブの作成時に指定した場合、このワーカープールにワーカー レプリカが含まれます。
"workerpool2" ジョブの作成時に指定した場合、このワーカープールにパラメータ サーバーが含まれます。
"workerpool3" ジョブの作成時に指定した場合、このワーカープールにエバリュエータが含まれます。
"environment" 文字 cloud
"task" コードが実行されている特定のノードのタスクを記述します。この情報を使用して、分散ジョブ内の特定のワーカー用のコードを記述できます。このエントリは、次のキーを持つ辞書です。
"type" このタスクが実行されているワーカープールのタイプ。たとえば、"workerpool0" はプライマリ レプリカを参照します。
"index"

ゼロから開始するタスクのインデックス。たとえば、トレーニング ジョブに 2 つのワーカーが含まれている場合、この値は 1 つのワーカーの 0 と、もう 1 つのワーカーの 1 に設定されます。

"trial" 現在実行中のハイパーパラメータ調整トライアルの ID。ジョブのハイパーパラメータ調整を構成する際に、トレーニングするトライアル数を構成します。この値により、コードで実行されているトライアルを識別できます。ID は、トライアル番号を含む文字列値で、1 から始まります。
job

現在のトレーニング ジョブを作成するために提供した CustomJobSpec で、辞書として表現されます。

CLUSTER_SPEC の例

値の例は次のとおりです。


{
   "cluster":{
      "workerpool0":[
         "cmle-training-workerpool0-ab-0:2222"
      ],
      "workerpool1":[
         "cmle-training-workerpool1-ab-0:2222",
         "cmle-training-workerpool1-ab-1:2222"
      ],
      "workerpool2":[
         "cmle-training-workerpool2-ab-0:2222",
         "cmle-training-workerpool2-ab-1:2222"
      ],
      "workerpool3":[
         "cmle-training-workerpool3-ab-0:2222",
         "cmle-training-workerpool3-ab-1:2222",
         "cmle-training-workerpool3-ab-2:2222"
      ]
   },
   "environment":"cloud",
   "task":{
      "type":"workerpool0",
      "index":0,
      "trial":"TRIAL_ID"
   },
   "job": {
      ...
   }
}

TF_CONFIG の形式

Vertex AI は、CLUSTER_SPEC に加えて、分散トレーニング ジョブの各レプリカに TF_CONFIG 環境変数を設定します。Vertex AI は、単一レプリカ トレーニング ジョブに TF_CONFIG を設定しません。

CLUSTER_SPECTF_CONFIG は一部の値を共有しますが、形式は異なります。どちらの環境変数にも、TensorFlow で必要とするもの以外のフィールドも含まれます。

TensorFlow を使用した分散トレーニングは、ビルド済みコンテナを使用する場合と同じカスタム コンテナを使用できます。

TF_CONFIG 環境変数は、次の形式の JSON 文字列です。

TF_CONFIG フィールド
cluster

TensorFlow クラスタの説明。1 つ以上のタスク名(chiefworkerps、または master)を、タスクが実行されているネットワーク アドレスのリストにマッピングする辞書。特定のトレーニング ジョブでは、この辞書はすべての VM で同じです。

これは tf.train.ClusterSpec コンストラクタの有効な 1 番目の引数です。エバリュエータをジョブで使用していても、エバリュエータがトレーニング クラスタの一部と見なされることはないため、この辞書に evaluator がキーとして含まれることはありません。

task

この環境変数が設定されている VM のタスクの説明。特定のトレーニング ジョブでは、この辞書は VM によって違います。この情報を使用して、分散トレーニング ジョブの各 VM で実行するコードをカスタマイズできます。また、この情報を使用して、ハイパーパラメータ調整ジョブのさまざまなトライアルを実行する際に、トレーニング コードの動作を変更することもできます。

この辞書には、次の Key-Value ペアが含まれています。

task フィールド
type

この VM で実行されているタスクのタイプ。この値はワーカーでは worker、パラメータ サーバーでは ps、評価ツールでは evaluator に設定されます。ジョブのマスター ワーカーでは、値は chief または master のいずれかに設定されます。この 2 つの違いについては、このドキュメントの chiefmaster のセクションをご覧ください。

index

ゼロから開始するタスクのインデックス。たとえば、トレーニング ジョブに 2 つのワーカーが含まれている場合、この値は 1 つのワーカーの 0 と、もう 1 つのワーカーの 1 に設定されます。

trial

この VM で現在実行中のハイパーパラメータ調整トライアルの ID。このフィールドは、現在のトレーニング ジョブがハイパーパラメータ調整ジョブである場合にのみ設定されます。

ハイパーパラメータ調整ジョブの場合、Vertex AI によって毎回さまざまなハイパーパラメータで多数のトライアルが行われ、トレーニング コードが繰り返し実行されます。このフィールドには現在のトライアル番号が含まれます。トライアル番号は最初のトライアルで 1 から始まります。

cloud

Vertex AI で内部的に使用される ID。このフィールドは無視できます。

job

現在のトレーニング ジョブを作成するために提供した CustomJobSpec で、辞書として表現されます。

environment

文字 cloud

TF_CONFIG の例

次のサンプルコードでは、TF_CONFIG 環境変数をトレーニング ログに出力します。

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

ハイパーパラメータ調整ジョブがランタイム バージョン 2.1 以降で、マスター ワーカー、2 つのワーカー、パラメータ サーバーを使用している場合、このコードでは最初のハイパーパラメータ調整トライアル中に、ワーカーのいずれかで次のログを出力します。出力例では、簡潔にするために job フィールドを非表示にし、一部の ID を汎用値に置き換えています。

{
  "cluster": {
    "chief": [
      "training-workerpool0-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "training-workerpool2-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "training-workerpool1-[ID_STRING_1]-0:2222",
      "training-workerpool1-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

TF_CONFIG を使用するタイミング

TF_CONFIG は分散トレーニング ジョブにのみ設定されます。

トレーニング コードで直接 TF_CONFIG 環境変数を操作する必要はありません。TF_CONFIG 環境変数にアクセスするのは、TensorFlow の分散戦略と Vertex AI の標準のハイパーパラメータ調整ワークフロー(両方とも以降のセクションで説明します)がジョブで機能しない場合のみです。

分散トレーニング

Vertex AI では、TensorFlow が分散トレーニングに必要とする仕様を拡張するように、TF_CONFIG 環境変数を設定します。

TensorFlow を使用して分散トレーニングを実行するには、tf.distribute.Strategy API を使用します。特に、Keras API を MultiWorkerMirroredStrategy または ParameterServerStrategyジョブのパラメータ サーバーを指定する場合)と組み合わせて使用することをおすすめします。ただし、TensorFlow は現在、これらの戦略に対して試験運用版のサポートのみを提供していることに注意してください。

これらの分散戦略では、環境変数 TF_CONFIG を使用して、トレーニング ジョブの各 VM にロールを割り当て、VM 間の通信を容易にします。TensorFlow が環境変数 TF_CONFIG を処理するため、トレーニング コードで直接アクセスする必要はありません。

環境変数 TF_CONFIG を直接解析するのは、トレーニング ジョブを実行する各 VM の動作をカスタマイズする場合のみです。

ハイパーパラメータ調整

ハイパーパラメータ調整ジョブを実行すると、Vertex AI から各トライアルのトレーニング コードにさまざまな引数が渡されます。トレーニング コードでは、現在どのトライアルを実行中かということを必ずしも認識する必要はありません。さらに、Cloud Console でハイパーパラメータ調整ジョブの進行状況をモニタリングできます。

必要に応じて、コードで TF_CONFIG 環境変数の task フィールドの trial フィールドから、現在のトライアル番号を読み取ることができます。

次のステップ