Nunchaku 4비트 확산 추론을 Diffusers에 적용하다

대부분의 백엔드는 가중치 전용입니다. 이는 가중치를 낮은 정밀도로 저장하고 컴퓨팅 시점에 높은 정밀도로 역양자화(dequantize)한다는 의미입니다. 이 방식은 메모리 사용량을 크게 줄여주지만, 보통 추론 속도를 빠르게 만들지는 못하며 오히려 약간의 지연 시간 오버헤드를 추가할 수도 있습니다.

인기 있는 Nunchaku 추론 엔진의 양자화 방법인 SVDQuant는 다른 접근 방식을 취합니다. 이는 주요 트랜스포머 레이어를 4비트 가중치 및 활성화(W4A4)로 실행하여 메모리를 줄이는 동시에 노이즈 제거 루프를 가속화합니다. 자세한 내용은 아래에 다루지만, 이전까지 이러한 체크포인트를 사용하려면 별도의 추론 라이브러리가 필요했습니다.

현재 Diffusers에서는 kernels 패키지 덕분에 from_pretrained()를 호출하는 것만으로 Nunchaku 체크포인트를 로드할 수 있으며, 로컬 CUDA 컴파일이 필요하지 않습니다. 또한, 동반되는 diffuse-compressor 툴킷을 사용하면 새로운 아키텍처를 직접 양자화하고 일반 Diffusers 리포지토리로 게시할 수 있습니다.

• Nunchaku Lite 시작하기
• 배경: SVDQuant 및 Nunchaku
• Nunchaku Lite 소개
• Diffusers에서의 네이티브 로딩
• 더 높은 속도와 낮은 메모리 확보
• 벤치마크
• 자신만의 모델 양자화하기
• 사용 가능한 체크포인트
• 결론
• 감사의 말씀

먼저, 필수 요구 사항을 설치합니다. 최신 버전의 Diffusers와 Hugging Face kernels 패키지가 필요합니다:

pip install -U diffusers transformers accelerate kernels bitsandbytes

그런 다음 다른 Diffusers 모델과 마찬가지로 사전 양자화된 파이프라인을 로드합니다:

