🤖 machine learning

Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention

이 논문은 NVFP4 양자화를 선형 투영을 넘어 안정적인 옵티마이저 상태, Root/Muon 연산, 혼합 정밀도 어텐션까지 확장하는 모듈형 프레임워크인 Full-Stack FP4를 소개하며, 이를 통해 단일 RTX 5090에서 상당한 메모리 및 속도 이득과 함께 BF16에 근접한 사전 학습 성능을 달성합니다.

원저자: Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

게시일 2026-08-11

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 디지털 뇌가 말하고, 쓰고, 생각하는 법을 가르치려 한다고 상상해 보십시오. 이 "뇌"는 거대 언어 모델(LLM)이며, 학습하기 위해서는 산더미 같은 정보를 처리해야 합니다. 하지만 문제가 하나 있습니다. 더 많이 배울수록 컴퓨터 메모리와 전기를 더 많이 갈구하게 된다는 점입니다. 이는 마치 아주 작은 오래된 노트북에서 고화질 영화를 재생하려는 것과 같습니다. 화면은 흐릿해지고 팬은 비명을 지릅니다. 이를 해결하기 위해 과학자들은 컴퓨터가 계산에 사용하는 숫자를 줄이려고 노력해 왔습니다. 매우 정밀하고 무거운 숫자(정교한 사진 같은) 대신, 더 작고 가벼운 숫자(픽셀화된 스케치 같은)로 바꾸는 것입니다. 이것을 "저정밀도 훈련(low-precision training)"이라고 부릅니다. 최근에는 "FP4"(4비트 부동 소수점)라는 새로운 종류의 아주 작은 숫자가 등장하여, 이 거대한 뇌들을 훨씬 더 작고 빠르게 만들 수 있다는 약속을 하고 있습니다. 하지만 문제는 이렇습니다. 뇌의 주요 부분들은 이 작은 숫자들을 감당할 수 있지만, 학습을 위해 기억해 두는 "기억 장치"나 문장을 이해하기 위해 수행하는 복잡한 수학 연산 같은 "보조" 부분들은 이토록 작은 숫자를 강요받으면 계속 충돌하거나 혼란에 빠진다는 것입니다. 이는 마치 레이싱 카 엔진을 가지고 있지만 자전거 체인으로 구동하려는 것과 같습니다. 엔진은 준비되었지만, 나머지 기계가 따라오지 못하는 상황입니다.

이 논문은 이 디지털 뇌들을 위한 숙련된 정비사 역할을 하는 Full-Stack FP4라는 영리하고 새로운 도구 상구를 소개합니다. 연구진은 뇌의 모든 부분에 똑같은 작은 숫자 규칙을 적용해서는 안 된다는 것을 깨달았습니다. 주요 연결 부위 같은 곳은 작은 숫자를 잘 처리할 수 있지만, 어제 배운 것을 기억하는 데 사용하는 "기억"과 같은 부분들은 매우 민민감하며 특별한 종류의 관리가 필요합니다. 팀은 뇌의 각 부분에 서로 다른 "레시피"를 사용하는 모듈형 시스템을 만들었습니다. 주요 연결 부위의 경우, 나머지 부분은 픽셀화된 블록을 사용하면서도 가장 중요한 세부 사항은 고해해상도 스케치로 유지하는 LoRA-SVD라는 기술을 사용했습니다. 기억의 경우, 숫자들이 작은 상자에 담길 때 길을 잃지 않도록 숫자를 압축하기 전에 "매끄럽게 다듬는" 방법을 발명했습니다. 뇌가 적절한 단어에 집중하도록 돕는 수학 연산(어텐션)의 경우, 가장 민감한 부분은 고해상도로 유지하고 나머지는 픽셀화된 상태로 두기로 결정했습니다.

