サブロウ丸

サブロウ丸

主にプログラミングと数学

サーベイ: PipeDream: Generalized Pipeline Parallelism for DNN Training

https://dl.acm.org/doi/abs/10.1145/3341301.3359646?casa_token=L-sKQKrRoE4AAAAA%3AYKo9NPdnPyG6IouMN5jfTHTCYFAGORDxen32GKAteeSG-ROhqx_OX-hVOfuyHiVBXLLJH0RPujhFPEk

@inproceedings{narayanan2019pipedream,
  title={PipeDream: generalized pipeline parallelism for DNN training},
  author={Narayanan, Deepak and Harlap, Aaron and Phanishayee, Amar and Seshadri, Vivek and Devanur, Nikhil R and Ganger, Gregory R and Gibbons, Phillip B and Zaharia, Matei},
  booktitle={Proceedings of the 27th ACM Symposium on Operating Systems Principles},
  pages={1--15},
  year={2019}
}

paper:

背景

  • モデルの巨大化が進行中(自然言語処理モデルGPT-3は1750億のパラメタを持つ)
  • 単一のメモリにモデルが載らないケースも
  • データ並列はGPU数が多くなると通信オーバーヘッドが90%近くになり、性能向上が頭打ちになっている

どんなもの?

  • モデル並列型のパイプライン処理システムPipeDreamの開発

先行研究と比べてどこがすごい?

  • モデルの自動分割手法を提案
  • 1F1B-pipelineスケジューリングを提案、ハードウェア使用率を向上させた

参考

  • 最も単純なモデル分割(一つのバッチを順にworkerが処理する)

  • GPipe(複数の(マイクロ)バッチをパイプラインに流して集約)

技術や手法のキモはどこ

  • モデルを動的計画法により自動分割
    • 事前にsingle CPU/GPU上でプロファイルしたデータをもとに最適化
    • 計算機がhierarchicalトロポジであることを想定すると、straight-pipeline処理の中でdata-parallelを行った方が良い箇所も提案してくれる

動的計画法によるモデル分割の補足資料はこちら。 ただ、これだけ見ても分からないと思うので、元論文と照らし合わせたときに少しでも理解の足しになれば嬉しいです。

  • forwardとbackwardを交互に行う1F1Bスケジューリングを提案
    • mini-batchのforwardとbackwardを交互に, worker間で非同期に行う
    • (利点)pipelineのflushの回数を減らす、それによりハードウェア使用率が向上、通信時間をoverwrapできる
    • (欠点)モデルパラメタの同期を行わないので、workerが持つ最新のモデルパラメタは基本的にバラバラ
      • backward時に正確に勾配を計算するにはforwardで用いたモデルパラメタが必要だが、backwardを行うときにはすでに他のmini-batchによるモデル更新が行われている
      • forwardとbackwardが"あるタイミングにおけるモデルパラメタ"で確実に行われることを保証するためにモデルパラメタセットのストックを行う(weight stashingと呼んでいる)
        • メモリ使用量が増えそうだが、そもそもworkerはモデルの一部しか持っていないため、複数のパラメタセットを持っていたとしてもメモリ使用量はそれほどでもないという主張
  • PipeDreamもGPipeと同様に、micro-batchをパイプラインに流してmini-batch単位で同期を取ることも可能
    • 主張としては、GPipeよりもより早い段階でbackwardを行えるので、(backwardに必要な)活性化状態量や特徴マップを早く捨てることができてその点はメモリ効率が良い

どうやって有効だと検証した?

  • 画像分類、翻訳タスクを目的としてさまざまなモデルで訓練し
    • データ並列よりも少ない計算時間、少ないメモリフットプリントを達成
  • 競合のpipelineアプローチであるGPipeよりも1.7倍のスループットを達成 (GNMT-16モデルの学習 in 16GPU)
    • GpipeとPipeDreamと同じ分割方法 (PipeDreamは分割アルゴリズムも提案している)
    • GpipeはCluster-AとCluster-Bで55%、71%のスループット低下
    • pipeline flushが頻繁に行われることが主な原因?
    • しかし、これはあくまでスループットの比較であって、学習時間あたりのモデル性能がどうか、という話は報告されていない
  • また、メモリ使用量はDPとほぼ変わらず (パラメタをバッチごとに保存する必要があるため)

議論はある?

  • モデルの自動分割アルゴリズムについて
    • hierarchicalトポロジ上での設定に限定されており、多くの計算機性能、通信バンド幅を持つような非対称のシステムには適用できない
    • 計算量がモデルの層数Nの3乗であり、モデルをスケールさせると計算量が増大する(しかし、現実的な設定+少しの工夫により問題はなさそう)
  • 1F1Bスケジュールのために、モデルパラメタセットを最大pipline上でactiveなmini-batch数コ用意する必要があり、メモリをその分使用するので、モデル側のスケーリングの限界が早い
  • モデルの分割がうまくいかなくて、worker間の負荷にばらつきがある場合はスループット効率が悪くなり、非同期パラメタ更新のうまみが生かせず、勾配の陳腐化が発生している分だけ学習が遅くなるかも

次に読むべき論文は?

  • Huang, Yanping, et al. "Gpipe: Efficient training of giant neural networks using pipeline parallelism." Advances in neural information processing systems 32 (2019).

ちなみに...