NVIDIA B200에서 MXFP8 및 NVFP4 양자화를 활용한 확산 모델 추론 가속화

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

핵심 요약

NVIDIA Blackwell(B200) 아키텍처에서 지원하는 마이크로스케일링 포맷인 MXFP8과 NVFP4를 활용해 이미지 및 비디오 확산 모델의 추론 효율성을 극대화했다. Flux.1-Dev, LTX-2, QwenImage 모델을 대상으로 테스트한 결과, BF16 대비 MXFP8은 최대 1.26배, NVFP4는 최대 1.68배의 속도 향상을 달성했다. 단순히 전체 모델을 양자화하는 대신 특정 레이어를 제외하는 선택적 양자화(Selective Quantization)와 CUDA Graphs를 적용해 정확도 손실을 최소화하면서 CPU 오버헤드를 줄였다. LPIPS 지표를 통해 양자화 모델이 생성한 이미지의 시각적 품질이 원본과 유사함을 검증했으며, 실무에서 즉시 적용 가능한 TorchAO 기반의 구현 코드를 포함한다.

배경

NVIDIA Blackwell 아키텍처 GPU (B200 등), PyTorch Nightly 버전 및 TorchAO 라이브러리, Diffusers 라이브러리에 대한 기본 지식

대상 독자

NVIDIA Blackwell GPU 기반으로 고성능 이미지/비디오 생성 서비스를 구축하려는 ML 엔지니어 및 최적화 연구자

의미 / 영향

이 기술은 고해상도 생성 모델의 운영 비용을 획기적으로 낮추고 추론 처리량을 높여, Blackwell 아키텍처 기반 클라우드 인프라의 가치를 극대화합니다. 특히 NVFP4와 같은 초저정밀도 포맷이 시각적 품질을 유지하면서도 실용 가능하다는 것을 입증함으로써 차세대 양자화 표준의 방향성을 제시합니다.

섹션별 상세

확산 모델의 높은 메모리 및 연산 요구사항을 해결하기 위해 Blackwell 아키텍처 전용 마이크로스케일링 포맷을 도입했다. MXFP8은 8비트 블록 스케일링을 통해 품질 저하 없는 가속을 제공하며, NVFP4는 4비트 포맷으로 이론상 최대 처리량과 최소 메모리 점유율을 실현한다. 이를 통해 B200 GPU에서 기존 BF16 대비 현저한 성능 이득을 얻을 수 있다.

선택적 양자화 기법을 사용하여 모델의 정확도와 추론 속도 사이의 최적의 균형을 찾았다. 가중치나 활성화 형상이 너무 작아 양자화 이득보다 오버헤드가 큰 레이어(min(M, K, N) < 1024)나 정확도에 민감한 임베딩 레이어 등을 양자화 대상에서 제외했다. Flux.1-Dev 테스트 결과, 선택적 양자화가 전체 양자화보다 낮은 LPIPS 점수(더 높은 품질)와 더 빠른 속도를 동시에 기록했다.

python

from diffusers import DiffusionPipeline, TorchAoConfig, PipelineQuantizationConfig
import torch
from torchao.prototype.mx_formats.inference_workflow import (
    NVFP4DynamicActivationNVFP4WeightConfig,
)

config = NVFP4DynamicActivationNVFP4WeightConfig(
    use_dynamic_per_tensor_scale=True,
    use_triton_kernel=True,
)
pipe_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(config)}
)
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    quantization_config=pipe_quant_config
).to("cuda")
pipe.transformer.compile_repeated_blocks(fullgraph=True)

TorchAO의 NVFP4 양자화 설정을 Diffusers 파이프라인에 적용하고 지역 컴파일을 활성화하는 예시

작은 배치 사이즈에서 발생하는 CPU 오버헤드를 해결하기 위해 CUDA Graphs를 결합한 'reduce-overhead' 컴파일 모드를 적용했다. QwenImage 모델의 배치 사이즈 1 환경에서 CUDA Graphs 적용 시 GPU 트레이스의 공백(bubbles)이 사라지며 약 1.81배의 추가 속도 향상이 확인됐다. 이는 실시간 서비스 환경에서 지연 시간을 단축하는 데 결정적인 역할을 한다.

양자화에 따른 이미지 품질 저하를 객관적으로 평가하기 위해 LPIPS 지표를 활용한 벤치마크를 수행했다. Flux.1-Dev의 경우 MXFP8은 0.11, NVFP4는 0.44의 평균 LPIPS를 기록하여 시각적으로 원본과 매우 유사한 결과를 보였다. 반면 QwenImage는 양자화에 더 민감하게 반응하여 추가적인 알고리즘 최적화의 필요성을 시사했다.

실무 Takeaway

NVIDIA B200 환경에서 Flux.1-Dev 모델에 NVFP4 양자화를 적용하면 BF16 대비 메모리 사용량을 약 32% 절감하고 추론 속도를 1.59배 향상시킬 수 있다.
양자화 적용 시 모든 레이어를 변환하지 말고, 행렬 크기가 1024 미만인 레이어는 제외하는 선택적 양자화 전략을 통해 정확도와 성능의 최적점을 확보해야 한다.
배치 사이즈 1의 실시간 추론 최적화가 필요하다면 torch.compile의 reduce-overhead 모드를 활성화하여 CPU 오버헤드를 제거해야 한다.

언급된 리소스

GitHub실험 재현 코드 저장소

API DocsTorchAO 문서: MXFP8

API DocsTorchAO 문서: NVFP4