DeepSpeed完全ガイド：ZeRO最適化と大規模モデル訓練

はじめに

GPT-4、LLaMA、Falconのような大規模言語モデルを単一のGPUで訓練することは、単純に不可能です。GPT-3の1750億パラメータをfp16で保存するには約350GBのGPUメモリが必要で、オプティマイザーの状態（Adam）を含めると1.4TBになります。Microsoftはこの問題を解決するためにDeepSpeedを開発しました。

DeepSpeedはZeRO（Zero Redundancy Optimizer）最適化技術を中心に据え、通常のGPUクラスタで数千億〜数兆パラメータのモデルの訓練を可能にします。このガイドでは、DeepSpeedのコアコンセプトから本番設定まですべてを網羅します。

1. DeepSpeedの概要

1.1 背景と動機

ディープラーニングモデルの大規模化に伴い、3つのボトルネックが生じます：

メモリ不足：モデルパラメータ、勾配、オプティマイザーの状態がGPU VRAMを超過
計算ボトルネック：単一GPU FLOPSの制限
通信ボトルネック：マルチGPU設定でのグラジェント同期オーバーヘッド

従来のデータ並列（DP）では各GPUがモデルの完全なコピーを保持する必要があり、メモリ問題を解決できません。モデル並列は実装が複雑で非効率になりがちです。

DeepSpeedは両方のアプローチの長所を組み合わせ、独自のZeRO最適化によって革命的なメモリ効率を実現します。

1.2 PyTorchとの統合

DeepSpeedは既存のPyTorchコードへの変更を最小限に抑えます。主な変更は単一の deepspeed.initialize() 呼び出しで、エンジン作成、分散初期化、オプティマイザーラッピングを同時に処理します。

import deepspeed

# 既存のPyTorchコード
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# DeepSpeedへの切り替え（最小限の変更）
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

1.3 インストール

pip install deepspeed

# CUDAカーネルビルドと共にインストール
DS_BUILD_OPS=1 pip install deepspeed

# インストール確認
ds_report

2. ZeRO最適化：段階別分析

ZeRO（Zero Redundancy Optimizer）はデータ並列訓練のメモリ冗長性を排除するコア技術です。標準的なデータ並列では、N個のGPUが全て同一のモデルコピーを維持します——これは極端なメモリの無駄遣いです。ZeROはこの冗長性を段階的に排除します。

訓練中のメモリ使用量分析

Ψ個のモデルパラメータに対して、混合精度（fp16）訓練でのGPU当たりメモリは：

コンポーネント	サイズ	説明
パラメータ（fp16）	2Ψ バイト	順方向/逆方向パス
勾配（fp16）	2Ψ バイト	逆方向パスの結果
マスターパラメータ（fp32）	4Ψ バイト	オプティマイザー用
Adamモメンタム（fp32）	4Ψ バイト	第1モーメント
Adam分散（fp32）	4Ψ バイト	第2モーメント
合計	16Ψ バイト

7Bパラメータモデルには約112GB、70Bモデルには約1.1TBが必要です。

2.1 ZeROステージ1：オプティマイザー状態の分割

ZeRO-1は最大のメモリコンポーネントであるオプティマイザー状態をN個のGPUに分割します。

Adam状態：fp32パラメータ + モメンタム + 分散 = 12Ψ バイト
N個のGPU使用時：GPU当たりのオプティマイザー状態 = 12Ψ/N バイト

メモリ節約：

ベースライン: GPU当たり16Ψ バイト
ZeRO-1: GPU当たり 4Ψ + 12Ψ/N バイト
(N=64の場合: ≈ 4.2Ψ バイト、~4倍の節約)

更新時、各GPUは割り当てられたパラメータシャードのみを更新し、All-gatherオペレーションで完全なパラメータを復元します。

{
  "zero_optimization": {
    "stage": 1
  }
}

2.2 ZeROステージ2：勾配の分割

ZeRO-2はオプティマイザー状態に加えて勾配も分割します。

逆伝播中、各GPUは自分のパラメータシャードに対応する勾配のみを蓄積します。Reduce-scatterオペレーションで各GPUの割り当てセグメントの勾配を収集します。

