Uma Implementação Distribuída e Paralela em Dados do Otimizador Shampoo Distribuído em PyTorch para Treinamento de Redes Neurais em 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
Resumo
O Shampoo é um algoritmo de otimização online e estocástico que pertence à família de métodos AdaGrad para o treinamento de redes neurais. Ele constrói um pré-condicionador em bloco diagonal, onde cada bloco consiste em uma aproximação grosseira do produto de Kronecker para o AdaGrad de matriz completa aplicado a cada parâmetro da rede neural. Neste trabalho, fornecemos uma descrição completa do algoritmo, bem como das otimizações de desempenho que nossa implementação utiliza para treinar redes profundas em grande escala no PyTorch. Nossa implementação permite um treinamento distribuído rápido em multi-GPU com paralelismo de dados, distribuindo a memória e a computação associadas aos blocos de cada parâmetro por meio da estrutura de dados DTensor do PyTorch e realizando uma operação primitiva AllGather nas direções de busca calculadas em cada iteração. Esse aprimoramento significativo de desempenho nos permite alcançar, no máximo, uma redução de 10% no tempo de execução por etapa em comparação com métodos adaptativos de gradiente baseados em escalonamento diagonal padrão. Validamos nossa implementação realizando um estudo de ablação no treinamento do ResNet50 no ImageNet, demonstrando a superioridade do Shampoo sobre receitas de treinamento padrão com 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.