신경망 대규모 학습을 위한 분산 Shampoo 최적화기의 분산 데이터 병렬 PyTorch 구현

초록

Shampoo는 신경망 훈련을 위한 AdaGrad 계열의 온라인 및 확률적 최적화 알고리즘입니다. 이 알고리즘은 블록 대각 행렬 형태의 전처리기를 구성하며, 각 블록은 신경망의 각 매개변수에 대해 완전 행렬 AdaGrad의 대략적인 크로네커 곱 근사로 이루어져 있습니다. 본 연구에서는 이 알고리즘의 완전한 설명과 더불어, PyTorch에서 대규모 심층 신경망을 훈련하기 위해 우리의 구현이 활용하는 성능 최적화 기법들을 제시합니다. 우리의 구현은 PyTorch의 DTensor 데이터 구조를 통해 각 매개변수의 블록과 관련된 메모리 및 계산을 분산시키고, 각 반복에서 계산된 탐색 방향에 대해 AllGather 연산을 수행함으로써 빠른 다중 GPU 분산 데이터 병렬 훈련을 가능하게 합니다. 이러한 주요 성능 개선을 통해 우리는 표준 대각 스케일링 기반 적응형 그래디언트 방법 대비 최대 10%의 단계별 벽시계 시간 성능 저하만을 달성할 수 있었습니다. 우리는 ImageNet ResNet50 훈련에 대한 절제 연구를 수행하여 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