xDiT

xDiT는 DiT(Diffusion Transformer) 계열 모델을 여러 대의 GPU·여러 대의 컴퓨터에 걸쳐 대규모 병렬로 배포해 추론할 수 있게 해주는 확장 가능한(scalable) 추론 엔진이에요. 이미지·비디오 생성의 최신 흐름을 이끌고 있는 DiT에서는 입력 컨텍스트 길이가 늘어날수록 Attention 연산의 계산 비용이 이차함수적으로(quadratically) 증가하기 때문에, 온라인 서비스의 실시간(real-time) 요구를 맞추려면 멀티 GPU·멀티 머신 배포가 필수예요. xDiT는 바로 이런 병렬 추론을 위해 설계된 엔진으로, DiT에 특화된 효율적인 병렬화 기법들과 단일 GPU 성능 최적화 기능을 함께 제공해요.

huggingface의 diffusers에 구현된 DiT 모델들을 간단한 래퍼(wrapper)로 감싸 하이브리드 병렬 구현에 적응시킬 수 있는 API를 제공하며, Flux, Stable Diffusion 3(SD3), SANA, PixArt, HunyuanVideo, Wan2.x, CogVideoX 등 많은 모델을 지원해요.

출처: 문서

본문

핵심 기능

xDiT는 DiT를 위한 네 가지 병렬화 방식을 제공하며, 이를 하이브리드로 조합해 네트워크 하드웨어 특성에 맞게 통신 패턴을 최적화할 수 있어요.

Sequence Parallelism (USP) — 우리가 제안한 통합 시퀀스 병렬 방식이에요. DeepSpeed-Ulysses와 Ring-Attention을 결합했어요. (논문)
PipeFusion — 시퀀스 레벨 파이프라인 병렬화로, TeraPipe와 유사하지만 디퓨전 모델의 입력 시간적 중복(temporal redundancy) 특성을 활용해요. (논문, NeurIPS 2025 채택)
Data Parallel — 여러 프롬프트를 처리하거나 단일 프롬프트에서 여러 이미지를 이미지 단위로 병렬 생성해요.
CFG Parallel (Split Batch) — classifier-free guidance(CFG) 사용 시 상수 병렬도 2로 활성화돼요.

이 밖에도 참고용으로 Tensor Parallelism과 DistriFusion 병렬 전략도 구현되어 있어요. VAE 모듈에는 Out-of-Memory(OOM) 문제를 막기 위한 병렬 구현 DistVAE를 제공해요.

Cache Acceleration — TeaCache, First-Block-Cache, DiTFastAttn 등 디퓨전 모델의 서로 다른 스텝 간 계산 중복을 활용해 단일 GPU 추론을 가속하는 캐시 기법을 지원해요.

Computing Acceleration — 병렬화와 독립적으로 단일 GPU에서 성능을 높이는 최적화예요. torch.compile, onediff 같은 컴파일 가속 기술과 다양한 Attention 최적화 라이브러리를 사용해요.

설치

기본 설치와 flash attention 포함 설치는 다음과 같아요.

pip install xfuser  # Basic installation
pip install "xfuser[flash-attn]"  # With flash attention

소스에서 설치할 때는 다음 명령을 사용해요.

pip install -e .
# Or optionally, with flash attention
pip install -e ".[flash-attn]"

flash_attn은 xDiT와 함께 설치할 수 있는 선택적 라이브러리이며, 실행 시 설치되어 사용 가능한 Attention 백엔드를 자동으로 감지해 가장 빠른 것을 선택해요. GPU 아키텍처에 따라 cuDNN, FAv3, FAv4, Transformer Engine FP8, AITER 등을 추천해요. AMD GPU(MI300X 이상)에서는 AITER를 권장해요.

사용 예시

Model runner를 사용하면 대부분의 지원되는 디퓨전 모델을 벤치마킹·프로파일링 지원과 함께 실행할 수 있어요.

xdit --model FLUX.1-dev \
    --prompt "A cat running in a garden" \
    --ulysses_degree 8

예제 스크립트로 실행할 수도 있어요. ./examples/ 디렉토리에서 모델 타입, 모델 디렉토리, 병렬 옵션을 쉽게 수정할 수 있어요.

bash examples/run.sh

하이브리드 병렬을 구성할 때 중요한 것은 모든 병렬 degree의 곱이 디바이스 수와 일치해야 한다는 점이에요. CFG, PipeFusion, 시퀀스 병렬을 조합해 귀여운 강아지 이미지를 생성하는 예시는 다음과 같아요. 여기서 ulysses_degree * pipefusion_parallel_degree * cfg_degree(use_cfg_parallel) == 8 (디바이스 수)예요.

torchrun --nproc_per_node=8 \
examples/pixartalpha_example.py \
--model models/PixArt-XL-2-1024-MS \
--pipefusion_parallel_degree 2 \
--ulysses_degree 2 \
--num_inference_steps 20 \
--warmup_steps 0 \
--prompt "A cute dog" \
--use_cfg_parallel

⚠️ PipeFusion(그리고 DistriFusion)을 적용할 때는 warmup_steps를 설정해야 해요. 보통 num_inference_steps보다 작은 값으로 두는데, warmup 스텝은 병렬 실행이 불가능해 직렬 실행으로 느려지므로 효율에 영향을 줘요. PixArt 모델에서는 warmup 0이 효과가 없는 것을 관찰했어요.

HTTP 서비스를 띄워 이미지를 생성하는 것도 가능해요. (Text-to-Image HTTP 서비스)

xDiT는 ComfyUI 기반으로 멀티 GPU 병렬 워크플로우를 구현한 상용 프로젝트 TACO-DiT에서도 사용되어 Flux.1의 성능 문제를 해결했어요. xDiT용 ComfyUI 플러그인도 제공돼요. (xdit-comfyui-private)

지원 모델

다음은 xDiT가 공식 지원하는 주요 DiT 모델 목록이에요 (추천 모델 및 문서/숫자 표기 그대로):

비디오: StepVideo, HunyuanVideo (1.5), ConsisID-Preview, CogVideoX (1.5), Mochi-1, Latte, Wan2.1 / Wan2.2, Wan2.2-Distilled, LTX-2 / LTX-2.5, MiniMax-H3
이미지: HunyuanDiT-v1.2, Z-Image Turbo, Flux / Flux 2 / Flux 2 klein / Flux Kontext, Qwen Image / Image-Edit, Krea2-Raw / Turbo, Ideogram 4, PixArt-alpha / Sigma, Stable Diffusion 3, SANA / SANA Sprint, SDXL

병렬 메서드 지원 여부(CFG/SP/PipeFusion/TP/MR)는 모델마다 다르며, 자세한 내용은 원문 및 성능 리포트에서 확인할 수 있어요.

참고 (Limitations)

모델마다 요구하는 diffusers 버전이 다를 수 있어요. 최신 모델일수록 더 새로운 버전을 요구하거나 main 브랜치에서 설치해야 할 수 있어요. 예: Flux ≥ 0.35.2, Flux 2 ≥ 0.36.0, Ideogram 4 ≥ 0.39.0.
검증된 diffusers 버전에 없는 모델은 여러 버전을 시도해보아야 해요.
Cache 기법(TeaCache, First-Block-Cache)은 현재 USP를 사용하는 FLUX 모델에서만 지원되며, PipeFusion에는 적용할 수 없어요.

더 알아보기 (Learn more)

GitHub 저장소