Device Combinations Bug of Flux Quantization With Bitsandbytes

[CyberVy]

Describe the bug

Quantizing Flux with Bitsandbytes and moving the pipeline to another cuda device instead of cuda:0 will cause a device combination bug.

Everything works fine when the pipeline works in cuda:0.

Reproduction

This reproduction shows how the components are quantized.
And now I've updated a easier reproduction in the latest comment.

# First, load a quantized Flux pipeline
# The way loading the pipeline is from https://huggingface.co/docs/diffusers/quantization/bitsandbytes?bnb=4-bit

import torch
from diffusers import FluxPipeline

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import FluxTransformer2DModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(load_in_4bit=True,bnb_4bit_compute_dtype=torch.float16)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True,bnb_4bit_compute_dtype=torch.float16)

transformer_4bit = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_4bit,
    text_encoder_2=text_encoder_2_4bit,
    torch_dtype=torch.float16,
).to("cuda:1")

print([str(item.device) for item in list(pipe.components.values()) if hasattr(item,"device")])
# The output is ['cuda:1', 'cuda:1', 'cuda:1', 'cuda:1'], which means the components are moved to cuda:1 successfully.

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]
# An error below will be raised.

Logs

...
/usr/local/lib/python3.10/dist-packages/bitsandbytes/functional.py in is_on_gpu(tensors)
    467 
    468     if len(gpu_ids) > 1:
--> 469         raise RuntimeError(
    470             f"Input tensors need to be on the same GPU, but found the following tensor and device combinations:\n {[(t.shape, t.device) for t in tensors]}",
    471         )

RuntimeError: Input tensors need to be on the same GPU, but found the following tensor and device combinations:
 [(torch.Size([393216, 1]), device(type='cuda', index=1)), (torch.Size([1, 256]), device(type='cuda', index=1)), (torch.Size([1, 3072]), device(type='cuda', index=1)), (torch.Size([12288]), device(type='cuda', index=1)), (torch.Size([16]), device(type='cuda', index=0))]

System Info

• 🤗 Diffusers version: 0.32.2
• Platform: Linux-6.6.56+-x86_64-with-glibc2.35
• Running on Google Colab?: Yes
• Python version: 3.10.12
• PyTorch version (GPU?): 2.5.1+cu121 (True)
• Flax version (CPU?/GPU?/TPU?): 0.8.5 (gpu)
• Jax version: 0.4.33
• JaxLib version: 0.4.33
• Huggingface_hub version: 0.28.1
• Transformers version: 4.48.3
• Accelerate version: 1.2.1
• PEFT version: 0.14.0
• Bitsandbytes version: 0.45.2
• Safetensors version: 0.4.5
• xFormers version: not installed
• Accelerator: Tesla T4, 15360 MiB
  Tesla T4, 15360 MiB
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help?

@sayakpaul @DN6

[CyberVy]

RuntimeError: Input tensors need to be on the same GPU, but found the following tensor and device combinations:
 [(torch.Size([393216, 1]), device(type='cuda', index=1)), (torch.Size([1, 256]), device(type='cuda', index=1)), (torch.Size([1, 3072]), device(type='cuda', index=1)), (torch.Size([12288]), device(type='cuda', index=1)), (torch.Size([16]), device(type='cuda', index=0))]

It is very strange that a small tensor whose size is [16] is in cuda:0.

[sayakpaul]

Cc: @matthewdouglas

[CyberVy]

I've uploaded the quantized model to huggingface, here is the easier reproduction.

from diffusers import FluxPipeline
import torch
pipe = FluxPipeline.from_pretrained("eramth/flux-4bit",torch_dtype=torch.float16).to("cuda:1")

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]

[CyberVy]

Cc: @matthewdouglas

@sayakpaul
Thanks for your response.
However, I think maybe it's not the bug of bitsandbytes.
It seems like Diffusers works on cuda:0 by default when quantization, but after moving the pipeline to the other cuda device, something still works on cuda:0.

[CyberVy]

@sayakpaul

According to my test, I found the culprit.

This snippet is from diffusers/models/modeling_utils.py

if hf_quantizer is None:
    param_device = "cpu"
    # TODO (sayakpaul,  SunMarc): remove this after model loading refactor
    else:
        param_device = torch.device(torch.cuda.current_device())

This line causes the bug.

 param_device = torch.device(torch.cuda.current_device())

torch.cuda.current_device() is 0 by default, so torch.device(torch.cuda.current_device()) is always cuda:0.

When I change param_device into the other cuda device, diffusers will work on that device with no bug.

 param_device = torch.device("cuda:1")

It works on cuda:1.
However, it only works on cuda:1, moving to other device will cause the same bug.

I've also checked the code of bitsandbytes.
This snippet is from bitsandbytes/nn/modules.py

out = bnb.matmul_4bit(x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state)

The device of quant_state.code is always on the param_device which may lead an error for bitsandbytes.functional.is_on_gpu.
It seems like there is something wrong withquant_state which is passed into by Diffusers.

[sayakpaul]

Thanks for the investigation. Can you check #10604 and if installing diffusers from the PR branch solves the problem?

[CyberVy]

Thanks for the investigation. Can you check #10604 and if installing diffusers from the PR branch solves the problem?

@sayakpaul
I've tried that branch, but the same bug still exists.

[sayakpaul]

I think we need to discuss these kinds of CUDA device related movements. Which is why I cc'd @matthewdouglas

[matthewdouglas]

Hi folks,

I'm going to open an issue and a PR over on bitsandbytes for this. I can reproduce this and it is an issue to be fixed on the bitsandbytes side. This happens because bitsandbytes.nn.Params4bit.to() does not move the quant_state.code tensor. Everything else appears to be moved correctly.

[CyberVy]

Hi folks,

I'm going to open an issue and a PR over on bitsandbytes for this. I can reproduce this and it is an issue to be fixed on the bitsandbytes side. This happens because bitsandbytes.nn.Params4bit.to() does not move the quant_state.code tensor. Everything else appears to be moved correctly.

@matthewdouglas Thank you very much.

[sayakpaul]

@matthewdouglas I see that the PR that is fixing the issue has been merged. Would you suggest installing bitsandbytes from the main and retrying?

I am also going to close this issue as we agreed that this is not a diffusers specific bug.

[matthewdouglas]

At the moment installing from a sdist like main is a little complicated as it also requires building the CUDA code separately.

I'm expecting to put out a patch release very soon (within the next few days), but if needed in the meantime wheels from our CI can be installed:

• bdist_wheel_ubuntu-latest_x86_64
• bdist_wheel_windows-latest_x86_64