NF4 quantized flux models with loras

[hamzaakyildiz]

Is there any update here ? With nf4 quantized flux models, i could not use any lora

Update: NF4 serialization and loading are working fine. @DN6 let's brainstorm how we can support it more easily? This would help us unlock doing LoRAs on the quantized weights, too (cc: @BenjaminBossan for PEFT). I think this will become evidently critical for larger models.

transformers has a nice reference for us to follow. Additionally, accelerate has: https://huggingface.co/docs/accelerate/en/usage_guides/quantization, but it doesn't support NF4 serialization yet.

Cc: @SunMarc for jamming on this together.

Originally posted by @sayakpaul in #9165 (comment)

[sayakpaul]

Without any reproducible snippet and detailing what you have already tried, we cannot do much.

[hamzaakyildiz]

sorry for the confusion, i thought working with loras on quantized models is an ongoing work. I misunderstood, sorry 🙏

here is what i am trying to do

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

from diffusers import FluxTransformer2DModel, DiffusionPipeline, FluxPipeline
from transformers import T5EncoderModel
import torch
dtype = torch.bfloat16

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

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

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

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

orig_pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype)
pipeline = FluxPipeline.from_pipe(
    orig_pipeline, transformer=transformer_4bit, text_encoder_2=text_encoder_2_4bit, torch_dtype=dtype
)

adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
pipeline.load_lora_weights(adapter_id)

pipeline.enable_model_cpu_offload()

prompt = "A mystic cat with a sign that says hello world!"
image = pipeline(prompt, guidance_scale=3.5, num_inference_steps=8).images[0]

this is the error i get.
RuntimeError Traceback (most recent call last)
Cell In[1], line 39
34 pipeline = FluxPipeline.from_pipe(
35 orig_pipeline, transformer=transformer_4bit, text_encoder_2=text_encoder_2_4bit, torch_dtype=dtype
36 )
38 adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
---> 39 pipeline.load_lora_weights(adapter_id)
41 pipeline.enable_model_cpu_offload()
43 prompt = "A mystic cat with a sign that says hello world!"

File ~/.pyenv/versions/3.9.16/envs/jupyter/lib/python3.9/site-packages/diffusers/loaders/lora_pipeline.py:1856, in FluxLoraLoaderMixin.load_lora_weights(self, pretrained_model_name_or_path_or_dict, adapter_name, **kwargs)
1849 transformer_norm_state_dict = {
1850 k: state_dict.pop(k)
1851 for k in list(state_dict.keys())
1852 if "transformer." in k and any(norm_key in k for norm_key in self._control_lora_supported_norm_keys)
1853 }
1855 transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
-> 1856 has_param_with_expanded_shape = self.maybe_expand_transformer_param_shape_or_error(
1857 transformer, transformer_lora_state_dict, transformer_norm_state_dict
1858 )
1860 if has_param_with_expanded_shape:
1861 logger.info(
1862 "The LoRA weights contain parameters that have different shapes that expected by the transformer. "
1863 "As a result, the state_dict of the transformer has been expanded to match the LoRA parameter shapes. "
1864 "To get a comprehensive list of parameter names that were modified, enable debug logging."
1865 )

File ~/.pyenv/versions/3.9.16/envs/jupyter/lib/python3.9/site-packages/diffusers/loaders/lora_pipeline.py:2359, in FluxLoraLoaderMixin.maybe_expand_transformer_param_shape_or_error(cls, transformer, lora_state_dict, norm_state_dict, prefix)
2356 parent_module = transformer.get_submodule(parent_module_name)
2358 with torch.device("meta"):
-> 2359 expanded_module = torch.nn.Linear(
2360 in_features, out_features, bias=bias, dtype=module_weight.dtype
2361 )
2362 # Only weights are expanded and biases are not. This is because only the input dimensions
2363 # are changed while the output dimensions remain the same. The shape of the weight tensor
2364 # is (out_features, in_features), while the shape of bias tensor is (out_features,), which
2365 # explains the reason why only weights are expanded.
2366 new_weight = torch.zeros_like(
2367 expanded_module.weight.data, device=module_weight.device, dtype=module_weight.dtype
2368 )

File ~/.pyenv/versions/3.9.16/envs/jupyter/lib/python3.9/site-packages/torch/nn/modules/linear.py:105, in Linear.init(self, in_features, out_features, bias, device, dtype)
103 self.in_features = in_features
104 self.out_features = out_features
--> 105 self.weight = Parameter(
106 torch.empty((out_features, in_features), **factory_kwargs)
107 )
108 if bias:
109 self.bias = Parameter(torch.empty(out_features, **factory_kwargs))

File ~/.pyenv/versions/3.9.16/envs/jupyter/lib/python3.9/site-packages/torch/nn/parameter.py:46, in Parameter.new(cls, data, requires_grad)
42 data = torch.empty(0)
43 if type(data) is torch.Tensor or type(data) is Parameter:
44 # For ease of BC maintenance, keep this path for standard Tensor.
45 # Eventually (tm), we should change the behavior for standard Tensor to match.
---> 46 return torch.Tensor.make_subclass(cls, data, requires_grad)
48 # Path for custom tensors: set a flag on the instance to indicate parameter-ness.
49 t = data.detach().requires_grad(requires_grad)

RuntimeError: Only Tensors of floating point and complex dtype can require gradients

[sayakpaul]

Can you try something similar to https://github.com/huggingface/diffusers/tree/main/examples/research_projects/flux_lora_quantization#inference?

I will give it a broader look to see if load_lora_weights() just works, too. Cc: @BenjaminBossan anything obvious to remember to make a quantized model work so that it can load LoRAs?

