分散ディープラーニングの実装において、異なるGPUを使用するためのフレームワーク統合が重要です。特にTensorFlowやPyTorchなどの多様なフレームワークを一つのプラットフォームで利用できるHorovodは便利です。この記事では、Horovodを使った分散環境のセットアップ方法について説明します。

1. 分散学習アーキテクチャ: PSとRing-Allreduce

Parameter Server (PS) アーキテクチャ PSアーキテクチャでは、クラスタ内のノードはパラメータサーバーとワーカーに分けられます。パラメータサーバーがモデルのパラメータを管理し、ワーカーが勾配を計算します。しかし、通信の不均衡が問題となる場合があります。

Ring-Allreduce アーキテクチャ Ring-Allreduceでは、各デバイスがリング状につながり、中央集約点がありません。各ワーカーがミニバッチ訓練を行い、他のワーカーに勾配を伝達します。これにより帯域幅の最適化が可能で、並列性が向上します。

2. 環境準備

Ubuntu 18.04、NVIDIAドライバ、CUDA、cuDNN、PyTorch、TensorFlowをインストール。 SSHによるパスワード不要ログイン設定とNFSファイル共有システムの構築。 OpenMPIとNCCLの導入。 Horovodのインストールと設定。

注意点として、各ツールのバージョン互換性を確認することが重要です。

import argparse
import torch.nn as nn
from torchvision import datasets, transforms
import horovod.torch as hvd

# 学習設定
parser = argparse.ArgumentParser(description='分散MNISTトレーニング')
parser.add_argument('--batch-size', type=int, default=64)
parser.add_argument('--epochs', type=int, default=10)
args = parser.parse_args()

hvd.init()
torch.manual_seed(args.seed)

if torch.cuda.is_available():
    torch.cuda.set_device(hvd.local_rank())
    torch.cuda.manual_seed(args.seed)

train_dataset = datasets.MNIST('/data/', train=True, download=True,
                               transform=transforms.Compose([
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.1307,), (0.3081,))
                               ]))
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler)

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.fc1 = nn.Linear(160, 10)
    
    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = x.view(-1, 160)
        x = self.fc1(x)
        return F.log_softmax(x)

model = SimpleNet()

if torch.cuda.is_available():
    model.cuda()

optimizer = optim.SGD(model.parameters(), lr=0.01 * hvd.size())

hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

def train(epoch):
    model.train()
    for data, target in train_loader:
        if torch.cuda.is_available():
            data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

for epoch in range(1, args.epochs + 1):
    train(epoch)

上記スクリプトの実行コマンド：

nohup mpirun -np 2 -H node_1:1,node_2:1 -bind-to none -map-by slot -x NCCL_DEBUG=INFO python /home/user/mnist_train.py > log.txt 2>&1 &