SDXL dreambooth LoRA can not train text_encoders in "train_dreambooth_lora_sdxl.py"

[shyammarjit]

Describe the bug

While enabling --train_text_encoder in the train_dreambooth_lora_sdxl.py script, it initializes two text encoder parameters but its require_grad is False. Due to this, the parameters are not being backpropagated and updated.

It also shows a warning:
/home/smarjit/miniconda3/envs/diffusers/lib/python3.11/site-packages/torch/utils/checkpoint.py:31: UserWarning: None of the inputs have requires_grad=True. Gradients will be None warnings.warn("None of the inputs have requires_grad=True. Gradients will be None")

I am also attaching the output of the text_encoder LoRA parameters. Here we can see that text_model.encoder.layers.11.self_attn.v_proj.lora_linear_layer.up.weight weights are 0 as they are being initialized to 0 and also it's requires_grad is False.

Reproduction

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export OUTPUT_DIR=""
export INSTANCE_DIR=""

accelerate launch train_dreambooth_lora_sdxl.py
--pretrained_model_name_or_path=$MODEL_NAME
--instance_data_dir=$INSTANCE_DIR
--output_dir=$OUTPUT_DIR
--mixed_precision="fp16"
--instance_prompt="a photo of sks {subject}"
--resolution=1024
--train_batch_size=1
--gradient_accumulation_steps=4
--learning_rate=1e-4
--lr_scheduler="constant"
--lr_warmup_steps=0
--max_train_steps=1000
--seed="0"
--use_8bit_adam
--gradient_checkpointing
--enable_xformers_memory_efficient_attention
--train_text_encoder \

Note: Here subject is the dataset name.

Logs

No response

System Info

• diffusers version: 0.21.0.dev0
• Platform: Linux-5.4.0-155-generic-x86_64-with-glibc2.31
• Python version: 3.11.5
• PyTorch version (GPU?): 2.0.1+cu117 (True)
• Huggingface_hub version: 0.17.1
• Transformers version: 4.33.1
• Accelerate version: 0.22.0
• xFormers version: 0.0.21
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

Questions on LoRA text encoder training with SDXL: @williamberman, @patrickvonplaten, and @sayakpaul

[sayakpaul]

Slightly confused about this as we patch the LoRA-affected modules in the text encoder and pass the parameters to optimize accordingly:

diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py

Line 792 in 19edca8

text_lora_parameters_one = LoraLoaderMixin._modify_text_encoder(

So, it should not be having this behavior.

@williamberman could you comment a bit here?

[AmericanPresidentJimmyCarter]

This is due to checkpointing being enabled. See this issue for how to fix it: kohya-ss/sd-scripts#323

[shyammarjit]

After modifications from kohya-ss/sd-scripts#323 issue, I have added two additional lines in train_dreambooth_lora_sdxl.py and created a pull request on the branch shyammarjit:sdxl_text_encoder_grad

if args.train_text_encoder:
   text_encoder_one.train()
   text_encoder_two.train()
   # set top parameter requires_grad = True for gradient checkpointing works
+  text_encoder_one.text_model.embeddings.requires_grad_(True)
+  text_encoder_two.text_model.embeddings.requires_grad_(True)

[shyammarjit]

Although the above changes solved the warning issue & the gradient is also computed. At the same time, it brings up another issue for the SDXL config. In SDXL, we have two text encoders. The first text encoder has 11 layers. At each layer, we have a total of 8 LoRA matrices for KQVO (4*2_LoRA_matric). At the end layer of the first text encoder, the gradient of these 8 weight matrices is None and these weight matrices are not getting updated.

I have added line no 1170 to 1176. Here we can see that 8 weight matrices of text-encoder 1 have grad=None

Please have a look at https://github.com/shyammarjit/diffusers/tree/sdxl_text_encoder_bug

Who can help?

Questions on LoRA text encoder training with SDXL: @williamberman, @patrickvonplaten, and @sayakpaul

[sayakpaul]

Are you using a machine with multiple GPUs?

[shyammarjit]

no

[patrickvonplaten]

Is this one still open?