Using LoRA will deactivate xFormers, even if it is enabled #3551

takuma104 · 2023-05-24T17:34:56Z

Describe the bug

Discussed in #3437 (comment) . It appears that memory usage increases significantly when using LoRA, even in environments using xFormers. I investigated the cause using this script. The results suggest that even in environments where xFormers is enabled, effectively resulting in the same situation as if xFormers had been deactivated.

As a solution, it seems good to use LoRAXFormersAttnProcessor instead of LoRAAttnProcessor if xFormers is enabled in this part.

diffusers/src/diffusers/loaders.py

Lines 275 to 286 in a94977b

    
           if isinstance( 
        
               attn_processor, (AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0) 
        
           ): 
        
               cross_attention_dim = value_dict["add_k_proj_lora.down.weight"].shape[1] 
        
               attn_processor_class = LoRAAttnAddedKVProcessor 
        
           else: 
        
               cross_attention_dim = value_dict["to_k_lora.down.weight"].shape[1] 
        
               attn_processor_class = LoRAAttnProcessor 
        
           attn_processors[key] = attn_processor_class( 
        
               hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=rank 
        
           )

What do you think?

Reproduction

https://gist.github.com/takuma104/e2139bda7f74cd977350e18500156683

Logs

{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 3837}
{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 3837}
{"width": 512, "height": 768, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 5878}
{"width": 512, "height": 768, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 5880}
{"width": 512, "height": 512, "batch": 2, "xformers": "OFF", "lora": "OFF", "mem_MB": 5505}
{"width": 512, "height": 512, "batch": 2, "xformers": "OFF", "lora": "ON", "mem_MB": 5507}
{"width": 512, "height": 768, "batch": 2, "xformers": "OFF", "lora": "OFF", "mem_MB": 9589}
{"width": 512, "height": 768, "batch": 2, "xformers": "OFF", "lora": "ON", "mem_MB": 9591}
{"width": 512, "height": 512, "batch": 4, "xformers": "OFF", "lora": "OFF", "mem_MB": 8842}
{"width": 512, "height": 512, "batch": 4, "xformers": "OFF", "lora": "ON", "mem_MB": 8844}
{"width": 512, "height": 768, "batch": 4, "xformers": "OFF", "lora": "OFF", "mem_MB": 17011}
{"width": 512, "height": 768, "batch": 4, "xformers": "OFF", "lora": "ON", "mem_MB": 17013}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 2806}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 3837}
{"width": 512, "height": 768, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 3125}
{"width": 512, "height": 768, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 5880}
{"width": 512, "height": 512, "batch": 2, "xformers": "ON", "lora": "OFF", "mem_MB": 3243}
{"width": 512, "height": 512, "batch": 2, "xformers": "ON", "lora": "ON", "mem_MB": 5507}
{"width": 512, "height": 768, "batch": 2, "xformers": "ON", "lora": "OFF", "mem_MB": 3780}
{"width": 512, "height": 768, "batch": 2, "xformers": "ON", "lora": "ON", "mem_MB": 9591}
{"width": 512, "height": 512, "batch": 4, "xformers": "ON", "lora": "OFF", "mem_MB": 4317}
{"width": 512, "height": 512, "batch": 4, "xformers": "ON", "lora": "ON", "mem_MB": 8844}
{"width": 512, "height": 768, "batch": 4, "xformers": "ON", "lora": "OFF", "mem_MB": 5392}
{"width": 512, "height": 768, "batch": 4, "xformers": "ON", "lora": "ON", "mem_MB": 17013}

System Info

diffusers version: 0.16.1
Platform: Linux-5.19.0-41-generic-x86_64-with-glibc2.35
Python version: 3.10.11
PyTorch version (GPU RTX3090): 2.0.1+cu117 (True)
Huggingface_hub version: 0.14.1
Transformers version: 4.29.1
Accelerate version: 0.19.0
xFormers version: 0.0.20
Using GPU in script?: True
Using distributed or parallel set-up in script?: Nope

The text was updated successfully, but these errors were encountered:

takuma104 · 2023-05-24T17:45:57Z

By the way, when I tried it with the latest main a94977b, the log became a bit strange. Even though xformers is not enabled, the log seems to suggest that it is enabled.

