대규모 모델 학습·추론 최적화 시리즈 1편: Mixed Precision Training (AMP)

"학습 속도를 2배 이상 올리고, VRAM도 절약할 수 있다?"
Mixed Precision Training은 FP16/BF16과 FP32를 적절히 섞어 쓰는 학습 기법입니다.

---

 

AMP

 

Mixed Precision Training이란?

기존 딥러닝 학습은 모든 연산을 FP32(32-bit float)로 수행했습니다.
하지만 FP16(16-bit float) 연산은 속도가 더 빠르고, 메모리 사용량도 절반입니다.

Mixed Precision Training은 다음을 결합합니다:

• FP16/BF16로 대부분의 연산 (행렬 곱, convolution 등)
• FP32로 중요한 부분 (loss 계산, weight update, scale 유지)

결과:

성능 거의 그대로 유지 + 속도 1.5~3배 향상 + 메모리 절감

---

왜 필요한가?

• 대규모 모델 학습 시 GPU 메모리 부족 문제 해결
• 더 큰 batch size로 학습 가능 → 안정적 gradient & 더 빠른 convergence
• 최신 GPU (Ampere, Hopper)는 FP16/TF32/Tensor Core 가속 최적화

---

동작 원리

1. 모델 & 데이터 FP16 변환
2. Loss Scaling
   • FP16은 표현 가능한 값 범위가 작아 underflow 위험
   • 따라서 loss를 일정 배수로 scale → gradient 계산 후 다시 scale down
3. Weight Update
   • FP32 master weight 유지 → 정밀도 손실 최소화

---

PyTorch 예시 (Automatic Mixed Precision)

 

import torch
from torch.cuda.amp import autocast, GradScaler

model = get_model().cuda()
optimizer = torch.optim.AdamW(model.parameters())
scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    with autocast():  # FP16 연산 적용
        output = model(data.cuda())
        loss = loss_fn(output, target.cuda())
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

 

autocast()가 알아서 연산을 FP16/FP32로 구분해 실행
GradScaler()가 loss scaling을 자동 적용

---

BF16 (bfloat16) vs FP16

항목	FP16	BF16
표현 범위	좁음 → underflow 위험	FP32와 같은 exponent → 안정적
속도	매우 빠름	FP16보다 약간 느림
권장	Ampere 이상 GPU	Hopper, TPU, 최신 GPU

---

장점

1. 학습 속도 1.5~3배 ↑
2. 메모리 사용량 절반 ↓
3. 더 큰 모델 / batch 학습 가능
4. FP32 성능 거의 유지

---

단점 & 주의점

1. FP16 underflow → loss scaling 필수
2. 일부 연산은 FP16 지원 안 됨 (softmax, batchnorm 등은 FP32 유지)
3. 너무 큰 learning rate → overflow 위험

---

실제 효과 (벤치마크)

모델FP32 학습 속도AMP 학습 속도메모리 절감

ResNet-50	500 img/s	950 img/s	-45%
BERT-Large	1.0 step/s	2.1 step/s	-47%

AMP 적용으로 속도 2배 가까이 향상 + VRAM 절반 절약

---

실무 적용 팁

• AMP + Gradient Accumulation 같이 사용 → batch size 한계 극복
• 최신 GPU라면 BF16 우선 (안정적, scaling 불필요)
• Torch 2.x에서는 torch.compile()과 AMP 같이 쓰면 추가 속도 향상

---

결론

Mixed Precision Training은 속도·메모리 최적화의 사실상 표준입니다.
AMP를 적용하지 않는 건 GPU 자원을 낭비하는 것과 다름없습니다.
최신 GPU + AMP 조합으로 학습 속도, 모델 사이즈, 비용 모두 개선하세요.

 

 

 

다음 편은 아래에 있습니다!

https://machineindeep.tistory.com/73

대규모 모델 학습·추론 최적화 시리즈 2편: Gradient Accumulation