多層チェックポイント処理による早期段階からの頻繁な保存で大規模な AI トレーニング ジョブを最適化

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

基盤モデルのトレーニング インフラストラクチャが数万台に及ぶアクセラレータにスケールされると、これらの価値の高いリソースを効率的に活用することが非常に重要になります。特に、クラスタが大きくなるほどハードウェア障害の頻度が高くなり（数時間ごと）、以前に保存したチェックポイントからの復元に時間がかかります（最大 30 分）。その結果、トレーニングの進行が大幅に遅くなります。チェックポイントは、モデルのトレーニングの進行を特定の時点で保存した状態を表し、中間モデルの重み付けとその他のパラメータのセットで構成されます。

Google は最近、AI Hypercomputer に多層チェックポイント処理を導入しました。AI Hypercomputer は、Google が 10 年以上にわたって培ってきた AI に関する専門知識を活用した統合型スーパーコンピューティング システムです。このソリューションは、クラスタ内メモリ（RAM）とレプリケーションを含む複数のストレージ階層、そして Google Cloud Storage を利用することで大規模なトレーニング ジョブの ML グッドプットを向上させます（たとえば、TPU v5p 上の 35,000 チップのワークロードで 6.59% 向上）。これにより、トレーニング ジョブの進行状況の喪失を最小限に抑え、平均復旧時間（MTTR）を改善できます。このソリューションは、MaxText をリファレンス アーキテクチャとして使用する JAX と、PyTorch / GPU を利用する NeMo と互換性があります。

このアーキテクチャでは、最大規模のモデルと非常に大規模なノードクラスタ全体で最適な頻度でチェックポイントを取得（チェックポイントの保存は 5 分未満まで劣線形でスケール）し、数千のノードを含むクラスタに 1 分未満で復元できます。

グッドプットが向上すると、インフラストラクチャ費用の削減に直接的な効果があります。たとえば、アクセラレータ チップを使用したモデルのトレーニングが完了するまでに 1 か月かかるとします。比較的小規模なトレーニング ワークロードでも、最適なチェックポイント処理による費用削減効果は大きなものです。1 時間あたり 88 ドルの費用がかかる VM（a3-highgpu-8g）1,000 台で 1 週間のトレーニング ジョブを行う場合、このトレーニング タスクのグッドプットが 6.5% 向上すると、インフラストラクチャ費用を約 100 万ドル節約できる可能性があります。

障害が多いほど多くのチェックポイント処理が必要

確率論的には、クラスタのサイズが大きくなるにつれて、トレーニング ジョブの平均故障間隔（MTBF）が短くなり、障害が頻繁に発生するようになります。したがって、基盤モデルを開発する場合、トレーニング ジョブの進行の喪失をできるだけ回避できるよう、より頻繁にチェックポイントを取得することが重要です。以前は、Google Kubernetes Engine（GKE）のお客様は 30 分に 1 回しかチェックポイントを書き込むことができず（保存先は Cloud Storage）、最後に保存されたチェックポイントを読み取ってクラスタ内のすべてのノードに配布するまでに最長で 30 分待たなければなりませんでした。

多層チェックポイント処理ではデータがノードのメモリ（RAM）に非同期で書き込まれ、クラスタ内で定期的に複製されて Cloud Storage にバックアップされるため、チェックポイントの書き込みを大幅に高速化し、保存の頻度を高めることができます。障害が発生した場合、近くのノードのインメモリ チェックポイントからデータを取得することで、ジョブの進行状況を迅速に復元できます。近くのノードの RAM にチェックポイント データが存在しない場合、チェックポイントは Cloud Storage バケットのバックアップからダウンロードされます。このソリューションを使用すると、チェックポイントの書き込みレイテンシはクラスタ内のノード数に応じて増加せず、一定のままになります。読み取りも一定で、独立してスケールされるため、チェックポイントの読み込みが高速になり MTTR が短くなります。

アーキテクチャの詳細

この多層チェックポイント処理ソリューションのコンセプトは、ML トレーニング ジョブがチェックポイントの保存とチェックポイントからの復元に使用する「魔法」のような単一のローカル ファイル システム ボリュームを提供することです。RAM ディスクレベルの読み取り / 書き込み速度と、Cloud Storage に付随するデータの耐久性を兼ね備えているため、「魔法」と言えます。

このソリューションを有効にすると、ML ジョブが認識できる唯一のストレージ階層がローカル ボリューム（ノード ストレージ）になります。そこに書き込まれたチェックポイントは、クラスタ内で 1 つまたは 2 つ以上のピアノードに自動的にレプリケートされ、Cloud Storage に定期的にバックアップされます。