ベースライン: GPU当たり16Ψ バイト
ZeRO-2: GPU当たり 2Ψ + (2Ψ + 12Ψ)/N バイト
(N=64の場合: ≈ 2.2Ψ バイト、~8倍の節約)

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

2.3 ZeROステージ3：パラメータの分割

ZeRO-3はモデルのパラメータ自体も分割します。これは最も強力なステージですが、実装が最も複雑です。

動作の仕組み：

各GPUは常に全パラメータの1/Nのみを保持
順方向パス：各レイヤーの実行直前にAll-gatherで完全なパラメータを収集
計算後：パラメータを即座に解放
逆方向パス：同様にオンデマンドでパラメータを収集

ZeRO-3: GPU当たり (2Ψ + 2Ψ + 12Ψ)/N バイト
N=64の場合: 16Ψ/64 ≈ 0.25Ψ バイト (~64倍の節約！)

トレードオフ： All-gather通信が各レイヤーで発生し、通信オーバーヘッドが増加します。ただし、パイプライン実行との重複により実際の遅延は小さくなります。

{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

ZeRO-3 Pythonの初期化：

import deepspeed
from transformers import AutoModelForCausalLM

# ZeRO-3ではモデル作成方法が重要
with deepspeed.zero.Init(config_dict_or_path="ds_config.json"):
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# パラメータは最初から分割され、メモリピークを防止

3. ZeRO-OffloadとZeRO-Infinity

3.1 ZeRO-Offload

ZeRO-Offloadはオプティマイザーの状態と勾配をCPUメモリに移動させ、GPUメモリの負荷を大幅に削減します。

GPU：fp16パラメータ + fp16勾配（順方向/逆方向）
CPU：fp32マスターパラメータ + Adam状態（オプティマイザー更新）

これにより、単一GPUでも数十億パラメータのモデルの訓練が可能になります。

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

パラメータオフロード（ZeRO-3 + Offload）：

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    }
  }
}

3.2 ZeRO-Infinity（NVMeオフロード）

ZeRO-InfinityはNVMe SSDを活用して、実質的に無制限のモデルサイズをサポートします。

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 4,
      "fast_init": false
    },
    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 5,
      "buffer_size": 1e8
    },
    "aio": {
      "block_size": 1048576,
      "queue_depth": 8,
      "thread_count": 1,
      "single_submit": false,
      "overlap_events": true
    }
  }
}

帯域幅の考慮事項：

GPU-CPU帯域幅：PCIe 4.0 x16で~32 GB/s
CPU-NVMe帯域幅：NVMe SSDで~7 GB/s
NVMeはCPUメモリより遅く訓練オーバーヘッドが増加しますが、以前は不可能だったモデルサイズを実現します

4. DeepSpeed設定の完全ガイド

4.1 ds_config.jsonの基本構造

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": 1.0,

  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "bf16": {
    "enabled": "auto"
  },

  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  },

  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },

  "activation_checkpointing": {
    "partition_activations": false,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": null,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  },

  "wall_clock_breakdown": false,
  "steps_per_print": 100
}

4.2 混合精度設定の詳細

fp16設定：

{
  "fp16": {
    "enabled": true,
    "auto_cast": false,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "consecutive_hysteresis": false,
    "min_loss_scale": 1
  }
}

loss_scale: 0：動的ロススケーリングを有効化
initial_scale_power: 16：初期ロススケール = 2^16 = 65536
loss_scale_window：スケール増加前の成功ステップ数
hysteresis：スケール減少前のオーバーフロー回数

bf16設定（Ampere以降のGPU）：

{
  "bf16": {
    "enabled": true
  }
}

bf16はfp16より広いダイナミックレンジを持ち、ロススケーリングが不要です。A100/H100に推奨されます。

4.3 勾配累積とバッチサイズ

バッチサイズの関係：

train_batch_size = train_micro_batch_size_per_gpu x gradient_accumulation_steps x world_size

{
  "train_batch_size": 2048,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 64
}

この設定を8つのGPUで使用：4 x 64 x 8 = 2048のグローバルバッチサイズ。

"auto"を使用するとTransformers Trainerが値を自動的に設定します。

4.4 学習率スケジューラーオプション

{
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "last_batch_iteration": -1,
      "total_num_steps": 100000,
      "warmup_min_lr": 0,
      "warmup_max_lr": 3e-4,
      "warmup_num_steps": 2000,
      "warmup_type": "linear"
    }
  }
}

