frompretrained()로 구현하는 VRAM 50% 절감과 1.7초 추론

RTX 5090에서 1024px 이미지를 생성할 때 BF16(Bfloat16) 파이프라인이 사용하는 24GB의 피크 VRAM은 Nunchaku Lite를 통해 12GB로 절반가량 줄어든다. 이를 통해 고해상도 이미지 생성 시 발생하는 메모리 부족 문제를 해결할 수 있다.

구현 과정에서 로컬 CUDA 컴파일 과정이 제거되었다. 별도의 추론 엔진 설치나 커스텀 파이프라인 클래스 정의 없이 4비트 가중치-활성화(W4A4) 추론을 적용할 수 있으며, 필요한 NVFP4(엔비디아 4비트 부동소수점) 커널은 kernels 패키지를 통해 최초 사용 시 자동으로 다운로드된다.

환경 구축을 위해 diffusers와 huggingface_hub 패키지를 설치한다.

pip install diffusers huggingface_hub

이후 `from_pretrained()` 메서드를 호출하면 지정된 체크포인트의 모델을 즉시 로드하여 추론에 사용할 수 있다.

pipe = DiffusionPipeline.from_pretrained("rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder")

실제 성능 측정 결과 RTX 5090에서 1024x1024 이미지 한 장을 생성하는 데 걸리는 시간은 약 1.7초다. 이는 BF16 파이프라인 대비 피크 VRAM 사용량을 50% 감소시키면서 추론 지연 시간은 약 30% 개선한 수치다.

SVDQuant의 이상치 처리와 Nunchaku Lite의 런타임 패칭

SVDQuant는 확산 트랜스포머의 가중치와 활성화 값에 존재하는 대규모 이상치를 분리해 처리하여 4비트 양자화의 정확도 손실을 해결한다. 활성화 값의 이상치를 가중치 쪽으로 이동시킨 뒤, 가중치 행렬의 핵심 부분을 16비트 저차원 분기로 따로 떼어내고 나머지 잔차 부분만 4비트로 양자화한다.

Nunchaku Lite는 Diffusers 환경에서 이러한 체크포인트를 바로 사용할 수 있게 하는 통합 경로를 제공한다. 모델 체크포인트를 로드하기 전, 기존 Diffusers 모델의 `nn.Linear` 모듈을 SVDQ 또는 AWQ 전용 선형 레이어로 교체하는 런타임 패칭을 수행한다. 이를 통해 사용자는 표준 모델 구조를 유지하면서 양자화된 가중치를 처리하는 전용 커널을 사용할 수 있다.

구체적인 양자화 설정은 `config.json` 파일 내의 `quantization_config` 블록을 통해 제어한다. 이 블록에는 양자화가 적용된 모듈의 경로, 양자화 스킴, 런타임 레이어 클래스가 명시된다. SVDQuant의 4비트 가중치와 4비트 활성화를 사용하는 경우 `SVDQW4A4Linear` 클래스를, AWQ의 4비트 가중치와 16비트 활성화를 사용하는 경우 `AWQW4A16Linear` 클래스를 지정한다.

"quantization_config": {

"modules": ["transformer.single_transformer_block.0.attn.qkv"],

"scheme": "svdq",

"layer_class": "SVDQW4A4Linear"

}

`transformer.single_transformer_block.0.attn.qkv`와 같이 구체적인 모듈 경로를 지정함으로써 모델 전체가 아닌 필요한 레이어에만 선택적으로 양자화 기법을 적용할 수 있다.

기존 방식과 달라진 지점

Nunchaku Lite는 아키텍처 전용 퓨즈드 커널을 제외하는 대신 범용성을 확보했으며, 오리지널 엔진보다 속도는 느리지만 기본적으로 약 30%의 속도 향상을 제공한다. 퓨즈드 실행 경로(fused execution path)는 QKV 프로젝션이나 GELU/MLP 커널처럼 서로 다른 연산 단계를 하나로 묶어 메모리 입출력 횟수를 줄이는 방식이다. 오리지널 Nunchaku는 이 방식으로 성능을 극대화했으나, 모델마다 개별적인 통합 작업이 필수적이었다.