연구진이 30억 개의 파라미터를 가진 모델(중간 크기의 디지털 뇌)을 사용하여 640억 개의 단어로 학습 테스트를 진행했을 때, 결과는 인상적이었습니다. "Full-Stack FP4" 뇌는 전통적인 고정밀 방식만큼 거의 똑같이 잘 학습했습니다. 두 방식의 학습 점수 차이는 **0.838%**에 불과했으며, 이는 거의 눈에 띄지 않는 수준입니다. 실제로, 본 적 없는 질문에 답하는 능력을 테스트했을 때, 새로운 작은 숫자 버전은 훨씬 적은 메모리를 사용했음에도 불구하고 오히려 정답을 맞히는 능력이 약간 더 좋았습니다. 하나의 강력한 그래픽 카드(RTX 5090)에서 이 새로운 방식은 특정 작업에 대해 뇌의 학습 과정을 2.5배에서 2.8배 더 빠르게 만들었으며, 표준 고정밀 방식보다 메모리를 38%에서 42% 적게 사용했습니다.

연구진은 이 방식이 30억 파라미터 규모의 모델까지는 잘 작동하지만, 더 큰 뇌에서도 작동할지 혹은 훨씬 더 긴 학습 시간 동안 작동할지는 아직 증명되지 않았다는 점을 주의 깊게 언급했습니다. 또한, 모든 것에 작은 숫자를 무작정 던져서는 안 된다는 것도 발견했습니다. 만약 이러한 특별한 레시피 없이 가장 민감한 부분에 작은 숫자를 강제로 적용하려고 하면 학습이 무너집니다. 하지만 현재로서는, 이 "Full-Stack" 접근 방식은 모든 시스템에 하나의 크기를 강요하는 대신, 뇌의 각 부분에 실제로 필요한 수준의 정밀도를 다르게 적용함으로써 우리가 더 똑똑하고, 빠르고, 저렴한 AI를 구축할 수 있음을 보여줍니다.

기술 요약: Full-Stack FP4: 양자화된 투영(Projections), 옵티마이저 및 어텐션을 이용한 안정적인 LLM 사전 학습

1. 문제 정의

최근 NVIDIA Blackwell GPU의 발전으로 네이티브 NVFP4(4비트) 텐서 코어(Tensor Core) 지원이 실질적인 목표가 되었습니다. 기존 연구들은 트랜스포머 선형 투영(Linear Projections)에 대한 안정적인 4비트 사전 학습을 성공적으로 입증했지만, 다른 핵심 학습 모듈에서는 여전히 상당한 격차가 존재합니다. 현재 대부분의 NVFP4 레시피는 옵티마이저 상태(Optimizer States), 옵티마이저 연산 및 어텐션 메커니즘을 위해 더 높은 정밀도(예: BF16 또는 FP8)를 유지합니다.

주요 과제는 이러한 모듈들이 단일한 양자화 규칙으로는 해결할 수 없는 서로 다른 수치적 구조를 가지고 있다는 점입니다:

선형 투영(Linear Projections): 학습이 진행됨에 따라 그래디언트 신호는 약해지는 반면 양자화 노이즈는 일정하게 유지되어, 지속적인 저정밀도 학습의 효용성을 저하시킵니다.
AdamW 상태(States): 2차 모멘텀(Second-moment momentum)은 비음수(non-negative)이며, 헤비 테일(heavy-tailed) 특성을 갖고, 지속적이며, 분모 값에 민감하여 직접적인 4비트 저장은 왜곡을 초래하기 쉽습니다.
Root/Muon 옵티마이저: 이들은 반복적인 뉴턴-슐츠(Newton–Schulz) 행렬 곱에 의존하며, 저정밀도 반복 오차가 누적되어 수렴을 불안정하게 만들 수 있습니다.
어텐션(Attention): 이 연산자는 소프트맥스에 민감한 경로(예: $PV$, $dOV^\top$ )와 그렇지 않은 경로(예: $QK^\top$ )를 모두 포함하고 있어 정밀도 조절이 필요합니다.

결과적으로, 핵심 구성 요소에 대한 혼합 정밀도 폴백(fallback)에 의존하지 않고 투영, 옵티마이저, 어텐션 전체에 걸쳐 안정적이고 네이티브한 NVFP4 사전 학습을 가능하게 하는 통합된 모듈형 프레임워크가 부족한 실정입니다.

2. 방법론: Full-Stack FP4

저자들은 각 학습 모듈의 구체적인 수치적 어려움을 해결하기 위해 설계된 네 가지 독립적이고 결합 가능한 레시피로 구성된 모듈형 프레임워크인 Full-Stack FP4를 제시합니다.

