Hugging Face, Diffusers에 Nunchaku 4-bit Diffusion 추론 통합

Hugging Face, Diffusers에 Nunchaku 4-bit Diffusion 추론 통합

Hugging Face는 Nunchaku Lite를 Diffusers 라이브러리에 통합하여, 사용자가 4-bit 가중치 및 활성화(W4A4)를 사용하여 대규모 확산 트랜스포머(diffusion transformers)를 실행할 수 있도록 했습니다. 이 통합을 통해 VRAM 요구 사항을 최대 50%까지 줄이고, BF16 베이스라인 대비 디노이징 지연 시간을 약 30% 개선하여 소비자용 GPU에서도 고성능 이미지 생성이 가능해졌습니다.

SVDQuant 및 Nunchaku 엔진

Nunchaku Lite는 확산 트랜스포머의 가중치와 활성화 모두에서 발견되는 대규모 아웃라이어(outliers)를 처리하도록 설계된 양자화 방법인 SVDQuant를 기반으로 합니다. 표준 4-bit 양자화와 달리, SVDQuant는 아웃라이어를 가중치로 이동시켜 격리하고, 각 가중치 행렬의 가장 까다로운 부분을 작은 16-bit 저차원(low-rank) 브랜치로 표현합니다. 남은 잔차(residual)는 4-bit로 양자화됩니다.

참조용 Nunchaku CUDA 추론 엔진은 저차원 다운 프로젝션(low-rank down projection)을 양자화 커널과 결합하고, 저차원 업 프로젝션(low-rank up projection)을 4-bit 연산 커널과 결합한 융합 커널(fused kernels)을 사용하여 이 프로세스를 최적화하며, 이를 통해 16-bit 브랜치에 대한 메모리 액세스 오버헤드를 효과적으로 제거합니다.

Diffusers 내 Nunchaku Lite 통합

Nunchaku Lite는 Diffusers 내에서 간소화된 통합 경로를 제공하여, 별도의 추론 엔진이나 로컬 CUDA 컴파일 없이 from_pretrained()를 통해 Nunchaku 스타일의 체크포인트를 로드할 수 있게 합니다. 이는 nn.Linear 모듈을 런타임 SVDQ/AWQ 선형 레이어로 패치함으로써 구현됩니다.

Nunchaku Lite는 두 가지 주요 커널 제품군을 사용합니다:

  • svdq_w4a4: 대부분의 연산이 발생하는 트랜스포머 어텐션 및 MLP 프로젝션에 사용됩니다. INT4 및 NVFP4 변형으로 제공됩니다.
  • awq_w4a16: 적응형 정규화(adaptive normalization) 및 변조 프로젝션(예: FLUX adanorm_single / adanorm_zero)과 같이 메모리 제한적이고 정밀도에 민감한 레이어에 사용되며, 4-bit 가중치와 16-bit 활성화를 활용합니다.

Nunchaku Lite가 기존 Nunchaku 엔진의 아키텍처별 융합 실행 경로를 포함하지는 않지만, 여전히 약 30%의 속도 향상과 동일한 VRAM 절감 효과를 제공합니다.

하드웨어 지원 및 성능

Nunchaku Lite 커널에 대한 지원은 GPU 세대 및 정밀도에 따라 다릅니다:

Scheme Precision Supported GPUs
svdq_w4a4 nvfp4 Blackwell (RTX 50 series, RTX PRO 6000, B200)
svdq_w4a4 int4 Turing / Ampere / Ada (RTX 30 & 40 series, A100, L40S)
awq_w4a16 int4 Turing / Ampere / Ada (RTX 30 & 40 series, A100, L40S)

Volta 및 Hopper GPU는 현재 지원되지 않습니다.

벤치마크

NVIDIA RTX PRO 6000 (Blackwell)에서 1024x1024 해상도로 ERNIE-Image-Turbo 체크포인트를 사용하여 측정됨:

Configuration Full pipeline Denoise loop Peak VRAM Speedup
BF16 baseline 3.00 s 2.86 s 31.1 GB 1.0x
Nunchaku Lite NVFP4 2.27 s 2.13 s 20.6 GB 1.35x
Nunchaku Lite NVFP4 + torch.compile 1.68 s 1.53 s 20.6 GB 1.8x
Nunchaku Lite NVFP4 + NF4 text encoder 2.29 s 2.13 s 16.0 GB 1.35x

Nunchaku Lite를 torch.compile과 결합하면 엔드 투 엔드 속도가 1.8배 증가하며, bitsandbytes NF4 양자화 텍스트 인코더를 추가하면 피크 VRAM이 16.0 GB로 더욱 감소합니다.

diffuse-compressor를 이용한 커스텀 모델 양자화

Nunchaku Lite는 아키텍처에 구애받지 않습니다. diffuse-compressor 툴킷을 사용하면 사용자가 자신의 Diffusers 모델을 보정(calibrate), 양자화, 패키징 및 게시할 수 있습니다. 워크플로우는 다음과 같습니다:

  1. Inspection: 범용 스캐너가 SVDQ W4A4 대상(반복되는 트랜스포머 블록)과 AWQ W4A16 대상(변조 선형 레이어)을 식별합니다.
  2. Quantization: SVDQuant를 실행하여 양자화된 체크포인트를 생성합니다 (int4 또는 nvfp4 지원).
  3. Packaging: 양자화된 트랜스포머를 베이스 파이프라인과 결합하고 transformer/config.jsonnunchaku_lite 구성을 생성합니다.
  4. Verification: DiffusionPipeline.from_pretrained()를 통해 모델을 로드하고 최종 결과를 Hugging Face Hub에 푸시합니다.

구조적 재작성 (Structural Rewrites)

최대 성능을 위해 일부 모델은 별도의 Q, K, V 프로젝션을 단일 to_qkv 모듈로 결합하는 것과 같은 구조적 재작성이 필요합니다. 범용 diffuse-compressor 경로가 이러한 변경 사항을 추론할 수는 없지만, 융합 연산을 활성화하기 위해 rootonchair/nunchaku-lite에서 제공하는 모델별 타겟 구성 및 런타임 어댑터를 통해 구현할 수 있습니다.

바로 사용 가능한 체크포인트

Hub에서 다음과 같은 여러 사전 양자화된 체크포인트를 사용할 수 있습니다:

  • ERNIE-Image-Turbo: NF4 텍스트 인코더를 포함한 INT4 및 NVFP4 변형으로 제공됩니다.
  • Krea 2 Turbo: NVFP4로 제공됩니다.
  • lite-infer: 추가적인 Nunchaku Lite 체크포인트 모음입니다.

Sources