大規模モデル分散学習実践：DeepSpeed ZeRO、FSDPとマルチGPU最適化

AI与大数据

概要

  • 大規模モデル学習のメモリボトルネック：7BモデルFP16では28GBパラメータ+14GB勾配+14GBオプティマイザ=56GBが必要、単一A100 80GBでは学習不可
  • DeepSpeed ZeRO-3はパラメータ分割によりメモリ使用量を8倍削減、64GPUで70Bモデルの学習が可能
  • PyTorch FSDPはMeta公式ソリューション、PyTorchエコシステムとシームレスに統合、2026年には主流の選択肢に
  • 3D並列（データ+テンソル+パイプライン）は超大規模モデル学習の標準、Megatron-LMがデファクトスタンダード
  • 本記事では単一マシンマルチGPUから千GPUクラスタまでの完全な学習ソリューションを提供、Slurmスケジューリングと障害復旧も含む

目次

大規模モデル学習のメモリボトルネック

学習メモリの内訳

┌──────────────────────────────────────────────────────────────┐
│              大規模モデル学習メモリ内訳                          │
│                                                                │
│  モデルパラメータ(Weights):     2 × params × dtype_size       │
│  勾配(Gradients):               2 × params × dtype_size       │
│  オプティマイザ状態(Adam):      12 × params × dtype_size (FP32)│
│  活性化値(Activations):         batch × seq_len × hidden × layers│
│                                                                │
│  7BモデルFP16学習メモリ:                                       │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ パラメータ:     14GB                                  │    │
│  │ 勾配:           14GB                                  │    │
│  │ オプティマイザ: 56GB (FP32のmとv)                     │    │
│  │ 活性化値:       8-16GB (batchとseq_lenに依存)         │    │
│  │ ─────────────────────────                            │    │
│  │ 合計:           92-100GB ❌ 単一A100 80GBを超過       │    │
│  └──────────────────────────────────────────────────────┘    │
└──────────────────────────────────────────────────────────────┘

モデル規模別学習要件

モデル パラメータ数 学習メモリ(FP16+Adam) 最低GPU構成 推奨構成
1.8B 1.8B 26GB 1×A100 40GB 1×A100 80GB
7B 7B 100GB 2×A100 80GB 4×A100 80GB
13B 13B 186GB 4×A100 80GB 8×A100 80GB
70B 70B 1TB 16×A100 80GB 32×H100 80GB
671B 671B 9.6TB 128×H100 80GB 256×H100 80GB

DeepSpeed ZeRO：3段階メモリ最適化

ZeRO 3段階の原理

┌──────────────────────────────────────────────────────────────┐
│              ZeRO 3段階メモリ最適化                             │
│                                                                │
│  ZeRO-1: オプティマイザ状態分割                                │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ 各GPUはオプティマイザ状態(mとv)の1/Nのみを保存         │    │
│  │ メモリ削減: 4× (オプティマイザが最大の割合)            │    │
│  │ 通信量: DDPと同じ                                     │    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  ZeRO-2: オプティマイザ+勾配分割                              │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ 各GPUはオプティマイザ状態+勾配の1/Nのみを保存          │    │
│  │ メモリ削減: 8×                                        │    │
│  │ 通信量: DDPと同じ(勾配reduce-scatter)                 │    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  ZeRO-3: オプティマイザ+勾配+パラメータ分割                   │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ 各GPUは全状態の1/Nのみを保存                           │    │
│  │ メモリ削減: N× (N=GPU数)                              │    │
│  │ 通信量: 1.5×増加(all-gatherパラメータが必要)          │    │
│  └──────────────────────────────────────────────────────┘    │
└──────────────────────────────────────────────────────────────┘

DeepSpeed ZeRO-3学習設定

{
  "train_batch_size": 128,
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "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
  },
  "bf16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "prescale_gradients": false,
  "wall_clock_breakdown": false
}

ZeRO各段階のメモリ比較（7Bモデル、8×A100 80GB）

段階 パラメータ 勾配 オプティマイザ 活性化値 合計 GPU当たり
DDP 14GB 14GB 56GB 12GB 96GB 96GB
ZeRO-1 14GB 14GB 7GB 12GB 47GB 47GB
ZeRO-2 14GB 1.75GB 7GB 12GB 35GB 35GB
ZeRO-3 1.75GB 0.22GB 0.88GB 12GB 15GB 15GB

DeepSpeed起動スクリプト

deepspeed --num_gpus=8 \
    --num_nodes=2 \
    --master_addr=$MASTER_ADDR \
    --master_port=29500 \
    train.py \
    --model_name_or_path Qwen/Qwen2.5-7B \
    --data_path ./data/train.jsonl \
    --bf16 \
    --deepspeed ds_config_z3.json \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --learning_rate 2e-5 \
    --num_train_epochs 3 \
    --warmup_ratio 0.03 \
    --lr_scheduler_type cosine \
    --weight_decay 0.01 \
    --save_strategy epoch \
    --output_dir ./output \
    --gradient_checkpointing true \
    --logging_steps 10 \
    --report_to tensorboard

PyTorch FSDP：Meta公式ソリューション

FSDP vs DeepSpeed ZeRO

項目 DeepSpeed ZeRO-3 PyTorch FSDP
開発元 Microsoft Meta
統合度 独立ライブラリ PyTorchネイティブ
パラメータ分割 ✅ ✅
勾配分割 ✅ ✅
オプティマイザ分割 ✅ ✅
CPUオフロード ✅ ✅
混合精度 ✅ ✅
パイプライン並列 ✅ (DeepSpeed PP) ⚠️ (手動設定が必要)
デバッグのしやすさ 中 高(PyTorchネイティブ)
コミュニティの活発さ ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐

