[LoRA] text encoder: read the ranks for all the attn modules #8324

elias-gaeros · 2024-05-30T17:59:25Z

What does this PR do?

The enumeration of LoRA adapters for gathering their ranks in LoraLoaderMixin.load_lora_into_text_encoder is modified in order to:

In addition to out_proj, read the ranks of adapters for q_proj, k_proj, and v_proj
Allow missing adapters (UNet already supports this)

Motivation

Some resized LoRAs have a different rank for each adapter module. Currently, only the ranks of the out_proj attention modules are gathered in the rank dictionary. Using a LoraConfig with missing rank_patterns result in the following errors:

RuntimeError: Error(s) in loading state_dict for CLIPTextModel:
        size mismatch for text_model.encoder.layers.0.self_attn.k_proj.lora_A.default_0.weight: copying a param with shape torch.Size([8, 768]) from checkpoint, the shape in current model is torch.Size([3, 768]).
        size mismatch for text_model.encoder.layers.0.self_attn.k_proj.lora_B.default_0.weight: copying a param with shape torch.Size([768, 8]) from checkpoint, the shape in current model is torch.Size([768, 3]).

UNet supports missing adapter modules, but the text encoder doesn't, resulting in:

  File "/home/elias/repos/diffusers/src/diffusers/loaders/lora.py", line 1363, in load_lora_weights
    self.load_lora_into_text_encoder(
  File "/home/elias/repos/diffusers/src/diffusers/loaders/lora.py", line 572, in load_lora_into_text_encoder
    rank[rank_key] = text_encoder_lora_state_dict[rank_key].shape[1]
                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^
KeyError: 'text_model.encoder.layers.11.self_attn.out_proj.lora_B.weight'

encoder.layers.11.*.lora_B.weight is often all zeroes when the layer is skipped during fine-tuning, in this case it might be desirable to not store them in LoRA safetensors.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? I could'nt find an open issue or forum thread about the problems
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul is the author of the lines affected by this PR.

* In addition to out_proj, read the ranks of adapters for q_proj, k_proj, and v_proj * Allow missing adapters (UNet already supports this)

sayakpaul · 2024-05-30T23:17:59Z

Thanks for your contributions!

Could you also add a test to ensure the robustness of the changes?

HuggingFaceDocBuilderDev · 2024-05-30T23:25:42Z

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.

younesbelkada

Clean, thanks !

elias-gaeros · 2024-05-31T14:47:24Z

Thanks for your contributions!

Could you also add a test to ensure the robustness of the changes?

@sayakpaul I may need some guidance for writing the tests. Testing this change requires a specially resized LoRA.
What about:

Using test_sdxl_1_0_lora as a template,
~~Making a PR to hf-internal-testing/sdxl-1.0-lora with sd_xl_offset_example-resized-lora_1.0.safetensors? or to different repository?~~ Actually no, it doesn't have LoRA for the TE layers.

Alternatively the test could drop layers and ranks in the state_dict, but the output will be nonsensical without orthogonalization.

sayakpaul · 2024-06-01T05:59:03Z

Sure.

@sayakpaul I may need some guidance for writing the tests. Testing this change requires a specially resized LoRA.

So, we could add a utility to create a peft LoRA config that tests this use case specifically. And then add an actual test case for it. We could add the utility and the test case to this file.

WDYT?

elias-gaeros · 2024-06-06T14:36:10Z

So, we could add a utility to create a peft LoRA config that tests this use case specifically. And then add an actual test case for it. We could add the utility and the test case to this file.

WDYT?

I added a test that takes inspiration from test_simple_inference_with_text_unet_lora_save_load since it's is one of the few tests that actually use LoraLoaderMixin.load_lora_into_text_encoder.

A LoraConfig is used for assigning different ranks to each projection layer. After initializing the adapters, the state_dict is extracted and half of the layers are discarded before reloading them using load_lora_weights.

sayakpaul · 2024-06-06T15:07:30Z

tests/lora/utils.py

+            # Discard half of the adapters.
+            rng = np.random.default_rng(0)
+            key2adapters = {k: k.rsplit(".", 2)[0] for k in state_dict.keys()}
+            adapters = list(set(key2adapters.values()))
+            adapters = set(rng.choice(adapters, size=len(adapters) // 2, replace=False))
+            state_dict = {k: state_dict[k] for k, adapter in key2adapters.items() if adapter in adapters}


I think we should keep this behavior truly deterministic, i.e., don't rely on rng.choice() and just manually remove the adapters. I prefer this to be a slightly better for testing.

I also think that the test should:

Run a single inference pass after the LoRAs have been added to the text encoder. Then we compare the outputs to the non-LoRA case.

Run another inference pass with adapters discarded and compare the outputs to the original LoRA outputs.

I updated the test, now removing all adapters from text_model.encoder.layers.4.

sayakpaul

Thanks for the tests. I left some comments.

…rministic

sayakpaul

Just a minor comment. Rest looks really good to me. Thanks for taking care of it!

tests/lora/utils.py

sayakpaul · 2024-06-08T02:30:34Z

Thanks so much for the iterations! Will merge once the CI is green.

elias-gaeros · 2024-06-16T16:46:16Z

What is keeping this from being merged?

sayakpaul · 2024-06-16T16:58:18Z

That will be on me. I forgot to merge it. Let the CI run once again and I will merge.

sayakpaul · 2024-06-16T17:00:01Z

Ah looks like need to run a couple of quality related formatting. Keeping it open and will take care of it once I am back to my keyboard.

sayakpaul · 2024-06-16T19:08:46Z

I am unable to locally reproduce the LoRA related test failures :/ Seems like some NumPy version mismatch bug. Can quickly open a PR to pin the NumPy version to be down under 2 so that we can unblock the PRs. @yiyixuxu @DN6 WDYT?

sayakpaul · 2024-06-18T20:11:14Z

The failing test is completely unrelated and sorry for the delay on my end. Thank you for your amazing work!

* [LoRA] text encoder: read the ranks for all the attn modules * In addition to out_proj, read the ranks of adapters for q_proj, k_proj, and v_proj * Allow missing adapters (UNet already supports this) * ruff format loaders.lora * [LoRA] add tests for partial text encoders LoRAs * [LoRA] update test_simple_inference_with_partial_text_lora to be deterministic * [LoRA] comment justifying test_simple_inference_with_partial_text_lora * style --------- Co-authored-by: Sayak Paul <[email protected]>

[LoRA] text encoder: read the ranks for all the attn modules

01c5cf2

* In addition to out_proj, read the ranks of adapters for q_proj, k_proj, and v_proj * Allow missing adapters (UNet already supports this)

yiyixuxu requested a review from sayakpaul May 30, 2024 23:08

sayakpaul requested a review from younesbelkada May 30, 2024 23:15

younesbelkada approved these changes May 31, 2024

View reviewed changes

elias-gaeros and others added 2 commits May 31, 2024 19:06

ruff format loaders.lora

cb5bff3

Merge branch 'main' into lora_te_read_all_ranks

57f3fa7

[LoRA] add tests for partial text encoders LoRAs

a74b84c

elias-gaeros force-pushed the lora_te_read_all_ranks branch from 8dae060 to a74b84c Compare June 6, 2024 14:42

sayakpaul reviewed Jun 6, 2024

View reviewed changes

sayakpaul and others added 3 commits June 6, 2024 20:41

Merge branch 'main' into lora_te_read_all_ranks

535d40d

[LoRA] update test_simple_inference_with_partial_text_lora to be dete…

eabe622

…rministic

Merge branch 'main' into lora_te_read_all_ranks

bb806e6

sayakpaul approved these changes Jun 7, 2024

View reviewed changes

tests/lora/utils.py Show resolved Hide resolved

elias-gaeros and others added 2 commits June 7, 2024 16:55

[LoRA] comment justifying test_simple_inference_with_partial_text_lora

f9a4375

Merge branch 'main' into lora_te_read_all_ranks

d4be978

Merge branch 'main' into lora_te_read_all_ranks

ce708e0

sayakpaul added 2 commits June 16, 2024 19:29

style

8c3c264

Merge branch 'main' into lora_te_read_all_ranks

ff4f4e6

Merge branch 'main' into lora_te_read_all_ranks

84c2e72

sayakpaul merged commit 298ce67 into huggingface:main Jun 18, 2024
14 of 15 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[LoRA] text encoder: read the ranks for all the attn modules #8324

[LoRA] text encoder: read the ranks for all the attn modules #8324

elias-gaeros commented May 30, 2024 •

edited

Loading

sayakpaul commented May 30, 2024

HuggingFaceDocBuilderDev commented May 30, 2024

younesbelkada left a comment

elias-gaeros commented May 31, 2024 •

edited

Loading

sayakpaul commented Jun 1, 2024

elias-gaeros commented Jun 6, 2024

sayakpaul Jun 6, 2024

sayakpaul Jun 6, 2024 •

edited

Loading

elias-gaeros Jun 6, 2024

sayakpaul left a comment

sayakpaul left a comment

sayakpaul commented Jun 8, 2024

elias-gaeros commented Jun 16, 2024

sayakpaul commented Jun 16, 2024

sayakpaul commented Jun 16, 2024

sayakpaul commented Jun 16, 2024

sayakpaul commented Jun 18, 2024 •

edited

Loading

[LoRA] text encoder: read the ranks for all the attn modules #8324

[LoRA] text encoder: read the ranks for all the attn modules #8324

Conversation

elias-gaeros commented May 30, 2024 • edited Loading

What does this PR do?

Motivation

Before submitting

Who can review?

sayakpaul commented May 30, 2024

HuggingFaceDocBuilderDev commented May 30, 2024

younesbelkada left a comment

Choose a reason for hiding this comment

elias-gaeros commented May 31, 2024 • edited Loading

sayakpaul commented Jun 1, 2024

elias-gaeros commented Jun 6, 2024

sayakpaul Jun 6, 2024

Choose a reason for hiding this comment

sayakpaul Jun 6, 2024 • edited Loading

Choose a reason for hiding this comment

elias-gaeros Jun 6, 2024

Choose a reason for hiding this comment

sayakpaul left a comment

Choose a reason for hiding this comment

sayakpaul left a comment

Choose a reason for hiding this comment

sayakpaul commented Jun 8, 2024

elias-gaeros commented Jun 16, 2024

sayakpaul commented Jun 16, 2024

sayakpaul commented Jun 16, 2024

sayakpaul commented Jun 16, 2024

sayakpaul commented Jun 18, 2024 • edited Loading

elias-gaeros commented May 30, 2024 •

edited

Loading

elias-gaeros commented May 31, 2024 •

edited

Loading

sayakpaul Jun 6, 2024 •

edited

Loading

sayakpaul commented Jun 18, 2024 •

edited

Loading