QAT로 정확도와 성능을 높이는 gpt-oss 파인튜닝

Reading Time: 5 minutes

대규모 오픈소스 파운데이션 모델 공개는 AI 커뮤니티에 늘 큰 기대를 불러일으킵니다. 새로운 아키텍처적 혁신과 기능이 함께 등장하기 때문입니다. OpenAI 연구소가 GPT-2 이후 처음으로 공개한 오픈소스 모델 패밀리 gpt-oss 역시 예외가 아니었습니다. 이 모델은 전문가 혼합(MoE) 아키텍처, 128K 컨텍스트 길이, 조정 가능한 심층 추론 기능을 제공하며, 가장 큰 변형인 gpt-oss-120B는 오픈 벤치마크에서 OpenAI의 비공개 모델인 o3, o4와 비슷한 성능을 보여주었습니다.

하지만 뛰어난 벤치마크 성능에도 불구하고, 대부분의 파운데이션 모델은 실제 프로덕션 환경에 적용하려면 후처리 학습(post-training) 기법이 필요합니다. 특히 오류 허용 범위가 극도로 낮은 헬스케어와 금융 산업에서는 이러한 과정이 필수적입니다. OpenAI가 gpt-oss를 MXFP4 정밀도로 네이티브 공개한 것은 업계 최초의 시도로, 파인튜닝 측면에서 새로운 과제를 던졌습니다.

이번 글에서는 gpt-oss를 대상으로 한 파인튜닝 워크플로우를 소개합니다. FP4의 성능 이점을 유지하면서도 학습 이후 정확도를 회복하는 방법을 다루며, 그 핵심은 다음과 같습니다.

BF16 업캐스팅 모델에 대한 지도형 파인튜닝(SFT) 수행
NVIDIA TensorRT Model Optimizer를 활용한 양자화 인지 학습(QAT) 적용

SFT와 QAT를 활용한 gpt-oss 파인튜닝

네이티브 FP4 형식에서 모델을 학습하는 새로운 기법들이 최근 주목받고 있습니다. 이 방식은 정확도 회복보다는 학습 시간 최적화 측면에서 큰 가능성을 보여주고 있습니다. 그러나 gpt-oss 파인튜닝에서는 네이티브 MXFP4 정밀도가 아직 안정적인 정확도를 보장하지 못하고 있습니다. 이로 인해 파인튜닝 과정이 까다로운데, 안정적인 그래디언트 누적을 보장하기 위해 먼저 더 높은 정밀도로 업캐스팅해야 합니다.

업캐스팅이 완료되면, 더 높은 정밀도의 체크포인트를 활용해 첫 번째 SFT(Supervised Fine-Tuning)를 수행합니다. 이후 QAT(Quantization-Aware Training)를 적용해 모델을 다시 FP4 정밀도로 변환하면서, 작업 특화 성능을 회복할 수 있습니다(Figure 2). 이 접근법은 SFT로 작업별 행동을 강화하고, QAT로 가중치를 목표 저정밀 형식에 적응시켜 정렬과 성능을 동시에 확보한 상태로 배포할 수 있게 합니다.

*그림 1. gpt-oss의 QAT 및 SFT 워크플로우: 업캐스팅된 체크포인트부터 양자화 준비 모델까지*

이 레시피의 전체 코드는 Model Optimizer 저장소에서 확인할 수 있습니다. 이번 학습 예시는 Hugging Face의 gpt-oss-recipes에 포함된 파인튜닝 예제를 기반으로, QAT와 권장 컴포넌트를 통합하도록 수정한 것입니다. 주요 단계는 다음과 같습니다.

1. MXFP4 체크포인트를 BF16/FP16으로 업캐스팅: Hugging Face의 Transformers 라이브러리를 사용하면 손쉽게 MXFP4 체크포인트를 BF16/FP16으로 변환할 수 있습니다. 업캐스팅을 통해 그래디언트가 더 안정적으로 수렴하며, 이후 QAT를 적용해 FP4로 다시 양자화할 때 정확도를 효과적으로 회복할 수 있습니다. 이 과정은 전체를 FP4로 파인튜닝했을 때 얻는 계산 효율이 미미하고, 파인튜닝이 사전 학습에 비해 훨씬 적은 토큰을 사용한다는 점에서 합리적인 절충안이 됩니다.

