AI ハイパーコンピュータにおける GPU ベースの分散トレーニング ワークロードのオーケストレーション

※この投稿は米国時間 2024 年 12 月 14 日に、Google Cloud blog に投稿されたものの抄訳です。

AI に関して、大規模言語モデル（LLM）と ML が業界全体を次のレベルに引き上げています。しかしモデルやデータセットが大規模になると、開発者はモデルを効率的にトレーニングするために、複数のコンピューティング ホストで複数の AI アクセラレータ（GPU、TPU など）にまたがる分散環境が必要になります。そのため、オーケストレーション、リソース管理、スケーラビリティの問題が生じることがあります。

そこで Google Cloud では、AI ハイパーコンピュータ アーキテクチャの一部として、堅牢な一連の GPU および TPU リソースと高度なオーケストレーション ツールを提供し、分散型の大規模なトレーニングを簡素化しています。このブログでは、Google Cloud の GPU アクセラレータで利用可能なオーケストレーション ツールについて説明し、ML ワークフローの合理化とスケーリングを行う方法を紹介します。

適切なアクセラレータ ファミリーを選択する

分散トレーニングでは、適切な GPU を選択することが重要です。Google Cloud の特化型のマシン ファミリーは、パフォーマンスや費用対効果に関するさまざまなニーズに合わせたソリューションを提供しています。NVIDIA H100 GPU と NVIDIA H200 GPU（予定）を搭載した A3 マシン シリーズは、GPU 間帯域幅が優れており、大規模なトレーニング ワークフローに最適です。対照的に、NVIDIA A100 GPU を搭載した A2 マシン シリーズは、合理化されたシングルノード トレーニングなど、ノード間通信を最小限に抑える必要のあるシナリオを対象として設計されています。さらに、NVIDIA L4 GPU を搭載した汎用の G2 マシン ファミリーでは、推論やテストのワークロードに必要な柔軟性が得られます。

また、大規模なトレーニングのニーズを満たすため、複数の GPU 利用モデルが用意されています。

• 確約利用割引（CUD）: 長期的な利用契約を購入いただくことで、費用が大幅に削減され、容量は保証されます。

• Dynamic Workload Scheduler（DWS）: 2 つのモードがあり、保証が必要なワークロードや、開始時間に柔軟性があるワークロードなど、さまざまなワークロード向けに設計されています。容量は一定期間利用可能で、手頃な正規価格が設定されています。

• オンデマンド利用: 最も柔軟です。前払いは不要ですが、容量の可用性は保証されません。

• Spot VM: 費用を大幅に削減できますが、プリエンプティブルです。復元力があり停止が許容されるジョブ設計が必要です。

分散トレーニングをさらに加速するために、Google Cloud におけるパワフルな 3 つのオーケストレーション戦略、すなわち Google Kubernetes Engine（GKE）、Cluster Toolkit、Vertex AI カスタム トレーニング パイプラインを詳しく見てみましょう。それぞれのアプローチに独自の強みがあり、Google Cloud のパワフルなインフラストラクチャを活用して、ML プロジェクトを迅速かつスケーラブルに推進できます。

Google Cloud の高度なオーケストレーション ツールが、リソースの最適化、複雑性の低減、ML の取り組みにおける優れたパフォーマンスの実現にどのように役立つのかについて理解を深めていただけるよう、それぞれのオプションを詳しく説明します。

オプション 1: GKE による統合ワークロード管理

堅牢なプラットフォーム チームを持つ企業は多くの場合、管理をシンプルにするために、カスタム トレーニングを含むすべてのワークロードを実行する統合環境を必要としています。その点では GKE が適しています。GKE の柔軟性とスケーラビリティにより、単一のプラットフォーム上で多様なワークロードが処理されます。プラットフォーム チームは GKE を使用することで一元管理し、可視性を確保しながら、リソース使用率を最適化して管理を合理化できます。

GKE で実行される ML ワークロードをオーケストレートする方法は次のとおりです。