サポートされるスケジューラータイプ：

LRRangeTest
OneCycle
WarmupLR
WarmupDecayLR
WarmupCosineLR（カスタム）

5. DeepSpeed + PyTorch統合

5.1 基本的な訓練ループ

import torch
import deepspeed
from torch.utils.data import DataLoader

def main():
    # 分散環境を初期化
    deepspeed.init_distributed()

    # モデルを作成
    model = GPT2LMHeadModel.from_pretrained("gpt2")

    # DeepSpeedエンジンを初期化
    model_engine, optimizer, train_loader, lr_scheduler = deepspeed.initialize(
        model=model,
        training_data=train_dataset,
        config="ds_config.json"
    )

    # 訓練ループ
    for epoch in range(num_epochs):
        for batch in train_loader:
            input_ids = batch["input_ids"].to(model_engine.device)
            labels = batch["labels"].to(model_engine.device)

            # 順方向パス
            outputs = model_engine(input_ids=input_ids, labels=labels)
            loss = outputs.loss

            # 逆方向パス（DeepSpeedがスケーリング/累積を処理）
            model_engine.backward(loss)

            # パラメータ更新（勾配クリッピングを含む）
            model_engine.step()

    print(f"Step: {model_engine.global_steps}, Loss: {loss.item():.4f}")

5.2 チェックポイントの保存と読み込み

# チェックポイントを保存
def save_checkpoint(model_engine, save_dir, tag=None):
    # DeepSpeedフォーマットで保存（分散状態を含む）
    model_engine.save_checkpoint(save_dir, tag=tag)

# 例：1000ステップごとに保存
if model_engine.global_steps % 1000 == 0:
    save_checkpoint(model_engine, "./checkpoints", tag=f"step_{model_engine.global_steps}")

# チェックポイントを読み込み
_, client_sd = model_engine.load_checkpoint("./checkpoints", tag="step_1000")

# ZeRO-3下で標準PyTorchフォーマットで重みを抽出
if args.zero_stage == 3:
    state_dict = model_engine._zero3_consolidated_16bit_state_dict()
    if model_engine.local_rank == 0:
        torch.save(state_dict, "model_weights.pt")

5.3 Transformers Trainer統合

HuggingFace TransformersでDeepSpeedを統合する最もシンプルな方法：

from transformers import TrainingArguments, Trainer, AutoModelForCausalLM

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    fp16=True,
    deepspeed="ds_config.json",  # DeepSpeed設定ファイルへのパス
    logging_steps=10,
    save_steps=1000,
)

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

trainer.train()

コマンドライン実行：

deepspeed --num_gpus=8 train.py \
    --deepspeed ds_config.json \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --dataset_name openwebtext \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 3 \
    --fp16

5.4 Accelerate統合

from accelerate import Accelerator
from accelerate.utils import DeepSpeedPlugin

deepspeed_plugin = DeepSpeedPlugin(
    zero_stage=2,
    gradient_accumulation_steps=8,
    gradient_clipping=1.0,
)

accelerator = Accelerator(
    mixed_precision="fp16",
    deepspeed_plugin=deepspeed_plugin,
)

model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)

for batch in train_dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

6. DeepSpeedパイプライン並列

6.1 パイプライン並列の概念

パイプライン並列はモデルレイヤーを複数のGPUに順次配置します。GPU 0が初期レイヤー、GPU 1が中間レイヤー、GPU Nが最終レイヤーを担当します。

1F1Bスケジュール（One Forward, One Backward）： パイプラインバブルを最小化します。各GPUは1マイクロバッチの順方向パスの後、即座に逆方向パスを実行します。

6.2 PipelineModuleの使用

from deepspeed.pipe import PipelineModule, LayerSpec

class GPT2PipelineModel(PipelineModule):
    def __init__(self, num_layers, hidden_size, num_heads, vocab_size):
        layers = [
            LayerSpec(EmbeddingLayer, vocab_size, hidden_size),
        ]
        for i in range(num_layers):
            layers.append(LayerSpec(TransformerBlock, hidden_size, num_heads))
        layers.append(LayerSpec(LMHead, hidden_size, vocab_size))

        super().__init__(
            layers=layers,
            loss_fn=cross_entropy_loss,
            num_stages=4,       # パイプラインステージ（= GPUの数）
            topology=PipeDataParallelTopology(num_pp=4, num_dp=2)
        )

