[LoRA fusion tests] SDXL LoRA tests with fusion when using `peft`

[sayakpaul]

Currently, the following test is failing:

RUN_SLOW=1 pytest tests/lora/test_lora_layers_peft.py::LoraSDXLIntegrationTests::test_sdxl_1_0_fuse_unfuse_all

Error:

self = <test_lora_layers_peft.LoraSDXLIntegrationTests testMethod=test_sdxl_1_0_fuse_unfuse_all>

    def test_sdxl_1_0_fuse_unfuse_all(self):
        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
        text_encoder_1_sd = copy.deepcopy(pipe.text_encoder.state_dict())
        text_encoder_2_sd = copy.deepcopy(pipe.text_encoder_2.state_dict())
        unet_sd = copy.deepcopy(pipe.unet.state_dict())
    
        pipe.load_lora_weights(
            "davizca87/sun-flower", weight_name="snfw3rXL-000004.safetensors", torch_dtype=torch.float16
        )
    
        fused_te_state_dict = pipe.text_encoder.state_dict()
        fused_te_2_state_dict = pipe.text_encoder_2.state_dict()
        unet_state_dict = pipe.unet.state_dict()
    
        for key, value in text_encoder_1_sd.items():
>           self.assertTrue(torch.allclose(fused_te_state_dict[key], value))
E           KeyError: 'text_model.encoder.layers.0.self_attn.k_proj.weight'

When I investigated fused_te_state_dict with the following:

targeted_layers = list(filter(lambda x: "text_model.encoder.layers.0.self_attn" in x, fused_te_state_dict.keys()))
print(targeted_layers)

I got:

['text_model.encoder.layers.0.self_attn.k_proj.base_layer.weight', 'text_model.encoder.layers.0.self_attn.k_proj.base_layer.bias', 'text_model.encoder.layers.0.self_attn.k_proj.lora_A.default_0.weight', 'text_model.encoder.layers.0.self_attn.k_proj.lora_B.default_0.weight', 'text_model.encoder.layers.0.self_attn.v_proj.base_layer.weight', 'text_model.encoder.layers.0.self_attn.v_proj.base_layer.bias', 'text_model.encoder.layers.0.self_attn.v_proj.lora_A.default_0.weight', 'text_model.encoder.layers.0.self_attn.v_proj.lora_B.default_0.weight', 'text_model.encoder.layers.0.self_attn.q_proj.base_layer.weight', 'text_model.encoder.layers.0.self_attn.q_proj.base_layer.bias', 'text_model.encoder.layers.0.self_attn.q_proj.lora_A.default_0.weight', 'text_model.encoder.layers.0.self_attn.q_proj.lora_B.default_0.weight', 'text_model.encoder.layers.0.self_attn.out_proj.base_layer.weight', 'text_model.encoder.layers.0.self_attn.out_proj.base_layer.bias', 'text_model.encoder.layers.0.self_attn.out_proj.lora_A.default_0.weight', 'text_model.encoder.layers.0.self_attn.out_proj.lora_B.default_0.weight']

So the error makes sense to me. @BenjaminBossan @younesbelkada could you take a look here?

[younesbelkada]

@sayakpaul do you know since when the test is failing to identify which commit caused the bug? it will be much easier to understand what happened

[sayakpaul]

@younesbelkada I can confirm that with peft 0.6.2 it works as expected but not with 0.7.1.

[BenjaminBossan]

Yes indeed, the key is now

+ text_model.encoder.layers.0.self_attn.k_proj.base_layer.weight
- text_model.encoder.layers.0.self_attn.k_proj.weight

We probably have to adjust the mapping :-/

[younesbelkada]

thanks @sayakpaul @BenjaminBossan for investigating! @BenjaminBossan what do you mean by the mapping? here we should just tweak that specific test right?

[BenjaminBossan]

Thanks for discovering this and alerting us @sayakpaul. Hopefully, we can get the CI to work correctly so that we will detect these issues immediately the next time.

I have investigated further and discussed with Younes. The reason why this fails is that we made a change in PEFT 0.7.0 that slightly re-arranges where weights of the base model are stored when loading the adapter. The LoRA weights are not affected. This change is backwards compatible, i.e. it's possible to load existing adapters, but some weights may end up in a different attribute. Therefore, when comparing the keys, they need to be remapped. I made some amendments for the test and it now passes:

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import copy
+import importlib
 import os
 import tempfile
 import time
@@ -24,6 +25,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
 from huggingface_hub.repocard import RepoCard
+from packaging import version
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
 from diffusers import (
@@ -1983,10 +1985,26 @@ class LoraSDXLIntegrationTests(unittest.TestCase):
         fused_te_2_state_dict = pipe.text_encoder_2.state_dict()
         unet_state_dict = pipe.unet.state_dict()
 
+        peft_ge_070 = version.parse(importlib.metadata.version("peft")) >= version.parse("0.7.0")
+
+        def remap_key(key, sd):
+            # some keys have moved around for PEFT >= 0.7.0, but they should still be loaded correctly
+            if (key in sd) or (not peft_ge_070):
+                return key
+
+            # instead of linear.weight, we now have linear.base_layer.weight, etc.
+            if key.endswith(".weight"):
+                key = key[:-7] + ".base_layer.weight"
+            elif key.endswith(".bias"):
+                key = key[:-5] + ".base_layer.bias"
+            return key
+
         for key, value in text_encoder_1_sd.items():
+            key = remap_key(key, fused_te_state_dict)
             self.assertTrue(torch.allclose(fused_te_state_dict[key], value))
 
         for key, value in text_encoder_2_sd.items():
+            key = remap_key(key, fused_te_2_state_dict)
             self.assertTrue(torch.allclose(fused_te_2_state_dict[key], value))

Ideally, we would also compare that the outputs haven't changed, but this test doesn't include that. Therefore, I just did a few generations and they look correct to me:

I think if the base weights were loaded incorrectly, we wouldn't see these results.

[sayakpaul]

Great! Thanks so much for this! Would you maybe like to open a PR here?

Cc: @DN6

[patrickvonplaten]

I don't understand how this is not backwards breaking. It's said that:

The LoRA weights are not affected.

but in #6216 we need to remap essentially all text encoder keys for PEFT 7.0.

[patrickvonplaten]

This means we now have support two different PEFT formats in diffusers which really is not great. What was the PR in PEFT that causes us to merge this change to better understand?

[BenjaminBossan]

but in #6216 we need to remap essentially all text encoder keys for PEFT 7.0.

In that PR, the weights are loaded correctly, even without the mapping. It's just that when we compare the state dicts before and after, some keys have different names. The re-mapping is not required to load the weights correctly, it only exists to make the test pass.

This means we now have support two different PEFT formats in diffusers which really is not great.

If I'm not missing something, that is not necessary.

What was the PR in PEFT that causes us to merge this change to better understand?

huggingface/peft#1106