Una Implementación Distribuida y Paralela de Datos en PyTorch del Optimizador Shampoo Distribuido para el Entrenamiento de Redes Neuronales a Gran Escala

A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale

September 12, 2023

Autores: Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, Michael Rabbat

cs.AI

Resumen

Shampoo es un algoritmo de optimización en línea y estocástico que pertenece a la familia de métodos AdaGrad para el entrenamiento de redes neuronales. Construye un precondicionador de bloques diagonales donde cada bloque consiste en una aproximación gruesa mediante el producto de Kronecker al AdaGrad de matriz completa para cada parámetro de la red neuronal. En este trabajo, proporcionamos una descripción completa del algoritmo, así como las optimizaciones de rendimiento que nuestra implementación aprovecha para entrenar redes profundas a gran escala en PyTorch. Nuestra implementación permite un entrenamiento distribuido y paralelo de datos en múltiples GPU de manera rápida, distribuyendo la memoria y el cálculo asociados con los bloques de cada parámetro mediante la estructura de datos DTensor de PyTorch y realizando una primitiva AllGather en las direcciones de búsqueda calculadas en cada iteración. Esta mejora significativa en el rendimiento nos permite lograr, como máximo, una reducción del 10% en el tiempo de reloj por paso en comparación con los métodos de gradiente adaptativo basados en escalado diagonal estándar. Validamos nuestra implementación realizando un estudio de ablación en el entrenamiento de ResNet50 en ImageNet, demostrando la superioridad de Shampoo sobre las recetas de entrenamiento estándar con un ajuste mínimo de hiperparámetros.

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.