tdrussell/diffusion-pipe

diffusion-pipe – 확산 모델용 파이프라인 병렬 학습 프레임워크

개요

DeepSpeed의 파이프라인 병렬 처리를 사용하여 여러 GPU에서 대규모 확산 모델(이미지 변환, 텍스트-이미지, 텍스트-비디오 등)을 학습시키기 위한 오픈소스 학습 스크립트입니다.
최근 확산 모델 제품군(SDXL, Flux, Stable Diffusion 3, Qwen-Image, Cosmos, Hunyuan, MiniMax H3 등)의 긴 목록을 지원하며 이미지 및 비디오 데이터 모두에서 작동합니다.

주요 기능

기능	중요성
파이프라인 병렬 처리	단일 모델을 여러 GPU에 분할하여 한 카드에 맞지 않는 모델도 학습할 수 있습니다.
하이브리드 데이터 + 파이프라인 병렬 처리	나머지 GPU를 데이터 병렬 처리에 사용하여 배치 크기를 자동으로 확장합니다.
TensorBoard 로깅 및 평가 지표	학습 중 손실, VRAM 사용량, 검증 점수를 모니터링할 수 있습니다.
체크포인트 및 재개	전체 학습 상태(가중치, 옵티마이저, 데이터로더)를 저장하고, 설정을 다시 실행하지 않고도 모든 체크포인트에서 다시 시작할 수 있습니다.
잠재 변수 및 텍스트 임베딩 사전 캐싱	무거운 VAE 및 텍스트 인코더 작업을 전처리 단계로 이동하여 메인 모델용 VRAM을 확보합니다.
통합 이미지/비디오 파이프라인	동일한 코드 경로가 정지 이미지와 비디오 프레임을 처리하며, 비디오 모델의 경우 선택적으로 오디오 처리도 지원합니다.
확장 가능한 모델 인터페이스	새 확산 모델을 추가하려면 작은 하위 클래스 구현만 필요합니다.
낮은 VRAM 트릭	예제 구성은 8비트 AdamW, 블록 스와핑, 활성화 체크포인트, CUDA 메모리 할당기 트릭을 사용하여 단일 24GB GPU에서 140억 파라미터 모델을 학습시키는 방법을 보여줍니다.

시작하기

리포지토리 클론(하위 모듈 포함)

git clone --recurse-submodules https://github.com/tdrussell/diffusion-pipe

Conda 환경 생성 (Python 3.12 필요)

conda create -n diffusion-pipe python=3.12
conda activate diffusion-pipe

GPU/CUDA 버전에 맞는 PyTorch 설치 (예: RTX 4090의 경우 torch==2.9.0+cu128).

나머지 Python 종속성 설치

pip install -r requirements.txt
# 많은 모델에서 권장 (선택 사항)
pip install flash-attn

Cosmos 비디오가 필요한 경우 README에 설명된 환경 변수 해킹을 사용하여 transformer_engine을 설치합니다.

학습 작업 실행

examples/의 예제 TOML 구성 파일(기본: examples/main_example.toml)을 편집합니다. 데이터 세트 경로, 출력 디렉터리, 학습률 등을 조정합니다.
DeepSpeed로 실행합니다. 예:
```
NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1 \
deepspeed --num_gpus=4 train.py --deepspeed --config examples/hunyuan_video.toml
```
*--num_gpus*는 DeepSpeed가 사용할 GPU 수를 지정합니다. 구성의 pipeline_stages는 각 모델이 분할되는 파이프라인 파티션 수를 제어합니다.
재개하려면 --resume_from_checkpoint <폴더>를 추가하고 정확히 동일한 구성 파일을 재사용합니다.

데이터 준비

미디어 파일(이미지 또는 비디오)을 하나 이상의 디렉터리에 구성합니다.
각 미디어 파일에 대해 동일한 기본 이름의 .txt 캡션 파일을 만듭니다.
이미지는 Pillow 호환 형식, 비디오는 ImageIO 호환 형식입니다 (WebP 비디오는 지원되지 않음).
스크립트는 각 데이터 세트 디렉터리 내의 cache/ 하위 폴더에 VAE 잠재 변수와 CLIP 스타일 텍스트 임베딩을 캐시합니다.

출력 레이아웃

각 실행에 대해 output_dir 아래에 새 하위 폴더가 생성됩니다.
내부에는 다음이 포함됩니다:
- DeepSpeed 체크포인트 폴더(global_stepXXXX).
- safetensors 가중치와 PEFT 어댑터 JSON이 포함된 모델 스냅샷(epoch1/, epoch2/ 등).
- 재현성을 위한 원래 diffusion-pipe TOML 구성.
- 손실/검증 곡선용 TensorBoard 이벤트 파일.

플랫폼 참고 사항

DeepSpeed의 Windows 지원이 제한적이므로 Windows 네이티브 학습은 작동하지 않을 가능성이 높습니다. Windows 호스트를 사용해야 하는 경우 WSL 2를 사용하세요.

관련 프로젝트

작성자는 대규모 언어 모델용 유사한 파이프라인 병렬 트레이너인 qlora-pipe도 유지 관리합니다.

결론: diffusion-pipe는 단일 GPU 메모리 예산을 초과하는 최첨단 확산 모델을 학습해야 하는 연구자와 엔지니어에게 실용적이고 문서화가 잘 된 툴킷입니다. 파이프라인 병렬 처리, 캐싱, 체크포인트 관리의 무거운 작업을 처리하므로 데이터와 하이퍼파라미터에 집중할 수 있습니다.