Add PEFT to advanced training script #6294

apolinario · 2023-12-22T13:31:24Z

What does this PR do?

Adds PEFT to the advanced training script.

Some questions are open still regarding PEFT integration and whether we should change more things in the script to accommodate/adapt the textual inversion training that takes place in this script:

Should we change something in the way we freeze all parameters except for the token embeddings for textual inversion here: https://github.com/huggingface/diffusers/blob/add-peft-to-advanced-training-script/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py#L1259?
- Related: would this filtering of trainable parameters work if all but token embeddings params were frozen in the way we do it in the previous step? https://github.com/huggingface/diffusers/blob/add-peft-to-advanced-training-script/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py#L1354
For pivotal tuning, we "Pivot Halfway" meaning that we can stop the textual inversion training at a % of the steps. I don't see how PEFT affects that but flagging in case someone sees something there: https://github.com/huggingface/diffusers/blob/add-peft-to-advanced-training-script/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py#L1678

Closes #6118

HuggingFaceDocBuilderDev · 2023-12-22T13:42:37Z

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

Thanks, it looks overall good, I left one open question about handling additional unfrozen params

younesbelkada · 2023-12-22T15:40:00Z

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

-        text_lora_parameters_two = LoraLoaderMixin._modify_text_encoder(
-            text_encoder_two, dtype=torch.float32, rank=args.rank
+        text_lora_config = LoraConfig(
+            r=args.rank, init_lora_weights="gaussian", target_modules=["q_proj", "k_proj", "v_proj", "out_proj"]


This should be correct with respect to the alpha issue we discussed offline @sayakpaul right? Related: #6225

Thanks for pointing to that PR! Adding the lora alpha fixed the problem and now the script is working! (before it was giving jumbled results)

without alpha

with alpha

Holy cow 🐮

Amazing! 🤩

younesbelkada · 2023-12-22T15:41:05Z

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py


        if args.train_text_encoder:
            text_encoder_one = accelerator.unwrap_model(text_encoder_one)
-            text_encoder_lora_layers = text_encoder_lora_state_dict(text_encoder_one.to(torch.float32))
+            text_encoder_lora_layers = get_peft_model_state_dict(text_encoder_one.to(torch.float32))


If you train extra parameters that you keep unfrozen for text encoder, you need to add them in modules_to_save=[xxx] when defining the lora config for the text encoder

Got it! But I'm not training extra parameters with that operation. With args.train_text_encoder we're doing regular text encoder training.

However, with args.train_text_encoder_ti (that is mutually exclusive with args.train_text_encoder) then the goal is to freeze all but the token_embedding of the text encoder and train the text embeddings for the new tokens introduced in the model.

This is where is taking place - and was working prior to adding PEFT elsewhere: https://github.com/huggingface/diffusers/pull/6294/files#diff-24abe8b0339a563b68e03c979ee9e498ab7c49f3fd749ffb784156f4e2d54d90R1249

So, only thing that is outside of the peft paradigm currently is what's happening when args.train_text_encoder_ti is True, yeah?

Yes, which makes sense because it is not training an adapter per se

Thanks @apolinario for explaining ! makes sense

sayakpaul

Fantastic work.

To answer your questions, here's what I think:

https://github.com/huggingface/diffusers/blob/add-peft-to-advanced-training-script/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py#L1259

I think to be on the safe side, we could tackle the token embedding part after we're done handling train_text_encoder. Which is what is happening now. So, that's good. I would maybe remove the parameter upcasting part because we're already doing in later in the script.

Also, args.train_text_encoder_ti and args.train_text_encoder ==> can both be set to True?

For the sub-question, I think you meant this line:

diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

Line 1361 in 38aece9

    
           text_lora_parameters_one = list(filter(lambda p: p.requires_grad, text_encoder_one.parameters()))

Given that's true, yeah I think this step should be able to pick that up. But just to be sure, I'd maybe right out print the param names.

For pivotal tuning, we "Pivot Halfway" meaning that we can stop the textual inversion training at a % of the steps. I don't see how PEFT affects that but flagging in case someone sees something there:

No, it's should affect anything related to peft.

sayakpaul · 2023-12-23T03:05:53Z

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

@@ -37,6 +37,8 @@
 from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
 from huggingface_hub import create_repo, upload_folder
 from packaging import version
+from peft import LoraConfig


Will need to peft as a dependency in the requirements.txt.

sayakpaul · 2023-12-23T03:06:16Z

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

-        unet_lora_parameters.extend(attn_module.to_out[0].lora_layer.parameters())
+    unet_lora_config = LoraConfig(
+        r=args.rank,
+        lora_alpha=args.rank,


Very important!

sayakpaul · 2023-12-23T03:06:39Z

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

+    # Make sure the trainable params are in float32.
+    if args.mixed_precision == "fp16":
+        models = [unet]
+        if args.train_text_encoder:
+            models.extend([text_encoder_one, text_encoder_two])
+        for model in models:
+            for param in model.parameters():
+                # only upcast trainable parameters (LoRA) into fp32
+                if param.requires_grad:
+                    param.data = param.to(torch.float32)


Another important one!

sayakpaul · 2023-12-23T03:12:51Z

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py


        if args.train_text_encoder:
            text_encoder_one = accelerator.unwrap_model(text_encoder_one)
-            text_encoder_lora_layers = text_encoder_lora_state_dict(text_encoder_one.to(torch.float32))
+            text_encoder_lora_layers = get_peft_model_state_dict(text_encoder_one.to(torch.float32))


So, only thing that is outside of the peft paradigm currently is what's happening when args.train_text_encoder_ti is True, yeah?

apolinario · 2023-12-23T06:19:55Z

Also, args.train_text_encoder_ti and args.train_text_encoder ==> can both be set to True?

Not yet. We plan supporting training both the text encoder and textual inversion - but that's for a future version of the script

Apply #6355 fix

* Fix ProdigyOPT in SDXL Dreambooth script * style * style * Add PEFT to Advanced Training Script * style * style * ✨ style ✨ * change order for logic operation * add lora alpha * style * Align PEFT to new format * Update train_dreambooth_lora_sdxl_advanced.py Apply huggingface#6355 fix --------- Co-authored-by: multimodalart <[email protected]>

apolinario added 4 commits December 22, 2023 05:55

Fix ProdigyOPT in SDXL Dreambooth script

6338ad5

style

565416c

style

3338ce0

Add PEFT to Advanced Training Script

d62076a

apolinario requested a review from younesbelkada December 22, 2023 13:31

Merge branch 'main' into add-peft-to-advanced-training-script

965b40a

apolinario added 3 commits December 22, 2023 07:42

style

9b910bd

style

566aaab

✨ style ✨

a837033

younesbelkada reviewed Dec 22, 2023

View reviewed changes

apolinario added 3 commits December 22, 2023 13:09

change order for logic operation

2bfdcab

add lora alpha

b03aa10

style

38aece9

apolinario marked this pull request as ready for review December 23, 2023 01:53

sayakpaul reviewed Dec 23, 2023

View reviewed changes

sayakpaul mentioned this pull request Dec 25, 2023

[bug fixes - advanced sdxl lora script] - add left lora changes from canonical script and fix bug in snr gamma #6327

Closed

Align PEFT to new format

daa7566

apolinario mentioned this pull request Dec 27, 2023

🛠️ Fix widget format for advanced training #6355

Closed

multimodalart and others added 2 commits December 27, 2023 07:54

Merge branch 'main' into add-peft-to-advanced-training-script

f6844d3

Update train_dreambooth_lora_sdxl_advanced.py

0f9427f

Apply #6355 fix

apolinario merged commit 645a62b into main Dec 27, 2023
16 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add PEFT to advanced training script #6294

Add PEFT to advanced training script #6294

apolinario commented Dec 22, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Dec 22, 2023

younesbelkada left a comment

younesbelkada Dec 22, 2023

apolinario Dec 23, 2023

apolinario Dec 23, 2023 •

edited

Loading

sayakpaul Dec 23, 2023

younesbelkada Dec 25, 2023

younesbelkada Dec 22, 2023

apolinario Dec 22, 2023 •

edited

Loading

sayakpaul Dec 23, 2023

apolinario Dec 23, 2023

younesbelkada Dec 25, 2023

sayakpaul left a comment

sayakpaul Dec 23, 2023

sayakpaul Dec 23, 2023

sayakpaul Dec 23, 2023

sayakpaul Dec 23, 2023

apolinario commented Dec 23, 2023

Add PEFT to advanced training script #6294

Add PEFT to advanced training script #6294

Conversation

apolinario commented Dec 22, 2023 • edited Loading

What does this PR do?

HuggingFaceDocBuilderDev commented Dec 22, 2023

younesbelkada left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apolinario Dec 23, 2023 • edited Loading

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

apolinario Dec 22, 2023 • edited Loading

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

sayakpaul 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

apolinario commented Dec 23, 2023

apolinario commented Dec 22, 2023 •

edited

Loading

apolinario Dec 23, 2023 •

edited

Loading

apolinario Dec 22, 2023 •

edited

Loading