혼합 정밀도 훈련
혼합 정밀도 훈련
혼합 정밀도(mixed precision) 훈련은 계산 집약적인 forward와 backward pass를 fp16 또는 bf16로 수행하면서, 옵티마이저 업데이트를 위해 fp32 복사본의 가중치를 유지합니다. 이를 통해 계산이 빨라지고 가중치와 활성화 메모리가 줄어들며 훈련 안정성이 유지됩니다.
출처: 문서
본문
완전 정밀도(fp32) 훈련은 모든 것을 32비트로 저장하고 계산합니다. 혼합 정밀도는 계산 집약적인 forward와 backward pass에서 fp16 또는 bf16을 사용하면서, 옵티마이저 업데이트를 위해 fp32 복사본의 가중치를 유지합니다. 계산이 더 빨라지고, 가중치와 활성화 메모리가 줄어들며, 훈련 안정성이 유지됩니다.
┌─────────────────────────────────────────────────────┐
│ MIXED PRECISION TRAINING LOOP │
│ │
│ fp32 master weights ──cast──▶ fp16/bf16 │
│ ▲ │ │
│ │ FORWARD (autocast) │
│ │ matmuls in fp16/bf16 │
│ │ reductions stay fp32 │
│ │ │ loss │
│ │ LOSS SCALE ×S ──fp16 │
│ │ │ │
│ │ BACKWARD │
│ │ grads in fp16/bf16 │
│ │ │ │
│ │ UNSCALE ÷S ──fp16 │
│ │ check inf/nan │
│ │ cast grads → fp32 │
│ └────────────────────────── optimizer.step │
└─────────────────────────────────────────────────────┘
혼합 정밀도 훈련을 활성화하려면 bf16() 또는 fp16()을 True로 설정합니다. 둘 다 16비트 타입이지만, bf16은 fp32와 동일한 지수 범위를 가져 거의 오버플로우가 발생하지 않습니다. Ampere 이상 GPU(A100, H100)에서는 bf16을 사용하고, V100이나 T4 같은 구형 하드웨어에서는 fp16으로 대체하세요.
[!WARNING] 모델을 fp32로 로드해야 합니다. 그렇지 않으면 autocast가 no-op이 됩니다. bf16이나 fp16으로 로드하면 옵티마이저가 업데이트할 fp32 마스터 복사본이 남지 않습니다.
from transformers import TrainingArguments
args = TrainingArguments(..., bf16=True)
args = TrainingArguments(..., fp16=True)
모델이 bf16/fp16에서 수치적으로 안정적이라면 혼합 정밀도를 건너뛰고 bf16/fp16으로 직접 로드하고 훈련할 수 있습니다. 이렇게 하면 메모리에 fp32 가중치 복사본이 생기지 않습니다.
tf32
tf32는 Ampere GPU에서 matmul에 23비트 대신 10비트 가수(mantissa)를 사용하는 계산 모드입니다. 특히 bf16/fp16과 함께 사용하면 속도 향상을 얻을 수 있습니다. PyTorch는 Ampere 이상 GPU에서 matmul에 기본적으로 tf32를 활성화하지만, TrainingArguments에서 명시적으로 설정하면 PyTorch 버전이나 환경 기본값과 무관하게 활성화됩니다.
from transformers import TrainingArguments
args = TrainingArguments(..., bf16=True, tf32=True)
다음 단계
- 커스텀 퓨전 커널로 훈련을 가속화하는 방법은 Kernels 가이드를 참조하세요.
- 추가 처리량을 위해 forward와 backward pass를 컴파일하는 방법은 torch.compile 가이드를 참조하세요.