diffusers/v0.43.0/pt/optimization/data_parallelism.md

Paralelismo de dados

Paralelismo de dados replica o mesmo modelo em cada GPU e fornece prompts diferentes para cada GPU. Ele aumenta o throughput quando o modelo já cabe em uma única GPU e mantém a latência de um prompt único inalterada. Se ainda não fez, experimente Inferência em lote em uma única GPU primeiro. Consulte Inferência distribuída para compará-lo com as outras estratégias.

Você pode configurar o paralelismo de dados com Accelerate ou diretamente com PyTorch Distributed.

Accelerate

O Accelerate configura os processos e atribui uma GPU a cada um deles.

Instale o Accelerate com o seguinte comando.

uv pip install accelerate

Crie um PartialState em um arquivo chamado run_distributed.py para configurar o ambiente distribuído. Mova o DiffusionPipeline para o PartialState.device para atribuir uma GPU a cada processo.

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)

No mesmo arquivo, use split_between_processes como um gerenciador de contexto para dividir os prompts entre os processos.

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

Execute o script com accelerate launch e defina --num_processes para o número de GPUs.

accelerate launch --num_processes=2 run_distributed.py

Para um exemplo completo, veja este script de inferência multi-GPU minimalista. Para saber mais, consulte o guia Inferência distribuída com Accelerate.

Distribuído do PyTorch

Você também pode configurar o paralelismo de dados diretamente com torch.distributed, sem o Accelerate.

Chame init_process_group para criar o ambiente distribuído. torchrun define o rank de cada processo e o número total de processos (o world size) por meio de variáveis de ambiente, então você não precisa passá-los manualmente.

Cada processo precisa de sua própria GPU e de seu próprio prompt. torchrun define a variável de ambiente LOCAL_RANK com o índice da GPU do processo em seu nó, portanto, use-a para escolher a GPU. Use get_rank, que conta os processos em todos os nós, para escolher o 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()

Salve o script como run_distributed.py e execute-o com torchrun. Defina --nproc-per-node para o número de GPUs. A lista prompts tem um prompt por processo, então adicione um prompt para cada GPU que você usar.

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