2BP: 2-Fase Backpropagation

2BP: 2-Stage Backpropagation

May 28, 2024

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

cs.AI

Samenvatting

Naarmate Deep Neural Networks (DNN's) in omvang en complexiteit toenemen, overschrijden ze vaak het geheugencapaciteit van een enkele accelerator, wat het verdelen van modelparameters over meerdere accelerators noodzakelijk maakt. Pipeline-parallelisme is een veelgebruikte strategie voor het verdelen van grote DNN's tijdens training. Huidige implementaties van pipeline-parallelisme worden echter onbedoeld beperkt door de automatische differentiatietools die door ML-frameworks worden aangeboden. Dit artikel introduceert 2-staps backpropagatie (2BP). Door de backward propagation-stap op te splitsen in twee afzonderlijke fasen, kunnen we de inactieve rekentijd verminderen. We hebben 2BP getest op verschillende modelarchitecturen en pipelining-schema's, waarbij in alle gevallen een toename in doorvoersnelheid werd bereikt. Met 2BP konden we een 1,70x hogere doorvoersnelheid bereiken in vergelijking met traditionele methoden bij het trainen van een LLaMa-achtige transformer met 7 miljard parameters verdeeld over 4 GPU's.

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.