🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

本論文は、異種メモリ階層全体にオプティマイザ状態を動的に分散し、高価な前処理計算を GPU 訓練のクリティカルパスから分離することでオーバーヘッドを削減し収束を加速し、大規模言語モデルに対する実用的でスケーラブルな 2 次最適化を可能にするランタイムシステム「Asteria」を導入する。

原著者： Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

公開日 2026-05-18

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で超知能なロボット（大規模言語モデル）に物語を書かせることを想像してみてください。そのためには、ロボットの間違いを見て改善方法を教える「教師」（オプティマイザ）が必要です。

現在、ほとんどの教師はAdamWと呼ばれるシンプルな方法を使用しています。これは宿題をチェックし、いくつかの間違いを見つけ、小さく素早く修正する学生のようなものです。これは速く、よく機能しますが、本当に上手になるには（何百万もの例を含む）膨大な練習が必要です。

第二階最適化（SOAP や Shampoo など）と呼ばれる、より賢く高度な教師方法があります。この教師は、何が間違っていたかだけでなく、間違いの「形状」を分析して問題を深く理解します。これははるかに速く学習し、より少ない例で済みます。しかし、大きな問題があります。この賢い教師は非常に重く、複雑な数学を行うために莫大なメモリと計算能力を必要とし、コンピュータをクラッシュさせたり、単純な教師の方が実際には仕事を完了するほど遅く動作させたりすることがよくあります。

Asteria の登場：「スマートな物流マネージャー」

オックスフォード大学の研究者たちは、Asteriaと呼ばれる新しいシステムを構築しました。Asteria は新しい教師ではなく、スマートな教師が学校を壊すことなく最終的に仕事をこなせるよう、教室全体を再編成する天才的な物流マネージャーとして考えてください。

以下は、Asteria がシンプルなアナロジーを用いて解決する 3 つの最大の課題です。

1. 「狭い部屋の家具」問題（メモリ）

課題： スマートな教師は、巨大で複雑なチャート（行列）を頭（GPU メモリ）に保持する必要があります。標準的なコンピュータでは、十分なスペースがありません。キングサイズのベッド、ダイニングテーブル、ワードローブをスタジオアパートに詰め込もうとするようなものです。部屋がいっぱいになり、教師は辞めざるを得なくなります。
Asteria の解決策： Asteria は折りたたみ家具の専門家のように機能します。教師が一度にすべてを手に持っている必要はないと理解しています。

最も活動的なチャートの部分を GPU（机）に保持します。
重く、頻繁に使用されない部分は CPU（廊下）やハードドライブ（ガレージ）に移します。
重要なのは、重い部分を必要な時にだけ机に戻すことです。これにより、スマートな教師はスーパーコンピュータと同じように、小さなラップトップでも作業できます。

2. 「渋滞」問題（速度）

課題： 数ステップごとに、スマートな教師はチャートを更新するために、巨大なパズルを解くような大規模で複雑な計算を行う必要があります。これは時間がかかり、クラス全体の作業を停止させます。単一の荷物を荷下ろししている間に、配送トラックが高速道路全体を塞いでいるようなものです。GPU（コンピュータの脳）は計算を待ってアイドル状態になります。
Asteria の解決策： Asteria は並列アセンブリラインを導入します。

重い数学を行うためにメインのクラスを停止する代わりに、Asteria は「重労働」をバックオフィスの作業員チーム（CPU）に送ります。
メインのクラス（GPU）がロボットに教え続ける間、バックオフィスの作業員は背景で複雑なパズルを静かに解いています。
クラスが新しいチャートを必要とする頃には、バックオフィスはすでに完了しており、それらをスライドさせて渡しています。メインのクラスは決して停止する必要がありません。「渋滞」は消えます。

3. 「グループチャット」問題（分散学習）

課題： 複数のコンピュータで同時に学習する場合、通常、次のステップに進む前に全員が正確に同じ情報で合意するために停止する必要があります。これは、全員が最も遅い人の返信を待つまで誰も入力できないグループチャットのようなものです。
Asteria の解決策： Asteria は**「十分良い」方針**を使用します。

全員が完全に同期するのを待つ代わりに、コンピュータが短時間、わずかに「古くなった（わずかに時代遅れの）」情報で作業することを許可します。
更新は本当に必要な時だけ送信されます。
これは、5 分おきにプロジェクト全体を停止するのではなく、後で更新をキャッチアップしながら互いに作業し続けることを信頼するチームのように、グループを速く動かし続けます。

結果：彼らは何を見つけましたか？

研究者たちは、強力な単一コンピュータ（Nvidia DGX Spark）と大規模なコンピュータクラスター（Nvidia GH200）を含む実際のハードウェアで Asteria をテストしました。

