값비싼 Diffusion model을 받드는 저비용 MLOps

디퓨전 모델을 '디퓨전 프로세스를 쓰는 멀티모달 이미지 생성'으로 정의한다. 이미지 생성의 역사는 사실적인 얼굴을 만든 StyleGAN에서 임의 프롬프트로 이미지를 만드는 오픈AI의 DALL-E로 이어진다. 핵심은 싱글모달과 멀티모달의 구분이다. 사진 한 종류만 학습하는 GAN은 싱글모달이지만, 디퓨전 모델은 사진과 그 사진을 설명하는 텍스트가 짝(pair)을 이루는 학습 데이터를 필요로 한다.

디퓨전 프로세스 자체를 보면, GAN이 랜덤 숫자를 한 번에 이미지로 바꾸는 것과 달리 디퓨전은 노이즈 낀 이미지를 반복적으로 걷어내며(디노이징) 또렷하게 만든다. 학습 시 1000번, 생성 시 25~50번 스텝을 반복하기에 GAN보다 느리다. 그래서 단점이 분명하다. 생성이 느려 GPU 비용이 크고, 이미지-텍스트 페어 데이터가 GAN보다 훨씬 많이 필요하다. 오픈AI가 멀티모달의 가능성을 보이자 4억 장의 이미지를 모아 학습하기 시작했는데, 이는 4년 전 데이터보다 5100배 이상 많은 양으로 멋진 결과 뒤에 막대한 데이터·연산·비용이 든다.

선배 연구자들이 이 모델을 어떻게 발전시켜 왔는지 짚는다. 2022년 초 뮌헨대학교와 Stability AI, Runway 등 오픈소스 기여자들이 모여 모델을 학습시켰는데, 그 규모는 A100 GPU 기준 약 15만 시간(AWS 환산 약 67억 원)에 달했다. 그렇게 학습한 모델을 소스 코드와 웨이트까지 공개한 것이 Stable Diffusion이고, 구글의 DreamBooth 같은 후속 연구로 이어지며 생태계가 빠르게 커졌다.

Stable Diffusion은 그 어떤 오픈소스 프로젝트보다 빠르게 스타를 모았다. 당시 V100에서 50스텝이 약 10초 걸렸는데, FlashAttention이 어텐션의 다수 IO 과정을 없애 속도를 약 2배 높여 2~3초로 줄였다. 또 프롬프트에 단어를 더해 부분 수정을 가능하게 한 연구, 이를 학습 데이터 제작에 활용해 사진을 대화하듯 수정하는 InstructPix2Pix, 텍스트가 아닌 참조 이미지로 생성을 제어하는 ControlNet까지 가속의 흐름을 정리한다.

본격적인 최적화 실무다. 먼저 GPU는 CUDA 코어, 빠르지만 작은 SRAM, 크지만 느린 HBM으로 구성되며 데이터의 잦은 왕복이 병목인데, FlashAttention은 이 IO 왕복을 줄인 것이다. 발표자 팀의 첫 최적화는 지식 증류(knowledge distillation)다. 모델 크기를 줄이는 보통의 증류 대신 구글·스탠퍼드는 반복 스텝 수를 줄이는 '스텝 다이어트'를 택했다. 35번 반복을 3번으로 줄이면 10배 빠른 모델이 되어 10초 추론을 1초로 만든다. 팀은 6번 증류에 약 6개월을 들여 거의 같은 퀄리티를 10배 빠르게 내는 모델을 얻어, 퀄리티 중시 유저에게는 프로·빠른 경험 중시 유저에게는 무료로 제공하는 전략을 쓸 수 있게 됐다.

둘째는 TensorRT를 통한 연산 최적화다. 코사인 계산처럼 IO 왕복이 잦은 연산자들을 하나로 합치는(operator fusion) 라이브러리로, 모델을 ONNX로 번역한 뒤 컴파일한다. 자동으로 못 바꾸는 비효율 코드는 오픈AI의 Triton 컴파일러로 복잡한 CUDA 프로그래밍을 간편하게 작성해 레이턴시를 다시 절반가량 줄였다. 마지막 서빙은 베어메탈이 아닌 PC 기반 서버 클러스터로, 두 서비스가 하나의 모델을 공유하도록 구성하고 엔비디아의 Triton Inference Server를 썼다. 앙상블 모델 스케줄러가 그래프 의존성을 계산해 의존성 없는 모델을 동시에 실행함으로써 같은 시간에 훨씬 많은 추론을 수행한다. 정리하면 디스틸레이션과 TensorRT로 레이턴시를 크게 줄여 디퓨전 모델을 저비용으로 효율적으로 서빙한 시행착오를 공유하며 마친다.