Single-GPU Training Step Review

학습 병렬화를 이해하려면 먼저 병렬화하지 않은 한 번의 학습 step을 정확히 잡아야 한다.

1 micro-batch training data

->

2 forward loss + activations

->

3 backward gradients

->

4 optimizer updated weights

한 optimizer step은 보통 다음 순서로 진행된다.

micro-batch -> forward -> loss -> backward -> gradients -> optimizer update

forward는 입력에서 loss까지 값을 만든다. 이때 backward에서 다시 쓸 activation도 저장한다.

backward는 loss에서 시작해 모델을 거꾸로 지나가며 각 parameter의 gradient를 만든다.

optimizer는 gradient를 사용해 parameter를 업데이트한다.

parameters + gradients + optimizer states -> updated parameters

이 카드에서 중요한 점은 DP나 ZeRO가 갑자기 새로운 학습 과정을 만드는 것이 아니라는 것이다. 둘 다 이 기본 step 위에 올라간다.

병렬화가 건드리는 위치

학습 병렬화는 이 흐름 중 일부를 여러 GPU에 나눠 맡기는 일이다.

데이터를 나눌 것인가?
parameter를 나눌 것인가?
gradient를 나눌 것인가?
optimizer state를 나눌 것인가?
layer를 나눌 것인가?

DP는 먼저 데이터부터 나눈다. ZeRO는 DP 위에서 gradient, optimizer state, parameter의 중복 저장을 줄인다. TP와 PP는 계산 구조 자체를 더 적극적으로 나눈다.

확인

• backward pass는 forward와 같은 방향으로 진행되는가, 반대 방향으로 진행되는가?
• optimizer update 전에 반드시 준비되어야 하는 것은 무엇인가?
• DP나 ZeRO도 결국 어떤 기본 학습 step 위에 올라가는가?