Amazon EKS NVRx Training、GPU障害からの復旧を数分から数秒に短縮
Amazon EKSはNVIDIA NVRxを再現可能なトレーニングスタックに統合し、注入されたGPU障害からおよそ10〜17秒で復旧した。Amazon EKS NVRx trainingの設計は、16〜64基のH100 GPUを対象とした一部テストで、チェックポイント効率も99%超に維持した。これらの結果は、信頼性の高い復旧にはコンテナの再起動やKubernetesジョブ全体の再構築が必要だという高コストな前提に疑問を投げかける。
このシステムは、PyTorch Fully Sharded Data Parallel（FSDP）と、3つの独立したNVRx機能を組み合わせる。非同期チェックポイントはストレージへの書き込みをトレーニングループから切り離す。プロセス内再起動は、Pythonプロセスを置き換えずに分散状態を再構築する。ft_launcherコンポーネントは、より深刻な障害後に既存ジョブ内で新しいワーカーを起動する。
したがって重要な対比は、AWSと別のクラウドプロバイダーの競争ではない。アプリケーションを理解した復旧と、インフラストラクチャだけに依存する復旧の比較である。Kubernetesは引き続きスケジューリングとノードレベルの障害を担う一方、NVRxはトレーニングプロセスに近い層で障害を処理する。AWSは、この分担によって健全だが高価なGPUが障害を起こしたピアを待つ時間を大幅に削減できるとしている。
Amazon EKS NVRx Training、ジョブ内部へ復旧を移す
中心的な変化は、ワーカー障害がもはやコンテナのライフサイクル全体に及ぶイベントになる必要がないことだ。
AWSとNVIDIAは、Amazon EKS、自己管理型GPUノードグループ、PyTorch FSDPを中心にリファレンス環境を構築した。両社の公開ベンチマークでは、80 GBメモリのNVIDIA H100 GPUを8基搭載するp5.48xlargeインスタンスを使用した。
テストされたクラスターは、2〜8ノード、すなわち16〜64基のGPUに拡張された。各インスタンスは高帯域幅通信向けに32個のElastic Fabric Adapterインターフェースも公開している。Amazon FSx for Lustreは、トレーニングポッド間で共有されるチェックポイントストレージを提供した。
NVIDIA Resiliency Extensionの略称であるNVRxは、PyTorchワークロードに復旧およびチェックポイント機能を追加するPythonパッケージだ。PyTorchのフォーク、カスタムカーネル、再コンパイルは必要ない。チームはトレーニングフレームワーク全体を置き換えることなく、必要な機能を個別に追加できる。
このモジュール型設計は重要である。チェックポイント性能と障害復旧は別の問題だからだ。あるワークロードではプロセス復旧を必要とせず、保存の高速化だけが求められるかもしれない。別のワークロードでは、既存のチェックポイント実装を維持しながらプロセスクラッシュへの保護が必要になる場合がある。
リファレンスアーキテクチャは、各層をその障害スコープに応じて扱う。プロセス内再起動は、Pythonインタープリターが生きている状態で発生する例外や通信ハングを処理する。ft_launcherは、SIGKILL、メモリ不足による終了、一部のOSレベルのハングなどを処理する。
Kubernetesは、ノード全体を失わせる障害に対する外側の層として残る。このアプローチは、入れ子状の復旧ゾーンに似ている。各メカニズムは、障害が一つ下の層の境界を越えた場合にのみ介入する。
トレーニングポッドは、ピア探索のためにヘッドレスKubernetes Servicesを使用する。ワーカーは固定IPアドレスではなくDNSを通じて互いを見つける。この構成により、オペレーターがジョブ設定を書き換えなくても、代替ワーカーが復帰しやすくなる。
AWSは以前、PyTorchツールを使用したEKS上のelastic distributed trainingについて説明している。NVRxの取り組みは、復旧ループをさらに狭めるものだ。個々のrankが障害、停止、消失した際にも、実行中のジョブの生産性を維持することに焦点を当てている。
この違いが、本記事の中心的な緊張関係を生む。Kubernetesはインフラストラクチャを復旧できるが、インフラ復旧にはモデル状態、プロセスグループ、チェックポイントのタイミングに関する詳細な知識がない。NVRxは、こうした判断をトレーニングアプリケーション内に持ち込む。
ブロッキング型チェックポイントがウォールタイムの約40%を消費していた
最初の性能問題はGPU計算ではなかった。チェックポイントデータがストレージに到達するまで、すべてのrankが待機する時間だった。
同期チェックポイントでは、必要なモデルおよびオプティマイザー状態が書き込まれるまでトレーニングが停止する。分散FSDPジョブでは、この停止は参加するすべてのrankに影響する。健全なGPUは割り当てられたままだが、書き込み中は順伝播も逆伝播も実行しない。
AWSによると、同期チェックポイントではスケーリングテストにおけるトレーニング効率は57%から61%にとどまった。ストレージ書き込みには約275秒かかり、その時間は16〜64基のGPU間でおおむね一定だった。したがって、計算リソースを追加してもストレージ起因の停止は解消されなかった。
NVRxの非同期チェックポイントは書き込み経路を変える。トレーニングプロセスは状態をCPU上にステージングし、その処理を常駐バックグラウンドプロセスに渡す。その後、ストレージI/Oが継続する間に、メインプロセスは次のトレーニングステップへ戻る。
実装ではTorchAsyncCheckpointとそのasync_save()メソッドを使用する。別の保存を始める前、または終了する前に、アプリケーションは未完了の処理を完了させる。この調整により、未完了のチェックポイントが次の処理と気付かれないまま衝突することを防ぐ。
FSDPのローカル状態辞書はこの設計を強化する。各rankが自身のシャードを書き込むため、all-gather操作と、単一のrank-zeroによる書き込みボトルネックを回避できる。PyTorchのFSDPドキュメントでは、参加ワーカー間にパラメーターを分散する、より広範なシャーディングモデルを説明している。
AWSによれば、チェックポイント間隔を1,000ステップとした場合、NVRxの非同期チェックポイントは2ノードで99.2%のトレーニング効率に達した。8ノードでは効率が99.8%に達した。同期方式との比較では、8ノード時に60.3%だった。
これらの数値はベンダー報告のベンチマーク結果であり、すべてのモデルやストレージ構成での保証ではない。それでも、その背景にある仕組みは単純だ。計算時間がストレージ書き込み時間より長ければ、バックグラウンド処理は有用なトレーニング作業の背後にほぼ完全に隠せる。
AWSがチェックポイント頻度を上げたとき、その限界は明確になった。8ノードで100ステップごとに1回チェックポイントを取得した場合、同期方式の効率は14.7%まで低下した。非同期方式も低下したが、29.6%とより高い水準を維持した。
理由はタイミングにある。100回のトレーニングステップには約280秒かかる一方、チェックポイント書き込みには約275秒かかった。次のストレージ処理を隠すための余剰計算時間は、ほとんどなかった。
これがNVRx非同期チェックポイントの真の限界である。非同期I/Oは計算の背後に書き込みを隠せるが、ストレージを無限に高速化することはできない。保存要求がファイルシステムの完了速度と同じ頻度で到着すると、最終的にキューが負荷を生む。
この制約があっても、この機能はチームによるチェックポイント間隔の選び方を変える。同期システムでは、保存のたびに目に見えるアイドル時間が発生するため、チェックポイント回数を減らしがちだ。保存頻度が低ければ、障害後に失われるトレーニング量は増える。
非同期チェックポイントはこのトレードオフを緩和する。書き込みの間に十分な計算がある場合、チームはより頻繁に保存できる。より短い間隔はロールバック距離を減らし、I/OのオーバーラップはGPU投資をより多く維持する。
結果は単に高速なチェックポイントAPIではない。定常状態の効率と、復旧可能な進捗の間にある異なるバランスだ。このバランスは、ジョブが長期化し、障害を起こしやすいコンポーネントが増えるほど価値を増す。
プロセス内再起動、Kubernetesのみの復旧モデルに異議を唱える
最速の復旧経路ではPythonプロセスを維持し、障害で損なわれた分散リソースだけを再構築する。
リファレンス実装では、NVRxがメインのトレーニング関数をプロセス内再起動コントローラーでラップする。サポート対象の例外が発生すると、ラッパーは実行中の試行を中断し、次の呼び出しを準備する。この一連の処理を通じて、外側のPythonプロセスは生存し続ける。
NVRxはまず、損傷したPyTorch分散プロセスグループを中止する。フライトレコーダーのトレースを収集し、NCCLバックエンドを停止して、無効なグループを破棄できる。NCCLは、GPU間の集団操作に用いるNVIDIAの通信ライブラリである。
次に、ヘルスチェックが各rankに関連するリソースを検査する。これらのチェックは、GPU、NVLink接続、ネットワークインターフェース、rank障害の繰り返しを対象にできる。再試行コントローラーは再起動試行回数を制限し、何個のアクティブなrankが生き残る必要があるかを決定する。
システムは生き残ったrankを連続したグループに再割り当てし、新しいランデブーを開始する。ラップされたトレーニング関数はFSDPモデルを再作成し、最新のチェックポイントを読み込み、処理を再開する。Pythonインタープリターと、ラップされた関数の外側にあるオブジェクトは引き続き利用できる。
この手法はソフト障害を対象とする。例として、処理されないアプリケーション例外や、ウォッチドッグが検出できるNCCLハングが挙げられる。ブロックされたネイティブコールのすべてから、Python例外が確実に発生するとは想定していない。
その代わり、進捗ウォッチドッグがPythonバイトコード操作間の活動を記録する。別の監視スレッドが共有状態を確認し、あるrankの進捗が停止すると再起動を要求できる。その後、システムは参加ワーカー間で中断を調整する。
AWSはこのアプローチをft_launcherおよびベースラインのKubernetes復旧と比較した。この実験では、p5.48xlargeノード2台、H100 GPU 16基、FSDP下のLlama 3.1 8Bを使用した。ジョブは2,000ステップ実行され、500ステップごとに保存された。
研究者らは同じスケジュールを用い、各実行に5件の決定論的な障害を注入した。NVRxのプロセス内再起動は、コンテナを再起動せずに障害ごと約10秒で復旧した。AWSは、トレーニングgoodputを31%、インフラストラクチャgoodputを87%と測定した。
トレーニングgoodputは、有効なトレーニング進捗を生む時間を測定する。インフラストラクチャgoodputは、割り当てられたインフラストラクチャが稼働し利用可能な時間を測定する。両者の差には、ロールバックやチェックポイントの読み込みなど、モデルを進めない作業が含まれる。
ベースラインのKubernetes復旧では、注入された障害1件あたり約270秒を要した。報告された実験では、トレーニングgoodputは11.5%、インフラストラクチャgoodputは35.8%だった。これは比較対象がコンテナ起動の最適化をはるかに超えることを意味する。
AWSによると、1つのrankの障害が生き残ったrankで通信タイムアウトを引き起こした。その後、ポッドが同期せずに再起動し、タイムアウトの反復サイクルとCrashLoopBackOffの挙動につながった。オーケストレーターは、分散トレーニンググループがどのように一体で復旧すべきかを理解しないままコンテナを復旧した。
アプリケーションを理解した復旧には、その欠けているコンテキストへアクセスできる。進捗がいつ停止したか、どのプロセスグループが無効になったか、どのチェックポイントからトレーニングを再開できるかを把握している。Kubernetesが認識するのはポッド状態とノードの健全性であり、FSDPトレーニングステップの完全なセマンティクスではない。
これはKubernetes復旧を不要にするものではない。停止したノードでは、インタープリター、CUDA状態、ローカルプロセスを維持できない。ノードの置き換えは引き続きクラスター層の役割であり、復旧したワーカーにも永続的なチェックポイントデータが必要になる。
負担を受けるのは、ポッド再起動だけを障害耐性ポリシーとしているチームだ。このアプローチは依然として単純だが、復旧時間帯に大量のアクセラレーター時間を浪費する可能性がある。クラスターが大きいほど、1件の障害で本来健全な多数のワーカーがアイドル状態になるため、コストは増幅される。
NVRx Fault Tolerance、ソフト障害とハード障害を分ける
単一の再起動メカニズムですべての障害をカバーすることはできないため、NVRx fault toleranceはプロセスを維持する復旧とワーカー置換を分離する。
ft_launcher コンポーネントは、プロセス内再起動では耐えられない障害を処理します。これには SIGKILL、メモリ不足による強制終了、使用可能な Python インタープリタが残らない障害が含まれます。使い慣れたランデブーの概念を維持しつつ、torchrun を置き換えます。
各トレーニング rank は、分散初期化後に RankMonitorClient を作成します。クライアントはトレーニング中にハートビートを送信します。rank ごとのモニターサーバーは、通常時および初期起動時に設定されたタイムアウトとこれらのシグナルを照合します。
AWS の構成では、rank のハートビートタイムアウトに 900 秒を使用していました。初回のモデル読み込みにより長い時間を許容するため、初期ハートビートのタイムアウトは 1,200 秒でした。5 秒のモニター間隔によって、ランチャーがワーカーの状態を確認する頻度を制御していました。
これらの値は汎用的な推奨値ではなく、構成例です。ハートビートのタイムアウトは、シグナル間で発生し得る正当な最長遅延を上回る必要があります。短すぎる場合、低速なチェックポイント処理やモデル初期化がワーカー障害と見なされる可能性があります。
ワーカーが停止または応答しなくなると、ft_launcher は残りのワーカーを終了します。GPU メモリを回収し、再度ランデブーを実行して、同じジョブ内で新しいプロセスを開始します。新しいワーカーは、直近のチェックポイントから状態を復元します。
launcher guide では、NVIDIA の NeMo RL スタックにおける同様の一般的なパターンが示されています。このより広範な統合は、NVRx が EKS 専用のユーティリティではなく、再利用可能なレジリエンス層として意図されていることを示唆しています。
AWS は、ft_launcher による障害注入 1 件あたりの復旧時間を約 17 秒と測定しました。報告された実行では、トレーニング goodput は 25.5%、インフラストラクチャ goodput は 85.9% に達しました。これはプロセス内復旧より低速でしたが、270 秒の Kubernetes ベースラインよりははるかに高速でした。
この差は、各方式がどれだけの状態を保持できるかを反映しています。プロセス内再起動ではインタープリタと外側のプロセスを維持します。ft_launcher は新しいワーカーを作成し、分散状態を初期化し、モデルを再構築して、チェックポイントを再読み込みする必要があります。
より大規模な環境では、チェックポイントの読み込みが復旧時間の大部分を占める場合があります。プロセス作成を高速化しても、モデルおよびオプティマイザのシャードを読み込む必要はなくなりません。そのため、共有ファイルシステムのスループットは依然としてフォールトトレランス設計の一部です。
Amazon EKS の NVRx トレーニングアーキテクチャでは、GPU ノードと同じ Availability Zone にある FSx for Lustre SCRATCH_2 ファイルシステムを使用していました。この配置は、チェックポイント読み込み時のレイテンシ低減を目的としています。復旧後、すべてのワーカーが同じ永続化済み状態に到達できます。
NVRx の非同期チェックポイントは、両方の再起動経路と直交する機能です。書き込み関連のアイドル時間を削減し、失われる可能性のある進捗量を制御します。再起動層は、障害後にワーカーがどれだけ早く復帰するかを決定します。
この分離により運用者の選択肢は増えますが、同時にポリシー設計も必要になります。どの例外でプロセス内復旧を発動できるか、安全なリトライ回数はいくつか、どのヘルスチェックで rank を除外すべきかを決めなければなりません。ハートビートとランデブーのタイムアウトも設定する必要があります。
NVIDIA は NVRx project を実験的で活発に開発中のプロジェクトと位置付けています。ドキュメントでは、機能とインターフェースが変更される可能性について警告しています。本番チームは、バージョン選定とアップグレードテストをレジリエンス計画の一部として扱うべきです。
AWS はベンチマークの再現に NVRx 0.4.1 を使用しました。投稿では、現在のデプロイメント向けに、更新されたランチャー構成とともにバージョン 0.6.0 を推奨しています。このバージョン差は重要です。フォールトトレランスシステムは、重要な起動および復旧経路に直接関与するためです。
復旧不能なジョブを繰り返し再起動する場合、失敗した復旧メカニズムは自動化がない場合より悪い結果になり得ます。リトライ上限、最小 world size、障害カウンターは無制限のループを防ぎます。それでも運用者には、復旧成功と繰り返される障害を区別するアラートが必要です。
99% という結果には重要な条件がある
このベンチマークは強力なメカニズムを裏付けていますが、すべての分散トレーニングワークロードで 99% の効率を確立するものではありません。
AWS は、H100 ベースの p5 インスタンス上で、PyTorch FSDP を用いる Llama 3.1 8B という主要なモデル構成を 1 件テストしました。EFA ネットワークと FSx for Lustre ストレージを使用しています。モデルサイズ、ストレージパス、チェックポイント形式、ステップ時間が異なれば、オーバーラップ可能な時間枠も変化します。
99% という数値は、選定された間隔での非同期チェックポイント効率に適用されます。繰り返し障害が発生する環境でのエンドツーエンドの goodput を示すものではありません。障害注入テストでは、トレーニング goodput はプロセス内復旧で 31%、ft_launcher で 25.5% にとどまりました。
この低い結果は、チェックポイント測定と矛盾するものではありません。両者は別の問いに答えています。非同期効率は通常のトレーニング中のチェックポイントオーバーヘッドを測定する一方、goodput には障害注入、ロールバック、読み込み、その他の復旧作業が含まれます。
チェックポイント頻度にも避けられない限界があります。100 ステップごとに 1 回のチェックポイントでは、非同期トレーニングの効率は 99% ではなく 29.6% に達しました。計算時間と書き込み時間が近かったため、ストレージシステムはほぼ継続的に使用されていました。
メモリ負荷にも注意が必要です。非同期チェックポイントは、即時の GPU 操作の外部にデータをステージングし、バックグラウンドプロセスに書き込みの責任を持たせます。チームは実際の state dictionary を用いて、CPU メモリ、キュー深度、ストレージのバックログを測定すべきです。
復旧対象範囲も別の制約です。オペレーティングシステムがワーカーを強制終了した場合やノードが消失した場合、プロセス内再起動は役に立ちません。ft_launcher は停止したプロセスを置き換えられますが、それでもジョブ、クラスターネットワーク、ランデブーサービス、チェックポイントストアに依存します。
ノード喪失は依然として Kubernetes の課題です。リージョン規模のストレージ中断や破損したチェックポイントは、すべての復旧層を同時に無力化する可能性があります。このアーキテクチャは複数の一般的な障害コストを削減しますが、共有依存関係を取り除くものではありません。
障害検知が誤検知を生むこともあります。長いコンパイルフェーズ、データ読み込みの停止、ファイルシステムの停止が、厳しすぎるハートビートタイムアウトを超えるかもしれません。その場合、ランチャーは正常なワーカーを再起動し、有効な進捗を破棄します。
チームは自動復旧を有効化する前に、ワークロード固有のタイムアウトデータを必要とします。最長のモデル初期化、チェックポイント、検証、データ入力の間隔を取得すべきです。テストには通常のトレーニングステップ内の障害だけでなく、それらのフェーズ中の障害も含める必要があります。
AWS のベンチマークでは、決定論的に障害を注入しました。この手法は再現可能な比較を支援しますが、本番障害はそれほど整然としていません。実際のクラスターでは、ネットワーク劣化、ストレージの低速化、熱問題、プロセスクラッシュが同時に発生する場合があります。
reference implementation は、チームが構成を再現するための有用な出発点を提供します。異なるインスタンスファミリーとモデル規模で再現することで、報告された利点がどの程度移植可能かを判断できます。
運用上の複雑さが最後のトレードオフです。このスタックには Kubernetes Jobs、DNS ベースのピア検出、EFA リソース、共有ストレージ、NVRx ラッパー、モニタークライアント、複数のタイムアウト層が含まれます。各コンポーネントが新たな構成領域を生み出します。
大規模な GPU フリートが 1 つの rank の障害後に数分間アイドル状態になる場合、この複雑さはなお正当化できるかもしれません。一方、より小規模なジョブでは、より単純な Pod 再起動戦略を受け入れられる可能性があります。重要なのはベンチマークの名声ではなく、復旧コストに障害頻度を掛け合わせた計算です。
設計を評価するチームは、トレーニングとインフラストラクチャの両方の goodput を追跡すべきです。モデルが古いチェックポイントを繰り返し再読み込みしていても、GPU 使用率だけは健全に見える場合があります。有用なダッシュボードには、完了ステップ、ロールバック距離、チェックポイントの鮮度、再起動原因を表示する必要があります。
エンジニアリング組織には、これらの実験に関する永続的な記録も必要です。検索可能な engineering knowledge base は、タイムアウトの変更、障害トレース、ベンチマーク結果を結び付けられます。この文脈は、失敗した復旧構成をチームが繰り返すことを防ぐ助けになります。
Amazon EKS NVRx ベンチマーク後に注目すべき点
次に必要な証拠は、この設計がより大規模なモデル、実際の障害、変化する NVRx リリースにわたって優位性を維持できるかを示すものです。
最初の指標は、64 基の H100 GPU を超える独立した再現です。AWS は 2 ノードから 8 ノードまでのスケーリングをテストしましたが、クラスタ規模が大きくなると障害頻度と協調コストも増加します。数百基のアクセラレータにわたる結果は、ランデブーとチェックポイント読み込みの限界をより明確に示すでしょう。
より大規模なテストでは、平均復旧時間以上を報告すべきです。まれに発生する 5 分間の復旧が長時間実行の経済性を支配し得るため、分布も重要です。レポートにはテールレイテンシ、再起動失敗回数、インシデントごとの失われた進捗を含める必要があります。
2 つ目の指標は、主要なトレーニングフレームワーク内での採用です。NVRx はすでに PyTorch ベースのワークロードと接続され、NVIDIA のより広範なソフトウェアスタックにも登場しています。より多くのネイティブ統合が実現すれば、チームが保守するカスタムラッパーとランチャーコードを減らせます。
フレームワークによる採用は、NVRx のフォールトトレランスが標準的なアプリケーション層になり得るという根拠を強めます。統合が断片化すれば、特に各フレームワークで異なるタイムアウト、チェックポイント、ランデブーロジックが必要な場合、その根拠は弱まります。
3 つ目の指標は、制御されていない本番障害からの証拠です。比較には決定論的な障害注入が必要ですが、現場のインシデントでは実験室で再現されることが少ない組み合わせが試されます。運用者は、GPU エラー、NCCL の停止、OOM 強制終了、ノード喪失について、それぞれ復旧率を公開すべきです。
復旧の成功は、単にプロセスを再起動する以上の意味を持つべきです。ジョブは有効な状態を復元し、正しい更新を継続して生成し、サイレントなチェックポイント破損を回避しなければなりません。繰り返し復旧後のモデル収束は、復旧速度と同じ程度に精査されるべきです。
現在 Amazon EKS NVRx トレーニングを検討するチームにとって、実践的な最初のステップは、制御されたシャドーベンチマークです。実際のモデル、チェックポイントサイズ、ファイルシステム、ステップ時間を使用してください。1 つの障害スケジュールの下で、同期保存、非同期保存、ランチャー復旧、Kubernetes のみの復旧を比較します。
次に、観測された計算時間とストレージ時間を基にチェックポイント間隔を調整します。チェックポイントにかかる時間が保存間隔とほぼ同じなら、非同期オーバーラップは不完全なままです。計算によってより広い時間枠が得られる場合、報告された 99% 効率の実現可能性は高まります。
復旧ポリシーは、保守的なリトライ上限から始めるべきです。すべての再起動理由を記録し、診断トレースを保持してください。同じ rank またはチェックポイントで繰り返し失敗するジョブには、無限の復旧ループではなくエスカレーションが必要です。
AWS と NVIDIA の取り組みは、見過ごしがたい結論を示しています。分散トレーニングの信頼性を、インフラストラクチャだけの課題にとどめることはできません。アプリケーションは、進捗、チェックポイントの有効性、プロセスグループの状態を、オーケストレーターにはできない形で理解しています。
Amazon EKS は引き続き、不可欠なスケジューリングおよびノード置換の基盤を提供します。NVRx はその境界内で、より高速な応答を追加します。両者を組み合わせることで、複数の重要な障害クラスにおいて、4 分間の復旧サイクルから秒単位の再起動へ進む、信頼できる道筋が得られます。
現在の未解決の問いは、概念的なものではなく運用上のものです。チームは、自身のモデル、ストレージシステム、実際の障害パターンで、これらの利点を再現できるでしょうか。次の Amazon EKS NVRx トレーニングデプロイメントは、そのテストに導かれるべきです。