{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 2818}
{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 3312}
{"width": 512, "height": 768, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 3136}
{"width": 512, "height": 768, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 4683}
{"width": 512, "height": 512, "batch": 2, "xformers": "OFF", "lora": "OFF", "mem_MB": 3254}
{"width": 512, "height": 512, "batch": 2, "xformers": "OFF", "lora": "ON", "mem_MB": 4444}
{"width": 512, "height": 768, "batch": 2, "xformers": "OFF", "lora": "OFF", "mem_MB": 3791}
{"width": 512, "height": 768, "batch": 2, "xformers": "OFF", "lora": "ON", "mem_MB": 7186}
{"width": 512, "height": 512, "batch": 4, "xformers": "OFF", "lora": "OFF", "mem_MB": 4329}
{"width": 512, "height": 512, "batch": 4, "xformers": "OFF", "lora": "ON", "mem_MB": 6707}
{"width": 512, "height": 768, "batch": 4, "xformers": "OFF", "lora": "OFF", "mem_MB": 5403}
{"width": 512, "height": 768, "batch": 4, "xformers": "OFF", "lora": "ON", "mem_MB": 12192}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 2818}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 3312}
{"width": 512, "height": 768, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 3136}
{"width": 512, "height": 768, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 4683}
{"width": 512, "height": 512, "batch": 2, "xformers": "ON", "lora": "OFF", "mem_MB": 3254}
{"width": 512, "height": 512, "batch": 2, "xformers": "ON", "lora": "ON", "mem_MB": 4444}
{"width": 512, "height": 768, "batch": 2, "xformers": "ON", "lora": "OFF", "mem_MB": 3791}
{"width": 512, "height": 768, "batch": 2, "xformers": "ON", "lora": "ON", "mem_MB": 7186}
{"width": 512, "height": 512, "batch": 4, "xformers": "ON", "lora": "OFF", "mem_MB": 4329}
{"width": 512, "height": 512, "batch": 4, "xformers": "ON", "lora": "ON", "mem_MB": 6707}
{"width": 512, "height": 768, "batch": 4, "xformers": "ON", "lora": "OFF", "mem_MB": 5403}
{"width": 512, "height": 768, "batch": 4, "xformers": "ON", "lora": "ON", "mem_MB": 12192}

takuma104 · 2023-05-24T18:09:15Z

I have created a patch. It seems to be working as intended. The behavior when xFormers is not enabled remains unchanged.
main...takuma104:diffusers:lora-xformers

{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 2818}
{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 3311}
{"width": 512, "height": 768, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 3135}
{"width": 512, "height": 768, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 4682}
{"width": 512, "height": 512, "batch": 2, "xformers": "OFF", "lora": "OFF", "mem_MB": 3252}
{"width": 512, "height": 512, "batch": 2, "xformers": "OFF", "lora": "ON", "mem_MB": 4444}
{"width": 512, "height": 768, "batch": 2, "xformers": "OFF", "lora": "OFF", "mem_MB": 3790}
{"width": 512, "height": 768, "batch": 2, "xformers": "OFF", "lora": "ON", "mem_MB": 7185}
{"width": 512, "height": 512, "batch": 4, "xformers": "OFF", "lora": "OFF", "mem_MB": 4327}
{"width": 512, "height": 512, "batch": 4, "xformers": "OFF", "lora": "ON", "mem_MB": 6706}
{"width": 512, "height": 768, "batch": 4, "xformers": "OFF", "lora": "OFF", "mem_MB": 5402}
{"width": 512, "height": 768, "batch": 4, "xformers": "OFF", "lora": "ON", "mem_MB": 12191}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 2817}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 2819}
{"width": 512, "height": 768, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 3135}
{"width": 512, "height": 768, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 3137}
{"width": 512, "height": 512, "batch": 2, "xformers": "ON", "lora": "OFF", "mem_MB": 3252}
{"width": 512, "height": 512, "batch": 2, "xformers": "ON", "lora": "ON", "mem_MB": 3254}
{"width": 512, "height": 768, "batch": 2, "xformers": "ON", "lora": "OFF", "mem_MB": 3790}
{"width": 512, "height": 768, "batch": 2, "xformers": "ON", "lora": "ON", "mem_MB": 3791}
{"width": 512, "height": 512, "batch": 4, "xformers": "ON", "lora": "OFF", "mem_MB": 4327}
{"width": 512, "height": 512, "batch": 4, "xformers": "ON", "lora": "ON", "mem_MB": 4329}
{"width": 512, "height": 768, "batch": 4, "xformers": "ON", "lora": "OFF", "mem_MB": 5402}
{"width": 512, "height": 768, "batch": 4, "xformers": "ON", "lora": "ON", "mem_MB": 5403}

sayakpaul · 2023-05-25T04:08:49Z

