NoLoCo: Метод обучения больших моделей с низким уровнем коммуникации без использования all-reduce

NoLoCo: No-all-reduce Low Communication Training Method for Large Models

June 12, 2025

Авторы: Jari Kolehmainen, Nikolay Blagoev, John Donaghy, Oğuzhan Ersoy, Christopher Nies

cs.AI

Аннотация

Обучение крупных языковых моделей обычно осуществляется с использованием методов оптимизации на кластерах, содержащих десятки тысяч акселераторов, которые взаимодействуют через высокоскоростные соединения. Масштабирование таких кластеров является дорогостоящим и может стать непрактичным, что накладывает ограничения на размер моделей, которые можно обучить. Несколько недавних исследований предложили методы обучения, требующие меньшего объема коммуникаций, что позволяет избежать необходимости в высокосвязанных вычислительных кластерах. Эти современные методы обучения с низким уровнем коммуникаций по-прежнему используют этап синхронизации параметров модели, который при выполнении на всех репликах модели может стать затратным в условиях низкоскоростной сети. В данной работе мы предлагаем новый метод оптимизации, NoLoCo, который не требует явной синхронизации всех параметров модели в процессе обучения и, как следствие, не использует коллективные коммуникации. NoLoCo неявно синхронизирует веса модели с помощью новой модификации оптимизатора Нестерова, частично усредняя веса модели с весами, случайно выбранными из другой реплики. Мы предоставляем как теоретический анализ сходимости для предложенного оптимизатора, так и эмпирические результаты обучения языковых моделей. Мы тестируем NoLoCo на широком диапазоне количества акселераторов и размеров моделей, от 125 миллионов до 6,8 миллиардов параметров. Наш метод требует значительно меньших затрат на коммуникации по сравнению с полностью распределенным параллельным обучением данных или даже широко используемым методом обучения с низким уровнем коммуникаций, DiLoCo. Этап синхронизации оценивается как на порядок быстрее, чем операция all-reduce, используемая в DiLoCo для обучения на нескольких сотнях акселераторов через интернет. Мы также не используем глобальные блокирующие коммуникации, что сокращает время простоя акселераторов. По сравнению с DiLoCo, мы также наблюдаем до 4% более быструю скорость сходимости для широкого диапазона размеров моделей и количества акселераторов.

English

Training large language models is generally done via optimization methods on clusters containing tens of thousands of accelerators, communicating over a high-bandwidth interconnect. Scaling up these clusters is expensive and can become impractical, imposing limits on the size of models that can be trained. Several recent studies have proposed training methods that are less communication intensive, avoiding the need for a highly connected compute cluster. These state-of-the-art low communication training methods still employ a synchronization step for model parameters, which, when performed over all model replicas, can become costly on a low-bandwidth network. In this work, we propose a novel optimization method, NoLoCo, that does not explicitly synchronize all model parameters during training and, as a result, does not require any collective communication. NoLoCo implicitly synchronizes model weights via a novel variant of the Nesterov momentum optimizer by partially averaging model weights with a randomly selected other one. We provide both a theoretical convergence analysis for our proposed optimizer as well as empirical results from language model training. We benchmark NoLoCo on a wide range of accelerator counts and model sizes, between 125M to 6.8B parameters. Our method requires significantly less communication overhead than fully sharded data parallel training or even widely used low communication training method, DiLoCo. The synchronization step itself is estimated to be one magnitude faster than the all-reduce used in DiLoCo for few hundred accelerators training over the internet. We also do not have any global blocking communication that reduces accelerator idling time. Compared to DiLoCo, we also observe up to 4% faster convergence rate with wide range of model sizes and accelerator counts.