Mixed Precision: autocast 와 GradScaler

메모리 절반, 속도 두 배 (자주)

Mixed-precision training 은 대부분 op 를 float16 또는 bfloat16 으로 돌리고 numerically sensitive op (softmax, normalization, loss reduction) 작은 set 만 float32 유지. Ampere+ GPU 와 Apple Silicon 에서 보통 1.5–2x speedup AND activation 메모리 약 절반.

modern PyTorch API 는 torch.amp (옛 torch.cuda.amp 는 deprecated). 두 재료:

autocast(device_type='cuda', dtype=...) — op 별 정밀도 선택하는 context manager.
GradScaler(device_type) — fp16 에만 필요. backward 전 loss 를 scale-up 해서 좁은 fp16 range 에서 gradient underflow 방지, optimizer step 전 unscale.

fp16 vs bf16

bfloat16 — float32 와 같은 exponent range, mantissa 정밀도 적음. Ampere+ NVIDIA GPU (A100, RTX 30/40, H100) 와 Apple Silicon. GradScaler 불필요. 추천 modern default.
float16 — 좁은 exponent range. GradScaler 필요. 옛 GPU (V100, T4, RTX 20-series) 때문에 여전히 관련.

실제로 wrap 하는 거

forward pass 와 loss 계산 wrap. backward 는 autocast 블록 밖. optimizer step 도 밖.

Code

bfloat16 mixed precision — modern recipe·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 mixed precision — 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

autocast 가 실제로 뭘 고르는지 inspect·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

어떤 training loop든 잡기. 한 epoch 을 float32 로, 그 다음 bfloat16 autocast 로 시간. Ampere+ hardware 에서 GradScaler 없이 1.3–2x speedup 봐야 함. loss curve 도 봐 — bf16 가 loss 더 noisy 하거나 나쁘게 만들면, 보통 명시적 float32 cast 필요한 numerically sensitive op.