AWS、Amazon EKSとNVRxによる分散学習の耐障害性ソリューションを公開 I/O待機削減と多層リカバリを実現
要点
- 米Amazon Web Services(AWS)は2026年9月16日、Amazon EKS上でNVIDIA Resiliency Extension(NVRx)を用いたPyTorch FSDP分散学習の耐障害性手法を公開した。
- 状態保存をバックグラウンド処理する非同期チェックポイントにより、クラスタ実行時間の最大40%を占めていたI/Oブロックを解消する。
- 一時的な障害から数秒で復旧する「インプロセス再起動」と、メモリ不足やハング時に同一ジョブ内でワーカーを再生成する「ft_launcher」を導入している。
- NVIDIA H100 GPUを搭載した2ノードから8ノード規模の検証環境におけるベンチマーク結果と、再現可能なコード一式が提供されている。
リード文
米Amazon Web Services(AWS)は2026年9月16日、公式ブログ「AWS ML Blog」において、コンテナオーケストレーションサービス「Amazon Elastic Kubernetes Service(Amazon EKS)」上で、NVIDIAの耐障害性拡張ライブラリ「NVIDIA Resiliency Extension(NVRx)」を統合した分散学習手法を公開した。機械学習フレームワーク「PyTorch」の分散学習機能であるFully Sharded Data Parallel(FSDP、モデルや勾配を複数GPUへ分散配置する技術)と組み合わせることで、大規模学習時の中断リスクや待機時間を大幅に抑制する。NVIDIA H100 GPUを用いたマルチノード環境での検証結果とともに、再現用コードも公開されている。
本文
大規模分散学習における障害連鎖とI/O待機の課題
大規模な分散深層学習ジョブは、数十台規模のノードにまたがり数時間から数日間にわたって連続稼働する。このような大規模かつ長時間の環境では、ネットワークの寸断、メモリエラー、ソフトウェアの例外処理、インフライベントなどによるワーカーの中断が統計的に避けられないという。
従来の環境では、単一のGPU障害が発生すると、GPU間通信ライブラリであるNVIDIA Collective Communication Library(NCCL)のタイムアウトが正常なワーカーへ連鎖的に伝播してしまう。その結果、ポッドが非同期にクラッシュや再起動を繰り返し、学習の進捗が完全に停止した状態で高価なGPU稼働時間だけを消費し続ける事態が発生していた。
さらに、モデルの状態を保存する同期型チェックポイント処理も深刻なアイドル時間を生み出していた。状態の書き込みが完了するまで全GPUの計算プロセス(ランク)がI/O待ちでブロックされ、今回の記事で示されたクラスタ構成では、総実行時間(wall time)の最大40%がチェックポイント保存の待機に費やされていたと説明されている。
既存コードを保持したまま導入可能な「NVRx」の構造
今回示されたソリューションは、アプリケーション層の耐障害性を担うNVRxと、GPUスケジューリングや広帯域ネットワーク、共有ストレージを提供するAmazon EKSインフラを組み合わせる。
NVRxは pip install nvidia-resiliency-ext コマンドで導入可能なPythonレイヤーとして提供されている。PyTorch本体のフォーク(派生版)やカスタムカーネル、再コンパイルを必要とせず、既存のFSDPスクリプトに通常のインポートを追加するだけで動作する。モデル構造や学習ループのコードに手を加える必要がなく、各機能は個別に導入可能となっている。記事では「非同期チェックポイント」「インプロセス再起動」「ft_launcher」の3機能の実装が示された。
耐障害性を実現する3つのコア機能
記事で紹介されている具体的な機能は以下の通りである。
-
非同期チェックポイント(TorchAsyncCheckpoint):
従来のtorch.saveに代わりasync_save()を呼び出すことで、状態辞書(state dict)をバックグラウンドプロセスに渡して即座に学習処理へ復帰させる。次の保存処理の直前にfinalize_async_save()を呼び出して前回の書き込みをコミットする設計だ。FSDPのLOCAL_STATE_DICTと併用することで、各ランクが自身のシャード(分割データ)を直接書き込むため、GPU間での全データ集約(all-gather)やランク0への負荷集中を伴わない。 -
インプロセス再起動(inprocess.Wrapper):
学習関数をラッパーで包むことで、未処理の例外やNCCLのハングといった一時的なソフト障害が発生しても、Pythonプロセス自体を終了させずに回復する。NVRxはアクティブなプロセスグループを破棄し、GPU、GPU間高速接続技術であるNVLink、ネットワークインターフェースカード(NIC)の健全性をランクごとにチェックする。その後、生存ワーカー同士で再度接続(ランデブー)を行い、最新チェックポイントからラップされた関数を再実行する。インタプリタやCUDAメモリアロケータ、外部スコープのオブジェクトが維持されるため、コンテナのライフサイクルに影響を与えず数秒で復旧できる。 -
ft_launcherによるインジョブ再起動:
インプロセスでは対応できないSIGKILLシグナル、メモリ不足による強制終了(OOM Kill)、OSレベルのハングといったハード障害に対処するバイナリである。各ランクでRankMonitorClientを実行し、CLIで指定されたタイムアウトに基づいてハートビート(生存信号)を監視する。停止や失速を検知すると、ランチャーは残存ワーカーを終了してGPUメモリを解放し、同一ジョブ内で新しいワーカーを自動再生成(respawn)して最新チェックポイントから学習を再開する。
障害クラスに応じた多層防御アーキテクチャ
NVRxの各リカバリ層は独立しており、発生する障害の規模や種類に合わせて適用範囲を選択できる。一時的なソフト障害にはインプロセス再起動、プロセスの強制終了を伴うハード障害には ft_launcher、物理ノードの完全な喪失にはAmazon EKSなどのクラスタオーケストレーターがそれぞれ対応する多層防御の構造をとっている。
補足
本稿で検証されたH100 GPU環境(2ノードから8ノード構成)におけるベンチマーク結果および検証用の全コードは、再現可能な形で公開されている。