2. SFT 수행: 목적에 맞는 파인튜닝 데이터셋을 사용해, 업캐스팅된 BF16 또는 FP16 모델에 대해 양자화를 적용하지 않은 상태에서 지도형 파인튜닝을 진행합니다.

3. TensorRT Model Optimizer를 사용한 양자화: BF16으로 파인튜닝된 모델을 mtq.quantize() 함수를 사용해 양자화합니다. 이 함수는 모델을 PTQ(Post-Training Quantization)나 QAT(Quantization-Aware Training)에 적합하도록 준비합니다.

import modelopt.torch.quantization as mtq

config = mtq.MXFP4_MLP_WEIGHT_ONLY_CFG

# Define forward loop for calibration
def forward_loop(model):
    for data in calib_set:
        model(data)

# quantize the model and prepare for QAT
model = mtq.quantize(model, config, forward_loop)

4. FP4로 양자화된 모델 파인튜닝: 두 번째 파인튜닝 단계는 작은 학습률(예: Adam 옵티마이저를 사용한 1e-5)로 진행되며, 바로 이 과정이 QAT(Quantization-Aware Training) 단계입니다.

# QAT with a regular finetuning pipeline
train(model, train_loader, optimizer, scheduler, ...)

QAT의 최적 하이퍼파라미터와 학습 시간은 상황에 맞게 조정할 수 있는 변수입니다. 다만 Step 1을 건너뛰고 바로 QAT를 진행하면 정확도가 낮게 나오는 경향이 있었습니다. 따라서 먼저 고정밀 파인튜닝을 수행한 뒤 QAT를 적용하는 방법을 권장합니다. 만족스러운 수렴 결과를 얻은 후에는 Model Optimizer API를 활용해 모델을 표준 PyTorch 체크포인트로 내보낼 수 있으며, 이를 통해 공개 벤치마크와 커스텀 태스크에 대해 검증할 수 있습니다.

MXFP4 QAT 파인튜닝이 gpt-oss에 미치는 영향

위의 QAT 파인튜닝 워크플로우 효과를 보여주기 위해 두 가지 다운스트림 평가 작업을 분석했습니다. 첫 번째는 OpenAI Cookbook의 다국어 데이터셋을 활용한 비영어권 추론 능력 향상이며, 두 번째는 Amazon의 FalseReject 데이터셋을 활용한 안전한 사용자 프롬프트에 대한 불필요한 거부 감소입니다. gpt-oss는 초기 상태에서 이 두 작업에서 각각 16%, 30%의 점수를 기록하며 개선 여지가 있음을 보여주었습니다(Figure 3). 그러나 이번 레시피를 적용한 뒤 두 작업 모두에서 98%의 통과율을 기록하며, 매우 큰 성능 향상을 달성했습니다.

*그림 2. gpt-oss-20b의 원본 모델과 SFT, PTQ, QAT 적용 후 검증 통과율 비교*

이번 결과는 gpt-oss 파인튜닝에서 업캐스팅과 QAT를 적용해 정확도를 회복하는 접근법이 효과적임을 잘 보여줍니다. 그러나 여전히 작업별 성능을 더 끌어올릴 기회는 남아 있습니다. NVIDIA Blackwell 아키텍처의 등장과 함께 NVFP4라는 새로운 FP4 형식이 도입되었는데, 이는 학습과 추론 효율을 모두 고려해 설계되었습니다. NVFP4를 QAT와 결합하면 정확도 회복을 한층 더 강화할 수 있는 가능성이 열리게 됩니다(Figure 3).

*그림 3. MXFP4 체크포인트 기반의 배포 옵션 흐름도 — 직접 배포, QAT를 통한 MXFP4 변환, 또는 과제 특화 정확도를 높이기 위한 NVFP4 변환 포함*

NVFP4는 개발자가 2세대 NVIDIA Transformer Engine의 특화 명령어를 활용할 수 있도록 하며, 최대 15 PFLOPs에 달하는 FP4 NVIDIA Blackwell Ultra 연산을 활용해 더 높은 모델 정확도를 제공합니다. 또한 E4M3 FP8 스케일링 정밀도는 “가짜 양자화(fake quantization)” 과정에서 효과적으로 작동하여, 순전파 단계에서 발생하는 양자화 오류를 줄여줍니다. 이를 통해 원래 모델 가중치가 목표 정밀도에 더 쉽게 적응할 수 있습니다. 앞서 소개한 MXFP4 레시피를 NVFP4에 맞게 적용하려면 단 한 줄만 수정하면 됩니다. 아래 코드 스니펫처럼 간단히 변경할 수 있습니다.

