Elamur Research⌕
☰
home / papers / 2020 / Memory-Efficient Pipeline-Parallel DNN Training

Memory-Efficient Pipeline-Parallel DNN Training

Paper · 2020 · Deepak Narayanan, Amar Phanishayee, Kaiyu Shi et al.

Summary

Very large language models often cannot fit their weights and activations on one accelerator, so training must be spread across devices. Earlier pipeline-parallel systems force a choice. GPipe keeps one weight copy but stalls by periodically draining the pipeline. PipeDream avoids the stalls but may store as many weight copies as there are stages. PipeDream-2BW resolves this with double-buffered weight updates (2BW). Gradients from many microbatches are summed into one update. Microbatches already in the pipeline finish on the previous "shadow" weights, while newly admitted ones use the fresh weights, so each worker holds only two versions. The update uses gradients that are exactly one step stale. A variant, PipeDream-Flush, gives up some speed for even less memory and standard update semantics. A memory-aware planner chooses stage sizes, replication, batch size and activation recomputation, and keeps heavy traffic inside servers. On GPT and BERT models, training runs up to 20× faster than a tuned model-parallel baseline with comparable accuracy. This matters because it makes billion-parameter training feasible on commodity GPU clusters.

Key points

  • Memory per worker: 2BW keeps two weight versions per worker, versus up to d versions in PipeDream for a pipeline of depth d. It also removes GPipe's pipeline flushes.
  • Speedups:
    • 1.3× to 20× faster end-to-end on GPT models than an optimized model-parallel (Megatron-style) baseline. The largest gains appear on AWS multi-server setups with limited inter-node bandwidth.
    • Up to 3.2× faster than GPipe.
  • Model size: about 30 billion parameters trained on 64 GPUs with 16 GB each. For comparison, Megatron-LM trained 8.3B parameters on 8 GPUs with 32 GB each.
  • Convergence: Adam with 2BW tracks the training loss of vanilla Adam at equal batch size, and downstream finetuning accuracy is similar.
  • Planner cost: the planner considers only O(m²) (width, depth) configurations for m workers, compared with O(n³m²) for PipeDream's planner. It finishes in under a second, whereas PipeDream's partitioner needs roughly 30 minutes on large models.
  • Limitation: accumulating gradients inside the pipeline raises the minimum batch size the system supports.

Interesting ideas

  • I1. Delay updates by a fixed amount. Accept a constant, uniform one-step staleness in exchange for bounded memory. Because every stage sees the same delay, no new hyperparameter tuning is needed. This could apply to any asynchronous or streaming training setup.
  • I2. Exploit repetition to shrink the search. Models built from repeated blocks, such as stacks of transformer layers, can be split into equal, equally replicated stages. This turns an exponential partitioning problem into a small exhaustive search, a useful trick for any placement or scheduling task over regular structures.
  • I3. Treat memory savers as tunable options. Activation recomputation should be chosen jointly with throughput rather than always switched on, since always recomputing cost about 1.33× in speed. This is a reminder to put memory/speed trade-offs inside the cost model.
  • I4. Plan around network topology. Placing communication-heavy splits on fast intra-server links, and only lighter traffic across servers, explains much of the gain on cheap cloud clusters. It is relevant for teams training on rented, unevenly connected hardware.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Stanford UniversityStanford · United States506 works · 330,052 citations🏢MicrosoftRedmond · United States378 works · 161,630 citations🏢DatabricksSan Francisco · United States2 works · 217 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksComputer Vision and Pattern Recognition · 3,486 worksHardware and Architecture · 167 works

🧭 Topics 1Research subjects and languages the work is about.Efficient models & compression · 297 works

Cite

Cite as Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen and Matei Zaharia. 2020. Memory-Efficient Pipeline-Parallel DNN Training. ICML 2021.

The paper itself, not this page. BibTeX from the library's record of the work. arXiv 2006.09503

@inproceedings{narayanan2020memory,
  title = {{Memory-Efficient Pipeline-Parallel DNN Training}},
  author = {Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen and Matei Zaharia},
  booktitle = {ICML 2021},
  year = {2020},
  eprint = {2006.09503},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2006.09503}
}

References