# ds_configにパイプライン設定を追加
pipeline_config = {
    "train_batch_size": 64,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 8,
    "pipeline": {
        "seed_layers": True,
        "activation_checkpoint_interval": 1
    }
}

6.3 パイプライン + ZeROの組み合わせ

model_engine, _, _, _ = deepspeed.initialize(
    model=pipeline_model,
    config={
        "train_micro_batch_size_per_gpu": 2,
        "gradient_accumulation_steps": 8,
        "zero_optimization": {"stage": 1},  # ZeRO-1と組み合わせ
        "fp16": {"enabled": True}
    }
)

# パイプライン訓練
for step in range(num_steps):
    loss = model_engine.train_batch()
    if model_engine.global_steps % 100 == 0:
        print(f"Step {model_engine.global_steps}, Loss: {loss.item():.4f}")

7. テンソル並列

7.1 Megatron-DeepSpeed統合

テンソル並列は行または列次元に沿って行列乗算を分割し、個々のレイヤーを複数のGPUに分散します。

# Megatron-DeepSpeed設定
megatron_config = {
  "tensor_model_parallel_size": 4,
  "pipeline_model_parallel_size": 2,
  "data_parallel_size": 4,  # world_size / (tp_size x pp_size)
}

ds_config.jsonのテンソル並列設定：

{
  "tensor_parallel": {
    "tp_size": 4,
    "mpu": null,
    "tp_grain_size": 64
  }
}

起動コマンド：

deepspeed --num_nodes=4 --num_gpus=8 \
    pretrain_gpt.py \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 4 \
    --num-layers 96 \
    --hidden-size 12288 \
    --num-attention-heads 96 \
    --seq-length 2048 \
    --global-batch-size 1024 \
    --train-iters 500000 \
    --deepspeed \
    --deepspeed_config ds_config.json

8. 活性化チェックポイント（勾配チェックポイント）

8.1 コンセプトとトレードオフ

活性化チェックポイントは順方向パス中に中間活性化を破棄し、逆伝播中に再計算します。

メモリ節約：活性化メモリをO(layers)からO(sqrt(layers))に削減
速度コスト：約33%の追加計算オーバーヘッド

{
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true,
    "contiguous_memory_optimization": true,
    "number_checkpoints": 4,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  }
}

# PyTorchレベルでも有効化
from deepspeed.runtime.activation_checkpointing import checkpointing

class TransformerLayer(nn.Module):
    def forward(self, x):
        return checkpointing.checkpoint(self._forward, x)

    def _forward(self, x):
        # 実際の計算
        return self.attention(self.norm1(x)) + x

9. DeepSpeedでのMixture of Experts（MoE）

9.1 MoEアーキテクチャの概要

Mixture of Expertsは各入力トークンを一部の「エキスパート」サブネットワークのみにルーティングし、計算量を一定に保ちながらパラメータ数を増加させます。

from deepspeed.moe.layer import MoE

class MoETransformerBlock(nn.Module):
    def __init__(self, hidden_size, num_experts, top_k=2):
        super().__init__()
        self.attention = MultiHeadAttention(hidden_size)
        self.moe_layer = MoE(
            hidden_size=hidden_size,
            expert=FeedForward(hidden_size, hidden_size * 4),
            num_experts=num_experts,
            ep_size=1,           # エキスパート並列サイズ
            k=top_k,             # 有効化するエキスパート数
            capacity_factor=1.25,
            eval_capacity_factor=2.0,
            min_capacity=4,
            use_residual=False
        )

    def forward(self, x, attention_mask=None):
        x = x + self.attention(x, attention_mask)
        x, _, _ = self.moe_layer(x)
        return x

9.2 エキスパート並列設定

{
  "zero_optimization": {
    "stage": 2
  },
  "moe_expert_parallel_size": 4,
  "moe": {
    "enabled": true,
    "ep_size": 4,
    "moe_param_group": true
  }
}

