DeMo: Otimização de Momento Desacoplado

DeMo: Decoupled Momentum Optimization

November 29, 2024

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

cs.AI

Resumo

O treinamento de grandes redes neurais geralmente requer o compartilhamento de gradientes entre aceleradores por meio de interconexões especializadas de alta velocidade. Inspirados nos princípios de processamento de sinais de decomposição em frequência e compactação de energia, demonstramos que a sincronização completa dos estados do otimizador e dos parâmetros do modelo durante o treinamento é desnecessária. Ao desacoplar as atualizações de momento e permitir divergências controladas nos estados do otimizador entre os aceleradores, alcançamos uma convergência aprimorada em comparação com otimizadores de ponta. Apresentamos o {De}coupled {Mo}mentum (DeMo), um otimizador fusível e um algoritmo paralelo de dados que reduz os requisitos de comunicação entre aceleradores em várias ordens de magnitude. Isso possibilita o treinamento de grandes redes neurais mesmo com largura de banda de rede limitada e hardware heterogêneo. Nosso método é agnóstico em relação à topologia e independente de arquitetura, e suporta treinamento distribuído síncrono por clock escalável com sobrecarga computacional e de memória negligenciável. Resultados empíricos mostram que modelos treinados com DeMo correspondem ou excedem o desempenho de modelos equivalentes treinados com AdamW, eliminando a necessidade de interconexões de alta velocidade ao pré-treinar modelos de fundação em larga escala. Uma implementação de referência de código aberto em PyTorch está disponível no GitHub em 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