Buckets:
| # Data parallelism | |
| [Data parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=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](../using-diffusers/batched_inference) on one GPU first. Refer to [Distributed inference](../training/distributed_inference) to compare it with the other strategies. | |
| You can set up data parallelism with [Accelerate](https://huggingface.co/docs/accelerate/index) or with [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) directly. | |
| ## Accelerate | |
| Accelerate sets up the processes and assigns a GPU to each one. | |
| Install Accelerate with the following command. | |
| ```bash | |
| uv pip install accelerate | |
| ``` | |
| Create a [PartialState](https://huggingface.co/docs/accelerate/main/en/package_reference/state#accelerate.PartialState) in a file called `run_distributed.py` to set up the distributed environment. Move the [DiffusionPipeline](/docs/diffusers/pr_14972/en/api/pipelines/overview#diffusers.DiffusionPipeline) to the `PartialState.device` to assign a GPU to each process. | |
| ```py | |
| 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](https://huggingface.co/docs/accelerate/main/en/package_reference/state#accelerate.PartialState.split_between_processes) as a context manager to split the prompts between processes. | |
| ```py | |
| 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. | |
| ```bash | |
| accelerate launch --num_processes=2 run_distributed.py | |
| ``` | |
| > [!TIP] | |
| > For a complete example, see this [minimal multi-GPU inference script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba). To learn more, see the [Distributed inference with Accelerate](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) guide. | |
| ## PyTorch Distributed | |
| You can also set up data parallelism with `torch.distributed` directly, without Accelerate. | |
| Call [init_process_group](https://pytorch.org/docs/stable/distributed.html#torch.distributed.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. | |
| ```py | |
| 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. | |
| ```bash | |
| torchrun --nproc-per-node 2 run_distributed.py | |
| ``` | |
Xet Storage Details
- Size:
- 4.01 kB
- Xet hash:
- 9803f5b47db43f3c23e2384984d6e45714e8a17c0203542e87910a61a7169a79
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.