マルチGPU分散学習完全ガイド: DDP、FSDP、DeepSpeed

# NCCLデバッグログの有効化
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

# 通信インターフェースの指定（マルチノード環境で重要）
export NCCL_SOCKET_IFNAME=eth0

# InfiniBandの無効化（必要に応じて）
export NCCL_IB_DISABLE=1

# P2P（Peer-to-Peer）通信レベルの設定
export NCCL_P2P_LEVEL=NVL  # NVLinkを使用

GPU 0: モデルコピー + データシャード 0 → 勾配 0 ─┐
GPU 1: モデルコピー + データシャード 1 → 勾配 1 ─┤── AllReduce ──→ 平均勾配
GPU 2: モデルコピー + データシャード 2 → 勾配 2 ─┤
GPU 3: モデルコピー + データシャード 3 → 勾配 3 ─┘

GPU 0: レイヤーの重み行列の上半分  ─┐
                                     ├── 結合 → レイヤー出力
GPU 1: レイヤーの重み行列の下半分  ─┘

GPU 0: レイヤー 1-8    │ マイクロバッチ 1 → │ マイクロバッチ 2 → │ ...
GPU 1: レイヤー 9-16   │                    │ マイクロバッチ 1 → │ マイクロバッチ 2 → │ ...
GPU 2: レイヤー 17-24  │                    │                    │ マイクロバッチ 1 → │ ...
GPU 3: レイヤー 25-32  │                    │                    │                    │ マイクロバッチ 1 → │

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)

    # モデルの作成とDDPラッピング
    model = nn.Linear(1024, 512).to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # データセットとDataLoader（DistributedSamplerが必須）
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    optimizer = optim.Adam(ddp_model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)  # 各エポックで呼び出してデータシャッフルを保証
        for batch in dataloader:
            inputs, targets = batch
            inputs = inputs.to(rank)
            targets = targets.to(rank)

            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = loss_fn(outputs, targets)
            loss.backward()  # ここで勾配AllReduceが自動的に実行される
            optimizer.step()

    cleanup()

# 実行: torchrun --nproc_per_node=4 train_script.py

accumulation_steps = 4

for i, batch in enumerate(dataloader):
    inputs, targets = batch[0].to(rank), batch[1].to(rank)

    # 最後の蓄積ステップでのみ勾配を同期
    context = ddp_model.no_sync if (i + 1) % accumulation_steps != 0 else nullcontext
    with context():
        outputs = ddp_model(inputs)
        loss = loss_fn(outputs, targets) / accumulation_steps
        loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 単一ノード、4GPU
torchrun --nproc_per_node=4 train.py

# マルチノード（2ノード、各8GPU）
# ノード0上:
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
    --master_addr="192.168.1.1" --master_port=29500 train.py
# ノード1上:
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 \
    --master_addr="192.168.1.1" --master_port=29500 train.py

import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

def train_fsdp(rank, world_size):
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    # Transformerモデルの作成
    model = Transformer(model_args).to(rank)

    # Mixed Precisionの設定
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,    # パラメータをbfloat16で格納
        reduce_dtype=torch.float32,     # 勾配のリデュースをfloat32で実行
    )

    # 各Transformerレイヤーにfsdpを適用（レイヤーレベルのシャーディング）
    for layer in model.layers:
        fully_shard(layer, mp_policy=mp_policy)

    # トップレベルのモデルにもFSDPを適用
    fully_shard(model, mp_policy=mp_policy)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        for batch in dataloader:
            inputs = batch["input_ids"].to(rank)
            labels = batch["labels"].to(rank)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

{
  "zero_optimization": {
    "stage": 1,
    "reduce_bucket_size": 5e8
  }
}

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

{
  "zero_optimization": {
    "stage": 3,
    "contiguous_gradients": true,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_prefetch_bucket_size": 1e7,
    "stage3_param_persistence_threshold": 1e5,
    "reduce_bucket_size": 1e7,
    "sub_group_size": 1e9,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    }
  }
}

{
  "train_batch_size": 64,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 2e8,
    "allgather_bucket_size": 2e8
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 1e-4,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 1e-4,
      "warmup_num_steps": 1000,
      "total_num_steps": 50000
    }
  }
}

from accelerate import Accelerator

accelerator = Accelerator()

# 既存コードからの唯一の変更
model, optimizer, dataloader, scheduler = accelerator.prepare(
    model, optimizer, dataloader, scheduler
)

for batch in dataloader:
    optimizer.zero_grad()
    outputs = model(batch["input_ids"])
    loss = loss_fn(outputs, batch["labels"])
    accelerator.backward(loss)  # loss.backward()の代わり
    optimizer.step()
    scheduler.step()

$ accelerate config

# 質問に回答すると設定ファイルが自動生成される
# - 分散学習タイプ（マルチGPU、マルチノード、TPUなど）
# - GPU数
# - 混合精度の使用有無
# - DeepSpeed / FSDPの使用と設定

compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
num_machines: 1
num_processes: 4
mixed_precision: bf16
use_cpu: false

compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
deepspeed_config:
  zero_stage: 2
  gradient_accumulation_steps: 4
  offload_optimizer_device: none
  offload_param_device: none
