xDiT

xDiT é um motor de inferência projetado para a implantação paralela em larga escala de Diffusion Transformers (DiTs). O xDiT oferece um conjunto de abordagens paralelas eficientes para Diffusion Models, além de acelerações de kernels de GPU.

Existem quatro métodos paralelos suportados no xDiT, incluindo Unified Sequence Parallelism, PipeFusion, paralelismo CFG e paralelismo de dados. Os quatro métodos paralelos no xDiT podem ser configurados de forma híbrida, otimizando os padrões de comunicação para melhor se adequar ao hardware de rede subjacente.

A otimização ortogonal à paralelização foca em acelerar o desempenho de uma única GPU. Além de utilizar bibliotecas bem conhecidas de otimização de Atenção, aproveitamos tecnologias de aceleração por compilação, como torch.compile e onediff.

A visão geral do xDiT é mostrada a seguir.

pip install xfuser

Aqui está um exemplo de uso do xDiT para acelerar a inferência de um modelo do Diffusers.

 import torch
 from diffusers import StableDiffusion3Pipeline

 from xfuser import xFuserArgs, xDiTParallel
 from xfuser.config import FlexibleArgumentParser
 from xfuser.core.distributed import get_world_group

 def main():
+    parser = FlexibleArgumentParser(description="xFuser Arguments")
+    args = xFuserArgs.add_cli_args(parser).parse_args()
+    engine_args = xFuserArgs.from_cli_args(args)
+    engine_config, input_config = engine_args.create_config()

     local_rank = get_world_group().local_rank
     pipe = StableDiffusion3Pipeline.from_pretrained(
         pretrained_model_name_or_path=engine_config.model_config.model,
         dtype=torch.float16,
     ).to(f"cuda:{local_rank}")
    
# do anything you want with pipeline here

+    pipe = xDiTParallel(pipe, engine_config, input_config)

     pipe(
         height=input_config.height,
         width=input_config.height,
         prompt=input_config.prompt,
         num_inference_steps=input_config.num_inference_steps,
         output_type=input_config.output_type,
         generator=torch.Generator(device="cuda").manual_seed(input_config.seed),  # or "mps", "xpu", "cpu"
     )

+    if input_config.output_type == "pil":
+        pipe.save("results", "stable_diffusion_3")

if __name__ == "__main__":
    main()

Como você pode ver, precisamos apenas usar xFuserArgs do xDiT para obter os parâmetros de configuração e passar esses parâmetros, junto com o objeto de pipeline da biblioteca Diffusers, para xDiTParallel para completar a paralelização de um pipeline específico no Diffusers.

Os parâmetros de runtime do xDiT podem ser visualizados na linha de comando usando -h, e você pode consultar este exemplo de uso para mais detalhes.

O xDiT precisa ser iniciado usando o torchrun para suportar suas capacidades de paralelismo multi-nó e multi-GPU. Por exemplo, o seguinte comando pode ser usado para inferência paralela em 8 GPUs:

torchrun --nproc_per_node=8 ./inference.py --model models/FLUX.1-dev --data_parallel_degree 2 --ulysses_degree 2 --ring_degree 2 --prompt "A snowy mountain" "A small dog" --num_inference_steps 50

Modelos suportados

Um subconjunto dos modelos do Diffusers é suportado no xDiT, como Flux.1, Stable Diffusion 3, etc. Os modelos mais recentes suportados podem ser encontrados aqui.

Benchmark

Testamos diferentes modelos em várias máquinas, e aqui estão alguns dos dados de benchmark.

Flux.1-schnell

Stable Diffusion 3

HunyuanDiT

Métricas de desempenho mais detalhadas podem ser encontradas em nossa página no GitHub.

Referência

xDiT-project

USP: Uma abordagem unificada de paralelismo de sequência para IA generativa de contexto longo

PipeFusion: Paralelismo de Pipeline de Patch Deslocado para Inferência de Modelos de Transformer de Difusão