attention.py only supports one GPU

[Marcophono2]

Model/Pipeline/Scheduler description

Hello!

In attention.py the lines

                _ = xformers.ops.memory_efficient_attention(
                    torch.randn((1, 2, 40), device="cuda"),
                    torch.randn((1, 2, 40), device="cuda"),
                    torch.randn((1, 2, 40), device="cuda"),
                )

use always gpu 1 ("cuda:0"). That is fine if there is only one gpu. But if you have > 1 gpu all gpus use the vram of gpu 1. At the end the performance of gpu 1 is decreasing with every additional gpu parking its load in the gpu 1 vram.

I found a solution for my three gpus in form of three different conda environments where I replaced device="cuda" with device="cuda:0", device="cuda:1" and device="cuda:2" in the three different attention.py files in the three environments. Surely if would have only advantages if one could overgive the desired cuda in form of

pipe1.enable_xformers_memory_efficient_attention("cuda:1")

i.E. if using xFormers.
Or is there already a simple way to overwrite the cuda setting?

Best regards
Marc

Open source status

• The model implementation is available
• The model weights are available (Only relevant if addition is not a scheduler).

Provide useful links for the implementation

No response

[averad]

Quick Link:
attention.py

Code Referenced (class BasicTransformerBlock(nn.Module) line 313:

    def _set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool):
        if not is_xformers_available():
            print("Here is how to install it")
            raise ModuleNotFoundError(
                "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
                " xformers",
                name="xformers",
            )
        elif not torch.cuda.is_available():
            raise ValueError(
                "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only"
                " available for GPU "
            )
        else:
            try:
                # Make sure we can run the memory efficient attention
                _ = xformers.ops.memory_efficient_attention(
                    torch.randn((1, 2, 40), device="cuda"),
                    torch.randn((1, 2, 40), device="cuda"),
                    torch.randn((1, 2, 40), device="cuda"),
                )
            except Exception as e:
                raise e
            self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
            self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers

[camenduru]

hi @Marcophono2 I am using from threading import Thread for each gpu pipe maybe helps

idea:

pipe0 = StableDiffusionPipeline.from_pretrained(model_folder).to("cuda:0")
pipe1 = StableDiffusionPipeline.from_pretrained(model_folder).to("cuda:1")

Thread(target=generate0_with_pipe0, args=(width, height, prompt, negative_prompt)).start()
Thread(target=generate1_with_pipe1, args=(width, height, prompt, negative_prompt)).start()

[Marcophono2]

Hi @camenduru ! I also use threading in the same way as you described. But in this case the same attention.py file from the SD folder is used. And there you can only set one cuda.
By the way: Meanwhile I am on a Ubuntu with GUI. When I wrote the code I rent regualary a computer instance with 2 or 4 RTX 3090. But that were Ubuntu server setups and I really had problems at the beginning to not simply let run different instanzes in different comand boxes. That was the only reason to switch to threading. Okay, it also works but I could imagine to go back again to the command boxes. There is no really good reason for it but old habits but I cannot exclude that there is a significant higher cpu load. When I rent that instanzes also with 8x 3090 and let them run all near 99%, the cpu load was about 1%. In my case I am at 300% with three GUIs. Strange.

P.S.: 300% for the sum of the loads of every core. I have a 5900x, so 12 physical and 12 virtual cores. Enough resources left.

[patrickvonplaten]

Interesting, just to better understand - what exactly is the use case?
To use stable diffusion on multiple GPUs in inference? Or Training? Also are we talking about data parallelism (split batch over multiple GPUs) or model parallelism (split model layers over multiple GPUs)

[Marcophono2]

@patrickvonplaten: Inference only. So I cannot split the job of creating an image. I only can create images simultanously on 3 GPUs. Independing from each other.

[MaxFBurg]

It is problematic for training using data parallelism too: using 2 GPUs, the # Make sure we can run the memory efficient attention check allocates memory on the first gpu 2 times (for both checks), limiting the batch size one can use for multi gpu training.

[patrickvonplaten]

@MaxFBurg,

Could you maybe provide a code example that would help me debug this problem? E.g. a bash script that runs an official training script with torch distributed

[patrickvonplaten]

I see @MaxFBurg , yes you are 100% right here

[patrickvonplaten]

Opening a PR to not do this automatically

[MaxFBurg]

Thank you @patrickvonplaten ! sorry I did not reply before, I did not see the notification

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.