추론 속도 및 메모리 최적화를 위한 Stable Diffusion XL (SDXL) 최적화

Hugging Face는 Stable Diffusion XL (SDXL)을 특히 소비자용 GPU에서 추론에 더 실용적으로 만들기 위한 일련의 간단한 최적화 방법을 상세히 설명했습니다. 낮은 정밀도(lower precision), 메모리 효율적인 어텐션(memory-efficient attention), 그리고 다양한 오프로딩(offloading) 기술을 구현함으로써 메모리 사용량을 28GB에서 11.47GB로 줄이고, 추론 지연 시간을 72.2초에서 약 10.3초로 단축할 수 있습니다.

추론 속도 최적화

확산 모델(diffusion models)에서 속도 최적화는 매우 중요합니다. 이미지 생성의 반복적인 특성상 원하는 결과를 얻기 위해 여러 번의 실행이 필요할 때가 많기 때문입니다. 다음 기술들은 지연 시간을 줄이는 데 중점을 둡니다:

낮은 정밀도 (fp16)

표준 float32 (fp32) 대신 float16 (fp16)을 사용하면 메모리 사용량을 절반으로 줄이고 계산 속도를 높일 수 있습니다. 이는 fp16이 더 좁은 범위의 부동 소수점을 캡처하며 최신 GPU 하드웨어에서 더 잘 지원되기 때문입니다. 테스트된 SDXL 파이프라인에서 fp16으로 전환하면 메모리 사용량이 21.7GB로 줄어들고 추론 시간이 72.2초에서 14.8초로 단축되었습니다.

메모리 효율적 어텐션 (SDPA)

트랜스포머의 어텐션 블록은 입력 시퀀스 길이에 따라 메모리 요구 사항이 제곱으로 증가하기 때문에 메모리 병목 현상을 일으킬 수 있습니다. PyTorch 2.0은 Flash Attention과 메모리 효율적 어텐션(xFormers)의 융합된 구현을 제공하는 Scaled Dot Product Attention (SDPA)을 도입했습니다. 🤗Diffusers에서 SDPA는 PyTorch ≥ 2.0에 대해 기본적으로 활성화되어 있어, 21.7GB의 메모리 사용량을 유지하면서 추론 시간을 11.4초로 더욱 단축합니다.

torch.compile을 이용한 JIT 컴파일

PyTorch 2.0의 torch.compile API를 사용하면 PyTorch 코드를 최적화된 커널로 적시(just-in-time, JIT) 컴파일할 수 있습니다. SDXL UNet을 torch.compile(mode="reduce-overhead" 사용)으로 감싸면 추론 시간이 10.2초로 더욱 향상됩니다. 첫 번째 컴파일 실행은 느리지만, 이후 호출은 훨씬 빠르다는 점에 유의하십시오.

모델 메모리 사용량(Footprint) 줄이기

SDXL은 이전 Stable Diffusion 모델보다 약 3배 더 크기 때문에(3.5B 파라미터 UNet), 이를 VRAM에 맞추는 것이 주요 과제입니다. 다음과 같은 기술들로 메모리 사용량을 줄일 수 있습니다:

CPU 오프로딩

모델 오프로딩은 GPU에서 활발하게 필요하지 않은 파이프라인 구성 요소를 CPU로 이동시킵니다:

• Model CPU Offloading: 텍스트 인코더와 VAE는 CPU에 유지하면서 UNet을 GPU 메모리에 로드합니다. 이를 통해 메모리 사용량을 20.2GB로 줄였습니다.
• Sequential CPU Offloading: 개별 UNet 서브모듈의 가중치를 CPU로 오프로딩하고, 순전파(forward pass) 직전에만 GPU로 로드합니다. 이는 메모리를 19.9GB로 줄여주지만 지연 시간을 67초로 크게 증가시킵니다.

VAE Slicing

Variational Autoencoder (VAE)는 잠재 변수(latents)를 이미지로 디코딩하며, 이 과정은 배치 크기에 따라 메모리 사용량이 증가합니다. VAE slicing은 입력 텐서를 더 작은 슬라이스로 나누어 여러 단계에 걸쳐 디코딩합니다. 이 최적화는 지연 시간에 큰 영향을 주지 않으면서 메모리 사용량을 15.4GB로 줄여줍니다.

계산 결과 캐싱

SDXL은 프롬프트에서 임베딩을 계산하기 위해 두 개의 텍스트 인코더를 사용합니다. 이러한 임베딩은 역확산(reverse diffusion) 과정 내내 일정하게 유지되므로, 미리 계산하여 캐싱할 수 있습니다. 임베딩이 생성되면 텍스트 인코더와 토크나이저를 GPU 메모리에서 제거할 수 있으며, SDPA 및 fp16과 결합하면 메모리 사용량이 21.9GB가 됩니다.

Tiny Autoencoder (TAESD)

표준 VAE를 madebyollin의 Tiny Autoencoder(약 10MB)와 같은 증류된(distilled) 버전으로 교체하면 메모리 사용량이 15.6GB로 줄어들고 추론 지연 시간이 감소합니다. 하지만 Tiny Autoencoder는 세밀한 디테일을 생략할 수 있으므로 주로 이미지 미리보기용으로 권장됩니다.

성능 요약

프롬프트당 4개의 이미지를 생성하는 A100 GPU(40 GB)에서 수행된 테스트 결과, 메모리와 지연 시간 사이의 트레이드오프는 다음과 같습니다:

기술	메모리 (GB)	추론 지연 시간 (ms)
최적화되지 않은 파이프라인	28.09	72200.5
fp16	21.72	14800.9
fp16 + SDPA (기본값)	21.72	11413.0
기본값 + torch.compile	21.73	10296.7
기본값 + model CPU offload	20.21	16082.2
기본값 + sequential CPU offload	19.91	67034.0
기본값 + VAE slicing	15.40	11232.2
기본값 + VAE slicing + sequential CPU offload	11.47	66869.2
기본값 + 사전 계산된 텍스트 임베딩	21.85	11909.0
기본값 + Tiny Autoencoder	15.48	10449.7