혼합 정밀도 학습

Mixed precision training; 혼합 정밀도 학습

연산 대부분은 16비트(또는 8비트) 부동소수점으로 하고, 수치에 민감한 부분만 32비트로 남겨 정확도를 유지하면서 학습 속도와 메모리를 아끼는 기법

혼합 정밀도 학습(mixed precision training)은 딥 러닝 모델을 학습할 때 행렬 곱 같은 무거운 연산은 FP16이나 bfloat16 같은 낮은 정밀도로 하고, 가중치 갱신·손실 계산·큰 합산처럼 오차에 민감한 부분은 FP32(단정밀도)로 남기는 방법이다. 2017년 바이두와 엔비디아 연구진이 FP16만으로는 깨지는 학습을 세 가지 장치(FP32 마스터 가중치, 손실 스케일링, FP32 누산)로 살려 낸 논문을 내면서 표준 기법이 되었다.

얻는 것은 두 가지다. 첫째, 활성값·가중치를 16비트로 저장하면 메모리 사용량과 메모리 대역폭이 거의 절반으로 준다. 논문은 이 방법으로 모델 메모리 사용량을 "거의 2배" 줄였다고 보고했다. 둘째, 엔비디아 볼타(V100) 이후 GPU의 텐서 코어는 FP16·BF16 행렬 곱을 FP32보다 훨씬 빠르게 처리하므로 연산 시간이 준다. 대형 언어 모델 사전 학습에는 BF16 혼합 정밀도를 쓰고, 엔비디아 호퍼·블랙웰 세대부터는 FP8·FP4까지 내려가는 연구가 이어지고 있다. 수 형식 자체의 비교는 bfloat16 문서에 있다.

왜 그냥 FP16으로는 안 되는가

IEEE 754 반정밀도(FP16)는 지수 5비트, 가수 10비트라서 표현 범위가 좁다. 가장 큰 수는 65504, 가장 작은 정규수는 $2^{-14}\approx 6.1\times 10^{-5}$, 비정규수까지 써도 $2^{-24}\approx 6\times 10^{-8}$ 아래는 0이 된다. 학습 중 생기는 수치 문제는 세 가지로 나뉜다.

문제	현상	대책
기울기 언더플로	크기가 $2^{-24}$보다 작은 기울기가 0으로 사라짐. 논문의 Multibox SSD 실험에서 활성값 기울기의 상당수가 FP16 표현 범위 아래에 있었다.	손실 스케일링
갱신량 소실	가중치에 비해 갱신량 $\eta g$가 너무 작으면(가중치의 약 $2^{-11}$배 미만) FP16 덧셈에서 반올림으로 사라짐	FP32 마스터 가중치
누산 오차·오버플로	긴 내적, 소프트맥스·정규화의 합산을 FP16으로 하면 오차가 쌓이거나 65504를 넘음	FP32로 누산, 일부 연산은 FP32로 실행

논문은 FP32 마스터 가중치 없이 FP16으로만 갱신하면 중국어 음성 인식 모델의 정확도가 약 80% 상대 손실을 입었다고 보고했고, Multibox SSD 검출기는 손실 스케일링 없이는 발산했지만 스케일 8로 FP32와 같은 정확도(mAP 77.1% 대 76.9%)를 냈다.

세 가지 핵심 기법

FP32 마스터 가중치

가중치의 원본은 FP32로 따로 보관한다. 순전파·역전파 때는 이를 FP16(또는 BF16)으로 복사해 쓰고, 옵티마이저가 계산한 갱신량은 FP32 원본에 더한다. $w_{t+1}^{(32)}=w_t^{(32)}-\eta \cdot \mathrm{fp32}(g_t^{(16)}),w_{t+1}^{(16)}={\mathrm{round}}_{16}(w_{t+1}^{(32)})$ 작은 갱신량이 FP32 쪽에 계속 쌓이므로 언젠가 16비트 사본에도 반영된다. 대가로 가중치 메모리는 오히려 늘지만(FP32 사본 + 16비트 사본), 학습 메모리의 대부분은 활성값이 차지하므로 전체로는 이득이다. Adam 옵티마이저를 쓰면 파라미터당 16비트 가중치 2바이트, 16비트 기울기 2바이트, FP32 사본·1차 모멘트·2차 모멘트 12바이트로 모두 16바이트가 든다. 이 16바이트를 GPU 사이에 나눠 갖는 것이 ZeRO다.

