36. DeepSpeed ZeROで巨大モデルを訓練！メモリシャーディングの威力

スポンサーリンク

36. DeepSpeed ZeROで巨大モデルを訓練！メモリシャーディングの威力

こんにちは！今回は、単一GPUの限界を超えて巨大なLLMをトレーニングできる革新的な技術、DeepSpeed ZeROについて学んでいきます。

前回、データ並列化（DDP）でトレーニングを高速化する方法を学びましたよね。でも、DDPには1つの大きな制限がありました：モデル全体が各GPUに収まる必要があるんです。

今回紹介するDeepSpeed ZeROは、この問題を解決します。GPU全体でモデル状態をシャーディング（分割）することで、単一GPUには決して収まらないモデルをトレーニングできるようになるんです！

はじめに：単一GPU制限を超えたスケーリング

DDPの限界

前のレッスンで、DDPが複数のGPU全体で異なるバッチを処理することでトレーニングを高速化する方法を見ました。

DDPはスループット（処理速度）に最適です。しかし、メモリには役立ちません。

なぜでしょうか？DDPはすべてのGPUでモデル全体を複製するためです。

DDPの問題：
– 各デバイスが、すべてのパラメータ、勾配、オプティマイザー状態の完全なコピーを保持
– モデルが1つのGPUに収まらない場合、DDPを実行する4つのGPUにも収まらない

DeepSpeed ZeROの解決策

それがDeepSpeed ZeROが解決する問題なんです！

ZeROのアプローチ：
– すべてを複製する代わりに、GPU全体でモデル状態をシャード（分割）
– 各デバイスは、パラメータ、勾配、オプティマイザー状態の一部のみを保持
– 必要なときに必要なものを再構築するために調整

これはモデル並列化と呼ばれます：データだけでなく、モデル自体を分散するんですね。

結果：
より高価なハードウェアを購入したり、コードを劇的に変更したりすることなく、単一GPUに決して収まらないモデルをトレーニングできます！

DeepSpeedとは何か？

DeepSpeedは、大規模な深層学習モデルをトレーニングするためにMicrosoftが開発したオープンソース最適化ライブラリです。

目的：
分散トレーニングをより効率的にする
– メモリ使用量の削減
– トレーニング速度の向上

これを、トレーニングコードとPyTorchの間に位置し、マルチGPUトレーニングに必要な複雑な調整を処理するレイヤーと考えてください。

DeepSpeedには複数の最適化テクニックが含まれていますが、最も重要なのがZeRO（Zero Redundancy Optimizer）なんです。

ZeROコンセプト：GPU全体でメモリをシャーディング

ZeROの背後にある洞察はシンプルですが強力です：

トレーニング中に保存するもののほとんどは、すべてのGPUで複製する必要はない

メモリの内訳

以前のレッスンで学んだように、トレーニング中、GPUは次を保持します：

  1. モデルパラメータ：重みとバイアス
  2. 勾配：バックプロパゲーション中に計算された導関数
  3. オプティマイザー状態：モーメンタムと分散項（Adamの場合）
  4. 活性化：バックプロップのために保存された中間出力

従来のDDPの冗長性

従来のトレーニングでは、これらすべてがすべてのGPUに完全に存在します。

例：4つのGPUでDDPを実行
– すべてのコンポーネント（パラメータ、勾配、オプティマイザー状態、活性化）の4つのコピー
– これは大規模に冗長

ZeROの解決策

ZeROは、シャーディングを通じてその冗長性を排除します：

  • これらのコンポーネントをデバイス全体に分割
  • 各GPUがそれらの一部のみを保存

GPU全体で戦略的シャーディングを使用することで、同じ総メモリ量ではるかに大きなモデルを収めることができます。いくつかの通信オーバーヘッドを犠牲にして。

問題は：どれだけシャードするか？

これがZeROの3つのステージが登場する場所です！

ZeROの3つのステージ

ZeROには、徐々に攻撃的になる3つのステージがあります。各ステージは、メモリ節約と通信オーバーヘッドの間の異なるトレードオフを提供します。

Stage 1：オプティマイザー状態のシャーディング

