【マルチGPU】DeepSpeed ZeRO分散学習セットアップガイド（実践編）

はじめに

大規模言語モデル（LLM）のファインチューニングや大規模なモデルの学習を行う際、単一のGPUではメモリ容量が不足したり、学習に膨大な時間がかかったりする問題に直面します。この課題を解決するための強力な技術が、分散学習です。特に、Microsoftが開発したDeepSpeedライブラリとその中核技術であるZeRO（Zero Redundancy Optimizer）は、複数のGPUに効率的に計算とメモリ負荷を分散させ、従来では扱えなかった巨大なモデルを訓練可能にしました。

本記事は実践編として、実際にDeepSpeed ZeROを利用してマルチGPU環境で分散学習をセットアップする手順を、具体的な設定ファイルやコード例とともに詳細に解説します。ZeROの各ステージ（Stage 1/2/3）の違いを理解し、ご自身のハードウェアリソースと学習タスクに最適な構成を選択・実装できることを目指します。

前提条件・必要な環境

以下の環境を前提として進めます。異なる環境の場合、パスやコマンドが若干変更になる可能性があります。

• ハードウェア: 複数台のNVIDIA GPUを搭載したサーバー（例: 2x / 4x / 8x NVIDIA A100/V100）。
• OS: Ubuntu 20.04 LTS または 22.04 LTS。
• ドライバーとCUDA: NVIDIA Driver (>=470), CUDA Toolkit 11.7 または 11.8。
• Python: 3.8 以上。
• 必須ライブラリ:
  • PyTorch (CUDA対応版, 例: torch==2.0.1+cu117)
  • Transformers (Hugging Face)
  • DeepSpeed
  • その他、学習スクリプトに依存するライブラリ（datasets, accelerate等）
• ネットワーク: マルチノードで実行する場合は、高速なインターコネクト（InfiniBand）が理想的ですが、単一ノード内のマルチGPUでも効果を発揮します。

手順1: 環境構築とDeepSpeedのインストール

まず、PyTorchとDeepSpeedをインストールします。CUDAバージョンに合ったPyTorchを先にインストールすることが重要です。

# 例: CUDA 11.7 用のPyTorchとDeepSpeedをインストール
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
pip install deepspeed
# オプションで、ビルド済みのwheelが存在しない場合は、ソースからビルドします（時間がかかります）
# DS_BUILD_CPU_ADAM=1 DS_BUILD_FUSED_ADAM=1 DS_BUILD_UTILS=1 pip install deepspeed

インストール後、以下のコマンドでDeepSpeedが正しく認識されるか確認します。

python -c "import deepspeed; print(f'DeepSpeed version: {deepspeed.__version__}')"

手順2: ZeROの基本概念とステージの選択

ZeROは、オプティマイザの状態（Optimizer States）、勾配（Gradients）、モデルパラメータ（Parameters）という3つの主要メモリ要素をGPU間で分割・共有することで、冗長性をゼロ（Zero）に近づけます。使用するリソースとトレードオフに応じて、3つのステージから選択します。

• ZeRO Stage 1: オプティマイザの状態のみを分割。実装が簡単で通信オーバーヘッドが小さい。大きなバッチサイズが取れる場合のメモリ削減に有効。
• ZeRO Stage 2: オプティマイザの状態と勾配を分割。Stage 1よりメモリ効率が向上。最もバランスが取れており、一般的に推奨されるステージ。
• ZeRO Stage 3: オプティマイザの状態、勾配、さらにモデルパラメータ自体も分割。最大のメモリ削減効果があり、数十億パラメータ規模のモデルを少ないGPU数で扱える。ただし、通信オーバーヘッドが大きく、設定が複雑。

初めての方はZeRO Stage 2から始めることをお勧めします。

手順3: DeepSpeed設定ファイル（ds_config.json）の作成

DeepSpeedの動作は、JSON形式の設定ファイルで詳細に制御します。プロジェクトのルートディレクトリにds_config.jsonを作成します。

ZeRO Stage 2 の基本設定例:

{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "bf16": {
    "enabled": false
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 5e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 5e-5,
      "warmup_num_steps": 1000
    }
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": 4,
  "gradient_clipping": 1.0,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": 4,
  "steps_per_print": 50,
  "wall_clock_breakdown": false
}

主な設定項目の説明:

• fp16: 混合精度訓練を有効にし、メモリ使用量と計算速度を改善。
• zero_optimization: ここが核心。"stage": 2を指定。overlap_commやcontiguous_gradientsはパフォーマンス向上のためのオプション。
• train_micro_batch_size_per_gpu: 各GPUが一度に処理するバッチサイズ。GPUメモリに収まる最大サイズに設定。
• gradient_accumulation_steps: 勾配を蓄積するステップ数。実効バッチサイズ = micro_batch * gradient_accumulation_steps * GPU数。