손실 스케일링

역전파는 연쇄 법칙에 따른 곱셈이므로, 손실 $L$에 상수 $S$를 곱하면 모든 기울기가 똑같이 $S$배가 된다. 그래서 $S\cdot L$로 역전파해 작은 기울기를 FP16 표현 범위 안으로 끌어올린 뒤, 가중치 갱신 직전에 FP32에서 $1/S$를 곱해 되돌린다. ${\mathrm{\nabla }}_w(S\cdot L)=S\cdot {\mathrm{\nabla }}_{w}L\Rightarrow g=\frac{1}{S}\mathrm{fp32}({\mathrm{\nabla }}_w(SL))$ 논문은 8부터 32K까지 고정 스케일로 여러 모델을 학습했고, 오버플로만 나지 않으면 스케일을 크게 잡아서 손해 볼 것이 없다고 했다. 실무에서는 동적 손실 스케일링을 쓴다.

• 큰 값(예: $2^{16}$)에서 시작한다.
• 매 단계 기울기에 inf나 NaN이 있는지 검사한다. 있으면 그 단계의 가중치 갱신을 건너뛰고 스케일을 줄인다(예: 절반).
• 정해진 횟수(예: 2000단계) 동안 오버플로가 없으면 스케일을 키운다(예: 2배).

이렇게 하면 표현 범위의 위쪽 끝 가까이에 기울기 분포를 붙여 둘 수 있다. PyTorch의 GradScaler 기본값이 정확히 이것이다(초기값 65536, 증가 2배, 감소 0.5배, 증가 간격 2000단계). BF16은 지수 비트가 FP32와 같아 언더플로 문제가 거의 없으므로 손실 스케일링이 필요 없다.

FP32 누산과 연산별 정밀도

16비트 곱의 합을 16비트로 누적하면 오차가 커진다. 텐서 코어는 FP16 곱을 FP32로 누산한 뒤 결과를 16비트로 내보내며, 논문도 이 FP32 누산이 없으면 일부 모델이 기준 정확도에 못 미쳤다고 했다. 배치 정규화 통계나 소프트맥스 같은 큰 합산(reduction)도 FP32로 해야 한다. 이 판단을 연산마다 자동으로 해 주는 것이 프레임워크의 자동 혼합 정밀도(automatic mixed precision, AMP)다.

PyTorch 자동 혼합 정밀도(torch.amp)

파이토치는 torch.autocast 문맥 안에서 연산마다 정밀도를 골라 준다. PyTorch 2.14 문서 기준으로 CUDA에서 FP16으로 바뀌는 연산은 matmul, mm, bmm, addmm, linear, conv1d~3d, LSTMCell 등 행렬 곱 계열이고, FP32로 올리는 연산은 exp, log, pow, softmax·log_softmax, layer_norm, group_norm, cross_entropy, mse_loss, norm, cumsum 등 범위가 넓거나 합산이 큰 연산이다. 목록에 없는 연산은 입력 자료형을 따른다.

autocast의 기본 자료형은 CUDA에서 torch.float16, CPU에서 torch.bfloat16이다. 예전 API인 torch.cuda.amp.autocast와 torch.cuda.amp.GradScaler는 폐기 예정(deprecated)이 되었고 torch.amp.autocast("cuda", ...), torch.amp.GradScaler("cuda", ...)로 바뀌었다.

FP16 학습 루프

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(512, 2048), nn.GELU(), nn.Linear(2048, 10)).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler("cuda")          # 동적 손실 스케일링
loss_fn = nn.CrossEntropyLoss()

