NNCF와 🤗 Optimum을 사용한 Intel CPU용 Stable Diffusion 최적화

Overview

Hugging Face는 Neural Network Compression Framework (NNCF) 양자화 인식 훈련과 Token Merging을 결합하여 Intel CPU용 Stable Diffusion을 최적화하는 방법을 제시했으며, 이로 인해 추론 속도가 최대 5.1x 빨라지고 모델 크기가 원래 PyTorch 체크포인트의 0.25x로 감소했습니다.

Stable Diffusion 최적화

Stable Diffusion 파이프라인의 UNet 구성 요소는 가장 계산 비용이 높은 부분이므로 이를 최적화하면 큰 속도 향상을 얻을 수 있습니다. 전통적인 사후 훈련 8비트 양자화는 이 모델에 잘 작동하지 않는데, 이는 픽셀 수준 예측 작업이 매개변수 변경에 매우 민감하고, 수백만 개의 샘플로 훈련되어 모델에 중복성이 거의 없기 때문입니다. 정확도를 유지하려면 Quantization‑Aware Training (QAT)와 같은 보다 정교한 양자화 방법이 필요합니다.

최적화 워크플로

우리는 Pokemon 데이터셋(svjack/Stable-Diffusion-Pokemon-en)에서 파인튜닝된 Stable Diffusion 모델에서 시작했습니다. Diffusers 텍스트‑to‑이미지 파인튜닝 예제를 사용하여 NNCF 기반 QAT를 훈련 스크립트에 통합하고, 원본 모델이 교사 역할을 하는 지식 증류 손실을 추가했으며, 훈련 안정성을 위해 모델 파라미터(양자화기 제외)에 지수 이동 평균(EMA)을 적용했습니다. 그래디언트 체크포인팅과 EMA 모델을 RAM에 유지함으로써 전체 최적화를 24 GB VRAM을 가진 단일 GPU에서 하루 미만으로 4096번의 반복을 수행할 수 있었습니다.

Quantization‑Aware Training을 넘어

양자화만으로도 모델 풋프린트, 로드 시간, 메모리 소비 및 지연이 감소합니다. 우리는 8비트 양자화와 Token Merging (ToME) 방법을 결합했는데, 이는 self‑attention 블록 이전에 중복 토큰을 병합하여 계산을 줄입니다. 결합된 워크플로는 위에서 설명한 지식 증류, EMA, 그래디언트 체크포인팅을 포함합니다. 다시 Pokemon‑fine‑tuned 모델에서 시작하여, 양자화 위에 병합 비율 0.4의 ToME를 적용했습니다. 결과 모델은 클라이언트 또는 에지 CPU에서의 추론을 목적으로 합니다.

결과

PyTorch 기준을 OpenVINO FP32로 변환하면 1.9x 속도 향상이 있었습니다. 8비트 양자화를 추가하면 PyTorch 대비 3.9x 속도 향상이 되고 모델 풋프린트가 원래 체크포인트의 0.25x로 감소했습니다. 양자화 위에 Token Merging을 쌓으면 풋프린트를 동일한 0.25x 수준으로 유지하면서 5.1x 추론 속도 향상을 달성했습니다. 모든 측정은 3세대 Intel® Xeon® Scalable 프로세서와 Intel® Deep Learning Boost 기술을 갖춘 Hugging Face Spaces CPU 업그레이드 인스턴스에서 OpenVINO 2022.3을 사용하여 수행되었으며, 기본 50 추론 단계를 사용했습니다. 블로그에서는 단계 수가 적을수록 속도는 빨라지지만 이미지 품질에 영향을 줄 수 있으며, 단계 수와 스케줄러를 실험해 볼 것을 권장합니다.

from optimum.intel import OVStableDiffusionPipeline

# Load and compile the pipeline for performance.
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Generate an image.
prompt = "a drawing of a green pokemon with red eyes"
output = pipe(prompt, num_inference_steps=50, output_type="pil\)).images[0]
output.save("image.png

The training and quantization code are available in the Optimum Intel repository, a demonstration notebook is provided, and optimized models can be found on the Hugging Face Hub under the OpenVINO organization. A live demo runs on Hugging Face Spaces.

범용 Stable Diffusion 모델은 어떻게 되나요?

Pokemon 모델에 대한 시연 워크플로는 modeste한 훈련 자원으로도 상당한 최적화가 가능함을 보여줍니다. 범용 Stable Diffusion 모델을 처음부터 훈련하는 것은 비용이 많이 들지만, 충분한 예산과 하드웨어가 있다면 동일한 접근 방식을 적용할 수 있습니다. 주의할 점은 Token Merging이 모델 용량을 줄인다는 점인데, 따라서 더 복잡한 데이터셋에서는 최적화 중에 낮은 병합 비율을 사용하는 것이 좋습니다.

For further reading on complementary approaches for 4th‑generation Intel Xeon CPUs, see the related Hugging Face blog post on Stable Diffusion inference with Intel.