ZeRO(ゼロ)とは

ZeROとは、複数のGPUでAI訓練するとき、各GPUが同じものを丸ごと抱え込む「重複」をなくし、モデルの状態を分担して持たせることでメモリを節約する手法です。ふつうのやり方では、GPUが何枚あっても全部がそっくり同じデータを抱えます。ZeROはその無駄に着目し、保管する中身を手分けして持たせることで、同じハードでもより大きなモデルを訓練できるようにします。

英語表記:Zero Redundancy Optimizer(ZeRO)

ZeROの仕組み

AIの訓練では、モデルのパラメータに加えて、勾配やオプティマイザの状態といった「学習のための作業データ」もメモリを大きく占める要素です。複数GPUで手分けする一般的なやり方では、これらを全GPUがまるごと同じだけ複製するため、重複が膨らみます。ZeROが行うのは、これらを段階的に分けていくことです。まずオプティマイザの状態を分け(ステージ1)、次に勾配も分け(ステージ2)、さらにパラメータまで分けます(ステージ3)。各GPUは自分の担当ぶんだけを保管し、計算に必要なときだけ他から集める。こうして重複を減らしつつ、通信は最小限に抑えるのが狙いです。

テンソル並列やFSDPとの違い

テンソル並列パイプライン並列は、モデルの「計算」そのものを複数GPUに割る考え方です。これに対してZeROは、計算の分け方ではなく、各GPUが抱える「保管データ」の重複をなくす点に特徴があります。データ並列の手軽さを保ったまま、メモリだけを節約するイメージです。なお、よく似た発想をPyTorchに取り込んだ仕組みがFSDPで、両者は同じ「状態を分担して持つ」という考え方を土台にしています。

経営から見たZeROの意味

ZeROは、同じGPU枚数でより大きなモデルを訓練できるようにする、メモリ節約の定番技術です。Microsoftが2019年に発表したDeepSpeedという開発基盤の中心機能で、当時1750億規模に迫る大規模モデルの訓練を後押ししました。AI開発で最も高くつく資源のひとつがGPUのメモリです。その壁を下げるZeROは、巨大モデルを「現実的な台数」で動かすための重要なピースになっています。

Topic名前そのものが設計思想を表している

「ZeRO」はZero Redundancy Optimizerの略で、直訳すると「冗長ゼロの最適化」です。複数GPUで訓練する従来のやり方は、GPUが10枚あれば10個のまったく同じコピーを抱えるという、よく考えると大きな無駄を含んでいました。ZeROは「なぜ同じものを何個も持つのか」という素朴な疑問に正面から答え、重複を限りなくゼロに近づけます。技術の名前が、そのまま「無駄をなくす」という発想を言い表している、すっきりとした命名です。

ZeROに関するよくある質問

ZeROとテンソル並列は何が違いますか?
テンソル並列はモデルの計算そのものを複数GPUに割ります。ZeROは計算の分け方ではなく、各GPUが抱える保管データ(パラメータや勾配など)の重複をなくす点が違いで、データ並列の手軽さを保ったままメモリを節約します。
ZeROの3つのステージとは何ですか?
節約する対象を段階的に増やす設定です。ステージ1はオプティマイザの状態、ステージ2はそれに加えて勾配、ステージ3はさらにパラメータまでを各GPUに分担させます。番号が大きいほどメモリを節約できます。
ZeROとFSDPの関係は?
FSDPはPyTorchが、ZeROと同じ「状態を分担して持つ」発想(特にステージ3)を取り込んだ仕組みです。ZeROはMicrosoftのDeepSpeed由来で、FSDPはPyTorchに標準搭載されている、という違いがあります。

ZeROに関連する記事