1. GKE クラスタとノードプールのプロビジョニング予約（CUD または DWS カレンダー）があり、Terraform を使用する場合は、クラスタ プロビジョニング テンプレートの手順に沿って操作し、パラメータ ファイル（terraform.tfvars）を指定します。

次に以下のコマンドを実行して、GKE クラスタとノードプールをプロビジョニングします。

さらに、Cluster Toolkit には A3 または A3 Mega GKE クラスタとノードプールをプロビジョニングするための、Terraform ベースのブループリントのサンプルが含まれています。

gcloud コマンドを使用する場合は、こちらのチュートリアルに沿って操作し、A3 / A3 Mega VM で GKE クラスタとノードプールを作成します。

DWS Flex の場合は、次の gcloud コマンドを使用することで DWS 対応のノードプールを作成できます。

2. GPU と A3 TCPX / A3 Mega TCPXO の直接通信を有効にして初期ベンチマーク テストを実施する以下の手順に沿って GPUDirect for TCPX / TCPXO ライブラリのインストールと NCCL の構成を行い、テスト ワークロードをデプロイして、初期ベンチマーク テストを実施します。

2 つの A3 Mega ノードについて、allgather ベンチマーク テストの出力を検証します。

上に示したベンチマーク出力の表において、最初の列はメッセージ サイズ、右側の algbw 列と busbw 列は GPU あたりの帯域幅です。通常は、メッセージ サイズが最大のイン / アウトオブ プレイス busbw 列（ハイライト表示された列）を使用してノード間の帯域幅を求めます。A3 Mega ノードの場合、GPU あたり 185～190 GB/s の範囲になると予想されます。これは、8 個の NVIDIA H100 GPU と 8 個の NIC による A3 Mega ノードに関し、ノード間ネットワーク帯域幅が 1,600 gbps に近いことを示している可能性があります。

NCCL テストをノード数 2 から 8、16、32 などに拡張することで、ノード間のネットワーク パフォーマンスが適切な範囲内にあることと、すべてのノードが健全であることを確認できます。

3. 分散トレーニング バッチ ワークロードを構成する統合 API によって k8s ジョブのグループをユニットとして管理するための Kubernetes ネイティブ API である JobSet を使用すると、分散 HPC（MPI など）と AI / ML トレーニング ワークロード（PyTorch、Jax、Tensorflow など）を Kubernetes にデプロイできます。

次の例は、GPUDirect-TCPX を使用した A3 の JobSet yaml マニフェストを示しています。

a. 主な JobSet 構成要素

b. トレーニング ジョブ設定（pytorch メインコンテナを含む）

c. gcsfuse、tcpx（A3 high）、tcpxo（A3 Mega）RxDM コンテナ

d. NCCL 環境変数

DWS バッチ ワークロードについては、Kueue と JobSet の設定を統合した A3 Mega ベースの例を参照してください。

最後に、A3 Mega で Megatron LM（Llama2）トレーニングを行う方法については、Helmchart の例を参照してください。

オプション 2: Cluster Toolkit 経由の Slurm

Slurm は広く普及しているハイ パフォーマンス コンピューティング（HPC）ジョブ スケジューラです。学術界や産業界の研究者によって使用されている Slurm は、使い慣れたセマンティクスで LLM トレーニングのオーケストレーションを行うための堅牢なソリューションを提供します。Google Cloud における Slurm は Cluster Toolkit（旧称 Cloud HPC Toolkit）によってサポートされます。これは Google Cloud に HPC、AI、ML のワークロードをデプロイするプロセスを簡素化するオープンソース ソフトウェアです。高度なカスタマイズと拡張が可能で、大規模な LLM トレーニング向けのインフラストラクチャのデプロイなど、さまざまなユースケースのデプロイニーズに対応するように設計されています。

1. A3-high クラスタと A3-mega クラスタのプロビジョニング公開ドキュメントの構成手順を参照して Cluster Toolkit をインストールします。サポートされている Go、Terraform、Packer のバージョンなど、前提条件に注意してください。