# To do NVFP4 PTQ/QAT, simply replace MXFP4_MLP_WEIGHT_ONLY_CFGconfig = mtq.NVFP4_MLP_WEIGHT_ONLY_CFG
# Alternatively, do weight-activation quantization for better performanceconfig = mtq.NVFP4_MLP_ONLY_CFG

gpt-oss 파인튜닝 레시피 적용 후 MXFP4와 NVFP4의 검증 손실을 비교해 본 결과, NVFP4 버전이 일관되게 더 나은 수렴을 보였습니다. 특히 이 작업들에서 NVFP4는 검증 손실이 MXFP4 대비 2~3% 더 낮게 나타났습니다(Figure 4). 이러한 성능 향상은 심층 추론, 더 엄격한 기준, 오류 허용 범위가 낮은 다운스트림 작업과 같은 까다로운 환경에서도 유리하게 작용할 수 있습니다.

*그림 4. NVFP4와 MXFP4 포맷에서 SFT, PTQ, QAT 적용 시 다국어 및 FalseReject 과제의 검증 손실 비교*

앞으로 NVIDIA TensorRT-LLM에서 gpt-oss의 NVFP4 지원이 추가되면, 개발자들은 NVFP4를 손쉽게 활용할 수 있게 될 것입니다. 또한 NVIDIA는 다른 오픈소스 추론 프레임워크에서도 gpt-oss NVFP4 지원을 우선적으로 추진하고 있습니다. 그때까지는 MXFP4 기반의 SFT + QAT 워크플로우가 입증된 방법으로 남아 있습니다. 하지만 NVFP4 기반 gpt-oss가 완전히 지원되면, 동일한 레시피로 NVIDIA Blackwell에서 한층 더 높은 정확도를 확보할 수 있을 것입니다.

파인튜닝된 모델 배포

이 레시피 실행을 마친 후에는 Model Optimizer 저장소에서 제공하는 편의 스크립트를 사용해 BF16으로 학습된 체크포인트를 MXFP4 형식으로 변환할 수 있습니다.

python examples/gpt-oss/convert_oai_mxfp4_weight_only.py --model_path qat_model_dir/ --output_path qat_model_mxfp4/

이 레시피를 통해 생성된 MXFP4 체크포인트는 SGLang, TensorRT-LLM, vLLM 등 업스트림 프레임워크에서 테스트를 완료했습니다. 아래 명령어를 사용하면 TensorRT-LLM 1.1.0rc1을 통해 배포할 수 있습니다.

# Use trtllm-serve to host endpoint
trtllm-serve qat_model_mxfp4/ --tokenizer <tokenizer_path> --max_batch_size <max_batch_size> --max_num_tokens <max_num_tokens> --max_seq_len <max_seq_len> --tp_size <tp_size> --pp_size <pp_size> --host 0.0.0.0 --kv_cache_free_gpu_memory_fraction 0.95

결론

gpt-oss 파인튜닝의 핵심 과제는 FP4에서 정확도를 회복하면서도, 저정밀 배포가 제공하는 효율성을 유지하는 것입니다. BF16으로 업캐스팅한 뒤 SFT를 수행하고, 이어서 QAT를 적용하는 접근법은 이 간극을 메우며, 가중치를 저정밀 환경에 적응시켜 모델이 프로덕션에서 안정적이면서도 효율적으로 동작하도록 만듭니다.

실제로 이 방식은 정확도를 회복할 뿐만 아니라 작업 특화 성능을 강화하여 사용자 경험, 안전성, 다운스트림 애플리케이션에서의 활용성을 크게 향상시킵니다. 이는 곧 더 높은 서비스 품질과 더 나은 투자 수익률(ROI)로 이어집니다. 앞으로 NVFP4는 더 빠른 수렴과 더 엄격한 기준, 심층 추론 같은 까다로운 환경에서도 여유 있는 성능을 제공할 것이며, TensorRT-LLM과 다른 프레임워크에서 gpt-oss NVFP4 지원이 추가되면 이러한 장점은 더욱 확대될 것입니다. 지금 바로 NVIDIA Model Optimizer 리포지토리에서 제공되는 완전한 SFT + QAT 레시피로 시작해 보세요.