小型マシンでも動作します： 以前はこのスマートな教師のメモリ要件を処理できなかった単一マシンで、大規模言語モデル（10 億パラメータ）を学習させることができました。
リアルタイムで速い： 「渋滞」を隠すため、数学がより難しくても、実際のクロック時間は短縮され、学習が完了します。
エネルギーを節約します： コンピュータの脳（GPU）を忙しく保ち、計算を待ってアイドル状態にさせないことで、Asteria は古い重厚な方法と比較して、同じ結果を達成するために実際には少ない総エネルギーを使用します。
品質は失われません： モデルは「ネイティブ」（最適化されていない）なスマートな教師と同じように学習しますが、はるかに速く、安価に到達します。

まとめ：
Asteria は新しい数学の公式を発明するわけではありません。代わりに、コンピュータがその数学を実行する方法を再考します。重労働をメインの作業から分離し、利用可能なすべての記憶領域を賢く使い、コンピュータを非同期で動作させます。これにより、「理論的には素晴らしいが実際には不可能な」方法を、次世代の AI を学習するための実用的なツールに変えます。

技術サマリー：Asteria – スケーラブルな LLM 学習のためのランタイム編成型 2 次最適化

1. 問題定義

2 次最適化手法（Shampoo、SOAP、KL-Shampoo など）は、AdamW などの 1 次最適化器と比較して、大規模言語モデル（LLM）の学習において優れた統計的効率性と収束速度を提供します。しかし、その大規模な実用的展開は、深刻なシステムレベルの障壁によって阻まれています。

メモリ制約: 2 次手法は、大規模な共分散行列（ $L$ および $R$ ）の保存と、密な固有値分解の実行を必要とし、 $O(N^2)$ のメモリ使用量と $O(N^3)$ の計算コストをもたらします。メモリ制約のある環境（ユニファイドメモリアーキテクチャを備えた単一ノードワークステーションなど）では、これらの状態が物理メモリを枯渇させ、Out-of-Memory（OOM）エラーを引き起こすか、モデルサイズやバッチサイズの削減を強要します。
レイテンシの急増: 逆平方根（例： $L^{-1/4}$ ）の周期的な計算には、立方コストの行列演算が伴います。既存の実装では、これらの同期更新が GPU 学習パイプラインをブロックし、分散学習（例：FSDP）の効率性に不可欠な計算と通信の細かなオーバーラップを妨げる深刻な周期的なレイテンシの急増を引き起こします。
同期オーバーヘッド: 現在の分散 2 次手法は、すべてのランクにわたる厳格で同期されたグローバル状態更新（AllGather）に依存しています。これは、高速なノード内通信と低速なノード間同期を結合し、システムが非同期でトポロジーを認識した陳腐化（staleness）境界を利用することを妨げています。

DTensor ベースのシャーディングなどの既存の解決策は、リソース制約のある単一ノード設定や、グローバル同期の通信オーバーヘッドが利点を上回る状況において、これらの問題に対処できていません。

2. 手法：Asteria ランタイム

Asteria は、PyTorch FSDP の上に構築されたハードウェア・ソフトウェア協調設計のランタイムシステムです。これは、3 つの中核メカニズムを通じて、2 次最適化ロジックを重要な GPU 学習パスから分離します。

A. アーキテクチャ適応型非対称メモリティアリング

状態を均質なベクトルとして扱う 1 次最適化器とは異なり、Asteria は 2 次状態の非対称なライフサイクルを認識します。

GPU: 学習中に頻繁に更新されるfactor_matrices（クリフォック因子）を保存します。
ホスト（CPU/UVM）: inv_factor_matrices（逆因子）を保存します。これらは GPU メモリの圧力を回避するために、CPU からのアクセス可能なメモリに維持されます。
NVMe（オプション）: 容量制限のある実行のためのステージングティアを提供し、シリアライズされた計算または JIT ページングを介して状態をスワップアウト可能にします。
メカニズム: ランタイムはホスト側の因子行列のスナップショットを捕捉し、非同期 CPU ワーカーを派遣して逆平方根を計算し、結果をホストメモリに書き戻します。必要に応じて、GPU はユニファイド仮想メモリ（UVM）ビューまたは明示的なプリフェッチを介してこれらにアクセスし、冗長なホストからデバイスへのコピーを回避します。

B. フック編成型シャドウ状態パイプライン

レイテンシの急増を排除するため、Asteria は $O(N^3)$ の逆平方根計算を重要な GPU パスから除去します。

シャドウストリーム: 低優先度の CUDA ストリームが、主要な学習ストリームをブロックすることなく、GPU 側の状態ステージング（プリフェッチとページング）を処理します。
非同期実行: 高価な逆平方根更新は、CPU ワーカースレッド上で非同期に実行されます。
フック統合: 軽量な PyTorch フック（register_forward_hook、register_full_backward_pre_hook）がスケジューリング信号として機能します。これらは、準備完了の状態転送の排出と、すぐに消費されることになる逆因子のプリフェッチをトリガーし、GPU がメイン学習ループを停止させることなく新鮮な前処理条件（preconditioners）を入手できるようにします。

