제로 버블 파이프라인 병렬화

Zero Bubble Pipeline Parallelism

November 30, 2023

저자: Penghui Qi, Xinyi Wan, Guangxing Huang, Min Lin

cs.AI

초록

파이프라인 병렬화는 대규모 분산 학습의 핵심 요소 중 하나이지만, 그 효율성은 불가피하다고 여겨졌던 파이프라인 버블로 인해 저하됩니다. 본 연구에서는 우리가 아는 한 동기식 학습 의미론 하에서 파이프라인 버블을 성공적으로 제로로 만드는 최초의 스케줄링 전략을 소개합니다. 이러한 개선의 핵심 아이디어는 역전파 계산을 입력에 대한 그래디언트를 계산하는 부분과 파라미터에 대한 그래디언트를 계산하는 부분으로 나누는 것입니다. 이 아이디어를 바탕으로, 우리는 기존 방법들을 크게 능가하는 새로운 파이프라인 스케줄을 수작업으로 설계했습니다. 또한, 특정 모델 구성과 메모리 제한에 따라 최적의 스케줄을 자동으로 찾는 알고리즘을 개발했습니다. 더 나아가, 진정한 제로 버블을 달성하기 위해 옵티마이저 단계에서의 동기화를 우회하는 새로운 기술을 도입했습니다. 실험 평가 결과, 우리의 방법은 유사한 메모리 제한 하에서 1F1B 스케줄보다 최대 23% 더 높은 처리량을 보였습니다. 메모리 제약이 완화되면 이 수치는 31%까지 더 높아질 수 있습니다. 우리는 이러한 결과가 파이프라인 병렬화의 진정한 잠재력을 활용하는 데 있어 큰 진전을 이뤘다고 믿습니다. 우리는 인기 있는 Megatron-LM 저장소를 기반으로 한 구현을 https://github.com/sail-sg/zero-bubble-pipeline-parallelism에서 오픈소스로 공개했습니다.

English

Pipeline parallelism is one of the key components for large-scale distributed training, yet its efficiency suffers from pipeline bubbles which were deemed inevitable. In this work, we introduce a scheduling strategy that, to our knowledge, is the first to successfully achieve zero pipeline bubbles under synchronous training semantics. The key idea behind this improvement is to split the backward computation into two parts, one that computes gradient for the input and another that computes for the parameters. Based on this idea, we handcraft novel pipeline schedules that significantly outperform the baseline methods. We further develop an algorithm that automatically finds an optimal schedule based on specific model configuration and memory limit. Additionally, to truly achieve zero bubble, we introduce a novel technique to bypass synchronizations during the optimizer step. Experimental evaluations show that our method outperforms the 1F1B schedule up to 23% in throughput under a similar memory limit. This number can be further pushed to 31% when the memory constraint is relaxed. We believe our results mark a major step forward in harnessing the true potential of pipeline parallelism. We open sourced our implementation based on the popular Megatron-LM repository on https://github.com/sail-sg/zero-bubble-pipeline-parallelism.