[remixer-dec]

I'm having the same issue RuntimeError: Only Tensors of floating point and complex dtype can require gradients despite the top-level function called with torch.no_grad.
I've used https://github.com/HighCWu/flux-4bit with the previous versions of diffusers (0.31.0) and everything worked just fine, loras worked as expected

[sayakpaul]

Do you have a reproducible snippet that worked?

[remixer-dec]

No, only used it as a part of a bigger project, and I'm not sure how much different the implementation of 4-bits is in that repo from codebase that was merged into diffusers to add support for nf4, have to mention that the loras themselves were not quantized and it just worked out of the box.

[sayakpaul]

I'm not sure how much different the implementation of 4-bits is in that repo from codebase that was merged into diffusers to add support for nf4

Think most of it was repurposed from the stuff I shared in #9165 (comment) which is even older than the first commit in that repo, but I could be mistaken.

In any case, a minimal reproducible snippet would help.

I've used https://github.com/HighCWu/flux-4bit with the previous versions of diffusers (0.31.0) and everything worked just fine, loras worked as expected

I think if you install diffusers==0.31.0 (when bitsandbytes quantization was introduced in diffusers) and try to do LoRAs, it would work fine.

[sayakpaul]

Tracking here: #10550.

Turned out that even with 0.31.0 diffusers, loading LoRA into a quantized base model from diffusers will lead to errors.

[hamzaakyildiz]

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

from diffusers import FluxTransformer2DModel, DiffusionPipeline, FluxPipeline
from transformers import T5EncoderModel
import torch
import gc
dtype = torch.bfloat16

transformer = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    torch_dtype=dtype,
)

pipeline = FluxPipeline(
    transformer=transformer, 
    text_encoder_2=None, 
    vae=None,
    scheduler=None,
    text_encoder=None, 
    tokenizer=None,
    tokenizer_2=None
)

adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
pipeline.load_lora_weights(adapter_id)
pipeline.fuse_lora()

pipeline.transformer.save_pretrained("fused_transformer")

del pipeline
del transformer
gc.collect()
torch.cuda.empty_cache()


quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

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

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

transformer_4bit = FluxTransformer2DModel.from_pretrained(
    "./fused_transformer",
    quantization_config=quant_config,
    torch_dtype=dtype,
)

orig_pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype)
pipeline = FluxPipeline.from_pipe(
    orig_pipeline, transformer=transformer_4bit, text_encoder_2=text_encoder_2_4bit, torch_dtype=dtype
)

pipeline.to('cuda')

prompt = "A mystic cat with a sign that says hello world!"
image = pipeline(prompt, guidance_scale=3.5, num_inference_steps=8).images[0]

it gives an error of missing keys as loading the transformer_4bit. it works properly without loading the lora.

ValueError: Cannot load <class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> from ./fused_transformer because the following keys are missing:
transformer_blocks.17.ff_context.net.0.proj.bias, transformer_blocks.1.ff.net.0.proj.weight, single_transformer_blocks.34.attn.to_k.weight, transformer_blocks.10.attn.add_q_proj.weight,
....
....
transformer_blocks.3.norm1.linear.weight, transformer_blocks.14.ff_context.net.0.proj.bias, transformer_blocks.15.ff.net.2.bias, single_transformer_blocks.32.attn.to_k.weight, transformer_blocks.4.attn.add_v_proj.bias, transformer_blocks.5.attn.add_k_proj.weight, single_transformer_blocks.0.proj_mlp.weight.
Please make sure to pass low_cpu_mem_usage=False and device_map=None if you want to randomly initialize those weights or else make sure your checkpoint file is correct.

Additionally, loading the full transformer block and LoRA (approximately 24GB of VRAM) with an additional max_shard_size (default 10GB) to save the model to disk, resulting in a total VRAM requirement of (24GB + max_shard_size). This can be controlled with:

pipeline.transformer.save_pretrained("fused_transformer", max_shard_size="{x}GB")

[sayakpaul]

it gives an error of missing keys as loading the transformer_4bit. it works properly without loading the lora.

You have to call unload_lora_weights() before saving the transformer.

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

from diffusers import FluxTransformer2DModel, DiffusionPipeline, FluxPipeline
from transformers import T5EncoderModel
import torch
import gc
dtype = torch.bfloat16

transformer = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    torch_dtype=dtype,
)

pipeline = FluxPipeline(
    transformer=transformer, 
    text_encoder_2=None, 
    vae=None,
    scheduler=None,
    text_encoder=None, 
    tokenizer=None,
    tokenizer_2=None
)

adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
pipeline.load_lora_weights(adapter_id, adapter_weights=0.125)
pipeline.fuse_lora()
pipeline.unload_lora_weights()

pipeline.transformer.save_pretrained("fused_transformer")

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, FluxTransformer2DModel
from diffusers import FluxPipeline
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig, T5EncoderModel
import torch

dtype = torch.bfloat16
quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

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

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

transformer_4bit = FluxTransformer2DModel.from_pretrained(
    "fused_transformer",
    quantization_config=quant_config,
    torch_dtype=dtype,
)

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    transformer=transformer_4bit, 
    text_encoder_2=text_encoder_2_4bit, 
    torch_dtype=dtype
).to("cuda")

prompt = "A mystic cat with a sign that says hello world!"
image = pipeline(prompt, guidance_scale=3.5, num_inference_steps=8).images[0]

Works for me.

[hamzaakyildiz]

okay. it's working now. thank you.

[yiyixuxu]

closing as fixed! thanks @sayakpaul !