Hugging Face와 Google Cloud가 JAX + Cloud TPU v5e 조합으로 Stable Diffusion XL 추론을 2.4배 비용 효율로 개선하고 4초 내 1024×1024 이미지 4장 생성

Context

Stable Diffusion XL은 이전 버전 대비 UNet 컴포넌트가 3배 커져서 메모리 요구사항과 추론 시간이 급증했다. 프로덕션 배포 시 높은 계산 비용과 지연 시간이 주요 과제였다.

Technical Solution

• JAX의 JIT(Just-In-Time) 컴파일 도입: 정적 모양(static shapes)이 필요한 이미지 생성 작업의 특성을 활용해 첫 실행 시 최적화된 TPU 바이너리를 생성하고 이후 호출에서 재사용
• XLA 컴파일러 기반 병렬화: JAX의 pmap을 사용해 SPMD(Single-Program Multiple-Data) 프로그램 구현으로 여러 TPU 칩에서 동시 실행 (예: 8칩 TPU에서 1칩이 이미지 1개 생성할 시간에 이미지 8개 생성)
• 모델 파라미터 정밀도 최적화: 32비트 파라미터를 bfloat16으로 변환해 메모리 절감 및 계산 속도 향상 (단, 스케줄러 상태는 float32 유지)
• Cloud TPU v5e 인스턴스 활용: 1칩, 4칩, 8칩, 256칩 규모의 다양한 구성으로 사용 사례에 맞춘 확장성 제공
• Hugging Face Diffusers 라이브러리 통합: FlaxStableDiffusionXLPipeline을 통해 JAX 기반 추론 파이프라인 간소화

Impact

• Cloud TPU v5e는 TPU v4 대비 2.4배 높은 성능/달러 달성 (TPU v5e-4: 배치 크기 4일 때 21.46 vs TPU v4-8: 9.05)
• SDXL 이미지 4장(1024×1024) 생성 시간: 전체 4초 (실제 생성 시간 2.3초, 형식 변환 및 통신 포함)
• 배치 크기 4 기준 지연시간: TPU v5e-4에서 2.33초 (TPU v4-8: 2.16초)
• Cloud TPU v5e는 TPU v4 대비 절반 이하의 비용

Key Takeaway

생성형 AI 모델 배포에서 고정 입출력 모양 특성을 활용한 JIT 컴파일과 전문화된 하드웨어(TPU) + 최적화된 소프트웨어 스택(JAX) 조합은 단순한 성능 향상을 넘어 비용 효율성까지 대폭 개선할 수 있다. 정적 구조의 워크로드 특성 파악과 그에 맞는 기술 스택 선택이 성능과 경제성의 핵심이다.