[lora] `exclude_modules` won't consider the modules names in `target_modules`

[sayakpaul]

In #11806, we introduced automatic extraction of exclude_modules and passing it to the following function to prepare LoraConfig kwargs:

diffusers/src/diffusers/utils/peft_utils.py

Line 153 in 425a715

def get_peft_kwargs(

The problem with it is that if the target_modules includes a module name where a substring from exclude_modules is present, the module from target_modules becomes invalid for peft. So, for example, if target_modules has single_transformer_blocks.0.proj_out and exclude_modules has proj_out, we will see a problem. See below

Expand

oading adapter weights from state_dict led to unexpected keys found in the model: single_transformer_blocks.0.proj_out.lora_A.default_0.weight, single_transformer_blocks.0.proj_out.lora_B.default_0.weight, single_transformer_blocks.1.proj_out.lora_A.default_0.weight, single_transformer_blocks.1.proj_out.lora_B.default_0.weight, single_transformer_blocks.2.proj_out.lora_A.default_0.weight, single_transformer_blocks.2.proj_out.lora_B.default_0.weight, single_transformer_blocks.3.proj_out.lora_A.default_0.weight, single_transformer_blocks.3.proj_out.lora_B.default_0.weight, single_transformer_blocks.4.proj_out.lora_A.default_0.weight, single_transformer_blocks.4.proj_out.lora_B.default_0.weight, single_transformer_blocks.5.proj_out.lora_A.default_0.weight, single_transformer_blocks.5.proj_out.lora_B.default_0.weight, single_transformer_blocks.6.proj_out.lora_A.default_0.weight, single_transformer_blocks.6.proj_out.lora_B.default_0.weight, single_transformer_blocks.7.proj_out.lora_A.default_0.weight, single_transformer_blocks.7.proj_out.lora_B.default_0.weight, single_transformer_blocks.8.proj_out.lora_A.default_0.weight, single_transformer_blocks.8.proj_out.lora_B.default_0.weight, single_transformer_blocks.9.proj_out.lora_A.default_0.weight, single_transformer_blocks.9.proj_out.lora_B.default_0.weight, single_transformer_blocks.10.proj_out.lora_A.default_0.weight, single_transformer_blocks.10.proj_out.lora_B.default_0.weight, single_transformer_blocks.11.proj_out.lora_A.default_0.weight, single_transformer_blocks.11.proj_out.lora_B.default_0.weight, single_transformer_blocks.12.proj_out.lora_A.default_0.weight, single_transformer_blocks.12.proj_out.lora_B.default_0.weight, single_transformer_blocks.13.proj_out.lora_A.default_0.weight, single_transformer_blocks.13.proj_out.lora_B.default_0.weight, single_transformer_blocks.14.proj_out.lora_A.default_0.weight, single_transformer_blocks.14.proj_out.lora_B.default_0.weight, single_transformer_blocks.15.proj_out.lora_A.default_0.weight, single_transformer_blocks.15.proj_out.lora_B.default_0.weight, single_transformer_blocks.16.proj_out.lora_A.default_0.weight, single_transformer_blocks.16.proj_out.lora_B.default_0.weight, single_transformer_blocks.17.proj_out.lora_A.default_0.weight, single_transformer_blocks.17.proj_out.lora_B.default_0.weight, single_transformer_blocks.18.proj_out.lora_A.default_0.weight, single_transformer_blocks.18.proj_out.lora_B.default_0.weight, single_transformer_blocks.19.proj_out.lora_A.default_0.weight, single_transformer_blocks.19.proj_out.lora_B.default_0.weight, single_transformer_blocks.20.proj_out.lora_A.default_0.weight, single_transformer_blocks.20.proj_out.lora_B.default_0.weight, single_transformer_blocks.21.proj_out.lora_A.default_0.weight, single_transformer_blocks.21.proj_out.lora_B.default_0.weight, single_transformer_blocks.22.proj_out.lora_A.default_0.weight, single_transformer_blocks.22.proj_out.lora_B.default_0.weight, single_transformer_blocks.23.proj_out.lora_A.default_0.weight, single_transformer_blocks.23.proj_out.lora_B.default_0.weight, single_transformer_blocks.24.proj_out.lora_A.default_0.weight, single_transformer_blocks.24.proj_out.lora_B.default_0.weight, single_transformer_blocks.25.proj_out.lora_A.default_0.weight, single_transformer_blocks.25.proj_out.lora_B.default_0.weight, single_transformer_blocks.26.proj_out.lora_A.default_0.weight, single_transformer_blocks.26.proj_out.lora_B.default_0.weight, single_transformer_blocks.27.proj_out.lora_A.default_0.weight, single_transformer_blocks.27.proj_out.lora_B.default_0.weight, single_transformer_blocks.28.proj_out.lora_A.default_0.weight, single_transformer_blocks.28.proj_out.lora_B.default_0.weight, single_transformer_blocks.29.proj_out.lora_A.default_0.weight, single_transformer_blocks.29.proj_out.lora_B.default_0.weight, single_transformer_blocks.30.proj_out.lora_A.default_0.weight, single_transformer_blocks.30.proj_out.lora_B.default_0.weight, single_transformer_blocks.31.proj_out.lora_A.default_0.weight, single_transformer_blocks.31.proj_out.lora_B.default_0.weight, single_transformer_blocks.32.proj_out.lora_A.default_0.weight, single_transformer_blocks.32.proj_out.lora_B.default_0.weight, single_transformer_blocks.33.proj_out.lora_A.default_0.weight, single_transformer_blocks.33.proj_out.lora_B.default_0.weight, single_transformer_blocks.34.proj_out.lora_A.default_0.weight, single_transformer_blocks.34.proj_out.lora_B.default_0.weight, single_transformer_blocks.35.proj_out.lora_A.default_0.weight, single_transformer_blocks.35.proj_out.lora_B.default_0.weight, single_transformer_blocks.36.proj_out.lora_A.default_0.weight, single_transformer_blocks.36.proj_out.lora_B.default_0.weight, single_transformer_blocks.37.proj_out.lora_A.default_0.weight, single_transformer_blocks.37.proj_out.lora_B.default_0.weight. 

Reproducer:

from diffusers import DiffusionPipeline
import torch 

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
).to("cuda")
pipe_kwargs = {
    "prompt": "{trigger_word} A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 28,
    "max_sequence_length": 512,
}
pipeline.load_lora_weights("glif/l0w-r3z")
image = pipeline(**pipe_kwargs).images[0]

