NVRxを活用したAmazon EKS上でのフォールトトレラントな分散学習：大規模AIモデルの信頼性向上
大規模なAIモデルの訓練は、多くの場合、数十のノードにまたがり、数時間から数日間にわたる分散学習ジョブとして実行されます。このようなスケールと期間では、ネットワークの問題、メモリのエラー、ソフトウェアの例外、インフラストラクチャのイベントなど、何らかの障害が発生することは避けられません。単一のGPU障害がNVIDIA Collective Communication Library (NCCL) のタイムアウトを引き起こし、健全なワーカーへの伝播、ポッドのクラッシュと非同期での再起動など、連鎖的な問題を引き起こす可能性があります。これにより、高価なGPU時間を消費しながらも、訓練の進捗がゼロになるという事態が発生します。同期チェックポインティングは、すべてのランクがI/Oによってブロックされるため、アイドル時間の原因となり、本記事で示されるクラスタサイズでは、総実行時間の最大40%を占める可能性がありました。本レポートでは、NVIDIA Resiliency Extension (NVRx) をPyTorch Fully Sharded Data Parallel (FSDP) トレーニングに統合し、Amazon Elastic Kubernetes Service (Amazon EKS) 上でこれらの課題を解決する手法について解説します。
大規模分散学習における耐障害性の課題とNVRxのソリューション
大規模分散学習では、長期にわたるGPUリソースの占有と多数のノード利用により、計算時間の損失が大きな課題となります。特に、単一のGPU障害がNVIDIA Collective Communication Library (NCCL) のタイムアウトを引き起こし、クラスター全体のトレーニング進行を停止させることが頻繁に発生します。また、従来の同期チェックポインティングは、I/O処理中に全てのランクの処理を停止させるため、大規模なGPUクラスターではウォールタイムの最大40%を消費する可能性があり、訓練効率を著しく低下させます。
NVIDIA Resiliency Extension (NVRx) は、pip install nvidia-resiliency-ext でインストール可能なPythonレイヤーとして提供され、既存のPyTorchスクリプトにフォールトトレランス機能を追加します。NVRxは、カスタムカーネルやPyTorchのフォーク、再コンパイルを必要とせず、FSDPスクリプトに通常のインポートとして組み込むことができます。NVRxは主に以下の3つの機能でこれらの課題に対処します:
-
非同期チェックポインティング (Async Checkpointing): I/Oとトレーニング処理をオーバーラップさせることで、チェックポイント保存によるアイドル時間を大幅に削減します。これにより、より頻繁なチェックポイント保存が可能となり、障害発生時の作業損失を最小限に抑えます.
-
インプロセス再開 (In-Process Restart): 一時的なソフトフォルト（未処理の例外やNCCLハングなど）が発生した場合に、Pythonプロセスを終了させることなく、数秒でトレーニングを回復させます。NVRxはアクティブなプロセスグループを中止し、各ランクのヘルスチェック（GPU、NVLink、NIC）を実行し、生き残ったランクを再ランデブーさせ、最新のチェックポイントからラップされた関数を再実行します。インタプリタ、CUDAアロケータ、および外側スコープのオブジェクトは維持されます。
-
インジョブ再開 (
ft_launcherによるIn-Job Restart): SIGKILL、OOM (Out-Of-Memory) キル、OSレベルのハングなど、インプロセス再開では捕捉できないハードフォルトに対して機能します。ft_launcherバイナリは、これらのケースでワーカーがクラッシュした場合に自動的に再起動します。これにより、ジョブ全体の再起動を必要とせず、トレーニングを継続できます。
これらの機能はそれぞれ独立して採用可能であり、既存のFSDPスクリプトに対してモデルやトレーニングコードに大きな変更を加えることなく導入できます.
Amazon EKSとNVRxによる堅牢な分散学習環境の構築
NVRxのフォールトトレランスプリミティブは、高性能なマルチノードGPUワークロード向けに設計されたAmazon EKSベースのトレーニング環境と組み合わせて使用されます. Amazon EKSは、GPUスケジューリング、高帯域幅ネットワーキング、およびチェックポイント永続化のための共有ストレージといったインフラストラクチャ基盤を提供します.
アーキテクチャ概要では、2ノードから8ノード構成のp5インスタンスを使用するEKSクラスターと、EFA (Elastic Fabric Adapter) によるインターコネクト、FSx for Lustre、そしてトレーニングポッド内のNVRxコンポーネントが示されています.
- Amazon EKS: Kubernetesのコントロールプレーン、ポッドスケジューリング、ジョブライフサイクル管理を担当します. EKSは、コンテナ化されたアプリケーションをオーケストレーションするための堅牢なプラットフォームを提供し、MLOpsにおいて重要な役割を果たします.
- Amazon EC2 P5 インスタンス: 最新のGPUインスタンスであるH100 GPUを搭載しており、大規模なディープラーニングモデルの訓練に求められる高い計算能力を提供します.
- Elastic Fabric Adapter (EFA): GPUノード間の低遅延かつ高スループットなネットワーク通信を可能にし、分散学習における通信ボトルネックを解消します.
- FSx for Lustre: チェックポイントの永続化と高速な読み書きのために利用される高性能パラレルファイルシステムです. 特に、GPUノードと同じアベイラビリティゾーンにFSx for Lustreを配置することで、リカバリ時の読み取りレイテンシを最小化し、チェックポイントのロード時間を最適化することが重要です.これにより、リカバリ時間は主にチェックポイントのロードによって支配されるため、この最適化は大規模なリカバリにおいて極めて重要となります.
この組み合わせにより、NVRxがアプリケーションレベルのレジリエンス（非同期チェックポインティング、インプロセス再開、インジョブ再開）を処理し、EKSが安定したインフラストラクチャを提供することで、費用対効果が高く、信頼性の高い分散学習環境が実現されます.ベンチマーク結果では、H100 GPUを2ノードから8ノードスケールで利用し、これらの機能が有効な場合に効率が向上することが示されています.
開発者・エンジニア視点での考察
-
選択的導入による段階的耐障害性強化の機会: NVRxの非同期チェックポインティング、インプロセス再開、ft_launcherといった機能は、既存のPyTorch FSDPスクリプトに個別に導入できる柔軟性を提供します. これは、開発者が既存のモデルやトレーニングロジックに大きな変更を加えることなく、最もクリティカルな障害シナリオから順に耐障害性を強化できることを意味します。例えば、I/Oボトルネックが顕著な場合は非同期チェックポインティングから、NCCLハングなどのソフトフォルトが多い場合はインプロセス再開から導入するなど、段階的なアプローチでリスクを管理し、開発コストを最適化できます。
-
リカバリ時間最適化におけるストレージ戦略の重要性: 分散学習の耐障害性設計において、チェックポイントの「保存頻度」だけでなく、「チェックポイントのロード時間」がリカバリ時間を支配するという洞察は重要です. 特に大規模モデルでは、チェックポイントサイズが大きくなるため、Amazon FSx for Lustreのような高性能な共有ファイルシステムをGPUノードと同じアベイラビリティゾーンに配置し、低遅延でのデータアクセスを確保する戦略は、実質的なダウンタイムを最小化する上で不可欠です. 開発者は、単にチェックポイントを永続化するだけでなく、その「ロードパス」にも着目し、全体のリカバリ時間を評価するべきです。
-
Kubernetesを通じたMLOps基盤の標準化と自動化: NVRxとAmazon EKSの組み合わせは、分散学習のインフラストラクチャをKubernetes上に標準化し、MLOpsの運用効率を大幅に向上させる可能性を秘めています. GPUスケジューリング、高帯域幅ネットワーキング (EFA)、そして永続ストレージ (FSx for Lustre) の管理がEKSによって統合されることで、開発者はインフラの詳細から解放され、モデル開発とトレーニングロジックの改善に集中できます. さらに、NVRxによる自動的な障害検出と回復は、手動での介入を減らし、大規模なAIワークロードの継続的な統合・継続的なデリバリー (CI/CD) パイプラインをより堅牢なものにします.
Source / 元記事
この記事について
この記事は、公開されているニュース、論文、公式発表、RSSフィードなどをもとに、AIが要約・補足調査・考察を行って作成しています。
元記事の完全な翻訳・逐語的な要約ではなく、AIによる背景説明や開発者向けの考察を含みます。
重要な技術仕様・価格・提供状況などは、必ず元記事または公式情報をご確認ください。
