Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

本論文は、Transformer モデルのトレーニングにおけるメモリ制約を克服するため、フォワード・バックワード・更新フェーズのメモリ使用量変動を巧みに利用してオプティマイザ状態を CPU と GPU 間で動的に転送する「Deep Optimizer States」という手法を提案し、DeepSpeed と統合することで最先端の手法に比べて 2.5 倍高速な反復処理を実現したことを示しています。

要約

巨大な AI を動かすための「賢い引越し」作戦

～Deep Optimizer States（深いオプティマイザ状態）の仕組みをわかりやすく解説～

皆さん、最近の AI（大規模言語モデル）がどれほど巨大になっているかご存知でしょうか？パラメータ（AI の知識の断片）の数が数百億、数千億にも及ぶようになり、これらを学習させるには、スーパーコンピュータのような巨大なメモリが必要です。

しかし、現実には「メモリ不足」という壁にぶつかることがよくあります。まるで、**「本屋の棚（GPU メモリ）が小さすぎて、全巻揃った百科事典（AI モデル）を一度に置けない」**ような状態です。

この論文は、その問題を解決するための新しい「引越し作戦」を提案しています。

---

🏠 従来の問題：「狭い部屋」と「遅い廊下」

AI を学習させるには、大きく分けて 3 つの作業が必要です。

1. 読み込み（フォワード）：本を読んで理解する。
2. 振り返り（バックワード）：間違えたところを反省し、メモを取る。
3. 修正（アップデート）：メモをもとに、知識（パラメータ）を書き換える。

【従来のやり方（DeepSpeed など）】
GPU（高性能な作業台）のメモリが足りないため、知識の「書き換え作業（オプティマイザの状態）」を、CPU（普通のデスクトップ PC）のメモリに預けていました。

• 問題点 1：廊下が狭い（PCIe 帯域の限界）
  GPU と CPU の間には「廊下（PCIe 接続）」がありますが、これが狭く、一度に運べる荷物の量に限界があります。
• 問題点 2：作業台が空いているのに、作業員が待っている
  知識を書き換える作業は、CPU が行います。しかし、CPU は GPU に比べて作業速度が遅い上、GPU が「読み込み」や「振り返り」をしている間、CPU は「書き換え」を待たなければなりません。逆に、CPU が書き換えをしている間、GPU は「廊下を待って」いるだけで、作業台が空転しています。
  **「廊下で荷物を運んでいる間、作業台も、作業員も、どちらもダラダラと待っている」**という非効率な状態でした。

---

🚀 新提案：「Deep Optimizer States」の「賢い引越し」

この論文が提案する「Deep Optimizer States」は、**「荷物を細かく分け、作業台と廊下をフル活用する」**という戦略です。

1. 荷物を「小分け」にする（サブグループ化）

巨大な百科事典を「1 冊ごと」ではなく、「1 章ごと」に細かく分けます。

• 従来のやり方：「1 章分」を CPU に全部送って書き換え、終わってから GPU に戻す。
• 新しいやり方：「1 章」を CPU で書き換えながら、同時に「次の 1 章」を GPU で書き換える。

2. 「廊下」と「作業台」を同時に使う（オーバーラップ）

これが最大のポイントです。

• CPU が「A 章」を書き換え中：その間に、GPU は「B 章」の書き換えをスタートします。
• 廊下（PCIe）の活用：CPU が書き換えた「A 章」を GPU に戻す作業と、GPU が書き換えた「B 章」を CPU に送る作業を、同時に行います。

まるで、**「狭い廊下を、行きと帰りのトラックがすれ違いながら、絶えず荷物を運んでいる」**ようなイメージです。これにより、廊下が空転する時間が激減します。

3. 作業場所を「動的」に決める

「どの章をどこで書き換えるか」は、その時の廊下の混雑具合や、作業台の空き状況を見て、AI が瞬時に判断します。

• 「廊下が空いているなら、GPU で処理しよう」
• 「CPU の方が速いなら、CPU に任そう」
  このように、**「その瞬間の状況に合わせて、最適な場所へ荷物を振り分ける」**ことで、全体のスピードを最大化します。

---

🍳 料理の例えで理解しよう

【従来のやり方】

• 料理人（GPU）：材料を切る（フォワード）→ 鍋で炒める（バックワード）→ 待機。
• 助手（CPU）：待機 → 調味料を混ぜる（アップデート）→ 料理人に渡す。
• 結果：料理人が待っている間、助手も待っている。廊下（調味料の受け渡し）も空いている。

