기술 요약: TetraJet-v2

1. 문제 제기

대규모 언어 모델 (LLM) 사전 학습은 점점 더 과도한 비용이 소요되어 저정밀도 완전 양자화 학습 (FQT) 의 필요성이 대두되고 있습니다. 4 비트 포맷인 NVFP4 와 같은 포맷은 FP8 및 BF16 대비 상당한 효율성 향상을 제공하지만, 이 정밀도에서 안정적이고 손실에 가까운 학습을 달성하는 것은 여전히 중요한 과제입니다. 이전의 완전 4 비트 학습 시도들은 심각한 성능 저하를 겪었습니다.

저자들은 단순한 양자화 오차와 구별되는 완전 FP4 학습을 방해하는 두 가지 주요 최적화 병목 현상을 식별했습니다:

1. 가중치 진동 (Weight Oscillation): 고정밀 마스터 가중치에 대한 작은 업데이트가 양자화된 표현이 인접한 버킷 (예: 0 과 0.5 사이) 사이에서 반복적으로 뒤집히게 하는 지배적인 실패 모드입니다. 이 현상은 특히 학습률이 감소함에 따라 최적화를 방해하여 모델이 안정적인 지역 최소값으로 수렴하는 것을 방지합니다.
2. 이상치 특징 (Outlier Features): 소수의 활성화 채널이 거대한 크기를 보여 FP4 가 정확하게 표현할 수 없는 동적 범위를 생성합니다. 이는 동일한 양자화 그룹 내의 더 작은 값들의 정밀도를 파괴하는 높은 스케일링 인자로 이어집니다.

NVIDIA 의 NVFP4 레시피와 같은 기존 방법들은 종종 혼합 정밀도 (일부 블록을 BF16 으로 유지) 에 의존하거나, 처음부터 시작하는 사전 학습 환경에서 이러한 최적화 역학을 완전히 해결하지 못합니다.

2. 방법론: TetraJet-v2

TetraJet-v2 는 모든 선형 계층에서 활성화, 가중치, 그리고 기울기에 NVFP4 포맷을 적용하도록 설계된 엔드 투 엔드 4 비트 FQT 방법입니다. 이 방법론은 세 가지 핵심 구성 요소를 통해 식별된 병목 현상을 해결합니다:

A. NVFP4 선형 계층을 위한 편향 없는 더블-블록 양자화

NVFP4 의 수치적 제약 (8 비트 스케일링을 사용하는 16 개 요소 그룹) 내에서 정확도를 극대화하기 위해, 저자들은 더블-블록 양자화 방식을 제안합니다:

• 구조: 128 개 요소로 구성된 외부 블록은 전역적으로 스케일링되는 반면, 16 개 요소로 구성된 내부 블록은 로컬로 스케일링됩니다. 이는 텐서 전체에 걸쳐 최대 절대값을 계산하는 것을 피하여 텐서당 스케일링에 비해 하드웨어 친화성과 정확도를 향상시킵니다.
• 편향 없는 기울기 추정: 선형 계층 설계는 다음과 같이 편향 없는 기울기 추정을 보장합니다:
  • 역전파에서 텐서를 정렬하여 ($X$ 대신 $\hat{X}$를 추정).
  • 역전파에서 **확률적 반올림 (Stochastic Rounding)**을 사용하여 $E[\text{round}(x)] = x$를 보장하고, 기울기 기대값의 편향을 방지.
  • 더 높은 정확도를 위해 $1 \times 16$ 가중치 그룹 모양 (비정사각형) 을 채택하며, 이는 전치된 가중치 행렬에 대한 특정 정렬을 요구.

B. OsciReset: 진동 억제

가중치 진동을 해결하기 위해, 저자들은 LLM 사전 학습 중 이 병목 현상을 억제하도록 설계된 최초의 알고리즘인 OsciReset을 도입합니다.

• 감지: 이 방법은 가중치의 최적화 궤적을 추적하여, 양자화된 가중치가 이동한 거리 ($dist_Q$) 와 마스터 가중치가 이동한 거리 ($dist_M$) 의 비율로 정의된 "진동 위험" 지표를 계산합니다. 높은 위험은 마스터 가중치 업데이트가 작음에도 불구하고 가중치가 버킷을 반복적으로 뒤집고 있음을 나타냅니다.
• 메커니즘: 학습률이 낮고 진동 위험이 임계값을 초과할 때, 진동하는 요소의 마스터 가중치는 현재 양자화 버킷의 중앙으로 재설정됩니다.
• 장점: 가중치를 동결하거나 강력한 정규화를 사용하는 이전 방법들 (이는 전역 최적화에 해를 끼침) 과 달리, OsciReset 은 성능을 즉시 저하시키지 않으면서 진동에 덜 취약한 새로운 궤적에서 가중치가 계속 최적화되도록 허용합니다.

C. OutControl: 이상치 제어

이상치를 관리하기 위해 TetraJet-v2 는 하이브리드 전략을 사용합니다:

• 무작위 하다마르 변환 (RHT): 역전파 계산 ($dX$및$dW$) 에 적용되어 이상치 에너지를 채널 간에 더 고르게 분산시켜 양자화 분산을 완화합니다.
• 정적 이상치 유지: 이상치 채널이 구조적으로 일관성 (특정 채널이 항상 높은 노름을 가짐) 이 있음을 인식하여, 짧은 보정 기간 동안 상위 $p\%$(예: 10%) 의 이상치 채널을 식별합니다. 이러한 채널들은 정방향 및 역방향 전파 모두에서 더 높은 정밀도 (예: FP8) 로 유지되며, 네트워크의 나머지는 FP4 에 머뭅니다. 이 정적 선택은 배치별 동적 조정의 계산 오버헤드를 피합니다.

3. 주요 기여

1. 편향 없는 더블-블록 양자화: LLM 을 위한 편향 없는 기울기 추정과 최적의 블록 모양 ($1 \times 16$) 을 갖춘 새로운 NVFP4 선형 계층 설계로, 기존 NVFP4 레시피보다 우수한 성능을 발휘합니다.
2. OsciReset: 사전 학습 중 LLM 의 가중치 진동을 억제하는 최초의 효과적인 알고리즘으로, 저비트 학습의 주요 불안정성 원인을 해결합니다.
3. OutControl: 역전파를 위한 RHT 와 정방향 및 역방향 전파를 위한 정적 이상치 유지를 결합하여 정확도를 보존하는 혼합 정밀도 알고리즘입니다.
4. 종합 평가: OLMo-2 모델 (최대 370M 파라미터, 212B 토큰) 에 대한 광범위한 사전 학습 실험을 통해 전체 스택의 유효성을 입증했습니다.

4. 실험 결과

저자들은 처음부터 학습된 OLMo-2 모델 (70M, 150M, 370M 파라미터) 에서 TetraJet-v2 를 평가했습니다.

• 성능 격차 축소: 제안된 모든 방법을 통합한 TetraJet-v2-full 은 모델 크기에 걸쳐 전체 정밀도 (BF16) 학습 대비 성능 격차를 평균 51.3% 줄였습니다.
• 퍼플렉시티 (PPL): 212B 토큰으로 학습된 370M 모델에서 TetraJet-v2-full 은 검증 PPL 을 25.11로 달성했으며, 이는 BF16 의 23.70 과 NVIDIA NVFP4 기준의 26.20 과 비교됩니다.
• 하류 작업: 이 방법은 ARC, BoolQ, MMLU 등 일련의 벤치마크에서 모든 FP4 방법 중 가장 높은 평균 성능을 달성했으며, 370M 모델은 평균 43.60% 의 정확도에 도달했습니다.
• 효율성: CUDA 커널 구현은 RTX 5090 하드웨어에서 FP8 (TransformerEngine) 대비 엔드 투 엔드 속도 향상 1.67 배를 달성하여, 진동 억제 및 이상치 제어의 추가 복잡성에도 불구하고 실질적인 성능 향상을 입증했습니다.
• 애블레이션 연구:
  • OsciReset 을 제거하면 심각한 성능 저하와 손실의 반등이 발생했습니다.
  • OutControl 은 정방향 및 역방향 전파 모두에서 정확도를 향상시켰습니다.
  • 이 방법은 하이퍼파라미터 선택 (예: 진동 임계값, 시작 시간) 에 강건함이 입증되었습니다.

5. 중요성 및 주장

이 논문은 TetraJet-v2 가 LLM 의 완전 저비트 (4 비트) 사전 학습을 실현 가능하게 하고 고정밀도 학습과 경쟁력 있게 만드는 중요한 단계라고 주장합니다.

• 최적화 vs 표현: 저자들은 저비트 학습의 성능 격차가 단순히 양자화 오차 (표현) 에만 기인하는 것이 아니라, 이전 방법들이 체계적으로 해결하지 못한 최적화 문제 (진동 및 이상치) 에 의해 결정적으로 주도된다고 강조합니다.
• 일반화 가능성: Vision Transformers 를 위해 설계된 이전 진동 완화 기술들 (종종 미세 조정이나 특정 하이퍼파라미터 튜닝에 의존) 과 달리, OsciReset 은 최소한의 튜닝으로 처음부터 시작하는 LLM 사전 학습에 적용 가능한 간단하고 효과적인 방법으로 제시됩니다.
• 실용성: FP8 대비 1.67 배의 속도 향상을 달성하면서 BF16 대비 정확도 격차를 50% 이상 좁힘으로써, 이 작업은 완전 양자화 학습이 계산 효율성과 모델 품질을 모두 제공할 수 있음을 보여주며, 기초 모델 연구의 하드웨어 및 에너지 비용을 줄일 가능성을 시사합니다.

저자들은 실험의 규모 (최대 370M 파라미터) 에 대한 한계와 추가 정확도 향상을 위한 잠재적인 미래 방향으로 탐구한 FP6 $\times$ FP4 와 같은 혼합 정밀도 전략에 대한 현재 주류 하드웨어 지원의 부재를 인정합니다.