for x, y in loader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad(set_to_none=True)
    # 순전파와 손실 계산만 autocast 안에 둔다
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(x)                    # Linear는 FP16으로 실행
        loss = loss_fn(logits, y)            # cross_entropy는 FP32로 실행
    scaler.scale(loss).backward()            # S * loss 로 역전파
    scaler.unscale_(optimizer)               # 기울기를 1/S 로 되돌림 (클리핑 전에 필요)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)                   # inf/NaN이 있으면 이 단계를 건너뜀
    scaler.update()                          # 스케일 조정

모델 파라미터는 FP32 그대로 두고(model.half()를 부르지 않는다), autocast가 연산 직전에 16비트 사본을 만든다. 이 사본이 FP32 마스터 가중치 기법의 구현이다. 역전파는 autocast 밖에서 부르며, 각 연산의 역전파는 순전파 때 쓴 자료형으로 실행된다. 그래디언트 클리핑을 하려면 반드시 unscale_ 뒤에 해야 원래 크기의 기울기 노름을 잴 수 있다.

BF16 학습 루프

for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        loss = loss_fn(model(x.cuda()), y.cuda())
    loss.backward()          # GradScaler 불필요
    optimizer.step()

BF16은 FP32와 지수 범위가 같아 스케일러가 필요 없다. 대신 가수가 7비트뿐이라 FP16보다 반올림 오차가 크다. PyTorch 문서는 BF16으로 사전 학습된 모델 대부분이 FP16의 최대값 65504 범위 안에서 동작하지 못해 기울기가 오버플로할 수 있다고 경고한다. 즉 BF16으로 학습된 모델을 FP16으로 미세 조정하는 것은 위험하다.

어떤 정밀도를 고를까

방식	손실 스케일링	메모리	하드웨어	주의할 점
FP32	불필요	기준	모든 GPU	가장 느림
TF32	불필요	FP32와 같음	엔비디아 암페어 이후	저장은 FP32, 텐서 코어 내부에서 입력 가수를 10비트로 줄여 곱함. 코드 변경 없이 켤 수 있음
FP16 혼합 정밀도	필요	활성값 약 절반	볼타 이후	오버플로·언더플로 관리, 스케일러 필수
BF16 혼합 정밀도	불필요	활성값 약 절반	암페어·TPU·최신 CPU	가수 정밀도가 낮음, 구형 GPU 미지원
FP8 혼합 정밀도	텐서·블록별 스케일 필요	행렬 곱 입력 약 1/4	호퍼 이후	전용 라이브러리 필요, 스케일 관리가 복잡

TF32는 엔비디아 A100(암페어)에서 도입된 텐서 코어용 형식으로, FP32와 같은 8비트 지수와 FP16과 같은 10비트 가수를 쓴다. PyTorch는 1.7~1.11에서는 행렬 곱 TF32를 기본으로 켰다가 1.12부터 기본으로 끈다. torch.set_float32_matmul_precision("high")로 켤 수 있다.

실무 선택 기준은 대체로 다음과 같다. 암페어 이후 GPU나 TPU라면 BF16이 첫 선택이다. 스케일러가 필요 없고 발산 위험이 적다. V100·T4처럼 BF16 텐서 코어가 없는 GPU에서는 FP16과 GradScaler를 쓴다. FP16으로 학습한 기존 모델을 이어 학습하거나 추론할 때는 FP16을 그대로 쓰는 편이 안전하다.

FP8 혼합 정밀도

2022년 엔비디아·Arm·인텔 연구진은 E4M3(지수 4비트, 가수 3비트)과 E5M2(지수 5비트, 가수 2비트) 두 FP8 형식을 제안하고, 순전파의 가중치·활성값에는 E4M3, 역전파의 기울기에는 E5M2를 쓰는 조합으로 최대 1750억 매개변수 GPT 계열 모델까지 16비트 학습과 같은 수준의 결과를 냈다. 같은 해 발표된 엔비디아 H100(호퍼)의 4세대 텐서 코어가 FP8을 지원하며, 엔비디아는 층마다 FP8과 16비트 사이를 고르고 스케일을 관리하는 소프트웨어 겸 하드웨어 기능을 "트랜스포머 엔진"이라 불렀다.

