[딥러닝] Loss Scaling 기법 + Mixed Precision

2024.02.05 - [프로그래밍 공부/ML 및 DL 관련 이론] - [논문] 임베디드에서의 float32와 float 16, 그리고 int8 비교(양자화)

[논문] 임베디드에서의 float32와 float 16, 그리고 int8 비교(양자화)

 

지난 시간의 논문 리뷰에서 양자화에 대해 설명했었다.

이때 Loss Scaling에 대해서 넘어갔었는데, 여기서는 Loss Scaling이 뭔지에 대해서 설명하기로 한다.

 

float 32와 float16의 본질적인 차이점을 알고 있어야 하므로, 위 글에 설명해 놓은 차이점을 보면 좋다. 아주 상세하게 설명해놓았다.

 

float32를 float 16으로 변화할 때 문제점

양자화 기법에서 float 32를 float 16으로 변화시킨다는 개념인 건 직관적으로 알 수 있다.

 

그러면 이렇게 변환 시켰을 때 발생하는 문제는 뭘까?

 

https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html

 

 

간단하다. float32의 경우, 최대 값이 2^32( (2 − 2^(−23)) × 2^127 ≈ 3.4028235 × 10^38. )이지만, float 16의 경우에는 2^16까지 가능하므로 최대한 가능한 최종 범위가 65536까지 밖에 안 된다.

그렇기에 훈련을 시켰을 때,gradient가 너무 크면 누적오차가 계속해서 발생한다는 문제점이 있다.

a = np.array([0.123456789121212,2,3], dtype=np.float16)
print("16bit: ", a[0])

a = np.array([0.123456789121212,2,3], dtype=np.float32)
print("32bit: ", a[0])

b = np.array([0.123456789121212121212,2,3], dtype=np.float64)
print("64bit: ", b[0])

• 16비트: 0.1235
• 32비트: 0.12345679
• 64비트: 0.12345678912121212

 

0.0000000596046 (highest, for values closest to 0) to 32 (lowest, for values in the range 32768-65536).

 

따라서 해당 범위를 Shifting 해주어서, 32비트의 범위처럼 사용할 수 있게 한다.

무슨 말이냐?

위 그래프에서 보듯이, fp16이 버틸 수 있는 지수 range는 -24 정도까지가 한계이다.(log 스케일이기 때문에)

그래서 그 아래의 지수들은 0이 되어버린다.

그런데 값을 보면, -24 부터 값이 많다. 그러면 어떻게 해야 하느냐?

 

간단하다. Shift 연산을 해서 값을 올리면 된다. 즉, -24 아래의 범위들을 -24 이상으로 전부 끌고 올라온다.

위의 그래프 범위에서는 -60부분의 끝자락을 denorm의 끝자락으로 맞추는 형태가 될 것이다.

 

이렇게 Shift 한 값을 float32와 같이 사용하기 때문에, 같이 사용한다는 의미로 Mixed Precision이라 한다.

 

그렇게 값을 올려도 차이가 없는 건가 하면, 다음 규칙에 따라 진행되기 때문에 차이가 없다.

1. FP32 Weight에 대한 FP16 copy weight을 만들기

 

2. FP16 copy weight을 이용해 forward pass(순전파)를 진행.

 

3. forwad pass로 계산된 FP16 prediction 값을 FP32로 casting

 

4. FP32 prediction을 이용해 FP32 loss를 계산하고, 여기에 scaling factor를 곱한다.

 

5. scaled FP32 loss를 FP16으로 casting

 

6. scaled FP16 loss를 이용하여 backward propagation(역전파)를 진행하고 gradient 계산.

 

7. FP 16 gradient를 FP32로 casting 한 후, scaling factor로 다시 나눈다.(chain rule에 의해서 모든 gradient는 같은 크기로 스케일링)

 

8. FP32 gradient를 이용해 FP32 weight을 업데이트

 

한 마디로 FP32 weight은 계속 저장해 두고, FP16 copy weihgt을 만들어서 순전파, 역전파 진행.

오차를 구할 때만 shift를 하고 나서, 다시 shift를 해서 업데이트 한다. 

FP16 copy weight으로 얻은 gradient로 FP32 weight을 업데이트 하기 때문에, 모델의 성능에는 거의 지장이 없다.(곱하거나 나누서 나오는 차이 정도)

 

출처 : https://bo-10000.tistory.com/32

 

따라서 이런 편리한 방법을 쓰면 오히려 성능이 더 좋아질 수 있다는 연구 결과였다.

 

mixed precision 기법은 pytorch에서 제공하고 있으므로, 시간이 난다면 한 번 시도 해보자.

단, Nan loss가 나는 경우도 있으므로 주의하자. 값의 범위를 생각하여 scaling factor를 잘 정해줘야 한다.