[論文レビュー] torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models

Chiheon Kim, Heungsub Lee+6 more|arXiv (Cornell University)|Apr 21, 2020

Advanced Neural Network Applications参考文献 26被引用数 32

ひとこと要約

GPytorchライブラリを用いた PyTorch GPipe風のマイクロバッチパイプライン並列性とチェックポイント機構を実装し、イーガー実行環境で非常に大規模なモデルを効率的に訓練することを可能にします。

ABSTRACT

We design and implement a ready-to-use library in PyTorch for performing micro-batch pipeline parallelism with checkpointing proposed by GPipe (Huang et al., 2019). In particular, we develop a set of design components to enable pipeline-parallel gradient computation in PyTorch's define-by-run and eager execution environment. We show that each component is necessary to fully benefit from pipeline parallelism in such environment, and demonstrate the efficiency of the library by applying it to various network architectures including AmoebaNet-D and U-Net. Our library is available at https://github.com/kakaobrain/torchgpipe .

研究の動機と目的

• 単一デバイスの能力を超える巨大なニューラルネットワークの訓練を動機づけ、実現する。
• チェックポイント機構を備えたGPipe風のパイプライン並列性を実現する、すぐに使える PyTorch ライブラリを提供する。
• define-by-run/eager PyTorch における効率的なパイプライン並列性を達成するために、いくつかの最適化コンポーネントの必要性を示す。
• AmoebaNet-D や U-Net のようなアーキテクチャで性能向上を示す。
• パイプライン並列性における非逐次モデルとスキップ接続の取り扱いを論じる。

提案手法

• ネットワークを、パラメータ集合が互いに素な分割の連続として表現する。
• メモリ使用量を削減するために、勾配チェックポイントを用いたマイクロバッチ・パイプライン並列性を適用する。
• デバイス間でフォワードタスクをスケジュールするための決定論的なクロックサイクルを導入する。
• フォークとジョインによる後方の依存関係を提供し、バックパスの自動微分グラフを正しく強制する。
• デフォルト以外の CUDA ストリームを使用してコピーと計算の同時実行を可能にする。
• スキップ接続とチェックポイント再計算のために、ポータルと共有メモリ自動微分関数を導入する。

実験結果

リサーチクエスチョン

• RQ1 GPipe風のパイプライン並列性を PyTorch の define-by-run/eager モードでどのように効率的に実現できるか？
• RQ2 パイプライン並列性を用いて巨大モデルを訓練する際、ハイストループットとメモリ効率を達成するために不可欠なコンポーネントは何か？
• RQ3 非逐次モデルやスキップ接続をパイプライン並列訓練にどのように適合させられるか？
• RQ4 提案されたコンポーネントがスループット、デバイス利用率、メモリ使用量に与える影響は何か？
• RQ5 torchgpipe は PyTorch で AmoebaNet-D や U-Net のような大規模アーキテクチャにおいて GPipe と競争力を発揮できるか？

主な発見

• 各最適化コンポーネント（決定論的クロックサイクル、後方の Fork/Join 依存関係、非デフォルトストリーム、ポータル）は、測定可能な速度向上を提供する。
• 組み合わせたコンポーネントは、報告された U-Net 実験でベースラインに対してほぼ2倍の速度向上をもたらす。
• 非デフォルトストリームはコピーと計算の同時実行を可能にし、利用率を向上させる。
• スキップ接続が存在する場合、ポータルは不要なコピーを削減し、メモリ圧力を低減し、タイムラインの効率を向上させる。
• AmoebaNet-D と U‑Net のベンチマークは、PyTorch における効果的なパイプライン並列性と一致するスループットとメモリ特性を示している。