torchgpipe

torchgpipe

PytorchでのGPIPE実装。 TPUではなくCUDA向けに最適化されています。

 from torchgpipe import GPipe
model = nn . Sequential ( a , b , c , d )
model = GPipe ( model , balance = [ 1 , 1 , 1 , 1 ], chunks = 8 )
output = model ( input )

gpipeとは何ですか？

GPIPEは、Google Brainが発行したスケーラブルなパイプラインパラレズムライブラリであり、大規模なメモリを消費するモデルの効率的なトレーニングを可能にします。論文によると、GPIPEは8xデバイス（TPU）を使用して25倍の大きなモデルをトレーニングし、4xデバイスを使用してモデルを3.5倍高速にトレーニングできます。

GPIPE：パイプライン並列処理を使用した巨大なニューラルネットワークの効率的なトレーニング

GoogleはGPIPE上の557mパラメーターを備えたAmoebanet-Bをトレーニングしました。このモデルは、Imagenet分類ベンチマーク（2019年5月時点での最先端のパフォーマンス）で、84.3％のTOP-1および97.0％のトップ5精度を達成しました。

GPIPEは（a）パイプラインの並列性と（b）バックプロパゲーション中の順方向伝播の自動再構成を使用するため、大きなモデルのトレーニングを活用します。 （b）は、Pytorchコミュニティのよく知られている用語に従って、チェックポイントと呼びます。

パイプライン並列性

GPIPEはモデルを複数のパーティションに分割し、各パーティションを別のデバイスに配置して、より多くのメモリ容量を占有します。また、ミニバッチを複数のマイクロバッチに分割して、パーティションを可能な限り並行して機能させます。

チェックポイント

チェックポイントは各パーティションに適用され、モデルによるメモリ全体の消費を最小限に抑えます。前方伝播中、パーティション間の境界にあるテンソルのみが記憶されています。他のすべての中間テンソルは揮発され、必要に応じてバックプロパゲーション中に再計算されます。

使用法

現在、Torchgpipeには次の環境が必要です。

• Python 3.6+
• Pytorch 1.1+

Torchgpipeを使用するには、Pypi経由でインストールします。

$ pip install torchgpipe

gpipeでモジュールをトレーニングするには、 torchgpipe.GPipeでラップするだけです。 GPIPEが自動的にモジュールを連続したレイヤーを使用してパーティションに分割するため、モジュールはnn.Sequentialでなければなりません。 balance引数は、各パーティションのレイヤー数を決定します。 chunks引数は、マイクロバッチの数を指定します。入力、出力、および中間テンソルは、 TensorまたはTuple[Tensor, ...]でなければなりません。

以下の例コードは、4つのレイヤーを持つモジュールをそれぞれ4つのパーティションに分割する方法を示しています。このコードは、ミニバッチを8つのマイクロバッチに分割します。

 from torchgpipe import GPipe

model = nn . Sequential ( a , b , c , d )
model = GPipe ( model , balance = [ 1 , 1 , 1 , 1 ], chunks = 8 )

for input in data_loader :
    output = model ( input )

ドキュメント

API参照を含む詳細については、torchgpipe.readthedocs.ioにアクセスしてください。

ベンチマーク

詳細とその他のベンチマークは、torchgpipe.readthedocs.ioで入手できます。

RESNET-101精度ベンチマーク

バッチサイズ	torchgpipe	Nn.Datapar Allel	Goyal et al。
256	21.99±0.13	22.02±0.11	22.08±0.06
1k	22.24±0.19	22.04±0.24	n/a
4k	22.13±0.09	n/a	n/a

GPIPEは、追加のハイパーパラメーターチューニングを導入しないように透過的でなければなりません。透明性を検証するために、Goyal et al。

u-net（b、c）メモリベンチマーク

実験	u-net（b、c）	パラメーター	メモリの使用
ベースライン	（6、72）	362.2m	20.3ギブ
パイプライン-1	（11、128）	2.21b	20.5ギブ
パイプライン-2	（24、128）	4.99b	43.4ギブ
パイプライン-4	（24、160）	7.80b	79.1ギブ
パイプライン-8	（48、160）	15.82b	154.1ギブ

この表は、GPIPEがscaling u-netモデルをどのように促進するかを示しています。ベースラインは、パイプラインの並列化やチェックポイントのないベースラインを示し、パイプライン-1 、 -2 、 -4 、 -8は、対応する数のパーティションでモデルがGPIPEでトレーニングされていることを示します。

ここでは、簡略化されたU-Netアーキテクチャを使用しました。モデルのサイズは、それぞれレイヤーとフィルターの数に比例するハイパーパラメータBおよびCによって決定されます。

u-net（5、64）速度ベンチマーク

実験	スループット	スピードアップ
ベースライン	28.500/s	1×
パイプライン-1	24.456/s	0.858×
パイプライン-2	35.502/s	1.246×
パイプライン-4	67.042/s	2.352×
パイプライン-8	88.497/s	3.105×

スキップ接続で効率を確認するために、さまざまな数のデバイスでU-Netのスループットを測定しました。いくつかの長いスキップ接続があるため、U-Netを使用することを選択しました。

amoebanet-d（18、256）速度ベンチマーク

実験	スループット	torchgpipe	Huang et al。
n = 2、m = 1	26.733/s	1×	1×
n = 2、m = 4	41.133/s	1.546×	1.07×
n = 2、m = 32	47.386/s	1.780×	1.21×
n = 4、m = 1	26.827/s	1.006×	1.13×
n = 4、m = 4	44.543/s	1.680×	1.26×
n = 4、m = 32	72.412/s	2.711×	1.84×
n = 8、m = 1	24.918/s	0.932×	1.38×
n = 8、m = 4	70.065/s	2.625×	1.72×
n = 8、m = 32	132.413/s	4.966×	3.48×

（ N ：パーティションの数、 M ：マイクロバッチの数）

この表は、Huang et al。論文のKをnに置き換えたことに注意してください。

メモ

このプロジェクトは機能的ですが、インターフェイスはまだ確認されていません。すべてのパブリックAPIは、v0.1.0まで警告なしに変更される場合があります。

著者とライセンス

Torchgpipeプロジェクトは、Heungsub Lee、Myungryong Jeong、およびChiheon KimによってKakao Brainによって開発され、Sungbin Lim、Ildoo Kim、Woonhyuk Baek、Boogeon Yoonの助けがあります。 3節BSDライセンスの下で配布されています。

引用

このライブラリをプロジェクトと調査に適用する場合は、コードを引用してください。

 @article{kim2020torchgpipe,
    title={torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models},
    author={Chiheon Kim and Heungsub Lee and Myungryong Jeong and Woonhyuk Baek and Boogeon Yoon and Ildoo Kim and Sungbin Lim and Sungwoong Kim},
    year={2020},
    eprint={2004.09910},
    archivePrefix={arXiv}
}