Problem with Flux Schnell bfloat16 multiGPU

[OlegRuban-ai]

Describe the bug

Hello! I set device_map='balanced' and get images generated in 2.5 minutes (expected in 12-20 seconds), while in pipe.hf_device_map it shows that the devices are distributed like this:

{
 "transformer": "cuda:0",
 "text_encoder_2": "cuda:2",
 "text_encoder": "cuda:0",
 "vae": "cuda:1"
                }

I have 3 video cards 3090 Ti 24 GB and I can’t run it on them.

I also tried this way:
pipe.transformer.to('cuda:2')
pipe.text_encoder.to('cuda:2')
pipe.text_encoder_2.to('cuda:1')
pipe.vae.to('cuda:0')

What is the best way to launch it so that generation occurs on the GPU and quickly?

Reproduction

            pipe = FluxPipeline.from_pretrained(
                path_chkpt,
                torch_dtype=torch.bfloat16,
                device_map='balanced',
            )

Logs

No response

System Info

ubuntu 22.04 3 GPU: 3090 TI 24 GB

accelerate==0.30.1
addict==2.4.0
apscheduler==3.9.1
autocorrect==2.5.0
chardet==4.0.0
cryptography==37.0.2
curl_cffi
diffusers==0.30.0
beautifulsoup4==4.11.2
einops
facexlib>=0.2.5
fastapi==0.92.0
hidiffusion==0.1.6
invisible-watermark>=0.2.0
numpy==1.24.3
opencv-python==4.8.0.74
pandas==2.0.3
pycocotools==2.0.6
pymystem3==0.2.0
pyyaml==6.0
pyjwt==2.6.0
python-multipart==0.0.5
pytrends==4.9.1
psycopg2-binary
realesrgan==0.3.0
redis==4.5.1
sacremoses==0.0.53
selenium==4.2.0
sentencepiece==0.1.97
scipy==1.10.1
scikit-learn==0.24.1
supervision==0.16.0
tb-nightly==2.14.0a20230629
tensorboard>=2.13.0
tomesd
transformers==4.40.1
timm==0.9.16
yapf==0.32.0
uvicorn==0.20.0

spacy==3.7.2
nest_asyncio==1.5.8
httpx==0.25.0

torchvision==0.15.2

insightface==0.7.3
psutil==5.9.6
tk==0.1.0
customtkinter==5.2.1
tensorflow==2.13.0
opennsfw2==0.10.2
protobuf==4.24.4
gfpgan==1.3.8

Who can help?

No response

[OlegRuban-ai]

Now device_map can only be balanced and does not support loading the dictionary with distribution over the GPU, but when I set it to balanced, the transformer is sent to the CPU, which may be the reason for the severe slowdown. It doesn’t fit into one GPU, and fp8 doesn’t yet support smaller formats. Please help solve the problem.

[asomoza]

in both examples your using the transformer which is the bigger model with a text encoder in the same GPU, the transformer can fit (alone) in a 24GB GPU if you're not using a GUI so the best solution here is to change the distribution so that both text encoders are in the same GPU and leave the transformer alone in another.

Also with this PR you can use a sharded transformer Flux model.

[OlegRuban-ai]

in both examples your using the transformer which is the bigger model with a text encoder in the same GPU, the transformer can fit (alone) in a 24GB GPU if you're not using a GUI so the best solution here is to change the distribution so that both text encoders are in the same GPU and leave the transformer alone in another.

Also with this PR you can use a sharded transformer Flux model.

The problem is that when I put it like this:

        pipe.transformer.to('cuda:0')
        pipe.text_encoder.to('cuda:1')
        pipe.text_encoder_2.to('cuda:1')
        pipe.vae.to('cuda:2')

then I run into the error RuntimeError('Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)')

when I put it like this:
pipe.transformer.to('cuda:0')
pipe.text_encoder.to('cuda:0')
pipe.text_encoder_2.to('cuda:1')
pipe.vae.to('cuda:2')

CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 23.69 GiB total capacity; 22.29 GiB already allocated; 8.75 MiB free; 22.30 GiB reserved in total by PyTorch)

[asomoza]

oh, this means that in the second configuration it fails before the inference because it OOMs when loading the model, it doesn't mean that it works with the second configuration better than the first one.

One solution is to separate the stages like in the PR (get text embeddings, denoising, decode latents) so the tensors are in the same device.

I'm not sure if this is supposed to be done automatically in the pipeline when using multi GPUs so ccing @sayakpaul for more insights.

Also I'm curious if you're going to get better performance that just using one GPU with cpu offloading, you can just use one 3090 with flux in bfloat16 if you have all the 24 GB available.

[sayakpaul]

In this case, it's best to use enable_model_cpu_offload() and run it in a single GPU:

pipe = FluxPipeline.from_pretrained(
    path_chkpt,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

Make sure you install accelerate from source. I am also not sure if bfloat16 is a preferred data-type for 3090.

In any case, the slow-down is somewhat expected because of the data movements.

[OlegRuban-ai]

In this case, it's best to use enable_model_cpu_offload() and run it in a single GPU:

pipe = FluxPipeline.from_pretrained(
    path_chkpt,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

I tried this option, but the generation speed of 80 seconds per image does not suit me. Are there any options to speed it up?

[sayakpaul]

Did you try with FP16?

[OlegRuban-ai]

Did you try with FP16?

Yes

[sayakpaul]

Hmm, then there are a couple of things to keep in mind.

Let's first summarize what we have:

1> device_map="balanced" won't solve the latency problem because of the data movements.
2> We cannot use a single 24GB card with Flux because not all the modules would fit on the VRAM.
3> enable_model_cpu_offload() doesn't meet your inference latency requirements (cc: @asomoza do you observe similar trends in timing too?).

We can perform quantization with optimum.quanto to solve 2. https://gist.github.com/AmericanPresidentJimmyCarter/873985638e1f3541ba8b00137e7dacd9 shows a thorough example. But that might lead to a slowdown in latency because the compute dtype is still FP16/BF16.

Next option would be to try out torchao, more specifically, the int8wo (as in int8 weight-only) quantization. This will help us fit the modules in memory, hopefully.

Now, to solve the latency problem we can make use of torch.compile() as in discussed here.

Would you like to give this a try?

Code (haven't tested):

from torchao.quantization import int8_weight_only, quantize_

... # pipeline loading code

quantize_(pipeline.transformer, int8_weight_only())
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)

You can potentially apply the same scheme to other modules of the pipeline, too and see if they help. Note that torch.compile() support may not be there yet for the CLIP and T5 model.

You can also try out the other quantization schemes from torchao which you can find from the repo.

We will be integrating NF4 support soon through bitsandbytes. Hopefully, that should solve your issues better. See: #9174

[OlegRuban-ai]

I used this code:

'''#from torchao.quantization import int8_weight_only, quantize_

pipeline loading code

quantize_(pipe.transformer, int8_weight_only())
pipe.transformer.to(memory_format=torch.channels_last)
pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
'''

Unsupported: hasattr ConstDictVariable to

from user code:
File "/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py", line 164, in new_forward
args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py", line 363, in pre_forward
return send_to_device(args, self.execution_device), send_to_device(
File "/usr/local/lib/python3.10/dist-packages/accelerate/utils/operations.py", line 148, in send_to_device
if is_torch_tensor(tensor) or hasattr(tensor, "to"):

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

1. In fp16, the generation speed is the same low (80 sec/it), but RAM and VRAM consumption is minimal

2. Doesn’t work via torch.compile()

3. transformer and text_encoder must be on the same video card, but together they do not fit into 24 GB

4. Quantization failed. There is hope that it will be possible to implement this model in diffusers: https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4

Thank you for your time.

[sayakpaul]

For torch.compile() and quantization things, please always use the nightlies.

transformer and text_encoder must be on the same video card, but together they do not fit into 24 GB

What if you first keep them on CPU, quantize, and then move to GPU?

Quantization failed. There is hope that it will be possible to implement this model in diffusers:

See here: #9165

[asomoza]

there's something wrong in your env though, I can run Flux dev with bfloat16 (don't have space to test schnell right now) in a single 3090 (not TI) without a GUI and just using enable_model_cpu_offload and with 20 steps it takes 29 seconds to generate

Taking 80 seconds using schnell with 4 steps with a 3090 TI means that something else is wrong.

Since you had OOMs is not that you're using RAM instead of VRAM like when you use Windows so my guess is that you're using really slow RAM or maybe you're swapping to disk?

As a reference, with bfloat16 or float16 and offloading to cpu you'll need more than 32GB of RAM, in my case I use more than 50GB with what I have loaded from before, so take that into consideration.

[yiyixuxu]

@OlegRuban-ai
can you confirm what's your ourput for this script?

from diffusers import DiffusionPipeline
import torch
path_chkpt = "black-forest-labs/FLUX.1-schnell"
pipe = DiffusionPipeline.from_pretrained(
    path_chkpt,
    torch_dtype=torch.bfloat16,
    device_map='balanced',
)

print(pipe.hf_device_map)

the device map you pasted here #9195 (comment) is not what we would expect, the algorithm should have automatically put vae and text_encoder into the same gpu

{
 "transformer": "cuda:0",
 "text_encoder_2": "cuda:2",
 "text_encoder": "cuda:0",
 "vae": "cuda:1"
                }

[OlegRuban-ai]

@yiyixuxu
This is what I get in the output. It is impossible to set parameters directly with a dictionary, as was previously possible in diffusers, because it only accepts device_map='balanced'

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

as you can see, the transformer gets to the CPU, although I have one 3090 video card that remains completely free