FSDP学習コード

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.distributed.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B",
    torch_dtype=torch.bfloat16,
    device_map={"": local_rank},
)

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

auto_wrap_policy = transformer_auto_wrap_policy(
    transformer_layer_names=["Qwen2DecoderLayer"],
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=mp_policy,
    auto_wrap_policy=auto_wrap_policy,
    device_id=local_rank,
    use_orig_params=True,
)

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)

for epoch in range(num_epochs):
    for batch in dataloader:
        batch = {k: v.to(local_rank) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

3D並列：データ+テンソル+パイプライン

3D並列アーキテクチャ

┌──────────────────────────────────────────────────────────────┐
│              3D並列アーキテクチャ                               │
│                                                                │
│  データ並列(DP): 4グループ × グループあたり4GPU                 │
│  ┌──────────────────────────────────────────────────────┐    │
│  │  DPグループ0    DPグループ1    DPグループ2    DPグループ3│    │
│  │  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────┐│    │
│  │  │ TP=2     │  │ TP=2     │  │ TP=2     │  │TP=2  ││    │
│  │  │ PP=2     │  │ PP=2     │  │ PP=2     │  │PP=2  ││    │
│  │  │ ┌──┬──┐ │  │ ┌──┬──┐ │  │ ┌──┬──┐ │  │┌──┬─┐││    │
│  │  │ │G0│G1│ │  │ │G4│G5│ │  │ │G8│G9│ │  ││12│13│││    │
│  │  │ ├──┼──┤ │  │ ├──┼──┤ │  │ ├──┼──┤ │  │├──┼──┤││    │
│  │  │ │G2│G3│ │  │ │G6│G7│ │  │ │10│11│ │  ││14│15│││    │
│  │  │ └──┴──┘ │  │ └──┴──┘ │  │ └──┴──┘ │  │└──┴──┘││    │
│  │  └──────────┘  └──────────┘  └──────────┘  └──────┘│    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  TP=テンソル並列(レイヤ内分割)  PP=パイプライン並列(レイヤ間分割)│
│  DP=データ並列(データ分割)  合計16GPU                          │
└──────────────────────────────────────────────────────────────┘

3D並列選択ガイド

モデル規模 GPU数 推奨並列戦略 説明
7B 8 DP=8 データ並列のみで十分
13B 16 DP=8, TP=2 テンソル並列を追加
70B 64 DP=8, TP=4, PP=2 3D並列
671B 256 DP=16, TP=8, PP=2 3D並列+ZeRO

マルチノード学習と障害復旧

Slurm起動スクリプト

#!/bin/bash
#SBATCH --job-name=llm-train
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=8
#SBATCH --mem=512G
#SBATCH --time=72:00:00
#SBATCH --partition=gpu-a100

export MASTER_ADDR=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n1)
export MASTER_PORT=29500
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_IB_HCA=mlx5

srun torchrun \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --nproc_per_node=8 \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train.py \
    --model_name_or_path Qwen/Qwen2.5-7B \
    --deepspeed ds_config_z3.json \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --learning_rate 2e-5 \
    --num_train_epochs 3 \
    --output_dir ./output \
    --save_strategy steps \
    --save_steps 500 \
    --save_total_limit 5

学習障害復旧

from transformers import Trainer

class FaultTolerantTrainer(Trainer):
    def _save_checkpoint(self, model, trial, metrics=None):
        super()._save_checkpoint(model, trial, metrics)
        self._save_training_state()

    def _save_training_state(self):
        state = {
            "global_step": self.state.global_step,
            "epoch": self.state.epoch,
            "rng_state": torch.cuda.get_rng_state().cpu().numpy().tolist(),
        }
        with open(f"{self.args.output_dir}/training_state.json", "w") as f:
            json.dump(state, f)

    def _load_training_state(self):
        state_path = f"{self.args.output_dir}/training_state.json"
        if os.path.exists(state_path):
            with open(state_path) as f:
                state = json.load(f)
            return state
        return None

学習パフォーマンス最適化チェックリスト

最適化項目 効果 設定
Flash Attention 2 学習速度+30% attn_implementation="flash_attention_2"
勾配チェックポイント メモリ-40% gradient_checkpointing=True
BF16混合精度 メモリ-50% bf16=True
NCCL IB通信 通信遅延-60% NCCL_IB_DISABLE=0
シャードパラメータプリフェッチ 通信計算オーバーラップ stage3_prefetch_bucket_size
勾配蓄積 大batchと同等 gradient_accumulation_steps=8

まとめと関連記事

大規模モデル分散学習の核心はメモリ最適化と通信最適化です。DeepSpeed ZeRO-3はパラメータ分割によりメモリをN倍削減し、FSDPはPyTorchネイティブソリューションでデバッグが容易、3D並列は超大規模モデルの標準です。

学習のポイント振り返り：

  1. ZeRO-3はメモリ最適化の究極のソリューションですが、通信量は1.5倍増加します
  2. FSDPはPyTorchネイティブソリューションで、2026年には主流の選択肢となっています
  3. 3D並列：7BはDP、13BはTP追加、70B+はDP+TP+PP
  4. Flash Attention 2 + 勾配チェックポイントは学習の必須設定です
  5. 障害復旧とチェックポイントは長時間学習に不可欠です

関連記事：

権威ある参考文献：

ブラウザローカルツールを無料で試す →

#大模型分布式训练#DeepSpeed ZeRO#FSDP训练#Megatron-LM#多GPU训练优化#2026