DiPaCo: Composición de Rutas Distribuidas

DiPaCo: Distributed Path Composition

March 15, 2024

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

cs.AI

Resumen

Los avances en el aprendizaje automático (ML) han sido impulsados por el escalamiento de modelos de redes neuronales. Este escalamiento ha sido posible gracias a hazañas de ingeniería cada vez más ambiciosas, necesarias para acomodar enfoques de ML que requieren una comunicación de alto ancho de banda entre dispositivos que trabajan en paralelo. En este trabajo, proponemos una arquitectura modular y un enfoque de entrenamiento co-diseñados para modelos de ML, denominados Composición de Rutas Distribuidas (DiPaCo). Durante el entrenamiento, DiPaCo distribuye el cálculo mediante rutas a través de un conjunto de módulos compartidos. Junto con una optimización inspirada en SGD Local (DiLoCo) que mantiene los módulos sincronizados con una comunicación drásticamente reducida, nuestro enfoque facilita el entrenamiento en trabajadores heterogéneos y con conexiones deficientes, con un diseño que garantiza robustez frente a fallos e interrupciones de los trabajadores. En el momento de la inferencia, solo es necesario ejecutar una única ruta para cada entrada, sin necesidad de compresión del modelo. Consideramos este enfoque como un primer prototipo hacia un nuevo paradigma de aprendizaje a gran escala, uno que es menos sincrónico y más modular. Nuestros experimentos en el ampliamente utilizado benchmark C4 muestran que, para la misma cantidad de pasos de entrenamiento pero en menos tiempo de reloj, DiPaCo supera el rendimiento de un modelo de lenguaje denso de transformadores con 1.000 millones de parámetros al elegir una de las 256 rutas posibles, cada una con un tamaño de 150 millones de parámetros.

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.