GPipe : Entraînement efficace de réseaux de neurones géants utilisant le parallélisme en pipeline

Augmenter la capacité des réseaux de neurones profonds est reconnu comme une approche efficace pour améliorer la qualité du modèle dans plusieurs tâches d'apprentissage automatique. Dans de nombreux cas, augmenter la capacité du modèle au-delà de la limite de mémoire d'un seul accélérateur nécessite le développement d'algorithmes ou d'infrastructures spécifiques. Ces solutions sont souvent propres à une architecture et ne se transfèrent pas à d'autres tâches. Pour répondre au besoin de parallélisme de modèles efficace et indépendant des tâches, nous introduisons GPipe, une bibliothèque de parallélisme en pipeline qui permet de faire évoluer n'importe quel réseau pouvant être exprimé comme une séquence de couches. En pipelineant différentes sous-séquences de couches sur des accélérateurs séparés, GPipe offre la flexibilité d'augmenter efficacement à grande échelle une variété de réseaux différents. De plus, GPipe utilise un nouvel algorithme de pipelining de division de lots, résultant en une accélération presque linéaire lorsque le modèle est partitionné sur plusieurs accélérateurs. Nous démontrons les avantages de GPipe en entraînant des réseaux neuronaux à grande échelle sur deux tâches différentes avec des architectures de réseau distinctes : (i) Classification d'images : nous entraînons un modèle AmoebaNet de 557 millions de paramètres et atteignons une précision top-1 de 84,4 % sur ImageNet-2012, (ii) Traduction automatique neuronale multilingue : nous entraînons un unique modèle Transformer de 6 milliards de paramètres, 128 couches, sur un corpus couvrant plus de 100 langues et obtenons une qualité supérieure à tous les modèles bilingues.