L'entraînement de modèles de langage à communication efficace s'adapte de manière fiable et robuste : Lois d'échelle pour DiLoCo

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

March 12, 2025

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

cs.AI

Résumé

Alors que nous passons à des modèles d'apprentissage automatique de plus en plus massifs, les exigences fréquentes de synchronisation inhérentes aux approches parallèles sur les données créent des ralentissements significatifs, posant un défi critique pour une mise à l'échelle ultérieure. Des travaux récents ont développé une approche (DiLoCo) qui relâche les exigences de synchronisation sans compromettre la qualité du modèle. Cependant, ces travaux n'analysent pas attentivement comment le comportement de DiLoCo évolue avec la taille du modèle. Dans cette étude, nous examinons le comportement des lois d'échelle de DiLoCo lors de l'entraînement de LLM (modèles de langage de grande taille) sous un budget de calcul fixe. Nous nous concentrons sur la manière dont les facteurs algorithmiques, y compris le nombre de répliques de modèle, les hyperparamètres et le budget de tokens, affectent l'entraînement de manière prévisible via les lois d'échelle. Nous constatons que DiLoCo s'adapte de manière prévisible et robuste à la taille du modèle. Lorsqu'il est bien réglé, DiLoCo s'adapte mieux que l'entraînement parallèle sur les données avec l'augmentation de la taille du modèle, et peut surpasser l'entraînement parallèle sur les données même pour des modèles de petite taille. Nos résultats mettent en évidence un ensemble plus large d'avantages de DiLoCo que ceux précédemment documentés, y compris des tailles de lots optimales accrues, une meilleure généralisation en aval avec l'échelle, et une amélioration de la perte d'évaluation pour un budget de tokens fixe.

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.