PyTorch 分散トレーニング
モデルが大きくなり、データが増えるにつれて、単一GPUではトレーニング要件を満たせなくなります。
分散トレーニングは、複数のGPUや複数台のマシンによる並列計算により、トレーニング時間を大幅に短縮できます。
このセクションでは、PyTorch の分散トレーニング技術について詳しく説明します。DataParallel、DistributedDataParallel、混合精度分散トレーニングを含みます。
1. 分散トレーニングの基礎
1.1 なぜ分散トレーニングが必要か
深層学習モデルの規模は年々拡大し、トレーニング時間も増加しています。分散トレーニングの中心的な目的は次のとおりです。計算タスクを複数の計算ユニットに分散し、許容可能な時間内に大規模モデルとデータのトレーニングを完了する。
分散トレーニングは主に2つの問題を解決します。
- GPUメモリ不足：大規模モデルのパラメータ、オプティマイザの状態、勾配には大量のGPUメモリが必要
- トレーニング時間が長すぎる：単一GPUでの大規模モデルのトレーニングには数週間から数か月かかる可能性がある
1.2 分散トレーニングの2大モード
分散トレーニングは主に2つのモードに分けられます。
| モード | 原理 | 適用シナリオ |
|---|---|---|
| データ並列（Data Parallel） | 各計算ユニットが完全なモデルを保持し、異なる計算ユニットが異なるデータを処理する | モデルが単一GPUのメモリに収まる場合。データ量が大きくトレーニングを高速化する必要がある場合。小規模チームでマシンリソースが限られている場合 |
| モデル並列（Model Parallel） | モデルを複数の計算ユニットに分割し、各計算ユニットはモデルの一部のみを保持する | モデルが非常に大きく単一GPUに収まらない場合。専用クラスターでのマルチマシントレーニング |
実際の本番環境では、データ並列が最も一般的な方法です。PyTorch は DataParallel と DistributedDataParallel の2つの実装を提供しています。
2. DataParallel 使用ガイド
2.1 単一マシン複数GPU：DataParallel
DataParallel（DP）これは PyTorch が提供する最も簡単なマルチGPUトレーニング方法です。数行のコードを変更するだけで、単一マシン上の複数のGPUを利用してトレーニングできます。
例
import torch
import torch.nn as nn
import torch.optim as optim
# ── モデル定義 ──────────────────────────────────────
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(128, 256),
nn.ReLU(),
nn.Linear(256, 256),
nn.ReLU(),
nn.Linear(256, 10)
)
def forward(self, x):
return self.net(x)
# 利用可能なGPU数を確認
print(f"利用可能なGPU数: {torch.cuda.device_count()}")
# モデルを作成してGPUに移動
model = SimpleModel()
# ── 方法1：DataParallelを使用（最も簡単な方法）───────
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
# GPUに移動（DataParallelを使用する場合、device_ids は自動的に処理されます）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 損失関数とオプティマイザ
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# ── トレーニングループ ──────────────────────────────────────
def train_epoch_dp(model, loader, criterion, optimizer, device):
model.train()
total_loss = 0
correct = 0
total = 0
for inputs, labels in loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
# ダミーデータ
train_loader = [
(torch.randn(32, 128), torch.randint(0, 10, (32,))) for _ in range(10)
]
# トレーニング開始
for epoch in range(3):
loss, acc = train_epoch_dp(model, train_loader, criterion, optimizer, device)
print(f"Epoch {epoch+1}: Loss={loss:.4f}, Acc={acc:.4f}")
print(DataParallel トレーニング完了！)
import torch.nn as nn
import torch.optim as optim
# ── モデル定義 ──────────────────────────────────────
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(128, 256),
nn.ReLU(),
nn.Linear(256, 256),
nn.ReLU(),
nn.Linear(256, 10)
)
def forward(self, x):
return self.net(x)
# 利用可能なGPU数を確認
print(f"利用可能なGPU数: {torch.cuda.device_count()}")
# モデルを作成してGPUに移動
model = SimpleModel()
# ── 方法1：DataParallelを使用（最も簡単な方法）───────
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
# GPUに移動（DataParallelを使用する場合、device_ids は自動的に処理されます）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 損失関数とオプティマイザ
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# ── トレーニングループ ──────────────────────────────────────
def train_epoch_dp(model, loader, criterion, optimizer, device):
model.train()
total_loss = 0
correct = 0
total = 0
for inputs, labels in loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
# ダミーデータ
train_loader = [
(torch.randn(32, 128), torch.randint(0, 10, (32,))) for _ in range(10)
]
# トレーニング開始
for epoch in range(3):
loss, acc = train_epoch_dp(model, train_loader, criterion, optimizer, device)
print(f"Epoch {epoch+1}: Loss={loss:.4f}, Acc={acc:.4f}")
print(DataParallel トレーニング完了！)
2.2 DataParallel の動作原理
DataParallel の処理フローは次のとおりです。
- モデル複製：モデルを各GPUにコピーする
- データ分散：バッチデータを各GPUに均等に割り当てる
- 並列計算：各GPUが独立して順伝播を計算する
- 勾配集約：全GPUの勾配をメインGPUに集約する
- パラメータ更新：メインGPUがパラメータを更新後、他のGPUに同期する
例
# DataParallel の主要パラメータの説明
model = nn.DataParallel(
module, # 並列化するモデル（必須）
device_ids=[0, 1, 2, 3], # 使用するGPUデバイスID、デフォルトでは全GPUを使用
output_device=0, # 出力結果を集約するデバイス、デフォルトは1番目のGPU
dim=0 # データ分割の次元、デフォルトはバッチ次元
)
# 現在のモデルが配置されているデバイスを確認
print(f"モデルが配置されているデバイス: {next(model.parameters()).device}")
# 実際に使用されているデバイスを確認
print(f"使用されているGPU数: {model.device_ids if hasattr(model, 'device_ids') else 'N/A'}")
model = nn.DataParallel(
module, # 並列化するモデル（必須）
device_ids=[0, 1, 2, 3], # 使用するGPUデバイスID、デフォルトでは全GPUを使用
output_device=0, # 出力結果を集約するデバイス、デフォルトは1番目のGPU
dim=0 # データ分割の次元、デフォルトはバッチ次元
)
# 現在のモデルが配置されているデバイスを確認
print(f"モデルが配置されているデバイス: {next(model.parameters()).device}")
# 実際に使用されているデバイスを確認
print(f"使用されているGPU数: {model.device_ids if hasattr(model, 'device_ids') else 'N/A'}")
DataParallel はシンプルで使いやすいですが、2つの主要な問題があります。1）メインGPUのメモリ負荷が大きい（勾配を集約する必要がある）。2）GPU間の通信効率が低い。そのため、公式では DistributedDataParallel の使用を推奨しています。
3. DistributedDataParallel 詳細解説
3.1 DDP の基礎
DistributedDataParallel（DDP）これは PyTorch が推奨する分散トレーニング方式です。DataParallel と比較して、DDP には以下の利点があります。
- 各GPUが独立して勾配を計算し、メインGPUでの集約が不要
- 効率的な勾配同期アルゴリズム（Ring AllReduce）を使用
- マルチマシン・マルチGPUトレーニングをサポート
- トレーニング速度が速く、GPUメモリの使用効率が高い
3.2 DDP による単一マシン複数GPUトレーニング
例
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os
# ── 分散環境の初期化 ────────────────────────────
def setup(rank, world_size):
"""分散環境を設定する"""
# 現在のプロセスが使用するGPUを設定
torch.cuda.set_device(rank)
# プロセスグループを初期化
dist.init_process_group(
backend="nccl", # NCCL バックエンドを使用（GPU推奨）
init_method="env://", # 環境変数による初期化方法
world_size=world_size, # 総プロセス数
rank=rank # 現在のプロセスのランク
)
def cleanup():
"""分散環境をクリーンアップする"""
dist.destroy_process_group()
# ── データローダー（分散対応）───────────────────────
def get_distributed_loader(batch_size, world_size):
"""分散対応のデータローダーを作成する"""
from torch.utils.data import DataLoader, DistributedSampler
# ダミーデータセット
dataset = torch.utils.data.TensorDataset(
torch.randn(100, 3, 32, 32),
torch.randint(0, 10, (100,))
)
# DistributedSampler がデータを自動的に分割する
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True
)
loader = DataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=2,
pin_memory=True
)
return loader
# ── モデル定義 ──────────────────────────────────────
class ImageClassifier(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
)
self.classifier = nn.Linear(64, 10)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
# ── DDP トレーニング関数 ──────────────────────────────────
def train_ddp(rank, world_size, epochs=3):
"""分散トレーニングのメイン関数"""
# 初期化
setup(rank, world_size)
# モデルを作成して対応するGPUに移動
model = ImageClassifier().cuda(rank)
# DDPモデルにラップ
ddp_model = DDP(
model,
device_ids=[rank], # 現在のプロセスが使用するデバイスを指定
output_device=rank,
find_unused_parameters=False # 未使用のパラメータを検出するかどうか
)
# 損失関数とオプティマイザ
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(ddp_model.parameters(), lr=1e-3)
# データローダーを取得
loader = get_distributed_loader(batch_size=16, world_size=world_size)
# 訓練ループ
for epoch in range(epochs):
# 各 epoch の開始時に epoch 番号を設定（データをシャッフルするため）
loader.sampler.set_epoch(epoch)
for batch_idx, (inputs, labels) in enumerate(loader):
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
outputs = ddp_model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if rank == 0 and batch_idx % 5 == 0:
print(f"Epoch {epoch+1}, Batch {batch_idx}, Loss: {loss.item():.4f}")
# クリーンアップ
if rank == 0:
print(トレーニング完了！)
cleanup()
# 注意：実際の実行時にはコマンドライン引数で rank と world_size を指定する必要があります
# このコードは各プロセスで独立して実行する必要があります
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os
# ── 分散環境の初期化 ────────────────────────────
def setup(rank, world_size):
"""分散環境を設定する"""
# 現在のプロセスが使用するGPUを設定
torch.cuda.set_device(rank)
# プロセスグループを初期化
dist.init_process_group(
backend="nccl", # NCCL バックエンドを使用（GPU推奨）
init_method="env://", # 環境変数による初期化方法
world_size=world_size, # 総プロセス数
rank=rank # 現在のプロセスのランク
)
def cleanup():
"""分散環境をクリーンアップする"""
dist.destroy_process_group()
# ── データローダー（分散対応）───────────────────────
def get_distributed_loader(batch_size, world_size):
"""分散対応のデータローダーを作成する"""
from torch.utils.data import DataLoader, DistributedSampler
# ダミーデータセット
dataset = torch.utils.data.TensorDataset(
torch.randn(100, 3, 32, 32),
torch.randint(0, 10, (100,))
)
# DistributedSampler がデータを自動的に分割する
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True
)
loader = DataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=2,
pin_memory=True
)
return loader
# ── モデル定義 ──────────────────────────────────────
class ImageClassifier(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
)
self.classifier = nn.Linear(64, 10)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
# ── DDP トレーニング関数 ──────────────────────────────────
def train_ddp(rank, world_size, epochs=3):
"""分散トレーニングのメイン関数"""
# 初期化
setup(rank, world_size)
# モデルを作成して対応するGPUに移動
model = ImageClassifier().cuda(rank)
# DDPモデルにラップ
ddp_model = DDP(
model,
device_ids=[rank], # 現在のプロセスが使用するデバイスを指定
output_device=rank,
find_unused_parameters=False # 未使用のパラメータを検出するかどうか
)
# 損失関数とオプティマイザ
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(ddp_model.parameters(), lr=1e-3)
# データローダーを取得
loader = get_distributed_loader(batch_size=16, world_size=world_size)
# 訓練ループ
for epoch in range(epochs):
# 各 epoch の開始時に epoch 番号を設定（データをシャッフルするため）
loader.sampler.set_epoch(epoch)
for batch_idx, (inputs, labels) in enumerate(loader):
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
outputs = ddp_model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if rank == 0 and batch_idx % 5 == 0:
print(f"Epoch {epoch+1}, Batch {batch_idx}, Loss: {loss.item():.4f}")
# クリーンアップ
if rank == 0:
print(トレーニング完了！)
cleanup()
# 注意：実際の実行時にはコマンドライン引数で rank と world_size を指定する必要があります
# このコードは各プロセスで独立して実行する必要があります
3.3 分散トレーニングの起動
PyTorch の分散トレーニングは、torchrun 或 torch.distributed.launch起動する必要があります。以下は一般的な起動方法です：
例
# 方法1：torchrun を使用（推奨、PyTorch 2.0+）
# ファイル：train_ddp.py
"""
使用方法：
torchrun --nproc_per_node=4 train_ddp.py
"""
# 方法2：python -m を使用（従来の方法）
"""
python -m torch.distributed.launch --nproc_per_node=4 train_ddp.py
"""
# 方法3：複数ノード・複数GPUでの起動
# マシン1（マスターノード）
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 train_ddp.py
# マシン2（ワーカーノード）
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 train_ddp.py
# パラメータの説明：
# --nproc_per_node: 各ノードのGPU数
# --nnodes: ノード数
# --node_rank: 現在のノード番号（0から開始）
# --master_addr: マスターノードのIPアドレス
# --master_port: マスターノードのポート番号
# ファイル：train_ddp.py
"""
使用方法：
torchrun --nproc_per_node=4 train_ddp.py
"""
# 方法2：python -m を使用（従来の方法）
"""
python -m torch.distributed.launch --nproc_per_node=4 train_ddp.py
"""
# 方法3：複数ノード・複数GPUでの起動
# マシン1（マスターノード）
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 train_ddp.py
# マシン2（ワーカーノード）
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 train_ddp.py
# パラメータの説明：
# --nproc_per_node: 各ノードのGPU数
# --nnodes: ノード数
# --node_rank: 現在のノード番号（0から開始）
# --master_addr: マスターノードのIPアドレス
# --master_port: マスターノードのポート番号
3.4 完全な DDP トレーニングスクリプト
例
# 完全なDDPトレーニングスクリプト（train_ddp.pyとして保存）
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import argparse
import os
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1, help=torchrun によって自動的に渡されます)
parser.add_argument("--epochs", type=int, default=10)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--lr", type=float, default=1e-3)
return parser.parse_args()
class TrainModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(128, 10)
)
def forward(self, x):
return self.net(x)
def setup(rank, world_size):
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def cleanup():
dist.destroy_process_group()
def train(rank, world_size, args):
setup(rank, world_size)
# モデルを作成
model = TrainModel().cuda(rank)
model = DDP(model, device_ids=[rank])
# オプティマイザと損失関数
optimizer = optim.Adam(model.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss()
# データセット
dataset = torch.utils.data.TensorDataset(
torch.randn(1000, 3, 32, 32),
torch.randint(0, 10, (1000,))
)
# DistributedSampler は各プロセスが異なるデータを見るようにします
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True
)
loader = DataLoader(
dataset,
batch_size=args.batch_size,
sampler=sampler,
num_workers=2,
pin_memory=True
)
# トレーニングループ
for epoch in range(args.epochs):
# 重要：各 epoch で epoch を設定し、データの分割をシャッフルします
sampler.set_epoch(epoch)
model.train()
epoch_loss = 0
for inputs, labels in loader:
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
# すべてのプロセスを同期し、各プロセスが現在の epoch を完了することを確認します
dist.barrier()
if rank == 0:
avg_loss = epoch_loss / len(loader)
print(f"Epoch {epoch+1}/{args.epochs}, Loss: {avg_loss:.4f}")
cleanup()
if __name__ == "__main__":
args = parse_args()
# 環境変数から world_size を取得
world_size = int(os.environ["WORLD_SIZE"])
# 環境変数から rank を取得（torchrun によって自動設定）
rank = int(os.environ["RANK"])
if torch.cuda.is_available():
train(rank, world_size, args)
else:
print("分散トレーニングを実行するには CUDA 環境が必要です")
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import argparse
import os
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1, help=torchrun によって自動的に渡されます)
parser.add_argument("--epochs", type=int, default=10)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--lr", type=float, default=1e-3)
return parser.parse_args()
class TrainModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(128, 10)
)
def forward(self, x):
return self.net(x)
def setup(rank, world_size):
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def cleanup():
dist.destroy_process_group()
def train(rank, world_size, args):
setup(rank, world_size)
# モデルを作成
model = TrainModel().cuda(rank)
model = DDP(model, device_ids=[rank])
# オプティマイザと損失関数
optimizer = optim.Adam(model.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss()
# データセット
dataset = torch.utils.data.TensorDataset(
torch.randn(1000, 3, 32, 32),
torch.randint(0, 10, (1000,))
)
# DistributedSampler は各プロセスが異なるデータを見るようにします
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True
)
loader = DataLoader(
dataset,
batch_size=args.batch_size,
sampler=sampler,
num_workers=2,
pin_memory=True
)
# トレーニングループ
for epoch in range(args.epochs):
# 重要：各 epoch で epoch を設定し、データの分割をシャッフルします
sampler.set_epoch(epoch)
model.train()
epoch_loss = 0
for inputs, labels in loader:
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
# すべてのプロセスを同期し、各プロセスが現在の epoch を完了することを確認します
dist.barrier()
if rank == 0:
avg_loss = epoch_loss / len(loader)
print(f"Epoch {epoch+1}/{args.epochs}, Loss: {avg_loss:.4f}")
cleanup()
if __name__ == "__main__":
args = parse_args()
# 環境変数から world_size を取得
world_size = int(os.environ["WORLD_SIZE"])
# 環境変数から rank を取得（torchrun によって自動設定）
rank = int(os.environ["RANK"])
if torch.cuda.is_available():
train(rank, world_size, args)
else:
print("分散トレーニングを実行するには CUDA 環境が必要です")
重要なポイント：各 epoch で呼び出す必要があります
sampler.set_epoch(epoch)、データの分割がシャッフルされることを確認します。そうしないと、複数のプロセスが同じデータを見ることになります。
4. 分散トレーニングにおける混合精度
4.1 DDP + AMP の組み合わせ
分散トレーニングと混合精度トレーニングを組み合わせることで、トレーニング速度をさらに向上させることができます。PyTorch 2.0+ の DDP は AMP と完全に統合されています。
例
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler
class AMPModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(128, 10)
)
def forward(self, x):
return self.net(x)
def train_ddp_amp(rank, world_size):
"""分散 + 混合精度トレーニング"""
# 初期化
torch.cuda.set_device(rank)
dist.init_process_group("nccl", rank=rank, world_size=world_size)
# モデルを作成
model = AMPModel().cuda(rank)
model = DDP(model, device_ids=[rank])
# オプティマイザと GradScaler
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
criterion = nn.CrossEntropyLoss()
# トレーニングループ
model.train()
for inputs, labels in loader:
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
# autocast を使用して精度を自動的に切り替えます
with autocast(device_type='cuda'):
outputs = model(inputs)
loss = criterion(outputs, labels)
# scaler を使用して勾配のスケーリングを処理します
scaler.scale(loss).backward()
# 勾配クリッピング（最初に unscale が必要）
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
dist.destroy_process_group()
# 起動コマンド
# torchrun --nproc_per_node=4 train_ddp_amp.py
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler
class AMPModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(128, 10)
)
def forward(self, x):
return self.net(x)
def train_ddp_amp(rank, world_size):
"""分散 + 混合精度トレーニング"""
# 初期化
torch.cuda.set_device(rank)
dist.init_process_group("nccl", rank=rank, world_size=world_size)
# モデルを作成
model = AMPModel().cuda(rank)
model = DDP(model, device_ids=[rank])
# オプティマイザと GradScaler
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
criterion = nn.CrossEntropyLoss()
# トレーニングループ
model.train()
for inputs, labels in loader:
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
# autocast を使用して精度を自動的に切り替えます
with autocast(device_type='cuda'):
outputs = model(inputs)
loss = criterion(outputs, labels)
# scaler を使用して勾配のスケーリングを処理します
scaler.scale(loss).backward()
# 勾配クリッピング（最初に unscale が必要）
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
dist.destroy_process_group()
# 起動コマンド
# torchrun --nproc_per_node=4 train_ddp_amp.py
4.2 分散トレーニングにおける勾配同期
DDP は次を使用しますRing AllReduceアルゴリズムを使用して勾配を同期します。その動作原理は次のとおりです：
- 各 GPU は勾配を n 個に分割します（n は GPU 数）
- 各 GPU は隣接する GPU と 1 つの勾配を交換します
- n 回繰り返すと、各 GPU は完全な勾配を持ちます
例
# DDP 勾配同期関連の設定
# 1. 勾配バケットのバッファサイズを設定
model = DDP(
model,
device_ids=[rank],
gradient_as_bucket_view=True, # メモリを節約（PyTorch 1.8+）
broadcast_buffers=False, # buffers をブロードキャストせず、通信を削減
bucket_cap_mb=25 # 勾配バケットサイズ（MB）
)
# 2. モデルパラメータを手動で同期（必要に応じて）
# 一部のシナリオでは、モデルパラメータを手動で同期する必要があります
def sync_params(model):
for param in model.parameters():
dist.broadcast(param, src=0)
# 3. 勾配同期状態を確認
# DDP は勾配同期を自動的に処理するため、手動操作は不要です
# ただし、次の方法で確認できます
for name, param in model.named_parameters():
if param.grad is not None:
print(f"{name}: grad requires_grad={param.grad.requires_grad}")
# 1. 勾配バケットのバッファサイズを設定
model = DDP(
model,
device_ids=[rank],
gradient_as_bucket_view=True, # メモリを節約（PyTorch 1.8+）
broadcast_buffers=False, # buffers をブロードキャストせず、通信を削減
bucket_cap_mb=25 # 勾配バケットサイズ（MB）
)
# 2. モデルパラメータを手動で同期（必要に応じて）
# 一部のシナリオでは、モデルパラメータを手動で同期する必要があります
def sync_params(model):
for param in model.parameters():
dist.broadcast(param, src=0)
# 3. 勾配同期状態を確認
# DDP は勾配同期を自動的に処理するため、手動操作は不要です
# ただし、次の方法で確認できます
for name, param in model.named_parameters():
if param.grad is not None:
print(f"{name}: grad requires_grad={param.grad.requires_grad}")
5. モデル並列とパイプライン
5.1 シンプルなモデル並列
モデルが大きすぎて単一の GPU に収まらない場合、モデルを複数の GPU に分割する必要があります。
例
# モデル並列の例：モデルの異なる層を異なる GPU に配置
class ModelParallel(nn.Module):
def __init__(self, num_gpus=2):
super().__init__()
self.num_gpus = num_gpus
# 最初の部分：GPU 0
self.features1 = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
).to(f"cuda:0")
# 2番目の部分：GPU 1
self.features2 = nn.Sequential(
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
).to(f"cuda:1")
# 分類器：GPU 0
self.classifier = nn.Linear(128, 10).to(f"cuda:0")
def forward(self, x):
# データは異なる GPU 間で転送されます
x = x.to("cuda:0")
x = self.features1(x)
x = x.to("cuda:1")
x = self.features2(x)
x = x.to("cuda:0")
x = self.classifier(x)
return x
# 使用
model = ModelParallel(num_gpus=2)
output = model(torch.randn(1, 3, 32, 32))
print(f"出力デバイス: {output.device}")
class ModelParallel(nn.Module):
def __init__(self, num_gpus=2):
super().__init__()
self.num_gpus = num_gpus
# 最初の部分：GPU 0
self.features1 = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
).to(f"cuda:0")
# 2番目の部分：GPU 1
self.features2 = nn.Sequential(
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
).to(f"cuda:1")
# 分類器：GPU 0
self.classifier = nn.Linear(128, 10).to(f"cuda:0")
def forward(self, x):
# データは異なる GPU 間で転送されます
x = x.to("cuda:0")
x = self.features1(x)
x = x.to("cuda:1")
x = self.features2(x)
x = x.to("cuda:0")
x = self.classifier(x)
return x
# 使用
model = ModelParallel(num_gpus=2)
output = model(torch.randn(1, 3, 32, 32))
print(f"出力デバイス: {output.device}")
5.2 パイプラインパラレリズム（Pipeline Parallelism）
PyTorch は、次を提供しています。torch.distributed.pipeline.syncモジュール（PP パッケージ）を使用してパイプライン並列を実装し、モデルを層ごとに分割して計算パイプラインを形成します。
例
# パイプライン並列の例（torch-pipeline 依存関係のインストールが必要）
# pip install torch-pipeline
"""
# 注意：PyTorch ネイティブのパイプラインは特定のバージョンまたはサードパーティライブラリを使用する必要があります
# ここでは概念を示し、実装を簡略化します
from torch.distributed.pipeline.sync import Pipe
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(128, 256)
self.layer2 = nn.Linear(256, 256)
self.layer3 = nn.Linear(256, 10)
def forward(self, x):
x = torch.relu(self.layer1(x))
x = torch.relu(self.layer2(x))
x = self.layer3(x)
return x
# モデルを層ごとに異なる GPU に分割
# layer1 -> GPU 0
# layer2 -> GPU 1
# layer3 -> GPU 2
model = SimpleModel()
model.layer1 = model.layer1.cuda(0)
model.layer2 = model.layer2.cuda(1)
model.layer3 = model.layer3.cuda(2)
# Pipeline を使用してモデルを接続
# 注意：実際に使用するには、対応するバージョンの pipeline ライブラリをインストールする必要があります
# from torch.distributed.pipeline.sync import Pipe
# model = Pipe(model, chunks=4)
"""
print("パイプライン並列には torch-pipeline ライブラリの追加インストールが必要です")
print("pip install torch-pipeline-sync")
# pip install torch-pipeline
"""
# 注意：PyTorch ネイティブのパイプラインは特定のバージョンまたはサードパーティライブラリを使用する必要があります
# ここでは概念を示し、実装を簡略化します
from torch.distributed.pipeline.sync import Pipe
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(128, 256)
self.layer2 = nn.Linear(256, 256)
self.layer3 = nn.Linear(256, 10)
def forward(self, x):
x = torch.relu(self.layer1(x))
x = torch.relu(self.layer2(x))
x = self.layer3(x)
return x
# モデルを層ごとに異なる GPU に分割
# layer1 -> GPU 0
# layer2 -> GPU 1
# layer3 -> GPU 2
model = SimpleModel()
model.layer1 = model.layer1.cuda(0)
model.layer2 = model.layer2.cuda(1)
model.layer3 = model.layer3.cuda(2)
# Pipeline を使用してモデルを接続
# 注意：実際に使用するには、対応するバージョンの pipeline ライブラリをインストールする必要があります
# from torch.distributed.pipeline.sync import Pipe
# model = Pipe(model, chunks=4)
"""
print("パイプライン並列には torch-pipeline ライブラリの追加インストールが必要です")
print("pip install torch-pipeline-sync")
6. 分散トレーニングのベストプラクティス
6.1 パフォーマンス最適化のヒント
| 最適化項目 | テクニック | 説明 |
|---|---|---|
| 通信の最適化 | NCCL バックエンドを使用 | NCCL は GPU に最適な選択肢です。Gloo より高速です。 |
| データ読み込み | num_workers と pin_memory を設定 | num_workers は CPU コア数に設定します。pin_memory は CPU から GPU への転送を高速化します。 |
| 勾配同期 | bucket_cap_mb を調整 | 大規模モデルでは bucket を大きくできます。小規模モデルでは小さくして応答性を高められます。 |
| 混合精度 | DDP + AMP | 通信量を削減し、トレーニング速度を向上 |
6.2 よくある問題と解決策
| 問題 | 原因 | 解決策 |
|---|---|---|
| 通信が遅い | ネットワーク帯域幅が不足 | InfiniBand または高速ネットワークを使用 |
| メモリ不足（VRAM） | batch size が大きすぎる | batch size を小さくするか、勾配累積を使用 |
| 負荷の不均衡 | モデルの分割が不適切 | 層の分布を調整し、できるだけ均等にします |
| 同期エラー | 勾配が NaN | 勾配クリッピングとスケーリングの設定を確認 |
6.3 完全なマルチマシン・マルチGPUトレーニングスクリプト
例
# 完全な複数ノード・複数GPUトレーニングスクリプトのテンプレート
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import argparse
import os
class TrainingConfig:
def __init__(self):
self.local_rank = -1
self.epochs = 30
self.batch_size = 32
self.lr = 1e-3
self.num_workers = 4
self.gradient_clip = 1.0
self.use_amp = True
def init_distributed(config):
"""分散環境の初期化"""
# 環境変数から構成を取得
local_rank = config.local_rank
# NCCL 設定
os.environ["NCCL_DEBUG"] = "WARN"
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")
return local_rank
def create_dataloader(dataset, config, world_size, rank):
"""データローダーの作成"""
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True,
drop_last=True
)
loader = DataLoader(
dataset,
batch_size=config.batch_size,
sampler=sampler,
num_workers=config.num_workers,
pin_memory=True,
persistent_workers=True
)
return loader, sampler
def train_epoch(model, loader, criterion, optimizer, scaler, config, rank):
"""1 エポックのトレーニング"""
model.train()
total_loss = 0
correct = 0
total = 0
for inputs, labels in loader:
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
if config.use_amp and scaler is not None:
# 混合精度トレーニング
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
scaler.step(optimizer)
scaler.update()
else:
# 通常のトレーニング
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
optimizer.step()
total_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()
config = TrainingConfig()
config.local_rank = args.local_rank
rank = init_distributed(config)
world_size = dist.get_world_size()
print(f"プロセス {rank}/{world_size} の初期化が完了しました")
# モデルを作成
model = YourModel().cuda(rank)
model = DDP(model, device_ids=[rank])
# オプティマイザ
optimizer = optim.Adam(model.parameters(), lr=config.lr)
# GradScaler（混合精度を使用する場合）
scaler = GradScaler() if config.use_amp else None
# データセット（あなたのデータセットに置き換えてください）
dataset = YourDataset()
loader, sampler = create_dataloader(dataset, config, world_size, rank)
criterion = nn.CrossEntropyLoss()
# トレーニングループ
for epoch in range(config.epochs):
sampler.set_epoch(epoch)
loss, acc = train_epoch(model, loader, criterion, optimizer, scaler, config, rank)
# rank 0 でのみ出力
if rank == 0:
print(f"Epoch {epoch+1}/{config.epochs}: Loss={loss:.4f}, Acc={acc:.4f}")
dist.destroy_process_group()
# 起動方法：
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 train.py
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 train.py
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import argparse
import os
class TrainingConfig:
def __init__(self):
self.local_rank = -1
self.epochs = 30
self.batch_size = 32
self.lr = 1e-3
self.num_workers = 4
self.gradient_clip = 1.0
self.use_amp = True
def init_distributed(config):
"""分散環境の初期化"""
# 環境変数から構成を取得
local_rank = config.local_rank
# NCCL 設定
os.environ["NCCL_DEBUG"] = "WARN"
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")
return local_rank
def create_dataloader(dataset, config, world_size, rank):
"""データローダーの作成"""
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True,
drop_last=True
)
loader = DataLoader(
dataset,
batch_size=config.batch_size,
sampler=sampler,
num_workers=config.num_workers,
pin_memory=True,
persistent_workers=True
)
return loader, sampler
def train_epoch(model, loader, criterion, optimizer, scaler, config, rank):
"""1 エポックのトレーニング"""
model.train()
total_loss = 0
correct = 0
total = 0
for inputs, labels in loader:
inputs = inputs.cuda(rank, non_blocking=True)
labels = labels.cuda(rank, non_blocking=True)
optimizer.zero_grad()
if config.use_amp and scaler is not None:
# 混合精度トレーニング
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
scaler.step(optimizer)
scaler.update()
else:
# 通常のトレーニング
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
optimizer.step()
total_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()
config = TrainingConfig()
config.local_rank = args.local_rank
rank = init_distributed(config)
world_size = dist.get_world_size()
print(f"プロセス {rank}/{world_size} の初期化が完了しました")
# モデルを作成
model = YourModel().cuda(rank)
model = DDP(model, device_ids=[rank])
# オプティマイザ
optimizer = optim.Adam(model.parameters(), lr=config.lr)
# GradScaler（混合精度を使用する場合）
scaler = GradScaler() if config.use_amp else None
# データセット（あなたのデータセットに置き換えてください）
dataset = YourDataset()
loader, sampler = create_dataloader(dataset, config, world_size, rank)
criterion = nn.CrossEntropyLoss()
# トレーニングループ
for epoch in range(config.epochs):
sampler.set_epoch(epoch)
loss, acc = train_epoch(model, loader, criterion, optimizer, scaler, config, rank)
# rank 0 でのみ出力
if rank == 0:
print(f"Epoch {epoch+1}/{config.epochs}: Loss={loss:.4f}, Acc={acc:.4f}")
dist.destroy_process_group()
# 起動方法：
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 train.py
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 train.py
分散トレーニングの鍵は、1）適切な並列モード（データ並列 vs モデル並列）を選択すること、2）より良いパフォーマンスを得るために DataParallel の代わりに DDP を使用すること、3）混合精度トレーニングを組み合わせてさらに高速化すること、4）データローダーと通信パラメータを適切に設定することです。
7. 分散トレーニングのモニタリングとデバッグ
7.1 よく使う監視ツール
例
# 分散トレーニング監視コード
def monitor_training(rank, world_size):
"""分散トレーニングの状態を監視"""
# 1. プロセスグループの状態を確認
print(f"Rank: {rank}, World Size: {world_size}")
print(f"Backend: {dist.get_backend()}")
print(f"Rank: {dist.get_rank()}")
print(f"CUDA Device: {torch.cuda.current_device()}")
# 2. GPUメモリ使用量を監視
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
allocated = torch.cuda.memory_allocated(i) / 1024**2
reserved = torch.cuda.memory_reserved(i) / 1024**2
print(fGPU {i}: 割り当て済み {allocated:.1f} MB, 予約済み {reserved:.1f} MB)
# 3. 勾配同期の監視
# DDP は自動的に同期時間を記録します
# 環境変数を設定することで詳細ログを有効にできます
# os.environ["NCCL_DEBUG"] = "INFO"
# トレーニングループ内で定期的に呼び出します
monitor_training(rank, dist.get_world_size())
def monitor_training(rank, world_size):
"""分散トレーニングの状態を監視"""
# 1. プロセスグループの状態を確認
print(f"Rank: {rank}, World Size: {world_size}")
print(f"Backend: {dist.get_backend()}")
print(f"Rank: {dist.get_rank()}")
print(f"CUDA Device: {torch.cuda.current_device()}")
# 2. GPUメモリ使用量を監視
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
allocated = torch.cuda.memory_allocated(i) / 1024**2
reserved = torch.cuda.memory_reserved(i) / 1024**2
print(fGPU {i}: 割り当て済み {allocated:.1f} MB, 予約済み {reserved:.1f} MB)
# 3. 勾配同期の監視
# DDP は自動的に同期時間を記録します
# 環境変数を設定することで詳細ログを有効にできます
# os.environ["NCCL_DEBUG"] = "INFO"
# トレーニングループ内で定期的に呼び出します
monitor_training(rank, dist.get_world_size())
7.2 デバッグのヒント
- まず単一マシンでテストし、その後複数マシンに拡張します
- 使用
os.environ["NCCL_DEBUG"] = "INFO"詳細な通信ログを表示する - すべてのマシンの時刻同期を確保する（NTPを使用）
- ファイアウォールがNCCL通信ポートを許可しているか確認する
その他の拡張分散トレーニングは複雑さを増しますが、もたらす性能向上は顕著です。大規模モデルのトレーニングには、ほぼ必須の技術手段です。