技術概要：TrainMover

問題定義

大規模な大規模言語モデル（LLM）のトレーニングジョブは、しばしば数万の GPU にまたがり、数週間から数ヶ月にわたって実行されるため、中断の影響を非常に受けやすい。これらの中断は、ハードウェアの異常（例：GPU の故障、過熱）、ソフトウェアの競合（例：OS ノイズ、ガベージコレクション）、ネットワークの障害、および管理イベント（例：ドライバの更新、リソースの再配分）に起因する。

これらの中断に対処する既存のソリューションは、重大な非効率性を抱えている：

1. 停止・再スケジューリング・再初期化: 業界標準では、ジョブを停止し、故障したノードを交換し、最新のチェックポイントから再初期化する。このプロセスには、チェックポイントの読み込み、NCCL グループの再インスタンス化、およびコールドウォームアップに起因する甚大なダウンタイム（例：8,192 GPU で約 6.5 分）が生じる。
2. ランタイム再構成: Oobleck や Parcae などのシステムは、ジョブを停止せずにクラスターを弾力的に拡張しようとする。しかし、これらも依然としてクリティカルパスのボトルネックに直面している：新しいマシン（ジョイナー）をオンライン化するには、再初期化とグローバル同期が必要であり、これによりジョブ全体が数十秒から数百秒にわたって停止してしまう。

クラスターサイズが拡大する（例：64K または 128K GPU へ）につれ、中断の頻度が増加し、回復に伴うダウンタイムが支配的なボトルネックとなり、実効トレーニング時間比率（ETTR）を劇的に低下させ、数百万の GPU 時間を浪費する。

手法

TrainMover は、マシン交換時のダウンタイムを最小化し、メモリオーバーヘッドを排除するように設計された耐障害性 LLM トレーニングランタイムである。これは、新しいマシン（ジョイナー）の初期化を進行中のトレーニングジョブのクリティカルパスから切り離すことでこれを達成する。システムは、トレーニングと重複する「準備フェーズ」と、最小限のダウンタイムで実行される「切り替えフェーズ」という 2 フェーズのワークフローで動作する。

中核的な手法は、以下の 3 つの主要な技術に依存している：

1. 通信を伴わないサンドボックス化されたウォームアップ

初期化中のグローバルなブロッキングを回避するため、TrainMover は、ジョイナーがメインのトレーニングループに参加する前に独立して「シャドウイテレーション」を実行できるサンドボックス環境を導入する。

• メカニズム: 記録レイヤが、メインジョブの最初の数イテレーション中の集合通信呼び出し（例：all-reduce）を傍受し、結果として生じるテンソルを保存する。
• 再生: ジョイナーがサンドボックスに入ると、モデル状態をロードし、シャドウイテレーションを実行する。他のマシンと通信する代わりに、サンドボックスは通信呼び出しを傍受し、事前に記録されたテンソルを再生する。
• 利点: これにより、アクティブなマシンの支援やグローバル同期を必要とすることなく、すべての暗黙的な初期化ステップ（JIT コンパイル、CUDA コンテキストの作成、メモリ割り当て、データローダのセットアップ）がトリガーされる。

2. 2 フェーズのデルタベース通信グループ設定

従来の集合通信ライブラリ（CCL）の設定（例：NCCL）は静的であり、メンバーシップが変更されるとグローバルな teardown と再構築が必要となる。TrainMover は、動的な 2 フェーズのアプローチを導入する：

• フェーズ 1（準備）: クリティカルパスに属さないセットアップステップ（ネットワークのブートストラップ、トポロジーの発見、新しいリンクの接続確立）をバックグラウンドで実行し、進行中のトレーニングと重複させる。これには、必要な接続変更のみを特定する「デルタ」トポロジーの計算が含まれる。
• フェーズ 2（切り替え）: ジョイナーが準備できると、システムはデルタ更新を実行し、離脱者からジョイナーへの特定のマシン間接続（例：RDMA キューペア）のみを切り替える。このフェーズは極めて高速であり、グローバル同期や追加の GPU メモリを必要としない。

3. 汎用スタンバイ設計

