분산 LLM 학습 11 - Pipeline Parallel의 기본과 Stage 분할 감각
모델을 레이어 단위로 여러 stage에 나누는 순간 계산 분할뿐 아니라 idle time과 stage imbalance가 핵심 문제가 된다
pipeline parallel은 무엇을 나누는가
tensor parallel이 레이어 내부 연산을 분할한다면, pipeline parallel은 레이어 묶음 자체를 여러 stage로 나눈다. 예를 들어 48-layer transformer를 4개 stage로 나누면 각 stage가 12개 layer씩 담당하는 식이다.
이 전략은 모델 전체를 한 GPU에 두기 어려울 때 유용하다. 하지만 단순히 레이어를 등분한다고 끝나지 않는다.
가장 큰 문제는 idle time이다
pipeline parallel의 핵심 적은 bubble이다. 첫 micro-batch가 stage 1을 통과해 stage 2로 넘어갈 때까지 뒤쪽 stage는 빈다. 마지막 micro-batch가 앞 stage를 빠져나간 뒤에도 일부 stage는 다시 빈다.
즉 pipeline은 메모리를 나누는 대신 유휴 시간을 끌고 온다.
왜 micro-batch가 중요한가
pipeline은 큰 batch를 micro-batch로 쪼개서 stage들이 동시에 다른 조각을 처리하도록 만든다. micro-batch 수가 너무 적으면 bubble이 커지고, 너무 많으면 scheduling 오버헤드와 activation 관리가 복잡해진다.
따라서 pipeline parallel은 항상 다음 질문과 함께 간다.
- stage 수를 몇 개로 둘 것인가
- micro-batch 수를 얼마나 둘 것인가
- 각 stage의 계산량이 균형적인가
stage imbalance가 왜 위험한가
한 stage가 다른 stage보다 느리면 전체 pipeline 속도는 그 stage에 맞춰진다. 특히 embedding, final projection, 특정 attention-heavy layer가 비대칭이면 균형이 쉽게 무너진다.
그래서 단순한 레이어 수 균등 분할보다는 실제 계산량과 메모리 사용량을 기준으로 stage를 나누는 것이 더 중요하다.
pipeline parallel이 맞는 상황
- 모델 depth가 크다
- 레이어 구조가 비교적 규칙적이다
- stage 간 통신보다 메모리 절감 효과가 크다
반대로 node 간 지연이 크고 stage 균형이 좋지 않다면 pipeline은 운영 난이도만 올릴 수 있다.
다음 글에서는 pipeline schedule을 더 구체적으로 본다. GPipe, 1F1B처럼 자주 나오는 스케줄이 bubble과 activation 메모리에 어떤 차이를 만드는지가 핵심이다.