少ない費用で AI をトレーニング: 弾力性のあるトレーニングとチェックポイント処理の最適化で ML グッドプットを改善

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

大規模な AI トレーニングにかかる費用を削減したいと考える組織は多いことでしょう。数千ものアクセラレータを使って数週間にわたり実施される一般的な PyTorch LLM トレーニングのワークロードの場合、ML グッドプットを 1% 向上させると、100 万ドル以上の費用削減につながる可能性があります1。したがって、ML グッドプットの向上は、効率性とモデルの反復処理速度の両方の観点から、モデル トレーニングにおける重要な目標となります。

しかし、ML のグッドプットを向上させるには、いくつかの課題があります。たとえば、中断によって最新のチェックポイントからの再起動が頻繁に発生する、インラインのチェックポイント処理の遅延によりトレーニングが中断される、オブザーバビリティの制限が理由で障害が検出しづらい、といった課題が挙げられます。こうした問題が、製品化までの時間（TTM）とトレーニング コストの大幅な増加につながっています。Arxiv で公開されているこちらの論文を含め、複数の業界出版物がこれらの課題を取り上げています。

ML グッドプットの向上

ML グッドプットを向上させるには、中断を伴うイベントがトレーニング ワークロードの進行に与える影響を最小限に抑える必要があります。ジョブをすばやく再開するには、ジョブを自動的にスケールダウンするか、障害が発生したリソースと予備容量のリソースをスワップします。Google Cloud では、これを弾力性のあるトレーニングと呼んでいます。Google Cloud ではさらに、チェックポイント処理時のワークロードの中断を減らせるほか、障害発生時に最も近い利用可能なストレージのロケーションからチェックポイントの読み込みを行う際の読み込み速度を向上できます。これらの機能を非同期チェックポイント処理と多層チェックポイント処理と呼びます。

以下は、これらの手法が提供するエンドツーエンドの修復ワークフローを使ってトレーニングの ML グッドプットを向上させる流れを示した図です。図中には、9 つのノードからなるワークロードに対して 3 方向のデータ並列処理（DP）と 3 方向のパイプライン並列処理（PP）を行う例が示されています。また、障害の種類と予備容量に応じた複数の修復アクションも記載されています。

特定のワークロードの修復ポリシーをカスタマイズできます。たとえば、ホットスワップとスケールダウンのどちらの修復戦略を採用するのか選択したり、チェックポイント処理の頻度を構成したりすることが可能です。スーパーバイザー プロセスは、診断サービスから障害、パフォーマンス低下、ストラグラーのシグナルを受信します。スーパーバイザーは、ポリシーを使用してこれらのイベントを管理します。修正可能なエラーの場合、スーパーバイザーはジョブ内での再起動をリクエストし、可能であればローカルのチェックポイントから復元します。修正不可能なハードウェア障害の場合、ホットスワップによって障害のあるノードを置き換え、ピアのチェックポイントから復元できます。利用可能な余剰リソースがない場合は、システムをスケールダウンできます。これらのメカニズムにより、よりレジリエンスが高く、リソースの変更に適応しやすいトレーニングを実現できます。代わりのノードが利用可能になると、トレーニングは自動的にスケールアップされ、GPU の使用率が最大化されます。スケールダウンやスケールアップの最中は、ユーザー定義のコールバックを使って学習率やバッチサイズなどのハイパーパラメータを調整できます。修復ポリシーの設定には Python スクリプトを使用できます。

ML グッドプットの最適化に使用できる主な手法を詳しく見ていきましょう。

弾力性のあるトレーニング

弾力性のあるトレーニングは、ワークロードの障害検出と軽減機能を有効にすることで、LLM トレーニングのレジリエンスを強化します。これにより、GPU のリセット、ノードのホットスワップ、ワークロードのデータ並列規模のスケールダウンによる障害のあるノードの使用回避といった修復戦略を自動的に実行できるため、ジョブの中断時間を短縮し、ML グッドプットを向上させることができます。さらに、弾力性のあるトレーニングでは、代わりのノードが利用可能になったときにデータ並列レプリカの自動スケールアップが可能になり、トレーニングのスループットが最大化されます。

弾力性のあるトレーニングの手法の実践例については、こちらの短い動画をご覧ください。