C. トポロジー認識型有界陳腐化コヒーレンス

分散学習において、Asteria は厳格なグローバル同期を、選択的コヒーレンスプロトコルに置き換えます。

陳腐化の追跡: CoherenceRegistry が各前処理条件ブロックの鮮度を追跡します。設定された陳腐化予算を超えたブロックのみが同期されます。
階層的同期: 同期は 2 段階で行われます。まずノードローカルなプロセスグループ内で行われ、次にノード代表者間で行われます。これにより、クロスノードトラフィックを最小化し、通信をホストアクセス可能なメモリに留めることで、不要な CPU から GPU への往復を回避します。
有界バリア: ランタイムは、CPU ジョブが進行中に前処理条件のバージョンが過度に陳腐化した場合、最適化器が完了を待つことを保証するローカルバリアを強制し、鮮度とスループットのバランスを取ります。

3. 主要な貢献

ランタイムの分離: Asteria は、2 次最適化を実用的にするためにアルゴリズムを単純化するのではなく、ランタイムレベルでの状態管理、バックグラウンド計算、同期を再考することによって可能になることを実証します。
メモリスケーラビリティ: 従来のシャーディングが不可能な、単一ノードのメモリ制約ハードウェア（例：128GB ユニファイドメモリを備えた単一 GB10 GPU）での 2 次学習を可能にします。
レイテンシの隠蔽: 逆平方根計算をホストへ移動し、シャドウパイプラインを使用することで、Asteria は周期的なレイテンシの急増を平坦化し、1 次手法に匹敵する滑らかなハードウェア実行プロファイルを実現します。
エネルギー効率: システムはエネルギー損失のトレードオフを改善し、GPU のアイドル時間を削減し、ホスト側計算のオーバーラップを最適化することで、単位エネルギーあたりの損失削減量を増大させます。

4. 実験結果

著者らは、OLMo 学習フレームワークを使用して、2 つのプラットフォーム上で OLMo ファミリーモデル（660M、1B、7B パラメータ）における Asteria を評価しました。

単一ノード（DGX Spark、GB10 GPU）:
- スループット: Asteria-SOAP と Asteria-KL-Shampoo は、ネイティブ実装で見られた周期的なステップ時間の急増を排除しました（ピークステップ時間を約 96 秒から約 66.5 秒に削減）。
- エネルギー: ネイティブの 2 次手法は AdamW よりも約 17-20% 多くの総 SoC エネルギーを消費していましたが、Asteria はこのギャップを大幅に縮小しました（例：Asteria-KL-Shampoo は AdamW よりも約 7.9% しか多くのエネルギーを消費しませんでした）。
- 効率性: Asteria-KL-Shampoo は、正規化された損失削減効率（3.3010）で最高を記録し、ネイティブの 2 次手法と AdamW の両方を上回りました。
マルチノード（GH200 クラスタ）:
- 収束: Asteria 変種は、ネイティブの 2 次手法に近い収束軌道を維持し、目標損失レベルに到達するまでのウォールクロック時間で AdamW を大幅に上回りました。
- 陳腐化許容性: 陳腐化予算（ $S$ ）を変化させた実験により、 $S$ を 1 から 5 に増加させることで、最終モデルの品質にほとんど影響を与えずに学習時間を大幅に短縮できることが示されました。
- 強スケーリング: 2 ノードから 16 ノードへの 7B モデルのスケーリングにおいて、Asteria はネイティブのベースラインと比較して、より高いスピードアップと低いステップあたり時間を達成し、状態維持の相対コストが増加しても効果的であることを証明しました。

5. 意義と主張

本論文は、実用的な 2 次 LLM 学習の主な障壁は、曲率認識更新の数学的コストではなく、2 次状態維持と現代の学習スタックとの間のシステムミスマッチであると主張しています。

問題の再定義: 著者らは、2 次最適化はアルゴリズム、ランタイムスケジューリング、および異種メモリシステムを含む協調設計問題と見なされるべきであると主張しています。
実現可能性: 彼らは、大規模な集約 GPU メモリが 2 次学習の根本的な前提条件ではないことを実証しています。ライフサイクルを認識した配置と段階的アクセスにより、これらの手法は厳しいメモリ予算でも実行可能です。
実用性: 前処理条件の鮮度と通信を、厳格なグローバル状態ではなく、ランタイムによって管理されるリソースとして扱うことで、Asteria は 2 次最適化の統計的約束と現実世界での展開可能性の間のギャップを狭めます。

本論文は、Asteria がアルゴリズムの単純化を必要とすることなく、より良い収束とエネルギー効率を提供し、LLM 学習のための実行可能なシステム選択として 2 次最適化を実現すると結論付けています。