Декуплированная оптимизация моментов: Новая парадигма для обучения крупных нейронных сетей

Обучение крупных нейронных сетей обычно требует обмена градиентами между акселераторами через специализированные высокоскоростные соединения. Опираясь на принципы обработки сигналов, такие как частотное разложение и сжатие энергии, мы демонстрируем, что синхронизация полных состояний оптимизатора и параметров модели в процессе обучения не является необходимой. Разделяя обновления момента и позволяя контролируемое расхождение в состояниях оптимизатора между акселераторами, мы достигаем улучшенной сходимости по сравнению с современными оптимизаторами. Мы представляем {De}coupled {Mo}mentum (DeMo), объединенный оптимизатор и алгоритм параллельной обработки данных, который снижает требования к межакселераторной связи на несколько порядков. Это позволяет обучать крупные нейронные сети даже при ограниченной пропускной способности сети и неоднородном оборудовании. Наш метод не зависит от топологии и архитектуры и поддерживает масштабируемое распределенное обучение с синхронизацией тактовых импульсов при незначительных затратах вычислительных ресурсов и памяти. Эмпирические результаты показывают, что модели, обученные с помощью DeMo, соответствуют или превосходят производительность эквивалентных моделей, обученных с помощью AdamW, устраняя необходимость в высокоскоростных соединениях при предварительном обучении крупномасштабных фундаментальных моделей. Открытая реализация на базе PyTorch опубликована на GitHub по адресу https://github.com/bloc97/DeMo.