チェックポイント処理の最適化

チェックポイント処理の方法が最適でない場合、トレーニング中に不要なオーバーヘッドが生じたり、中断によって以前のチェックポイントが復元される場合にトレーニングの生産性が大幅に低下したりする可能性があります。非同期チェックポイント処理用のプロセスを定義し、トレーニング状態を GPU の高帯域幅メモリからホストメモリに迅速にオフロードするように最適化することで、これらの問題の影響を大幅に軽減できます。ジョブの中断率や非同期のオーバーヘッドなどの要因に基づいてチェックポイントの頻度を調整することは非常に重要です。最適な間隔は、ワークロードやクラスタのサイズによって数時間から数分までさまざまです。トレーニング オペレーション中のチェックポイントのオーバーヘッドと、予期しない中断時の計算損失の両方を最小限に抑えられる頻度が最適なチェックポイントの頻度です。

頻繁なチェックポイント処理の要件を満たすための堅牢な方法は、ローカルノード ストレージ（ローカル SSD など）、同じクラスタ内のピアノード ストレージ、Google Cloud Storage の 3 つのレベルのストレージを活用することです。この多層チェックポイント処理のアプローチでは、保存および復元のオペレーション中に、ホスト ネットワーク インターフェースまたは NCCL（NVIDIA Collective Communications Library）を介して、これらのストレージ階層間でデータを自動的に複製します。これにより、システムは最速でアクセス可能なストレージ オプションを使用することができます。非同期チェックポイント処理と多層ストレージ戦略を組み合わせることで、高い生産性を維持し、計算の進行状況の損失を最小限に抑えながら、復元時間の短縮とトレーニング ワークフローのレジリエンス強化を実現できます。

最適化されたチェックポイント処理手法の実演は、こちらの短い動画からご覧ください。

これらの ML グッドプット改善手法は、障害シグナリングとジョブ内再起動機能を提供する NVIDIA Resiliency Extension のほか、前述のチェックポイント関連の最適化の一部をサポートする PyTorch の分散チェックポイント処理に最近追加された機能を活用しています。さらに、これらの機能は Google Kubernetes Engine（GKE）と NVIDIA NeMo トレーニング フレームワークに統合され、コンテナ イメージに事前パッケージ化されているのに加え、ML グッドプット最適化レシピで利用できるため、簡単にデプロイできます。

弾力性のあるトレーニングの実例

1,024 個の A3 Mega GPU アクセラレーテッド インスタンス（NVIDIA Hopper 上に構築）を使用した最近の社内ケーススタディでは、これらの手法を組み合わせることで、ワークロード ML グッドプットが 80% 以上から 90% 以上に向上しました。すべてのワークロードが同等の恩恵を受けるわけではありませんが、各手法が特定の指標の改善と ML グッドプットの向上に与える効果がこの表で確認できます。

まとめ

まとめると、弾力性のあるトレーニングとチェックポイント処理の最適化、簡単なデプロイ オプションは、大規模な PyTorch トレーニング ワークロードの ML グッドプットを最大化するために重要な戦略です。上述のケーススタディからわかるように、これらの要素は ML グッドプットの大幅な向上と効率化に大きく貢献します。これらの機能は、Python スクリプトを使用してカスタマイズおよびコンポーズできます。現在 Google Cloud で PyTorch GPU トレーニング ワークロードを実行している場合は、ML グッドプット最適化レシピを試してみることをおすすめします。このレシピは、弾力性のあるトレーニングとチェックポイント処理に推奨される基本構成を提供します。ぜひお試しいただき、フィードバックをお寄せください。

^{Google Cloud 内のさまざまなチームと個人がこの取り組みに貢献しました。このプロジェクトを成功に導いたコンポーネントの開発に多大な貢献をしてくれた Jingxin Ye、Nicolas Grande、Gerson Kroiz、Slava Kovalevskyi に深く感謝します。また、協力パートナーの Jarek Kazmierczak、David Soto、Dmitry Kakurin、Matthew Cary、Nilay Goyal、Parmita Mehta にも感謝します。}

^{1. 20,000 個の GPU を使用するジョブを 8 週間以上にわたって実行した場合の A3 Ultra の想定料金}

-グループ プロダクト マネージャー、Deepak Patil