[LoRA] allow loras to be loaded with low_cpu_mem_usage.

[sayakpaul]

What does this PR do?

huggingface/peft#1961 added the ability to set low_cpu_mem_usage while loading LoRAs. This can be quite helpful in speeding up the loading of LoRAs that are large and have many layers.

#8953 is a good example where this feature could be beneficial.

Benchmarking code

from diffusers import DiffusionPipeline 
import torch 
import time 
import fire

def main(ckpt_id: str, lora_id: str, low_cpu_mem_usage: bool = False):
    pipeline = DiffusionPipeline.from_pretrained(ckpt_id, torch_dtype=torch.bfloat16).to("cuda")
    
    for _ in range(10):
        start_time = time.time()  
        pipeline.load_lora_weights(lora_id, low_cpu_mem_usage=low_cpu_mem_usage)
        end_time = time.time()  
        pipeline.unload_lora_weights()
        elapsed_time = end_time - start_time 
        
        print(f"Iteration {_ + 1}: Load Lora weights took {elapsed_time:.6f} seconds with {low_cpu_mem_usage=}")

if __name__ == "__main__":
    fire.Fire(main)

Results:

Iteration 1: Load Lora weights took 13.924374 seconds with low_cpu_mem_usage=True
Iteration 2: Load Lora weights took 1.621597 seconds with low_cpu_mem_usage=True
Iteration 3: Load Lora weights took 1.612010 seconds with low_cpu_mem_usage=True
Iteration 4: Load Lora weights took 1.670260 seconds with low_cpu_mem_usage=True
Iteration 5: Load Lora weights took 1.664858 seconds with low_cpu_mem_usage=True
Iteration 6: Load Lora weights took 1.482521 seconds with low_cpu_mem_usage=True
Iteration 7: Load Lora weights took 1.633697 seconds with low_cpu_mem_usage=True
Iteration 8: Load Lora weights took 1.593326 seconds with low_cpu_mem_usage=True
Iteration 9: Load Lora weights took 1.503672 seconds with low_cpu_mem_usage=True
Iteration 10: Load Lora weights took 1.566633 seconds with low_cpu_mem_usage=True

---

Iteration 1: Load Lora weights took 33.370373 seconds with low_cpu_mem_usage=False
Iteration 2: Load Lora weights took 3.937800 seconds with low_cpu_mem_usage=False
Iteration 3: Load Lora weights took 4.364943 seconds with low_cpu_mem_usage=False
Iteration 4: Load Lora weights took 4.303800 seconds with low_cpu_mem_usage=False
Iteration 5: Load Lora weights took 4.154818 seconds with low_cpu_mem_usage=False
Iteration 6: Load Lora weights took 3.869319 seconds with low_cpu_mem_usage=False
Iteration 7: Load Lora weights took 4.153911 seconds with low_cpu_mem_usage=False
Iteration 8: Load Lora weights took 4.275074 seconds with low_cpu_mem_usage=False
Iteration 9: Load Lora weights took 4.395445 seconds with low_cpu_mem_usage=False
Iteration 10: Load Lora weights took 4.071344 seconds with low_cpu_mem_usage=False

The feature currently needs the user to install peft and transformers from the source. So, I suggest we wait until both the libraries have made stable releases to merge this PR.

Once Ben reviews the PR, will request for a review from Yiyi.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

I think we need to have a low_cpu_mem_usage tag in the load_adapter() method too: https://github.com/huggingface/diffusers/blob/28f9d84549c0b1d24ef00d69a4c723f3a11cffb6/src/diffusers/loaders/lora_pipeline.py#L371C30-L371C42

If so, then this PR would be contingent on that. We could, however, use a combination of inject_adapter_in_model() and set_peft_model_state_dict() to mimic the same thing, I assume. I wouldn't personally prefer that because load_adapter() has been there for a while in diffusers.

Could you please clarify that bit? On the PEFT side, we have low_cpu_mem_usage on load_adapter but that's not the method being used here (just has the same name), right? Is this method coming from transformers (i.e. here)?

[sayakpaul]

that's not the method being used here (just has the same name), right? Is this method coming from transformers (i.e. here)?

Yes, that is correct.

[sayakpaul]

@BenjaminBossan when I used:

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(
    "TheLastBen/The_Hound", 
    weight_name="sandor_clegane_single_layer.safetensors", 
    low_cpu_mem_usage=True
)

