2BP: Backpropagation in Due Fasi
2BP: 2-Stage Backpropagation
May 28, 2024
Autori: Christopher Rae, Joseph K. L. Lee, James Richings
cs.AI
Abstract
Man mano che le Reti Neurali Profonde (DNN) crescono in dimensioni e complessità, spesso superano la capacità di memoria di un singolo acceleratore, rendendo necessario lo sharding dei parametri del modello su più acceleratori. Il parallelismo pipeline è una strategia di sharding comunemente utilizzata per l'addestramento di grandi DNN. Tuttavia, le attuali implementazioni del parallelismo pipeline sono involontariamente limitate dagli strumenti di differenziazione automatica forniti dai framework di machine learning. Questo articolo introduce la retropropagazione a 2 stadi (2BP). Suddividendo il passo di propagazione all'indietro in due fasi separate, è possibile ridurre il tempo di inattività del calcolo. Abbiamo testato 2BP su varie architetture di modelli e schedulazioni di pipeline, ottenendo aumenti di throughput in tutti i casi. Utilizzando 2BP, siamo riusciti a ottenere un aumento di 1,70x nel throughput rispetto ai metodi tradizionali durante l'addestramento di un trasformatore simile a LLaMa con 7 miliardi di parametri su 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.