Bringing Nunchaku 4-bit Diffusion Inference to Diffusers

🖼️ 인포그래픽

Bringing Nunchaku 4-bit Diffusion Inference to Diffusers 내용을 설명하는 본문 이미지

🖼️ 4컷 인포그래픽

💡 한 줄 요약

Nunchaku Lite는 SVDQuant 기반 4비트 가중치·활성값 연산을 Diffusers의 일반 모델 로딩 방식에 통합해, 별도 추론 엔진이나 로컬 CUDA 컴파일 없이 VRAM 사용량과 추론 지연을 함께 줄인다.

📌 핵심 요약

최신 텍스트-이미지 모델은 BF16 기준으로 보통 20~30GB의 VRAM을 요구하며, 기존 가중치 전용 양자화는 메모리는 줄여도 계산 시 역양자화 때문에 추론 속도를 높이지 못할 수 있다.
SVDQuant는 활성값의 이상치를 가중치 쪽으로 이동하고, 가중치 행렬의 어려운 부분을 16비트 저랭크 분기로 분리한 뒤 나머지를 4비트로 양자화해 W4A4 연산을 구현한다.
Nunchaku Lite 체크포인트는 일반 Diffusers 저장소와 동일하게 from_pretrained()로 불러오며, 필요한 CUDA 커널은 kernels 패키지를 통해 허브에서 자동으로 내려받는다.
Nunchaku Lite는 주요 어텐션·MLP 투영에 SVDQ W4A4를, 정밀도에 민감한 정규화·변조 투영에는 AWQ W4A16을 적용하며, 원본 Nunchaku의 모델별 융합 경로 없이도 약 30%의 속도 향상과 VRAM 절감을 제공한다.
RTX PRO 6000 벤치마크에서 BF16 파이프라인은 3.00초와 31.1GB를 기록한 반면, Nunchaku Lite에 torch.compile을 적용한 구성은 1.68초와 20.6GB를 기록했으며 NF4 텍스트 인코더 구성은 최대 VRAM을 16.0GB까지 낮췄다.

🧩 주요 포인트

가중치뿐 아니라 활성값도 4비트로 계산하므로, 단순 저장 공간 축소를 넘어 디노이징 구간의 실제 연산 시간까지 줄인다.
기존 Diffusers 모델 구조를 그대로 유지하므로 스케줄러, LoRA, 오프로딩, torch.compile 같은 후속 기능과의 호환성을 보존하면서 도입 장벽을 낮춘다.
모델별 융합 최적화를 포기한 범용 통합 경로이지만, 컴파일과 텍스트 인코더 양자화를 조합해 실행 속도와 메모리 병목을 각각 추가로 완화할 수 있다.

🧠 상세 정리

1. 기존 확산 모델과 가중치 전용 양자화의 한계

대형 확산 트랜스포머는 이미지뿐 아니라 영상·오디오·텍스트 생성에도 활용되지만, 최신 텍스트-이미지 모델을 BF16으로 불러오려면 보통 20~30GB의 VRAM이 필요해 다수의 소비자용 GPU에서 실행하기 어렵다. Diffusers는 이미 bitsandbytes, GGUF, torchao, Quanto 등 여러 양자화 백엔드를 제공하지만, 대다수는 가중치만 저정밀도로 저장했다가 계산할 때 고정밀도로 복원하는 가중치 전용 방식이다. 이 방식은 저장 및 로딩 메모리를 크게 줄여도 실제 계산은 높은 정밀도로 수행하기 때문에 추론 속도를 높이지 못하는 경우가 많으며, 역양자화 과정에서 작은 지연이 추가될 수도 있다. 글은 이러한 한계를 보완하기 위해 가중치와 활성값을 함께 4비트로 처리하는 Nunchaku 계열 기술을 Diffusers에 직접 통합하는 방법을 설명한다.

2. Nunchaku Lite의 설치와 기본 사용 방식

사용자는 최신 Diffusers와 transformers, accelerate, kernels, bitsandbytes를 설치한 뒤 일반 모델과 동일하게 파이프라인의 from_pretrained()를 호출해 사전 양자화된 체크포인트를 불러올 수 있다. 별도 파이프라인 클래스나 독립적인 Nunchaku 추론 엔진은 필요하지 않으며, NVFP4 커널은 처음 사용할 때 kernels 패키지를 통해 허브의 Nunchaku Lite 커널 페이지에서 내려받으므로 로컬 CUDA 컴파일도 요구되지 않는다. 제시된 ERNIE-Image-Turbo 예시는 Nunchaku NVFP4 트랜스포머와 bitsandbytes NF4 텍스트 인코더를 조합하고, 8단계로 1024×1024 이미지를 생성한다. 글에 따르면 RTX 5090에서 이 예시는 약 1.7초의 생성 시간과 약 12GB의 최대 메모리 사용량을 기록했으며, 비교 대상인 BF16 파이프라인은 약 24GB를 사용한다. 다만 NVFP4 체크포인트는 RTX 50 시리즈, RTX PRO 6000, B200 등 NVIDIA Blackwell GPU가 필요하며, 지원되는 이전 세대 GPU에서는 INT4 변형을 사용해야 한다.

