PyTorch AMP - 1 : Mixed Precision Training 논문 리뷰

PyTorch AMP

최근 발표되는 Tensor core(https://computing-jhson.tistory.com/10)와 같은 딥러닝 하드웨어에서는 기존의 single precision floating point(FP32) data type이 아닌, half precision floating point(FP16) data type과 같은 좀 더 적은 개수의 bits를 사용하는 데이터 타입을 이용해 layer weights, tensor 등의 데이터를 표현하고 있다. 이를 통해 메모리 사용량 및 데이터 전송량을 줄이면서, 회로를 최적화하여 딥러닝 학습 및 추론 속도를 향상을 달성하고 있다. 많은 딥러닝 프레임워크에서 이러한 low precision 연산을 딥러닝 학습에 지원하고 있는데, PyTorch 프레임워크에서는 AMP package를 통해 NVIDIA gpu 상에서 FP16 연산을 이용한 학습을 지원한다.

 

AMP는 Automatic Mixed Precision의 약자로, 딥러닝 네트워크 학습 시 자동으로 FP16 혹은 FP32 데이터타입을 혼합하여 사용하도록 돕는다. 사실 딥러닝 네트워크 학습에 사용되는 모든 데이터를 FP16로 표현할 경우 네트워크 정확도가 떨어질 수 있다. 이는 FP16 데이터 타입의 한계로, FP16은 최대 65,504까지 표현가능하며, 최소로 표현가능한 정밀도는 2^(-24)이기 때문이다. 이 때문에 큰 값을 가지거나 높은 정밀도를 요구하는 계산(sigmoid, tanh 연산, loss function 등)은 기존의 FP32 데이터를 사용하고, 작은 값을 가지거나 높은 정밀도를 요구하지 않는 계산(convolution, pooling 등의 연산)에서는 FP16 데이터를 사용한다. 이렇게 필요에 따라 두 데이터타입을 혼합해 사용하는 학습을 mixed precision learning이라고 하며, AMP는 이를 자동으로 해준다.

 

오늘은 AMP의 기본이 되는 아이디어를 제안한 "Mixed Precision Training"[1]에 대하여 정리해보고자 한다. Mixed precision training 자체는 이전에도 제안되었지만, 이 논문은 이것 뿐만 아니라 FP32 데이터 타입 네트워크 학습에 사용되었던 hyper-parameter를 조절하지 않고도 FP32 네트워크만큼의 정확도를 유지할 수 있도록 추가적인 학습 방법을 제안한다.

 

Fig 1. PyTorch 로고

 

논문 요약

• 제안 하는 것
  • Methodology for training deep neural networks using half-precision floating point numbers, without losing model accuracy or having to modify hyperparameters
  • FP16 weights, activation, gradient를 사용하는 딥러닝 네트워크 학습 방법 제안
  • 이때, 모델 정확도를 유지하면서 hyper-parameter tuning도 필요없게 하고자 함
• 문제 상황
  • FP16 데이터를 통한 학습은 계산 속도 향상, 저장 공간 감소 및 에너지 사용량 감소의 장점이 있지만 데이터의 정밀도가 낮아짐
  • 이에 따라 모델 정확도가 떨어지거나 hyper-parameter를 tuning하여 새로운 학습 전략을 세워야 함
• 문제 상황을 해결하기 위해 방법 3가지를 제안함
  • Weights의 경우 FP32으로 저장. 높은 정밀도를 이용해 작은 gradient도 weights 업데이트에 반영될 수 있도록 함. 딥러닝 계산 시 사용되는 weights는 FP32에서 FP16으로 변환하여 사용
  • Loss-scaling 기법 제안. Back-propagation 중 작은 값을 가지는 gradient가 0이 되는 것을 방지함
  • FP16 연산의 중간 결과(partial sum)를 FP32 output에 저장함. 실제 메모리에 저장할 때는 다시 FP16으로 변환하여 저장. 중간 결과에 의한 오차를 줄여 전체 정확도를 높이고자 함.
• 장점
  • 메모리 사용량 감소, 데이터 전송 속도 증가, 에너지 사용량 감소
  • Tensor core와 같은 FP16 가속 하드웨어를 이용할 경우 계산 속도 향상

Fig 2. Mixed precision training 방법 [1]

논문이 제안하는 내용을 정리하자면 Fig 2.와 같다. Fig 2.는 deep learing training의 한 iteration을 묘사하는 그림이다. 이 논문에서는 deep learning layer(주로 convolution과 fully connected layer)의 입력 tensor(activation), 출력 tensor(activation), weights, gradient(activation grad)를 FP16 데이터타입을 이용해 저장한다. 따라서 forward pass와 backward pass의 계산이 FP16 연산이 되어 빠르게 계산할 수 있다. 다만 실제 weights(master-weights)는 FP32로 유지하여 gradient에 의해 weight update 시 좀 더 높은 정밀도로 weights를 업데이트할 수 있도록 한다.

 

 

 

제안 1. Master Copy Of Weights

딥러닝 네트워크의 weights(master-weights라고 함)를 FP32 데이터 타입으로 저장한다. 이를 통해 FP16이 표현하지 못하는 작은 gradient까지 weight update에 반영하여, 기존의 FP32 네트워크 학습과 거의 동일한 학습 패턴을 보이고자 한다.

 

Fig 2.는 SSD network 학습 중 계산된 모든 gradient의 정밀도 기준으로 histogram으로 표현한 것이다. FP16로 표시할 수 없는 정밀도를 가진 gradient는 0이 되는데, 반 이상의 gradient 값(Become zero in FP16 영역에 속하는 값)이 0이 되는 것을 확인할 수 있다. 따라서 이러한 gradient update가 제대로 되기 위해서 네트워크의 실제 weights를 FP32 데이터 타입으로 저장하고, FP32 연산으로 weight update를 진행한다.

 

앞서 말했듯, 딥러닝 레이어에서의 forward와 backward는 FP16 연산자를 사용하므로, FP32 master weights를 FP16 weights 변환하여 저장해놓아야 하는 단점이 있다. 논문에서도 언급하는데, 이 경우에는 50%의 추가적인 메모리 공간(FP32+FP16 둘 다 저장해야 하기에)을 차지하게 된다. 또한 매 iteration마다 업데이트된 master weights를 FP16으로 타입 변환하여 다시 저장해야 하기에 overhead가 발생한다. 다만 학습에서 weights의 크기보다 input & output tensors(activation)이 훨씬 크기에 전체 메모리 사용량은 대략적으로 50%에 수렴한다고 한다.

Fig 2. SSD network 학습 중 계산된 gradient의 값의 정밀도에 따른 histogram [1]

 

 

 

제안 2. Loss Scaling

제안 1과 같은 맥략으로 gradient는 FP16으로 표현하기에는 많은 수의 gradient가 0으로 표현되는 문제가 발생한다. Fig 2. 에서 히스토그램 중 "Become zero in FP16"에 속한 bin들에 속한 값들은 0이 된다. 하지만 만약 저 히스토그램으로 "FP16 Representable range"에 들어가도록 right shift 시키면 FP16으로 모든 gradient를 표현할 수 있다. 그 아이디어가 제안 2. loss scaling이다.

 

Loss에 큰 값 α를 곱한다(즉 α만큼 scaling한다). Back-propagation 과정에서 chain rule에 의해 loss에 곱해진 scalar α의 값만큼 모든 gradient value에 똑같은 크기 α만큼 곱해진 결과가 나온다. 단순히 loss에 α를 곱하기만 했을 뿐인데 모든 gradient가 α만큼 scaling된 결과가 된다. 이는 Fig 2.의 histrogram이 오른쪽으로 log_2(α)만큼 right shift만 결과를 만들어낸다. 따라서 적당한 α값을 선택하여 loss를 scaling하면 모든 gradient를 FP16으로 표현할 수 있다. 실제 weight update시에는 gradient를 FP32로 변환한 후 α으로 나눠준 값을 이용해 update하면 된다.

 

scaling factor α를 선택하는 것이 성능 상에 중요한 영향을 미치는 또 다른 hyper-parameter가 되는 한계가 있다. 가장 쉬운 방법은 실험적으로 α 값 하나를 선택하는 것이다. α 선택은 FP16이 표현할 수 있는 값의 범위인 65,504를 넘기지 않도록 선택하기만 하면 된다고 한다. 개인적인 생각인데 FP16은 값이 커질수록 정밀도가 작아지기에 α가 커서 scaled gradient의 값이 커지면 오차가 쌓이는 것 아닌가 하는 생각도 든다.

 

 

 

제안 3. Arithematic Precision

FP16 연산(convolution이나 matrix multiplicatin 등)의 partial sum을 FP32 타입으로 설정하여, 연산 정확도를 높이겠다는 아이디어이다. 

 

FP16 연산의 경우 FP32에 비해 정밀도의 한계에 의해 어쩔수 없이 약간의 오차가 생기는데, 오차가 있는 partial sum들의 reduction은 더 큰 오차를 만든다. 따라서 partial sum을 FP32로 나타내어 오차를 제거하고, 최종 결과만을 FP16로 변환해 저장하면 오차를 줄일 수 있다. 예를 들어 1.0001을 만번 더하는 연산에서, 1.0001을 FP16으로 저장하면 1이 되기에 만 번 더한 결과는 10000이 되는데 비해, FP32로 저장하여 만 번 더하면 10001이 된다. 이처럼 딥러닝 연산에서 partial sum들의 reduction이 많기에 partial sum만을 FP32로 나타내면 오차를 줄일 수 있다.

 

Partial sum은 결국 register에 저장되어 있을 것이기에 메모리 접근 속도의 제한도 크게 없을 것 같으며, FP16연산보다 FP32연산이 더 오래걸리겠지만 덧셈 연산이기에 크게 overhead가 발생하지 않을 것이라고 예상된다.

 

 

 

성능

FIg 3. ILSVRC12 classificion model에 대한 mixed precision learning의 결과 [1]

Fig 3.에 따르면 다양한 딥러닝 모델에 대하여 mixed precision learning의 정확도가 떨어지지 않는 것을 확인할 수 있다. 앞선 3가지 제안이 성능 정확도 유지에 도움이 된 것으로 파악된다.

 

Fig 4. bigLSTM training 과정. loss scaling의 영향을 확인할 수 있음 [1]

Fig 4.는 language modeling task에 사용된 bigLSTM의 학습 과정을 나타낸 그림이다. Loss scale이 적용되지 않을 경우(loss scale 1) 성능이 converge하지 않지만, loss scale이 128인 경우 FP32 학습과 비슷한 양상을 보인다. Loss scale이 안정적인 학습에 기여한다는 것을 알 수 있다.

 

 

 

논문에 대한 개인 생각

FP32이 아닌 mixed precision 연산을 딥러닝 네트워크에 적용해도 성능을 유지할 수 있는 3가지 방법을 제안하는 중요한 논문이다. Mixed precision (혹은 FP16) learning은 계산 속도 향상 및 에너지 사용량을 줄일 수 있는 방법이지만, 네트워크의 정확도를 떨어뜨릴 수 있다. 하지만 여기서 제안하는 3가지 기법은 이러한 한계를 극복하게 해줄 수 있는 좋은 참고 자료가 될 것이다. Tensor core, TPU 등과 같은 저정밀도 연산을 사용하는 딥러닝 프로세서가 발전하는데 큰 역할을 할 것이라고 생각한다.

 

다만 quantization 기법 등 FP32뿐만 아니라 저정밀도 데이터타입만으로 딥러닝 네트워크를 구성해도 학습이 된다는 것을 보이는 연구가 진행되고 있다. 따라서 FP16 연산만을 통해 학습한 결과와 비교해보면 더 좋은 인사이트를 얻어볼 수 있을 것 같다.

 

 

 

Reference

[1] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. 2018. Mixed precision training. In International Conference on Learning Representations.

[2] https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#mptrain

[3] https://pytorch.org/docs/stable/amp.html