El entrenamiento de modelos de lenguaje eficiente en comunicación escala de manera confiable y robusta: Leyes de escalamiento para DiLoCo

Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo

March 12, 2025

Autores: Zachary Charles, Gabriel Teston, Lucio Dery, Keith Rush, Nova Fallen, Zachary Garrett, Arthur Szlam, Arthur Douillard

cs.AI

Resumen

A medida que escalamos hacia modelos de aprendizaje automático más masivos, las frecuentes demandas de sincronización inherentes a los enfoques de paralelismo de datos generan ralentizaciones significativas, lo que plantea un desafío crítico para un mayor escalamiento. Trabajos recientes desarrollan un enfoque (DiLoCo) que relaja las demandas de sincronización sin comprometer la calidad del modelo. Sin embargo, estos estudios no analizan detenidamente cómo cambia el comportamiento de DiLoCo con el tamaño del modelo. En este trabajo, estudiamos el comportamiento de las leyes de escalamiento de DiLoCo al entrenar modelos de lenguaje grandes (LLMs) bajo un presupuesto fijo de cómputo. Nos enfocamos en cómo factores algorítmicos, incluyendo el número de réplicas del modelo, los hiperparámetros y el presupuesto de tokens, afectan el entrenamiento de maneras que pueden predecirse con precisión mediante leyes de escalamiento. Encontramos que DiLoCo escala de manera predecible y robusta con el tamaño del modelo. Cuando está bien ajustado, DiLoCo escala mejor que el entrenamiento con paralelismo de datos en función del tamaño del modelo, y puede superar al entrenamiento con paralelismo de datos incluso en tamaños de modelo pequeños. Nuestros resultados muestran un conjunto más amplio de beneficios de DiLoCo de lo que se había documentado previamente, incluyendo tamaños de lote óptimos más grandes, una mejor generalización en tareas posteriores con el escalamiento y una mejora en la pérdida de evaluación para un presupuesto fijo de tokens.

English

As we scale to more massive machine learning models, the frequent synchronization demands inherent in data-parallel approaches create significant slowdowns, posing a critical challenge to further scaling. Recent work develops an approach (DiLoCo) that relaxes synchronization demands without compromising model quality. However, these works do not carefully analyze how DiLoCo's behavior changes with model size. In this work, we study the scaling law behavior of DiLoCo when training LLMs under a fixed compute budget. We focus on how algorithmic factors, including number of model replicas, hyperparameters, and token budget affect training in ways that can be accurately predicted via scaling laws. We find that DiLoCo scales both predictably and robustly with model size. When well-tuned, DiLoCo scales better than data-parallel training with model size, and can outperform data-parallel training even at small model sizes. Our results showcase a more general set of benefits of DiLoCo than previously documented, including increased optimal batch sizes, improved downstream generalization with scale, and improved evaluation loss for a fixed token budget.