DeMo: Optimización de Momento Desacoplado

DeMo: Decoupled Momentum Optimization

November 29, 2024

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

cs.AI

Resumen

El entrenamiento de grandes redes neuronales generalmente requiere compartir gradientes entre aceleradores a través de interconexiones especializadas de alta velocidad. Basándonos en los principios de procesamiento de señales de descomposición en frecuencia y compresión de energía, demostramos que sincronizar estados completos de optimización y parámetros del modelo durante el entrenamiento es innecesario. Al separar las actualizaciones de momento y permitir una divergencia controlada en los estados del optimizador entre aceleradores, logramos una convergencia mejorada en comparación con optimizadores de última generación. Presentamos Decoupled Momentum (DeMo), un optimizador fusionado y un algoritmo de paralelismo de datos que reduce los requisitos de comunicación entre aceleradores en varios órdenes de magnitud. Esto permite el entrenamiento de grandes redes neuronales incluso con ancho de banda de red limitado y hardware heterogéneo. Nuestro método es agnóstico a la topología e independiente de la arquitectura, y admite un entrenamiento distribuido síncrono a nivel de reloj con una sobrecarga computacional y de memoria despreciable. Los resultados empíricos muestran que los modelos entrenados con DeMo igualan o superan el rendimiento de modelos equivalentes entrenados con AdamW, al tiempo que eliminan la necesidad de interconexiones de alta velocidad al preentrenar modelos de base a gran escala. Se ha publicado una implementación de referencia de código abierto en PyTorch en GitHub en 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