Thanks for the investigation.

I took three snippets from the logs you posted above for the following settings: "width": 512, "height": 512, "batch": 1.

So that I better understand what the suspect looks like, what am I looking for here?

In the vanilla case:

{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 3837}
{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 3837}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 2806}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 3837}

We see for LoRA ons, even when it shows xformers on, the memory usage is still the same as the setting when xformers is off and LoRA is on. This is faulty. Correct?

From here:

{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 2818}
{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 3312}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 2818}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 3312}

What are we looking for here? You mentioned

Even though xformers is not enabled, the log seems to suggest that it is enabled.

With xformers and LoRA both being off, we have 2818 MBs. Even when xformers is on it's the same. So, it means xformers wasn't off in the first place, correct?

From here:

{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "OFF", "mem_MB": 2818}
{"width": 512, "height": 512, "batch": 1, "xformers": "OFF", "lora": "ON", "mem_MB": 3311}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "OFF", "mem_MB": 2817}
{"width": 512, "height": 512, "batch": 1, "xformers": "ON", "lora": "ON", "mem_MB": 2819}

Things actually don't seem off.

Is my understanding correct?

Regarding the solution you proposed -- I am okay with it. I think we should consider adding a test suite for this case to make the behaviour more robust.

Also, seeking suggestions from @pcuenca @williamberman @patrickvonplaten.

takuma104 · 2023-05-25T15:04:39Z

Is my understanding correct?

Yes, correct. I'm starting to think that the reason why it doesn't seem to change whether xFormers is On/Off when LoRA is OFF in versions other than 0.16.1, is because the AttnProcessor2_0 for torch2.0 is operating and the memory usage doesn't look much different from xFormers. It might be good to see what is actually being used by the Unet's AttnProcessor, so I'll write some code and check it out later.

Regarding the solution you proposed -- I am okay with it. I think we should consider adding a test suite for this case to make the behaviour more robust.

Thanks! I just open draft PR for it #3556.

williamberman · 2023-05-25T16:32:40Z

Is the summary just that the non-xformers LoRA attention is being used when xformers is enabled?

takuma104 · 2023-05-25T17:07:27Z

Is the summary just that the non-xformers LoRA attention is being used when xformers is enabled?

Yes! In my above environment where both LoRA and xFormers are OFF, the following passes, so it seems that the reason the memory usage is low is not because xFormers is enabled, but because AttnProcessor2_0 is being used.

for _, module in pipe.unet.named_modules():
    if isinstance(module, Attention):
        assert isinstance(module.processor, AttnProcessor2_0)

patrickvonplaten · 2023-05-26T09:40:04Z

We should add a AttnProcessor2_0 LoRA class here actually

sayakpaul · 2023-05-26T12:01:06Z

Inline with #3464 @patrickvonplaten

sayakpaul · 2023-05-26T12:03:49Z

Closing with #3556

takuma104 added the bug Something isn't working label May 24, 2023

takuma104 mentioned this issue May 24, 2023

Support Kohya-ss style LoRA file format (in a limited capacity) #3437

Merged

3 tasks

takuma104 mentioned this issue May 25, 2023

Fix to apply LoRAXFormersAttnProcessor instead of LoRAAttnProcessor when xFormers is enabled #3556

Merged

1 task

sayakpaul closed this as completed May 26, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Using LoRA will deactivate xFormers, even if it is enabled #3551

Using LoRA will deactivate xFormers, even if it is enabled #3551

takuma104 commented May 24, 2023

takuma104 commented May 24, 2023

takuma104 commented May 24, 2023

sayakpaul commented May 25, 2023

takuma104 commented May 25, 2023

williamberman commented May 25, 2023

takuma104 commented May 25, 2023 •

edited

Loading

patrickvonplaten commented May 26, 2023

sayakpaul commented May 26, 2023

sayakpaul commented May 26, 2023

Using LoRA will deactivate xFormers, even if it is enabled #3551

Using LoRA will deactivate xFormers, even if it is enabled #3551

Comments

takuma104 commented May 24, 2023

Describe the bug

Reproduction

Logs

System Info

takuma104 commented May 24, 2023

takuma104 commented May 24, 2023

sayakpaul commented May 25, 2023

takuma104 commented May 25, 2023

williamberman commented May 25, 2023

takuma104 commented May 25, 2023 • edited Loading

patrickvonplaten commented May 26, 2023

sayakpaul commented May 26, 2023

sayakpaul commented May 26, 2023

takuma104 commented May 25, 2023 •

edited

Loading