혼합 정밀도: 자동 형변환과 GradScaler

메모리는 절반, 속도는 두 배 가까이

혼합 정밀도 학습은 대부분의 연산을 float16이나 bfloat16으로 실행하고, 소프트맥스·정규화·손실 축소처럼 수치 오차에 민감한 일부 연산만 float32로 유지해. Ampere 이후 GPU와 Apple Silicon에서는 보통 1.5~2배 빨라지고 활성화 메모리는 약 절반으로 줄어.

현대적인 PyTorch API는 torch.amp야. 예전의 torch.cuda.amp는 사용 중단됐어. 구성 요소는 두 가지야:

autocast(device_type='cuda', dtype=...): 연산마다 알맞은 정밀도를 선택하는 문맥 관리자야.
GradScaler(device_type): fp16에서만 필요해. 역전파 전에 손실 크기를 키워서 좁은 fp16 범위에서 기울기가 언더플로되지 않게 하고, 옵티마이저가 갱신하기 전에 원래 크기로 되돌려.

fp16과 bf16

bfloat16: float32와 같은 지수 범위를 갖지만 가수부 정밀도는 낮아. Ampere 이후 NVIDIA GPU(A100, RTX 30/40, H100)와 Apple Silicon에서 쓸 수 있어. GradScaler는 필요하지 않아. 현대적인 기본 선택으로 권장해.
float16: 지수 범위가 좁아서 GradScaler가 필요해. V100, T4, RTX 20 시리즈 같은 이전 세대 GPU에서는 여전히 중요해.

어디를 감싸야 할까

자동 형변환 문맥으로 순전파와 손실 계산을 감싸. 역전파와 옵티마이저 갱신은 그 문맥 밖에서 실행해.

Code

bfloat16 혼합 정밀도: 현대적인 구성법·python

import torch
import torch.nn as nn
from torch.amp import autocast

model = nn.Linear(1024, 1024).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.bfloat16):
        out = model(x)
        loss = loss_fn(out, y)

    loss.backward()        # no GradScaler needed for bf16
    optimizer.step()

float16 혼합 정밀도: GradScaler 필요·python

import torch
import torch.nn as nn
from torch.amp import autocast, GradScaler

model = nn.Linear(1024, 1024).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
scaler = GradScaler('cuda')

for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.float16):
        out = model(x)
        loss = loss_fn(out, y)

    scaler.scale(loss).backward()      # scale up to avoid underflow
    scaler.step(optimizer)              # unscale + step (skips on inf/nan)
    scaler.update()                     # adjust scale based on overflow

자동 형변환이 실제로 뭘 고르는지 살펴보기·python

import torch
from torch.amp import autocast

x = torch.randn(8, 8, device='cuda')

with autocast(device_type='cuda', dtype=torch.bfloat16):
    y_matmul = x @ x        # bfloat16 (fast)
    y_softmax = torch.softmax(x, dim=-1)  # promoted to float32 (numerically sensitive)
    print(y_matmul.dtype, y_softmax.dtype)
# torch.bfloat16 torch.float32

# autocast keeps a built-in op→dtype map. You can usually trust it,
# but for custom ops you may need to specify behavior explicitly.

Exercise

학습 반복문 하나를 골라 한 에포크는 float32로, 다음 에포크는 bfloat16 자동 형변환으로 실행 시간을 재 봐. Ampere 이후 하드웨어라면 GradScaler 없이도 보통 1.3~2배 빨라져. 손실 곡선도 함께 확인해. bf16에서 손실이 더 불안정하거나 나빠진다면 수치 오차에 민감한 연산을 명시적으로 float32로 변환해야 할 가능성이 커.