Pipeline Parallelism
마지막 수정:
Pipeline Parallelism은 모델의 layer들을 여러 GPU 또는 노드에 나눠 배치하는 방법이다.
16 layers -> 4 pipeline stages L0 L1 L2 L3 L4 L5 L6 L7 L8 L9 L10 L11 L12 L13 L14 L15 S0 -> act -> S1 -> act -> S2 -> act -> S3 S3 -> grad act -> S2 -> grad act -> S1 -> grad act -> S0 (p - 1) / m m forward activations wait for backward after Bi, activation i can be freed warmup + cooldown remain activations / PP each stage owns fewer layers PP x (activations / PP) ~= activations several microbatches are stored at once 1F1B start backward earlier to free activations sooner V0: L0-L1 V4: L8-L9 V1: L2-L3 V5: L10-L11 V2: L4-L5 V6: L12-L13 V3: L6-L7 V7: L14-L15 smaller chunks -> tighter schedule more stage boundaries and sends TP는 한 layer 내부의 tensor 계산을 나눈다. PP는 layer 깊이 방향으로 모델을 나눈다.
TP: one layer inside -> split tensor dimensions
PP: whole model depth -> split layer groups
예를 들어 16개 layer를 4개 stage로 나누면 다음과 같다.
Stage 0: layers 0-3
Stage 1: layers 4-7
Stage 2: layers 8-11
Stage 3: layers 12-15
각 stage는 자기 layer의 parameters만 저장하고 계산한다. 그래서 모델 weight 자체가 너무 커서 한 노드의 TP로도 감당하기 어려울 때 PP가 필요하다.
PP의 통신은 TP와 다르다
TP는 layer 안에서 all-reduce, all-gather, reduce-scatter 같은 collective가 자주 발생한다.
TP:
layer 내부에서 자주 통신
노드 간으로 커지면 성능 저하가 큼
PP는 stage 사이에서 activation과 gradient를 넘긴다.
forward:
Stage 0 -> activation -> Stage 1 -> activation -> Stage 2
backward:
Stage 2 -> grad activation -> Stage 1 -> grad activation -> Stage 0
그래서 보통 빠른 단일 노드 안에서는 TP/SP를 쓰고, 노드 사이로 모델을 더 키울 때 PP를 쓴다.
naive PP와 bubble
microbatch 하나만 넣으면 PP는 순차적으로 실행된다.
time 1: Stage 0 works
time 2: Stage 1 works
time 3: Stage 2 works
time 4: Stage 3 works
한 stage가 계산할 때 다른 stage는 기다린다. 이 idle time을 pipeline bubble이라고 한다.
bubble = pipeline stage가 계산하지 못하고 기다리는 시간
stage 수가 많아질수록 bubble도 커진다.
bubble ratio without microbatches ~= p - 1
여기서 p는 pipeline stage 수다.
microbatch로 pipeline 채우기
PP는 batch를 여러 microbatch로 나눠서 pipeline을 채운다.
global batch = microbatch size x number of microbatches
microbatch는 실제로 더 작은 batch 조각이다. 하지만 여러 microbatch의 gradient를 누적한 뒤 optimizer step을 한 번만 하면 global batch size는 유지된다.
for mb in microbatches:
forward(mb)
backward(mb)
accumulate gradients
optimizer.step()
microbatch가 있으면 mb0가 다음 stage로 넘어간 뒤 이전 stage는 mb1을 처리할 수 있다.
time 1: S0 mb0
time 2: S0 mb1 | S1 mb0
time 3: S0 mb2 | S1 mb1 | S2 mb0
이렇게 pipeline이 차면 여러 stage가 동시에 일한다.
microbatch 수를 m개로 늘리면 bubble의 상대적 비율은 대략 줄어든다.
bubble ratio with m microbatches ~= (p - 1) / m
AFAB: all forward, all backward
가장 단순한 pipeline schedule은 AFAB다.
AFAB = All Forward, All Backward
모든 microbatch의 forward를 먼저 수행하고, 그 뒤 모든 backward를 수행한다.
forward mb0
forward mb1
forward mb2
...
backward mb2
backward mb1
backward mb0
장점은 구현이 쉽다는 것이다. 일반적인 학습 코드처럼 forward를 먼저 쌓고 backward를 나중에 시작한다.
하지만 activation memory가 커진다.
왜 PP만으로 activation memory가 줄지 않는가
PP는 layer를 나누므로, microbatch 하나 기준으로 stage 하나가 저장하는 activation은 줄어든다.
one microbatch activation per stage ~= activations / PP
하지만 첫 backward가 시작되기 전까지 여러 microbatch의 forward activation을 저장해야 한다.
PP microbatches x (activations / PP) ~= activations
즉 한 microbatch당 activation은 작아졌지만, 여러 microbatch activation을 동시에 들고 있어서 절감이 상쇄된다.
AFAB에서는 이 문제가 더 크다. 모든 forward를 끝내기 전까지 backward를 시작하지 않으므로, 많은 microbatch activation을 지울 수 없다.
AFAB:
pipeline bubble은 줄일 수 있음
하지만 activation 저장량이 커짐
다음 단계는 forward와 backward를 섞어서 activation을 빨리 해제하는 schedule이다.
1F1B = one forward, one backward
1F1B: one forward, one backward
1F1B는 pipeline이 어느 정도 찬 뒤 forward와 backward를 번갈아 수행하는 schedule이다.
warmup:
pipeline을 채우기 위해 forward를 몇 개 먼저 실행
steady state:
forward 하나, backward 하나를 번갈아 실행
cooldown:
남은 backward를 마무리
AFAB와 가장 큰 차이는 backward를 더 빨리 시작한다는 점이다.
AFAB:
F0 F1 F2 F3 ... 이후 B3 B2 B1 B0
1F1B:
pipeline warmup 이후 F와 B를 섞음
Backward가 끝난 microbatch의 activation은 바로 버릴 수 있다.
mb0 forward activation 저장
mb0 backward 완료
-> mb0 activation release
그래서 1F1B는 AFAB보다 activation memory를 줄인다.
AFAB:
많은 microbatch activation이 backward 전까지 대기
1F1B:
각 stage가 가능한 빨리 backward를 시작해서 activation을 해제
하지만 bubble이 완전히 사라지는 것은 아니다. pipeline 처음에는 채워야 하고, 마지막에는 비워야 한다.
remaining bubble:
warmup bubble + cooldown bubble
그래도 steady state에서는 대부분의 stage가 동시에 일한다.
Stage 0: forward on new microbatch
Stage 1: forward/backward
Stage 2: forward/backward
Stage 3: backward on older microbatch
따라서 1F1B의 핵심은 처리량만이 아니라 activation lifetime을 줄이는 데 있다.
Interleaved 1F1B: virtual pipeline stages
1F1B에도 warmup과 cooldown bubble은 남는다. stage 수가 많거나 microbatch 수가 충분하지 않으면 이 빈 공간이 여전히 아깝다.
Interleaved 1F1B는 하나의 물리 GPU가 하나의 큰 stage만 맡지 않게 한다. 대신 layer를 더 작은 chunk로 쪼개고, 각 GPU가 여러 virtual pipeline stage를 맡는다.
기본 PP가 이런 배치라면:
GPU 0: layers 0-3
GPU 1: layers 4-7
GPU 2: layers 8-11
GPU 3: layers 12-15
interleaved PP는 이렇게 배치할 수 있다.
GPU 0: layers 0-1, 8-9
GPU 1: layers 2-3, 10-11
GPU 2: layers 4-5, 12-13
GPU 3: layers 6-7, 14-15
물리 GPU는 4개 그대로지만, pipeline stage는 더 잘게 나뉜다.
physical stages = 4
virtual stages = 8
이렇게 하면 각 stage의 작업 단위가 작아지고, forward/backward 조각을 더 촘촘하게 끼워 넣을 수 있다. 그래서 warmup/cooldown에서 생기는 bubble을 더 줄일 수 있다.
대신 공짜는 아니다.
장점:
bubble 감소
stage 작업이 더 촘촘하게 interleave됨
비용:
스케줄이 복잡해짐
stage boundary가 늘어 activation send/recv가 더 자주 발생할 수 있음
각 GPU가 여러 layer chunk를 관리해야 함
따라서 interleaving은 “PP를 더 잘게 쪼개서 GPU idle time을 줄이는 scheduling 기법”으로 보면 된다.
확인
- PP는 TP와 어떤 축이 다른가?
- PP에서 stage 사이에 주로 오가는 tensor는 무엇인가?
- pipeline bubble은 왜 생기는가?
- microbatch를 늘리면 bubble 비율이 왜 줄어드는가?
- AFAB가 activation memory를 많이 쓰는 이유는 무엇인가?
- 1F1B는 AFAB와 비교해 activation을 언제 더 빨리 해제할 수 있는가?
- interleaved 1F1B에서 virtual pipeline stage는 무엇인가?