[PyTorch] AMP (Automatic Mixed Precision)

일반적인 딥러닝 모델은 연산과 가중치 저장에 FP32를 사용한다.
하지만 AMP를 적용하면 네트워크 내부의 연산 특성에 맞춰 FP32와 FP16 또는 BF16을 자동으로 섞어서 사용하게 된다.
즉, VRAM을 절약하고 학습 속도를 증가시킬 수 있다.

AMP는 성능 저하 없이 연산을 가속하기 위해 크게 두 가지 메커니즘을 사용한다.

1. Autocast

   프레임워크가 연산의 종류를 파악해서, 안전한 경우에만 텐서를 FP16으로 변환하여 계산한다.

   행렬 곱셈이나 convolution과 같이 연산량이 많은 부분은 텐서 코어를 활용하기 위해 FP16으로 casting해서 빠르게 처리한다.

   반면 Softmax, Normalization, Loss 계산 등 미세한 차이가 결과에 큰 영향을 미치는 연산은 FP32를 그대로 유지하여 수치적 안정성을 확보한다.

2. Loss Scaling

   딥러닝 모델은 레이어가 깊어질수록 역전파를 통해 전달되는 그래디언트 값이 매우 작아지는 경향이 있다. (Chain Rule에 의해 그래디언트가 계속 곱해지기 때문임)
   FP16은 표현할 수 있는 숫자의 범위가 FP32보다 훨씬 좁기 때문에, 이렇게 작은 그래디언트들을 FP16으로 표현하면 값이 0으로 소멸해 버리는 Underflow 현상이 발생할 수 있다.

   이를 방지하기 위해 역전파 직전에 Loss 값에 특정 상수를 곱해 그래디언트 값을 안전한 범위로 키워서 계산하고, 옵티마이저가 가중치를 업데이트하기 직전에 다시 원래 비율로 나누어주는 과정을 거친다.

실제 구현

AMP를 적용할 때는 단순히 autocast로 Forward Pass 구간을 감싸고, 역전파와 옵티마이저 업데이트 과정에 scaler를 적용해 주면 된다.

import torch
from torch.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# GradScaler 객체 생성 (FP16의 Underflow를 방지하기 위함)
scaler = GradScaler()

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()

    # autocast로 Forward 패스를 감쌈 (이 안에서 일어나는 연산은 자동으로 FP16/FP32가 혼합되어 계산됨)
    with autocast(device_type='cuda'):
        output = model(data)
        loss = loss_function(output, target)

    # scaler를 사용해 loss를 스케일링하고 역전파
    scaler.scale(loss).backward()

    # scaler를 통해 optimizer 업데이트 및 scaler 업데이트
    scaler.step(optimizer)
    scaler.update()