DeMo: Ottimizzazione del Momento Decoppiata

DeMo: Decoupled Momentum Optimization

November 29, 2024

Autori: Bowen Peng, Jeffrey Quesnelle, Diederik P. Kingma

cs.AI

Abstract

Allenare grandi reti neurali richiede tipicamente la condivisione dei gradienti tra acceleratori tramite interconnessioni specializzate ad alta velocità. Attingendo ai principi del trattamento del segnale di decomposizione in frequenza e compattazione dell'energia, dimostriamo che sincronizzare stati completi dell'ottimizzatore e parametri del modello durante l'allenamento non è necessario. Decoppiando gli aggiornamenti del momento e consentendo una divergenza controllata negli stati dell'ottimizzatore tra acceleratori, otteniamo una convergenza migliorata rispetto agli ottimizzatori all'avanguardia. Introduciamo DeMo (DeMomentum), un ottimizzatore fuso e un algoritmo parallelo ai dati che riduce i requisiti di comunicazione tra acceleratori di diversi ordini di grandezza. Ciò consente l'allenamento di grandi reti neurali anche con larghezza di banda di rete limitata e hardware eterogeneo. Il nostro metodo è agnostico alla topologia e indipendente dall'architettura e supporta l'allenamento distribuito sincronizzato dall'orologio con un sovraccarico computazionale e di memoria trascurabile. I risultati empirici mostrano che i modelli addestrati con DeMo eguagliano o superano le prestazioni di modelli equivalenti addestrati con AdamW, eliminando nel contempo la necessità di interconnessioni ad alta velocità durante il pre-addestramento di modelli di base su larga scala. Un'implementazione di riferimento open source in PyTorch è pubblicata su GitHub all'indirizzo https://github.com/bloc97/DeMo

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