DiPaCo: Распределенная композиция путей

DiPaCo: Distributed Path Composition

March 15, 2024

Авторы: Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Adhiguna Kuncoro, Yani Donchev, Rachita Chhaparia, Ionel Gog, Marc'Aurelio Ranzato, Jiajun Shen, Arthur Szlam

cs.AI

Аннотация

Прогресс в области машинного обучения (ML) был обеспечен масштабированием моделей нейронных сетей. Это масштабирование стало возможным благодаря все более героическим подвигам в инженерии, необходимым для адаптации подходов ML, требующих высокопропускной связи между устройствами, работающими параллельно. В данной работе мы предлагаем совместно разработанную модульную архитектуру и подход к обучению моделей ML, названный DIstributed PAth COmposition (DiPaCo). Во время обучения DiPaCo распределяет вычисления по путям через набор общих модулей. Вместе с оптимизацией, вдохновленной Local-SGD (DiLoCo), которая поддерживает синхронизацию модулей с существенно сокращенной коммуникацией, наш подход облегчает обучение на слабо связанных и гетерогенных рабочих узлах, с дизайном, который обеспечивает устойчивость к сбоям и прерываниям работы рабочих узлов. Во время вывода, для каждого ввода требуется выполнение только одного пути, без необходимости сжатия модели. Мы рассматриваем этот подход как первый прототип новой парадигмы масштабного обучения, менее синхронного и более модульного. Наши эксперименты на широко используемом бенчмарке C4 показывают, что при том же количестве шагов обучения, но меньшем времени настенных часов, DiPaCo превосходит производительность модели языкового трансформера с плотностью 1 миллиард параметров, выбирая один из 256 возможных путей, каждый из которых имеет размер 150 миллионов параметров.

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.