Interleaving Optimization and Backward Computation in Large-Scale AI Models to Improve Pipeline Parallelism

Explore whether optimization and backward-computation phases can be interleaved to improve resource utilization and efficiency in pipeline-parallel training.

The candidate will: 1. Review pipeline parallelism and distributed optimization algorithms. 2. Analyze dependencies between backward computation and optimization steps. 3. Identify feasible interleaving and scheduling strategies. 4. Design and implement a prototype scheduling mechanism. 5. Evaluate pipeline utilization, throughput, memory consumption, and convergence. 6. Compare the proposed approach with conventional execution schedules.

3 to 9 months according to the student's background.

• Knowledge of deep learning
• Knowledge of parallel programming
• Familiarity with distributed training and optimization

1. Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, & Matei Zaharia (2021). Efficient large-scale language model training on GPU clusters using megatron-LM. In International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2021, St. Louis, Missouri, USA, November 14-19, 2021 (pp. 58). ACM.
2. Penghui Qi, Xinyi Wan, Guangxing Huang, & Min Lin (2024). Zero Bubble (Almost) Pipeline Parallelism. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.