flux.1-dev device_map didn't work

[hznnnnnn]

I try to use device_map to use multiple gpu's, but it not worked, how can I use all my gpus?

[a-r-r-o-w]

What did you try? What did not work? What environment are your running in? Could you update the description with the minimal reproducible snippet that can replicate the behaviour you are facing, and run diffusers-cli env to tell us about your environment?

cc @sayakpaul

[hznnnnnn]

What did you try? What did not work? What environment are your running in? Could you update the description with the minimal reproducible snippet that can replicate the behaviour you are facing, and run diffusers-cli env to tell us about your environment?

cc @sayakpaul

I try to use device_map in FluxPipeline: pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16, use_safetensors=True, device_map="balanced"), but it not worked

[hznnnnnn]

What did you try? What did not work? What environment are your running in? Could you update the description with the minimal reproducible snippet that can replicate the behaviour you are facing, and run diffusers-cli env to tell us about your environment?

cc @sayakpaul

I run in ubantu, and I have 8 gpus

[sayakpaul]

Please post the error trace.

[hznnnnnn]

Please post the error trace.

RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle)
but I have 8 gpus, only used 2 gpus

[sayakpaul]

Provide the full error trace please.

[hznnnnnn]

Provide the full error trace please.

now I can worked, this is my code:
max_memory = {0:"20GB", 1:"20GB", 2:"20GB", 3:"20GB", 4:"20GB", 5:"20GB", 6:"20GB", 7:"20GB"}
pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16, use_safetensors=True, device_map="balanced", max_memory=max_memory)
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
#pipe.enable_sequential_cpu_offload()

prompt = "A real person, with long red hair, hot body, sexy, lovely, 8k"
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
output_type="pil",
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-dev.png")

But now it only one gpu is working

[hznnnnnn]

Provide the full error trace please.

this is full error trance:
(flux.1-dev) root@master:/home/omnisky/flux.1-dev/FLUX.1-dev# python test.py
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00, 1.10it/s]
Loading pipeline components...: 86%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏ | 6/7 [00:08<00:02, 2.48s/it]You set add_prefix_space. The tokenizer needs to be converted from the slow tokenizers
Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:09<00:00, 1.31s/it]
0%| | 0/50 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/home/omnisky/flux.1-dev/FLUX.1-dev/test.py", line 13, in
image = pipe(
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py", line 696, in call
noise_pred = self.transformer(
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/accelerate/hooks.py", line 169, in new_forward
output = module._old_forward(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/diffusers/models/transformers/transformer_flux.py", line 362, in forward
hidden_states = self.x_embedder(hidden_states)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/accelerate/hooks.py", line 169, in new_forward
output = module._old_forward(*args, **kwargs)
File "/root/anaconda3/envs/flux.1-dev/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 117, in forward
return F.linear(input, self.weight, self.bias)
RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle)

and code:
pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16, use_safetensors=True, device_map="balanced")
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
#pipe.enable_sequential_cpu_offload()

prompt = "A real person, with long red hair, hot body, sexy, lovely, 8k"
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
output_type="pil",
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-dev.png")

[sayakpaul]

I am able to run the following:

from diffusers import FluxPipeline 
import torch 


model_path = "black-forest-labs/FLUX.1-dev"
pipe = FluxPipeline.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, device_map="balanced"
)
print(pipe.hf_device_map) 
image = pipe(
    prompt="dog",
    height=1024,
    width=1024,
    guidance_scale=3.5,
    output_type="pil",
    num_inference_steps=50,
    max_sequence_length=512,
).images[0]

I am running this on two A100s.

This is how the device map looks like:

{'transformer': 0, 'vae': 0, 'text_encoder_2': 1, 'text_encoder': 1}

[sayakpaul]

Able to do do this with the following setup too:

pipe = FluxPipeline.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0:"24GB", 1:"24GB"}
)

[sayakpaul]

Also able to run with

pipe = FluxPipeline.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0:"16GB", 1:"16GB", 2:"24GB"}
)

[hznnnnnn]

Also able to run with

pipe = FluxPipeline.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0:"16GB", 1:"16GB", 2:"24GB"}
)

this is my code:

import torch
from diffusers import FluxPipeline
from accelerate import PartialState

model_path = "/home/omnisky/flux.1-dev/FLUX.1-dev"

pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0:"16GB", 1:"16GB", 2:"16GB", 3:"16GB"})
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
#pipe.enable_sequential_cpu_offload()

print(pipe.hf_device_map)

prompt = "A real person, with long red hair, hot body, sexy, lovely, 8k"
image = pipe(
    prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    output_type="pil",
    num_inference_steps=50,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-dev.png")

the device map is : {'transformer': 'cpu', 'text_encoder_2': 0, 'text_encoder': 1, 'vae': 2}
Why transformer worked in cpu? Is it my gpu that sucks?My gpu is Tesla p40, 8 gpus.

[hznnnnnn]

Also able to run with

pipe = FluxPipeline.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0:"16GB", 1:"16GB", 2:"24GB"}
)

I removed generator=torch.Generator("cpu").manual_seed(0), but it always show transformer in cpu

[sayakpaul]

Probably because of the low VRAM. May try increasing that to 20GB since it's supported?

[hznnnnnn]

Probably because of the low VRAM. May try increasing that to 20GB since it's supported?

max_memory={0:"16GB", 1:"16GB", 2:"16GB", 3:"16GB"}) this 16gb increasing to 20GB?