23 references, 7 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Megatron Repository. — Language Models are Few-Shot Learners https://github.com/nvidia/megatron-lm. NVIDIA Deep Learning Examples, BERT. https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/README. md\#results. OpenWebText Dataset. https://github.com/jcpeterson/openwebtext. Brown, T., Mann, B., Ryder, N.… ↗
  2. Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training Deep Nets with Sublinear Memory Cost. arXiv preprint arXiv:1604.06174, 2016. find
  3. Chilimbi, T. M., Suzue, Y., Apacible, J., and Kalyanaraman, K. Project Adam: Building an Efficient and Scalable Deep Learning Training System. In 11th USENIX Symposium on Operating Systems Design and Implementation (OSDI ’14), volume 14, pp. 571–582, 2014. find
  4. Dean, J., Corrado, G., Monga, R., Chen, K., Devin, M., Mao, M., Senior, A., Tucker, P., Yang, K., Le, Q. V., et al. Large Scale Distributed Deep Networks. In Advances in Neural Information Processing Systems, pp. 1223–1231, 2012. find
  5. Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv preprint arXiv:1810.04805, 2018.
  6. Fan, S., Rong, Y., Meng, C., Cao, Z., Wang, S., Zheng, Z., Wu, C., Long, G., Yang, J., Xia, L., et al. DAP-PLE: A Pipelined Data Parallel Approach for Training Large Models. In Proceedings of the 26th ACM SIG-PLAN Symposium on Principles and Practice of Parallel Programming, pp. 431–445, 2021. find
  7. Griewank, A. and Walther, A. Revolve: An Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation. ACM Transactions on Mathematical Software (TOMS), 26(1):19–45, 2000. find
  8. Harlap, A., Narayanan, D., Phanishayee, A., Seshadri, V., Devanur, N., Ganger, G., and Gibbons, P. PipeDream: Fast and Efficient Pipeline Parallel DNN Training. arXiv preprint arXiv:1806.03377, 2018.
  9. Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. In Advances in Neural Information Processing Systems, pp. 103–112, 2019. find
  10. Jain, A., Phanishayee, A., Mars, J., Tang, L., and Pekhimenko, G. Gist: Efficient Data Encoding for Deep Neural Network Training. In 2018 ACM/IEEE 45th Annual International Symposium on Computer Architecture (ISCA), pp. 776–789. find
  11. Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Gonzalez, J., Keutzer, K., and Stoica, I. Breaking the Memory Wall with Optimal Tensor Rematerialization. In Proceedings of Machine Learning and Systems 2020, pp. 497–511. 2020. find
  12. Jia, Z., Zaharia, M., and Aiken, A. Beyond Data and Model Parallelism for Deep Neural Networks. In Proceedings of the 2nd Conference on Machine Learning and Systems (MLSys), 2018.
  13. Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., et al. In-Datacenter Performance Analysis of a Tensor Processing Unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture, pp. 1–12, 2017.
  14. Lai, G., Xie, Q., Liu, H., Yang, Y., and Hovy, E. RACE: Large-scale ReAding Comprehension Dataset From Examinations. arXiv preprint arXiv:1704.04683, 2017.
  15. Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. PipeDream: Generalized Pipeline Parallelism for DNN Training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, pp. 1–15, 2019. find
  16. Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving Language Understanding by Generative Pre-training. URL https://s3-us-west-2.amazonaws. com/openai-assets/researchcovers/languageunsupervised/language understanding paper. pdf, 2018. ↗
  17. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language Models are Unsupervised Multi-task Learners. OpenAI Blog, 1(8):9, 2019. find
  18. Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. ZeRO: Memory Optimization Towards Training A Trillion Parameter Models. arXiv preprint arXiv:1910.02054, 2019. find
  19. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LM: Training Multi-Billion Parameter Language Models using GPU Model Parallelism. arXiv preprint arXiv:1909.08053, 2019. find
  20. Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding. 2019. In the Proceedings of ICLR.
  21. Xing, E. P., Ho, Q., Dai, W., Kim, J. K., Wei, J., Lee, S., Zheng, X., Xie, P., Kumar, A., and Yu, Y. Petuum: A New Platform for Distributed Machine Learning on Big Data. IEEE Transactions on Big Data, 1(2):49–67, 2015. find
  22. Yang, B., Zhang, J., Li, J., Ré, C., Aberger, C. R., and De Sa, C. PipeMare: Asynchronous Pipeline Parallel DNN Training. arXiv preprint arXiv:1910.05124, 2019. — [arXiv:1910.05124](https://arxiv.org/abs/1910.05124) find
  23. A. Planner, Additional Details For every possible configuration of width and depth, PipeDream-2BW’s planner explores the benefit of pipelining and each space-saving optimization. For example, with activation recomputation as a target memory-savings optimization, PipeDream-2BW considers three possibl… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/narayanan-2020-memory-efficient-pipeline-parallel-dnn/