기술 요약: ScaleSweep

문제 정의

NVIDIA Blackwell GPU의 등장은 FP4 E2M1 데이터 타입과 두 단계 스케일링(FP8 마이크로 블록 스케일 및 글로벌 FP32 스케일)을 결합한 저정밀도 양자화 방식인 NVFP4 포맷에 대한 네이티브 지원을 도입했습니다. 이 포맷은 정수 전용 포맷에 비해 메모리 및 대역폭 효율성이 뛰어나지만, 기존의 사후 훈련 양자화(PTQ) 방법들은 그 잠재력을 완전히 활용하는 데 어려움을 겪고 있습니다.

AbsMax 및 4/6 휴리스틱과 같은 현재의 NVFP4 초기화 전략은 최대 표현 가능 값에 기반한 단순한 규칙에 의존합니다. 본 논문에서 입증되었듯이, 이러한 방법들은 최적의 솔루션(구체적으로는 FP8 양자화된 최적의 FP32 블록 스케일)과 비교했을 때 눈에 띄는 성능 격차를 남깁니다. 표준 INT4 양자화와 달리, NVFP4의 두 단계 스케일링 구조와 부동 소수점 코드북은 기존의 INT4 최적화 기술(예: Quaot와 같은 회전 기반 접근 방식)을 효과적이지 않거나 오히려 해롭게 만듭니다. 결과적으로, 과도한 계산 비용을 발생시키지 않으면서 재구성 오차(reconstruction error)를 최소화하기 위해 NVFP4 탐색 공간을 효율적으로 탐색할 수 있는 전용 스케일 최적화 방법이 절실히 필요합니다.

방법론: ScaleSweep

저자들은 NVFP4 양자화에서 블록 스케일을 최적화하기 위한 단순하면서도 효율적인 방법인 ScaleSweep을 제안합니다. 핵심 혁신은 스케일 선택 문제를 휴리스틱한 추측에서 제약 조건이 있는 최적화 작업으로 변환하는 데 있습니다.

1. 최적화 목표

이 방법은 두 가지 주요 손실 함수를 대상으로 합니다:

• 평균 제곱 오차 (MSE): 원본 텐서와 양자화된 텐서 사이의 재구성 오차를 최소화합니다.
• 가중 평균 제곱 오차 (WMSE): 최종 출력에 큰 영향을 미치는 채널의 오차를 우선시하기 위해 중요도 점수(입력 활성화 노름 또는 Gram 행렬 대각 성분에서 유도됨)를 포함합니다.

2. 효율적인 스위핑을 위한 이론적 경계

단순한 접근 방식은 각 블록에 대해 126개의 양의 유한 표현 가능한 FP8 값 전체를 전수 조사하는 것이며, 이는 계산적으로 비효율적입니다. ScaleSweep은 이론적 하한 및 상한을 도출하여 스케일링 범위를 획기적으로 줄임으로써 이 문제를 해결합니다:

• 글로벌 스케일 ($S$): 블록 최적화를 위한 동적 범위를 보존하도록 조정된 4/6 전략과 유사한 휴리스틱을 사용하여 고정됩니다 ($S = \max|x_i| / (256 \cdot 6)$).
• 블록 스케일 상한: 저자들은 이론적 분석(Lemma 4.1)을 통해, 만약 스케일 $s$가 $\frac{12}{7} s_{base}$($s_{base}$는 FP8로 반올림된 블록의 최대 크기)를 초과하면, 양자화 손실이 $s/2$보다 낮아질 수 없음을 증명했습니다. 이는 상한을 $\frac{12}{7} s_{base}$로 설정합니다. FP8 비트 패턴 공간에서 이는 기본 스케일의 정수 표현으로부터 +7의 오프셋에 해당합니다.
• 블록 스케일 하한:
  • MSE의 경우, Lemma 4.3은 최적의 스케일이 $\frac{4}{5} s_{base}$에 의해 하한이 결정됨을 확립합니다. 비트 패턴 공간에서 이는 -3의 오프셋입니다.
  • WMSE의 경우, 큰 활성화 값과 관련된 가중치가 작을 수 있으므로, 하한은 경험적으로 $\frac{1}{2} s_{base}$로 설정되며, 이는 -8의 오프셋에 해당합니다.

3. 구현

