xDiT

xDiT est un moteur d’inférence conçu pour le déploiement parallèle à grande échelle des Diffusion Transformers (DiTs). xDiT propose un ensemble d’approches parallèles efficaces pour les Diffusion Models, ainsi que des accélérations de noyaux GPU.

xDiT prend en charge quatre méthodes de parallélisme, notamment Unified Sequence Parallelism, PipeFusion, le parallélisme CFG et le parallélisme de données. Les quatre méthodes de parallélisme de xDiT peuvent être configurées de manière hybride, afin d’optimiser les motifs de communication de sorte à s’adapter au mieux au matériel réseau sous-jacent.

L’optimisation orthogonale à la parallélisation se concentre sur l’accélération des performances sur un seul GPU. En plus de l’utilisation de bibliothèques d’optimisation de l’Attention bien connues, nous exploitons des technologies d’accélération par compilation telles que torch.compile et onediff.

L'aperçu de xDiT est présenté ci-dessous.

pip install xfuser

Voici un exemple d’utilisation de xDiT pour accélérer l’inférence d’un modèle 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()

Comme vous pouvez le constater, il suffit d’utiliser xFuserArgs de xDiT pour obtenir les paramètres de configuration, puis de transmettre ces paramètres ainsi que l’objet pipeline de la bibliothèque Diffusers à xDiTParallel afin de finaliser la parallélisation d’un pipeline spécifique dans Diffusers.

Les paramètres d’exécution de xDiT peuvent être consultés dans la ligne de commande à l’aide de -h, et vous pouvez consulter cet exemple d’utilisation pour plus de détails.

xDiT doit être lancé à l'aide de torchrun pour prendre en charge ses capacités de parallélisme multi-nœuds et multi-GPU. Par exemple, la commande suivante peut être utilisée pour une inférence parallèle sur 8 GPU :

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

Modèles pris en charge

Un sous-ensemble des modèles Diffusers est pris en charge par xDiT, notamment Flux.1, Stable Diffusion 3, etc. Les modèles les plus récents pris en charge sont disponibles ici.

Benchmark

Nous avons testé différents modèles sur diverses machines, et voici une partie des données de référence.

Flux.1-schnell

Stable Diffusion 3

HunyuanDiT

Des métriques de performance plus détaillées sont disponibles sur notre page GitHub.

Référence

xDiT-project

USP : une approche unifiée du parallélisme de séquences pour l'IA générative à long contexte

PipeFusion : parallélisme de pipeline de patchs déplacés pour l’inférence des modèles de diffusion Transformer