Hugging Face Diffusers가 Flax 지원을 추가해 Stable Diffusion을 Google TPU에서 8개 병렬 장치를 활용한 추론으로 단일 칩 속도로 8배 이미지 생성 가능

Context

Stable Diffusion 추론은 계산량이 많아 GPU 환경에서도 시간이 소요된다. Google TPU 서버는 8개의 TPU 가속기를 병렬로 구동할 수 있으나, 기존 Diffusers 라이브러리는 이를 활용하지 못했다.

Technical Solution

• Flax를 사용한 상태비저장 모델 구조 도입: 모델 파라미터를 모델 외부에 저장하고 별도 관리
• JAX의 pmap(parallel map) 함수로 파이프라인의 _generate 메서드를 8개 TPU 디바이스에 병렬화: 각 디바이스가 다른 입력 배치를 처리하도록 자동 샤딩
• 프롬프트 복제 및 입력 샤딩: 단일 프롬프트를 8배로 복제한 후 flax.jax_utils.replicate와 shard를 사용해 각 디바이스에 분산
• bfloat16 반정밀도 연산 활용: 메모리 효율성을 위해 bf16 가중치 버전 로드
• 비동기 디스패치 처리: JAX의 비동기 특성을 활용하되 block_until_ready()로 정확한 측정 시간 확보

Impact

단일 칩이 1개 이미지를 생성하는 시간에 8개 디바이스가 동시에 8개 이미지 생성 (총 실행 시간 6.82초로 측정됨).

Key Takeaway

Functional 프로그래밍 패러다임(Flax)과 분산 컴파일 추상화(pmap)를 조합하면, 기저 함수를 수정 없이 병렬 실행 코드로 변환 가능하므로, 다중 가속기 환경에서 개발자의 복잡도 증가 없이 확장성을 확보할 수 있다.