Un'implementazione distribuita e data-parallel in PyTorch dell'ottimizzatore Shampoo Distribuito per l'addestramento di reti neurali su larga scala
A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale
September 12, 2023
Autori: Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, Michael Rabbat
cs.AI
Abstract
Shampoo è un algoritmo di ottimizzazione online e stocastico appartenente alla famiglia di metodi AdaGrad per l'addestramento di reti neurali. Costruisce un precondizionatore a blocchi diagonali in cui ogni blocco consiste in un'approssimazione grossolana del prodotto di Kronecker rispetto alla versione a matrice completa di AdaGrad per ciascun parametro della rete neurale. In questo lavoro, forniamo una descrizione completa dell'algoritmo, nonché delle ottimizzazioni delle prestazioni che la nostra implementazione sfrutta per addestrare reti profonde su larga scala in PyTorch. La nostra implementazione consente un addestramento distribuito dati-parallelo veloce su più GPU, distribuendo la memoria e il calcolo associati ai blocchi di ciascun parametro tramite la struttura dati DTensor di PyTorch e eseguendo una primitiva AllGather sulle direzioni di ricerca calcolate a ogni iterazione. Questo significativo miglioramento delle prestazioni ci permette di ottenere al massimo una riduzione del 10% nel tempo di esecuzione per passo rispetto ai metodi adattivi basati su scalatura diagonale standard. Validiamo la nostra implementazione eseguendo uno studio di ablazione sull'addestramento di ResNet50 su ImageNet, dimostrando la superiorità di Shampoo rispetto alle ricette di addestramento standard con una minima regolazione degli iperparametri.
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.