Nemotron 3.5 Lightning의 NVFP4 양자화와 QAD

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

Nemotron 3.5 Lightning의 BF16 체크포인트를 NVFP4로 양자화하면 모델 크기를 66GB에서 22GB로 줄이고 최대 4배 높은 throughput을 확보할 수 있지만, 공격적인 PTQ만으로는 정확도 손실이 커질 수 있습니다. NVIDIA Model Optimizer의 QAD pipeline은 먼저 PTQ로 W4A16 Student를 만든 뒤, simulated quantization을 거친 Student의 logits를 고정된 BF16 Teacher의 logits에 KL Divergence로 맞추며 양자화 noise를 학습합니다. 공격적인 중간 체크포인트에서 median score recovery는 Checkpoint A 기준 96.33%에서 99.72%로, Checkpoint B 기준 95.84%에서 98.53%로 상승했습니다. 최종 체크포인트에서는 보수적인 PTQ 설정으로 전체 중앙값의 차이가 작았지만 Terminal-Bench v2.1과 SWE-Bench Multilingual 같은 agentic·coding benchmark에서 QAD가 PTQ를 앞섰습니다.

빠른 이해

새로운 점

PTQ 이후 Teacher의 출력 분포와 실제 양자화 noise를 함께 학습하는 QAD로 Nemotron 3.5 Lightning의 공격적인 NVFP4 압축을 가능하게 한 점입니다.

핵심 메커니즘

BF16 Teacher를 PTQ로 W4A16 NVFP4 Student로 변환하고, Student의 매 forward pass에 simulated quantization을 적용합니다. 동일 입력의 Teacher와 Student logits를 KL Divergence로 비교해 Student weight를 갱신하며, PTQ 보정 방식에 따라 quantization scale을 dynamic 또는 frozen 전략으로 처리합니다. 이 과정을 통해 BF16 66GB 모델을 22GB NVFP4 checkpoint로 압축하면서 양자화 noise에 적응한 출력 행동을 학습합니다.

핵심 수치

체크포인트 크기: BF16 66GB -> NVFP4 22GB- 기사 도입부 기준
추론 throughput: 최대 4배 향상- Nemotron 3.5 Lightning NVFP4 checkpoint 기준
Checkpoint A median score recovery: PTQ 96.33% -> QAD 99.72%- 동일한 21.19GB 체크포인트, 11개 benchmark 중 10개 개선
Checkpoint B median score recovery: PTQ 95.84% -> QAD 98.53%- 동일한 21.19GB 체크포인트
최종 Terminal-Bench v2.1: PTQ 22.05 -> QAD 25.84- BF16 24.44
최종 SWE-Bench Multilingual: PTQ 37.00 -> QAD 38.07- BF16 37.13

섹션별 상세

NVFP4 압축의 목표

Nemotron 3.5 Lightning의 BF16 체크포인트는 66GB이므로 메모리와 추론 처리량을 함께 개선하려면 더 공격적인 양자화가 필요합니다. 이 글은 많은 가중치를 4비트 NVFP4로 저장해 모델 크기를 22GB로 줄이고 최대 4배 높은 throughput을 확보하는 과정을 다룹니다. 일반적인 PTQ는 고정밀 모델을 낮은 비트 격자에 매핑하지만, 정밀도 손실이 커질수록 agentic·coding benchmark의 점수가 하락합니다. QAD는 이 손실을 학습 단계에서 보정해 더 낮은 메모리 사용량과 높은 처리량을 유지하면서 BF16 기준선에 가까운 정확도를 확보하도록 설계됐습니다.

두 단계 QAD 파이프라인

QAD의 첫 단계에서는 BF16 Nemotron 3.5 Lightning을 PTQ로 W4A16 NVFP4 Student 체크포인트로 변환합니다. 두 번째 단계에서는 동일한 입력을 고정된 BF16 Teacher와 NVFP4 Student에 통과시키고, Student의 forward pass에는 simulated quantization을 삽입해 실제 추론 때 발생할 양자화 noise를 학습 과정에 반영합니다. 이후 Teacher와 Student의 logits 분포 차이를 KL Divergence loss로 계산하고 Student 가중치를 갱신합니다. 이 입력·양자화·출력 분포 정렬 흐름 덕분에 Student는 다음 토큰 하나만 맞추는 대신 원래 모델의 출력 행동을 모방하며, 기사에 포함된 공정도는 PTQ checkpoint에서 최종 QAD checkpoint로 이어지는 이 구조를 나타냅니다.

BF16 Teacher에서 PTQ로 W4A16 NVFP4 Student를 만든 뒤 QAD로 증류해 최종 NVFP4 체크포인트를 생성하는 2단계 흐름입니다. — Diagram도식의 왼쪽에서는 Full-Precision BF16 모델이 고정된 Teacher와 PTQ 입력으로 동시에 사용됩니다. PTQ는 낮은 정밀도의 Student checkpoint를 만들고, 다음 단계에서 Student의 simulated quantization 출력과 Teacher의 출력 분포를 distillation loss로 비교해 Student를 갱신합니다. 최종 출력은 공격적인 양자화 뒤에도 기준선에 가까운 정확도를 회복한 QAD checkpoint입니다.

