dreambooth advanced: embedding retraction happening twice

[rvorias]

Describe the bug

diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

Line 1868 in 774f5c4

for idx, text_encoder in enumerate(text_encoders):

The embedding handler already takes care of accounting for both text encoders, no need to loop over this.
(You instantiate it with both text encoders and it will loop over them when this retract is called.)

Reproduction

Logs

System Info

Who can help?

[rvorias]

Another bug:

diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

Line 1280 in 774f5c4

param = param.to(dtype=torch.float32)

This creates a new param, which means std ema'ing in the EmbeddingHandler doesn't work.
You are smoothing (std ema) one parameter and training another.

[rvorias]

3rd bug:
Pivot halfway learning rate zeroing only has effect for one step.
I suspect this is because of accelerate.

This actually pivots correctly

pivoted = False
...
    for epoch in range(first_epoch, args.num_train_epochs):
        # if performing any kind of optimization of text_encoder params
        if args.train_text_encoder or args.train_text_encoder_ti:
            if epoch == num_train_epochs_text_encoder:
                print("PIVOT HALFWAY", epoch)
                pivoted = True

        unet.train()
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(unet):
                if pivoted:
                    optimizer.param_groups[1]["lr"] = 0.0
                    optimizer.param_groups[2]["lr"] = 0.0

Then, also only call retract when pivoted = False, because otherwise you are ema'ing the std and manipulating the embedding vectors.

[DN6]

@linoytsaban @apolinario Would either of you be able to take a look here?

[linoytsaban]

Thanks for opening an issue and pointing this out! Hopefully all of these issues are now solved in #6599