Why set_lora_device doesn't work

[West2022]

Describe the bug

When I load serverl loras with set_lora_device(), the GPU memory continues to grow, cames from 20G to 25G, this function doesn't work

Reproduction

for key in lora_list:
weight_name = key + ".safetensors"
pipe.load_lora_weights(lora_path, weight_name=weight_name, adapter_name=key, local_files_only=True)
adapters = pipe.get_list_adapters()
print(adapters)
pipe.set_lora_device([key], torch.device('cpu'))

Logs

No response

System Info

V100 32G
diffusers 0.32.0.dev0
torch 2.0.1+cu118
peft 0.12.0

Who can help?

No response

[sayakpaul]

The reproduction seems very incomplete. Can you please provide a fuller reproduction?

Also what versions of diffusers and peft are you using?

[West2022]

The reproduction seems very incomplete. Can you please provide a fuller reproduction?

Also what versions of diffusers and peft are you using?

I need to load multiple Loras and switch between different Loras. Each time, I load the used Lora onto the GPU through set_lora_device, while the unused ones are loaded onto the CPU. When initializing these Loras, after load_lora_weight, they are uniformly loaded onto the CPU, as shown in the code:

lora_list = ['aaa', 'bbb', 'ccc', 'ddd', 'eee']
for key in lora_list:
    weight_name = key + ".safetensors"
    pipe.load_lora_weights(lora_path, weight_name=weight_name, adapter_name=key, local_files_only=True)
    adapters = pipe.get_list_adapters()
    print(adapters)
    pipe.set_lora_device([key], torch.device('cpu'))

actually, it does save much GPU memory, but GPU memory continues to grow slowly. I understand that after calling pipe.set_lora_device([adapter], 'cpu'), GPU VRAM should not grow

before:
19911MiB / 32510MiB
after:
21279MiB / 32510MiB

diffusers and peft version:
diffusers 0.32.0.dev0
torch 2.0.1+cu118
peft 0.12.0

[West2022]

The reproduction seems very incomplete. Can you please provide a fuller reproduction?

Also what versions of diffusers and peft are you using?

There's some issues. If different Loras are loaded, some Loras contain text_encoder, while others do not, only contain unet, then set_lora_device () will report an key error

for component in self._lora_loadable_modules:
      model = getattr(self, component, None)
      if model is not None:
          for module in model.modules():
              if isinstance(module, BaseTunerLayer):
                    for adapter_name in adapter_names:
                        module.lora_A[adapter_name].to(device)
                        module.lora_B[adapter_name].to(device)
                        # this is a param, not a module, so device placement is not in-place -> re-assign
                        if hasattr(module, "lora_magnitude_vector") and module.lora_magnitude_vector is not None:
                            if adapter_name in module.lora_magnitude_vector:
                                module.lora_magnitude_vector[adapter_name] = module.lora_magnitude_vector[
                                    adapter_name
                                ].to(device)

Need to determine if the key is in the module.lora_A and module.lora_B

if adapter_name in module.lora_A:
    module.lora_A[adapter_name].to(device)
if adapter_name in module.lora_B:
    module.lora_B[adapter_name].to(device)

[sayakpaul]

Can you try with a more recent version of PyTorch and peft?

If different Loras are loaded, some Loras contain text_encoder, while others do not, only contain unet, then set_lora_device () will report an key error

Yeah this seems right. This also seems like a different issue. Would you maybe like to open a PR for this?

Cc: @BenjaminBossan

[BenjaminBossan]

Without a reproducer, it's a bit hard to check, but I think this change should solve the issue:

def set_lora_device(model, adapter_names, device):
    # copied from LoraBaseMixin.set_lora_device
    for module in model.modules():
        if isinstance(module, BaseTunerLayer):
            for adapter_name in adapter_names:
                # ADDED next 2 lines:
                if (adapter_name not in module.lora_A) and (adapter_name not in module.lora_embedding_A):
                    continue
                module.lora_A[adapter_name].to(device)
                module.lora_B[adapter_name].to(device)
                # this is a param, not a module, so device placement is not in-place -> re-assign
                if hasattr(module, "lora_magnitude_vector") and module.lora_magnitude_vector is not None:
                    if adapter_name in module.lora_magnitude_vector:
                        module.lora_magnitude_vector[adapter_name] = module.lora_magnitude_vector[adapter_name].to(
                            device
                        )

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[a-r-r-o-w]

If this is still an issue, please provide us a minimal code example to reproduce the error. I believe Ben provided a solution that might work, so please do give it a try and let us know.

If this goes stale again due to inactivity, will be marking as closed.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.