/ #ai #deep learning 

PyTorch + ABCI によるマルチノード分散学習のベストプラクティス

本記事では、産総研が提供するGPUクラウド計算基盤「ABCI」を用いて分散学習を行うためのシンプルな方法をまとめます。 以下のレポジトリに、マルチノード学習に対応した ABCI 向け学習コードの最小サンプルを公開しています。

https://github.com/yukara-ikemiya/abci-code-sample

大規模モデル学習に ABCI を活用する

AI Bridging Cloud Infrastructure (ABCI) は、国立研究開発法人 産業技術総合研究所（産総研）が構築・運用する世界初のオープンな人工知能処理向け大規模クラウド計算基盤です。 (https://abci.ai/ja/about_abci/)

→ ABCI利用ガイド

Python 環境の構築

ABCI では、学習環境の構築に Singularity コンテナの利用が推奨されています。Singularity コンテナ（SIFファイル）をビルドする方法はいくつかありますが、上記レポジトリではまず Docker イメージを作成し、それを SIF ファイルに変換するスクリプトを提供しています。

SIF ファイルを作成したら、ABCI 上の任意のディレクトリに配置して学習実行時に読み込みます。

学習コードの準備

PyTorch で分散学習を行う際、一般的には DistributedDataParallel (DDP) を用いてコードを書きますが、ここでは DDP をラップしてよりシンプルに実装できる HuggingFace Accelerate の活用を推奨します。コードが簡潔になるだけでなく、バグの混入も防ぎやすくなります。

HuggingFace Accelerate

モデル（最小限のオートエンコーダ）

import torch
from torch import nn

class SimpleModule(nn.Module):
    def __init__(self, dim_in, dim_hidden):
        super().__init__()

        self.dim_in = dim_in
        self.dim_hidden = dim_hidden

        self.net = nn.Sequential(
            nn.Linear(self.dim_in, self.dim_hidden),
            nn.Linear(self.dim_hidden, self.dim_in)
        )

    def forward(self, x):
        return self.net(x)

トイデータセット

import numpy as np
import torch
from torch.utils.data import Dataset


class DummyDataset(Dataset):
    def __init__(self, dim:int, num_data:int=10000):
        super().__init__()
        self.dim = dim
        self.num_data = num_data
    
    def get_item(self, idx):
        data = np.linspace(0, idx, self.dim)
        return torch.from_numpy(data.astype(np.float32))

    def __len__(self):
        return self.num_data

    def __getitem__(self, idx):
        return self.get_item(idx)

学習コード

import argparse

import torch
from torch.utils.data import DataLoader
from accelerate import Accelerator

from simple_module import SimpleModule
from dummy_dataset import DummyDataset

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--bs", type=int, default=50, help="batch size")
    parser.add_argument("--lr", type=float, default=0.0001, help="learning rate")
    parser.add_argument("--amp", type=str, default='fp16', help="autmatic mixed precision")
    return parser.parse_args()

def main():
    args = get_args()

    # Accelerator の初期化
    accelerator = Accelerator(mixed_precision=args.amp, split_batches=True)

    # モデル
    model = SimpleModule(dim_in=1000, dim_hidden=100)

    # データセット
    num_data = 10000
    dataset = DummyDataset(model.dim_in, num_data)
    dataloader = DataLoader(dataset, batch_size=args.bs, num_workers=4,
                            pin_memory=True, persistent_workers=True, shuffle=True)

    # オプティマイザ
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=[0.0, 0.99])

    # 分散学習用に準備
    model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

    for idx_e in range(200):
        loss_epoch = 0.
        for idx_d, x in enumerate(dataloader):
            # forward
            optimizer.zero_grad(set_to_none=True)
            y = model(x)

            # RMSE loss
            loss = ((x - y) ** 2).mean().sqrt()

            # backward
            accelerator.backward(loss)
            optimizer.step()

            loss_epoch += loss.detach()

        if accelerator.is_main_process:
            loss_epoch /= idx_d + 1
            print(f'Epoch {idx_e+1} : {loss_epoch}')

if __name__ == '__main__':
    main()

実行スクリプトの準備

HuggingFace Accelerator が内部で参照する環境変数が自動取得されない場合があります。 これに対処するため、以下のように Python コマンドをラッパー bash でラップし、OpenMPI の環境変数をコピーして利用します。

python.bash

#!/bin/bash

# HuggingFace Accelerate 用の環境変数を定義・コピー
export WORLD_SIZE=$OMPI_COMM_WORLD_SIZE
export RANK=$OMPI_COMM_WORLD_RANK
export LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK
export MASTER_PORT=11111
exec python "$@"

ノードあたりの GPU 数などの必要な設定の多くはスクリプトで自動検出できます。

train.bash

#!/bin/bash
#$-cwd

# モジュールの読み込み（ABCIの環境に応じて変更可能）
source /etc/profile.d/modules.sh
module load hpcx/2.12
module load singularitypro/3.11
module load cuda/11.6/11.6.2
module load nccl/2.11/2.11.4-1

# Singularity コンテナのパス
CONTAINER_PATH="/path/to/your/container.sif"

# ジョブID
JOB_NAME=$JOB_ID

# GPU タイプ（V100 または A100）の判定
GPU_INFO=$(nvidia-smi --query-gpu=gpu_name --format=csv)
if [[ $GPU_INFO =~ "V100" ]]; then
    NUM_GPUS_PER_NODE=4
elif [[ $GPU_INFO =~ "A100" ]]; then
    NUM_GPUS_PER_NODE=8
else
    readonly PROC_ID=$!
    kill ${PROC_ID}
fi

# GPU数の取得
GPUS_IN_ONE_NODE=$(nvidia-smi --list-gpus | wc -l)
NUM_GPU=$(expr ${NHOSTS} \* ${GPUS_IN_ONE_NODE})
echo "NUM_GPU = ${NUM_GPU}"

# MPI オプション
MPIOPTS="-np $NUM_GPU -N ${NUM_GPUS_PER_NODE} -x MASTER_ADDR=${HOSTNAME} -hostfile $SGE_JOB_HOSTLIST"

# マウント対象のディレクトリ
ROOT_SRC="/path/to/your/source/codes/"

# 設定値
batch_size=256 # 16 の倍数（2ノード構成時）
learning_rate=0.0001
amp=fp16

# 実行
mpirun $MPIOPTS \
    singularity exec --nv --pwd ${ROOT_SRC}/src/ -B ${ROOT_SRC} \
    ${CONTAINER_PATH} \
    ${ROOT_SRC}/job/python.bash ${ROOT_SRC}/src/train.py \
    --bs ${batch_size} \
    --lr ${learning_rate} \
    --amp ${amp}

最後に、ABCI 上のマルチノードでジョブを投入して実行します。

_run_job_on_abci_nodes.bash

# 2ノードでの学習実行ジョブ投入
qsub -j y -g gce12345 -l rt_AF=2 -l h_rt=0:30:00 ./job/train.bash