3. SVDQuant가 4비트 이상치를 처리하는 원리

표준 4비트 양자화는 확산 트랜스포머의 가중치와 활성값에 큰 이상치가 모두 존재하기 때문에 적용하기 어렵다. Nunchaku의 기반 기술인 SVDQuant는 활성값의 이상치를 가중치 쪽으로 이동한 다음, 각 가중치 행렬에서 처리하기 가장 어려운 부분을 작은 16비트 저랭크 분기로 표현하고 나머지 잔차를 4비트로 양자화한다. 이를 통해 주요 트랜스포머 계층에서 4비트 가중치와 4비트 활성값을 사용하는 W4A4 연산을 수행하면서도 저랭크 보정으로 필요한 표현력을 유지한다. Nunchaku 참조 엔진은 저랭크 하향 투영을 양자화 커널과 융합하고, 저랭크 상향 투영을 4비트 계산 커널과 융합한다. 이 융합은 별도의 16비트 분기를 실행할 때 발생할 수 있는 추가 메모리 접근 비용을 제거해 SVDQuant의 계산 경로를 빠르게 만든다.

4. 원본 Nunchaku와 범용 Nunchaku Lite의 차이

원본 Nunchaku 엔진은 QKV 투영이나 GELU·MLP 계산을 하나로 묶는 모델별 융합 실행 경로에서 상당한 속도 이득을 얻는다. 그러나 이러한 최적화는 각 아키텍처의 모듈 배치와 체크포인트 형식에 결합되어 있어, 새로운 모델 계열을 지원할 때마다 전용 통합 작업이 필요하다. Nunchaku Lite는 이와 달리 기본 Diffusers 모델의 체크포인트를 읽기 전에 관련 nn.Linear 모듈을 런타임 SVDQ 또는 AWQ 선형 계층으로 교체하는 범용 경로를 사용한다. 주요 계산이 집중되는 어텐션과 MLP 투영에는 저랭크 보정을 포함한 svdq_w4a4 계층을 적용하고, 적응형 정규화나 변조 투영처럼 메모리 대역폭과 정밀도에 민감한 부분에는 awq_w4a16 계층을 적용한다. 아키텍처별 융합 모듈을 사용하지 않기 때문에 원본 Nunchaku의 최대 가속에는 미치지 못하지만, 글에 따르면 기본 구현만으로도 같은 수준의 VRAM 절감을 유지하면서 약 30%의 속도 향상을 제공한다.

5. Diffusers 저장소 형식과 기존 기능 호환성

Nunchaku Lite 모델 저장소는 일반적인 Diffusers 저장소이며, 특별한 부분은 트랜스포머의 config.json에 포함되는 quantization_config 블록이다. 이 설정에는 양자화 방식, 계산 자료형, SVDQ와 AWQ의 정밀도, 그룹 크기, 저랭크 순위, 양자화 대상 모듈 목록이 기록된다. Diffusers는 이 정보를 이용해 어떤 모듈이 어떤 방식으로 양자화됐는지 판단하고, 해당 위치에 SVDQW4A4Linear 또는 AWQW4A16Linear 런타임 계층을 생성한다. 양자화된 모델은 고밀도 원본과 정확히 같은 모듈 구조를 유지하므로, 파이프라인의 나머지 구성 요소는 이를 일반 Diffusers 모델처럼 취급할 수 있다. 그 결과 스케줄러, LoRA 로딩 훅, 모델 오프로딩, torch.compile 같은 기존 기능을 별도의 전용 파이프라인 없이 그대로 사용할 수 있다.

6. GPU 지원 범위와 추가 최적화 조합

Nunchaku Lite는 체크포인트 정밀도와 GPU 세대에 따라 서로 다른 커널을 사용하며, NVFP4 기반 svdq_w4a4는 RTX 50 시리즈, RTX PRO 6000, B200을 포함한 Blackwell GPU를 대상으로 한다. INT4 기반 svdq_w4a4와 awq_w4a16은 Turing, Ampere, Ada 계열에 지원 대상으로 제시되며, 예시 장치에는 RTX 30·40 시리즈, A100, L40S가 포함된다. Volta와 Hopper GPU는 현재 해당 4비트 커널을 지원하지 않으며, 양자화기는 모델을 불러올 때 CUDA 기능을 검사해 잘못된 출력을 내는 대신 명확한 오류를 발생시킨다. 트랜스포머에 torch.compile을 적용하면 제시된 벤치마크의 종단 간 가속률이 1.35배에서 1.8배로 높아지고, 반복 블록만 컴파일하는 기능을 사용하면 컴파일 준비 시간을 줄일 수 있다. 또한 T5나 Qwen3 같은 대형 텍스트 인코더를 bitsandbytes NF4로 양자화하면 최대 VRAM을 약 22% 더 줄일 수 있으며, 더 작은 GPU에는 기존 CPU 또는 순차 오프로딩 기능을 함께 사용할 수 있다.

