xDiT

xDiT es un motor de inferencia diseñado para la implementación paralela a gran escala de Diffusion Transformers (DiTs). xDiT ofrece un conjunto de enfoques paralelos eficientes para Diffusion Models, así como aceleraciones de kernels de GPU.

xDiT admite cuatro métodos paralelos, entre ellos Unified Sequence Parallelism, PipeFusion, el paralelismo CFG y el paralelismo de datos. Los cuatro métodos paralelos de xDiT se pueden configurar de forma híbrida, optimizando los patrones de comunicación para adaptarse mejor al hardware de red subyacente.

La optimización ortogonal a la paralelización se centra en acelerar el rendimiento de una sola GPU. Además de utilizar bibliotecas de optimización de Atención bien conocidas, aprovechamos tecnologías de aceleración por compilación como torch.compile y onediff.

A continuación se muestra un resumen de xDiT.

pip install xfuser

Aquí hay un ejemplo de cómo usar xDiT para acelerar la inferencia de un modelo de 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 puede ver, solo necesitamos usar xFuserArgs de xDiT para obtener los parámetros de configuración y pasar estos parámetros junto con el objeto de pipeline de la biblioteca Diffusers a xDiTParallel para completar la paralelización de un pipeline específico en Diffusers.

Los parámetros de tiempo de ejecución de xDiT se pueden ver en la línea de comandos usando -h, y puedes consultar este ejemplo de uso para más detalles.

xDiT debe iniciarse usando torchrun para soportar sus capacidades de paralelismo multinodo y multigpu. Por ejemplo, el siguiente comando puede usarse para inferencia paralela en 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 compatibles

Un subconjunto de los modelos de Diffusers es compatible con xDiT, como Flux.1, Stable Diffusion 3, etc. Los modelos más recientes compatibles se pueden encontrar aquí.

Benchmark

Probamos diferentes modelos en varias máquinas y aquí hay algunos de los datos de las pruebas de rendimiento.

Flux.1-schnell

Stable Diffusion 3

HunyuanDiT

Se pueden encontrar métricas de rendimiento más detalladas en nuestra página de GitHub.

Referencia

xDiT-project

USP: un enfoque unificado de paralelismo de secuencias para la IA generativa de contexto largo

PipeFusion: Paralelismo de pipeline de parches desplazados para la inferencia de modelos de transformadores de difusión