[논문 리뷰] Zero Bubble Pipeline Parallelism

Penghui Qi, Xinyi Wan외 2명|arXiv (Cornell University)|2023. 11. 30.

Parallel Computing and Optimization Techniques인용 수 4

한 줄 요약

이 논문은 동기 학습에서 파이프라인 버블을 거의 제로로 줄이는 데 목적이 있는 새로운 스케줄링 전략인 Zero Bubble Pipeline Parallelism (ZB)을 제안한다. 이는 역전파 계산을 입력 그래디언트(B)와 파라미터 그래디언트(W)로 분리함으로써 달성된다. 더 세밀한 계산 그래프 재구성과 새로운 스케줄(ZB-V)을 도입함으로써, 1F1B 대비 버블 비율을 최대 75% 감소시키며, 유사한 메모리 제약 조건 하에서는 처리량을 23% 향상시키고, 메모리가 여유로운 경우 최대 31% 향상시킨다.

ABSTRACT

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.

연구 동기 및 목표

• 동기 학습 환경에서 지속적인 비효율성의 원인이 되는 파이프라인 버블 문제를 해결한다.
• 개선된 점이 있음에도 불구하고 여전히 비제로 버블 비율을 겪는 기존 스케줄링 전략(예: 1F1B)의 한계를 극복한다.
• 버블 비율을 최소화하면서도 메모리 제약 조건을 고려하고 대규모 모델 학습을 효율적으로 가능하게 하는 스케줄링 프레임워크를 설계한다.
• 최적화 의미론을 훼손하거나 기존 기준 방법을 초월한 메모리 사용 증가 없이, 버블이 없는 학습을 가능하게 한다.
• 모델 구성과 메모리 제약 조건에 기반해 최적의 스케줄을 자동 생성할 수 있는 알고리즘을 개발한다.

제안 방법

• 기존의 단일한 역전파 함수를 분해하여, 입력에 대한 그래디언트(B)와 파rameter에 대한 그래디언트(W)로 두 개의 독립적인 컴포onent으로 나누어, 전통적인 통합된 역전파 함수의 구조를 해체한다.
• 파이프라인 버블을 유발하는 의존성을 제거하기 위해 B와 W 계산의 순서를 재정렬한 새로운 파이프라인 스케줄 ZB-V를 도입하며, 순방향, 역방향, 파라미터 그래디언트 계산 시간이 동일할 경우 버블이 완전히 제거된다.
• 최적화기 단계 동안 새로운 동기화 바이패스 기법을 설계하여 유휴 시간을 제거하고, 추가로 버블을 줄인다.
• 주어진 메모리 및 모델 제약 조건 하에서 버블 비율을 최소화하는 최적의 스케줄을 찾는 자동 스케줄링 알고리즘을 개발한다.
• 레이어 수준 추상화가 아닌 함수 수준(F, B, W)의 계산 그래프 표현을 사용하여 더 세밀한 스케줄링과 최적화를 가능하게 한다.
• Megatron-LM 기반으로 구현하고 오픈소스화하여 기존 대규모 학습 프레임워크에의 통합을 가능하게 한다.

실험 결과

연구 질문

• RQ1 역전파 계산의 재구성으로 인해 동기 파이프라인 병렬 학습에서 파이프라인 버블을 완전히 제거할 수 있는가?
• RQ2 메모리 효율성과 표준 학습 의미론과의 호환성을 유지하면서도 최소한의 버블 비율을 달성할 수 있는 스케줄링 전략은 무엇인가?
• RQ3 역전파를 입력 그래디언트와 파라미터 그래디언트 컴포넌트로 분리하는 것이 파이프라인 효율성과 메모리 사용에 어떤 영향을 미치는가?
• RQ4 자동 스케줄링이 다양한 모델 크기와 메모리 예산 하에서 버블 비율을 얼마나 줄일 수 있는가?
• RQ5 마이크로배치 크기와 버블 비율 사이의 상호관계는 무엇이며, 이를 통해 최대 처리량을 최적화할 수 있는가?

주요 결과

• ZB-V는 32개 스테이지와 256개 마이크로배치를 가진 28.3B 모델에서 버블 비율 0.0274를 기록했으며, 1F1B의 0.2676 대비 89.8% 감소한 결과를 보였다.
• 유사한 메모리 제약 조건 하에서 ZB-V는 1F1B 대비 최대 23%의 처리량 향상을 기록했고, 메모리가 여유로운 경우 최대 31% 향상되었다.
• 모델 크기와 구성에 관계없이 ZB-V는 1F1B 대비 버블 비율을 최대 75% 감소시켰다.
• ZB-V는 1F1B와 동일한 메모리 프로파일을 유지하면서도 거의 제로 버블을 달성했으며, 유사한 성능를 위해 두 배의 메모리를 요구하는 ZB-H2와는 달리 이점을 가졌다.
• 자동 스케줄링 알고리즘이 히ュ리스틱 기반 베이스라인을 뛰어넘는 최적의 스케줄을 성공적으로 생성했으며, 특히 메모리 제약이 깐깐한 환경에서 뚜렷한 성능 향상을 보였다.
• 마이크로배치 크기를 늘리면 더 큰 모델에서는 처리량이 최대 8% 향상되지만, 더 작은 모델에서는 이익이 점점 줄어들며, 마이크로배치 크기와 버블 비율 사이의 트레이드오프가 존재함을 시사한다.