PETRA: Parallel End-to-end Training with Reversible Architectures
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915292143157248 |
|---|---|
| author | Rivaud, Stéphane Fournier, Louis Pumir, Thomas Belilovsky, Eugene Eickenberg, Michael Oyallon, Edouard |
| author_facet | Rivaud, Stéphane Fournier, Louis Pumir, Thomas Belilovsky, Eugene Eickenberg, Michael Oyallon, Edouard |
| contents | Reversible architectures have been shown to be capable of performing on par with their non-reversible architectures, being applied in deep learning for memory savings and generative modeling. In this work, we show how reversible architectures can solve challenges in parallelizing deep model training. We introduce PETRA, a novel alternative to backpropagation for parallelizing gradient computations. PETRA facilitates effective model parallelism by enabling stages (i.e., a set of layers) to compute independently on different devices, while only needing to communicate activations and gradients between each other. By decoupling the forward and backward passes and keeping a single updated version of the parameters, the need for weight stashing is also removed. We develop a custom autograd-like training framework for PETRA, and we demonstrate its effectiveness on CIFAR-10, ImageNet32, and ImageNet, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2406_02052 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | PETRA: Parallel End-to-end Training with Reversible Architectures Rivaud, Stéphane Fournier, Louis Pumir, Thomas Belilovsky, Eugene Eickenberg, Michael Oyallon, Edouard Machine Learning Reversible architectures have been shown to be capable of performing on par with their non-reversible architectures, being applied in deep learning for memory savings and generative modeling. In this work, we show how reversible architectures can solve challenges in parallelizing deep model training. We introduce PETRA, a novel alternative to backpropagation for parallelizing gradient computations. PETRA facilitates effective model parallelism by enabling stages (i.e., a set of layers) to compute independently on different devices, while only needing to communicate activations and gradients between each other. By decoupling the forward and backward passes and keeping a single updated version of the parameters, the need for weight stashing is also removed. We develop a custom autograd-like training framework for PETRA, and we demonstrate its effectiveness on CIFAR-10, ImageNet32, and ImageNet, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models. |
| title | PETRA: Parallel End-to-end Training with Reversible Architectures |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2406.02052 |