DDP(Distributed Data Parallel)による学習時間の高速化を確認してみる

はじめに

こんにちは．AI技術開発部の小幡です．私はPyTorch[1]の学習コードに，PyTorchLightning[2]を利用しています．たまに内部の挙動に戸惑うこともありますが，学習コードをきれいにラップしてくれるので，重宝しています．
このきれいにラップしてくれるという点に関連して，PyTorchのGPU並列化設定について触れていきます．代表的な設定※にDataParallel(DP)とDistributedDataParallel(DDP)があります．DDPに設定することで，学習の高速化が期待できます．
PyTorchLightningはこの設定変更を容易に行うことができます．DDP設定が楽だから，という理由でPyTorchLightningを使っているという方も多いのではないでしょうか．
しかし，実際にDDPのドキュメントを見てみると，

This tutorial is a gentle introduction to PyTorch DistributedDataParallel (DDP) which enables data parallel training in PyTorch. Data parallelism is a way to process multiple data batches across multiple devices simultaneously to achieve better performance. In PyTorch, the DistributedSampler ensures each device gets a non-overlapping input batch. 

https://pytorch.org/tutorials/beginner/ddp_series_theory.html

とDDPはDPと比較してパフォーマンス良いとしか書いておらず，「実際どれくらい効果があるんだ？」と思ってしまいます．
そこで今回はDDPによる高速化度合いをいくつかのモデルで確認してみることにします．

※PyTorchLightningでは，TPU向けの設定などその他の設定があり，strategyとして簡単に設定することができます．

DPとDDP

DataParallel(DP)とDistributedDataParallel(DDP)はPyTorchで学習を行う際の，GPUの挙動を設定するものです．下記がPyTorchチュートリアル資料[3]からの抜粋です．

DataParallel is an older approach to data parallelism. DP is trivially simple (with just one extra line of code) but it is much less performant. DDP improves upon the architecture in a few ways:

https://pytorch.org/tutorials/beginner/ddp_series_theory.html

First, DataParallel is single-process, multi-thread, and only works on a single machine, while DistributedDataParallel is multi-process and works for both single- and multi- machine training. DataParallel is usually slower than DistributedDataParallel even on a single machine due to GIL contention across threads, per-iteration replicated model,

https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

上記の通り，DPはシンプルな設定ですが，DDPのほうがパフォーマンスが良いと記載があります．DPとDDPについて解説をすると，それだけで1つの記事になってしまうので，ここでは割愛します．
公式ドキュメント[3]内のDDPに関する部分と解説記事[4],[5]を記載しておくので，ぜひ読んでみてください．

実験条件

環境

ハードウェア

ソフトウェア

学習設定

バッチサイズはモデルごとに最大値が異なりますが，バッチサイズは8で固定しています．
バッチサイズ8はモデルによってはかなり低めの値です．しかし，ViTでバッチサイズ16にすると私のPCが落ちてしまうので，今回は8で統一しました．今回はモデル間の比較ではなく，モデル内のDPとDDPの相対比較が目的なので，問題ありません．
また，Classificationタスクにしたのは，torchvisionのモデルが豊富だったためです．
CIFAR10を利用したのは，データセットの準備がコード内で完結するためです．一方で，画像サイズが小さすぎてViTで学習できないため，resizeで拡大しています．
また，実験を早くまわすため，limit_train_batches=0.05として，データセットの5%のみ学習に用います．こちらも，DPとDDPの相対比較が目的ですので，問題ありません．

モデル

モデルはtorchvision[6]からいくつかピックアップしました．
各種モデルの条件等を下記に列挙します．
学習済みの重みは利用していません．また，パラメータ数はPyTorchLightningで標準出力から確認しています．

測定方法

学習時間の測定には，PyTorchLightningのSimpleProfilerを用います．
PyTorchにもProfilerは存在していますが，そちらは実際の学習時間を計測するわけではないようです．
下記がPyTorchLightningの公式ドキュメント[7]から引用した注記です．

