[Training] Make officially maintained training scripts compatible with `torch.compile()`

[sayakpaul]

We recently made the DreamBooth LoRA SDXL training script fully compatible with torch.compile() in the following PRs: #6483, #6511.

Since there might be speedup benefits from torch.compile(), it makes sense to propagate these changes to the other training scripts too. We'd appreciate help from the community in this matter.

Below is a list of the training scripts where we'd like to integrate this:

• DreamBooth LoRA
  • Regular
  • LoRA
• ControlNet
  • SD
  • SDXL
• Text-to-image
  • SD
  • SD LoRA
  • SDXL
  • SDXL LoRA
• InstructPix2Pix
  • SD
  • SDXL
• T2I Adapters
  • SDXL

#6483 and #6511 can serve as a reference.

Feel free to open PRs targeting only one example at a time and tag me :)

[gzguevara]

@sayakpaul could you please explain how PR #6483 is connected with torch.compile() ? I can see that mostly the logic for predicting the noise residual has been changed and the function encode_prompt().

[sayakpaul]

That makes torch.compile() work during training. Otherwise, it errors out. That is why it's connected. Your entry point to using torch.compile() during training is accelerate. While running accelerate config you ask it to use torch.compile(). See my accelerate configuration here.

[rootonchair]

Hi @sayakpaul , I would like to work on this issue

[sayakpaul]

Feel free to pick one from the list and submit a PR. Thank you for showing interest!

[suvadityamuk]

I'm currently doing the one for ControlNet, but wanted to just confirm the fact that we're delegating the actual trigger of torch.compile() to accelerate, right? As such, we do not invoke torch.compile() explicitly anywhere in the training script, correct? I'm following #6483 and #6511 as mentioned.

[sayakpaul]

Correct!