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 открытую реализацию на PyTorch: 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