DiPaCo: Verteilte Pfadkomposition

DiPaCo: Distributed Path Composition

March 15, 2024

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

cs.AI

Zusammenfassung

Fortschritte im Bereich des maschinellen Lernens (ML) wurden durch die Skalierung neuronaler Netzwerkmodelle vorangetrieben. Diese Skalierung wurde durch immer heroischere Ingenieursleistungen ermöglicht, die erforderlich sind, um ML-Ansätze zu unterstützen, die eine hohe Bandbreite an Kommunikation zwischen parallel arbeitenden Geräten erfordern. In dieser Arbeit schlagen wir eine modular aufgebaute Architektur und Trainingsmethode für ML-Modelle vor, die als DIstributed PAth COmposition (DiPaCo) bezeichnet wird. Während des Trainings verteilt DiPaCo die Berechnung über Pfade durch eine Reihe von gemeinsam genutzten Modulen. Zusammen mit einer von Local-SGD inspirierten Optimierung (DiLoCo), die Module synchronisiert und die Kommunikation drastisch reduziert, erleichtert unser Ansatz das Training über schlecht verbundene und heterogene Arbeitskräfte, wobei das Design Robustheit gegenüber Arbeitsausfällen und Unterbrechungen gewährleistet. Zur Zeit der Inferenz muss nur ein einziger Pfad für jede Eingabe ausgeführt werden, ohne dass eine Modellkompression erforderlich ist. Wir betrachten diesen Ansatz als ersten Prototypen für ein neues Paradigma des groß angelegten Lernens, das weniger synchron und modularer ist. Unsere Experimente am weit verbreiteten C4-Benchmark zeigen, dass DiPaCo bei gleicher Anzahl von Trainingsschritten, aber weniger Wanduhrzeit, die Leistung eines 1 Milliarde Parameter umfassenden dichten Transformer-Sprachmodells übertrifft, indem einer von 256 möglichen Pfaden ausgewählt wird, wobei jeder Pfad eine Größe von 150 Millionen Parametern hat.

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.