2.1 LoRA-SVD 선형 투영

학습 후기 단계에서 그래디언트 신호의 퇴화를 완화하기 위해, 저자들은 LoRA-SVD를 제안합니다.

메커니즘: 이 방법은 가중치를 $W = W_{res} + \beta L_2 L_1$ 로 파라미터화합니다. 컴팩트한 주성분 부 subspace( $L_2 L_1$ )는 BF16으로 유지하면서, 지배적인 잔차( $W_{res}$ )는 NVFP4로 계산합니다.
역전파(Backward Pass): 부공간 제약 역전파(Subspace-constrained backward pass)는 Cholesky-QR 제약을 사용하여 BF16 저계수(low-rank) 브랜치의 업데이트가 주성분 부공간과 일치하도록 보장하는 한편, 잔차는 NVFP4 구성 요소를 업데이트합니다.
재정렬(Realignment): 부공간은 랜덤 SVD를 사용하여 주기적으로(2,048 스텝마다) 재정렬함으로써 BF16 부공간이 최적의 상태를 유지하도록 합니다.

2.2 양자화된 상태의 AdamW (Q-AdamW)

지속적이고 헤비 테일 특성을 가진 모멘텀 상태를 NVFP4로 저장할 때 발생하는 불안정성을 해결하기 위해:

1차 모멘텀 ( $m_t$ ): 확률적 반올림(Stochastic Rounding, SR) 및 적응형 4/6 양자화를 적용하기 전에 분포를 매끄럽게 하기 위한 하다마르 혼합(Hadamard mixing)을 적용합니다.
2차 모멘텀 ( $v_t$ ): NVFP4 저장을 위한 3단계 전처리 파이프라인을 활용합니다:
1. 제곱근(Square Root): 헤비 테일을 단축합니다.
2. 타일 평균(Tile Mean): 양수인 공통 성분(BF16에 저장)과 대략 0에 가까운 평균을 가진 잔차를 분리합니다.
3. 하다마르 변환(Hadamard Transform): 잔차의 국소적 에너지를 분산시켜 더욱 대칭적으로 만듭니다.
재구성(Reconstruction): 변환과 평균 분리를 역전하여 상태를 재구성한 후, 그 결과를 제곱합니다. 이 파이프라인은 SVD 기반의 부공간 연산을 피할 수 있게 해줍니다.

2.3 NVFP4 Root 옵티마이저 (Q-Root)

Root 옵티마이저에서 사용되는 뉴턴-슐츠(Newton–Schulz) 반복을 위해:

직접 NVFP4 실행: 혼합 정밀도 경로를 사용하는 대신, NS 행렬 곱을 직접 NVFP4로 실행합니다.
안정화: 형태 의존적 계수(특정 행렬 종횡비에 최적화된 $a, b, c$ )와 **p99 아웃라이어 클리핑(outlier clipping)**을 통해 안정성을 확보합니다 보장합니다. 이 접근 방식은 보조적인 혼합 정밀도 경로 없이도 연속적인 NS 단계에서의 오차 증폭을 제한합니다.

2.4 혼합 정밀도 어텐션 (Q-Attn)

정밀도와 효율성의 균형을 맞추기 위해:

정밀도 분리: 프레임워크는 견고한 경로( $QK^\top$ 및 $dS$)에는 NVFP4를 할당하고, 소프트맥스에 민감한 경로($PV$, $P^\top dO$ , $dOV^\top$ )에는 BF16을 유지합니다.
일관성: $Q, K, V$ 의 양자화된 뷰(view)를 순전파와 역전파 간에 재사용하여 텐서 일관성을 보장하고 불필요한 메모리 트래픽을 방지합니다.

3. 주요 기여

