2BP: 2-Stage Backpropagation
Authors: name: Christopher Rae, name: Joseph K. L. Lee, name: James Richings

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.