【Deep Optimizer States のやり方】

• 料理人（GPU）：材料を切る → 鍋で炒める → 次の鍋で炒め始める（同時に）。
• 助手（CPU）： 前の鍋の調味料を混ぜる（同時に）→ 料理人に渡す。
• 廊下：調味料を渡すトラックと、次の材料を運ぶトラックがすれ違いながら走り続ける。
• 結果：誰も待たず、廊下も常に使われている。料理が完成するまでの時間が劇的に短縮されます。

---

🌟 この技術のすごいところ

1. 2.5 倍のスピードアップ
   実験の結果、この新しい方法を導入することで、AI の学習速度が従来の最高水準の技術よりも2.5 倍速くなりました。
2. メモリが少なくても大丈夫
   巨大な GPU がなくても、CPU のメモリを上手に使うことで、大規模な AI の学習が可能になります。
3. 無駄な待ち時間がゼロに
   「運んでいる間、作業していない」というムダを徹底的に排除しました。

🎯 まとめ

この論文は、**「限られたリソース（メモリ）の中で、どうすれば AI の学習を最速で終わらせるか」という課題に対し、「荷物を細かく分け、作業台と廊下を同時にフル回転させる」**という賢い工夫を提案しました。

これにより、将来的には、より安価なハードウェアでも、巨大で高性能な AI を手軽に作れるようになるかもしれません。まるで、**「狭いキッチンでも、職人の動きを工夫することで、高級レストラン並みの料理を次々と作り出す」**ようなものです。

技術概要

Deep Optimizer States: 交差型オフロードを用いたトランスフォーマーモデルの拡張可能なトレーニングに向けた技術的概要

本論文は、大規模言語モデル（LLM）やトランスフォーマーモデルのトレーニングにおいて発生する「メモリ壁（Memory Wall）」の問題を解決し、トレーニング効率を大幅に向上させる新しい手法「Deep Optimizer States」を提案するものです。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

1. 背景と問題定義

問題の核心

近年、LLM のパラメータ数は数百億から数兆規模へと急増しています。しかし、GPU メモリ容量はこれに追いついておらず、モデルパラメータ、勾配、アクティベーション、そして特に**オプティマイザ状態（Optimizer State: 重み、モーメント、分散など）**をすべて GPU メモリに保持することが困難です。

既存手法の限界

現在の最先端手法（DeepSpeed Offload, ZeRO-Offload など）では、オプティマイザ状態をホストメモリ（CPU メモリ）にオフロードし、CPU-GPU 間のハイブリッド計算を行うことでメモリ不足を回避しています。しかし、以下のボトルネックが存在します。

1. PCIe 帯域幅の未活用とブロッキング: ホストと GPU 間のデータ転送（PCIe 経由）がボトルネックとなります。特に、CPU で計算された更新パラメータを GPU に転送する際、転送と計算が直列に実行され、PCIe リンクや GPU がアイドル状態になる時間が長くなります。
2. CPU の計算速度の遅さ: CPU の計算能力は GPU に比べて桁違いに遅く（本論文の実験環境では約 20 倍）、オプティマイザの更新ステップがトレーニング全体のボトルネックとなります。
3. 静的なオフロードの非効率性: 既存のハイブリッド手法（例：DeepSpeed TwinFlow）では、オプティマイザの一部を GPU に「静的」に配置し、残りを CPU に配置します。しかし、フォワード/バックワードパス中は GPU メモリがアクティベーションで占有され、オプティマイザ更新フェーズでは GPU メモリが空くという利用状況の揺らぎを有効活用できていません。

2. 提案手法：Deep Optimizer States

著者らは、トレーニングの各フェーズ（フォワード、バックワード、更新）における GPU メモリ使用量と PCIe 帯域幅の利用状況に大きな変動があることに着目し、これを動的に利用する**「交差型オフロード（Interleaved Offloading）」**を提案しました。

主要な設計原則

1. GPU と CPU による更新の交差実行（Interleaved Updates）:

   • オプティマイザ状態を小さな「サブグループ」に分割します。
   • 各イテレーションにおいて、一部のサブグループの更新を GPU で実行し、残りを CPU で実行します。
   • これにより、CPU が計算している間に GPU が別のサブグループの更新を行い、転送と計算を重畳（オーバーラップ）させます。

