【PyTorch】Distributed Data Parallel(DDP)の基本

2024.03.19

PyTorchで複数のGPU(マルチGPU)を使い学習を高速化する方法にDistributed Data Parallel(以下DDP)があります。本記事ではDDPの基本概念とチュートリアルに沿った実装について解説します。

DDPの基本的な概念

DDPはデータ並列（Data Parallel）という方法で学習を高速化する仕組みです。ごく簡単に言えば複数GPUで同時に学習を行うことで学習を高速化する分散学習という手法の1つになります。

複数GPUで学習を行うためにはtorch.multiprocessingモジュールなどを使って、学習プロセスを複数立ち上げることになります。

学習ではなく推論の場合、立ち上げたそれぞれのプロセスで違なるGPUにモデルを乗せて独立に推論を行えば良く、さほど難しくありません。（DDPを使わず実現可能です）

ではなぜDDPを使うのかと言うと、学習の場合にはそれぞれのGPU上のモデルのパラメータを同期する必要があるからです。複数のGPUで同期をせずそれぞれ独立に学習を行ってもそれはGPUの数分の別の学習結果が得られるだけになってしまいます。DDPを使うと同期を簡単に行うことができ、複数GPUのパワーを使って1つの学習結果を得ることが出来ます。（つまり高速化が可能）

より正確にはモデルのパラメータを同期するというより、各GPU上のモデルのパラメータが同一になるために必要な勾配の計算を自動的に行なってくれるというものになります。

実装の解説

チュートリアル

まずは公式のチュートリアルに沿う形で実装を始めることをおすすめします。特にDDPを初めての環境で動かす場合、動作確認も兼ねて一度チュートリアルコードを実行すると良いと思います。

Getting Started with Distributed Data Parallel — PyTorch Tutorials 2.2.0+cu121 documentation
pytorch.org

PyTorchは仕様変更も多々あるため、下手に二次情報を当たると上手くいかないということもあり得ます。

本記事では執筆時点、PyTorch 2.2.0でのチュートリアルを引用しつつ、私が実装時につまづいた点について補足していきます。

用語

チュートリアルコード上で変数名などに表れる以下の用語は重要なのでまず知っておきましょう。

  • rank: プロセス番号。0, 1, 2…という値が割り振られる前提。
  • world_size: 総プロセス数。

セットアップ

Python
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

DDPの開始、終了処理です。よく言われる「おまじない」のようなものですが、ここで重要なのはサンプル中に”gloo”となっているバックエンドの選択となります。

バックエンドによりパフォーマンスが変わるほか、環境により使えるものが異なり、Windowsでは”gloo”のみとなります。

パフォーマンスについてはいろいろ違いがあるそうですが、複数使える場合は実測で選ぶのが良いと思います。

なお、この処理はシングルGPU環境ではエラーになるためシングルGPU環境でも同じコードを動かす場合にはGPUの搭載数を見て処理をスキップするなどの工夫が必要になります。

モデルをDDP管理下に置く

Python
model = ToyModel().to(rank)
ddp_model = DDP(model, device_ids=)

DDPの要の部分です。このようにしてモデルを各GPU上に転送し、そのモデルをDDPの管理下に置くことでモデル間の同期が先ほど設定したバックエンド経由で行われることになります。

DataLoader

チュートリアルには記載がありませんが、実際の学習ではDataLoaderを準備する必要があります。各GPUに対してデータを分割して渡すには、DistributedSamplerを使います。

torch.utils.data — PyTorch 2.2 documentation
pytorch.org

サンプルにあるように、DataLoaderのsamplerに対してDistributedSamplerを渡してあげれば、各rankに対し分割されたデータを受け取ることが出来ます。

Python
sampler = DistributedSampler(dataset, rank=rank, shuffle=False)
loader = DataLoader(dataset, shuffle=(sampler is None), sampler=sampler)

DataLoader側のshuffleが(sampler is None)のように設定されていることがやや違和感があるかもしれませんが、shuffleはDistributedSamplerの側で設定するため、DataLoader側ではある種のお約束のようにこのような記法とするそうです。

同期ポイント

チュートリアルにも記載があるようにDDPではフォワードパス、バックワード等が自動的に同期ポイントになります。

In DDP, the constructor, the forward pass, and the backward pass are distributed synchronization points.

Skewed Processing Speeds より

