NoLoCo: Método de entrenamiento con baja comunicación sin reducción total para modelos grandes

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

June 12, 2025

Autores: Jari Kolehmainen, Nikolay Blagoev, John Donaghy, Oğuzhan Ersoy, Christopher Nies

cs.AI

Resumen

El entrenamiento de modelos de lenguaje de gran escala generalmente se realiza mediante métodos de optimización en clústeres que contienen decenas de miles de aceleradores, comunicándose a través de una interconexión de alto ancho de banda. Escalar estos clústeres es costoso y puede volverse poco práctico, imponiendo límites en el tamaño de los modelos que se pueden entrenar. Varios estudios recientes han propuesto métodos de entrenamiento que son menos intensivos en comunicación, evitando la necesidad de un clúster de computación altamente conectado. Estos métodos de entrenamiento de baja comunicación, considerados de vanguardia, aún emplean un paso de sincronización para los parámetros del modelo, lo cual, cuando se realiza en todas las réplicas del modelo, puede volverse costoso en una red de bajo ancho de banda. En este trabajo, proponemos un nuevo método de optimización, NoLoCo, que no sincroniza explícitamente todos los parámetros del modelo durante el entrenamiento y, como resultado, no requiere ninguna comunicación colectiva. NoLoCo sincroniza implícitamente los pesos del modelo mediante una variante novedosa del optimizador de momento de Nesterov, promediando parcialmente los pesos del modelo con otro seleccionado aleatoriamente. Proporcionamos tanto un análisis teórico de convergencia para nuestro optimizador propuesto como resultados empíricos del entrenamiento de modelos de lenguaje. Evaluamos NoLoCo en una amplia gama de cantidades de aceleradores y tamaños de modelos, entre 125M y 6.8B de parámetros. Nuestro método requiere significativamente menos sobrecarga de comunicación que el entrenamiento en paralelo de datos completamente fragmentados o incluso el método de entrenamiento de baja comunicación ampliamente utilizado, DiLoCo. Se estima que el paso de sincronización en sí es un orden de magnitud más rápido que el all-reduce utilizado en DiLoCo para el entrenamiento con unos pocos cientos de aceleradores a través de internet. Además, no tenemos ninguna comunicación global de bloqueo que reduzca el tiempo de inactividad de los aceleradores. En comparación con DiLoCo, también observamos una tasa de convergencia hasta un 4% más rápida en una amplia gama de tamaños de modelos y cantidades de aceleradores.

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.