マルチGPU学習（DataParallel, DistributedDataParallel）

PyTorchにおけるマルチGPU学習は、大規模なモデルやデータセットを効率的に扱うための重要な技術です。PyTorchでは主に以下の2つの方法でマルチGPU学習を実現します：

---

1. torch.nn.DataParallel（非推奨傾向）

概要

DataParallel は比較的簡単にマルチGPU化ができるインターフェースです。単一プロセスで複数のGPUにデータを自動で分割し、各GPUで並列計算を行い、最終的に結果を統合します。

使用方法

python
import torch
import torch.nn as nn

model = MyModel()
model = nn.DataParallel(model)
model = model.to('cuda')

特徴

• 簡単に既存のコードをマルチGPUに対応できる。

• 1プロセスで複数GPUを管理する。

• CPU→GPU間の同期やバッチ分割は内部で自動的に処理される。

問題点

• GPU間通信はCPU経由になるため、オーバーヘッドが大きい。

• スケーラビリティが限定的で、大規模分散には不向き。

• PyTorch開発チームは現在、DistributedDataParallel（DDP）の使用を推奨している。

---

2. torch.nn.parallel.DistributedDataParallel（推奨）

概要

DistributedDataParallel（DDP）は、複数のプロセスを用いて各GPUに1つのプロセスを割り当て、それぞれのプロセスがモデルの部分を処理します。より高効率かつスケーラブルな学習が可能です。

初期化の基本手順

python
import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.multiprocessing as mp

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)

def cleanup():
    dist.destroy_process_group()

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

    model = MyModel().to(rank)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    # 学習ループなど
    cleanup()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

特徴

• 各GPUに1プロセスを割り当てるプロセス並列。

• GPU間の通信はNCCL（NVIDIA Collective Communications Library）を使い高速に行われる。

• 精密な制御が可能で、分散学習（複数ノードでの訓練）にも対応。

• 大規模モデル・データにおける事実上の標準。

---

比較表

---

注意点とベストプラクティス

• DistributedDataParallelを使用する際は、DistributedSampler を併用してデータローダーにランクごとのデータを割り当てる必要があります。

• モデルのパラメータは必ず .to(rank) で各GPUに移動させてから DDP にラップする。

• DDPでは、各プロセスが独立した学習ループを持つため、ログ出力やモデル保存は特定ランク（通常はrank=0）でのみ行うように設計する必要があります。

---

まとめ

PyTorchでのマルチGPU学習にはDataParallelとDistributedDataParallelの2つの方法がありますが、性能・スケーラビリティの観点から現在はDistributedDataParallelの使用が推奨されています。実装はやや複雑ですが、大規模な学習を行う上での必須技術です。