pytorchでのmulti gpu対応 (DP, DDP, DeepSpeed, Accelerate)

最近SEANというnvidiaが出しているGANのモデルを勉強がてら実装してみたが、その中でDataParallelを使った処理があった。

DataParallelについてpytorchのドキュメントを見てみると、DistributedDataParallelを使えと書いてある。このあたりについて今までちゃんと調べたことはなかったため、良い機会ということで調べてメモ

DPとDDP

DataParallel(DP)

https://pytorch.org/docs/1.13/generated/torch.nn.DataParallel.html
ドキュメントにはこんなwarningが書いてある

It is recommended to use DistributedDataParallel, instead of this class, to do multi-GPU training, even if there is only a single node. See: Use nn.parallel.DistributedDataParallel instead of multiprocessing or nn.DataParallel and Distributed Data Parallel.

  • single gpuによる学習時と同じくプロセスが一つだけ
  • forward、backwardというgpuでの処理が必要な部分のみ各gpu上で実行されて結果がメインgpuへ戻される
  • モデルパラメータの更新はメインとなる一つのgpu上で実行されてから各gpuへ共有される

single gpuで学習する場合とほとんど変わらない形のイメージで、複数gpuの実行が必要な部分の情報共有や集約部分が拡張されたような感じ。マルチノードでの学習はできない。

  • pros
    • 実装が簡単
      • 元のmodelを model = DDP(model) のようにwrapすれば使える
  • cons
    • パラメータ共有やinput dataの扱いが非効率的
    • マルチノード学習は非対応

DistributedDataParallel(DDP)

https://pytorch.org/docs/1.13/notes/ddp.html
こちらにもwarningが書かれている

The implementation of torch.nn.parallel.DistributedDataParallel evolves over time. This design note is written based on the state as of v1.4.

実装はversoin毎に進化しているので注意とのこと

  • gpuの数だけプロセスが作られる
  • 各プロセスのdataloaderはプロセス毎に異なるデータを持つよう設定される
  • モデルパラメータの共有は最初の一回のみ
  • 各プロセス(=gpu)間でgradientが相互に送られて各々和を取ってパラメータ更新に使う

DPと比べ、最初から複数gpuを使って処理する前提で構築されているという印象を受ける形になった。わかりづらい部分としてはパラメータの共有に関する部分で、

  • 全gpuが同じ値で初期化されたパラメータを持つ(最初1回の共有)
  • 計算が進み、各gpuでgradientが得られる
    • 各gpuは異なるinputを与えて処理しているので当然gradientは異なる値
  • すべてのgpu間でgradientを共有
    • 例えばgpu:0であれば0以外の全てからgradientを受け取る、gpu:1であれば1以外の全てから...
    • 各gpu上で和を取れば、全gpuが同じgradientの和を共有しているため更新分は一致する
    • よって更新後のモデルパラメータも厳密に一致する

そのため、どれか一つのプロセスがデータをまとめて処理するような偏りは起こらないため、そこがボトルネックになるようなこともなく効率的

  • pros
    • 処理が効率的でボトルネックが発生しない
    • マルチノード学習に対応
  • cons
    • node addressの設定やsetup系の関数を呼び出すなどDDPための実装が必要
    • プロセスが別れているためDPに比べてメインメモリは多く消費する

ちなみにDPとDDPの比較については公式のtutorialページに書いてある
https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

処理フローをよりわかりやすく解説してくれていたqiita
https://qiita.com/fam_taro/items/df6061b589c3ccf86089

その他の実装について

DDP

上で書いた通りだが、マルチGPUの学習が行いたい場合はこれを使っておけば問題なく動かすことはできる。 ただし、最近はDDPをwrapしたフレームワークも出てきており、そちらを使っておく方がいいかもしれない。

DeepSpeed

近年の大規模モデルでは一つのgpuではモデル乗り切らずに学習が進められないものがある。そのような大きなモデルでも学習できるよう、各gpuが持つパラメータなどの値を効率化して省メモリ化して動かす方法としてZeROという論文が出ている https://arxiv.org/abs/1910.02054

それの実装としてpytorchのDDPをwrapする形になっているのがDeepSpeed
https://github.com/microsoft/DeepSpeed

gpuの消費メモリを抑えるために以下の3つをそのgpuで必要なものだけ持つ形するというもの

  • optimizer state (stage 1)
  • gradients (stage 2)
  • model parameters (stage 3)

optimizer stateとgradientsは各gpuで必要なもののみしか使う必要がないため、追加で通信して共有する必要もなく元のように全て持っておく場合に比べ通信のオーバーヘッド発生しない。なのでデフォルトでstage 2が使われる形となっている模様。stage 3のようにmodel parameterを削減した場合、step毎に他のgpuからモデルパラメータを追加で共有してもらう必要が出てくるため余計な通信が必要になる。

また、メインメモリへのoffloadなども可能。

1.4Bパラメータを超えるモデルだと、単純にDDPを適用するだけだとv100系のgpuでは学習不可能とのこと

runs out of memory for models with more than 1.4B parameters on current generation of GPUs with 32 GB memory.

  • pros
    • 巨大なモデルも学習可能
    • batch sizeを大きくできる
  • cons
    • DeepSpeedに則した実装が必要

FullyShardedDataParallel (FSDP)

https://pytorch.org/docs/1.13/fsdp.html

A wrapper for sharding Module parameters across data parallel workers. This is inspired by Xu et al. as well as the ZeRO Stage 3 from DeepSpeed

と書かれておりDeepSpeedと同様ZeROを実装したもの。DeepSpeedのpytorch公式版という感じか

tutorialページ
https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html

As of PyTorch 1.12, FSDP only offers limited support for shared parameters (for example, setting one Linear layer’s weight to another’s).

こういうことも書かれておりbeta版な印象

今後実装が進めばpytorchとの馴染みは当然よいはずなので、実装時の手間が少ないDeepSpeedといった感じになるかも？

Accelerate

huggingfaceが出した各種マルチノード学習の実装に対応したwrapperフレームワーク
https://huggingface.co/docs/accelerate/v0.18.0/en/index

以下のものに対応している

  • DDP (default)
  • Cloud TPU
  • DeepSpeed (beta)
  • pytorchのFSDP (beta)

見ての通り複数の実装の違いを吸収してくれるwraperとなっていて、accelerateで実装しておけば設定に応じて各種実装での実行を切り替えられる。

また、実装が大変簡易に済む点が大きい。DDPで動かす場合に必要だったsetup系の関数呼びだしなどは中でやってくれるため、accelerateのインスタンスを作ったらそれに対してbackwardなどを呼び出す形にするだけでOK。

DPと同等の実装の手間でDDPが使えるというイメージ。かつTPUも対応しているのでその辺を使いたい場合はだいぶ楽になりそう。

  • pros
    • 実装が楽
    • 複数の手法やデバイスが同じ実装で使える
  • cons
    • DeepSpeedやFSDP対応はbetaの状態で動かないケースがある

consに書いた点は注意が必要。例えばGANのような複数モデルを協調して動かす必要がある場合DeepSpeedのモードでは動かせなかった。ドキュメントにも今後破壊的なな変更入るかもしれないから注意、と書いてある。実装をやってみたのでこの点については別記事で書くことにする。

基本的にaccelerateを使うことで実装は楽にできるのでこれを使っておくのが手間という意味ではもっとも良さそうだった。少なくともDDPモードなら問題なくマルチノードで動かすことができた。