tdrussell/diffusion-pipe

해결하는 문제

단일 GPU에 담기에는 너무 큰 대규모 확산 모델의 학습을 가능하게 합니다. 파이프라인 병렬 처리를 활용하여 모델의 워크로드를 여러 GPU에 분산시킴으로써, 제한된 하드웨어에서도 하이엔드 모델 학습을 수행할 수 있도록 합니다.

작동 원리

이 프로젝트는 DeepSpeed를 기반으로 한 데이터 및 파이프라인 병렬 처리의 하이브리드 방식을 사용합니다. 모델을 여러 GPU에 걸쳐 단계별로 분할하며, 모델의 여러 인스턴스를 동시에 실행할 수 있습니다. 메모리와 속도를 최적화하기 위해 Huggingface Datasets 라이브러리를 사용하여 잠재 변수(latents)와 텍스트 임베딩을 디스크에 사전 캐싱하므로, 실제 학습 과정 중에 VAE와 텍스트 인코더를 VRAM에 유지할 필요가 없습니다.

대상 사용자

멀티 GPU 설정을 사용하여 대규모 이미지 및 비디오 확산 모델(Flux, SDXL, HunyuanVideo 등)을 학습하거나 미세 조정하려는 개발자 및 AI 연구자.

주요 특징

• 폭넓은 모델 지원: SDXL, Flux, LTX-Video, HunyuanVideo, Cosmos를 포함한 방대한 모델과 호환됩니다.
• 통합 프레임워크: 단일 파이프라인 내에서 이미지 및 비디오 모델을 모두 지원합니다.
• 메모리 효율성: AdamW8BitKahan 옵티마이저, 블록 스와핑, Unsloth 활성화 체크포인팅과 같은 VRAM 절감 기능을 포함합니다.
• 유연한 학습: 양자화된 모델에서 직접 LoRA를 학습하는 것을 지원하며, 중단된 세션을 재개하기 위한 전체 상태 체크포인팅을 제공합니다.