> [!abstract] 概要（arXiv abstract の日本語訳） > 大規模言語モデルの訓練はさまざまな領域でますます重要になっているが、頻発する障害によって妨げられ、大きな時間的・経済的コストを生じている。クラウドベース環境における現行の障害回復手法は、生じる多様かつ複雑なシナリオに十分対応できておらず、クラスタ全体のコスト影響を考慮せず個々のタスクのダウンタイム解消に狭く焦点を当てている。我々は Unicron を提案する。Unicron は大規模言語モデル訓練における効率的な自己修復(self-healing)のために設計されたワークロードマネージャである。Unicron はクラスタ内の複数の並行タスクにわたって障害関連コストを最小化することで訓練プロセスを最適化する。主要な機能として、追加オーバーヘッドなしのリアルタイム誤り識別のためのインバンド誤り検知、最適な再構成のための動的コスト認識プラン生成機構、および状態変化時のダウンタイムを削減する効率的な遷移戦略がある。128 GPU の分散クラスタに展開した結果、Unicron は最先端手法と比較して訓練効率を最大 1.9× 改善し、障害回復コストを大幅に削減して大規模言語モデル訓練の信頼性を高めた。 ## 論文情報 - **タイトル**: Unicron: Economizing Self-Healing LLM Training at Scale - **著者**: Tao He, Xue Li, Zhibin Wang, Kun Qian, Jingbo Xu, Wenyuan Yu, Jingren Zhou（Alibaba Group、Zhibin Wang は南京大学も兼任） - **媒体**: arXiv プレプリント (arXiv:2401.00134v1) - **発表年**: 2023-12-30（arXiv 投稿） - **連絡先**: [email protected] - **コード**: 公開情報なし ## 概要 Unicron は Megatron を基盤として構築された分散ワークロードマネージャで、LLM 訓練クラスタにおける障害回復の経済的効率を最大化する。インバンド誤り検知・コスト認識プラン生成・効率的状態遷移の三機構を統合し、クラスタ内の複数タスクにわたるリソース全体最適化を実現する。Alibaba Cloud の本番クラスタ分析から出発し、128 GPU の評価クラスタで Megatron 比 1.9× の累積 WAF 改善を実証した。 ## 問題設定 **背景となる障害統計**: Alibaba Cloud の分析では、最もリソース集約的なタスク(上位 5%)の異常終了率が 43.4% に達する。障害原因の内訳は NCCL タイムアウト 10.1%・ECC エラー 5.0%・NVLink エラー 4.5%・リンクフラッピング 3.9%・タスクハング 3.1% など。正常終了は 56.6% のみ。 **既存手法の問題点**: 障害回復の既存アプローチはそれぞれ部分的な解決に留まる。 - **チェックポイント系**: 再起動時の再計算コストを下げるが、複雑な障害シナリオには対応不足 - **弾性訓練・スケジューリング系** (Oobleck, Varuna, Bamboo): 訓練の継続性を保てるが Megatron と比べて著しく効率が低い（図3a: スループットで Megatron に大きく劣る） - **冗長計算・ホットスペア系**: 継続性には有効だが経済的・リソースコストが大きい **核心の問題認識**: 単一タスクのダウンタイムを最小化する方向は最適ではない。クラスタ全体の失われたスループット（WAF）を最小化する視点が欠けている。 **入力/出力**: - 入力: n 台のワーカーで m 個のタスクを並行訓練するクラスタ、各タスクのモデルサイズ・重み・最小リソース要件 - 出力: 障害時に WAF を最大化する最適なワーカー配分（再構成プラン） **コスト形式化**: $C_{\text{recovery}} = C_{\text{detection}} + C_{\text{transition}} + C_{\text{sub-healthy}}$ ## 提案手法 ### アーキテクチャ Unicron は Megatron への非侵入(non-intrusive)統合を原則とし、**Unicron Agent** と **Unicron Coordinator** の 2 コンポーネントを追加する。 **Figure 5: The system architecture of Unicron** ![[_attachments/arxiv-2401.00134/fig05-unicron-architecture.png]] (Figure 5. Unicron のシステムアーキテクチャ。各ノードの Unicron Agent は GPU を監視するモニタリングスレッドと In-memory チェックポイントを管理し、Unicron Coordinator の Status Monitor(etcd 実装の分散 KV)・Failure Handler・Plan Generator・Task set と通信する。) **Unicron Agent**（ノード単位）: - CPU 上の専用モニタリングスレッドを各 GPU に割り当て（GPU に追加負荷なし） - コーディネータとの常駐接続を維持（ノード可用性の検知） - GEMINI 方式のインメモリチェックポイントを管理し非同期でリモートストレージに転送 **Unicron Coordinator**（クラスタ単位）: - etcd 実装の分散 KV（Status Monitor）でプロセス状態を集約 - 障害検知時の適切な回復戦略を決定・指揮 - 最適再構成プランを生成（動的計画法） - タスクセット管理、クラウドサービスとの調整 ### 誤り検知（§4） Unicron はインバンド誤り検知を採用し、4 種の検知手法を重ねる（表1）。 | 検知手法 | 検知する障害状態 | 深刻度 | |---|---|---| | ノード健全性監視 | 接続断 | SEV 1 | | プロセス監視 | 異常終了 | SEV 2 | | 例外伝播 | CUDA・NVLink・ECC・DMA・GPU ドライバエラー等 | SEV 1〜2 | | オンライン統計監視 | NCCL タイムアウト・リンクフラッピング・タスクハング | SEV 2〜3 | **オンライン統計監視の設計**: 通常の訓練ではイテレーション完了時間が一定の範囲に収まる（図6: 緑点群）。モニタリングスレッドが平均イテレーション時間の 3× を超えたら障害と判定する。これにより Megatron デフォルトの 30 分タイムアウトを大幅に短縮する（実測: 3× D_iter ≈ 数分）。 ### エラーハンドリング（§4） 深刻度に応じて 3 段階の回復戦略を適用する。 **Figure 7: The error handling workflow in Unicron** ![[_attachments/arxiv-2401.00134/fig07-error-handling-workflow.png]] (Figure 7. Unicron のエラーハンドリングワークフロー。SEV3 障害はインプレース再試行、SEV2 はプロセス再起動、SEV1 はクラスタ再構成へとエスカレートする。ノード参加④・タスク完了⑤・タスク起動⑥もプラン再生成のトリガーとなる。) 1. **インプレース再試行（SEV 3）**: リンクフラッピング等の一時的な問題に対して元の場所で再試行。成功すれば即座に通常訓練に戻る 2. **プロセス再起動（SEV 2）**: CUDA エラー等はプロセスを再起動。GPU 数・並列設定は変えず、他の DP レプリカまたは最新チェックポイントから状態を復元。失敗すれば SEV 1 に格上げ 3. **クラスタ再構成（SEV 1）**: 障害ノードを隔離し最適再構成プランに従って資源を再配分。回復ノードや新規ノードの参加④・タスク完了⑤・タスク起動⑥も再構成トリガーとなる ### 最適再構成プラン生成（§5） **WAF（Weighted Achieved aggregate FLOP/s）**の定義: $F(t, x) = \begin{cases} w(t) \cdot T(t, x) & \text{if } (t, x) \vdash T_{\text{necessary}}(t) \\ 0 & \text{otherwise} \end{cases}$ - $T(t, x)$: タスク $t$ に $x$ ワーカーを割り当てた時の達成集約 FLOP/s（クラスタキャリブレーション + 自動実行プラン生成で推定） - $w(t)$: タスクの優先度重み（既定 1.0、推奨範囲 0.5〜2.0） - $T_{\text{necessary}}(t)$: タスクの最小リソース要件 **最適化問題**（式(3)）: $\arg\max_{x_1', \ldots, x_m'} \sum_{i=1}^{m} G(t_i, x_i')$ ここで $G(t_i, x_i') = F(t_i, x_i') \cdot D_{\text{running}}(n') - F(t_i, x_i) \cdot \mathbf{1}(t_i, x_i \to x_i') \cdot D_{\text{transition}}$ 遷移ペナルティ項は頻繁な再構成を抑制し、健全なタスクへの不要な再配分を防ぐ。 **解法**: 動的計画法で $O(mn^2)$ の時間計算量。実際には $m$ と $n$ がともに中程度のため実用的。ルックアップテーブルを事前計算することで再構成必要時の解法を $O(1)$ に短縮できる。 ### 遷移戦略（§6） **失敗イテレーションからの再開**: - **シナリオ #1（all-reduce 前の障害）**: 障害 DP ランクのマイクロバッチを残りランクにラウンドロビン再分配し、部分的に蓄積されたグラジェントを再利用して訓練を継続 - **シナリオ #2（all-reduce 後の障害）**: グラジェントがすでに還元済みか否かに応じて、省略または部分グラジェントの再計算 **状態移行の近傍原則**: 障害ノードの状態は、同じ DP グループの健全ランクからの複製を優先（DP レプリカが既に状態を保持）。不可能な場合はインメモリチェックポイントへフォールバック。 ## 新規性 | 先行研究 | 課題 | Unicron の解決 | |---|---|---| | Oobleck, Bamboo, Varuna | 弾性を持つが Megatron 比でスループット大幅劣化 | Megatron をそのまま継承し追加最適化 | | チェックポイント手法 | 個別タスクのダウンタイム削減に特化 | クラスタ全体の WAF 最大化を目標に | | SLURM, Kubernetes | LLM 訓練の内部を知らずブラックボックス扱い | LLM 訓練のセマンティクスを理解した自己修復 | | 既存の out-of-band 監視 | タイムアウトまで長時間ハング(30 分) | 3× D_iter(数分)での早期検知 | **主な新規性**: 1. クラスタ全体の複数タスクを統合する「コスト認識プラン生成」という目標設定 2. WAF という複数タスクのクラスタ効率を測る指標の定義と最適化 3. Megatron の訓練セマンティクス（グラジェント蓄積の構造）を利用した部分結果再利用型遷移戦略 ## 実験設定 - **プラットフォーム**: Alibaba Cloud。16 インスタンス × 8 NVIDIA A800(80GB) GPU = 合計 128 GPU。各インスタンスは 96 CPU コア、1,600 GB CPU メモリ。インスタンス内は NVSwitch 接続、インスタンス間は 4× 200Gbps Ethernet NIC - **ソフトウェア**: Megatron v23.08、PyTorch v1.12.1、CUDA v11.3 - **ワークロード**: GPT-3 モデル（1.3B・7B・13B・70B・175B パラメータ） - **ベースライン**: Megatron（ホットスペア付き）、Oobleck、Bamboo、Varuna - **障害トレース**: trace-a（実環境 8 週間、SEV 1 障害 10 件 + 他 33 件）、trace-b（trace-a の 20× 増幅、7 日間、SEV 1 障害 26 件 + 他 80 件） ## 実験結果 ### 誤り検知効率（§7.2、表2） | ケース | 検知手法 | Unicron | Megatron なし(D_timeout) | |---|---|---|---| | 1（ノード終了） | ノード健全性監視 | 5.6 秒 | 5.7 秒 | | 2（プロセス終了） | プロセス監視 | 1.8 秒 | 30 分 | | 3（例外スロー） | 例外伝播 | 0.3 秒 | 30 分 | | 4（性能劣化） | 統計監視 | 3× D_iter | 30 分 | ケース 2〜4 では既定の 30 分タイムアウトに対し、Unicron は数秒〜数分で検知する。 ### 遷移効率（§7.3、図9） SEV 1 障害時の遷移時間（障害検知から訓練再開まで）: - Megatron・Varuna: 最後のチェックポイントからの再起動が必要で遷移時間が長い - Oobleck・Bamboo: 動的再構成で遷移時間を短縮するが Unicron より長い - Unicron: 部分結果の再利用により遷移時間を最小化し、クラスタサイズが増加しても安定した遷移時間を維持 ### 訓練スループットと WAF（§7.4、図10） - **単一タスク（図10a/10b）**: Unicron のスループット・FLOP/s 比率は Megatron と同等（追加オーバーヘッドなし） - **複数タスク（図10c）**: 5 ケース全てで Unicron が最高の WAF を達成。動的計画法プランが均等配分・重み比例配分・サイズ比例配分の各ベースラインを上回る ### 総合訓練効率（§7.5、図11） - **trace-a**: Unicron は Megatron 比 1.2×、Bamboo 比 4.6×、Oobleck 比 3.7×、Varuna 比 4.8× の累積 WAF - **trace-b（高頻度障害）**: Unicron は Megatron 比 **1.9×**、Bamboo 比 4.8×、Oobleck 比 3.8×、Varuna 比 **5.8×** の累積 WAF。障害頻度が上がるほど Unicron の優位性が拡大する（Megatron は頻繁な障害で回復コストが指数的に増大するため） ## 考察 - **単一タスク最適化 vs. クラスタ全体最適化**: 既存手法は個別タスクのダウンタイムを最小化するが、Unicron はクラスタ全体の WAF を最大化するため、ある健全タスクを再配分することで全体効率が上がる場合に積極的に再構成する - **3D 並列の非線形性**: GPU 数と実効 FLOP/s の関係が非線形・非単調であり、GPU を増やしても性能が下がることがある（図4）。Unicron はこの特性を考慮した最適配分を行う - **trace-b での Megatron の劣化**: Megatron は通常動作では最高効率だが、高頻度障害下では回復コストが積み重なり、弾性を持つ手法への優位性が消える ## 強み / 弱点・課題 **強み**: - Megatron の既存最適化を完全に継承しつつ自己修復を追加（non-intrusive） - クラスタ全体の複数タスクを統合最適化する視点 - インバンド検知により CPU 上でのみ動作し GPU に追加負荷なし - 動的計画法の事前計算で再構成判断を $O(1)$ に短縮可能 **弱点・課題（論文が述べる限界）**: - 評価規模が 128 GPU 止まり（大規模クラスタでのスケーラビリティは未検証） - チェックポイント機構の最適化は本論文のスコープ外（GEMINI 方式をそのまま利用） - 障害統計の取得は Alibaba Cloud 固有で他クラウドへの一般化が未検証 - 弾性（GPU 数の変更）に伴う最適並列設定の推定（T(t, x) の精度）は将来課題 **読み取れる懸念**: - 高頻度障害(trace-b)の想定は実運用の 20× であり、実際の本番効果は trace-a(1.2×)が現実的な上界に近い - WAF 指標はタスク重みの設定に依存し、ユーザが適切な重みを設定できるかが前提 ## 関連 - 概念: [[耐障害LLM訓練]] / [[LLM分散学習]] / [[弾性LLM訓練]] / [[チェックポイント]] - エンティティ: [[Alibaba Group]] / [[Alibaba Cloud]] / [[Megatron-LM]] / [[Tao He (Alibaba)]] / [[Jingren Zhou]] - 関連 MOC: [[分散深層学習 - MOC]] ## 出典 - arXiv:2401.00134v1 — Tao He, Xue Li, Zhibin Wang, Kun Qian, Jingbo Xu, Wenyuan Yu, Jingren Zhou, "Unicron: Economizing Self-Healing LLM Training at Scale", 2023-12-30