torchgpipe

torchgpipe

Une implémentation GPipe dans Pytorch. Il est optimisé pour CUDA plutôt que pour TPU.

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

Qu'est-ce que Gpipe?

GPIPE est une bibliothèque de parallélisme de pipeline évolutive publiée par Google Brain, qui permet une formation efficace de grands modèles consommateurs de mémoire. Selon l'article, GPIPE peut former un modèle 25x plus grand en utilisant des appareils 8x (TPU) et former un modèle 3,5x plus rapidement en utilisant des appareils 4x.

Gpipe: formation efficace des réseaux de neurones géants utilisant le parallélisme du pipeline

Amoebanet-B a formé par Google avec des paramètres de 557 m sur gpipe. Ce modèle a atteint 84,3% TOP-1 et 97,0% de précision top-5 sur la référence de classification ImageNet (la performance de pointe en mai 2019).

GPIPE utilise (a) le parallélisme du pipeline et (b) la recomputation automatique de la propagation vers l'avant pendant la rétropropagation, exploite donc la formation d'un grand modèle. Nous appelons (b) le point de contrôle, suivant la terminologie bien connue de la communauté Pytorch.

Parallélisme du pipeline

GPIPE divise un modèle en plusieurs partitions et place chaque partition sur un autre appareil pour occuper plus de capacité de mémoire. Et il divise un mini-lot en plusieurs micro-lots pour que les partitions fonctionnent aussi parallèles que possible.

Point de contrôle

Le point de contrôle est appliqué à chaque partition pour minimiser la consommation globale de mémoire par un modèle. Pendant la propagation directe, seuls les tenseurs aux limites entre les partitions sont rappelés. Tous les autres tenseurs intermédiaires sont volatilisés et recomputés pendant la rétro-propagation lorsque cela est nécessaire.

Usage

Actuellement, TorchGpipe nécessite les environnements suivants:

• Python 3.6+
• Pytorch 1.1+

Pour utiliser TorchGpipe, installez-le via PYPI:

$ pip install torchgpipe

Pour entraîner un module avec gpipe, enveloppez-le simplement avec torchgpipe.GPipe . Votre module doit être nn.Sequential car gpipe divisera automatiquement le module en partitions avec des couches consécutives. L'argument balance détermine le nombre de couches dans chaque partition. L'argument chunks spécifie le nombre de micro-lots. Les tenseurs d'entrée, de sortie et intermédiaires doivent être Tensor ou Tuple[Tensor, ...] .

L'exemple de code ci-dessous montre comment diviser un module avec quatre couches en quatre partitions chacune ayant une seule couche. Ce code divise également un mini-lot en 8 micro-lots:

 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 )

Documentation

Visitez TorchgPipe.readthedocs.io pour plus d'informations, y compris les références API.

Référence

Les détails complets et plus de repères sont disponibles dans TorchgPipe.readthedocs.io.

RESNET-101 Précision Benchmark

Taille de lot	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 doit être transparent pour ne pas introduire un réglage hyperparamètre supplémentaire. Pour vérifier la transparence, nous avons reproduit le taux d'erreur top 1 de RESNET-101 sur ImageNet, comme indiqué dans le tableau 2 (c) de SGD à minibatch précis et précis par Goyal et al.

U-net (b, c) Benchmark de mémoire

Expérience	U-net (b, c)	Paramètres	Utilisation de la mémoire
base de base	(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

Le tableau montre comment GPIPE facilite la mise à l'échelle des modèles U-Net. La ligne de base indique la ligne de base sans parallélisme du pipeline ni de pointage de contrôle, et le pipeline-1 , -2 , -4 , -8 indique que le modèle est formé avec GPIPE avec le nombre correspondant de partitions.

Ici, nous avons utilisé une architecture U-Net simplifiée. La taille d'un modèle est déterminée par les hyperparamètres B et C qui sont proportionnelles au nombre de couches et de filtres, respectivement.

U-net (5, 64) Benchmark Speed

Expérience	Déborder	Accélérer
base de base	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 ×

Pour vérifier l'efficacité avec les connexions de saut, nous avons mesuré le débit de U-Net avec divers nombre d'appareils. Nous avons choisi d'utiliser U-Net car il a plusieurs connexions de saut longues.

AMOBANET-D (18, 256) Benchmark Speed

Expérience	Déborder	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 : nombre de partitions, m : nombre de micro-lots)

Le tableau montre la référence de vitesse reproduite sur Amobanet-D (18, 256), comme indiqué dans le tableau 2 de GPIPE par Huang et al. Notez que nous avons remplacé K dans le papier par n .

Notes

Ce projet est fonctionnel, mais l'interface n'est pas encore confirmée. Toutes les API publiques sont susceptibles de changer sans avertissement jusqu'à V0.1.0.

Auteurs et licences

Le projet Torchgpipe est développé par Heungsub Lee, Myungryong Jeong et Chiheon Kim à Kakao Brain, avec Sungbin Lim, Ildoo Kim, Woonhyuk Baek et Boogeon Yoon. Il est distribué en vertu de la licence BSD à 3 clauses.

Citation

Si vous appliquez cette bibliothèque à n'importe quel projet et recherche, veuillez citer notre 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}
}