TorchAO: Unsloth 및 Axolotl 통합과 PARQ를 통한 양자화 인식 학습(QAT) 가속화

섹션별 상세

01

TorchAO의 QAT는 학습 중 '가짜 양자화(Fake Quantization)'를 적용하여 모델이 양자화된 표현에 적응하도록 유도함으로써 사후 양자화(PTQ)의 정확도 손실을 방지한다. 순전파 시에는 양자화 수치를 모방하고 역전파 시에는 고정밀도 가중치를 유지하는 방식을 사용하며, LoRA와 결합 시 원본 가중치는 동결하고 어댑터만 학습시켜 효율성을 높인다.

python

from torchao.quantization import quantize_, Int4WeightOnlyConfig
from torchao.quantization.qat import QATConfig

# PTQ와 동일한 설정 사용
base_config = Int4WeightOnlyConfig(group_size=32)

# 준비 단계: 모델을 "가짜 양자화" 상태로 만들고 학습 준비
quantize_(model, QATConfig(base_config, step="prepare"))
train(model)

# 변환 단계: 모델을 실제 양자화하여 추론 준비
quantize_(model, QATConfig(base_config, step="convert"))

TorchAO를 사용하여 모델에 QAT를 적용하는 기본적인 3단계 흐름

QAT와 LoRA가 결합된 학습 구조 다이어그램 — Diagram동결된 원본 가중치(W)와 학습 가능한 LoRA 어댑터(A, B) 모두에 가짜 양자화(Fake Quant)가 적용되는 과정을 보여준다. 이 구조를 통해 학습 시 메모리를 절감하면서도 최종 양자화 모델의 성능 저하를 미리 방지할 수 있음을 설명한다.

02

Unsloth 프레임워크와의 통합을 통해 INT4 가중치 전용 QAT를 적용한 결과, Gemma3-12B 모델에서 PTQ 대비 성능 저하의 45.5%를 복구하는 성과를 거두었다. 이는 별도의 모델 구조 변경 없이 파인튜닝 단계에서 qat_scheme 플래그만 추가하여 구현 가능하며, ExecuTorch를 통한 스마트폰 배포 시에도 최대 70%의 정확도 복구율을 보여준다.

python

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/gemma3-12b-it",
    max_seq_length = 2048,
    load_in_16bit = True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 32,
    # fp8-int4, fp8-fp8, int4, int8-int4 등 지원
    qat_scheme = "int4",
)

Unsloth 프레임워크에서 qat_scheme 플래그를 사용하여 QAT를 활성화하는 방법

Unsloth INT4 QAT의 성능 복구율 벤치마크 차트 — ChartGemma3, Qwen, Llama 모델군에서 QAT를 적용했을 때 PTQ 대비 성능 저하가 얼마나 복구되는지 수치로 보여준다. Gemma3-4B의 경우 최대 66.9%의 성능 저하를 복구했음을 시각화한다.

03

Axolotl 통합은 멀티 GPU 환경에서의 대규모 모델(최대 72B) QAT를 지원하며, 특히 Blackwell GPU(B200)의 NVFP4 형식을 활용하여 메모리 사용량을 BF16 대비 1/4로 줄인다. 실험 결과 Gemma3-27B 모델에서 NVFP4 QAT를 통해 성능 저하의 71.6%를 복구하고 절대 정확도를 2.3% 향상시켰다.

Axolotl NVFP4 QAT (Gemma3-12B) 성능 비교 차트 — ChartBF16, NVFP4 PTQ, NVFP4 QAT 세 가지 설정을 비교하여 QAT가 PTQ 대비 모든 벤치마크(BBH, MMLU 등)에서 정확도를 유의미하게 향상시킴을 증명한다.

Axolotl NVFP4 QAT (Gemma3-27B) MMLU 세부 항목 성능 차트 — Chart27B 규모의 대형 모델에서도 QAT가 수학 및 STEM 분야의 정확도를 최대 71.6%까지 복구할 수 있음을 구체적인 벤치마크 수치로 제시한다.

04

새로운 PARQ(Piecewise-Affine Regularized Quantization) 기법은 모델 코드 수정 없이 옵티마이저의 step 함수 내에서 양자화를 직접 수행하여 2-3비트 초저비트 학습을 용이하게 한다. 3비트 per-row 모델로 학습된 Phi-4-mini-instruct 모델은 4비트 PTQ 모델과 유사한 정확도를 유지하면서도 메모리 사용량은 58% 수준으로 낮추고 추론 속도는 1.57배 향상시켰다.

python

from torchao.prototype.parq.api import QuantConfig, create_optimizer

def linear_filter_fn(module, fqn):
    return isinstance(module, torch.nn.Linear) and fqn.endswith("weight")

linear_config = QuantConfig(bitwidth=2, group_size=None)
quant_configs_and_filter_fns = [(linear_config, linear_filter_fn)]

# optimizer.step()에서 양자화를 수행하는 QuantOptimizer 적용
optimizer = create_optimizer(
    model,
    quant_configs_and_filter_fns,
    base_optimizer_cls=torch.optim.AdamW,
    base_optimizer_kwargs={"weight_decay": 1e-2},
    quant_per_channel=True,
)

PARQ를 사용하여 옵티마이저 수준에서 2비트 양자화 학습을 설정하는 예시

05

QAT로 최적화된 모델은 ExecuTorch를 통해 iPhone 15 Pro와 같은 모바일 기기에 즉시 배포 가능하며, 비트 수가 낮아질수록 디코딩 속도와 메모리 효율이 선형적으로 개선됨이 확인됐다. 향후에는 PPO, GRPO와 같은 강화학습 알고리즘에 QAT를 적용하여 추론 환경과 일치하는 학습 환경을 구축할 계획이다.

용어 해설

양자화 인식 학습(Quantization-Aware Training): — 모델을 학습하거나 파인튜닝하는 과정에서 양자화로 인한 오차를 미리 반영하여 학습하는 기법이다. 추론 시점에 적용될 양자화 수치 특성을 시뮬레이션함으로써, 사후 양자화(PTQ)에서 발생하는 정확도 저하를 최소화하고 모델이 낮은 정밀도에서도 높은 성능을 유지하도록 돕는다.
가짜 양자화(Fake Quantization): — 학습 중 순전파(Forward Pass) 단계에서 가중치나 활성함수 값을 양자화된 값으로 변환했다가 다시 고정밀도(BF16 등)로 복원하여 계산하는 방식이다. 이를 통해 역전파 시에는 고정밀도 그래디언트를 유지하면서도, 모델이 실제 양자화 환경에서 겪게 될 수치적 제약을 학습 과정에 포함시킬 수 있다.
엔비디아 4비트 부동소수점(NVFP4): — NVIDIA Blackwell 아키텍처(B200 등)에서 지원하는 새로운 4비트 부동소수점 데이터 형식이다. 기존 INT4 대비 더 넓은 동적 범위를 제공하여 LLM 양자화 시 정확도 손실을 줄이면서도 메모리 대역폭과 연산 속도를 획기적으로 개선할 수 있는 하드웨어 특화 형식이다.
조각별 아핀 정규화 양자화(PARQ): — TorchAO에서 새롭게 제안한 옵티마이저 기반의 QAT 알고리즘으로, 3비트 이하의 초저비트 양자화에 최적화되어 있다. 가중치를 양자화 그리드에 더 고르게 분산시키는 '신축 탄성 양자화' 함수를 사용하여, 모델 코드 수정 없이 옵티마이저 단계에서 효율적인 경량화를 구현한다.