[Paper Review] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

논문 링크: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

논문 정보

항목	내용
출판 시점	2018년
저자	Yanping Huang et al.
주제	Distributed training, model parallelism, pipeline parallelism

핵심 아이디어

GPipe는 큰 neural network를 여러 accelerator에 나누어 올리고, mini-batch를 micro-batch로 쪼개 pipeline처럼 흘려보내는 학습 방법이다.

GPipe의 핵심은 model을 layer 단위로 나누고 micro-batch를 pipeline으로 처리해, 여러 device를 동시에 바쁘게 만드는 것이다.

Forward와 backward가 micro-batch 단위로 흐르지만, gradient는 전체 mini-batch 기준으로 모아 동기적으로 update한다. 따라서 pipeline을 쓰면서도 학습 결과의 일관성을 유지하려 한다.

왜 중요한가

대형 모델 학습에서는 single GPU memory가 병목이 된다. GPipe는 model parallelism과 pipeline parallelism을 실용적인 training recipe로 정리한 중요한 논문이다.

수식으로 보는 핵심

GPipe는 model을 $K$ 개의 partition으로 나눈다.

F = F_K \circ F_{K-1} \circ \cdots \circ F_1

Mini-batch $B$ 는 $M$ 개의 micro-batch로 나뉜다.

B = \{B_1, B_2, \dots, B_M\}

각 micro-batch는 partition pipeline을 순서대로 통과한다. Forward는 다음처럼 표현할 수 있다.

h_{m,k} = F_k(h_{m,k-1})

여기서 $m$ 은 micro-batch index, $k$ 는 partition index다. Gradient는 micro-batch별로 계산한 뒤 mini-batch 기준으로 합친다.

\nabla W = \frac{1}{M}\sum_{m=1}^{M}\nabla W_m

GPipe의 핵심 수식은 micro-batch pipeline을 쓰더라도 parameter update는 전체 mini-batch gradient 평균으로 동기화한다는 점이다.

Pipeline bubble의 상대적 비율은 micro-batch 수가 많아질수록 줄어든다. 대략적으로 $M \gg K$ 일수록 device idle time이 작아진다.

대형 모델 학습의 병목

모델이 커지면 한 device의 memory에 parameter, activation, optimizer state를 모두 올릴 수 없다. Data parallelism은 같은 model replica를 여러 device에 복제하므로, model 자체가 너무 크면 해결책이 되지 못한다. GPipe는 model을 layer group 단위로 나누어 여러 accelerator에 배치한다.

Pipeline Parallelism

모델을 여러 partition으로 나누면 하나의 mini-batch가 첫 partition, 두 번째 partition, 세 번째 partition을 순서대로 지나간다. 하지만 mini-batch 하나를 통째로 보내면 앞 device가 끝난 뒤 다음 device가 일하고, 많은 device가 놀게 된다.

GPipe는 mini-batch를 micro-batch로 나누어 pipeline에 연속적으로 넣는다. 이렇게 하면 서로 다른 micro-batch가 서로 다른 partition에서 동시에 처리된다.

GPipe는 micro-batch를 사용해 model parallelism의 memory 이점과 pipeline의 throughput 이점을 함께 얻는다.

Synchronous Update

GPipe는 micro-batch마다 바로 parameter를 업데이트하지 않는다. 전체 mini-batch의 gradient를 누적한 뒤 한 번 update한다. 따라서 pipeline scheduling을 사용하면서도, 수학적으로는 큰 mini-batch를 학습한 것과 같은 형태를 유지한다.

Bubble 문제

Pipeline에는 채우는 구간과 비우는 구간이 있다. 초기에는 앞 partition만 일하고, 마지막에는 뒤 partition만 일한다. 이 비효율을 bubble이라고 한다. Micro-batch 수가 많아질수록 bubble의 상대적 비율은 줄지만, activation 저장과 scheduling overhead가 증가한다.

한계

GPipe는 layer를 어떻게 나눌지에 민감하다. Partition별 계산 시간이 불균형하면 느린 stage가 전체 throughput을 제한한다. 또 activation memory와 communication cost도 고려해야 한다. 현대 LLM 학습은 GPipe식 pipeline parallelism에 tensor parallelism, data parallelism, ZeRO류 optimizer sharding을 함께 조합한다.

읽을 포인트

data parallelism과 model parallelism의 차이
pipeline bubble이 생기는 이유
micro-batch 수가 throughput에 주는 영향
이후 Megatron-LM, PipeDream, modern LLM training과의 연결

정리

GPipe는 거대 neural network를 학습하기 위한 parallelism 전략을 명확히 정리한 논문이다. 핵심은 모델을 나누는 것만으로는 부족하고, micro-batch scheduling으로 device utilization을 높여야 한다는 점이다.