PTQ 레시피와 공격성 조절

PTQ 단계에서는 Mamba linear layer를 FP8 대신 W4A16으로 낮추고, lm_head는 W4A16으로 양자화하며 attention projection layer는 BF16에 남기는 조합을 시험했습니다. 보정에는 1,000개 샘플을 사용했고 시퀀스 길이는 8K부터 128K까지 비교했으며, four_over_six 레시피에서는 32K가 가장 좋은 PTQ 결과를 냈습니다. 가장 공격적인 four_over_six + NVFP4 KV 설정은 K와 V를 NVFP4 W4A4로 낮추고 Q 및 일부 batched matrix multiplication은 BF16으로 유지합니다. PTQ만으로 최종 품질을 맞추는 경우보다 QAD를 후속 단계로 두면 초기 정확도 회복 목표를 95~99%로 낮춰 메모리와 latency 이득을 먼저 확보할 여지가 생깁니다.

python

import modelopt.torch.quantization as mtq # define forward loop with datasetdef forward_loop(model): for batch in calib_dataloader: model(batch) # Quantize base model to NVFP4 to create the PTQ student checkpoint# Example uses W4A16_NVFP4_CFG for quantizationmodel = mtq.quantize(model, mtq.W4A16_NVFP4_CFG, forward_loop=forward_loop)

보정 데이터셋을 순회하는 forward loop을 정의한 뒤 NVIDIA Model Optimizer로 base model을 W4A16 NVFP4 Student 체크포인트로 양자화합니다.

Dynamic scale과 frozen scale

QAD 학습에서는 PTQ 보정 방식에 따라 quantization scale을 매 step 다시 계산하거나 처음 값으로 고정합니다. max-calibrated checkpoint는 현재 weight와 activation에서 scale을 다시 계산하는 dynamic scale 방식을 사용하고, MSE 또는 four_over_six checkpoint는 PTQ에서 오차를 최소화하도록 찾은 scale을 frozen scale로 유지합니다. Dynamic scale에서는 BF16 weight와 scale이 함께 변하지만, frozen scale에서는 scale을 재탐색하지 않고 weight만 업데이트하므로 값비싼 MSE 검색을 매 학습 step 반복하지 않습니다. 제공된 비교 도식은 두 방식 모두 activation에서 simulated quantization과 GEMM, KL loss를 거치지만 scale 갱신 경로가 dynamic 방식에만 존재한다는 차이를 나타냅니다.

Dynamic scale은 학습 중 scale을 매번 재계산하고 frozen scale은 PTQ에서 얻은 scale을 고정하는 QAD 학습 방식의 차이를 나타냅니다. — Diagram두 경로 모두 Activations, simulated quantization, GEMM, KL loss를 순서대로 통과하며 gradient는 FP weights로 돌아갑니다. 위쪽 dynamic scale 경로는 현재 weight에서 scale s(t)를 매 step 다시 계산하고, 아래쪽 frozen scale 경로는 PTQ 후 저장한 scale s*를 유지하면서 weight만 갱신합니다. 기사에서는 max-calibrated PTQ에는 dynamic scale을, MSE 기반 PTQ에는 frozen scale을 연결합니다.

중간 체크포인트의 정확도 회복

공격적인 W4A16 설정에서 BF16 기준 체크포인트는 65.85GB, PTQ와 QAD 체크포인트는 모두 21.19GB로 동일하게 맞춰 양자화 방식의 차이를 비교했습니다. Checkpoint A에서는 PTQ의 median score recovery가 96.33%였지만 200회 QAD 뒤 99.72%로 올라갔고, 11개 benchmark 중 10개가 개선됐습니다. AIME 2025는 BF16 대비 PTQ에서 3.70점 하락했지만 QAD 후 차이가 0.57점으로 줄었고, SciCode Subtask와 HLE는 BF16보다 소폭 높은 점수를 기록했습니다. 업데이트된 평가군의 Checkpoint B에서도 median score recovery가 PTQ 95.84%에서 QAD 98.53%로 상승했으며, AA v4.1 Index는 20.03에서 23.48로 높아졌습니다.

최종 체크포인트와 재현 방법