ダウンロードした github リポジトリを含め、Cluster Toolkit のインストールが完了したら、ブループリントのディレクトリ examples/machine-learning に移動します。このディレクトリには、A3 シリーズのマシンシェイプに基づいて H100 クラスタをデプロイするための 2 つのフォルダ、a3-highgpu-8g と a3-megagpu-8g があります。この例では、a3-megagpu-8g フォルダのブループリントを取り上げます。

Google Cloud Cluster Toolkit のブループリントは、デプロイするインフラストラクチャを記述した Infrastructure as Code（IaC）ドキュメントであり、概念としては Terraform や他の IaC ツールと似ています。a3-megagpu-8g ブループリントの場合、デプロイを制御するメインのファイルが 3 つあります。

1. slurm-a3mega-base.yaml - クラスタノード上の共通のホーム ファイルシステムに使用される Filestore インスタンスとともに、必要な VPC ネットワークを作成

2. slurm-a3mega-image.yaml - クラスタの定義に基づいてノードをプロビジョニングするために Slurm によって使用される、Compute Engine イメージ インスタンスを作成

3. slurm-a3mega-cluster.yaml - Slurm コントローラ（Slurm ジョブのメインのオーケストレーター）、Slurm ログインノード（ジョブ送信に使用されるホスト）、a3mega パーティション（クラスタ内の作業ノード）を含む、メインのクラスタ コンポーネントを設定

必要に応じて各ブループリント コンポーネントをカスタマイズできますが、deployment-base.yaml と deployment-image-cluster.yaml で作業環境の詳細を指定するだけで簡単に始めることができます。

2. GPUDirect-TCPXO に最適化された NCCL 通信を有効にするSlurm クラスタを作成したら、GPU ネットワーク上の NCCL 通信を最適化するために、チュートリアルに沿って GPUDirect-TCPXO を有効にします。環境を検証し、TCPXO プラグインが適切に読み込まれていることを確認するには、NCCL テストをビルドしてコンパイルします。次に sbatch run-nccl-tests.sh をログインノードから実行します。スクリプト内のノード数は、クラスタのノード数に合わせて変更してください。これにより、スクリプトで指定された GPU とノードにわたって分散 all_gather テストが実行されます。

意図したとおりに機能している場合、NCCL テストは、さまざまなメッセージ サイズで高速な帯域幅スループットを示す結果を出力するはずです。パフォーマンスの測定では一般的に、メッセージ サイズ値 4 Gb または 8 Gb を示す、出力表の 2 行目または最終行の busbw 値（GB/s）を使用します。TCPXO が有効なクラスタは、約 190 GB/s の busbw スループットをレポートするはずです。これらの指標について詳しくは、NVIDIA NCCL-tests リポジトリのパフォーマンスに関するページをご覧ください。

3. NeMo トレーニング ワークロードを実行するNeMo トレーニングのチュートリアルに沿って、以下の手順でサンプルの NeMo フレームワーク事前トレーニング ジョブを実行します。

ステップ 1:

• 必要な TCPXO 環境変数を設定した NeMo フレームワーク由来のコンテナを作成します

• フレームワーク ランチャー スクリプトとその他の補助的なファイルを作業ディレクトリにコピーする Slurm ジョブを発行します

ステップ 2:

• Python 仮想環境を構築し、NeMo フレームワークの Python パッケージの依存関係をインストールします

ステップ 3:

• このコマンドは、モックデータを入力として使用し、8 つのノードにわたって 5B パラメータの GPT3 モデルの分散トレーニングを 10 ステップ実行します

オプション 3: Vertex AI

マネージド インフラストラクチャの機能と、Llama 3.1 や Mistral などの主要なオープンモデルへのアクセスを求めるチームにとって、Vertex AI Model Garden とカスタム トレーニング ジョブ サービスは魅力的なオプションです。このフルマネージド ソリューションによって、オーケストレーションの負担がほとんどなくなり、エンドツーエンドの ML プラットフォーム運用が実現するため、開発者はモデルの開発とテストに注力できます。Vertex AI のエンドツーエンドのトレーニング サポートは、データの準備からデプロイまでのワークフローを統合し、プロセスをさらに簡素化します。

