Enabling gradient checkpointing for VAE

[aandyw]

PR for issue #1883

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Let's try to only add gradient checkpointing for the blocks - only there we can really save memory. It doesn't help much to add it to individual layers, such as self.conv_in

@williamberman could you also have a look?

[aandyw]

Let's try to only add gradient checkpointing for the blocks - only there we can really save memory. It doesn't help much to add it to individual layers, such as self.conv_in

@williamberman could you also have a look?

Got it. Thanks for the feedback.

[patrickvonplaten]

Note that gradient checkpointing doesn't necessarily mean that the model is training

[aandyw]

How exactly is fine-tuning the VAE done in train_text_to_image?

It seems like the issue #1883 wants to enable gradient checkpointing specifically for training VAE independently so should we add in some sort of argument?

[patrickvonplaten]

The PR looks in a much better shape now, thanks!

Could you try to revert the reformatting of unrelated functions and then let's make sure to always freeze the unet and vae of train_text_to_image just like before. Note that gradient checkpointing still helps even if the weights are frozen :-)

Good job! Think we're on the final stretch now

[aandyw]

The PR looks in a much better shape now, thanks!

Could you try to revert the reformatting of unrelated functions and then let's make sure to always freeze the unet and vae of train_text_to_image just like before. Note that gradient checkpointing still helps even if the weights are frozen :-)

Good job! Think we're on the final stretch now

Thanks for the feedback! I really appreciate it.

Apologies on the formatting. I'm using black and it seems to really want these functions formatted this way. I might have missed this but is there a suggested linter/formatter to use?

[patrickvonplaten]

@pie31415 could you try to use the newest black version, i.e.:

diffusers/setup.py

Line 84 in 4ae54b3

"black~=23.1",

- does this correspond to the version you have? :-)

[aandyw]

Fixed the issue with formatting. It was because of issues with line endings being CRLF.

However, some of the code formatting is still because of black (formatting on save).
e.g.

tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]

becomes

tile = z[
  :,
  :,
  i : i + self.tile_latent_min_size,
  j : j + self.tile_latent_min_size,
 ]

[patrickvonplaten]

Suggested change

vae.requires_grad_(True)

Think even without requires_grad_ this should help as gradients are passed down

[patrickvonplaten]

Cool this now looks good to me!

Last thing we should do is to add a simple test that makes sure that gradient checkpointing works. Maybe you can get some inspiration from:

diffusers/tests/models/test_models_unet_2d_condition.py

Line 125 in 07a0c1c

def test_gradient_checkpointing(self):

and add a similar test here:
https://github.com/huggingface/diffusers/blob/main/tests/models/test_models_vae.py

After that we can merge I think :-)

[williamberman]

wow super cool :) I learned how our gradient checkpointing works from this PR 😁 will approve once test is written

[aandyw]

Seem to be failing CI/CD tests for MacOS. Anyway to fix this atm?

[williamberman]

macos tests are unrelated, looks great

[patrickvonplaten]

Good job @pie31415