ScaleSweep은 FP8 비트 패턴 공간에서의 좁은 로컬 이웃(예: MSE의 경우 $[-3, 7]$, WMSE의 경우 $[-8, 7]$ 오프셋)으로 스위핑 범위를 제한함으로써, 대상 목적 함수를 최소화하는 후보 스케일을 식별합니다. 이 접근 방식은 베이스라인 양자화 연산자에 비해 무시할 수 있는 수준의 오버헤드를 도입하면서도 최적의 후보가 탐색 범위 내에 유지되도록 보장합니다.

주요 기여

1. 이론적 분석: 본 논문은 NVFP4 양자화에 대한 엄격한 이론적 분석을 제공하며, MSE 및 WMSE 목적 함수 하에서 최적의 블록 스크레일에 대한 수학적으로 정당화된 하한 및 상한을 도출합니다.
2. ScaleSweep 알고리즘: FP8 블록 스케일 최적화를 비트 패턴 공간의 조밀한 구간으로 제한하여, RTN(Round-to-Nearest) 및 GPTQ 파이프라인 모두에서 효율적인 스케일 선택을 가능하게 하는 새로운 캘리브레이션 방법입니다.
3. 경험적 검증: 점점 더 공격적인 양자화 설정(Weights, Weights+Activations, Weights+Activations+KV Cache, Weights+Activations+KV Cache+Query States)에 대해 Llama 및 Qwen 모델을 대상으로 포괄적인 평가를 수행했습니다.

실험 결과

Llama-3.1-8B, Llama-3.2-3B, Qwen3-8B, Qwen3-4B 모델을 사용하여 5가지 벤치마크(MMLU Pro, GSM8K, IFEval, HellaSwag, WinoGrande)에서 실험을 진행했습니다.

• 성능 향상: ScaleSweep은 모든 설정에서 기존 초기화 방법(AbsMax 및 4/6)보다 일관되게 우수한 성능을 보였습니다.
  • WA (Weights + Activations) 설정에서, ScaleSweep은 BF16 베이스라인과의 격차를 크게 좁혔습니다. 예를 들어, Qwen3-8B에서 **99.50%**의 복구율을 달양했습니다.
  • 가장 공격적인 구성인 WAKVQ (Weights + Activations + KV Cache + Query States) 설정에서도, ScaleSweep은 전체 정밀도 성능의 **>93%**를 보존합니다. 구체적으로 Llama-3.2-3B는 최소 93%의 복구율을 유지하며, Qwen3-4B는 최소 94.5%를 달성합니다.
• 강건성: 이 방법은 양자화 난이도가 높아짐에 따라 강건성을 입증하며, 베이스라인 방법이 더 크게 저하되는 상황에서도 높은 복구율을 유지합니다.
• 효율성: 제한된 스위핑 범위 덕분에 vLLM의 기본 NVFP4 연산자와 비교하여 계산 오버헤드가 무시할 수 있는 수준입니다.
• 캘리브레이션 데이터: GPTQ를 위한 캘리브레이션 데이터셋으로 GSM8K 훈련 분할을 사용하면 특히 수학적 추론 작업, 특히 공격적인 설정에서 성능이 더욱 향상됩니다.

의의 및 주장

본 논문은 ScaleSweep가 하드웨어 지원 부동 소수점 형식을 위한 저비트 양자화의 필수적인 진화라고 주장합니다. AbsMax와 같은 단순한 휴리스틱에서 벗어나 이론적 경계를 통해 FP8 스케일의 제한된 탐색 공간을 활용함으로써, ScaleSweep은 정확하고 엔드 투 엔드(end-to-end) 방식의 NVFP4 양자화를 가능하게 합니다.

저자들은 이 방법이 가장 공격적인 압축 시나리오(KV 캐시 및 쿼리 상태 포함)에서도 전체 정밀도 성능의 93% 이상을 보존할 수 있음을 강조합니다. 이는 스케일 초기화가 올바르게 최적화된다면, NVFP4가 재학습 없이도 효율적인 LLM 추론을 위한 실행 가능한 솔루션이 될 수 있음을 시사합니다. 본 연구는 휴리스틱 초기화와 NVFP4의 최적 솔루션 사이의 간극이 상당하며, 표적화되고 이론적 근거를 갖춘 스위핑이 매우 실용적이고 효과적임을 보여줍니다.