Paralelismo de datos

Paralelismo de datos replica el mismo modelo en cada GPU y asigna a cada GPU distintos prompts. Aumenta el rendimiento cuando el modelo ya cabe en una GPU y deja invariada la latencia de un prompt individual. Si aún no lo has hecho, prueba primero Inferencia por lotes en una GPU. Consulta Inferencia distribuida para compararlo con las demás estrategias.

Puede configurar el paralelismo de datos con Accelerate o directamente con PyTorch Distributed.

Accelerate

Accelerate configura los procesos y asigna una GPU a cada uno.

Instale Accelerate con el siguiente comando.

uv pip install accelerate

Cree un PartialState en un archivo llamado run_distributed.py para configurar el entorno distribuido. Mueva el DiffusionPipeline al PartialState.device para asignar una GPU a cada proceso.

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)

En el mismo archivo, use split_between_processes como un administrador de contexto para dividir los prompts entre procesos.

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")

Ejecuta el script con accelerate launch y establece --num_processes en el número de GPUs.

accelerate launch --num_processes=2 run_distributed.py

Para un ejemplo completo, consulte este script de inferencia multi-GPU minimalista. Para obtener más información, consulte la guía Inferencia distribuida con Accelerate.

PyTorch Distributed

También puede configurar el paralelismo de datos directamente con torch.distributed, sin Accelerate.

Llama a init_process_group para crear el entorno distribuido. torchrun establece el rango de cada proceso y el número total de procesos (el tamaño del mundo) a través de variables de entorno, por lo que no es necesario que los pases tú mismo.

Cada proceso necesita su propia GPU y su propio prompt. torchrun establece la variable de entorno LOCAL_RANK con el índice de la GPU del proceso en su nodo, por lo que se debe usar para seleccionar la GPU. Utilice get_rank, que cuenta los procesos en todos los nodos, para seleccionar el 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()

Guarda el script como run_distributed.py y ejecútalo con torchrun. Establece --nproc-per-node al número de GPUs. La lista prompts tiene un prompt por proceso, así que añade un prompt para cada GPU que uses.

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