prompt = "sandor clegane drinking in a pub"
image = pipe(
    prompt=prompt,
    num_inference_steps=30,
    width=1024,
    generator=torch.manual_seed(42),
    height=1024,
).images[0]
image.save("sandor.png")

It leads to:

Error trace

  File "/home/sayak/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 98, in forward
    hidden_states = gate * self.proj_out(hidden_states)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/peft/tuners/lora/layer.py", line 585, in forward
    result = result + lora_B(lora_A(dropout(x))) * scaling
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_CUDA_mm)

I investigated this a bit and confirmed that the LoRA params are kept on CPU which causes this failure. In case of low_cpu_mem_usage=False the LoRA parameters are on the expected device ("cuda" in the above example).

I further investigated why the tests added in this PR don't fail. That is because the state dict we're supplying to set_peft_model_state_dict() (here) -- the tensors of that state dict are already on the desired device. When I forcibly changed their device to a CPU and ran the tests on a GPU, the tests failed, and they complained about the same thing.

Possible to look into this?

[BenjaminBossan]

Hmm, I could not reproduce the issue :-/ I had to change the code slightly due to memory constraints:

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0: "24GB", 1: "20GB"},
)
# after loading the LoRA adapter {p.device for p in pipe.transformer.parameters()} returns:
# {device(type='cuda', index=0)}

Could this be the reason why it works for me?

I also tried this with a normal PEFT model that I moved to CUDA and then loaded with low_cpu_mem_usage and it worked.

[sayakpaul]

Yeah you need to be on the exact same setup to replicate this. We cannot assume people will do load_lora_weights() only in a specific manner.

You can perhaps use an SD LoRA:

diffusers/tests/lora/test_lora_layers_sd.py

Line 329 in 534848c

def test_a1111(self):

[BenjaminBossan]

Quick update, I couldn't reproduce with that model:

import torch
from diffusers import FluxPipeline, StableDiffusionPipeline

generator = torch.Generator().manual_seed(0)
pipe = StableDiffusionPipeline.from_pretrained("hf-internal-testing/Counterfeit-V2.5", safety_checker=None).to("cuda")
lora_model_id = "hf-internal-testing/civitai-light-shadow-lora"
lora_filename = "light_and_shadow.safetensors"
pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, low_cpu_mem_usage=True)
images = pipe(
    "masterpiece, best quality, mountain", output_type="np", generator=generator, num_inference_steps=2
).images

[sayakpaul]

@BenjaminBossan here's a minimal reproduction:

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("sayakpaul/tiny-flux-pipeline-with-lora", torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(
    "sayakpaul/tiny-flux-pipeline-with-lora", weight_name="pytorch_lora_weights.bin", low_cpu_mem_usage=True
)

prompt = "sandor clegane drinking in a pub"
image = pipe(prompt=prompt, num_inference_steps=30).images[0]
image.save("sandor.png")

I am on peft:main.

[BenjaminBossan]

Okay, got it now, thanks for the memory-friendly reproducer.

Indeed, if the LoRA weights on the model are on meta device, the device will be taken from the state_dict, not the base layer. I worked on a fix: huggingface/peft#2113

For the time being, you could add this snippet and it should fix the issue:

    if low_cpu_mem_usage:
        for module in model.modules():
            if hasattr(module, "_move_adapter_to_device_of_base_layer"):
                module._move_adapter_to_device_of_base_layer(adapter_name)

[sayakpaul]

@BenjaminBossan could you give this a review?

[sayakpaul]

@BenjaminBossan thanks!

Docstrings still contain TODO

Yeah will be resolved after Yiyi's approval.

Since PEFT version v0.13.1 is now released, the min PEFT version should be updated accordingly.

That's done.

In case you plan on no longer supporting older PEFT and transformer versions in the future: I would add a TODO comment to all those version checks that they can be removed once support for those older versions is dropped. If you plan to support them indefinitely, ignore this comment.

For now, we can ignore it.

An entry to the diffusers PEFT docs would be nice to have, especially since the name of the argument is not really intuitive.

Done.

[sayakpaul]

@yiyixuxu could you review this PR?

After the approval, will add docs and request for a review from Steven.

[yiyixuxu]

thanks! I left one comment, otherwise PR looks good to me

[sayakpaul]

Thanks, @yiyixuxu!

I have taken care of the TODOs in the docs too. @stevhliu could you review the related changes?

[stevhliu]

Thanks! Just one minor change that needs to be propagated :)