DiLoCo: Entrenamiento Distribuido de Modelos de Lenguaje con Baja Comunicación

Resumen

Los modelos de lenguaje de gran escala (LLM, por sus siglas en inglés) se han convertido en un componente crítico en muchas aplicaciones del aprendizaje automático. Sin embargo, los enfoques estándar para entrenar LLM requieren un gran número de aceleradores estrechamente interconectados, con dispositivos que intercambian gradientes y otros estados intermedios en cada paso de optimización. Si bien es difícil construir y mantener un único clúster de computación que albergue muchos aceleradores, podría ser más fácil encontrar varios clústeres de computación, cada uno con un número menor de dispositivos. En este trabajo, proponemos un algoritmo de optimización distribuida, Distributed Low-Communication (DiLoCo), que permite el entrenamiento de modelos de lenguaje en islas de dispositivos con conexiones deficientes. Este enfoque es una variante del promedio federado, donde el número de pasos internos es grande, el optimizador interno es AdamW y el optimizador externo es el momento de Nesterov. En el ampliamente utilizado conjunto de datos C4, demostramos que DiLoCo con 8 trabajadores tiene un rendimiento comparable al de la optimización completamente sincrónica, mientras comunica 500 veces menos. DiLoCo muestra una gran robustez frente a la distribución de datos de cada trabajador. También es robusto frente a la indisponibilidad de recursos con el tiempo y, viceversa, puede aprovechar sin problemas los recursos que se vuelven disponibles durante el entrenamiento.

English

Large language models (LLM) have become a critical component in many applications of machine learning. However, standard approaches to training LLM require a large number of tightly interconnected accelerators, with devices exchanging gradients and other intermediate states at each optimization step. While it is difficult to build and maintain a single computing cluster hosting many accelerators, it might be easier to find several computing clusters each hosting a smaller number of devices. In this work, we propose a distributed optimization algorithm, Distributed Low-Communication (DiLoCo), that enables training of language models on islands of devices that are poorly connected. The approach is a variant of federated averaging, where the number of inner steps is large, the inner optimizer is AdamW, and the outer optimizer is Nesterov momentum. On the widely used C4 dataset, we show that DiLoCo on 8 workers performs as well as fully synchronous optimization while communicating 500 times less. DiLoCo exhibits great robustness to the data distribution of each worker. It is also robust to resources becoming unavailable over time, and vice versa, it can seamlessly leverage resources that become available during training.

DiLoCo: Entrenamiento Distribuido de Modelos de Lenguaje con Baja Comunicación

DiLoCo: Distributed Low-Communication Training of Language Models

Resumen

Support