このページでは、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 Gbps) / 32 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"
|
カスタム コンテナのクラスタの説明。 クラスタの説明には、指定した各ワーカープールのレプリカ名のリストが含まれます。 |
|
"workerpool0"
|
すべての分散トレーニング ジョブで、最初のワーカープール内にプライマリ レプリカが 1 つあります。 | |
"workerpool1"
|
ジョブの作成時に指定した場合、このワーカープールにワーカー レプリカが含まれます。 | |
"workerpool2"
|
ジョブの作成時に指定した場合、このワーカープールにパラメータ サーバーが含まれます。 | |
"workerpool3"
|
ジョブの作成時に指定した場合、このワーカープールにエバリュエータが含まれます。 | |
"environment"
|
文字 cloud 。 |
|
"task"
|
コードが実行されている特定のノードのタスクを記述します。この情報を使用して、分散ジョブ内の特定のワーカー用のコードを記述できます。このエントリは、次のキーを持つ辞書です。 | |
"type"
|
このタスクが実行されているワーカープールのタイプ。たとえば、"workerpool0" はプライマリ レプリカを参照します。 |
|
"index"
|
ゼロから開始するタスクのインデックス。たとえば、トレーニング ジョブに 2 つのワーカーが含まれている場合、この値は 1 つのワーカーの |
|
"trial"
|
現在実行中のハイパーパラメータ調整トライアルの ID。ジョブのハイパーパラメータ調整を構成する際に、トレーニングするトライアル数を構成します。この値により、コードで実行されているトライアルを識別できます。ID は、トライアル番号を含む文字列値で、1 から始まります。 | |
job |
現在のトレーニング ジョブを作成するために提供した |
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 つ以上のタスク名( これは |
||||||||||
task |
この環境変数が設定されている VM のタスクの説明。特定のトレーニング ジョブでは、この辞書は VM によって違います。この情報を使用して、分散トレーニング ジョブの各 VM で実行するコードをカスタマイズできます。また、この情報を使用して、ハイパーパラメータ調整ジョブのさまざまなトライアルを実行する際に、トレーニング コードの動作を変更することもできます。 この辞書には、次の Key-Value ペアが含まれています。
|
||||||||||
job |
現在のトレーニング ジョブを作成するために提供した |
||||||||||
environment |
文字 |
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
フィールドから、現在のトライアル番号を読み取ることができます。
次のステップ
- トレーニング パイプラインを作成する。