Handling mixed precision for dreambooth flux lora training

[icsl-Jeon]

What does this PR do?

Hello 😄 Thank you for the awesome example!
Here, I want to make a PR that helped me train dreambooth LoRA successfully.

• Avoid dtype change for transfermer after log_validation (especially for fp16). For mixed training, the original code upcast fp16 to fp32 for mixed precision training. However, after switching pipeline dtype in log_validation, transformer dtype returns to fp16, which can lead to fp16 unscaling error. Actually, I had this problem when I use fp16 option. (For some reason, T5 yielded nan output in bf16, that is why I came to use fp16) train_dreambooth_lora_flux validation RuntimeError: Input type (float) and bias type (c10::BFloat16) should be the same #9476
• Use the same preset dtype for text encoders for validation. For some reason, the two text encoders were in fp32 in the pipeline.
• Add unwarp to access the config field of transformer. train_dreambooth_lora_flux.py distributed bugs #9161 (comment)

Fixes # (issue)

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? Please add a link to it if that's the case.
• 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?

linoytsaban @sayakpaul

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]

Thank you! Just a single comment.

@linoytsaban could you also give this a look?

[sayakpaul]

Why are we doing it?

We should keep the pipeline model-level components (such as text encoders, VAE, etc.) to a reduced precision no?

[icsl-Jeon]

text encoders, vae are already in reduced precision :)
As I described in the PR description, this will change dtype of transformers
For mixed precision training, transformer was upcast into fp32 if fp16 training.
But this changes back to fp16, which leads to fp16 unscale error in clip gradient.

[sayakpaul]

Would something like this work?
#9549 (comment)

[icsl-Jeon]

Thank you for the suggestion! But, in this thread, I was interested in unwanted switch of fp32 into fp16 after validation, not in the computation of T5 :)

[sayakpaul]

Ah okay. Can you provide an example command for us to verify this? Maybe @linoytsaban could give it a try?

[sayakpaul]

@icsl-Jeon a friendly reminder :)

[icsl-Jeon]

This could be reproduced with any launch commands in the README.
accelerate launch ... --mixed_precision="fp16" ..

I checked the lora precision with

            for name, param in transformer.named_parameters():
                if 'lora' in name:
                    print(f"Layer: {name}, dtype: {param.dtype}, requires_grad: {param.requires_grad}")

Hope this help you reproduce!

[sayakpaul]

Avoid dtype change for transfermer after log_validation (especially for fp16). For mixed training, the original code upcast fp16 to fp32 for mixed precision training. However, after switching pipeline dtype in log_validation, transformer dtype returns to fp16, which can lead to #6442 Actually, I had this problem when I use fp16 option. (For some reason, T5 yielded nan output in bf16, that is why I came to use fp16)

We could avoid it by running inference in autocast, no? Here's an example:

diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py

Line 215 in 61d3764

autocast_ctx = torch.autocast(accelerator.device.type)

[HuggingFaceDocBuilderDev]

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.

[sayakpaul]

LGTM actually!

@linoytsaban could you also review?

[icsl-Jeon]

@linoytsaban thank you in advance

[linoytsaban]

Thanks @icsl-Jeon, LGTM!

[icsl-Jeon]

Is there any action to be done for merge ?

[sayakpaul]

Will merge this once the CI run is complete. Thanks a ton!

[sayakpaul]

Thanks for your contributions!