FP8 혼합 정밀도 LLM 학습: Transformer Engine과 Delayed Scaling으로 H100·B200 학습 비용을 절반으로 줄이는 방법

요약

대형 언어 모델 사전 학습 비용의 핵심 결정 요인은 연산 정밀도다. BF16으로 H100 하나를 4주 동안 돌리면 FP8로 2주면 끝난다는 단순 계산이 맞다. 그러나 FP8은 "정밀도가 낮은 BF16"이 아니다. 표현 범위·그래디언트 안정성·행렬 곱 누적 방식이 근본적으로 다르기 때문에, 단순히 dtype만 바꾸면 학습이 발산하거나 품질이 하락한다.

NVIDIA Transformer Engine은 이 문제를 해결하기 위해 만들어진 라이브러리다. FP8 연산 단위(GEMM, LayerNorm, attention)를 블록별로 분리하고, 각 텐서에 스케일 팩터(scale factor)를 붙여 수치 범위를 동적으로 관리한다. DeepSeek-V3, Llama 3.3, Phi-4 등 2025~2026년의 주요 오픈 모델 대부분이 FP8 혼합 정밀도 학습을 채택했다.

핵심 요약:

• FP8에는 E4M3(순전파)과 E5M2(역전파) 두 서브포맷이 있다.
• 스케일 팩터를 실시간으로 업데이트하는 Delayed Scaling이 표준 운영 패턴이다.
• 전체 학습 비용을 BF16 대비 40~60% 절감할 수 있다.
• PyTorch 2.13+ 및 NVIDIA Transformer Engine으로 기존 코드 대비 최소 변경으로 도입 가능하다.

배경: 왜 BF16으로는 충분하지 않은가

정밀도 포맷의 역사

LLM 학습 정밀도의 진화는 가용 GPU 하드웨어와 함께 진행됐다.

BF16은 FP32와 같은 지수 비트를 가지므로 오버플로우 없이 대부분의 학습을 안정적으로 수행한다. 그러나 H100의 Tensor Core는 FP8 연산 시 2배의 FLOP/s를 제공한다. H100 SXM5 기준 BF16 989 TFLOP/s → FP8 1979 TFLOP/s. 이것이 FP8 학습의 동기다.

FP8이 어려운 이유

FP8 E4M3의 표현 범위는 [-448, 448]이다. LLM 학습에서 가중치·활성화·그래디언트의 절댓값은 이 범위를 훨씬 벗어난다. 단순히 dtype을 변경하면 오버플로우(NaN) 또는 언더플로우(0으로 소실)가 발생한다.

해결책은 텐서에 per-tensor 또는 per-block 스케일 팩터 s를 곱해서 FP8 범위 안으로 끌어들이는 것이다. x_fp8 = clip(x / s, -448, 448). 이 s 값을 어떻게 관리하느냐가 FP8 학습의 핵심이다.

Transformer Engine 아키텍처

설계 원칙

NVIDIA Transformer Engine은 세 가지 원칙으로 설계됐다.

1. 레이어 단위 FP8 캡슐화: GEMM, LayerNorm, Attention 등 각 연산자를 FP8 입출력으로 감싸고, 스케일 팩터 관리를 내부에 숨긴다.
2. Delayed Scaling: 이전 반복(iteration)의 통계로 다음 반복의 스케일을 결정한다. 실시간 abs-max 계산 오버헤드를 제거한다.
3. 고정밀 마스터 가중치: 옵티마이저 상태(momentum, variance)와 마스터 가중치는 FP32 또는 BF16으로 유지한다.

Delayed Scaling 동작 원리

Transformer Engine의 Delayed Scaling은 다음 흐름으로 작동한다.

1. amax 기록: 매 스텝 종료 시 해당 텐서의 절댓값 최대값 amax를 amax_history 배열에 저장한다.
2. 스케일 계산: 다음 스텝 시작 전, 지난 N 스텝(기본 1024)의 amax_history에서 최댓값을 꺼내 스케일 팩터를 계산한다. scale = fp8_max_val / (margin × amax_max).
3. 적용: 계산된 스케일로 텐서를 FP8로 캐스팅한다.

margin은 미래 변동에 대비한 여유 계수다. NVIDIA 권고값은 1.0. 학습이 불안정하면 낮추고, FP8 활용도(포화 없음)가 낮으면 올린다.

실시간 스케일링(즉시 계산)과의 차이: 즉시 계산은 텐서를 먼저 훑어 amax를 구한 뒤 캐스팅하므로 두 번 읽어야 한다. Delayed Scaling은 이전 반복 통계를 쓰므로 텐서를 한 번만 읽는다. 1000~4000 토큰/스텝 규모에서 처리량 차이가 5~10% 발생한다.

두 FP8 서브포맷의 선택 이유

E4M3 vs E5M2

FP8에는 두 가지 표준 서브포맷이 있다. IEEE 754에 준하는 NaN과 무한대 처리 방식이 다르고, 지수·가수 비트 배분도 다르다.

순전파에 E4M3를 쓰는 이유: 활성화·가중치값은 특별한 처리 없이 일정 범위 내에 있다. 상대 정밀도(가수 비트)가 더 중요하므로 E4M3가 적합하다.

역전파에 E5M2를 쓰는 이유: 그래디언트는 매우 작은 값이 많아 언더플로우가 발생하기 쉽다. 지수 비트가 많은 E5M2로 동적 범위를 넓혀 그래디언트 소실을 방지한다.

