Распределенная реализация оптимизатора Shampoo с параллельной обработкой данных в PyTorch для масштабного обучения нейронных сетей

Аннотация

Shampoo — это алгоритм онлайн- и стохастической оптимизации, относящийся к семейству методов AdaGrad для обучения нейронных сетей. Он строит блочно-диагональный предобуславливатель, где каждый блок представляет собой грубое приближение Кронекера к полноматричному AdaGrad для каждого параметра нейронной сети. В данной работе мы предоставляем полное описание алгоритма, а также оптимизации производительности, которые использует наша реализация для масштабного обучения глубоких сетей в PyTorch. Наша реализация обеспечивает быстрое распределенное обучение с использованием нескольких GPU и параллелизацией данных за счет распределения памяти и вычислений, связанных с блоками каждого параметра, с использованием структуры данных DTensor в PyTorch и выполнения примитива AllGather для вычисленных направлений поиска на каждой итерации. Это значительное улучшение производительности позволяет нам достичь снижения времени выполнения на шаг не более чем на 10% по сравнению со стандартными адаптивными градиентными методами, основанными на диагональном масштабировании. Мы проверяем нашу реализацию, проводя исследование на обучении модели ResNet50 на наборе данных ImageNet, демонстрируя превосходство Shampoo над стандартными подходами к обучению при минимальной настройке гиперпараметров.

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.

Распределенная реализация оптимизатора Shampoo с параллельной обработкой данных в PyTorch для масштабного обучения нейронных сетей

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

Аннотация

Support