Une implémentation distribuée et parallèle en données de l'optimiseur Shampoo distribué dans PyTorch pour l'entraînement de réseaux de neurones à grande échelle
A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale
September 12, 2023
Auteurs: Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, Michael Rabbat
cs.AI
Résumé
Shampoo est un algorithme d'optimisation en ligne et stochastique appartenant à la famille des méthodes AdaGrad pour l'entraînement des réseaux de neurones. Il construit un préconditionneur en blocs diagonaux où chaque bloc consiste en une approximation grossière par produit de Kronecker de la version en matrice complète d'AdaGrad pour chaque paramètre du réseau de neurones. Dans ce travail, nous fournissons une description complète de l'algorithme ainsi que des optimisations de performance que notre implémentation exploite pour entraîner des réseaux profonds à grande échelle dans PyTorch. Notre implémentation permet un entraînement distribué parallèle sur plusieurs GPU en répartissant la mémoire et les calculs associés aux blocs de chaque paramètre via la structure de données DTensor de PyTorch et en effectuant une primitive AllGather sur les directions de recherche calculées à chaque itération. Cette amélioration majeure des performances nous permet d'obtenir au maximum une réduction de 10 % du temps d'exécution par étape par rapport aux méthodes de gradient adaptatif standard basées sur la mise à l'échelle diagonale. Nous validons notre implémentation en réalisant une étude d'ablation sur l'entraînement de ResNet50 sur ImageNet, démontrant la supériorité de Shampoo par rapport aux recettes d'entraînement standard avec un réglage minimal des hyperparamètres.
English
Shampoo is an online and stochastic optimization algorithm belonging to the
AdaGrad family of methods for training neural networks. It constructs a
block-diagonal preconditioner where each block consists of a coarse Kronecker
product approximation to full-matrix AdaGrad for each parameter of the neural
network. In this work, we provide a complete description of the algorithm as
well as the performance optimizations that our implementation leverages to
train deep networks at-scale in PyTorch. Our implementation enables fast
multi-GPU distributed data-parallel training by distributing the memory and
computation associated with blocks of each parameter via PyTorch's DTensor data
structure and performing an AllGather primitive on the computed search
directions at each iteration. This major performance enhancement enables us to
achieve at most a 10% performance reduction in per-step wall-clock time
compared against standard diagonal-scaling-based adaptive gradient methods. We
validate our implementation by performing an ablation study on training
ImageNet ResNet50, demonstrating Shampoo's superiority over standard training
recipes with minimal hyperparameter tuning.