2. オーバーラップ中心のデータ移動:

   • 非同期転送（Asynchronous Transfers）を活用し、CPU がサブグループ $S_1, S_2$ の更新を計算している間に、GPU 用の次のサブグループ $S_3$ の状態をホストから非同期でプリフェッチします。
   • GPU での更新と、CPU での更新パラメータの FP32→FP16 変換・転送を並列に行います。
   • これにより、PCIe リンクをフルデュプレックス（H2D と D2H の同時転送）で効率的に利用します。

3. 効率的な勾配管理:

   • バックワードパスで生成された勾配（通常は FP16）を、GPU メモリが解放されたタイミングで GPU 上で FP32 に変換し、ホストに転送します。これにより、ホスト側での動的な精度変換に伴うメモリ割り当てオーバーヘッドを回避し、転送速度を向上させます。

4. パフォーマンスモデルに基づく動的スケジューリング:

   • GPU と CPU の更新速度、PCIe 帯域幅、変換コストを考慮した数理モデル（式 1）を提案し、どの程度の割合（ストライド $k$）のサブグループを GPU で更新するかを動的に決定します。これにより、計算と転送のバランスを最適化します。

3. 主要な貢献

1. 詳細な分析: オプティマイザオフロード時のトレーニング挙動を分析し、更新フェーズにおける GPU メモリ利用の急激な低下と PCIe リンクの未利用という重要な観察結果を提示しました。
2. 新しい設計原則の導入: 静的なオフロードではなく、サブグループ単位での動的な CPU/GPU 間での更新スケジュールと、非同期データ移動によるオーバーラップを実現しました。
3. パフォーマンスモデルとアルゴリズム: GPU でのオフロード頻度を最適化するパフォーマンスモデルと、それを基にした交差型オフロードアルゴリズムを提案しました。
4. 実装と統合: 広く利用されているトレーニングランタイム（DeepSpeed, Megatron-LM）に統合されたミドルウェア「Deep Optimizer States」を開発・実装しました。
5. 広範な評価: 最大 200 億パラメータのモデルを用いた実験により、既存手法に対する大幅な高速化を実証しました。

4. 実験結果

実験環境は、4 基の NVIDIA H100 GPU (80GB) と 192 コアの CPU を搭載したノードを使用しました。

• イテレーション時間の短縮:

  • オプティマイザ状態を完全に CPU にオフロードするシナリオにおいて、Deep Optimizer States は DeepSpeed ZeRO-3（CPU 更新のみ）と比較して最大 2.5 倍のイテレーション速度向上を実現しました。
  • 200 億パラメータモデルにおいて、バックワードパス中の非同期転送による 1.9 倍の加速と、更新フェーズの 60% 加速が寄与しました。

• 更新スループットの向上:

  • 更新スループット（1 秒あたりに更新されるパラメータ数）は、ZeRO-3 より平均で70% 向上しました。これは、GPU ベースの更新を 50% 程度動的に実行できるためです。

• メモリ効率性:

  • 既存の TwinFlow 手法（GPU に 50% のオプティマイザを静的配置）と比較して、Deep Optimizer States はGPU メモリ使用量を約 45% 削減しつつ、同等かそれ以上の速度（40% 高速）を達成しました。

• スケーラビリティ:

  • マイクロバッチサイズの増加や、CPU コア数の変化、データ並列度の拡大に対しても、提案手法は高いスケーラビリティを示しました。特に CPU 数が少ない環境では最大 3 倍の高速化が確認されました。

5. 意義と将来展望

• リソース制約のある環境でのトレーニング実現: 限られた GPU メモリを持つ単一ノードや小規模クラスターでも、大規模モデルの微調整（Fine-tuning）を効率的に行えるようになります。
• ハードウェア進化への対応: 次世代システム（例：NVIDIA Grace Hopper）では CPU と GPU 間の帯域幅がさらに向上しますが、計算と転送のオーバーラップを最適化する本手法の考え方は、より一層重要になります。
• 汎用性: 本手法は DeepSpeed だけでなく、他のハイブリッド並列化構成やドメイン固有モデルにも適用可能です。

結論として、Deep Optimizer States は、メモリ制約下での LLM トレーニングにおけるボトルネックである「オプティマイザ更新」を、動的なリソース割り当てと非同期転送によって解決し、トレーニング効率を劇的に改善する画期的なアプローチです。