2BP: Retropropagación en 2 Etapas

2BP: 2-Stage Backpropagation

May 28, 2024

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

cs.AI

Resumen

A medida que las Redes Neuronales Profundas (DNNs) aumentan en tamaño y complejidad, a menudo superan la capacidad de memoria de un solo acelerador, lo que requiere la fragmentación de los parámetros del modelo en múltiples aceleradores. El paralelismo de tubería (pipeline parallelism) es una estrategia de fragmentación comúnmente utilizada para entrenar DNNs grandes. Sin embargo, las implementaciones actuales de paralelismo de tubería están siendo limitadas involuntariamente por las herramientas de diferenciación automática proporcionadas por los frameworks de aprendizaje automático. Este artículo introduce la retropropagación en 2 etapas (2BP). Al dividir el paso de propagación hacia atrás en dos etapas separadas, podemos reducir el tiempo de cómputo inactivo. Probamos 2BP en varias arquitecturas de modelos y planes de tubería, logrando aumentos en el rendimiento en todos los casos. Utilizando 2BP, pudimos alcanzar un aumento de 1.70x en el rendimiento en comparación con los métodos tradicionales al entrenar un transformador similar a LLaMa con 7 mil millones de parámetros en 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.