Nunchaku Lite의 제네릭 경로는 모델의 구조적 재작성(structural rewrite)을 수행하지 않는다. 예를 들어 FLUX.1-dev 모델의 경우 개별 분리된 Q, K, V 프로젝션 모듈을 하나의 `to_qkv` 모듈로 통합해야 연산 경로를 단축할 수 있는데, Lite의 제네릭 경로는 이러한 변경을 실행하지 못한다.

이로 인해 오리지널 Nunchaku의 퓨즈드 연산자가 QKV 프로젝션, Q/K 정규화, 회전 위치 임베딩(rotary embeddings)을 한 번에 처리하는 것과 달리, Nunchaku Lite는 이를 각각 별개의 모듈에서 순차적으로 실행한다.

대신 Nunchaku Lite는 제네릭 스캐너(generic scanner)를 통해 대상 레이어를 결정한다. 반복되는 트랜스포머 블록 내부의 호환 가능한 선형 레이어는 SVDQ W4A4 대상으로, 인식된 변조 선형 레이어는 AWQ W4A16 대상으로 지정하고 나머지는 밀집 상태로 유지하여 VRAM 절감 수준을 유지하며 속도를 개선했다.

torch.compile과 NF4 텍스트 인코더를 통한 추가 최적화

torch.compile 적용 시 엔드투엔드 속도 향상 폭이 1.35배에서 1.8배로 증가한다. Nunchaku Lite는 양자화 모델이 기존 밀집 모델의 모듈 구조를 그대로 유지하므로 파이토치의 그래프 컴파일러인 torch.compile을 즉시 적용할 수 있다.

pipe.transformer = torch.compile(pipe.transformer)

텍스트 인코더의 양자화는 추가적인 메모리 효율을 제공한다. T5나 Qwen3 같은 대형 텍스트 인코더를 bitsandbytes 라이브러리의 NF4(NormalFloat 4)로 양자화하면 피크 VRAM 사용량을 약 22% 더 줄일 수 있다.

또한 `enable_model_cpu_offload`나 `enable_sequential_cpu_offload`와 같은 기존 Diffusers 오프로딩 도구와도 완전히 호환된다. Nunchaku Lite는 런타임에 레이어를 패칭하지만 하위 구조는 표준 Diffusers 모델과 동일하기 때문이다.

최종 최적화 조합을 적용해 NVIDIA RTX PRO 6000 Blackwell 환경에서 rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder 모델로 1024x1024 해상도 이미지를 생성했을 때 전체 파이프라인 속도는 1.68초를 기록했다.

GPU 세대별 정밀도 선택과 diffuse-compressor 워크플로우

NVIDIA Blackwell(RTX 50 시리즈, B200) GPU는 NVFP4 정밀도를 사용할 수 있으며, 이전 세대 하드웨어는 INT4 변형 모델을 사용해야 한다. Volta 및 Hopper GPU는 4비트 커널을 지원하지 않아 모델 로드 시 에러가 발생하며, 양자화 도구가 로드 시점에 CUDA Capability를 검사하여 이를 방지한다.

`diffuse-compressor` 툴킷은 Diffusers 모델을 위한 SVDQuant 전체 워크플로우를 제공한다. 사용자는 모델 구조 검사(inspect)와 데이터 분포 교정(calibrate)을 거쳐 양자화, 패키징, 발행 단계로 진행한다. FLUX.2 Klein 4B 모델의 트랜스포머에 SVDQuant를 적용하는 명령은 다음과 같다.

python -m diffuse_compressor.quantize --precision int4 --model flux-2-klein-4b

Blackwell 전용 가중치를 생성하려면 `--precision` 인자를 `nvfp4`로 변경한다. 양자화 완료 후 safetensors 파일을 배포 가능한 형태로 변환하는 과정은 아래 명령어로 수행한다.

python -m diffuse_compressor.convert --model_path outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors

이 툴킷을 통해 사용자는 특정 하드웨어에 최적화된 아키텍처를 직접 양자화하고 Hugging Face 저장소에 배포할 수 있다.

보유한 GPU가 Blackwell 세대라면 NVFP4를, 그 외 세대라면 INT4를 선택하고 torch.compile을 활성화하여 성능 최적화 범위를 확보하는 것이 실행 기준이다.

Nunchaku Lite는 Diffusers 생태계 내에서 복잡한 설정 없이도 고성능 양자화 추론을 가능하게 함으로써, 고해상도 이미지 생성의 진입 장벽을 낮추고 하드웨어 효율성을 극대화하는 실질적인 경로를 제시한다.