Een gedistribueerde data-parallelle PyTorch-implementatie van de gedistribueerde Shampoo-optimalisator voor het trainen van neurale netwerken op grote schaal

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

Samenvatting

Shampoo is een online en stochastisch optimalisatie-algoritme dat behoort tot de AdaGrad-familie van methoden voor het trainen van neurale netwerken. Het construeert een blokdiagonaal preconditioner waarbij elk blok bestaat uit een grove Kronecker-productbenadering van de volledige matrix AdaGrad voor elke parameter van het neurale netwerk. In dit werk bieden we een volledige beschrijving van het algoritme, evenals de prestatieoptimalisaties die onze implementatie benut om diepe netwerken op grote schaal te trainen in PyTorch. Onze implementatie maakt snelle multi-GPU gedistribueerde data-parallelle training mogelijk door het geheugen en de berekeningen die gepaard gaan met blokken van elke parameter te verdelen via PyTorch's DTensor-gegevensstructuur en een AllGather-primitief uit te voeren op de berekende zoekrichtingen bij elke iteratie. Deze belangrijke prestatieverbetering stelt ons in staat om hooguit een 10% prestatievermindering in de wandkloktijd per stap te bereiken in vergelijking met standaard diagonale schalingsgebaseerde adaptieve gradientmethoden. We valideren onze implementatie door een ablatiestudie uit te voeren op het trainen van ImageNet ResNet50, waarbij de superioriteit van Shampoo wordt aangetoond ten opzichte van standaard trainingsrecepten met minimale hyperparameterafstemming.

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.