Quanto 양자화가 Transformer 확산 파이프라인의 메모리를 절감합니다

TL;DR

Quantization with Hugging Face Quanto reduces the GPU memory needed for Transformer‑based diffusion pipelines (e.g., PixArt‑Sigma, Stable Diffusion 3, Aura Flow) from ~12 GB to as low as ~5 GB with only minor latency impact and negligible quality loss.

소개 – 왜 메모리가 확산 Transformer에 중요한가

Transformer 백본은 고해상도 텍스트‑투‑이미지 확산 모델에서 지배적인 아키텍처가 되었으며, 파라미터 수는 0.6 B에서 8 B까지 확장됩니다. 모델이 커질수록 GPU 메모리 사용량이 급격히 증가합니다; FP16으로 실행되는 전체 Stable Diffusion 3 추론은 18.8 GB를 차지합니다. 이러한 메모리 장벽은 일반 소비자용 GPU에서의 채택을 제한하고 빠른 실험을 방해합니다. 이 글에서는 Diffusers 라이브러리에 통합된 Quanto의 양자화 유틸리티가 시각적 품질을 유지하면서 메모리 사용량을 크게 줄일 수 있음을 보여줍니다.

Quanto를 사용한 양자화 기본

Quanto는 Hugging Face Optimum 내부에 포함된 PyTorch 기반 양자화 툴킷입니다. FP8, INT8, INT4와 같은 여러 저정밀 포맷으로 가중치 전용 양자화를 지원하며, 모든 Diffusers 모듈에 적용할 수 있습니다.

from optimum.quanto import freeze, qfloat8, quantize
from diffusers import PixArtSigmaPipeline
import torch

pipeline = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")

# Quantize the diffusion transformer to FP8
quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)

동일한 quantize/freeze 호출은 텍스트 인코더나 다른 서브 모듈에도 적용됩니다.

메모리 및 지연 시간 결과 – FP8 가중치 양자화

확산 Transformer만

배치 크기 메모리 (GB) 지연 시간 (s)
1 11.55 (FP8) vs 12.09 (FP16) 1.54 vs 1.20
4 11.55 (FP8) vs 12.09 (FP16) 5.11 vs 4.48

FP8 가중치는 메모리를 약 0.5 GB 절감하지만 지연 시간이 약간 증가합니다.

텍스트 인코더 양자화 추가

배치 크기 텍스트 인코더 양자화? 메모리 (GB) 지연 시간 (s)
1 아니오 11.55 1.54
1 예 5.36 1.60
4 아니오 11.55 5.11
4 예 5.36 5.14

확산 Transformer와 텍스트 인코더를 모두 양자화하면 메모리 사용량이 절반으로 줄어들면서 지연 시간은 거의 변하지 않습니다.

모델 전반에 걸친 일반성

저자들은 세 가지 파이프라인을 평가했습니다:

  • PixArt‑Sigma (0.61 B 파라미터)
  • Stable Diffusion 3 (medium) (2.03 B 파라미터, 텍스트 인코더 3개)
  • Aura Flow (6.84 B 파라미터)

PixArt‑Sigma와 Aura Flow에서는 텍스트 인코더를 양자화하면 항상 큰 메모리 절감 효과가 있었습니다. Stable Diffusion 3은 두 번째 텍스트 인코더(중간 CLIP 변형)를 양자화하면 품질이 저하되므로 선택적 양자화가 필요했습니다. 권장 전략은 다음과 같습니다:

  • 첫 번째 CLIP 인코더만 양자화하거나, 또는
  • 세 번째 T5 인코더만 양자화하거나, 또는
  • 첫 번째와 세 번째 인코더를 모두 양자화합니다.

SD‑3에 대한 대표적인 표(배치 1, 확산 Transformer는 항상 FP8)에서는 메모리가 8.20 GB(세 인코더 모두 양자화)에서 16.40 GB(양자화 없음)까지 나타납니다.

추가 발견

H100에서 bfloat16 vs fp16

NVIDIA H100 GPU에서 bfloat16을 INT8 또는 FP8 가중치와 함께 사용하면 지연 시간이 개선됩니다:

정밀도 양자화 메모리 (GB) 지연 시간 (s)
FP16 INT8 5.363 1.538
BF16 INT8 5.364 1.454
FP16 FP8 5.363 1.601
BF16 FP8 5.363 1.495

INT8 (qint8) 및 결합된 QKV 프로젝션

INT8 가중치는 FP8보다 빠르며, 특히 어텐션 QKV 프로젝션을 결합(fuse_qkv_projections())했을 때 더욱 그렇습니다. PixArt‑Sigma(배치 1)의 경우 지연 시간이 1.538 s(INT8, 결합 안 함)에서 1.504 s(INT8, 결합된 QKV)로 감소했습니다.

공격적인 압축을 위한 INT4 (qint4)

H100에서 bfloat16과 결합하면 INT4는 메모리를 크게 줄입니다(예: PixArt‑Sigma가 9.38 GB에서 3.06 GB로 감소). 그러나 연산이 여전히 bfloat16으로 수행되기 때문에 지연 시간은 약 7.6 s로 증가합니다. 품질 손실이 눈에 띄므로, 저자들은 품질 저하를 완화하기 위해 최종 프로젝션 레이어(proj_out)를 양자화에서 제외할 것을 권장합니다.

양자화된 Diffusers 모델 저장 및 로드

Quanto는 지속 저장 및 재로드가 가능한 모델 클래스를 제공합니다:

from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel, qfloat8

model = PixArtTransformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", subfolder="transformer"
)
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("pixart-sigma-fp8")  # 587 MB checkpoint

로드도 동일한 패턴을 따르며, 양자화된 Transformer는 DiffusionPipeline에 주입할 수 있습니다.

개발자를 위한 실용적인 팁

  • 모듈별로 양자화 유형을 혼합하여 사용합니다(예: 텍스트 인코더에 FP8, 확산 Transformer에 INT8)하여 메모리와 속도 사이의 균형을 맞춥니다.
  • Quanto 양자화를 Diffusers의 기존 메모리 절감 유틸리티(enable_model_cpu_offload() 등)와 결합하여 추가적인 감소를 얻습니다.
  • INT4 배포 시에는 이미지 품질을 유지하기 위해 항상 최종 프로젝션 레이어(exclude="proj_out")를 제외합니다.

결론

Hugging Face Quanto를 사용해 Transformer 기반 확산 파이프라인을 양자화하면 GPU 메모리 요구량을 최대 70 %까지 줄일 수 있습니다(약 12 GB에서 약 5 GB로). 지연 시간은 낮게 유지되고 시각적 품질도 대부분 유지됩니다. 이 기법은 여러 최신 모델에 적용 가능하며, 결과 체크포인트는 크게 작아집니다(예: 587 MB vs 2.44 GB). 개발자는 이러한 워크플로를 기존 Diffusers 파이프라인에 쉽게 통합하고, 다른 메모리 최적화 전략과도 결합할 수 있습니다.

감사

이 글에 대한 광범위한 검토를 해준 Pedro Cuenca에게 감사드립니다.

SUMMARY: Hugging Face 양자화(Quanto)는 Transformer 확산 모델의 GPU 메모리를 약 12 GB에서 약 5 GB로 줄이며 지연 시간과 품질 영향을 최소화합니다.

TITLE: Quanto 양자화가 Transformer 확산 파이프라인의 메모리를 절감합니다

Sources