# MoEパラメータグループを分割（重要！）
def create_moe_param_groups(model):
    parameters = {
        "params": [],
        "name": "parameters"
    }
    moe_parameters = {
        "params": [],
        "moe": True,
        "name": "moe_parameters"
    }
    for module_name, module in model.named_modules():
        if isinstance(module, MoE):
            moe_parameters["params"].extend(
                [p for n, p in module.named_parameters() if "expert" in n]
            )
        else:
            parameters["params"].extend(module.parameters(recurse=False))
    return [parameters, moe_parameters]

10. DeepSpeed推論

10.1 推論の最適化

DeepSpeed推論は訓練済みモデルの推論速度を大幅に向上させます。

import deepspeed
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# DeepSpeed推論エンジンに変換
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=1,               # テンソル並列サイズ
    dtype=torch.float16,     # 推論データ型
    checkpoint=None,
    replace_with_kernel_inject=True,   # 最適化カーネルを注入
    max_tokens=2048,
    replace_method="auto"
)

model = ds_engine.module

# 推論を実行
inputs = tokenizer("DeepSpeed is ", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7
    )
print(tokenizer.decode(outputs[0]))

10.2 テンソル並列推論

# マルチGPU推論（テンソル並列）
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=4,               # 4つのGPUに分散
    dtype=torch.float16,
    replace_with_kernel_inject=True,
    injection_policy={
        LlamaDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")
    }
)

10.3 ZeRO-Inference

単一または少数のGPUでの大型モデル推論：

# ZeRO-3ベースの推論（CPU/NVMeからパラメータをストリーミング）
ds_engine = deepspeed.init_inference(
    model=model,
    config={
        "tensor_parallel": {"tp_size": 1},
        "dtype": "fp16",
        "enable_cuda_graph": False,
        "zero": {
            "stage": 3,
            "offload_param": {
                "device": "cpu"
            }
        }
    }
)

11. DeepSpeed Chat（RLHF）

11.1 RLHFパイプラインの概要

DeepSpeed-Chatは完全なRLHF（Reinforcement Learning from Human Feedback）パイプラインを提供します。

3つの訓練ステージ：

SFT（教師あり微調整）
報酬モデル訓練
PPOアルゴリズムによるRLHF/PPO

11.2 SFTステージ

# Step 1: SFT
deepspeed main.py \
    --data_path Dahoas/rm-static \
    --data_split 2,4,4 \
    --model_name_or_path facebook/opt-1.3b \
    --per_device_train_batch_size 8 \
    --max_seq_len 512 \
    --learning_rate 9.65e-6 \
    --weight_decay 0.1 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --lr_scheduler_type cosine \
    --num_warmup_steps 0 \
    --seed 1234 \
    --zero_stage 2 \
    --deepspeed \
    --output_dir /output/sft

11.3 報酬モデルの訓練

from deepspeed_chat.reward_model import RewardModel

# 報酬モデル = SFTモデル + 線形ヘッド
reward_model = RewardModel(
    base_model=sft_model,
    tokenizer=tokenizer,
    num_padding_at_beginning=0
)

# 選好ペア（chosen, rejected）で訓練
def compute_reward_loss(reward_model, chosen_ids, rejected_ids):
    chosen_reward = reward_model(chosen_ids)
    rejected_reward = reward_model(rejected_ids)
    # 選好される回答がより高い報酬を持つべき
    loss = -torch.mean(torch.log(torch.sigmoid(chosen_reward - rejected_reward)))
    return loss

11.4 PPO訓練

from deepspeed_chat.ppo_trainer import DeepSpeedPPOTrainer

trainer = DeepSpeedPPOTrainer(
    rlhf_engine=rlhf_engine,
    args=args
)

for epoch in range(args.num_train_epochs):
    for step, batch in enumerate(prompt_dataloader):
        out = trainer.generate_experience(batch["input_ids"])
        # Actor/Criticモデルを更新
        actor_loss, critic_loss = trainer.train_rlhf(out)

12. 実践例：LLM訓練

12.1 完全なLlama-2微調整の例

import torch
import deepspeed
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq
)
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, TaskType