7. 지연 시간과 최대 VRAM 벤치마크

벤치마크의 모든 수치는 NVIDIA RTX PRO 6000에서 1024×1024 해상도로 측정됐다. BF16 기준 파이프라인은 전체 실행 3.00초, 디노이징 2.86초, 최대 VRAM 31.1GB를 기록했으며 이를 1.0배 기준으로 삼았다. Nunchaku Lite 구성은 전체 실행 2.27초, 디노이징 2.13초, 최대 VRAM 20.6GB로 1.35배의 속도를 보였고, torch.compile을 추가하면 각각 1.68초와 1.53초로 줄어들면서 최대 VRAM은 20.6GB로 유지되고 가속률은 1.8배가 됐다. NF4 텍스트 인코더를 결합한 구성은 전체 실행 2.29초와 디노이징 2.13초를 유지하면서 최대 VRAM을 16.0GB까지 낮췄다. 글은 이를 최대 약 50%의 VRAM 절감과 약 30%의 지연 개선으로 요약하며, 남은 실행 비용의 상당 부분은 추가 커널 호출에서 발생하고 torch.compile이 이를 완화한다고 설명한다. 이미지 품질 부분에는 같은 시드와 설정을 사용한 BF16 및 4비트 출력 비교가 제시되지만, 본문에는 별도의 정량 품질 지표가 제시되지 않는다.

8. 새 모델의 양자화·패키징·배포 절차

Nunchaku Lite의 Diffusers 지원은 특정 아키텍처에 종속되지 않으며, diffuse-compressor 도구는 모델 보정, 양자화, 패키징, 게시까지 이어지는 SVDQuant 작업 흐름을 제공한다. FLUX.2 Klein 4B 예제에서는 먼저 일반 스캐너로 반복 트랜스포머 블록의 호환 선형 계층, 인식된 변조 계층, 고밀도로 남을 외부 선형 계층을 검사하며, 예상 결과는 SVDQ 대상 100개, AWQ 대상 3개, 고밀도 외부 선형 계층 6개이고 누락 패턴이나 중복 이름이 없어야 한다. 다음 단계에서는 INT4 또는 Blackwell용 NVFP4 정밀도로 트랜스포머를 양자화하고, 결과를 세이프텐서 체크포인트로 저장한다. 변환기는 이 체크포인트를 기본 파이프라인의 다른 구성 요소와 결합하고 transformer/config.json에 Nunchaku Lite 설정을 기록하며, 선택적으로 텍스트 인코더를 NF4로 변환한다. 완성된 로컬 Diffusers 파이프라인은 일반 DiffusionPipeline.from_pretrained()로 불러와 고정 시드 이미지 생성으로 확인할 수 있다. 출력이 적절한지 검증한 뒤 허브에 게시하면 다른 사용자도 동일한 from_pretrained() 방식으로 체크포인트를 사용할 수 있다.

🧾 핵심 주장 / 시사점

Nunchaku Lite의 핵심 차이는 저정밀 가중치를 저장하는 데 그치지 않고 주요 연산의 활성값까지 4비트로 처리해, 메모리와 디노이징 속도를 동시에 개선한다는 점이다.
범용 모듈 교체 방식은 원본 Nunchaku의 모델별 융합 가속 일부를 포기하는 대신, Diffusers의 표준 저장소 구조와 기존 생태계 기능을 유지하며 새로운 아키텍처의 통합 부담을 줄인다.
벤치마크는 트랜스포머 양자화와 torch.compile이 주로 실행 시간을 줄이고, 텍스트 인코더 NF4 양자화가 추가 메모리를 줄이는 상호 보완적 최적화임을 보여준다.

✅ 액션 아이템

BF16 대비 Nunchaku Lite의 from_pretrained 로딩 경로에서 VRAM과 디노이징 지연이 함께 줄는지 동일 조건으로 비교한다.
어텐션·MLP의 SVDQ W4A4와 정규화·변조 투영의 AWQ W4A16 혼합이 화질 민감 구간에 미치는 영향을 점검한다.
torch.compile과 NF4 텍스트 인코더를 조합할 때 1.68초·20.6GB와 16.0GB 구간 중 우선 목표를 정한다.