Open MPI を用いて PyΤorch によるマルチノード分散学習を行う方法

目的

DistributedDataParallelを用いてマルチノード環境で機械学習を分散させる実装の際に色々つまづいたのでそのメモです。
この記事ではtorch.multiprocessは使いません。

概要

自分が主につまづいたのは以下の３点です
① Pythonスクリプト実行時に必要な環境変数
② 通信の初期化（dist.init_proecss_group）
③ ランクの扱い方

コード

dist_run.sh

#!/bin/sh

n_nodes={ノード数}
n_gpu_per_node={1ノードあたりのGPU数}

chmod +x ./dist_python.sh
mpirun -x LD_LIBRARY_PATH -n $((n_nodes*n_gpu_per_node)) -npernode $n_gpu_per_node --bind-to none ./dist_python.sh main.py {学習に必要なコマンドライン引数}  

dist_python.sh

#!/bin/sh

python $* \
    --rank ${OMPI_COMM_WORLD_RANK} \
    --local_rank ${OMPI_COMM_WORLD_LOCAL_RANK} \
    --world_size ${OMPI_COMM_WORLD_SIZE}

main.py

import torch
import torch.distributed as dist
from torchvision import datasets, transforms
from torch.nn.parallel import DistributedDataParallel as DDP
import os
import argparse

# コマンドライン引数
parser = argparse.ArgumentParser( )
parser.add_argument('--rank', type=int, default=0)
parser.add_argument('--local_rank', type=int, default=0)
parser.add_argument('--world_size', type=int, default=1)
parser.add_argument('--port', type=int, default=50000)
args = parser.parse_args()

if __name__ == '__main__':
def main():
    # 通信の初期化
    master_addr    = os.getenv('MASTER_ADDR', default='localhost')
    master_port    = os.getenv('MASTER_PORT', default=args.port)
    dist_url       = 'tcp://{}:{}'.format(master_addr, master_port)
    ngpus_per_node = torch.cuda.device_count()
    dist.init_process_group('nccl', init_method=dist_url, rank=args.local_rank, world_size=ngpus_per_node)
    print('tcp connected, URL: {}'.format(dist_url))

    # デバイスの指定
    device = torch.device('cuda', args.local_rank)
    print('In rank {}, device: [{}]'.format(args.rank, device))

    # 分散学習のためのDataLoader
    train_dataset = datasets.MNIST('./data',
                                   train=True,
                                   download=True,
                                   transform=transforms.ToTensor())
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=args.rank)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=args.batch_size,
                                               sampler=train_sampler)

    # モデルを宣言した後DDPに通す
    model = MyModel().to(device)
    model = DDP(model, device_ids=[args.local_rank])

    # ここに学習用のコードを書く

    dist.destroy_process_group()   #プロセスの破棄

dist_run.sh

大まかな流れとしては、dist_run.shを実行することによって、複数プロセスがたちあがり、各プロセスで./dist_python.sh main.pyが実行されます。

dist_python.sh

Pythonスクリプトを実行する際にコマンドライン引数としてOpen MPIの環境変数を各プロセスに渡しています。
mpirunを用いてPythonスクリプトを実行する際に、各プロセスに渡す必要のあるOpen MPIの環境変数は以下の3つだと思います。

• OMPI_COMM_WORLD_RANK（プロセスのランク）
• OMPI_COMM_WORLD_LOCAL_RANK（ノード内におけるプロセスのランク）
• OMPI_COMM_WORLD_SIZE（プロセスの総数）

これらの値はos.getenv等を用いて各プロセスがPythonスクリプト内で取得することも可能ですが、今回はコマンドライン引数として渡しました。

main.py

特筆すべき箇所について書きます。

コマンドライン引数

./dist_python.sh main.pyでOpen MPIの環境変数を受け取るために、argparserに rank local_rank world_sizeの３変数を設定する必要があります。

通信の初期化

分散学習を行う前に、通信に用いるホストアドレスとポート番号を指定し、それらをPyTorchに設定する必要があります。この時、dist.init_process_groupの引数であるrankにはローカルランクを、world_sizeには１ノードのGPU数を渡す必要があります。自分はここをrank=args.rank world_size=args.world_sizeとしてしまい、はまっていました。

デバイスの指定

args.local_rankを用いて各プロセスが自分のいるノード内のGPUを指定します。args.rankでGPUを指定するとエラーが起きます。

分散学習用のDataLoader

分散学習によって学習を高速化するためには、１イテレーションの間に各プロセスが異なるデータを処理する必要があります。そのために DistributedSampler を指定する必要があります。DistributedSamplerの引数ではrank=args.rank world_size=args.world_sizeとします。