[PEFT] Allow PEFT model dict to be loaded #25721

patrickvonplaten · 2023-08-24T09:38:47Z

In order to allow peft to be leveraged in diffusers without breaking changes we need to allow loading adapters directly from a loaded state_dict. The reason is that in diffusers we currently store LoRA checkpoints in a format that is different to the PEFT format so we cannot just pass the model_id. This PR allows the user to manually pass a loaded PEFT model checkpoint as well as a PEFT configuration, thus circumventing the need to pass a model id.

In pseudo code, the integration of transformers + PEFT in diffusers should then look as follows the "load_lora" function of diffusers.

def load_lora_into_text_encoder(cls, state_dict, network_alphas, text_encoder, prefix=None, lora_scale=1.0):
    peft_state_dict, peft_config = convert_to_peft_format(state_dict, ...)  # <- this function will take care of all the remapping necessary for the different formats
    text_encoder.load_adapter(peft_state_dict, peft_config=peft_config)

Note, there might be more changes we have to do to PEFT, transformers' PEFT integration to be sure that everything works as expected. E.g. I'm not yet sure how to pass network_alphas etc... to PEFT to make sure we get 1-to-1 the same result.

cc @younesbelkada @sayakpaul @BenjaminBossan @pacman100

src/transformers/lib_integrations/peft/peft_mixin.py

BenjaminBossan

Thanks for the PR.

I'm not very knowledgeable about diffusers, so will let others comment on the overall solution. My comments are just minor issues, not blockers.

BenjaminBossan · 2023-08-24T10:04:37Z

src/transformers/lib_integrations/peft/peft_mixin.py

@@ -28,6 +29,10 @@
    from accelerate.utils import get_balanced_memory, infer_auto_device_map


+if is_torch_available():


Why is this necessary?

BenjaminBossan · 2023-08-24T10:06:48Z

src/transformers/lib_integrations/peft/peft_mixin.py