def main():
    # モデルとトークナイザーをロード
    model_name = "meta-llama/Llama-2-7b-hf"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token

    # モデルをロード（DeepSpeedがデバイス配置を管理）
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map=None  # DeepSpeedが処理
    )

    # LoRA設定（オプション：パラメータ効率的な微調整）
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.1,
        bias="none"
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # データセットをロード
    dataset = load_dataset("tatsu-lab/alpaca", split="train")

    def preprocess(example):
        prompt = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
        tokens = tokenizer(prompt, truncation=True, max_length=512, padding="max_length")
        tokens["labels"] = tokens["input_ids"].copy()
        return tokens

    tokenized_dataset = dataset.map(preprocess, remove_columns=dataset.column_names)

    # DeepSpeedを使用したTrainingArguments
    training_args = TrainingArguments(
        output_dir="./llama2-finetuned",
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        learning_rate=2e-4,
        warmup_steps=100,
        logging_steps=10,
        save_steps=500,
        bf16=True,
        deepspeed="ds_zero2_config.json",
        remove_unused_columns=False,
        report_to="wandb",
        run_name="llama2-alpaca-deepspeed"
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True)
    )

    trainer.train()
    trainer.save_model("./llama2-finetuned-final")

if __name__ == "__main__":
    main()

ZeRO-2設定ファイル（ds_zero2_config.json）：

{
  "bf16": {
    "enabled": "auto"
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto",
      "total_num_steps": "auto"
    }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 100,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}

起動：

deepspeed --num_gpus=8 train_llama.py \
    --deepspeed ds_zero2_config.json

12.2 パフォーマンスベンチマーク

8x A100 80GBでの測定結果：

設定	最大モデルサイズ	スループット（tokens/sec）
ベースライン DDP	~7B	12,000
ZeRO-1	~14B	11,500
ZeRO-2	~35B	10,800
ZeRO-3	~175B	9,200
ZeRO-3 + CPUオフロード	~350B	4,100
ZeRO-Infinity	~1T以上	1,800

13. デバッグとトラブルシューティング

13.1 よくある問題と解決策

OOM（Out of Memory）エラー：

# メモリ使用量をプロファイリング
ds_config = {
    "memory_breakdown": True,
    "wall_clock_breakdown": True,
}

# 段階的な解決策：
# 1. ZeROステージを上げる（1 -> 2 -> 3）
# 2. CPUオフロードを有効化
# 3. 活性化チェックポイントを有効化
# 4. バッチサイズを削減
# 5. シーケンス長を削減

勾配オーバーフロー（fp16）：

{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 500,
    "initial_scale_power": 12,
    "hysteresis": 2,
    "min_loss_scale": 1e-8
  }
}

ZeRO-3パラメータ収集エラー：

# ZeRO-3でパラメータに直接アクセス
from deepspeed import zero

with zero.GatheredParameters(model.parameters()):
    # このブロック内でのみ完全なパラメータが利用可能
    weight_norm = model.weight.norm()

13.2 パフォーマンスチューニングのヒント

# 1. 通信と計算を重複させる
zero_config = {
    "overlap_comm": True,
    "contiguous_gradients": True,
}

# 2. バケットサイズを最適化（GPUメモリと帯域幅に合わせてチューニング）
zero_config["allgather_bucket_size"] = 5e8   # 500MB
zero_config["reduce_bucket_size"] = 5e8

# 3. ZeRO-3のプリフェッチ
zero_config["stage3_prefetch_bucket_size"] = 5e7
zero_config["stage3_param_persistence_threshold"] = 1e6

# 4. コンパイル最適化（PyTorch 2.0以降）
model = torch.compile(model)

まとめ

DeepSpeedは現代のLLM訓練に不可欠なツールです。ZeRO最適化により、単一GPUでは不可能な大規模モデルの訓練が合理的なコストで実現できます。

重要なポイント：

ZeRO-1/2：オプティマイザー/勾配の分割でデータ並列の効率を向上
ZeRO-3：パラメータも分割し、GPU数に比例したメモリ節約を実現
ZeRO-Offload：CPU/NVMeにオフロードし、小規模クラスタで大型モデルの訓練を可能に
パイプライン + テンソル並列：水平/垂直モデル分散で通信効率を最大化
DeepSpeed推論：カーネル最適化で2〜5倍の推論高速化

実際のプロジェクトでは、ZeRO-2 + CPUオフロードが良い出発点です。より大きなモデルにはZeRO-3に移行し、サービング用途にはDeepSpeed推論のカーネル注入を使用しましょう。