GEMM 누적 정밀도 문제

FP8 행렬 곱에서 가장 미묘한 부분은 부분합 누적(accumulation)이다.

H100/B200의 FP8 Tensor Core는 다음 방식으로 동작한다.

C_partial = A_fp8 @ B_fp8   # 16×16 타일 단위 계산
C_fp32 = sum(C_partial)     # FP32로 누적
C_out = cast_to_fp8(C_fp32) # 출력 포맷으로 변환

누적은 FP32다. 이것이 핵심이다. 타일 내 부분합만 FP8에서 이루어지고, 타일 간 합산은 FP32 정밀도로 이루어진다. 이 덕분에 수치 안정성이 BF16 학습과 큰 차이 없이 유지된다.

주의점: 타일 크기(16×16 = 256 원소)를 초과하는 누적은 모두 FP32다. 매우 큰 hidden size에서도 이 보호가 유지된다.

DeepSeek-V3 FP8 학습 사례

DeepSeek-V3(671B 파라미터, 37B 활성화 MoE)는 FP8 혼합 정밀도 학습을 체계적으로 적용한 최초의 프론티어 오픈 모델이다.

주요 설계 결정:

1. 순전파 FP8: 모든 Linear 레이어의 활성화·가중치를 E4M3으로 캐스팅. 어텐션 레이어는 BF16 유지(수치 민감성 때문).
2. 역전파 FP8: 그래디언트를 E5M2으로 캐스팅. 단, 가중치 업데이트(∂L/∂W)는 FP32.
3. 전문가 병렬 통신 FP8: MoE 전문가 간 all-to-all 통신에서도 FP8 압축 적용. 네트워크 대역폭 절감.
4. 마스터 가중치 BF16: 옵티마이저 상태와 마스터 가중치는 BF16으로 유지.

결과: 2048 H800 GPU에서 14.8T 토큰 학습 시 BF16 대비 약 40% 연산 비용 절감.

PyTorch 2.13으로 FP8 학습 도입하기

PyTorch 2.13은 torch.amp에 FP8 지원을 추가했다. Transformer Engine와 달리 PyTorch 네이티브 경로를 쓴다.

import torch
from transformer_engine.pytorch import fp8_autocast, DelayedScaling

# Transformer Engine 방식 (권장)
with fp8_autocast(enabled=True, fp8_recipe=DelayedScaling()):
    output = model(input)
    loss = criterion(output, target)
loss.backward()

# PyTorch 2.13 네이티브 방식 (실험적)
with torch.autocast(device_type="cuda", dtype=torch.float8_e4m3fn):
    output = model(input)

Transformer Engine 방식을 권장하는 이유: DelayedScaling 레시피가 amax 이력 관리, per-tensor 스케일, 그래디언트 스케일을 자동 처리한다. PyTorch 네이티브는 스케일 관리를 수동으로 해야 한다.

성능 수치와 운영 기준

처리량 향상

Open question: B200에서 tcgen05 FP8 경로를 쓸 때 실측 처리량은 모델·시퀀스 길이에 크게 의존한다. 공개 벤치마크가 아직 제한적이다.

품질 영향

PPL(perplexity) 기준으로 BF16 학습 대비 FP8는 일반적으로 0.1~0.5% 이내의 차이를 보인다. 단, 다음 상황에서 품질 저하가 발생할 수 있다.

• window_size가 너무 작아 스케일이 급격히 변하는 경우
• Attention softmax나 LayerNorm을 FP8로 처리하는 경우 (권장하지 않음)
• 그래디언트 클리핑 없이 초기 학습에서 폭발적 그래디언트 발생

도입 체크리스트

• [ ] GPU가 H100(Hopper, SM90) 이상인지 확인. A100은 FP8 Tensor Core 없음.
• [ ] CUDA 12.0+, cuDNN 9.0+, NVIDIA Transformer Engine 1.x 이상 설치
• [ ] fp8_recipe.amax_history_len = 1024 (기본값)로 시작
• [ ] 초기 1000 스텝은 BF16 warm-up → FP8 전환 (불안정한 초기 스케일 회피)
• [ ] 손실 곡선을 BF16 체크포인트와 비교해 2000 스텝 후 동등한지 확인
• [ ] 옵티마이저 상태(momentum, variance)는 FP32 유지
• [ ] amax_compute_algo를 most_recent로 설정해 빠른 변화 추적 (불안정한 학습 시)
• [ ] 그래디언트 클리핑 (max_norm=1.0) 필수 유지

References

• NVIDIA Transformer Engine: https://github.com/NVIDIA/TransformerEngine
• NVIDIA Transformer Engine 문서: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/
• "FP8 Formats for Deep Learning" (IEEE 2022): https://arxiv.org/abs/2209.05433
• DeepSeek-V3 Technical Report (FP8 학습 섹션): https://arxiv.org/abs/2412.19437
• NVIDIA H100 FP8 처리량 사양: https://www.nvidia.com/en-us/data-center/h100/
• "Microscaling Data Formats for Deep Learning" (OCP MX spec): https://arxiv.org/abs/2310.10537
• PyTorch FP8 지원 문서: https://pytorch.org/docs/stable/amp.html
• Llama 3.3 학습 보고서 (Meta, FP8 사용): https://arxiv.org/abs/2407.21783