LoRA-SVD 투영: 표준 NVFP4 대비 선형 전용 손실 격차를 1.40%에서 0.61%로 줄이는 병합 가능한 재파라미터화 기법을 제공하여, BF16 내의 컴팩트한 부공간을 보호합니다.
양ized-State AdamW: 변환된 상태 파이프라인을 사용하여 2차 모멘트 재구성 오차를 약 60% 줄임으로써, 통합된 사전 학습 실행 중 안정적인 NVFP4 AdamW 모멘텀 상태 저장을 최초로 입증했습니다.
NVFP4 Root: 형태 인식 계수와 클리핑을 통해 구현된 직접적인 저정밀도 뉴턴-슐츠 경로를 통해, 반복 오차를 약 10%에서 3%로 줄였습니다.
학습 가능한 혼합 정밀도 어텐션: 소프트맥스에 민감한 경로는 BF16으로 보호하면서 적합한 순전파 및 역전파 경로는 양자화하는 전략을 통해, 일관된 NVFP4 사전 학습 연산자를 가능하게 했습니다.

4. 실험 결과

저자들은 64B 토큰 동안 학습된 3B 파라미터 Transformer 모델로 프레임워크를 평가했습니다.

학습 손실(Training Loss): BF16 베이스라인은 2.267의 손실에 도달한 반면, Full-Stack FP4는 2.286에 도달하여 상대적 격차는 **0.838%**였습니다.
다운스트림 성능: 제로샷 평가 결과, Full-Stack FP4의 평균 퍼플렉시티(Perplexity)는 26.665(BF16의 26.675 대비)였으며, 평균 정확도는 BF16 베이스라인보다 0.10 퍼센트 포인트 낮았습니다.
어블레이션 연구(Ablation Studies):
- LoRA-SVD 단독 사용 시 선형 전용 손실 격차를 1.40%에서 0.61%로 줄였습니다.
- 모듈형 레시피는 독립적으로 활성화하거나 결합할 수 있으며, "Full-Stack"은 전체적인 범위를 나타냅니다.
네이티브 하드웨어 효율성 (RTX 5090):
- Root 가속: Full-Stack FP4는 Root 옵티마이저 단계에서 최적화된 BF16 대비 2.50배~2.83배, NVIDIA Transformer Engine(TE) 대비 1.83배~2.17배의 속도 향상을 달성했습니다.
- 메모리 절감: 프레임워크는 BF16 대비 AdamW 피크 메모리를 37.9%~42.5% 감소시켰습니다.
- 시간 절감: 줄어든 단계 시간과 BF16 손실을 맞추기 위해 필요한 토큰의 미세한 증가를 고려했을 때, 일치된 손실 시간(matched-loss time)은 56.9%~61.9% 절감되는 것으로 계산되었습니다.

5. 의의 및 한계점

의의:
본 논문은 투영, 옵티마이저 상태, 옵티마이저 연산, 어텐션을 모두 아우르는 NVFP4 사전 학습을 위한 "Full-Stack" 접근 방식이 가능하다는 것을 보여줍니다. 이는 네이티브 4비트 사전 학습이 BF16과 유사한 성능(손실 격차 1% 미만)을 달면서도, Blackwell 클래스 하드웨어에서 상당한 메모리 및 처리량 이점을 제공할 수 있음을 입증합니다. 또한, 실무자가 저정밀도를 점진적으로 도입할 수 있는 모듈형 레시피 세트를 제공합니다.

한계점:

규모: 실험은 3B 파라미터 및 64B 토큰 규모의 모델로 제한되었습니다. 논문은 7B 이상의 규모나 더 긴 학습 기간에서의 동작을 확립하지 않았습니다.
효율성 측정: 네이티브 하드웨어 효율성은 전체 분산 모델이 아닌 단일 RTX 5090 상의 4개 디코더 블록에서 측정되었습니다.
추정치: 시간 절감 계산은 고정된 배치 크스를 기반으로 한 교차 설정 추정치이며, 확보된 메모리를 사용하여 배치 크스를 늘렸을 때 얻을 수 있는 잠재적인 처리량 이득은 고려하지 않았습니다.
커널 퓨전(Kernel Fusion): 일부 커널(LoRA backward, attention)은 아직 완전히 퓨전되지 않았으며, 구현체는 아직 모든 FlashAttention-4 최적화를 활용하지 못하고 있습니다.
포맷 특이성: 레시피와 측정값은 NVFP4에 특화되어 있으며, MXFP4로의 전이 가능성을 주장하지 않습니다.