@@ -59,14 +64,15 @@ class PeftAdapterMixin:

    def load_adapter(
        self,
-        peft_model_id: str,
+        peft_model_id: Union[str, Dict[str, "torch.Tensor"]],


Should we choose a different name, now that a state dict can be passed? Alternatively, we could add another (optional) argument to pass the state dict, make peft_model_id optional, and do a check that exactly one of the two should be passed.

Yes makes sense to me! peft_model_id_or_state_dict sounds great - was just wondering about breaking change here. But think this function is not yet in a release so should probs be fine to change no?

BenjaminBossan · 2023-08-24T10:07:03Z

src/transformers/lib_integrations/peft/peft_mixin.py

@@ -75,7 +81,7 @@ def load_adapter(
        Requires peft as a backend to load the adapter weights.

        Args:
-            peft_model_id (`str`):
+            peft_model_id (`str` or dictionary of `torch.Tensor`):


The description below should also be adjusted.

BenjaminBossan · 2023-08-24T10:07:25Z

src/transformers/lib_integrations/peft/peft_mixin.py

        adapter_name: Optional[str] = None,
        revision: Optional[str] = None,
        token: Optional[str] = None,
        device_map: Optional[str] = "auto",
        max_memory: Optional[str] = None,
        offload_folder: Optional[str] = None,
        offload_index: Optional[int] = None,
+        peft_config: Dict[str, Any] = None,


Suggested change

peft_config: Dict[str, Any] = None,

peft_config: Optional[Dict[str, Any]] = None,

HuggingFaceDocBuilderDev · 2023-08-24T10:15:38Z

The documentation is not available anymore as the PR was closed or merged.

younesbelkada

Looks great! agreed also with @BenjaminBossan 's comments, I think we should maybe add an extra check and raise a proper error, for the naming, what about peft_model_id_or_state_dict (maybe that's too long) ?

younesbelkada · 2023-08-24T10:16:43Z

src/transformers/lib_integrations/peft/peft_mixin.py

-            raise ValueError(
-                f"adapter model file not found in {peft_model_id}. Make sure you are passing the correct path to the "
-                "adapter model."
+        if peft_config is None:


Suggested change

if peft_config is None:

if peft_config is None and isinstance(peft_model_id, str):

and add a check below that if peft_model_id is not a state dict, raise an error

pacman100

Thank you @patrickvonplaten for adding the support for passing state dict and config, LGTM!

younesbelkada

Thanks again !

ArthurZucker

Have one question regarding the peft_model_id being a dictionary, otherwise looks good

ArthurZucker · 2023-09-15T13:11:02Z

src/transformers/integrations/peft.py

+            peft_model_id (`str` or dictionary of `torch.Tensor`):
                The identifier of the model to look for on the Hub, or a local path to the saved adapter config file
                and adapter weights.


Doc needs to be updated if this can be a state_dict. The name of the argument is not intuitive but I guess wee need to keep it for Backward compatibility? Otherwise would rather have a new arg

I think for BC it's fine as the feature is quite new

Not really if it was part of the release 😅

src/transformers/integrations/peft.py

ArthurZucker

LGTM thanks for iterating and adding a test!

tests/peft_integration/test_peft_integration.py

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

* Allow PEFT model dict to be loaded * make style * make style * Apply suggestions from code review * address comments * fixup * final change * added tests * fix test * better logic for handling if adapter has been loaded * Update tests/peft_integration/test_peft_integration.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --------- Co-authored-by: younesbelkada <younesbelkada@gmail.com> Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

Allow PEFT model dict to be loaded

032bde0

patrickvonplaten force-pushed the allow_peft_state_dict_loading branch from c9e7fb8 to 032bde0 Compare August 24, 2023 09:49

patrickvonplaten added 2 commits August 24, 2023 09:49

make style

4f4deb7

make style

a26a7de

patrickvonplaten commented Aug 24, 2023

View reviewed changes

src/transformers/lib_integrations/peft/peft_mixin.py Outdated Show resolved Hide resolved

Apply suggestions from code review

15cf855

BenjaminBossan reviewed Aug 24, 2023

View reviewed changes

younesbelkada approved these changes Aug 24, 2023

View reviewed changes

pacman100 reviewed Aug 25, 2023

View reviewed changes

pacman100 approved these changes Aug 25, 2023

View reviewed changes

Merge remote-tracking branch 'upstream/main' into HEAD

e8304c8

younesbelkada mentioned this pull request Aug 25, 2023

[WIP][DRAFT] PEFT integration huggingface/diffusers#4780

Closed

8 tasks

Merge remote-tracking branch 'upstream/main' into HEAD

13ea15a

younesbelkada approved these changes Sep 15, 2023

View reviewed changes

younesbelkada changed the title ~~[WIP][PEFT] Allow PEFT model dict to be loaded~~ [PEFT] Allow PEFT model dict to be loaded Sep 15, 2023

younesbelkada requested a review from ArthurZucker September 15, 2023 12:51

ArthurZucker reviewed Sep 15, 2023

View reviewed changes

address comments

55d657b

younesbelkada requested a review from ArthurZucker September 15, 2023 13:27

younesbelkada added 5 commits September 15, 2023 13:33

fixup

13ca0ad

final change

d05ef41

added tests

3e4f816

fix test

0bcd6fa

better logic for handling if adapter has been loaded

123666b

ArthurZucker approved these changes Sep 15, 2023

View reviewed changes

tests/peft_integration/test_peft_integration.py Outdated Show resolved Hide resolved

Update tests/peft_integration/test_peft_integration.py

b47dd15

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

younesbelkada merged commit 0a55d9f into main Sep 15, 2023

younesbelkada deleted the allow_peft_state_dict_loading branch September 15, 2023 16:22

This was referenced Oct 25, 2023

TypeError: PeftAdapterMixin.load_adapter() got an unexpected keyword argument 'adapter_state_dict' huggingface/diffusers#5522

Closed

[core / PEFT ]Bump transformers min version for PEFT integration huggingface/diffusers#5579

Merged

ydshieh added a commit that referenced this pull request Oct 31, 2023

run doctest 2023-09-15 [PR #25721] torch 2.1.0 + not run idefics.md

df7dddd

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PEFT] Allow PEFT model dict to be loaded #25721

[PEFT] Allow PEFT model dict to be loaded #25721

patrickvonplaten commented Aug 24, 2023 •

edited

Loading

BenjaminBossan left a comment

BenjaminBossan Aug 24, 2023

BenjaminBossan Aug 24, 2023

patrickvonplaten Aug 24, 2023

BenjaminBossan Aug 24, 2023

BenjaminBossan Aug 24, 2023

HuggingFaceDocBuilderDev commented Aug 24, 2023 •

edited

Loading

younesbelkada left a comment

younesbelkada Aug 24, 2023

pacman100 left a comment

younesbelkada left a comment

ArthurZucker left a comment

ArthurZucker Sep 15, 2023

younesbelkada Sep 15, 2023

ArthurZucker Sep 15, 2023

ArthurZucker left a comment

		@@ -28,6 +29,10 @@
		from accelerate.utils import get_balanced_memory, infer_auto_device_map


		if is_torch_available():

	peft_config: Dict[str, Any] = None,
	peft_config: Optional[Dict[str, Any]] = None,

	if peft_config is None:
	if peft_config is None and isinstance(peft_model_id, str):

[PEFT] Allow PEFT model dict to be loaded #25721

[PEFT] Allow PEFT model dict to be loaded #25721

Conversation

patrickvonplaten commented Aug 24, 2023 • edited Loading

BenjaminBossan left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented Aug 24, 2023 • edited Loading

younesbelkada left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pacman100 left a comment

Choose a reason for hiding this comment

younesbelkada left a comment

Choose a reason for hiding this comment

ArthurZucker left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ArthurZucker left a comment

Choose a reason for hiding this comment

patrickvonplaten commented Aug 24, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Aug 24, 2023 •

edited

Loading