DiPaCo: Composição Distribuída de Caminhos

Resumo

O progresso no aprendizado de máquina (ML) tem sido impulsionado pela escalabilidade de modelos de redes neurais. Essa escalabilidade foi possibilitada por feitos de engenharia cada vez mais impressionantes, necessários para acomodar abordagens de ML que exigem comunicação de alta largura de banda entre dispositivos trabalhando em paralelo. Neste trabalho, propomos uma arquitetura modular e uma abordagem de treinamento co-desenhadas para modelos de ML, denominada DIstributed PAth COmposition (DiPaCo). Durante o treinamento, o DiPaCo distribui a computação por caminhos através de um conjunto de módulos compartilhados. Juntamente com uma otimização inspirada no Local-SGD (DiLoCo) que mantém os módulos sincronizados com comunicação drasticamente reduzida, nossa abordagem facilita o treinamento em trabalhadores mal conectados e heterogêneos, com um design que garante robustez a falhas e preempções de trabalhadores. No momento da inferência, apenas um único caminho precisa ser executado para cada entrada, sem a necessidade de qualquer compressão do modelo. Consideramos essa abordagem como um primeiro protótipo em direção a um novo paradigma de aprendizado em larga escala, menos síncrono e mais modular. Nossos experimentos no amplamente utilizado benchmark C4 mostram que, para a mesma quantidade de passos de treinamento, mas em menos tempo de relógio, o DiPaCo supera o desempenho de um modelo de linguagem transformer denso com 1 bilhão de parâmetros ao escolher um dos 256 caminhos possíveis, cada um com um tamanho de 150 milhões 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.

DiPaCo: Composição Distribuída de Caminhos

DiPaCo: Distributed Path Composition

Resumo

Support