GPUクラスターで分散学習

環境

• 親ノード＋GPUを搭載した子ノード x 20台
• OS：Rocky Linux 9.5
• ノード間通信：イーサネット

前提

今回はDistributed Data Parallel（DDP）をPyTorchを使って実装しました．

各ノードでCUDAが利用できる状態です．

また， torch.distributtedで推奨されているNCCLもインストールしています．

手順

環境変数設定

以下のパラメータを考慮し、スクリプト内や実行コマンドに指定します．

• MASTER_ADDR: マスターノードのIPアドレス
• MASTER_PORT: マスターノードが使うポート番号
• NODE_RANK: 各ノードのランク（マスターノードを0とし、他のノードを1, 2, ...と番号付け）
• WORLD_SIZE: クラスター全体のプロセス数（例: ノード数×プロセス数/ノード）

コードの準備

1. torch.distributedの初期化

各プロセスが通信できるよう、以下の情報をtorch.distributed.init_process_group()で指定します.
- backend（今回はnccl）
- init_method（通信方法、例えばtcp://IP:PORT）
- rank（プロセスのID）
- world_size（全プロセス数）

import os
import torch.distributed as dist

rank = int(os.environ['RANK'])
world_size = int(os.environ['WORLD_SIZE'])

dist.init_process_group(
    backend='nccl', 
    init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}", 
    rank=rank,
    world_size=world_size
)

2. データ分散

各プロセスにデータを分割して割り当てるためにtorch.utils.data.DistributedSamplerを使用します．
from torch.utils.data import DataLoader, DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=10, sampler=sampler)

3. モデルのDDPラップ

モデルをtorch.nn.parallel.DistributedDataParallelでラップします。
from torch.nn.parallel import DistributedDataParallel as DDP

model = SimpleModel().to(rank)
ddp_model = DDP(model, device_ids=[rank])

4. 訓練ループ

for epoch in range(5):
    sampler.set_epoch(epoch)  # シードを設定してデータシャッフル
    for data, target in dataloader:
        data, target = data.to(rank), target.to(rank)
        optimizer.zero_grad()
        output = ddp_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        print(f"Rank {rank}, Loss: {loss.item()}")

dist.destroy_process_group()

5. メイン

torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size, join=True)

実行

今回はスクリプトファイルを使って各ノードで実行するようにしました．

#!/bin/bash

# 各ノードの設定
# ノードのIPリストを自動生成
NODES=()
for i in $(seq -w 01 20); do
    NODES+=("node$i")
done
MASTER_ADDR="192.168.0.1"
MASTER_PORT=29500
WORLD_SIZE=20  # 全プロセス数（ノード数×プロセス数/ノード）
SCRIPT_PATH="/Path/to/train.py"

# 各ノードでコマンドを実行
for i in "${!NODES[@]}"; do
  NODE=${NODES[$i]}
  RANK=$i
  echo "Starting on $NODE with rank $RANK"

  ssh -tt $NODE "
    cd /home/takayama/Diffusion_model/DDPM_DDP/
    export MASTER_ADDR=$MASTER_ADDR 
    export MASTER_PORT=$MASTER_PORT
    export WORLD_SIZE=$WORLD_SIZE
    export RANK=$RANK
    export CUDA_VISIBLE_DEVICES=0
    torchrun --nproc_per_node=1 --nnodes=$WORLD_SIZE --node_rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT $SCRIPT_PATH
  " &
done

wait
echo "Finish"