分散トレーニング

このページでは、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]`）	エバリュエータ

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

{
  "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 の仕組みについては、Vertex AI の Reduction Server を使用した分散 GPU トレーニングの高速化をご覧ください。

前提条件

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 を使用したトレーニング

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

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
カスタムトレーニングを実行する場所に最も近いマルチリージョンを選択すると、レイテンシが短縮される可能性があります。
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 の作成ガイドをご覧ください。

Reduction Server を使用したトレーニングのベストプラクティス

マシンタイプとカウント

Reduction Server のトレーニングでは、各ワーカーがすべてのレデューサホストに接続する必要があります。ワーカーホスト上の接続数を最小限に抑えるには、レデューサホストに最も高いネットワーク帯域幅のマシンタイプを使用します。

リデューサホストには、32 Gbps 下り（外向き）帯域幅を実現する 16 個以上の vCPU を備えた汎用 N1/N2 VM が最適です（n1-highcpu-16、n2-highcpu-16 など）。N1/N2 VM の Tier 1 VM 帯域幅は、下り（外向き）の最大帯域幅が 50 Gbps～100 Gbps の範囲で増加するため、リデューサ VM ノードに適しています。

ワーカーとレデューサの下り（外向き）帯域幅の合計は同じにする必要があります。たとえば、8 個の a2-megagpu-16g VM をワーカーとして使用する場合、少なくとも 25 台の n1-highcpu-16 VM をリデューサとして使用する必要があります。

`(8 worker VMs * 100&nbsp;Gbps) / 32&nbsp;Gbps egress = 25 reducer VMs`.

小さいメッセージをまとめてバッチ処理する

Reduction Server は、集計するメッセージのサイズが大きい場合に最適に動作します。ほとんどの ML フレームワークでは、all-reduce を実行する前に、小さな勾配テンソルをバッチ処理するさまざまな手法が提供されています。

Horovod

Horovod は、Tensor Fusion を使用して、all-reduce で小さなテンソルをバッチ処理します。バッファが完全にいっぱいになり、バッファに対する all-reduce オペレーションが実行されるまで、テンソルは Fusion バッファで処理されます。Fusion バッファのサイズは、HOROVOD_FUSION_THRESHOLD 環境変数を設定することで調整できます。

HOROVOD_FUSION_THRESHOLD 環境変数の推奨値は 128 MB 以上です。この場合、HOROVOD_FUSION_THRESHOLD 環境変数を 134217728（128 * 1024 * 1024）に設定します。

PyTorch

PyTorch の DistributedDataParallel はバッチメッセージをグラデーションバケットとしてサポートします。バッチバケットのサイズを制御するには、DistributedDataParallel コンストラクタで bucket_cap_mb パラメータを設定します。デフォルトのサイズは 25 MB です。

ベストプラクティス: bucket_cap_mb の推奨値は 64（64 MB）です。

クラスタの環境変数

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_SPEC と TF_CONFIG は一部の値を共有しますが、形式は異なります。どちらの環境変数にも、TensorFlow で必要とするもの以外のフィールドも含まれます。

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

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

TF_CONFIG フィールド

cluster

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

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

task

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

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

`task` フィールド
`type`	この VM で実行されているタスクのタイプ。この値はワーカーでは `worker`、パラメータサーバーでは `ps`、評価ツールでは `evaluator` に設定されます。ジョブのマスターワーカーでは、値は `chief` または `master` のいずれかに設定されます。この 2 つの違いについては、このドキュメントの `chief` 対 `master` のセクションをご覧ください。
`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 から各トライアルのトレーニングコードにさまざまな引数が渡されます。トレーニングコードでは、現在どのトライアルを実行中かということを必ずしも認識する必要はありません。さらに、Google Cloud コンソールでハイパーパラメータ調整ジョブの進行状況をモニタリングできます。

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

次のステップ

トレーニングパイプラインを作成する。

分散トレーニング

コードの必要条件

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

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

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

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

前提条件

Reduction Server を使用したトレーニング

Reduction Server を使用したトレーニングのベスト プラクティス

マシンタイプとカウント

小さいメッセージをまとめてバッチ処理する

Horovod

PyTorch

クラスタの環境変数

CLUSTER_SPEC の形式

CLUSTER_SPEC の例

TF_CONFIG の形式

TF_CONFIG の例

TF_CONFIG を使用するタイミング

分散トレーニング

ハイパーパラメータ調整

次のステップ

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

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

Reduction Server を使用したトレーニングのベストプラクティス

`CLUSTER_SPEC` の形式

`CLUSTER_SPEC` の例

`TF_CONFIG` の形式

`TF_CONFIG` の例

`TF_CONFIG` を使用するタイミング