それでは、Vertex でシングルノードまたはマルチノードのファインチューニング / トレーニング ワークロードを実行する方法を見てみましょう。   

Vertex おけるシングルノード マルチ GPU のファインチューニング / トレーニングこちらのノートブックでは、Vertex AI SDK を使用した Llama 3.1 モデルのファインチューニングとデプロイを説明しています。このノートブックの例はすべて、Low-Rank Adaptation（LoRA）とパラメータ効率の高いファインチューニング（PEFT）手法を使用して、トレーニングとストレージの費用を削減しています。LoRA は PEFT のアプローチです。事前トレーニングされたモデルの重みを固定し、ファインチューニング中に、モデルの重みの変化を表すランク分解行列をトレーニングします。

Vertex AI におけるマルチノード分散ファインチューニング / トレーニング

Vertex のサンプルのトレーニング リポジトリでは、Vertex の A3 Mega（8 個の NVIDIA H100）でマルチノード分散トレーニングを開始する方法が例示されています。

NeMo の例では、事前トレーニング、継続的な事前トレーニング、教師ありファインチューニング（SFT）の方法が示されています。さらに NeMo では、AI アクセラレータ（この場合は A3 Mega）を評価するためによく使用されるアプローチとして、最適化トレーニングを行うことができます。ベンチマークには、エポック時間やステップ時間など、報告された指標を使用できます。NeMo はほとんどの NVIDIA GPU タイプで動作するため、所定のタスクに対して異なる AI チップを比較する際に役立ちます。A3 Mega ノードタイプを使用して Vertex でサンプルを実行する方法について、詳細をご確認ください。

launch.sh は、コマンド パラメータを指定して NeMo 分散トレーニングを開始するためのメインのエントリ ポイントです。

例:

launch.sh スクリプトの最後で、curl コマンドを使用して Vertex customJobs API を呼び出し、Vertex で NeMo トレーニング ジョブを開始します。

vertex-payload.json のジョブ構成は Nemo トレーニングを開始する curl コマンドの一部です。以下のように、リソース要件に関するジョブ仕様が含まれています。

ジョブ構成引数の「${TRANSFER_MODEL_CMD} ${LAUNCH_CMD}」は、ジョブ トレーニング スクリプトの全コンテンツを順に埋め込みます。これには A3 Mega で必要なすべての NCCL 環境も含まれます。一方、他の pytorch launch コマンドは Vertex CustomJob によって実行されます。

また、vertex-payload.json の「imageUri」パラメータとして、独自のカスタム ジョブコンテナを構築することもできます。その場合は Dockerfile を参照してください。

DIY 派の場合: カスタム トレーニング環境を構築する

多くの組織は、実践的なアプローチを好み、特定のオーケストレーション ツールやフレームワークを使用したいと考えるものです。そのような場合、Google Compute Engine は、独自のトレーニング環境を構築する基盤を提供します。開発者は、GPU、CPU、メモリ、ストレージの種類や数など、希望する仕様で仮想マシン（VM）を作成して構成できます。また、このきめ細かな管理機能により、特定のトレーニング ワークロードに合わせてインフラストラクチャを最適化し、好みのオーケストレーション ツールを統合できます。

このプロセスを容易にするためにサンプルのコード スニペットが用意されており、gcloud compute instance create と gcloud compute instance bulk create の API 呼び出しを使用して標準の A3 Mega インスタンスを作成、管理する方法が示されています。単一の VM を作成する必要がある場合でも、大規模なクラスタをプロビジョニングする必要がある場合でも、これらのリソースはインフラストラクチャのセットアップを合理化するために役立ちます。

まとめ

適切なオーケストレーション戦略に基づき、Google Cloud の堅牢かつ最先端の AI インフラストラクチャを活用すれば、トレーニング目標を達成し、ビジネス目標を実現できます。

分散トレーニングについて詳しくは、GKE の例、Cluster Toolkit の例、Vertex AI の例をご覧ください。

ー AI インフラストラクチャ フィールド ソリューション アーキテクト Rick（Rugui）Chen

ー AI インフラストラクチャ フィールド ソリューション アーキテクト Rob Vogelbacher