DeMo: 분리된 모멘텀 최적화

DeMo: Decoupled Momentum Optimization

November 29, 2024

저자: Bowen Peng, Jeffrey Quesnelle, Diederik P. Kingma

cs.AI

초록

대규모 신경망을 훈련하는 데는 일반적으로 가속기 간의 그래디언트를 전문화된 고속 인터커넥트를 통해 공유해야 합니다. 주파수 분해와 에너지 압축의 신호 처리 원리를 활용하여, 훈련 중 전체 옵티마이저 상태와 모델 매개변수를 동기화하는 것이 불필요함을 입증합니다. 모멘텀 업데이트를 분리하고 가속기 간 옵티마이저 상태에서 제어된 발산을 허용함으로써, 최신 옵티마이저와 비교하여 향상된 수렴을 달성합니다. 저희는 {De}coupled {Mo}mentum (DeMo)이라는 퓨즈드 옵티마이저와 데이터 병렬 알고리즘을 소개합니다. 이를 통해 가속기 간 통신 요구 사항을 여러 단계 줄여 대규모 신경망을 훈련할 수 있게 되었습니다. 이는 제한된 네트워크 대역폭과 이질적 하드웨어에서도 가능해졌습니다. 우리의 방법은 토폴로지에 구애받지 않으며 아키텍처에 독립적이며, 무시할 만한 계산 및 메모리 오버헤드로 확장 가능한 시계 동기 분산 훈련을 지원합니다. 경험적 결과는 DeMo로 훈련된 모델이 AdamW로 훈련된 동등한 모델의 성능을 뛰어넘거나 맞추는 것을 보여주며, 대규모 기반 모델 사전 훈련 시 고속 인터커넥트가 필요 없어졌음을 보여줍니다. GitHub(https://github.com/bloc97/DeMo)에 공개된 PyTorch 참조 구현이 있습니다.

English

Training large neural networks typically requires sharing gradients between accelerators through specialized high-speed interconnects. Drawing from the signal processing principles of frequency decomposition and energy compaction, we demonstrate that synchronizing full optimizer states and model parameters during training is unnecessary. By decoupling momentum updates and allowing controlled divergence in optimizer states across accelerators, we achieve improved convergence compared to state-of-the-art optimizers. We introduce {De}coupled {Mo}mentum (DeMo), a fused optimizer and data parallel algorithm that reduces inter-accelerator communication requirements by several orders of magnitude. This enables training of large neural networks even with limited network bandwidth and heterogeneous hardware. Our method is topology-agnostic and architecture-independent and supports scalable clock-synchronous distributed training with negligible compute and memory overhead. Empirical results show that models trained with DeMo match or exceed the performance of equivalent models trained with AdamW, while eliminating the need for high-speed interconnects when pre-training large scale foundation models. An open source reference PyTorch implementation is published on GitHub at https://github.com/bloc97/DeMo