import torch
from diffusers import ErnieImagePipeline
pipe = ErnieImagePipeline.from_pretrained(
...

커스텀 파이프라인 클래스나 별도의 추론 엔진이 필요하지 않으며, 로컬에서 컴파일할 것도 없습니다. NVFP4 커널은 처음 사용될 때 Nunchaku Lite 커널 페이지를 통해 Hub에서 다운로드됩니다. 이 체크포인트는 Nunchaku NVFP4 트랜스포머와 bitsandbytes NF4 텍스트 인코더를 결합하여, BF16 파이프라인의 약 24GB에 비해 피크 메모리 사용량이 약 12GB인 RTX 5090에서 약 1.7초 만에 1024x1024 이미지를 생성합니다. Nunchaku Lite 체크포인트 형식에 대한 더 자세한 내용은 공식 Diffusers 문서를 참조하십시오.

NVFP4 체크포인트는 NVIDIA Blackwell GPU(RTX 50 시리즈, RTX PRO 6000, B200)가 필요합니다. 이전 세대의 경우 INT4 변형을 사용하십시오. 자세한 내용은 아래 하드웨어 지원 표를 참조하십시오.

SVDQuant는 Nunchaku의 기반이 되는 양자화 방식이자 레퍼런스 CUDA 추론 엔진입니다. 표준 4비트 양자화는 가중치(weights)와 활성화 값(activations) 모두에 큰 이상치(outliers)를 포함하고 있기 때문에 확산 트랜스포머(diffusion transformers)에는 적용하기 어렵습니다. SVDQuant는 이러한 문제를 해결하기 위해 활성화 값의 이상치를 가중치로 옮기고, 각 가중치 행렬에서 가장 어려운 부분을 작은 16비트 저랭크 브랜치(low-rank branch)로 표현하며, 나머지 잔여값(residual)을 4비트로 양자화합니다. Nunchaku는 4비트 경로와 저랭크 브랜치를 위한 통합 커널(fused kernels)을 사용하여 이를 빠르게 만듭니다.

원래의 Nunchaku 엔진은 통합된 QKV 투영(QKV projections) 및 통합된 GELU/MLP 커널과 같은 모델별 통합 실행 경로에서 많은 속도를 얻습니다. 이러한 최적화는 각 아키텍처의 모듈 레이아웃과 체크포인트 형식에 연결되어 있기 때문에, 새로운 모델 패밀리를 지원하려면 일반적으로 모델별 통합 작업이 필요합니다.

Nunchaku Lite는 Diffusers 내의 새로운 통합 경로입니다. 이를 통해 Diffusers는 커스텀 파이프라인이나 별도의 추론 엔진 없이 Nunchaku 스타일의 체크포인트를 로드할 수 있습니다. 내부적으로 Nunchaku Lite는 관련 nn.Linear

체크포인트가 로드되기 전에 스톡 Diffusers 모델의 모듈에 런타임 SVDQ/AWQ 선형 레이어를 적용합니다. CUDA 커널은 kernels 패키지를 통해 Hub에서 가져옵니다.

두 가지 커널 계열이 사용됩니다:

: SVDQuant 저랭크 보정(low-rank correction)을 사용한 4비트 가중치 및 활성화 값. 이 레이어는 트랜스포머의 어텐션 및 MLP 투영에서 거의 모든 계산이 이루어지며, INT4 및 NVFP4 변형으로 사용할 수 있습니다. svdq_w4a4

: 적응형 정규화(adaptive normalization) 및 FLUX와 같은 변조 투영에 사용되는 16비트 활성화 값을 가진 4비트 가중치. awq_w4a16

adanorm_single

/adanorm_zero

또는 Qwen-Image 변조 레이어. 이 레이어들은 메모리 바운드(memory-bound)이며 정밀도에 민감하여, AWQ가 메모리와 공간을 절약하면서도 정밀도를 유지하는 데 적합합니다.

단점은 아키텍처별 융합 커널(architecture-specific fused kernels)과 모듈이 없기 때문에 Nunchaku Lite가 원래의 Nunchaku 엔진만큼 속도 향상을 달성할 수는 없다는 것입니다. 하지만, 기본적인 구현만으로도 약 30%의 속도 향상을 제공하며 동일한 수준의 VRAM 감소를 유지합니다.

Diffusers에서 bitsandbytes나 torchao를 사용해 본 적이 있다면 메커니즘이 익숙하게 느껴질 것입니다. Nunchaku Lite 모델 저장소는 일반적인 Diffusers 저장소입니다. 유일하게 특별한 부분은 트랜스포머의 config.json 안에 있는 quantization_config 블록입니다:

"quantization_config": {
"quant_method": "nunchaku_lite",
"compute_dtype": "bfloat16",
...

이 설정은 Diffusers에게 어떤 모듈이 양자화되었는지, 어떤 방식(scheme)을 사용했는지, 그리고 어떤 Nunchaku Lite 런타임 레이어(SVDQW4A4Linear 또는 AWQW4A16Linear)를 인스턴스화해야 하는지 알려줍니다.

양자화된 모델이 원래의 밀집(dense) 모델과 정확한 모듈 구조를 유지하기 때문에, 다운스트림에 있는 모든 것(스케줄러, LoRA 로딩 훅, 오프로딩, torch.compile)은 정상적인 Diffusers 모델로 인식합니다.

Nunchaku Lite는 GPU 세대와 체크포인트 정밀도에 따라 다른 커널 변형을 사용합니다:

Volta 및 Hopper GPU는 현재 4비트 커널에서 지원되지 않습니다. 양자화기(quantizer)는 로드 시간에 GPU의 CUDA 기능을 검증하고 잘못된 출력을 생성하는 대신 명확한 오류를 발생시킵니다.

Nunchaku Lite는 다른 Diffusers 메모리 및 속도 최적화와 결합할 수 있습니다.

torch.compile. 트랜스포머(transformer) 컴파일은 종단 간(end-to-end) 속도를 1.35배에서 1.8배로 향상시킵니다:

pipe.transformer.compile(fullgraph=True)
# 또는 더 빠른 컴파일을 위해 compile_repeated_blocks() 사용
pipe.transformer.compile_repeated_blocks(fullgraph=True)

양자화된 텍스트 인코더. 트랜스포머가 메모리 점유율이 큰 유일한 구성 요소는 아닙니다. T5나 Qwen3와 같은 텍스트 인코더만으로도 수 기가바이트를 차지할 수 있습니다. bitsandbytes NF4로 텍스트 인코더를 추가 양자화하면 벤치마크에서 최대 VRAM을 약 22% 줄일 수 있었습니다.

오프로딩(Offloading). enable_model_cpu_offload() 및 enable_sequential_cpu_offload()와 같은 Diffusers 오프로딩 도우미는 더 작은 GPU에 파이프라인을 맞추어야 할 경우 평소처럼 작동합니다.

아래의 모든 수치는 rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder를 사용하여 NVIDIA RTX PRO 6000 (Blackwell)에서 1024x1024 해상도로 측정되었습니다.

위에서 보았듯이 Nunchaku는 지연 시간(latency)을 약 30% 개선하면서도 최대 50%까지 피크 VRAM을 줄여줍니다. 나머지 오버헤드는 주로 추가적인 커널 실행(kernel launches)에서 발생하며, 이는 torch.compile이 완화할 수 있어 전체 파이프라인을 1.68초로 낮추거나 BF16 기준보다 1.8배 더 빠르게 만들 수 있습니다.

Diffusers에서의 Nunchaku Lite 지원은 아키텍처에 구애받지 않으며, diffuse-compressor 툴킷은 Diffusers 모델을 위한 엔드투엔드 SVDQuant 워크플로우(보정(calibrate), 양자화(quantize), 패키징(package), 게시(publish))를 제공합니다.

아래에서는 예시로 FLUX.2 Klein 4B를 양자화하는 과정을 설명합니다. 이 과정은 모델 검사, 트랜스포머 보정 및 양자화, 결과를 Diffusers 파이프라인으로 패키징, 그리고 Hub에 검증 후 게시하는 주요 단계를 다룹니다. 전체 튜토리얼에서는 모든 플래그를 자세히 다룹니다.

일반 스캐너(generic scanner)는 모델을 순회하며 무엇을 목표로 할지 결정합니다: 반복되는 트랜스포머 블록 스택 내부의 호환 가능한 선형 계층(linears)은 SVDQ W4A4 타겟이 되고, 인식된 변조 선형 계층(modulation linears)은 AWQ W4A16 타겟이 되며, 나머지는 밀집(dense) 상태로 유지됩니다.

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
--precision int4 --rank 32 --inspect-config

양자화하기 전에 항상 이 보고서를 읽으십시오. FLUX.2 Klein 4B의 경우, 예상 결과는 100개의 SVDQ 타겟, 3개의 AWQ 타겟, 그리고 6개의 밀집 외부 선형 계층(dense outer linears)이며, 누락된 패턴이나 중복 이름은 없습니다.

다음 명령어는 트랜스포머에 대해 SVDQuant를 실행하고 양자화된 체크포인트를 outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors에 작성합니다:

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
--precision int4 \
--output outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors

--precision int4를 nvfp4로 대체하여 Blackwell 네이티브 가중치(Blackwell-native weights)를 구축할 수 있습니다.

컨버터는 양자화된 트랜스포머와 기본 파이프라인의 다른 구성 요소들을 결합하고, 간결한 nunchaku_lite 구성을 transformer/config.json에 작성하며, 선택적으로 텍스트 인코더를 NF4로 변환할 수 있습니다.

Once the outputs look good, run pipe.push_to_hub("your-name/your-model-nunchaku-lite-int4")

다른 사용자는 위에서 보여준 것과 같은 from_pretrained() 패턴으로 로드할 수 있습니다.

일반적인 경로는 아키텍처가 구조적 재작성 없이 양자화될 수 있다고 가정합니다. 추가 속도 향상을 위해, 원래의 Nunchaku 엔진은 Diffusers 레이어 그룹을 융합된 모듈(fused modules)로 재작성합니다. 일반적인 경로는 이러한 변경 사항을 자체적으로 추론할 수 없습니다. 예를 들어, 별도의 Q, K, V 투영(projections)을 하나의 모듈로 결합하거나 융합된 투영을 여러 모듈에 걸쳐 분할하는 등의 작업이 불가능합니다.

FLUX.1-dev의 QKV 투영은 구체적인 예시입니다. Diffusers는 세 개의 별도 모듈을 정의합니다:

self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)

Nunchaku FLUX 모듈은 이 레이어들을 하나의 양자화된 to_qkv 모듈로 결합합니다:

to_qkv = fuse_linears([other.to_q, other.to_k, other.to_v])
self.to_qkv = SVDQW4A4Linear.from_linear(to_qkv, **kwargs)

이 그룹화된 모듈은 Nunchaku의 융합 연산자(fused operator)가 QKV 투영, Q/K 정규화(normalization), 그리고 회전 임베딩(rotary embeddings)을 함께 사용하기 때문에 필요합니다. 비교하자면, 기본 Diffusers 경로는 이들을 별도로 실행합니다:

query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
...

Nunchaku 경로는 그룹화된 투영, 정규화 모듈, 그리고 회전 임베딩을 하나의 융합 연산자에 제공합니다:

qkv = fused_qkv_norm_rottary(
hidden_states, attn.to_qkv, attn.norm_q, attn.norm_k, image_rotary_emb
)

이것은 일반적인 경로로는 추론할 수 없는 구조적 재작성입니다. Diffusers는 to_q, to_k, 그리고 to_v 매개변수 접두사를 가진 세 개의 목적지 모듈을 가지고 있는 반면, Nunchaku는 to_qkv 아래에 하나의 그룹화된 모듈을 가지고 있습니다. 모델별 대상 설정(target config) 또는 어댑터는 Q, K, V 파라미터가 출력 차원(output dimension)을 따라 해당 순서로 연결되어야 하며, 이를 to_qkv에 로드해야 한다고 명시해야 합니다.