Trained layers with LoRA will be merged into model?

[lavinal712]

If I train a LoRA, and want to save some layers in safetenser, like this:

config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            target_modules=TARGET_MODULE,
            lora_dropout=args.lora_dropout,
            bias=args.lora_bias,
            modules_to_save=SAVED_MODULE,
        )
       model  = get_peft_model(model, config)

Then I conduct:

model = PeftModel.from_pretrained(model, "model")
model = model.merge_and_unload()
model.save_pretrained(save_path)

Trained layers with LoRA will be merged into model? Or just thrown away?

[BenjaminBossan]

model.merge_and_unload() will merge the LoRA weights into the base weights. If you want to throw them away, use model.unload().

[lavinal712]

Well, I mean I trained some layers (not lora) with lora, and when I use merge_and_unload(), will the layers be updated?

[BenjaminBossan]

No, only the LoRA weights should be updated. Therefore, when merging, only the layers that have LoRA applied should change. All other layers should be the same.

[lavinal712]

thank, I will try to adjust my code to fix it.

[BenjaminBossan]

If something doesn't seem to be working as expected, feel free to paste the code so that I can take a look.

[lavinal712]

I write this to merge lora and trainable layers to model, but failed.

CONTROLNET_TARGET_MODULE = ["to_q", "to_k", "to_v", "to_out.0",
                            "conv", "conv_in", "conv_out", "conv_shortcut", "conv1", "conv2",
                            "proj", "proj_in", "proj_out", "time_emb_proj",
                            "linear_1", "linear_2", "ff.net.2"]
CONTROLNET_SAVED_MODULE = ["controlnet_cond_embedding.blocks.0", "controlnet_cond_embedding.blocks.1",
                           "controlnet_cond_embedding.blocks.2", "controlnet_cond_embedding.blocks.3",
                           "controlnet_cond_embedding.blocks.4", "controlnet_cond_embedding.blocks.5",
                           "controlnet_down_blocks.0", "controlnet_down_blocks.1", "controlnet_down_blocks.2",
                           "controlnet_down_blocks.3", "controlnet_down_blocks.4", "controlnet_down_blocks.5",
                           "controlnet_down_blocks.6", "controlnet_down_blocks.7", "controlnet_down_blocks.8",
                           "controlnet_down_blocks.9", "controlnet_down_blocks.10", "controlnet_down_blocks.11",
                           "controlnet_mid_block"]
...
    if args.use_lora:
        config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            target_modules=CONTROLNET_TARGET_MODULE,
            lora_dropout=args.lora_dropout,
            bias=args.lora_bias,
            modules_to_save=CONTROLNET_SAVED_MODULE,
        )
        controlnet = get_peft_model(controlnet, config)
        controlnet.print_trainable_parameters()
...

from collections import OrderedDict
import torch
import torch.nn as nn

import diffusers
from diffusers import UNet2DConditionModel, ControlNetModel
from peft import PeftModel
from safetensors.torch import load_file


def _tie_weights(source_state_dict, target_state_dict):
    for k in source_state_dict.keys():
        if k.startswith("base_model.model."):
            new_k = k[len("base_model.model."):]
            if new_k in target_state_dict.keys():
                print(new_k)
                target_state_dict[new_k] = source_state_dict[k]
    return target_state_dict


base_model = "F:/stable-diffusion-v1-5"
controlnet_save_path = "./save"

unet = UNet2DConditionModel.from_pretrained(
    base_model, subfolder="unet", torch_dtype=torch.float16
)

controlnet = ControlNetModel.from_unet(unet)
controlnet_state_dict = controlnet.state_dict()
lora_state_dict = load_file("controlnet/adapter_model.safetensors")
controlnet_state_dict = _tie_weights(lora_state_dict, controlnet_state_dict)
controlnet.load_state_dict(controlnet_state_dict)
controlnet = PeftModel.from_pretrained(controlnet, "controlnet")
controlnet = controlnet.merge_and_unload()
controlnet.save_pretrained(controlnet_save_path)

The output model does not have control ability. I have trained it for 60000 steps with batch size of 4.
I wonder how to fix it.

[lavinal712]

I tried tie weights wihtout merge, there is result:

[lavinal712]

Here you can see there is no condition in image. I trained the model for 40000 steps with batch size of 4.

[lavinal712]

here is after merge_and_unload() w/o tie

still no control

[BenjaminBossan]

I don't know much about controlnet, so not sure what I should be looking for. I assume you have some reference image but the generated images don't resemble that reference image.

Anyway, if the model does not deliver expected results even without merging, it means that merging is not the issue. Instead, it appears like the model is not trained correctly. Is that observation correct?

[lavinal712]

Thanks for your help! I find my lr wae set 1e-5, which suggests that the model may converge slowly. After I change the lr, the model come to have control ability.