torchgpipe

Torchgpipe

Eine GPIPE -Implementierung in Pytorch. Es ist eher für CUDA als für TPU optimiert.

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

Was ist gpipe?

GPIPE ist eine skalierbare Pipeline-Parallelismusbibliothek, die von Google Brain veröffentlicht wurde und die ein effizientes Training großer, speicherkonsumierender Modelle ermöglicht. Gemäß dem Papier kann GPIPE ein 25 -fach größeres Modell mit 8x -Geräten (TPU) trainieren und ein Modell 3.5x schneller mit 4 -fachen Geräten trainieren.

GPIPE: Effizientes Training von riesigen neuronalen Netzwerken mithilfe von Pipeline -Parallelität

Google trainierte AmoeBanet-B mit 557M-Parametern über Gpipe. Dieses Modell hat 84,3% TOP-1 und 97,0% der Top-5-Genauigkeit bei der ImageNet-Klassifizierungsbenchmark (der hochmodernen Leistung bis Mai 2019) erreicht.

GPIPE verwendet (a) Pipeline -Parallelität und (b) automatische Neukomputation der Vorwärtsverbreitung während der Backpropagation und nutzt daher das Training eines großen Modells. Wir bezeichnen (b) als Checkpointing nach der bekannten Terminologie in der Pytorch-Community.

Pipeline Parallelität

GPIPE spaltet ein Modell in mehrere Partitionen auf und legt jede Partition auf ein anderes Gerät, um mehr Speicherkapazität zu belegen. Und es teilt einen Mini-Batch in mehrere Mikro-Stapel, damit die Partitionen so parallel wie möglich funktionieren.

Checkpointing

Checkpointing wird auf jede Partition angewendet, um den Gesamtspeicherverbrauch durch ein Modell zu minimieren. Während der Vorwärtsverbreitung werden nur die Tensoren an den Grenzen zwischen Partitionen in Erinnerung bleiben. Alle anderen mittleren Tensoren werden flüchtig und bei Bedarf während des Backpropagation neu berechnet.

Verwendung

Derzeit erfordert Torchgpipe die folgenden Umgebungen:

• Python 3.6+
• Pytorch 1.1+

Um es mit Torchgpipe zu verwenden, installieren Sie es über PYPI:

$ pip install torchgpipe

Um ein Modul mit GPIPE zu trainieren, wickeln Sie es einfach mit torchgpipe.GPipe ein. Ihr Modul muss nn.Sequential sein, da GPIPE das Modul automatisch in Partitionen mit aufeinanderfolgenden Schichten aufteilt. Das balance bestimmt die Anzahl der Schichten in jeder Partition. Das Argument chunks gibt die Anzahl der Mikrobefehlungen an. Eingangs-, Ausgangs- und Zwischen -Tensoren müssen Tensor oder Tuple[Tensor, ...] .

Der folgende Beispielcode zeigt, wie ein Modul mit vier Ebenen in vier Partitionen aufgeteilt werden kann, die jeweils eine einzige Ebene haben. Dieser Code teilt auch einen Mini-Batch in 8 Mikrobefugnisse auf:

 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 )

Dokumentation

Besuchen Sie Torchgpipe.readthedocs.io für weitere Informationen, einschließlich der API -Referenzen.

Benchmarking

Die vollständigen Details und weitere Benchmarks sind in Torchgpipe.readthedocs.io erhältlich.

Resnet-101 Genauigkeit Benchmark

Chargengröße	Torchgpipe	nn.dataparallel	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 sollte transparent sein, um keine zusätzliche Hyperparameterabstimmung einzuführen. Um die Transparenz zu überprüfen, reproduzierten wir die Top-1-Fehlerrate von Resnet-101 auf ImageNet, wie in Tabelle 2 (c) der genauen, großen Minibatch-SGD von Goyal et al.

U-NET (B, C) Speicherbenchmark

Experiment	U-net (b, c)	Parameter	Speicherverbrauch
Grundlinie	(6, 72)	362,2 m	20.3 Gib
Pipeline-1	(11, 128)	2.21b	20.5 Gib
Pipeline-2	(24, 128)	4.99b	43.4 Gib
Pipeline-4	(24, 160)	7.80b	79.1 Gib
Pipeline-8	(48, 160)	15.82b	154.1 Gib

Die Tabelle zeigt, wie GPIPE die Skalierung von U-NET-Modellen ermöglicht. Baseline bezeichnet die Grundlinie ohne Pipeline -Parallelität oder checkPointing, und Pipeline -1 , -2 , -4 , -8 bedeutet, dass das Modell mit GPIPE mit der entsprechenden Anzahl von Partitionen trainiert wird.

Hier haben wir eine vereinfachte U-NET-Architektur verwendet. Die Größe eines Modells wird durch Hyperparameter B und C bestimmt, die proportional zur Anzahl der Schichten bzw. Filter sind.

U-NET (5, 64) Geschwindigkeits-Benchmark

Experiment	Durchsatz	Beschleunigen
Grundlinie	28.500/s	1 ×
Pipeline-1	24.456/s	0,858 ×
Pipeline-2	35,502/s	1,246 ×
Pipeline-4	67.042/s	2,352 ×
Pipeline-8	88.497/s	3.105 ×

Um die Effizienz mit Skip-Verbindungen zu überprüfen, haben wir den Durchsatz von U-NET mit verschiedenen Geräten gemessen. Wir haben uns für U-NET entschieden, da es mehrere lange Überspring-Verbindungen hat.

Amoebanet-D (18, 256) Geschwindigkeitsbenchmark

Experiment	Durchsatz	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 : Anzahl der Partitionen, M : Anzahl der Mikrobeflüsse)

Die Tabelle zeigt den reproduzierten Geschwindigkeits-Benchmark für Amoebanet-D (18, 256), wie in Tabelle 2 von GPIPE von Huang et al. Beachten Sie, dass wir K in der Zeitung durch n ersetzt haben.

Notizen

Dieses Projekt ist funktional, aber die Schnittstelle ist noch nicht bestätigt. Alle öffentlichen APIs können sich bis V0.1.0 ohne Vorwarnung ändern.

Autoren und Lizenzierung

Das Torchgpipe -Projekt wurde von Heingsub Lee, Myungryong Jeong und Chiheon Kim im Kakao Brain mit Sungbin Lim, Ildoo Kim, Woonhyuk Baek und Boogeon Yoons Hilfe entwickelt. Es wird unter der 3-Klausel-BSD-Lizenz verteilt.

Zitat

Wenn Sie diese Bibliothek auf Projekt und Forschung anwenden, zitieren Sie bitte unseren Code:

 @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}
}