これは1つのプロセスでDDP管理下に置いたモデルがフォワード計算をした場合、他のプロセスのフォワード計算が終わるまで、先に計算したプロセスは待つことを意味します。この仕組みにより、勝手に1つのプロセスだけ先に進んでしまうということを防ぐようになっています。

逆に1つのプロセスのみで計算をしたい場合は注意が必要です。普通に実装すると同期ポイントの計算をしたあと、他のプロセスの計算を待ち続けそこで停止していまいます。

その場合、他のプロセスでダミーの計算を回してあげるなどの対策を取る必要があります。学習をDDPで複数GPUで行った後、テストを1プロセスで行いたいなどの場合は注意しなければなりません。

また、同期ポイントはdist.barrier()関数にて手動で作ることも可能です。チュートリアルの例にあるように、1つのプロセスでモデルの保存など重い処理を行う場合に他のプロセスと進度を合わせるため適宜挿入する必要があります。

モデルの保存

基本的にはチュートリアルの通りで問題ありません。

Python
if rank == 0:
    # All processes should see same parameters as they all start from same
    # random parameters and gradients are synchronized in backward passes.
    # Therefore, saving it in one process is sufficient.
    torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)

コメントにあるように、DDPが上手く行っている場合には全GPUで同じモデルとなるため、1つのプロセスのみで保存すれば良いことになります。

注意点としては、モデルの保存時にGPU情報を除外するため以下のように一度CPUに移動してから保存するという方法を取る場合もあると思います。

Python
torch.save(ddp_model.to('cpu').state_dict(), CHECKPOINT_PATH)

ただこのようにモデルをCPUに移動するとDDPの管理下から外れてしまいます。学習の途中結果を一時保存してその先も学習を進める場合には以降GPU間での同期が取れなくなってしまいます。GPU上のままモデルを保存するか、再度DDPの設定を行う必要があります。

うまく動かない時の確認ポイント

基本的にはチュートリアルに沿って実装すれば問題ないはずですが、上手くいかない場合も多々あるかと思います。その場合はDDPの原理に沿って学習の各処理における値が正しいものになっているかをデバッガやログ等で順番に確認していきましょう。ポイントは、各GPU（プロセス）において、どのデータが同じで、どのデータが異なるか、ということです。

まずモデルの重みの初期値については、異なる値となっていて構いません。後ほどDDPにより同期されるためです。（同じ値で初期化しても大丈夫です。）

学習データについても、GPUにより違うものを入力します。データセットを各GPUに分割して高速に学習することが目的なので、同じデータが入力されていると意味がありません。これはDistributedSamplerを利用していれば自動的に実現されるはずです。

forward計算、loss計算についても異なるデータを入力するため、GPUごとに違う値となります。

次にbackward計算ですが、DDPではここでGPU間の同期が行われます。初回のbackwardに関しては各GPUでモデルが同じ重みになるような勾配が自動計算されます。2回目以降はモデルの重みが同期されるため、勾配も同じ値となります。

2回目以降でもGPU間で同じ位置の勾配を見ているのに違う値が入っていた場合、同期が上手くできていないことになります。DDPの基本的な設定部分を見直してみましょう。

最後にモデルの重みの更新ですが、初期重みは異なっていても1度目の重みの更新以降は重みがGPU間で共通となります。

なお勾配および重みの具体値は以下のようなコードで確認できます。

Python
# [0]の部分にはモデルのshapeに合わせて適切な値を入れてください
[x for x in ddp_model.parameters()][0].grad[0] # 勾配の確認
[x for x in ddp_model.parameters()][0][0] # 重みの確認

まとめると、以下のようになります。

学習ループ1回目学習ループ2回目以降
モデルの重み初期値GPUにより異なる–
学習データGPUにより異なるデータGPUにより異なるデータ
forward計算GPUにより異なる結果GPUにより異なる結果
lossGPUにより異なる結果GPUにより異なる結果
モデルの勾配GPU間で異なる結果GPU間で同じ結果
更新されたモデルの重みGPU間で同じ結果GPU間で同じ結果

マルチプロセス特有の注意点

DDPを利用したスクリプトは基本マルチプロセスになるため、排他処理などマルチプロセス特有の問題に気をつける必要があります。

DDPはGPU間のモデルの同期は行ってくれますが、それ以外のことはケアしてくれません。

例えば結果やログの出力など、ファイル出力周りは特に排他処理で気をつけなければならない部分です。ここはDDPの領域ではなく一般論としての排他処理に気をつけて実装をする必要があります。

コメント

タイトルとURLをコピーしました