最も保守的なステージ

何をシャードするか：
オプティマイザー状態のみ（モーメンタムと分散項）

仕組み：
– GPU全体でオプティマイザー状態を分割
– 各GPUは自分のスライスのみを保持
– パラメータと勾配はすべてのGPUに完全に複製される

メモリ節約：
– Adamオプティマイザーの場合、オプティマイザー状態はパラメータあたり8バイト
– N個のGPUでシャーディングすると、これを約N倍削減

通信オーバーヘッド：
最小限。各GPUはまだすべてのパラメータを保持し、逆伝播を実行できます。

使用するタイミング：
– モデルがほぼ収まるが、オプティマイザー状態がメモリを押しのけている
– 最小限の通信オーバーヘッドが必要
– PCIeのような遅いインターコネクトを使用している

ZeRO-1は、最も侵襲性が低く、ほぼすべてのマルチGPUセットアップで安全に使用できます。

Stage 2：勾配 + オプティマイザー状態のシャーディング

推奨されるスタート地点

何をシャードするか：
勾配とオプティマイザー状態の両方

仕組み：
– 逆伝播中に各GPUが勾配のスライスを計算
– それらをローカルに保存
– 他のGPUと共有しない（all-gatherまで）

メモリ節約：
– FP32勾配 + オプティマイザー状態 = パラメータあたり約12バイト
– N個のGPUにシャーディングすると、デバイスあたり約12/Nバイトに削減

具体例：
10億パラメータモデル、4つのGPU：
– 勾配 + オプティマイザー状態：約12 GB
– ZeRO-2でシャーディング：GPUあたり約3 GB

これは9 GBの節約です！

通信オーバーヘッド：
穏やか。GPUは順伝播と逆伝播中にローカルで作業でき、最後のall-gatherステップで勾配のみを同期します。

使用するタイミング：
– モデルが単一GPUに収まらない
– 速度とメモリの間の良好なバランスが必要
– ほとんどのLLMファインチューニングシナリオ

ZeRO-2は、ほとんどのケースで「ちょうど良い」ステージです。 DDPよりも大幅にメモリを節約し、ZeRO-3よりもはるかに高速に実行されます。

Stage 3：パラメータ + 勾配 + オプティマイザー状態のシャーディング

最も攻撃的なステージ

何をシャードするか：
すべて。パラメータ、勾配、オプティマイザー状態。

仕組み：
– 各GPUは、パラメータの完全なセットのスライスのみを保持
– 順伝播中、GPUはレイヤーごとに必要なパラメータをフェッチ
– 計算後、それらを破棄
– 逆伝播も同様

メモリ節約：
最大。デバイスあたりのメモリ使用量は、GPUの数でほぼ線形にスケールします。

通信オーバーヘッド：
大幅に増加。GPUが所有していないパラメータが必要になるたびに、別のデバイスからフェッチする必要があります。

インターコネクトの重要性：
– NVLink（A100/H100）: ZeRO-3はうまく機能
– PCIe（コンシューマーセットアップ）: ZeRO-2よりも著しく遅くなる可能性

使用するタイミング：
– ZeRO-2でもモデルが収まらない
– 高速インターコネクト（NVLinkまたはInfiniBand）がある
– メモリが絶対的なボトルネックで、速度をトレードする意思がある

ZeRO-3は強力ですが、追加のメモリ節約が絶対に必要でない限り、ZeRO-2から始めてください。

DeepSpeed ZeROの実装

DeepSpeedの美しさは、Hugging FaceのTrainerとどれだけシームレスに統合されるかです。

インストール

pip install deepspeed

Hugging Face Trainerとの統合

トレーニングループを書き直したり、モデルセットアップを変更したりする必要はありません！

Week 3でのトレーニングコード：

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-4,
    fp16=True,
    gradient_accumulation_steps=4,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()

DeepSpeed ZeROを有効にするには、1行追加：

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-4,
    fp16=True,
    gradient_accumulation_steps=4,
    deepspeed="ds_config_zero2.json",  # <-- これだけ！
)

他のすべては同じまま！モデル、データセット、LoRA構成、前処理。DeepSpeedは舞台裏ですべての分散調整を処理します。

