GPU TensorCore를 최대한 이용하기 위한 Mixed Precision Training
어떻게 training하느냐에 따라 모델의 성능이 달라진다….
모델 학습을 학습하기 위해 대부분의 AI framework는 FP32로 학습한다. 최근 모델 사이즈가 급격히 증가하면서 모델 학습에 필요한 연산량과 메모리 용량이 또한 급격하게 증가하고 있다.
Mixed Precision Training은 학습에 FP32 대신 FP16을 사용한다는 아이디어로학습 연산량과 메모리 용량을 줄이기 위해 제안되었다. A100 GPU의 경우, FP32의 성능이 19.5 Tflops인 반면, TensorCore로 실행되는 FP16의 성능은 312 Tflops로 무려 16배 큰 엄청난 성능 향상 효과를 볼 수 있다.
하지만 FP32를 FP16으로 바꿔서 학습시킬 경우 , 아래 그림과 같이 training loss가 중간에 증가하는 경향을 보인다. 그 이유는 back-propagation 과정에서 gradient를 계산할 때gradient 값이 너무 크거나, 너무 작은 경우 FP16이 표현할 수 있는 범위가 FP32에 비해 협소하여 오차가 발생하고 이 오차가 점점 누적되기 때문이다.
NVIDIA는 Baidu 연구진과 함께 위와 같은 문제점을 해결한 Mixed-Precision Training(ICLR 2018)을 제안하였다. 아래 그림은 Mixed Precision Training를 이용하여 학습 시 사용되는 precision을 나타낸다.
Mixed Precision Training의 기본 원리는 Master-Weight를 FP32에서 FP16으로 변환 후(float2half) Forward/Back-propagation을 FP16으로 연산한다. 이때 계산된 gradient(FP16)로 부터 weight를 업데이트할 때 FP32로 변환되어master weight를 업데이트한다.
NVIDIA는 이 때 FP16을 그대로 사용하면 앞에서 설명한 training loss가 커지는 현상이 발생하여 학습이 제대로 이루어지지 않기 때문에 이를 방지하기 위한 loss scaling 기법을 제안하였다.
Loss scaling 기법이란 아래 그림과 같이 FP16으로 gradient를 연산할 때 gradient가 매우 작을 때 FP16은 이를 표현하지 못하므로 ‘0’이 되는 부분 때문에 발생하는 오차를 방지하기 위해 FP16이 표현할 수 있는 범위로 gradient를 “shift” 해주는 것이다.
Gradient를 FP16이 표현할 수있는 범위로 shift하기 위해선 간단히 scaling factor(S)를 곱하면 가능하다.
- 1. Master weight(FP32)로부터 copy weight(FP16)을 만든다.
- 2. Copy weight(FP16)을 이용해 forward 연산을 실행한다.
- 3. Forward pass에서 계산된 Prediction 값(FP16)을 FP32로 casting한다.
- 4. FP32로 casting된 값으로 Loss를 계산하고, scaling factor를 곱한다.
- 5. FP32 loss를 FP16으로 casting한다.
- 6. scaled FP16 loss를 이용하여 backward propagation을 진행하고, gradient를 계산한다.
- 7. FP16 gradient를 FP32로 casting한고, 이를 scaling factor S로 다시 나눈다.
- 8. FP32 gradient를 이용해 FP32 weight를 update한다.
레퍼런스
[1] Training with Mixed Precision
[2] Mixed Precision Training
