[論文レビュー] Memory-Efficient Pipeline-Parallel DNN Training

Deepak Narayanan, Amar Phanishayee+3 more|arXiv (Cornell University)|Jun 16, 2020

Topic Modeling参考文献 22被引用数 62

ひとこと要約

PipeDream-2BW はメモリ効率の高いパイプライン並列性を導入し、ダブルバッファ付きの重み更新と自動パーティショニングを可能にする。巨大なトランスフォーマーモデルを同等の精度でより速く訓練できる。

ABSTRACT

Many state-of-the-art ML results have been obtained by scaling up the number of parameters in existing models. However, parameters and activations for such large models often do not fit in the memory of a single accelerator device; this means that it is necessary to distribute training of large models over multiple accelerators. In this work, we propose PipeDream-2BW, a system that supports memory-efficient pipeline parallelism. PipeDream-2BW uses a novel pipelining and weight gradient coalescing strategy, combined with the double buffering of weights, to ensure high throughput, low memory footprint, and weight update semantics similar to data parallelism. In addition, PipeDream-2BW automatically partitions the model over the available hardware resources, while respecting hardware constraints such as memory capacities of accelerators and interconnect topologies. PipeDream-2BW can accelerate the training of large GPT and BERT language models by up to 20$\times$ with similar final model accuracy.

研究の動機と目的

• アクセラレータのメモリ制約の下で十億パラメータ級のDNNを訓練する課題を動機づける。
• フラッシュ誘発の停止なしにスループットを向上させる、メモリ効率の高いパイプライン並列システムを提案する。
• デバイスメモリ制約を尊重しつつ異種インターコネクトにまたがって自動的にモデルをパーティショニングする。
• 大規模トランスフォーマーモデル（GPT/BERT）における収束とスループットの利点をベースラインと比較して示す。

提案手法

• 2BW（ダブルバッファ付き重み更新）を導入して、二つの重みバージョンを保持しマイクロバッチを積算することでメモリ使用量を削減する。
• 計画アルゴリズムを用いてDNN演算子を等分散に複製されたステージとしてパーティションし、マイクロバッチと勾配蓄積を割り当ててスループットとメモリを均等に平衡化する。
• 単一の重みバージョンでパイプラインフラッシュを行う PipeDream-Flush という変種を提供し、少ないメモリフットプリントを実現するがスループットコストを伴う。
• 1ステップ遅延での重み更新セマンティクスを定式化: W(t+1)=W(t)−η·∇f(W(t−1)) で、通常の更新と収束性の類似性を議論する。
• 活性化再計算を取り入れて活性化の保存を減らし、各GPUで実現可能なマイクロバッチサイズを増やす。
• AWS上の PyTorch で NVLink を用いた大規模 GPT/BERT モデルを評価し、Megatronスタイルのテンソルモデル並列性および GPipe と比較。

実験結果

リサーチクエスチョン

• RQ1 パイプライン並列訓練で十億パラメータ級モデルのメモリとスループットをどうバランスさせるか？
• RQ2 ダブルバッファ重み更新スキームはパイプラインフラッシュなしで高スループットと許容されるメモリフットプリントを達成できるか？
• RQ3 自動パーティショニング戦略は、収束を保ちながら異種インターコネクト上にトランスフォーマー風アーキテクチャを最適にマップできるか？
• RQ4 2BWとPipeDream-Flush のセマンティクスは、NLP大規模モデルでの vanilla オプティマの精度と収束性を保つか？

主な発見

• 最適化されたモデルパラレル baselines に対して、さまざまな GPT モデルでエンドツーエンドの速度向上を1.3xから20x達成。
• パイプラインフラッシュが不要でメモリ使用量が低いためGPipeより最大3.2x高速。
• 64台の8-GPUまたは16-GPU構成を用いて、Megatronスタイルの構成をメモリ制約下で上回る、約300億パラメータ級の大規模トランスフォーマーモデルを訓練可能。
• BERTとGPTの事前訓練タスク全体で、Vanilla Adamと比較して2BWは同様の訓練損失曲線と下流のファインチューニング精度を示す。
• PipeDream-Flush はスループットを犠牲にしてメモリフットプリントを低くする一方、2BW は二つの重みバージョンと活性化再計算で高いスループットを実現。
• メモリフットプリントの比較はGPipe がメモリ課題、PipeDream-Flush が中程度、PipeDream-2BW が中間のフットプリントを示す。