ジョブが再開されると、ノードで実行されているトレーニング ジョブの新しい部分に固有のチェックポイント データが ML ジョブで使用するローカル ボリュームに自動的に表示されます。バックグラウンドでは、必要なデータがクラスタ内の別のノードまたは Cloud Storage から取得されます。完全に書き込まれた最新のチェックポイントを（どこにあるかにかかわらず）見つける処理も、ML ジョブに対して透過的に行われます。

階層間でデータを移動するコンポーネントは「Replicator」と呼ばれ、ローカル ボリューム マウントを提供する CSI ドライバの一部としてすべてのノードで実行されます。

詳しく説明すると、Replicator は次の重要な役割を果たします。

• インテリジェンスの一元化: Cloud Storage のバックアップとクラスタ内の集合データを分析して、再起動時にジョブを復元するための最新のチェックポイント全体を特定します。さらに、すべてのノードでチェックポイントの保存が成功したことを検出し、古いデータを安全にガベージ コレクションの対象にできるタイミングを通知します。また、Cloud Storage にバックアップするチェックポイントを戦略的に判断します。

• 状況に応じたピアの選択: TPU と GPU の両方で使用される基盤となるネットワーク トポロジを認識しているため、状況に応じて各ノードのレプリケーション ピアを選択します。このとき、高帯域幅で低レイテンシの「近い」ピアが優先されますが、この「近い」ピアでは、相関する障害（同じ TPU スライス内または GPU スーパーブロック内での障害など）が発生するリスクが高い場合があります。そのため、ネットワーク オーバーヘッドがわずかに増加するものの、独立した障害に対する復元力が高い「遠い」ピアも選択します（別の GPU スーパーブロックに存在するピアなど）。データ並列処理のシナリオでは、同一のデータを持つピアが優先されます。

• データの自動重複排除: データ並列処理を行う場合、複数のノードで同一のトレーニング パイプラインが実行されるため、同一のチェックポイントが複数保存されます。Replicator がピアを選択することで、これらのノードが確実にペアリングされるため、実際のデータ レプリケーションは不要になります。代わりに、各ノードがピアのデータの完全性を検証します。そのために追加の帯域幅が消費されることはなく、レプリケーションは瞬時に行われ、ローカル ストレージの使用量が大幅に削減されます。ピアが正しく構成されていない場合、引き続き標準のチェックポイントのコピーが行われます。

• データ並列処理を前提とした巨大モデルモード: このモードは、最適化を行うだけでなく、ノード自体のチェックポイントとピアのデータの両方を格納するにはローカル ノード ストレージが十分ではない最大規模のモデルに対応します。このような場合、ML ジョブはデータ並列処理を想定するように Replicator を構成するため、必要なローカル ストレージが大幅に減ります。この仕組みを応用して、最新のチェックポイントを保存するノードではなく、専用のノードが Cloud Storage のバックアップを処理するシナリオにも対応できます。

• Cloud Storage 利用の最適化: データ重複排除を利用することで、一意のデータすべてが Cloud Storage に 1 回だけ保存されるようになるため、保存容量、帯域幅の使用量、関連費用を最適化できます。

• 自動ガベージ コレクション: Replicator は、すべてのノードにおけるチェックポイントの保存を継続的にモニタリングします。最新のチェックポイントがすべての場所で正常に保存されたことが確認されると、古いチェックポイントの削除が自動的に開始されます。このプロセスが完了するまで、Cloud Storage にバックアップされているチェックポイントは保持されます。

幅広いチェックポイント処理ソリューション

Google Cloud では、多様な AI トレーニングのニーズに対応するため、包括的なチェックポイント処理ソリューション ポートフォリオを提供しています。Cloud Storage への直接アクセスや Cloud Storage FUSE などのオプションは、よりシンプルなアプローチであり、小規模から中規模のワークロードに非常に効果的です。Lustre などの並列ファイル システムは、大規模なクラスタに高スループットを提供します。その一方、多層チェックポイント処理は、頻繁な保存と迅速な復元を必要とする、最も要求が厳しい最大規模（1,000 ノード超）のトレーニング ジョブ専用に構築されています。

多層チェックポイント処理は現在プレビュー版で、Cloud TPU 用の JAX と GPU 上の PyTorch に重点が置かれています。ユーザーガイドをご覧のうえ、ぜひ導入をご検討ください。ご質問やフィードバックがございましたら、担当アカウント チームまでお気軽にお問い合わせください。

ー シニア ソフトウェア エンジニアリング マネージャー、Jason Messer

ー スタッフ ソフトウェア エンジニア、Dmitry Kakurin