When using the PyTorch Profiler, wall clock time will not not be representative of the true wall clock time. This is due to forcing profiled operations to be measured synchronously, when many CUDA ops happen asynchronously. It is recommended to use this Profiler to find bottlenecks/breakdowns, however for end to end wall clock time use the SimpleProfiler.

https://pytorch-lightning.readthedocs.io/en/1.6.2/advanced/profiler.html

今回は，これに従いSimplerProfilerを使って計測していきます．

コード

ここまでの実験条件をコードにまとめると，下記のようになります．
今回は，epochあたりの学習時間のみを確認したいので，lossや精度は確認しませんでした．
実際，データsyncの関係で学習の挙動は変わるはずですが，今回は確認しませんでした．

import os
import pytorch_lightning as pl
from pytorch_lightning.profiler import SimpleProfiler
from torchvision.models import convnext_tiny, resnet18, mobilenet_v3_small, swin_v2_t
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR10

DIRPATH = '/data'

IMAGE_SIZE = [224, 224]
MAX_EPOCHS = 3


class BaseModel(pl.LightningModule):

    def __init__(self, model):
        super().__init__()
        self.model = model(num_classes=10, weights=None)
        self.loss = torch.nn.CrossEntropyLoss()

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

    def training_step(self, batch):
        x, y = batch
        logits = self(x)
        loss = self.loss(logits, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)


def main(model, strategy, filename, batch_size):

    train_dataset = CIFAR10(
        DIRPATH,
        train=True,
        download=True,
        transform=transforms.Compose(
            [
                transforms.Resize(size=IMAGE_SIZE),
                transforms.ToTensor()
            ]
        )
    )
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        num_workers=os.cpu_count(),
        pin_memory=True,
        persistent_workers=True
    )
    trainer = pl.Trainer(
        accelerator='gpu',
        devices=-1,
        max_epochs=MAX_EPOCHS,
        limit_train_batches=0.05,
        profiler=SimpleProfiler(dirpath=DIRPATH, filename=filename),
        strategy=strategy
    )

    trainer.fit(model, train_loader)


if __name__ == "__main__":

   
    strategies = ['dp', 'ddp']
    batch_size = 8

    # mobilenetの実行例
    for strategy in strategies:
        main(BaseModel(mobilenet_v3_small), strategy, filename, batch_size)

実験結果　

実験結果を表にまとめました．
学習時間/epochは, SimpleProfilerの`run_training_epoch`の`run_training_epoch(s)`の値です．
実験設定からも分かる通り，DPとDDPの学習時間はあくまで参考値になります．
また，DDPの場合は2GPUそれぞれのprofilerが出力されますが，ほとんど差がなかったので，片方のGPUの値のみ記載しています．

すべてのモデルでおおよそ3倍程度の高速化が確認できました．
もちろん，実際には前処理などの学習時間に影響及ぼす他の要因があるので，ここまで高速化するかはわかりません．
ですが，PyTorchLightningで簡単にDDPに設定できるなら，しておいて損はなさそうです．

おわりに

今回は，DDPによる高速化効果をいくつかのモデルで確認してみました．
おおよそ3倍程度の高速化が確認でき，DDPの効果を実施に体感することができました．
また，PyTorchのprofilerなどにあまり触れる機会がない部分について知ることができ，個人的に良い機会でした．
この記事が，PyTorchを触る方々に有益な情報になれば幸いです．

コニカミノルタは画像IoTプラットフォームFORXAIを通じて，お客様やパートナー様との共創を加速させ，技術・ソリューションの提供により人間社会の進化に貢献してまいります．

引用

[1 ] https://pytorch.org/

[2] https://www.pytorchlightning.ai/

[3] https://pytorch.org/tutorials/

[4] https://www.telesens.co/2019/04/04/distributed-data-parallel-training-using-pytorch-on-aws/

[5] https://qiita.com/fam_taro/items/df6061b589c3ccf86089

[6] https://pytorch.org/vision/stable/index.html

[7] https://pytorch-lightning.readthedocs.io/en/stable/