diffusion-pipe: 대규모 이미지 및 비디오 확산 모델을 위한 파이프라인 병렬 훈련 프레임워크

해결하는 문제

파이프라인 병렬성을 활용하여 단일 GPU에 들어가지 못하는巨大한 확산 모델의 훈련을 가능하게 합니다. 이미지와 비디오 생성 모델 모두를 위한 통합 프레임워크를 제공하여 서로 다른 아키텍처를 관리하는 복잡성을 줄입니다.

작동 방식

이 프로젝트는 DeepSpeed를 통해 데이터 및 파이프라인 병렬성의 하이브리드를 사용하여 모델을 여러 GPU에 분산합니다. VRAM 사용을 최적화하기 위해 디스크에 잠재 표현(latents)과 텍스트 임베딩을 사전 캐시하여 실제 훈련 과정에서 VAE와 텍스트 인코더를 언로드할 수 있도록 합니다. 또한 8-bit 최적화기, 블록 스왑, 활성화 체크포인팅과 같은 메모리 절약 기술을 지원합니다.

대상 사용자

Flux, SDXL, HunyuanVideo와 같은 대규모 이미지 및 비디오 확산 모델을 미세 조정하거나 훈련해야 하는 AI 연구원 및 개발자를 대상으로 하며, 단일巨大한 GPU의 VRAM이 부족한 경우에 적합합니다.

주요 기능

  • 광범위한 모델 지원: SDXL, Flux, LTX-Video, HunyuanVideo, Cosmos, Wan2.1과 같은 다양한 최신 모델과 호환됩니다.
  • 파이프라인 병렬성: 단일 GPU 메모리를 초과하는 모델을 처리하도록 특별히 설계되었습니다.
  • 통합 이미지/비디오 워크플로: 단일 훈련 스크립트에서 두 가지 모달리티를 모두 지원합니다.
  • 효율적인 캐싱: 멀티프로세스, 멀티GPU 사전 캐싱을 통해 임베딩과 잠재 표현을 미리 로드하여 훈련 속도를 높이고 메모리를 절약합니다.
  • 훈련 관리: Tensorboard 로깅, 홀드아웃 세트에서의 평가, 그리고 훈련 재개를 위한 상태 체크포인팅을 포함합니다.

Sources