Skip to main content

​
Mixed Precision 학습

혼합 정밀도(Mixed Precision) 학습은 FP32와 FP16(또는 BF16)을 혼합 사용하여 학습 속도를 약 2배 가속하고 GPU 메모리를 절약하는 기법입니다.

​
핵심 아이디어

순전파와 역전파의 대부분의 연산을 FP16으로 수행하되, 가중치 업데이트와 손실 계산은 FP32로 유지합니다. NVIDIA의 Tensor Core를 활용하면 FP16 행렬 연산이 FP32보다 훨씬 빠릅니다.

​
구현

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

device = torch.device('cuda')
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
scaler = GradScaler('cuda')  # 기울기 스케일링 (FP16 언더플로 방지)

for epoch in range(10):
    model.train()
    for batch_data, batch_labels in train_loader:
        batch_data = batch_data.to(device)
        batch_labels = batch_labels.to(device)

        optimizer.zero_grad()

        # autocast: FP16으로 순전파
        with autocast('cuda'):
            outputs = model(batch_data)
            loss = criterion(outputs, batch_labels)

        # 스케일된 역전파
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

​
FP16 vs BF16

속성FP32FP16BF16
비트 수321616
지수 비트858
가수 비트23107
범위넓음좁음FP32 동일
정밀도높음높음낮음
GradScaler 필요-필요불필요
지원 GPU모든 GPUVolta+Ampere+
# BF16 사용 (Ampere 이상 GPU, GradScaler 불필요)
with autocast('cuda', dtype=torch.bfloat16):
    outputs = model(batch_data)
    loss = criterion(outputs, batch_labels)

loss.backward()
optimizer.step()
NVIDIA A100, H100 등 Ampere 이상 GPU에서는 BF16을 사용하세요. GradScaler가 불필요하고 수치적으로 더 안정적입니다.

​
체크리스트

  • autocast와 GradScaler의 역할을 이해한다
  • FP16과 BF16의 차이를 설명할 수 있다
  • Mixed Precision을 기존 학습 루프에 적용할 수 있다

​
다음 문서

기울기 누적

Mixed Precision과 결합하여 메모리 효율 극대화

분산 학습

여러 GPU로 학습을 분산하는 방법