.to() cast in .load_lora_weights breaks on bitsandbytes quantized text_encoder's for SD/SDXL

[Teriks]

Describe the bug

This .to() cast on the text encoder:

diffusers/src/diffusers/loaders/lora_base.py

Line 421 in 9836f0e

text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)

Is invalid when working with an SD1.5 / SDXL pipeline that has a bitsandbytes quantization config used on the text encoder

perhaps something like this would fix?

if is_bitsandbytes_available():
    quant, is_4bit, _ = _check_bnb_status(text_encoder)
else:
    quant, is_4bit = False, False

if not quant:
    text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)
elif is_4bit:
    text_encoder.to(device=text_encoder.device)

This problem does not seem to affect flux / sd3, so I am not sure how this would affect other pipelines?

Reproduction

import torch
import transformers

import diffusers
import diffusers.quantizers.quantization_config as _qc

text_encoder = transformers.CLIPTextModel.from_pretrained(
    'stabilityai/stable-diffusion-xl-base-1.0', subfolder='text_encoder', variant='fp16',
    torch_dtype=torch.float16, quantization_config=_qc.BitsAndBytesConfig(load_in_8bit=True))

pipeline = diffusers.StableDiffusionXLPipeline.from_pretrained(
    'stabilityai/stable-diffusion-xl-base-1.0',
    variant='fp16',
    torch_dtype=torch.float16,
    text_encoder=text_encoder
)

pipeline.load_lora_weights('Norod78/sdxl-emoji-lora')

pipeline.to('cuda')

pipeline(prompt='test')

Logs

REDACT\diffusers\venv\Scripts\python.exe REDACT\diffusers\test.py 
WARNING:torchao.kernel.intmm:Warning: Detected no triton, on systems without Triton certain kernels will not work
`low_cpu_mem_usage` was None, now default to True since model is quantized.
Loading pipeline components...: 100%|██████████| 7/7 [00:00<00:00, 15.86it/s]
Traceback (most recent call last):
  File "REDACT\diffusers\test.py", line 18, in <module>
    pipeline.load_lora_weights('Norod78/sdxl-emoji-lora')
  File "REDACT\diffusers\src\diffusers\loaders\lora_pipeline.py", line 657, in load_lora_weights
    self.load_lora_into_text_encoder(
  File "REDACT\diffusers\src\diffusers\loaders\lora_pipeline.py", line 894, in load_lora_into_text_encoder
    _load_lora_into_text_encoder(
  File "REDACT\diffusers\src\diffusers\loaders\lora_base.py", line 430, in _load_lora_into_text_encoder
    text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)
  File "REDACT\diffusers\venv\Lib\site-packages\transformers\modeling_utils.py", line 3089, in to
    raise ValueError(
ValueError: You cannot cast a bitsandbytes model in a new `dtype`. Make sure to load the model using `from_pretrained` using the desired `dtype` by passing the correct `torch_dtype` argument.

Process finished with exit code 1

System Info

diffusers == 0.34.0.dev0

Who can help?

@yiyixuxu @sayakpaul @DN6 @asomoza

[Teriks]

Correction* this issue also exists for SD3, which hits the same code path when loading a LoRA, I think the fix I mentioned is correct, I will test more later today and create a pull request

[Teriks]

for fuse_lora to work in 8bit particularly with SD3 it must be something like this

The re-offloading code interferes with peft merge when the text encoder is loaded in 8bit,
SD3 has is_sequential_cpu_offload=True and the weights end up on meta and it errors out
in peft dequant if not circumvent, in the bitsandbytes integration

Not sure where else that applies (what pipelines) but the sequential offload is incompatible with the 8bit quant
in every case I think, since the module cannot move

# scale LoRA layers with `lora_scale`
scale_lora_layers(text_encoder, weight=lora_scale)

if is_bitsandbytes_available():
    quant, is_4bit, is_8bit = _check_bnb_status(text_encoder)
else:
    quant, is_4bit, is_8bit = False, False, False

if not quant:
    text_encoder.to(device=text_encoder.device, dtype=text_encoder.dtype)
elif is_4bit:
    text_encoder.to(device=text_encoder.device)

# Offload back if not in 8bit
if not is_8bit:
    if is_model_cpu_offload:
        _pipeline.enable_model_cpu_offload()
    elif is_sequential_cpu_offload:
        _pipeline.enable_sequential_cpu_offload()
# Unsafe code />

This results in noise images with 8bit + LoRA though, not sure why

[sayakpaul]

Thanks a lot for all the investigation. Maybe we could tackle it case by case? Let's start with the initial issue first?

[Teriks]

Thanks a lot for all the investigation. Maybe we could tackle it case by case? Let's start with the initial issue first?

I think the code above solves the initial issue, as much as it can be in diffusers. this seems to work for merging lora into weights loaded in 4bit, albeit with a good amount of quality loss

Something goes wrong with the merge when loading in 8bit, probably in peft. which causes an overflow in inference, resulting in black images or noise.

There are warnings from peft about precession issues if they are fused.

Does not seem to overflow with 4bit quantized weights, just 8bit. pretty sure that is a peft issue.

[sayakpaul]

For fusing LoRAs into quantized models, see some guidance:
https://github.com/huggingface/diffusers/tree/main/examples/research_projects/flux_lora_quantization#inference

In any case, feel free to open PRs if you think diffusers could benefit from them.

[Teriks]

For fusing LoRAs into quantized models, see some guidance: https://github.com/huggingface/diffusers/tree/main/examples/research_projects/flux_lora_quantization#inference

In any case, feel free to open PRs if you think diffusers could benefit from them.

Thanks, I see why it is problematic to merge the LoRa into the already quantized base model, I guess the approach is the fuse it into the base model first and then quantize the resulting weights