Gradient Accumulation - Multi-GPU training of SDXL-LoRA

[Dvirbeno]

Describe the bug

I'm not quite sure whether it's a bug or a discussion -
Looking at the example scripts for dreambooth training (SDXL-LoRA, or the more "traditional" training script), I see that the accumulation of gradients, via accelerate, in the beginning of the training loop, is set the accumulate gradients only for the unet part of the model:

  for step, batch in enumerate(train_dataloader):
      with accelerator.accumulate(unet):
          pixel_values = batch["pixel_values"].to(dtype=weight_dtype)

(This part is somewhat similar across this set of scripts)

I was wondering, is the text_encoder (+ text_encoder_2) intentionally left outside the accumulation? (I see accumulation of multiple modules is supported according to this PR: huggingface/accelerate#1708)
What will happen, in terms of gradient accumulation, when we'll use multiple GPUs for training?

It might be related to this issue:
#4046
and might be the solution to this issue:
#5012

Reproduction

As noted, I would like to understand if it's a familiar limitation, and if including the text-encoders in the accumulation is expected to "damage"/"break" something in the training procedure.

the reproduction is demonstrated in the issue
#5012

Logs

No response

System Info

Diffusers main branch
(https://github.com/huggingface/diffusers/tree/e7c0af5e71b50faf1f7af32e9e7827b87c905eae)
The discussed snippet can be found on the scripts under examples/dreambooth:
https://github.com/huggingface/diffusers/tree/main/examples/dreambooth

Who can help?

@sayakpaul @PATR

[sayakpaul]

We have rarely trained the text encoder in multi-GPU settings actually. Most of our large-scale training runs have only involved the UNet or the denoiser network for that matter.

I am not quite an expert on this topic. So, I will let the accelerate folks chime in here :D

Ccing @muellerzr

[muellerzr]

If it has gradients/being trained on, add it to accumulate :)

[Dvirbeno]

Thanks @muellerzr & @sayakpaul

So, if that's the case then there's a resolution that will enable fine-tuning these models using multiple GPUs? (Existing limitation according to: #5012 (comment))

But, @muellerzr, in the scenario of multiple GPUs, what happens practically when some of the modules (unet) are included within the accumulate, and some are not (text_encoder)?
How are the gradients aggregated for the other modules in such case?

In addition, even when we use a single GPU, and we configure accelerate to, say 4, gradient accumulation steps -
if the accumulate context wraps only a subset of the trained modules, is there any gradient accumulation for the modules that are not included as part of the accumulate clause?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[Dvirbeno]

Rephrasing and listing my questions:

1. When the training includes an accumulate context manager, which wraps only part of the modules that are learned as part of the loop:
   a. When using gradient_accumulation_steps > 1 (even on a single GPU), how will that affect the modules that are not included in the accumulation clause?
   b. When using multiple GPUs, how are the gradients aggregated for the modules not included in the accumulation?

2. If, according to the changes in accelerate we can include every trained module as part of the accumulate context manager - doesn't that resolve and enable fine-tuning (e.g. SDXL) using multi-GPU? (Existing limitation according to:LoRA can not train text encoder in "train_text_to_image_lora_sdxl.py"? #5012 (comment))?

Thanks in advance for your answers!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

Gently pinging @muellerzr for the further questions.

[muellerzr]

Thanks! Sorry I missed this @Dvirbeno

When the training includes an accumulate context manager, which wraps only part of the modules that are learned as part of the loop:
a. When using gradient_accumulation_steps > 1 (even on a single GPU), how will that affect the modules that are not included in the accumulation clause?

The gradients will be accumulated, but if you're on a multi-GPU system you'll see some pretty high slowdowns. See here for more information as to what's going on: https://muellerzr.github.io/blog/gradient_accumulation.html

b. When using multiple GPUs, how are the gradients aggregated for the modules not included in the accumulation?

Same as vanilla pytorch, the key here is communication bandwidth.

If, according to the changes in accelerate we can include every trained module as part of the accumulate context manager - doesn't that resolve and enable fine-tuning (e.g. SDXL) using multi-GPU?

Yes!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

Closing this due to inactivity. Feel free to reopen.