Eine verteilte datenparallele PyTorch-Implementierung des verteilten Shampoo-Optimierers für das Training von neuronalen Netzwerken im großen Maßstab

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

September 12, 2023

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

cs.AI

Zusammenfassung

Shampoo ist ein Online- und stochastischer Optimierungsalgorithmus, der zur AdaGrad-Familie von Methoden für das Training neuronaler Netzwerke gehört. Es konstruiert einen blockdiagonalen Preconditioner, bei dem jeder Block aus einer groben Kronecker-Produkt-Approximation des vollständigen Matrix-AdaGrad für jeden Parameter des neuronalen Netzwerks besteht. In dieser Arbeit liefern wir eine vollständige Beschreibung des Algorithmus sowie der Leistungsoptimierungen, die unsere Implementierung nutzt, um tiefe Netzwerke in großem Maßstab in PyTorch zu trainieren. Unsere Implementierung ermöglicht ein schnelles, verteiltes Datenparallel-Training auf mehreren GPUs, indem der Speicher und die Berechnung, die mit den Blöcken jedes Parameters verbunden sind, über PyTorchs DTensor-Datenstruktur verteilt werden und ein AllGather-Primitiv auf den berechneten Suchrichtungen bei jeder Iteration durchgeführt wird. Diese wesentliche Leistungssteigerung ermöglicht es uns, höchstens eine 10%ige Reduzierung der pro Schritt gemessenen Wanduhrzeit im Vergleich zu standardmäßigen diagonalen Skalierungsmethoden für adaptive Gradienten zu erreichen. Wir validieren unsere Implementierung durch eine Ablationsstudie zum Training von ImageNet ResNet50, die die Überlegenheit von Shampoo gegenüber Standard-Trainingsrezepten mit minimaler Hyperparameter-Optimierung demonstriert.

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.