FP8은 표현 범위가 너무 좁아 FP16처럼 손실 하나에 스케일 하나를 곱하는 것으로는 부족하다. 텐서마다(per-tensor), 나아가 블록마다 스케일을 따로 둔다. FP8 논문은 FP16 AMP처럼 오버플로가 날 때만 스케일을 줄이는 방식이 FP8에는 맞지 않는다고 했다. 대표 방식은 다음과 같다.

• 지연 스케일링(delayed scaling): 이전 몇 단계의 텐서 최댓값(amax) 기록으로 이번 단계의 스케일을 정한다. 엔비디아 Transformer Engine의 초기 방식이다.
• 온라인·미세 단위 스케일링: 딥시크-V3는 활성값을 1×128 타일, 가중치를 128×128 블록 단위로 그때그때 스케일을 계산해 모든 텐서에 E4M3를 썼고, BF16 기준 대비 상대 손실 오차가 0.25% 미만이었다고 보고했다. 이 논문은 FP8 혼합 정밀도 학습을 초대형 모델에서 처음 검증했다고 밝혔다.
• 마이크로스케일링(MX) 형식: 32개 원소마다 2의 거듭제곱 스케일 하나를 두는 블록 형식(MXFP8·MXFP6·MXFP4)이며, 엔비디아 블랙웰은 16개 원소마다 FP8 스케일을 두는 NVFP4도 지원한다. 자세한 형식은 bfloat16 문서 참고.

FP8 학습에서도 마스터 가중치·옵티마이저 상태·정규화·소프트맥스는 BF16이나 FP32로 남는다. 행렬 곱 입력만 8비트로 내리는 것이 기본이다.

함께 쓰는 기법과 주의점

• 그래디언트 체크포인팅: 활성값 메모리를 더 줄이려면 같이 쓴다. 재계산도 autocast 상태를 그대로 따른다.
• ZeRO·FSDP: 16바이트/파라미터의 모델 상태를 GPU 사이에 나눈다. PyTorch FSDP2는 MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)처럼 계산용 자료형과 기울기 통신용 자료형을 따로 정한다.
• 기울기 누적: 여러 마이크로배치의 기울기를 모을 때 스케일 조정과 inf 검사는 실제 갱신 단위(유효 배치)마다 한 번만 해야 한다.
• 손실 급등(loss spike): BF16·FP8에서도 어텐션 로짓이나 출력층 로짓이 커지면 불안정해진다. 이때 문제를 정밀도 탓으로만 보지 말고 학습률, 초기화, 정규화 위치를 함께 본다.
• 추론 정밀도와는 다른 문제다. 학습 뒤 가중치를 INT8·INT4로 줄이는 것은 신경망 양자화에서 다룬다.

역사

연도	사건
2017	엔비디아 V100(볼타)에 FP16 입력·FP32 누산 텐서 코어 탑재. 같은 해 10월 마이크러시비어스 등 〈Mixed Precision Training〉 공개(ICLR 2018)
2019	구글이 클라우드 TPU v2·v3 행렬 곱 장치의 bfloat16 곱셈·FP32 누산 방식을 공개 설명
2020	A100(암페어)에서 BF16·TF32 텐서 코어 지원
2021	PyTorch 1.10에서 장치 공통 API torch.autocast 도입
2022	FP8(E4M3·E5M2) 형식 제안, H100(호퍼)에서 FP8 텐서 코어 지원
2023	마이크로소프트·AMD·인텔·메타·엔비디아·퀄컴 연구진이 MX 형식 연구 발표
2024	딥시크-V3가 FP8 혼합 정밀도로 대형 MoE 모델 사전 학습
2025	엔비디아 블랙웰용 NVFP4 형식 소개

같이 보기

• bfloat16
• 부동소수점 표현 (IEEE 754 규격)
• 텐서 코어
• ZeRO (딥 러닝)
• 그래디언트 체크포인팅
• 그래디언트 클리핑
• 신경망 분산 학습
• 신경망 양자화
• 블록 부동소수점 양자화

참고 문헌

• PyTorch 2.14 문서: torch.amp
• 파이토치 한국어 튜토리얼: 자동 혼합 정밀도(Automatic Mixed Precision) 가이드
• DeepSpeed 설정 문서(fp16·bf16 항목)