Data parallelism

Data parallelism replicates the same model on each GPU and gives each GPU different prompts. It increases throughput when the model already fits on one GPU, and it leaves the latency of a single prompt unchanged. If you haven’t already, try Batch inference on one GPU first. Refer to Distributed inference to compare it with the other strategies.

You can set up data parallelism with Accelerate or with PyTorch Distributed directly.

Accelerate

Accelerate sets up the processes and assigns a GPU to each one.

Install Accelerate with the following command.

uv pip install accelerate

Create a PartialState in a file called run_distributed.py to set up the distributed environment. Move the DiffusionPipeline to the PartialState.device to assign a GPU to each process.

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)

In the same file, use split_between_processes as a context manager to split the prompts between processes.

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

Run the script with accelerate launch, and set --num_processes to the number of GPUs.

accelerate launch --num_processes=2 run_distributed.py

For a complete example, see this minimal multi-GPU inference script. To learn more, see the Distributed inference with Accelerate guide.

PyTorch Distributed

You can also set up data parallelism with torch.distributed directly, without Accelerate.

Call init_process_group to create the distributed environment. torchrun sets the rank of each process and the total number of processes (the world size) through environment variables, so you don’t pass them yourself.

Each process needs its own GPU and its own prompt. torchrun sets the LOCAL_RANK environment variable to the process’s GPU index on its node, so use it to pick the GPU. Use get_rank, which counts processes across all nodes, to pick the 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()

Save the script as run_distributed.py and run it with torchrun. Set --nproc-per-node to the number of GPUs. The prompts list has one prompt per process, so add a prompt for each GPU you use.

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