FP4 4비트 LLM 학습 논문 해설: BF16급 정확도로 학습 비용 낮추기

핵심 메시지

• 논문 'Optimizing LLM Training Using FP4 Quantization'은 학습 중 가중치와 활성값을 단 4비트로 양자화하고도 성능 손실을 거의 없앤다.
• LLM 학습 비용의 90% 이상이 행렬 곱에서 나오므로, 4비트 행렬 곱은 처리량·캐시·메모리 대역폭에서 큰 이득을 줄 잠재력이 있다.
• FP4(E2M1)는 표현 가능한 값이 16개뿐이라 소박하게 학습하면 실패한다 — 그래서 혼합 정밀도 설계가 필요하다.
• 핵심 삼박자: 가중치는 FP16 마스터 복사본으로 갱신 후 FP4로 재양자화, 활성값은 이상치 클램핑+희소 보정, 역전파는 미분가능 그래디언트 추정으로 처리한다.
• 라마 13억·70억·130억을 1000억 토큰으로 학습한 결과 곡선이 BF16과 거의 동일했고, 9개 벤치마크 평균 정확도도 FP4 54.95 대 BF16 54.44로 대등했다.

쉽게 이해하기

오늘날 대부분의 LLM은 가중치마다 16~32비트를 쓰는 부동소수점으로 학습된다. 정밀도가 높아 수치적 뉘앙스를 담을 수 있지만 연산·에너지·시간·메모리, 결국 돈이 많이 든다. 추론 단계에선 8비트·4비트 양자화가 이미 흔하고 2비트·이진 양자화까지 있지만, 이 영상이 다루는 논문은 학습 그 자체를 4비트(FP4)로 하겠다는 도전이다.

이것이 중요한 이유는 LLM 학습 비용의 90~95%가 행렬 곱 하나에서 나오기 때문이다. 4비트 숫자로 행렬 곱을 하면 GPU 코어당 곱셈 수가 늘고 캐시 활용이 좋아지며 메모리 대역폭이 줄어, 하드웨어만 받쳐주면 큰 속도 향상이 가능하다. 다만 함정이 있다. FP4는 양자화 오차가 커서 소박하게 학습하면 손실 곡선이 발산하며 실패한다.

정밀도 사다리를 보면 FP32는 약 40억 개 값을 표현하지만 무겁고, BF16은 메모리를 절반으로 줄여 A100·H100에서 빠르며, FP8은 값이 256개뿐이라 42조차 정확히 표현하지 못하고 40으로 근사한다. 통상 학습은 순전파를 FP8로, 가중치 갱신은 FP16 마스터 복사본으로 하고 다음 반복에서 다시 FP8로 양자화하는 식으로 돈다. 발표자는 활성값이 처음부터 형식의 동적 범위 중심에 놓이도록 가중치를 매개화해 사후 재스케일링을 없애는 uμP 방식을 특히 선호한다고 소개한다.

FP4는 차원이 다른 문제다. E2M1 FP4 형식은 표현 가능한 값이 16개에 불과해 서로 다른 원래 값들이 같은 값으로 뭉개지고, 많은 양자화 함수가 미분 불가능해 역전파조차 제대로 되지 않는다. 그래서 오랫동안 FP4 학습은 불가능하거나 불안정하다고 여겨졌다. 논문은 이를 혼합 정밀도로 푼다 — 행렬 곱은 FP4로, 가중치 갱신·옵티마이저 상태·그래디언트 누적처럼 민감한 부분은 FP8/FP16으로 둔다.

세 가지 핵심 기법이 맞물린다. 첫째, 가중치는 매 스텝 absmax로 FP4에 재양자화하되, 작은 갱신값이 FP4에서 0으로 반올림돼 학습이 깨지지 않도록 FP16 마스터 복사본에 갱신을 적용한다('속도는 FP4, 안정성은 FP16'). 둘째, 활성값은 이상치가 동적 범위를 늘려 나머지가 0으로 반올림되는 문제를 상위 0.1% 클램핑으로 좁히고, 잘라낸 잔차를 희소 행렬로 따로 고정밀 계산해 정보를 보존한다. 셋째, 역전파에서는 순전파의 하드 양자화 대신 매끄러운 미분가능 근사를 써서 값을 FP4 레벨로 '절벽에서 떨어지듯'이 아니라 '경사로를 미끄러지듯' 유도해 안정적인 그래디언트를 얻는다.

주요 인사이트

• 직관적 요약: 계산은 FP4로 빠르게, 가중치 갱신은 FP16으로 정밀하게 — 이 이중 구조가 초저정밀 학습을 실제로 작동하게 하는 핵심 아이디어다.
• 활성값은 미리 결괏값을 알 수 없어 가중치보다 양자화가 어렵고, 소수의 큰 이상치가 전체를 뭉개므로 클램핑과 희소 잔차 보정이 필수다.
• 널리 쓰이는 스트레이트 스루 추정기(STE)는 FP4처럼 낮은 비트폭에서 조악해 수렴이 나빠지므로, 논문은 미분가능 그래디언트 추정으로 대체했다.
• 라마 13억·70억·130억을 1000억 토큰으로 학습한 곡선이 BF16과 거의 같았고, 13B에서 9개 벤치마크 평균 정확도가 FP4 54.95 대 BF16 54.44로 오히려 근소하게 앞서기도 했다.
• 한계는 실제 FP4 하드웨어가 없어 FP8로 에뮬레이션했다는 점(그래서 아직 실속도 이득은 없고 오히려 느림)이지만, 네이티브 FP4를 지원하는 Nvidia 블랙웰이 나오면 FP8 대비 처리량을 두 배로 끌어올릴 수 있어 소규모 랩·스타트업·대학의 대규모 학습을 앞당길 수 있다.

자주 묻는 질문

관련 AI 소식