DeepSpeed構成ファイル

魔法はJSON構成ファイルで起こります。ここで、使用するZeROステージと通信を最適化する方法をDeepSpeedに伝えます。

典型的なZeRO-2構成（ds_config_zero2.json）：

{
  "train_batch_size": "auto",
  "gradient_accumulation_steps": "auto",
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "contiguous_gradients": true,
    "overlap_comm": true
  }
}

重要なパラメータ：

  • "stage": 2：ZeRO-2を有効化（勾配 + オプティマイザー状態をシャード）
  • "overlap_comm": true：計算と通信をオーバーラップして速度向上
  • "train_batch_size": "auto"：TrainingArgumentsから設定を読み取る

トレーニングの起動

Accelerateと同じように簡単です：

# デフォルトDeepSpeed構成で起動
deepspeed train.py

# または特定の構成ファイルを指定
deepspeed --deepspeed_config ds_config_zero2.json train.py

どのZeROステージを選ぶべきか？

意思決定フローチャート：

1. モデルが単一GPUに収まるか？

  • はい → DDPを使用（データ並列化）
  • いいえ → 次の質問へ

2. ZeRO-2で収まるか？（勾配 + オプティマイザー状態をシャード）

  • はい → ZeRO-2を使用（推奨）
  • いいえ → 次の質問へ

3. 高速インターコネクト（NVLink/InfiniBand）があるか？

  • はい → ZeRO-3を試す
  • いいえ → より小さなモデル、量子化、またはより多くのGPUを検討

経験則：
– ZeRO-1：オプティマイザー状態のみがボトルネックの場合
– ZeRO-2：ほとんどのLLMファインチューニングで最良のバランス
– ZeRO-3：ZeRO-2でも収まらない巨大なモデル

実際のメモリ節約

Llama 3.1 7B モデルの完全ファインチューニング（BF16）：

単一GPU（DDPなし）

  • パラメータ：14 GB（7B × 2バイト）
  • 勾配：14 GB
  • オプティマイザー状態：28 GB（7B × 4バイト × 2状態）
  • 活性化：約4 GB
  • 合計：約60 GB → 40GB GPUに収まらない

4つのGPUでDDP

  • 各GPUが完全なコピーを保持
  • GPUあたり：約60 GB
  • 結果：まだ収まらない

4つのGPUでZeRO-2

  • パラメータ：14 GB（各GPUで複製）
  • 勾配：3.5 GB（14 GB / 4）
  • オプティマイザー状態：7 GB（28 GB / 4）
  • 活性化：約4 GB
  • GPUあたり：約28.5 GB → 40GB GPUに収まる！

4つのGPUでZeRO-3

  • パラメータ：3.5 GB（14 GB / 4）
  • 勾配：3.5 GB（14 GB / 4）
  • オプティマイザー状態：7 GB（28 GB / 4）
  • 活性化：約4 GB
  • GPUあたり：約18 GB → さらに余裕がある

まとめ

今回の記事では、DeepSpeed ZeROを使った効率的なマルチGPUメモリシャーディングについて学びました。

重要なポイント：

  1. ZeROはメモリの冗長性を排除：
    • DDPはすべてを複製
    • ZeROはGPU全体でシャード
  2. 3つのステージ：
    • Stage 1：オプティマイザー状態のみ
    • Stage 2：勾配 + オプティマイザー状態（推奨）
    • Stage 3：すべて（最も攻撃的）
  3. 簡単な統合：
    • Hugging Face Trainerと1行で統合
    • JSON構成ファイルで設定
  4. 実用的なメモリ節約：
    • 70億パラメータモデルが40GB GPUに収まる
    • 単一GPUでは不可能だったトレーニングが可能に
  5. トレードオフ：
    • メモリ節約 vs 通信オーバーヘッド
    • インターコネクトの速度が重要

推奨戦略：
まずZeRO-2から始めて、必要に応じてZeRO-3にエスカレートする

単一GPUの天井を破りましょう！DeepSpeed ZeROで、巨大なモデルも訓練可能になるんですね！

コメント

タイトルとURLをコピーしました