diffusers/tests/lora/test_lora_layers_wanvace.py

Line 171 in 425a715

def test_lora_exclude_modules_wanvace(self):

doesn't exactly test this scenario, either.

Cc: @BenjaminBossan

[BenjaminBossan]

It feels like there is always another edge case :-/

I tried to run your example and got the same warning. However, inference seems to run (but I can't verify if the result is correct due to memory). Do you get a correct image?

I tried to isolate the issue for a unit test, would you say that this correctly captures the essence of the error?

import torch
from torch import nn
from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict


class Inner(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj_out = nn.Linear(3, 4)  # same name as outer

class Outer(nn.Module):
    def __init__(self):
        super().__init__()
        self.inner = Inner()
        self.proj_out = nn.Linear(5, 6)  # same name as inner
        self.proj_in = nn.Linear(7, 8)   # some other layer, not important


# passes
def test_no_exclude_modules():
    state_dict = {
        "base_model.model.inner.proj_out.lora_A.weight": torch.ones(8, 3) * 33,
        "base_model.model.inner.proj_out.lora_B.weight": torch.ones(4, 8) * 44,
        "base_model.model.proj_out.lora_A.weight": torch.ones(8, 5) * 55,
        "base_model.model.proj_out.lora_B.weight": torch.ones(6, 8) * 66,
    }
    model = Outer()
    config = LoraConfig(target_modules=["proj_out", "proj_in"])
    model = get_peft_model(model, config)
    set_peft_model_state_dict(model, state_dict)
    get_peft_model_state_dict(model)

    assert (model.base_model.model.inner.proj_out.lora_A["default"].weight == 33).all()
    assert (model.base_model.model.inner.proj_out.lora_B["default"].weight == 44).all()
    assert (model.base_model.model.proj_out.lora_A["default"].weight == 55).all()
    assert (model.base_model.model.proj_out.lora_B["default"].weight == 66).all()


# passes
def test_exclude_inner():
    state_dict = {
        "base_model.model.proj_out.lora_A.weight": torch.ones(8, 5) * 55,
        "base_model.model.proj_out.lora_B.weight": torch.ones(6, 8) * 66,
    }
    model = Outer()
    config = LoraConfig(target_modules=["proj_out", "proj_in"], exclude_modules=["inner.proj_out"])
    model = get_peft_model(model, config)
    set_peft_model_state_dict(model, state_dict)
    get_peft_model_state_dict(model)

    assert not hasattr(model.base_model.inner.proj_out, "lora_A")
    assert not hasattr(model.base_model.inner.proj_out, "lora_B")
    assert (model.base_model.model.proj_out.lora_A["default"].weight == 55).all()
    assert (model.base_model.model.proj_out.lora_B["default"].weight == 66).all()


# fails
def test_exclude_outer():
    state_dict = {
        "base_model.model.inner.proj_out.lora_A.weight": torch.ones(8, 3) * 33,
        "base_model.model.inner.proj_out.lora_B.weight": torch.ones(4, 8) * 44,
    }
    model = Outer()
    config = LoraConfig(target_modules=["inner.proj_out", "proj_in"], exclude_modules=["proj_out"])
    model = get_peft_model(model, config)
    set_peft_model_state_dict(model, state_dict)
    get_peft_model_state_dict(model)

    print(model)
    # INCORRECT: base_model.model.inner.proj_out should be a LoRA layer
    # PeftModel(
    #   (base_model): LoraModel(
    #     (model): Outer(
    #       (inner): Inner(
    #         (proj_out): Linear(in_features=3, out_features=4, bias=True)
    #       )
    #       (proj_out): Linear(in_features=5, out_features=6, bias=True)
    #       (proj_in): lora.Linear(
    #         (base_layer): Linear(in_features=7, out_features=8, bias=True)
    #         (lora_dropout): ModuleDict(
    #           (default): Identity()
    #         )
    #         (lora_A): ModuleDict(
    #           (default): Linear(in_features=7, out_features=8, bias=False)
    #         )
    #         (lora_B): ModuleDict(
    #           (default): Linear(in_features=8, out_features=8, bias=False)
    #         )
    #         (lora_embedding_A): ParameterDict()
    #         (lora_embedding_B): ParameterDict()
    #         (lora_magnitude_vector): ModuleDict()
    #       )
    #     )
    #   )
    # )
    # Next line fails with: AttributeError: 'Linear' object has no attribute 'lora_A'
    assert (model.base_model.model.inner.proj_out.lora_A["default"].weight == 33).all()
    assert (model.base_model.model.inner.proj_out.lora_B["default"].weight == 44).all()
    assert not hasattr(model.base_model.model, "proj_out")
    assert not hasattr(model.base_model.model.proj_out, "lora_A")

The last test case fails because inner.proj_out is not targeted with LoRA.

Overall, this whole situation is very messy because of the different checkpoint formats. We're basically trying to reverse engineer a correct LoraConfig based on a state_dict that 1) has no meta data and 2) was not created using PEFT. There is no guarantee that it's even possible for each and every case to reverse engineer a correct LoraConfig at all [1]. Solutions like exclude_modules are really just clumsy workarounds. I wonder if we have to rethink this problem on a more fundamental level, e.g. by providing a way that does not rely on string matching at all, since in theory, we know which modules should be targeted.

[1] In theory, it is always possible, since target_modules can be a regex, so we should be able to express every pattern. But creating such a regex correctly seems almost impossible.

[sayakpaul]

I tried to run your example and got the same warning. However, inference seems to run (but I can't verify if the result is correct due to memory). Do you get a correct image?

No, I don't get a correct image. If you do pipe.enable_model_cpu_offload() instead of the direct pipe.to("cuda"), you will be fine I think.

Overall, this whole situation is very messy because of the different checkpoint formats. We're basically trying to reverse engineer a correct LoraConfig based on a state_dict that 1) has no meta data and 2) was not created using PEFT. There is no guarantee that it's even possible for each and every case to reverse engineer a correct LoraConfig at all [1]. Solutions like exclude_modules are really just clumsy workarounds. I wonder if we have to rethink this problem on a more fundamental level, e.g. by providing a way that does not rely on string matching at all, since in theory, we know which modules should be targeted.

I hear you, but IMO, this is a very valid scenario, actually. One can very well just target all the transformer blocks (which would include proj_out-named modules) and exclude proj_out (i.e., the final layer of the network).

by providing a way that does not rely on string matching at all, since in theory, we know which modules should be targeted.

I do not have a clear idea yet but when crafting the final exclude_modules in LoraConfig, we could try to also account for the target_module names. If there are overlaps between the substrings (substrings should denote module names here IIUC) of exclude_modules and target_modules, we could account for it.

[BenjaminBossan]

but IMO, this is a very valid scenario, actually

I didn't mean to imply that the scenario is not valid, I mean that the way we deal with it is messy.

I do not have a clear idea yet

I'm just speculating at this point, but when we call

inject_adapter_in_model(lora_config, self, adapter_name=adapter_name, **peft_kwargs)

in diffusers, we already know the correct state_dict at that point. I wonder if we could extend the signature of inject_adapter_in_model to accept a state_dict. Right now, in PEFT, we basically do:

for key, module in model.named_modules():
    if self._check_target_module_exists(peft_config, key):  # <= here all the messy string matching happens
        # inject LoRA layer

After we have injected the LoRA layers this way, we load them from the state_dict in diffusers. This is why it's crucial that we have a perfect match between the state_dict and the reverse engineered target_modules. However, if we already know the state_dict, we could instead do this in PEFT:

for key, module in model.named_modules():
    if key in state_dict:
        # inject LoRA layer

That way, the state_dict itself becomes the single source of truth and reverse engineering the correct target_modules would be unnecessary. I don't know yet if that works reliably in all situation, but it could be worth a try? Apart of requiring the correct lora_config.target_modules to initiate PEFT this way, do we need it for anything else? If not, I could imagine this proposal to be viable.

[BenjaminBossan]

Hmm, something that might also work is target_modules = "|".join(rf"^{key}$" for key in state_dict) but it is a bit cursed.

[sayakpaul]

Thanks for giving this more thoughts!

That way, the state_dict itself becomes the single source of truth and reverse engineering the correct target_modules would be unnecessary. I don't know yet if that works reliably in all situation, but it could be worth a try? Apart of requiring the correct lora_config.target_modules to initiate PEFT this way, do we need it for anything else? If not, I could imagine this proposal to be viable.

I think this is a better solution because of its explicitness and probably a bit more future proof than the regex one. But it's really upto you.

[BenjaminBossan]

Okay, when I have a bit of time on my hands, I'll prototype something and then we can see if it's a viable solution.