DiPaCo : Composition de Chemins Distribuée

DiPaCo: Distributed Path Composition

March 15, 2024

Auteurs: Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Adhiguna Kuncoro, Yani Donchev, Rachita Chhaparia, Ionel Gog, Marc'Aurelio Ranzato, Jiajun Shen, Arthur Szlam

cs.AI

Résumé

Les progrès en apprentissage automatique (ML) ont été alimentés par la mise à l'échelle des modèles de réseaux neuronaux. Cette mise à l'échelle a été rendue possible par des prouesses techniques toujours plus impressionnantes, nécessaires pour accueillir des approches de ML qui requièrent une communication à haut débit entre des dispositifs fonctionnant en parallèle. Dans ce travail, nous proposons une architecture modulaire et une approche d'entraînement co-conçues pour les modèles de ML, baptisées DIstributed PAth COmposition (DiPaCo). Pendant l'entraînement, DiPaCo répartit le calcul par chemins à travers un ensemble de modules partagés. Associée à une optimisation inspirée de Local-SGD (DiLoCo) qui maintient les modules synchronisés avec une communication drastiquement réduite, notre approche facilite l'entraînement sur des travailleurs faiblement connectés et hétérogènes, avec une conception qui garantit la robustesse aux défaillances et interruptions des travailleurs. Au moment de l'inférence, un seul chemin doit être exécuté pour chaque entrée, sans nécessiter de compression du modèle. Nous considérons cette approche comme un premier prototype vers un nouveau paradigme d'apprentissage à grande échelle, moins synchrone et plus modulaire. Nos expériences sur le benchmark largement utilisé C4 montrent que, pour le même nombre d'étapes d'entraînement mais moins de temps écoulé, DiPaCo dépasse les performances d'un modèle de langage transformateur dense de 1 milliard de paramètres en choisissant l'un des 256 chemins possibles, chacun ayant une taille de 150 millions de paramètres.

English

Progress in machine learning (ML) has been fueled by scaling neural network models. This scaling has been enabled by ever more heroic feats of engineering, necessary for accommodating ML approaches that require high bandwidth communication between devices working in parallel. In this work, we propose a co-designed modular architecture and training approach for ML models, dubbed DIstributed PAth COmposition (DiPaCo). During training, DiPaCo distributes computation by paths through a set of shared modules. Together with a Local-SGD inspired optimization (DiLoCo) that keeps modules in sync with drastically reduced communication, Our approach facilitates training across poorly connected and heterogeneous workers, with a design that ensures robustness to worker failures and preemptions. At inference time, only a single path needs to be executed for each input, without the need for any model compression. We consider this approach as a first prototype towards a new paradigm of large-scale learning, one that is less synchronous and more modular. Our experiments on the widely used C4 benchmark show that, for the same amount of training steps but less wall-clock time, DiPaCo exceeds the performance of a 1 billion-parameter dense transformer language model by choosing one of 256 possible paths, each with a size of 150 million parameters.