手順4: 学習スクリプトの修正（Hugging Face Transformers連携）

Hugging FaceのTrainerクラスは、deepspeed引数を渡すだけで簡単にDeepSpeedを統合できます。以下は、カスタムトレーニングループを使用する場合の例です。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from datasets import load_dataset
import deepspeed

# 1. モデルとトークナイザーの読み込み
model_name = "microsoft/phi-2" # 例として小規模モデルを使用
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

# トークナイザーのパディングトークンを設定（ない場合）
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 2. データセットの準備（簡易例）
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)

dataset = load_dataset("your_dataset_name")
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# 3. DeepSpeed設定ファイルのパス
deepspeed_config = "./ds_config.json"

# 4. トレーニング引数の設定（DeepSpeed用）
training_args = TrainingArguments(
    output_dir="./output",
    num_train_epochs=3,
    per_device_train_batch_size=4,  # ds_config.jsonの`train_micro_batch_size_per_gpu`と一致させる
    gradient_accumulation_steps=4,   # ds_config.jsonの`gradient_accumulation_steps`と一致させる
    learning_rate=5e-5,
    fp16=True, # ds_config.jsonの`fp16:{"enabled":true}`と一致
    logging_dir="./logs",
    logging_steps=50,
    save_steps=500,
    deepspeed=deepspeed_config, # ここに設定ファイルパスを指定
)

# 5. Trainerの初期化と学習実行
from transformers import Trainer, DataCollatorForLanguageModeling

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)

trainer.train()

手順5: 学習の実行とログ確認

スクリプトをdeepspeedランチャーを用いて実行します。これにより、自動的に複数GPUにプロセスが分散されます。

# 単一ノード、全GPUを使用する場合
deepspeed --num_gpus=4 train_script.py

# 特定のGPUのみを使用する場合
CUDA_VISIBLE_DEVICES=0,1 deepspeed --num_gpus=2 train_script.py

# マルチノードの場合（例）
deepspeed --hostfile=hostfile --master_addr=主ノードIP --master_port=29500 train_script.py

実行すると、DeepSpeedが初期化され、ZeROのステージや使用メモリなどの情報がログに出力されます。ログ内の “OVERFLOW!” というメッセージは勾配のアンダーフロー/オーバーフローを意味し、fp16設定のloss_scaleに関連します。

トラブルシューティング

• エラー: “CUDA out of memory”
  • train_micro_batch_size_per_gpuを小さくする。
  • gradient_accumulation_stepsを増やして実効バッチサイズを維持しつつ、マイクロバッチを小さくする。
  • ZeROのステージを上げる（Stage 2 → Stage 3）。Stage 3では、zero_optimizationに"stage": 3と"offload_param"や"offload_optimizer"を追加してCPUメモリへのオフロードを検討。
• 学習速度が遅い
  • 通信がボトルネックの場合: allgather_bucket_sizeやreduce_bucket_sizeの値を調整（デフォルトの2e8から5e8などに増やす）。
  • NVLink接続されているGPUを使用する。
  • overlap_commがtrueになっているか確認。
• NaNやLossが発散する
  • fp16訓練では初期のLoss Scaleが重要。設定ファイルの"loss_scale": 0（動的スケーリング）は大抵うまくいく。initial_scale_powerを下げてみる（例: 16 → 15）。
  • 学習率が高すぎる可能性がある。少し下げてみる。
  • BF16（Ampereアーキテクチャ以降のGPU）が利用可能であれば、fp16の代わりにbf16: {"enabled": true}を使用すると、数値的に安定することが多い。
• DeepSpeedが設定ファイルを読み込まない
  • 設定ファイルのパスが正しいか確認。
  • JSONファイルの構文エラーがないか確認（カンマの付け忘れなど）。オンラインJSONバリデータで検証可能。

まとめ

本記事では、DeepSpeed ZeROを用いたマルチGPU分散学習の実践的なセットアップ手順を解説しました。要点をまとめます。

1. ZeROの各ステージ（1, 2, 3）は、メモリ削減効果と通信オーバーヘッドのトレードオフであり、まずはStage 2から始めるのが良い。
2. 学習を制御する核心はds_config.json設定ファイルであり、特にzero_optimizationセクションとバッチサイズ関連のパラメータが重要。
3. Hugging Face Transformersとの統合は非常に簡単で、TrainingArgumentsにdeepspeed引数を渡すだけで良い。
4. 実行はdeepspeedランチャーを使用し、メモリ不足や速度の問題には、バッチサイズ、バケットサイズ、オフロード設定などを段階的に調整して対処する。

DeepSpeedを活用することで、限られたGPUリソースでも大規模なLLMのファインチューニングや学習が現実的な時間で実行可能になります。本ガイドが、効率的な分散学習の実装の一助となれば幸いです。