mixed_precision: bf16
num_machines: 1
num_processes: 8

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
fsdp_config:
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_state_dict_type: SHARDED_STATE_DICT
mixed_precision: bf16
num_machines: 1
num_processes: 8

accelerate launch train.py

# 基本的なGPUステータスの確認
nvidia-smi

# 1秒間隔の自動更新モニタリング
watch -n 1 nvidia-smi

# 特定のGPUのみモニタリング
nvidia-smi --id=0,1

# プロセスごとのGPU使用量モニタリング
nvidia-smi pmon -i 0 -s um -d 1

# 継続的なデバイスモニタリング（1秒間隔）
nvidia-smi dmon -d 1

# CSV形式で主要メトリクスを出力（スクリプトのパースに便利）
nvidia-smi --query-gpu=index,name,temperature.gpu,utilization.gpu,utilization.memory,memory.used,memory.total,power.draw --format=csv,noheader,nounits

# 出力例:
# 0, NVIDIA A100-SXM4-80GB, 45, 98, 72, 65536, 81920, 285
# 1, NVIDIA A100-SXM4-80GB, 43, 95, 68, 61440, 81920, 278

dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,           # CPUコア数に合わせて調整
    pin_memory=True,          # GPU転送速度を向上
    prefetch_factor=2,        # 事前ロードするバッチ数
    persistent_workers=True,  # エポック間でワーカーを維持
)

import torch

# 現在のGPUメモリ使用量を確認
print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"Cached: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
print(f"Max Allocated: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

# メモリスナップショットの保存（詳細分析用）
torch.cuda.memory._record_memory_history()
# ... 学習コードを実行 ...
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# ステップ1: バッチサイズを削減
batch_size = 16  # → 8、4、2と段階的に減少を試みる

# ステップ2: Gradient Accumulationで実効バッチサイズを維持
gradient_accumulation_steps = 4  # batch_size * 4 = 実効バッチサイズ

# ステップ3: 混合精度を使用
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

with autocast(dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# ステップ4: Gradient Checkpointingを有効化
from torch.utils.checkpoint import checkpoint
# またはHugging Faceモデルの場合:
model.gradient_checkpointing_enable()

# ステップ5: メモリ割り当て設定の最適化
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

バッチサイズ削減 → Mixed Precision → Gradient Checkpointing
→ ZeRO Stage 2 → ZeRO Stage 3 → CPUオフローディング → NVMeオフローディング

# 原因1: タイムアウト値が短すぎる
# 対策: タイムアウトを延長
export NCCL_BLOCKING_WAIT=1

# Pythonでタイムアウトを設定
import datetime
dist.init_process_group(
    backend="nccl",
    timeout=datetime.timedelta(minutes=30)  # デフォルトは30分
)

# 原因2: GPU P2P通信の問題
# 対策: P2Pを無効化
export NCCL_P2P_DISABLE=1

# 原因3: ネットワークインターフェースの選択ミス（マルチノード）
# 対策: 正しいネットワークインターフェースを指定
export NCCL_SOCKET_IFNAME=eth0
export GLOO_SOCKET_IFNAME=eth0

# 原因4: Docker環境での共有メモリ不足
# 対策: Docker実行時にSHMサイズを増加
# docker run --shm-size=16g ...

# デバッグ用にNCCLログを有効化
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL
export TORCH_DISTRIBUTED_DEBUG=DETAIL

import torch.autograd.profiler as profiler

with profiler.profile(
    activities=[
        profiler.ProfilerActivity.CPU,
        profiler.ProfilerActivity.CUDA,
    ],
    schedule=profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=profiler.tensorboard_trace_handler('./log/profiler'),
    record_shapes=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(dataloader):
        if step >= 7:
            break
        train_step(model, batch)
        prof.step()

# 1. 勾配圧縮の使用（DDP）
from torch.distributed.algorithms.ddp_comm_hooks import (
    default_hooks as default,
    powerSGD_hook as powerSGD,
)
ddp_model.register_comm_hook(
    state=powerSGD.PowerSGDState(process_group=None, matrix_approximation_rank=1),
    hook=powerSGD.powerSGD_hook,
)

# 2. バケットサイズの調整（DDP）
ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=50)  # デフォルト25MB

# 3. overlap_commの有効化（DeepSpeed）
# ds_config.json内:
# "zero_optimization": { "overlap_comm": true }

モデルが単一GPUに収まるか?
├── はい → DDP（最速かつ最もシンプル）
│         ├── バッチサイズを増やしたい → ZeRO Stage 1/2 または FSDP SHARD_GRAD_OP
│         └── 十分 → DDPだけで十分
└── いいえ → モデルが全GPUの合計メモリに収まるか?
              ├── はい → FSDP (FULL_SHARD) または ZeRO Stage 3
              │         ├── PyTorchネイティブを好む → FSDP
              │         └── より多くのオプション/カスタマイズが必要 → DeepSpeed
              └── いいえ → ZeRO Stage 3 + CPU/NVMeオフローディング
                          または Pipeline Parallelism + Tensor Parallelismの組み合わせ