Mixed Precision Training

딥러닝 코드를 돌리다 보면 model capacity가 커짐에 따라 memory error가 빈번하게 일어나는 경우가 종종 있다. 설령 모델이 돌아간다고 해도 연산속도가 매우 오래걸리는 경우가 허다하다. batch size를 작게 설정하거나 input image를 crop하는 등의 어떠한 작업을 해주는 것이 아니라 모델 자체에서 보다 더 빠르고 메모리를 덜 차지하는 방향으로 연산을 하는 방법이 있을까?

오늘 소개할 아이디어는 FP16 비트로 학습을 계산하여 보다 적은 GPU 자원을 사용하는 Mixed Precision Trainining이다.

Mixed Precision Training

Deep Learning methodology가 발전하면서, NN의 사이즈를 늘리는 것이 더 좋은 퍼포먼스를 가져온다고 통상적으로 잘 알려져 있다. 하지만, 그렇게 될 경우 학습에 필요한 연산량과 메모리 또한 증가하게 된다. 따라서 학습시 memory requirement를 줄이는 것은 필수적이였으며, 2017년에 NVIDIA에서는 이를 해결해줄 "Mixed Precision Training"기법을 소개하였다.

We introduce methodology for training deep neural networks using half-precision floating point numbers, without losing model accuracy or having to modify hyperparameters. This nearly halves memory requirements and, on recent GPUs, speeds up arithmetic.

Mixed Precision Training을 설명하기 앞서, Floating Poing Precision에 대해 간단한 용어 정리부터 하겠다.

Floating Point Precision

Floating-point representation(부동소수점)이란 real number를 컴퓨터가 이해할 수 있는 근사값으로 표현하는 방법을 말하며 유효숫자를 나타내는 가수(mantissa)와 소수점의 위치를 풀이하는 지수(exponent)로 나누어서 표현한다.

컴퓨터에서는 보다 넓은 범위의 수를 나타낼 수 있어서 계산에 많이 이용되지만 근삿값으로 표현되며 연산속도가 느리다는 단점이 있다.

https://www.semanticscholar.org/paper/A-Study-on-Convolution-using-Half-Precision-Numbers-Seznec-Gac/7ff8896d21d574e499be4cc527fd44d9888b2d42

IEEE754 standard 기준에 따르면 실수를 저장하는 방법에는 크게 아래와 같은 종류가 있다.

FP16	1 sign bit, 5 exponent bits, 10 mantissa bits
FP32	1 sign bit, 8 exponent bits, 23 mantissa bits
FP64	1 sign bit, 11 expoent bits, 52 mantissa bits

표를 보면 이용하는 비트 수가 많을수록 더 높은 정확도로 실수를 저장할 수 있다. 예를 들어 FP32경우에는 최대 20억개의 다른 수를 표현할 수 있는가 하면 FP16은 최대 3만개의 수를 표현할 수 있다. 현재 대부분의 딥러닝 학습에서는 FP32을 사용하여 weight값을 저장하고 gradient을 계산한다. 만약, FP32 대신 FP16을 사용한다면 GPU의 자원을 보다 더 적게 사용할 수 있을 것이다.

하지만 FP16만을 사용해 학습을 진행한다면 FP32보다 연산속도가 빨라지는 대신 정밀도가 떨어질 수밖에 없다. 따라서 gradient가 과도하게 크거나 작을경우, 컴퓨터가 숫자를 저장할 때 오차가 발생하며 오차가 누적됨에 따라 결국 학습이 제대로 이루어지지 않을 것이다.

Figure 2b shows the histogram for the exponents of weight gradients for Mandarin speech recognition training with FP32 weights. The gradients are sampled every 4,000 iterations during training for all the layers in the model.

위 그림을 보면 파란색의 histogram bar들은 FP32로 학습을 해주었을때 gradient의 exponent value들의 분포를 보여준다. 만약 같은 학습을 FP16비트로만 진행을 한다면 그림의 빨간색 선을 기준으로 좌측부분은 (eponent value가 매우 적은 부분) 0이 될 것이다.

 

따라서 해당 논문에서는 FP32와 FP16을 함께 사용하는 mixed precision기법을 통해 이러한 문제점을 해결하였다.

Figure 3: Histogram of activation gradient values during the training of Multibox SSD network. Note that the bins on the x-axis cover varying ranges and there’s a separate bin for zeros. For example, 2% of the values are in the [2−34, 2−32) range, 2% of values are in the [2−24, 2−23) range, and 67% of values are zero.

다시 histogram 그림을 살펴보면 실제 표현가능한 FP16의 범위는 다음과 같기 때문에 -log75~-log24 사이의 magnitude를 가지는 gradient value들은 0으로 변할 것이다.  이를 방지하기 위해 scaling factor라고 불리는 특정 값을 곱하여 loss, gradient를 계산한 이후에 weight 값을 업데이트해준다. 이때 scaling factor는 오버플로우, 혹은 언더플로우가 일어나지 않게 실험적으로 잘 정해주게 된다.

 

Implementation

Implementation 방법은 위 그림과 같다. Mixed precision training에서 weight, activation 그리고 gradient들은 FP16으로 저장이 된다. FP32 network의 accuracy와 맞추기 위해 weight값의 FP32 master copy값이 유지가 되며 optimizer step에서 weight gradient에 의해 업데이트된다. 매 iteration마다 master weight의 FP16 copy가 forward와 backward pass에 사용되어 FP32 trainig에 필요한 storage를 절반으로 줄인다. 

다시 말하자면, FP32 weight값은 계속 저장해두면서, forward & backward propagation에서는 FP16 copy weight을 사용한다. 그리고 FP16 copy weight으로 얻은 gradient를 이용해 FP32를 update해준다.

 

Result

실험 결과, mixed precision을 사용했을 때 baseline과 비슷한 성능을 보다 더 빠르고 적은 메모리 사용으로 얻을 수 있음을 알 수 있다.

 

Pytorch Implementation

그렇다면 pytorch 환경에서 어떻게 mixed precision을 사용할 수 있을까? 매우 감사하게도 이미 pytorch에서는 해당 기법을 간단하게 적용할 수 있도록 AMP package를 제공하고 있다. 따라서 optimization이 이뤄지는 (training update가 이뤄지고 있는) Train 단계에서 scaler를 설정해주기만 하면 된다.

https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html

Automatic Mixed Precision — PyTorch Tutorials 1.10.1+cu102 documentation

https://pytorch.org/blog/accelerating-training-on-nvidia-gpus-with-pytorch-automatic-mixed-precision/

PyTorch

use_amp = True

net = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

start_timer()
for epoch in range(epochs):
    for input, target in zip(data, targets):
        with torch.cuda.amp.autocast(enabled=use_amp):
            output = net(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
        opt.zero_grad() # set_to_none=True here can modestly improve performance
end_timer_and_print("Mixed precision:")

 

 

[1] Micikevicius, Paulius, et al. "Mixed precision training." arXiv preprint arXiv:1710.03740 (2017).