Loss Scaling과 Master Weights

FP16 mixed precision training에서 대표적인 안정화 기법은 세 가지다.

1. FP32 master weights
2. loss scaling
3. FP32 accumulation

FP32 master weights

모델 forward/backward는 FP16으로 빠르게 돌릴 수 있다. 하지만 실제 weight update를 FP16 weight에 바로 누적하면 작은 update가 반영되지 않을 수 있다.

weight = 1.0
update = 0.00001

FP16 resolution이 이 update를 표현하기에 너무 거칠면, 더해도 weight가 그대로일 수 있다.

그래서 더 높은 precision의 master copy를 둔다.

FP16 weight:
  forward/backward에 사용

FP32 master weight:
  optimizer update에 사용

gradient는 1보다 훨씬 작은 값이 될 수 있다.

gradient = 0.0000001

이 값이 FP16에서 너무 작으면 0으로 underflow될 수 있다. loss scaling은 backward 전에 loss를 크게 만들어 gradient도 같이 키운다.

scaled loss = loss x scale
backward
scaled gradients 생성
gradients = scaled gradients / scale
optimizer step

scale을 다시 나누기 때문에 수학적으로는 원래 gradient와 같은 방향을 유지한다. 목적은 backward 중에 작은 값이 0으로 사라지는 것을 막는 것이다.

sum, average, reduction 같은 연산은 중간 결과가 커지거나 작아질 수 있다. 그래서 중간 누적은 FP32로 하고, 최종 결과만 낮은 precision으로 cast하는 방식이 자주 쓰인다.

FP16 inputs
-> accumulate in FP32
-> cast result back to FP16/BF16