Is it possible to inference using multiple GPUs?

[DevJunghun]

Hi, Thanks for sharing this library for using stable diffusion.
There is one questions I want to ask.

Like title, Is it possible to inference using multiple GPUs? If possible, how?
do you share doc about inference using multiple GPUs?

OS: Linux Ubuntu 20.04
GPU: RTX 4090 (24GB) * n
RAM: 72GB
Python: 3.9.16

assume i have two stable diffusion models (model 1, model 2)
ex) GPU 1 - using model 1, GPU 2 - using model 2

or

assume i have two request, i want to process both request parallel (prompt 1, prompt 2)
ex) GPU 1 - processing prompt 1, GPU 2 - processing prompt 2

I think.. this question can be solved by using thread and two pipes like below.. right?

p_01 = StableDiffusionPipeline.from_pretrained(model_01).to("cuda:0")  
p_02 = StableDiffusionPipeline.from_pretrained(model_02).to("cuda:1")  

Thread(target=generate_pipe01, args=(prompt, negative_prompt)).start()  
Thread(target=generate_pipe02, args=(prompt, negative_prompt)).start()

I'll be waiting for your good opinions. Thank you.

[patrickvonplaten]

Hey @DevJunghun,

here what I would recommend.

1.) Create a python file run_distributed.py that works in distributed mode. Note that we set the world_size here to 2 assuming that you want to run your code in parallel over 2 GPUs.

#!/usr/bin/env python3
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from diffusers import DiffusionPipeline

sd = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)


def run_inference(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # move to rank
    sd.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")


def main():
    world_size = 2
    mp.spawn(
        run_inference,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )


if __name__ == "__main__":
    main()

2. Having defined the script you can start it by just running:

torchrun run_distributed.py

[patrickvonplaten]

Note that when using PyTorch's distributed data loaders you have much more control over what data goes to which GPU:
https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html

[patrickvonplaten]

@sayakpaul @williamberman @pcuenca - it might be worth actually creating a quick doc page for this

[DevJunghun]

@patrickvonplaten Thanks for your kindness! Have a nice day :)

[sayakpaul]

Reopening this issue to keep a better track of the doc @patrickvonplaten mentioned in #2977 (comment).

@muellerzr do you have any recommendations for this? Anything, in particular, we need to know for the accelerate side to run distributed inference? Any relevant pointers would be very useful for us to ensure the doc we're putting together sheds light into the best practices :)

[muellerzr]

@sayakpaul using accelerate launch removes any CLI specifics + spawning that Patrick showed, and you can use the PartialState for anything else @patrickvonplaten showed (such as the new PartialState().process_index, which is better for this stuff) to specify what GPU something should be run on. And in regards to .to(rank) you can use state.device

You can use AcceleratorState or Accelerator as well, but PartialState was designed for this more utility-focused approach

So in full as code:

accelerate launch file.py

#!/usr/bin/env python3
from accelerate import PartialState
from diffusers import DiffusionPipeline

sd = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)


def main():
    # Initialize the distributed environment
    state = PartialState()

    # move to rank
    sd.to(state.device)

    if state.process_index == 0:
        prompt = "a dog"
    elif state.process_index == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")

if __name__ == "__main__":
    main()

[chikiuso]

is there any way I could inference with multi gpu on one single image with text prompt? thanks

[sayakpaul]

@chikiuso could you describe your use case?

[chikiuso]

Hi @sayakpaul , I have 4 rtx 3090 gpu installed on ubuntu server, I would like to inference a text prompt to image as fast as possible (not each gpu process one prompt), to use 4 gpu to process one single image at a time, is it possible? thanks.

[sayakpaul]

Still not clear to me.

Are you trying to generate four images for a given (single in this case) prompt?

[chikiuso]

Hi @sayakpaul Sorry for my bad English, I am trying to generate one single image with one single prompt at the same time. thanks.