特定の故障した役割が不明な予期せぬ故障に対処するため、TrainMover は LLM トレーニングランクの構造的対称性を活用する。

• 対称性: ほとんどのトレーニング役割（データ並列、テンソル並列）は同一である。パイプライン並列（PP）は、3 つの異なる役割タイプ（最初のステージ、中間のステージ、最後のステージ）のみを導入する。
• メカニズム: 単一の「汎用スタンバイ」マシンを事前にウォームアップし、任意の役割をカバーできるようにする。バックグラウンドで、異なる役割タイプ（最初、中間、最後）のウォームアップイテレーションを順次実行する。
• 利点: 予期せぬ故障が発生すると、スタンバイは即座にジョイナーとして昇格する。すでにサンドボックス化されたウォームアップとフェーズ 1 の CCL セットアップを完了しているため、どの特定のマシンが故障したかに関わらず、直接切り替えフェーズに入ることができる。

主要な貢献

1. 耐障害性ランタイムアーキテクチャ: TrainMover は、ウォームアップと CCL セットアップという重たい初期化の負担をクリティカルパスから外し、ジョイナーがバックグラウンドで準備している間もトレーニングジョブが中断なく継続できるようにする。
2. ゼロメモリオーバーヘッド: 新しいグループのために追加の GPU メモリを必要とする再構成システムとは異なり、TrainMover のデルタベースの CCL セットアップと慎重なメモリ管理（状態転送のために勾配バッファを再利用）は、移行中に追加の GPU メモリオーバーヘッドをゼロにする。
3. 役割非依存の回復: 汎用スタンバイ設計により、単一の事前にウォームアップされたマシンが故障した任意の役割を回復でき、各特定のランクごとに専用のスタンバイプールを必要としない。
4. スケーラビリティ: システムはスケールに非依存するように設計されており、ダウンタイムはクラスターの総サイズではなく、切り替えられる接続数（離脱者/ジョイナーのペアに局所的）によって決定される。

結果

1,024 GPU のテストベッドで実施され、64,000 超の GPU スケールに拡張された実験は、以下の結果を示している：

• ダウンタイム: TrainMover は、1,024 GPU スケールにおいて、予期せぬ中断と予期せぬ中断の両方で、一貫して約20 秒のダウンタイムを達成する。これは、約 300 秒を要する Megatron-LM（停止・再起動）や、100〜200 秒以上を要する Oobleck/Parcae などの再構成システムと比較して、大幅な削減である。
• 効率性: 64K GPU スケールにおいて、TrainMover は最良の代替手段（Parcae）と比較して、無駄な GPU 時間を**55%**削減すると予測され、標準的なチェックポイント再起動方法を大幅に上回る。
• リバランス: TrainMover は、128 から 1,024 GPU までのスケール全体で >0.97 の ETTR を維持しつつ、10 分ごとなどの頻繁なリソース再配分を3% 未満のトレーニングスループット損失でサポートする。
• ストレイガーの処理: ストレイガー GPU を伴うシナリオにおいて、TrainMover は高い効率を維持し（理想的なイテレーションごとのチェックポイントと比較して約 4.7% のみ損失）、再起動を遅延させる戦略や進捗を失う戦略を上回る。

重要性

本論文は、TrainMover が LLM トレーニングの拡張における重要なボトルネックに対処していると主張している。トレーニングクラスターが 10 万超の GPU に成長するにつれ、中断の頻度とダウンタイムのコストは許容できなくなるものとなる。回復時間を数分から数秒に短縮し、動的再構成に伴うメモリオーバーヘッドを排除することで、TrainMover は以下のことを可能にする：

• より高い ETTR: 高価な GPU インフラの利用率を最大化する。
• コスト削減: 無駄な計算時間の数百万ドルを節約する（64K スケールでは週あたり 140 万 GPU 時間の節約が予測される）。
• 運用の柔軟性: 長期間実行されるトレーニングジョブを停止することなく、頻繁なメンテナンス、セキュリティパッチの適用、リソースの再配分を可能にする。

著者は、彼らのアプローチが基盤となるモデルアーキテクチャやトレーニングアルゴリズムの変更を必要としないことを強調しており、これは生産規模の LLM トレーニングに対する実用的な解決策である。