2BP: 2-Stage Backpropagation
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.
66569b0a76582ebbd3a0a604
Christopher Rae
claimed_verified
2024-05-29T08:23:51.603Z
62ec07852780dfe4e2debd7b
/avatars/6391e82a07492b2d642c37d08c120013.svg
Christopher Rae
user
mayfieldmob
66569b0a76582ebbd3a0a605
Joseph K. L. Lee
claimed_verified
2024-05-31T06:25:52.519Z
65d9d0ae0e6ad24551caf7de
/avatars/baf7e8b25631bc9db0f907e474cf40cd.svg
Joseph Lee
user
josephleekl
66569b0a76582ebbd3a0a606
James Richings
2BP: 2-Stage Backpropagation
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.
Introducing 2-stage backpropagation enhances the throughput of pipeline parallelism in training large DNNs by reducing idle compute time.
Deep Neural Networks (DNNs) pipeline parallelism backward propagation 2-stage backpropagation (2BP) throughput LLaMa-like transformer