2BP: Retropropagação em Duas Etapas

2BP: 2-Stage Backpropagation

May 28, 2024

Autores: Christopher Rae, Joseph K. L. Lee, James Richings

cs.AI

Resumo

À medida que as Redes Neurais Profundas (DNNs) aumentam em tamanho e complexidade, elas frequentemente excedem a capacidade de memória de um único acelerador, exigindo o particionamento dos parâmetros do modelo em vários aceleradores. O paralelismo em pipeline é uma estratégia de particionamento comumente usada para treinar DNNs grandes. No entanto, as implementações atuais de paralelismo em pipeline estão sendo involuntariamente limitadas pelas ferramentas de diferenciação automática fornecidas pelos frameworks de aprendizado de máquina. Este artigo introduz a retropropagação em 2 estágios (2BP). Ao dividir o passo de propagação reversa em dois estágios separados, podemos reduzir o tempo ocioso de computação. Testamos o 2BP em várias arquiteturas de modelo e cronogramas de pipeline, alcançando aumentos de taxa de transferência em todos os casos. Usando o 2BP, conseguimos um aumento de 1,70x na taxa de transferência em comparação com métodos tradicionais ao treinar um transformador semelhante ao LLaMa com 7 bilhões de parâmetros em 4 GPUs.

English

As Deep Neural Networks (DNNs) grow in size and complexity, they often exceed the memory capacity of a single accelerator, necessitating the sharding of model parameters across multiple accelerators. Pipeline parallelism is a commonly used sharding strategy for training large DNNs. However, current implementations of pipeline parallelism are being unintentionally bottlenecked by the automatic differentiation tools provided by ML frameworks. This paper introduces 2-stage backpropagation (2BP). By splitting the backward propagation step into two separate stages, we can reduce idle compute time. We tested 2BP on various model architectures and pipelining schedules, achieving increases in throughput in all cases. Using 2BP, we were able to achieve a 1.70x increase in throughput compared to traditional methods when training a LLaMa-like transformer with 7 billion parameters across 4 GPUs.