출시용 NVFP4 레시피는 정확도를 우선한 보수적 양자화이므로 PTQ의 median score recovery가 99.24%까지 올라갔고 QAD는 98.97%를 기록해 전체 중앙값의 추가 회복 폭은 작았습니다. 대신 QAD는 Terminal-Bench v2.1에서 PTQ보다 3.79점, SWE-Bench Multilingual에서 1.07점, HLE에서 0.65점 높은 점수를 기록하는 등 agentic·coding benchmark에서 차이를 냈습니다. NVIDIA Model Optimizer는 megatron_lm_qad.yaml launcher로 BF16 Teacher import, PTQ Student 생성, Teacher 기반 증류, Hugging Face 체크포인트 export를 연결하며, Megatron-Bridge는 토큰화 작업까지 포함한 mbridge_qad.yaml 흐름을 제공합니다. 재현 설정에는 5e-6 constant learning rate, dropout 비활성화, gradient clipping 1.0, 2 nodes × 8 GPUs 또는 Megatron-Bridge의 8 nodes × 4 GPUs 구성이 포함됩니다.

bash

# from tools/launchersource .env-slurmuv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml --yes

Megatron-LM용 QAD YAML launcher를 호출해 PTQ와 증류를 포함한 전체 파이프라인을 실행합니다.

text

task_1: # quantize the BF16 teacher into the NVFP4 student script: common/megatron_lm/quantize/quantize.sh args: - --seq-length 32768 --max-position-embeddings 1048576 - --calib-size 32 environment: - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG - TP: "1" - EP: "4"

BF16 Teacher를 NVFP4 Student로 변환하는 PTQ 작업에서 32,768 시퀀스 길이와 32개 보정 샘플, MAMBA_MOE_NVFP4_CONSERVATIVE_CFG를 지정합니다.

text

task_2: # distill the NVFP4 student against the BF16 teacher script: common/megatron_lm/train/sft.sh args: - --seq-length 32768 --max-position-embeddings 1048576 - --micro-batch-size 1 --global-batch-size 16 - --train-samples 6400 # -> 400 iterations - --modelopt-enabled - --export-kd-teacher-load /cicd/megatron-lm-bf16/.../BF16-MCore - --lr 5.0e-6 --lr-decay-style constant --lr-warmup-samples 0 - --clip-grad 1.0 --weight-decay 0.0 - --attention-dropout 0.0 --hidden-dropout 0.0 environment: - DATASET: nvidia/Nemotron-Post-Training-Dataset-v2 - TP: "1" - EP: "4"

고정된 BF16 Teacher를 기준으로 NVFP4 Student를 증류하며, 학습률 5.0e-6, 400 iterations, gradient clipping 1.0과 Nemotron-Post-Training-Dataset-v2를 사용합니다.

용어 해설

양자화 인식 증류(Quantization-Aware Distillation): — 양자화 과정에서 발생하는 오차를 고려해 저정밀도 Student 모델을 학습하는 방법입니다. 먼저 PTQ로 BF16 Teacher를 양자화한 뒤, 추론 시뮬레이션을 거친 Student의 출력 분포를 고정된 Teacher의 출력 분포에 맞춥니다. KL Divergence를 손실로 사용해 공격적인 양자화 뒤에도 원래 모델의 동작을 회복하는 데 의미가 있습니다.
사후 학습 양자화(Post-Training Quantization): — 학습이 끝난 모델의 가중치와 활성값을 낮은 비트 형식으로 변환하는 방법입니다. 보정 데이터로 각 텐서의 범위나 오차를 측정한 뒤 BF16 같은 고정밀 표현을 W4A16 또는 NVFP4 형식으로 매핑합니다. 추가 학습 없이 모델 크기와 메모리 사용량을 줄일 수 있지만, 공격적인 설정에서는 정확도 손실이 커질 수 있습니다.
KL 발산(KL Divergence): — 두 확률 분포가 얼마나 다른지 측정하는 손실 함수입니다. QAD에서는 같은 입력에 대한 Teacher와 Student의 next-token logits를 확률 분포로 바꾼 뒤 Student 분포가 Teacher 분포를 따르도록 오차를 계산합니다. 정답 토큰 하나만 맞추는 대신 원래 모델의 전체 출력 행동을 보존하는 데 사용됩니다.
NVFP4: — NVIDIA가 사용하는 4비트 부동소수점 기반 저정밀도 형식으로, Nemotron 3.5 Lightning의 가중치와 일부 KV cache를 압축하는 데 적용됐습니다. 기사에서는 W4A16 설정으로 많은 가중치를 4비트로 저장해 BF16 체크포인트의 66GB 크기를 22GB 수준으로 줄였습니다. 낮은 정밀도에 따른 오차는 QAD 학습으로 보정합니다.

기술

Nemotron 3.5 Lightning
NVFP4
W4A16
Post-Training Quantization
Quantization-Aware Distillation
NVIDIA Model Optimizer
Megatron-LM
Megatron-Bridge
KL Divergence
Mamba
KV cache

언급된 리소스

문서NVIDIA Model Optimizer 4o6 yaml recipe

GitHubNVIDIA Model Optimizer Github Repo

API DocsNVIDIA Model Optimizer QAD Documentation

문서원문 링크