2BP : Rétropropagation en deux étapes

2BP: 2-Stage Backpropagation

May 28, 2024

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

cs.AI

Résumé

À mesure que les réseaux de neurones profonds (DNN) gagnent en taille et en complexité, ils dépassent souvent la capacité mémoire d'un seul accélérateur, nécessitant ainsi la fragmentation des paramètres du modèle sur plusieurs accélérateurs. Le parallélisme par pipeline est une stratégie de fragmentation couramment utilisée pour entraîner de grands DNN. Cependant, les implémentations actuelles du parallélisme par pipeline sont involontairement limitées par les outils de différenciation automatique fournis par les frameworks de machine learning. Cet article présente la rétropropagation en deux étapes (2BP). En divisant l'étape de rétropropagation en deux phases distinctes, nous pouvons réduire le temps d'inactivité des calculs. Nous avons testé 2BP sur diverses architectures de modèles et plans de pipelining, obtenant des augmentations de débit dans tous les cas. Grâce à 2BP, nous avons pu atteindre une augmentation de débit de 1,70x par rapport aux méthodes traditionnelles lors de l'entraînement d'un transformeur de type LLaMa avec 7 milliards de paramètres sur 4 GPU.

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.