diffusers doesn't save and load the LoraConfig, resulting wrong `lora_alpha` during inference.

[YiqunChen1999]

Describe the bug

When I run train_dreambooth_lora_sdxl.py, I find the StableDiffusionXLLoraLoaderMixin (and LoraLoaderMixin) doesn't save & load the LoraConfig.

The default lora_alpha is 8 in LoraConfig but is 1 in peft Linear. The lora_alpha is 8 during training but is 1 when I load the weights from a checkpoint as the XXLoraLoaderMixin cannot parse the correct lora_alpha from the saved Lora weights. The wrong lora_alpha gives the wrong scaling in the peft Linear layer.

I manually reset the scaling to get the correct result, but I think this is a bug and should be fixed.

Reproduction

Clone the diffusers repo and install it. Then prepare any one image in a folder.

In the example/dreambooth folder, run:

python train_dreambooth_lora_sdxl.py \
    --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
    --pretrained_vae_model_name_or_path madebyollin/sdxl-vae-fp16-fix \
    --instance_data_dir [PATH_TO_YOUR_IMAGE_DIR]  \
    --instance_prompt "<sks>" \  # try to overfit the given image
    --validation_prompt "<sks>" \
    --num_validation_images 1 \
    --validation_epochs 50 \
    --output_dir [OUTPUT_DIR] \
    --train_batch_size 1 \
    --max_train_steps 300 \
    --checkpointing_steps 100 \
    --learning_rate 2e-4  \
    --rank 64

For example, I try to overfit the first image. The second image is the validation image at epoch 250 (tensorboard). The third one is the result after the pipeline loads Lora weights from the saved checkpoint (epoch 300, also available in tensorboard).

Logs

No response

System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

• diffusers version: 0.25.0.dev0
• Platform: Linux-4.18.0-305.3.1.el8.x86_64-x86_64-with-glibc2.28
• Python version: 3.11.5
• PyTorch version (GPU?): 2.1.1 (True)
• Huggingface_hub version: 0.19.4
• Transformers version: 4.35.2
• Accelerate version: 0.25.0
• xFormers version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

Who can help?

@sayakpaul @patri

[sayakpaul]

@younesbelkada could you take a look here?

[PhilippeNguyen]

I think I'm getting similar problems with examples/text_to_image/train_text_to_image_lora.py.
The tensorboard validation results are good as the model is training, but the final validation (when the model is saved and reloaded) is bad.
Loading the checkpoints afterwards also gives poor results

[YiqunChen1999]

I think I'm getting similar problems with examples/text_to_image/train_text_to_image_lora.py. The tensorboard validation results are good as the model is training, but the final validation (when the model is saved and reloaded) is bad. Loading the checkpoints afterwards also gives poor results

@PhilippeNguyen My current solution is the following, hope it's helpful:
During the training, I save the LoraConfig, e.g.:

unet_lora_config.save_pretrained(os.path.join(args.output_dir, 'unet'))
# optional
text_lora_config.save_pretrained(os.path.join(args.output_dir,
                                              'text_encoder'))

During inference, I load the LoraConfig and reset the lora_alpha and scaling:

adapter_name = 'default'
pipe.load_lora_weights(lora_path, adapter_name=adapter_name)
reset_pipeline_lora_alpha(pipe,
                          adapter_name,
                          unet_lora_config_path,
                          text_encoder_lora_config_path)

The functions are defined as follows:

from peft import LoraConfig
from peft.tuners.lora.layer import LoraLayer
from diffusers import DiffusionPipeline


def reset_pipeline_lora_alpha(pipe: DiffusionPipeline,
                              adapter: str,
                              unet_lora_config_path: str = None,
                              text_encoder_lora_config_path: str = None):
    if unet_lora_config_path:
        lora_config = LoraConfig.from_pretrained(unet_lora_config_path)
        print(f'Loading unet adapter config from {unet_lora_config_path}')
        print(f'Reset adapter {adapter} lora alpha in unet to '
              f'{lora_config.lora_alpha}')
        reset_model_lora_alpha(pipe.unet, lora_config.lora_alpha, adapter)
    if text_encoder_lora_config_path:
        lora_config = LoraConfig.from_pretrained(text_encoder_lora_config_path)
        print(f'Loading text_encoder adapter config from '
              f'{text_encoder_lora_config_path}')
        print(f'Reset adapter {adapter} lora alpha in text_encoder to '
              f'{lora_config.lora_alpha}')
        reset_model_lora_alpha(pipe.text_encoder,
                               lora_config.lora_alpha,
                               adapter)
        if hasattr(pipe, 'text_encoder_2'):
            print(f'Reset adapter {adapter} lora alpha in text_encoder_2 to '
                  f'{lora_config.lora_alpha}')
            reset_model_lora_alpha(pipe.text_encoder_2,
                                   lora_config.lora_alpha,
                                   adapter)


def reset_model_lora_alpha(model: nn.Module, lora_alpha: float, adapter: str):
    for name, module in model.named_modules():
        if isinstance(module, LoraLayer):
            reset_lora_alpha(module, lora_alpha, adapter)


def reset_lora_alpha(lora_layer: LoraLayer, lora_alpha: float, adapter: str):
    # Modified from peft.tuners.lora.layer.LoraLayer.
    if adapter not in lora_layer.active_adapters:
        return

    lora_layer.lora_alpha[adapter] = lora_alpha
    if lora_layer.r[adapter] > 0:
        lora_layer.scaling[adapter] = (lora_layer.lora_alpha[adapter]
                                       / lora_layer.r[adapter])

The solution might not be the best, but it works for me.

[sayakpaul]

@YiqunChen1999 #6135 does this help?

[YiqunChen1999]

@sayakpaul Yes, it works for me. Thanks for your efforts!

[sayakpaul]

@YiqunChen1999 we had to close #6135. Does #6225 help?

[YiqunChen1999]

@sayakpaul Yes, it works well in my case!

[eliahuhorwitz]

@sayakpaul I am getting the same when running the examples/dreambooth/train_dreambooth_lora.py script. At training time It seems to be using the default lora_alpha of 8 without an argument that controls it. At inference, when calling pipe.load_lora_weights(lora_model_path) lora_alpha is set to be lora rank (extracted from the lora state_dict).

[sayakpaul]

Ah, now I see. Applying the changes from #6225 should fix that. Can you confirm?

[eliahuhorwitz]

It fixes it, but I am not sure this is the best solution. Why not allow the user to specify the lora_alpha as an argument and only if they don't default it to the rank?
The issue with the above suggestion is that one would need to save the alpha value and load it when loading the lora weights.

[sayakpaul]

To keep it harmonized with the previous method of training. #6135 did what you just said but we eventually decided not to that route for now.

[eliahuhorwitz]

Got it, the results I am getting when training on rank=alpha=16 are worse than rank=16, alpha=8. I think (but not sure, running experiments now) that higher alpha should result in lower LR. Maybe it's worth noting this somewhere in the readme