今回の論文
今回取り上げるのは、Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase、Yuxiong He による論文「ZeRO: Memory Optimizations Toward Training Trillion Parameter Models」です。2019 年 10 月に arXiv で公開された論文で、研究分野は大規模モデル学習、分散学習、学習インフラ最適化です。URL は https://doi.org/10.48550/arXiv.1910.02054 です。
この論文を選んだ理由は、巨大モデル学習を難しくしていた本当のボトルネックをかなり分解して見せているからです。単に「GPU が足りない」で終わらせず、何がメモリを食っているのか、どこに無駄な複製があるのか、何を分散すれば実用になるのかを整理しています。今でも LoRA や事前学習、社内向けの独自モデル学習基盤を考えるときの土台になる発想です。
どんな技術か
ZeRO は、巨大なニューラルネットワークを学習するときに、各 GPU が同じ学習状態を丸ごと複製して持ってしまう無駄を減らす技術です。
通常のデータ並列学習では、各 GPU が同じモデル重み、同じ勾配、同じ optimizer state を持ちます。実装は簡単ですが、モデルが大きくなるほど「計算より先にメモリが尽きる」問題が出ます。特に Adam のような optimizer では、重み本体より optimizer state のほうが重くなりやすいです。
ZeRO は、この重複をなくします。optimizer state だけを分散する段階、さらに gradient も分散する段階、最後は parameter 自体も分散する段階へと進み、できるだけ各 GPU が必要な分だけ持つようにします。
ひとことで言えば、ZeRO は「モデルを縦に切る model parallelism ではなく、学習状態の複製を減らして巨大モデルを回す技術」です。大規模学習の民主化に効いたのは、まさにこの視点でした。
課題
ZeRO が解決しようとしているのは、巨大モデル学習で GPU メモリが足りなくなる問題です。ただし、単純な容量不足ではなく、「並列化の仕組みのせいで不要な複製が発生している」ことが本質です。
何が難しいのかというと、学習時には重みだけでなく、勾配や optimizer state も持つ必要があるからです。たとえば Adam では、各パラメータについて 1 次モーメントと 2 次モーメントを保持します。mixed precision では fp16 の重みだけでなく、更新用の fp32 コピーも必要になり、メモリ負担はさらに増えます。
既存の方法にも限界がありました。データ並列は計算効率や実装容易性に優れますが、各 GPU が同じ学習状態を複製するため、モデルサイズが大きくなるとすぐに限界に当たります。逆に model parallelism や pipeline parallelism はメモリを節約できますが、通信や実装の複雑さが増え、モデル改修や専用実装が必要になりやすいです。
なぜこの課題を解く必要があるのかというと、モデルサイズの拡大は精度向上に直結しやすい一方で、訓練インフラのハードルが急激に上がるからです。研究でも事業でも、巨大モデルを試せるかどうかが、性能だけでなく開発速度に影響します。もし 10B 級以上の学習が一部の特殊な分散実装でしか回らないなら、改善仮説の検証速度が落ちます。
実際の AI システムでは、独自事前学習、継続学習、大規模 SFT、Embedding モデルの大規模再学習など、学習そのものが競争力になる場面があります。そこで「どれだけ少ないインフラで、どれだけ大きいモデルを安全に回せるか」は、かなり現実的な問題です。
用語解説
- データ並列（Data Parallelism）
- 同じモデルを複数 GPU に複製し、それぞれ別のミニバッチを処理して勾配を集約する学習方法です。ZeRO はこの枠組みの使いやすさを保ちつつ、複製されていた学習状態だけを削る方向で設計されています。
- optimizer state
- Adam などの optimizer が更新計算のために持つ内部状態です。巨大モデルでは重み本体以上にメモリを消費しやすく、ZeRO が最初に狙う主要な削減対象です。
- mixed precision training
- 計算や保存の一部を fp16 で行い、速度とメモリ効率を上げる学習方法です。ただし更新の安定性のために fp32 コピーも持つので、巨大モデルでは結局かなりのメモリを使います。ZeRO の必要性を理解する上で重要です。
- model parallelism
- 1 つのモデルを複数 GPU に分割して持たせる方法です。大きなモデルを載せられますが、層間通信や実装制約が増えます。ZeRO は、これに頼りすぎずに大きなモデルを回す代替軸として位置づけられます。
- activation checkpointing
- 中間活性を全部保持せず、一部を再計算してメモリを節約する方法です。ZeRO 本体は主に model state の重複削減を扱いますが、論文では activation 側の最適化とも組み合わせています。
技術の仕組み
ZeRO の重要な点は、「GPU メモリを食っているもの」を model state と residual state に分けて考えたことです。model state は optimizer state、gradient、parameter です。residual state は activation、temporary buffer、memory fragmentation などです。ZeRO はまず model state の重複を消し、その後に残るメモリ問題も別途最適化します。
基本アイデア
普通のデータ並列では、各 GPU が同じ parameter、gradient、optimizer state を全部持っています。これは計算の粒度が大きく実装も簡単ですが、メモリの観点ではかなり非効率です。
ZeRO はここで、「全部を各 GPU に持たせる必要はない」と考えます。更新に必要な状態を GPU 間で分散保持し、必要なタイミングだけ通信して使えば、データ並列の計算しやすさを残したままメモリを大幅に削れます。
まず何がメモリを圧迫しているのか
論文では、mixed precision の Adam で学習すると、パラメータ数を Ψ としたとき、fp16 の parameter と gradient でそれぞれ 2Ψ バイト、fp32 の parameter copy と momentum と variance でそれぞれ 4Ψ バイト必要になると整理しています。合計すると 16Ψ バイトです。
ここで重要なのは、巨大モデルでは optimizer state がかなり重いことです。単に「重みが大きい」のではなく、「更新のための付随データまで含めると一気に膨らむ」ので、まず複製を疑うべきだというのが ZeRO の出発点です。
モデル構造ではなく学習状態を分散する
ZeRO は新しい Transformer ブロックを提案する論文ではありません。既存のモデル構造を変えるのではなく、学習中に必要な状態の配置を変えます。この点が実務ではかなり大きいです。モデルアーキテクチャの改変を最小限にしつつ、より大きいモデルを扱えるからです。
ZeRO-DP の 3 段階
論文の中心は ZeRO-DP です。これはデータ並列のまま、各種 model state の重複を段階的に減らします。
1. Optimizer state を分散する
最初の段階では、optimizer state だけを GPU 間で分割保持します。これにより、論文では 4 倍のメモリ削減が得られ、通信量は通常のデータ並列と同程度だと説明されています。
この設計が効くのは、Adam のような optimizer が大きな内部状態を持つからです。重みは複製されていても、最初に optimizer state を削るだけでかなり効きます。
2. Gradient も分散する
次の段階では gradient も分割保持します。これにより論文では 8 倍のメモリ削減が得られ、ここでも通信量は通常のデータ並列と同程度に抑えられるとされています。
勾配は逆伝播後に全 GPU が同じものを持つ必要がありそうに見えますが、実際には更新に必要な形で分散管理できるため、ここにも重複削減余地があります。
3. Parameter も分散する
最後の段階では parameter 自体も分散します。論文では、データ並列数を Nd とすると、メモリ削減が Nd に比例して伸びるとしています。たとえば 64 GPU なら 64 倍の削減ポテンシャルがあり、通信量は 50% ほど増えると述べています。
この段階まで行くと、各 GPU はモデル全体を常時保持せず、必要な部分を動的に集めて計算し、終われば手放す方向に近づきます。ここが ZeRO のスケール性能の核です。
ZeRO-R で residual state も削る
論文では model state の削減だけでは足りないと見ています。重複を減らした結果、今度は activation や temporary buffer、memory fragmentation が次のボトルネックになるからです。
そこで ZeRO-R では、activation partitioning、必要に応じた CPU offload、temporary buffer の調整、fragmentation を避けるメモリ管理を行います。つまり ZeRO は単なる optimizer trick ではなく、「大規模学習のメモリ問題を層ごとに分解して潰す」設計です。
処理の流れ
実行時のイメージは、普段は各 GPU が自分の shard を持ち、forward や backward の必要なタイミングで必要な state を通信して集め、処理が終わったら再び分散状態へ戻す、という流れです。
この設計の狙いは、model parallelism のように常時細かい分割計算をするのではなく、データ並列の大きな計算粒度を保ちながら、保持状態だけ賢く分散することです。計算しやすさとメモリ効率の折衷ではなく、両方をなるべく取りにいこうとしています。
実験と結果
論文では、ZeRO が理論上メモリを削れるだけでなく、実際に大きなモデルを高スループットで学習できるかを検証しています。
何を検証したのか
主な検証ポイントは 3 つです。1 つ目は、既存手法よりどれだけ大きなモデルを回せるかです。2 つ目は、そのとき学習スループットが落ちないか、むしろ改善するかです。3 つ目は、GPU 台数を増やしたときにスケールするかです。
データセットや評価設定
この論文は、分類精度やベンチマーク正答率を競うタイプではなく、学習システムとしての性能を測っています。したがって中心になる指標は、学習可能なモデルサイズ、GPU あたりの throughput、総計算性能、スケーラビリティです。
つまり「何点上がったか」ではなく、「今まで載らなかった規模のモデルを、どれだけ効率よく学習できたか」を見る評価です。大規模学習基盤の論文としては自然な観点です。
100B 級モデル学習で大きく前進
論文では、ZeRO の実装群により 100B を超えるモデル学習を扱えたと報告しています。Abstract では、400 GPU 上で 100B 超のモデルを訓練し、15 PFLOPS の throughput を達成したと述べています。
また同じく Abstract では、当時の最先端手法に対して、扱えるモデルサイズが 8 倍、到達できる性能が 10 倍になったと主張しています。ここからも、単なるメモリ節約ではなく、学習規模そのものを一段押し上げた技術だったことがわかります。
学習速度も改善した
単に「大きいモデルが載った」だけでなく、速度面でも改善が出ています。論文では、400 台の NVIDIA V100 GPU で 100B 級モデルを回し、総計で 15 PFLOPS を達成したと報告しています。これは単純平均でも GPU あたり約 37.5 TFLOPS に相当し、同モデルサイズで従来手法比 10 倍の到達性能改善という主張につながっています。
これは、メモリ効率が良くなることで 1 GPU あたりにより大きいバッチを載せられ、結果としてハードウェアを遊ばせにくくなるからです。メモリ最適化がそのまま速度改善に効く、というのが面白い点です。
スケーラビリティでも有利
論文では 64 GPU から 400 GPU へ増やす範囲で super-linear speedup を観測したと報告しています。普通は GPU を増やすと通信オーバーヘッドで伸びが鈍りやすいですが、ZeRO では GPU 数が増えるほど 1 台あたりの保持メモリが減り、より大きいバッチを積めるため、性能が想定以上に伸びたと解釈できます。
この結果から言えるのは、ZeRO は単なる延命策ではなく、クラスタ全体の aggregate memory を活かして学習設計そのものを変える技術だということです。
13B を model parallelism なしで扱える意義
論文では、ZeRO により model parallelism や pipeline parallelism なしでも 13B パラメータ級モデルを扱えたとしています。これは当時の最大級モデルだった T5 11B より大きい規模です。
この点は実務上かなり重要です。巨大モデル学習の難しさは、GPU 台数だけでなく分散実装の複雑さにもあります。ZeRO の価値は、専用のモデル分割実装を細かく書かなくても、より大きい学習実験に入れることでした。
何に使える？
ZeRO は基盤技術なので派手なアプリ機能そのものではありませんが、学習コストと実装難易度を下げることで、多くの AI 開発に効きます。
独自 LLM や Embedding モデルの継続学習
独自データで継続学習したいとき、最初に当たるのはメモリ制約です。ZeRO の考え方を使うと、全部の GPU に同じ状態を載せる前提を崩せるため、より大きいバッチやより大きいモデルで学習しやすくなります。
特に、自社ナレッジ向け embedding モデルや中規模 LLM の追加学習では、最先端の超巨大基盤モデルほどではなくても、optimizer state の重さが問題になります。ZeRO はそうした現実的な規模でも効きます。
LoRA 以前の事前学習やフルファインチューニング基盤
LoRA が使えないケースや、全重み更新が必要な学習では、parameter 本体だけでなく optimizer state の重みが効いてきます。ZeRO はその負担を減らせるため、フルファインチューニング基盤を作るときの中核候補になります。
小規模なチームでも、GPU を増やす前に「複製している state を減らせないか」を見る発想はそのまま使えます。
マルチモーダルや音声モデルの大規模訓練
画像、動画、音声系のモデルでも、モデル本体が大きくなると同じ問題が起きます。Transformer ベースの encoder-decoder や diffusion 系学習でも、重みだけでなく更新状態のメモリが効いてきます。
ZeRO の設計思想はモダリティ依存ではないので、巨大な vision-language モデルや ASR モデルの学習基盤にも応用しやすいです。
学習クラスタの資源効率改善
学習基盤運用の観点では、「1 ジョブが使える最大モデルサイズ」を伸ばせるだけでなく、「同じクラスタでより重いジョブを通せる」ことにも意味があります。メモリが詰まっていて GPU が計算しきれていない状態を減らせるからです。
これは SaaS 企業の社内 MLOps 基盤でも重要です。モデル改善が遅い原因が GPU 調達ではなく、メモリ効率の悪い学習設定であるケースは珍しくありません。
開発や事業へのヒント
この論文から得られる大きなヒントは、AI プロダクトの競争力はモデル選定だけでなく、学習基盤の設計でも決まるということです。
まず「何を複製しているか」を可視化する
自分で AI アプリや学習基盤を作るなら、OOM が出たときに単純に GPU を増やす前に、parameter、gradient、optimizer state のどれが効いているかを見分けるべきです。ZeRO は、この分解をすると対策が具体化することを示しています。
これは推論基盤にも通じる発想です。KV cache、activation、weight shard など、どこに重複があるかを見れば、単なるハード増強以外の解き方が見えてきます。
実装しやすい並列方式を優先する
論文の実務的な価値は、データ並列の使いやすさをなるべく残している点です。事業としては、最高性能の理論値より、チームが継続運用できる実装のほうが重要な場面が多いです。
つまり「最も賢い並列化」ではなく、「最も回し続けやすい並列化」を選ぶ視点が必要です。ZeRO はそこにかなり明確な回答を出しています。
学習インフラはプロダクト速度に直結する
独自モデルを持つ事業では、1 回の学習ジョブを成立させるまでの準備コストが高いと、改善サイクル全体が遅くなります。ZeRO のような基盤技術で model parallelism 依存を減らせると、実験の立ち上げ速度や運用の安定性が上がります。
これは研究組織だけでなく、RAG の埋め込み再学習や社内モデル更新を継続したい SaaS にも効く考え方です。
今後注目すべき方向性
今後注目すべきなのは、単独のメモリ最適化ではなく、分散学習、低精度化、checkpointing、offloading、通信最適化を組み合わせる流れです。ZeRO はその基盤部分を切り開いた論文として理解すると実務に活かしやすいです。
また、フル学習だけでなく大規模ファインチューニングやマルチモーダル学習でも、state の持ち方を最適化する方向は今後も重要だと考えられます。これは論文内容を踏まえた将来方向の考察です。
限界
ZeRO にも限界はあります。まず、parameter まで分散する段階では通信設計が重要で、メモリが減るほど実装と実行の複雑さは増えます。論文でも最終段階では通信量が増える点を明示しています。
また、メモリ削減だけで学習が速くなるとは限りません。通信帯域、ノード間接続、バッチサイズ上限、optimizer、データ読み込みなど、別のボトルネックが強ければ効果は限定されます。
データ依存性や再現性の観点では、ZeRO 自体はモデル精度を直接上げる技術ではありません。学習可能な規模を広げる技術なので、精度向上は最終的にモデル設計やデータ品質に依存します。
実装面でも、単一 GPU や小規模環境では恩恵が薄い場合があります。クラスタ規模、ネットワーク構成、学習タスクによっては、運用複雑性のわりに差が出にくいこともあります。
さらに、論文はシステム論文なので、特定タスクでの最終精度を広く比較しているわけではありません。したがって、現代の実運用に入れる際は、対象モデル、クラスタ、学習目的に応じた再評価が必要です。
よくある質問
Q. ZeRO は model parallelism と何が違うのですか？
A. model parallelism はモデルの計算そのものを複数 GPU に分割しますが、ZeRO は主に学習状態の重複を減らすことでメモリを節約します。計算粒度をなるべく大きく保てるので、実装しやすさと効率のバランスが取りやすいのが違いです。
Q. ZeRO は推論にもそのまま使えますか？
A. この論文の中心は学習時のメモリ最適化です。ただし「重複する状態を減らす」という考え方自体は、推論時の weight sharding や KV cache 管理にも通じます。直接の適用というより、設計思想が近いと考えるのが正確です。
Q. どんなチームでも導入価値がありますか？
A. 単一 GPU や小さなモデル中心なら効果は限定的です。一方で、複数 GPU を使い、optimizer state が支配的になっている学習では導入価値が高いです。特にフルファインチューニングや独自事前学習の基盤では有力です。
Q. ZeRO を入れると精度も上がりますか？
A. ZeRO 自体が精度を直接押し上げるわけではありません。価値は、今までメモリ上載らなかったモデルサイズやバッチ設定を試せるようになることです。その結果として、より良い学習条件に到達できる可能性があります。
Q. 実務で最初に見るべきなのはどこですか？
A. まずは optimizer state がどれだけメモリを使っているかです。大規模学習では重み本体よりこちらが重いことがあります。そこを見ずに GPU 追加だけで解決しようとすると、費用対効果が悪くなりやすいです。
今日の学び
この論文は、巨大モデル学習で GPU メモリが不足する課題を扱いました。特に、parameter そのものだけでなく、gradient や optimizer state の重複が大きな原因だと整理しました。これに対して ZeRO は、学習状態を段階的に分散保持し、必要なときだけ通信して使う設計で解こうとしました。
ここから得られるヒントは、AI の大規模化では「何を持つか」だけでなく「どこにどう持つか」が重要だということです。プロダクト開発でも学習基盤でも、単純な GPU 増設より先に、重複している状態を減らす設計を考える価値があります。