Saved in:
| Main Authors: | , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2405.18047 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913366865346560 |
|---|---|
| author | Rae, Christopher Lee, Joseph K. L. Richings, James |
| author_facet | Rae, Christopher Lee, Joseph K. L. Richings, James |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_18047 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | 2BP: 2-Stage Backpropagation Rae, Christopher Lee, Joseph K. L. Richings, James Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing 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. |
| title | 2BP: 2-Stage Backpropagation |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2405.18047 |