Parallélisme des données

La parallélisation des données réplique le même modèle sur chaque GPU et attribue à chaque GPU des prompts différents. Elle augmente le débit lorsque le modèle tient déjà sur un seul GPU, et elle laisse inchangée la latence d’un prompt unique. Si vous ne l’avez pas encore fait, essayez d’abord l’inférence par lots sur un seul GPU. Reportez-vous à l’inférence distribuée pour la comparer aux autres stratégies.

Vous pouvez configurer le parallélisme de données avec Accelerate ou directement avec PyTorch Distributed.

Accelerate

Accelerate configure les processus et attribue un GPU à chacun d'entre eux.

Installez Accelerate avec la commande suivante.

uv pip install accelerate

Créez un PartialState dans un fichier nommé run_distributed.py pour configurer l’environnement distribué. Déplacez le DiffusionPipeline sur le PartialState.device pour attribuer une GPU à chaque processus.

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

Dans le même fichier, utilisez split_between_processes comme gestionnaire de contexte pour répartir les invites entre les processus.

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

Exécutez le script avec accelerate launch et définissez --num_processes sur le nombre de GPU.

accelerate launch --num_processes=2 run_distributed.py

Pour un exemple complet, consultez ce script minimal d’inférence multi-GPU. Pour en savoir plus, consultez le guide Inférence distribuée avec Accelerate.

PyTorch Distributed

Vous pouvez également configurer le parallélisme de données directement avec torch.distributed, sans Accelerate.

Appel init_process_group pour créer l’environnement distribué. torchrun définit le rang de chaque processus et le nombre total de processus (la taille du monde) via des variables d’environnement, vous n’avez donc pas besoin de les passer vous-même.

Chaque processus a besoin de son propre GPU et de son propre prompt. torchrun définit la variable d’environnement LOCAL_RANK sur l’indice du GPU du processus sur son nœud, il faut donc l’utiliser pour sélectionner le GPU. Utilisez get_rank, qui compte les processus sur tous les nœuds, pour sélectionner le prompt.

import os
import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")

    pipeline = DiffusionPipeline.from_pretrained(
        "Qwen/Qwen-Image", dtype=torch.bfloat16,
    ).to(device)

    prompts = ["a dog", "a cat"]
    prompt = prompts[rank]

    image = pipeline(prompt).images[0]
    image.save(f"{prompt.replace(' ', '_')}.png")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Enregistrez le script sous le nom run_distributed.py et exécutez-le avec torchrun. Définissez --nproc-per-node sur le nombre de GPU. La liste prompts contient un prompt par processus, ajoutez donc un prompt pour chaque GPU utilisé.

torchrun --nproc-per-node 2 run_distributed.py