Enable Gradient Checkpointing for UNet2DModel (New) by dg845 · Pull Request #7201 · huggingface/diffusers

dg845 · 2024-03-04T04:39:22Z

What does this PR do?

This PR supports gradient checkpointing for UNet2DModel.

Successor to PR #6718.

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?

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
@yiyixuxu

sayakpaul

Clean!

HuggingFaceDocBuilderDev · 2024-03-06T11:35:17Z

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.

github-actions · 2024-04-03T15:02:51Z

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.

yiyixuxu · 2024-12-03T04:15:53Z

sorry @dg845
don't know how I missed this PR! would you be able to resolve the conflicts and we'll merge once CI is green?

…ggingface#3675) * fix: assertion. * assertion fix.

* Fixing the global_step key not found * Apply suggestions from code review --------- Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

update

…point check and remove references to deprecated scale/lora_scale

… UNet2DConditionModel gradient checkpointing tests

dg845 · 2024-12-15T03:03:49Z

@yiyixuxu resolved the merge conflicts and updated the code to match the way gradient checkpointing is currently done in UNet2DConditionModel. The tests are passing locally for me. (Also, messed up the git history a little bit, whoops😅).

sayakpaul

Thanks a lot!

yiyixuxu · 2024-12-18T08:00:42Z

@hlky can you take a look here and help merge in?

hlky

Thanks! I've left a general comment and I'll look into the test failures, should be good to merge after CI turns green.

hlky · 2024-12-18T08:31:20Z

src/diffusers/models/unets/unet_2d.py

+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value


I notice we use hasattr(module, "gradient_checkpointing") in around half the cases of _set_gradient_checkpointing and module names in the other half, no need to change it for now but we could look at making it uniform at some point.

hlky

Thanks @dg845!

hlky · 2024-12-18T11:55:16Z

tests/models/unets/test_models_unet_2d.py

+        )
+
+    def test_effective_gradient_checkpointing(self):
+        super().test_effective_gradient_checkpointing(skip={"time_proj.weight"})


cc @sayakpaul time_proj.weight has no grad, other parameters are ok

* Port UNet2DModel gradient checkpointing code from huggingface#6718. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: Vincent Neemie <92559302+VincentNeemie@users.noreply.github.com> Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: hlky <hlky@hlky.ac>

* Port UNet2DModel gradient checkpointing code from #6718. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: Vincent Neemie <92559302+VincentNeemie@users.noreply.github.com> Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: hlky <hlky@hlky.ac>

dg845 added 2 commits March 3, 2024 20:33

Port UNet2DModel gradient checkpointing code from huggingface#6718.

3f71d82

make fix-copies

f5fd68d

dg845 mentioned this pull request Mar 4, 2024

Enable Gradient Checkpointing for UNet2DModel #6718

Closed

6 tasks

sayakpaul approved these changes Mar 6, 2024

View reviewed changes

Merge branch 'main' into unet-2d-model-grad-checkpointing3

e7229d7

github-actions bot added the stale Issues that haven't received updates label Apr 3, 2024

yiyixuxu added close-to-merge and removed stale Issues that haven't received updates labels Dec 3, 2024

sayakpaul and others added 7 commits December 14, 2024 17:42

[UniDiffuser test] fix one test so that it runs correctly on V100 (hu…

e22c6ca

…ggingface#3675) * fix: assertion. * assertion fix.

Fixing the global_step key not found (huggingface#3844)

67297d0

* Fixing the global_step key not found * Apply suggestions from code review --------- Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

More IPAdapter test fixes (huggingface#6888)

6b3782b

update

Merge branch 'main' into unet-2d-model-grad-checkpointing3

a1606b0

Use torch.is_grad_enabled instead of self.training for gradient check…

658a8de

…point check and remove references to deprecated scale/lora_scale

Update UNet2DModel gradient checkpointing tests to follow the current…

b36b6b5

… UNet2DConditionModel gradient checkpointing tests

make style and make quality

eb92a65

dg845 requested a review from sayakpaul December 15, 2024 03:03

sayakpaul approved these changes Dec 15, 2024

View reviewed changes

yiyixuxu requested a review from hlky December 18, 2024 08:00

hlky reviewed Dec 18, 2024

View reviewed changes

hlky added 2 commits December 18, 2024 11:47

Fix tests

e619ec8

-skip

5f88933

hlky approved these changes Dec 18, 2024

View reviewed changes

hlky added 2 commits December 18, 2024 12:12

Merge branch 'main' into unet-2d-model-grad-checkpointing3

6330f6a

Merge branch 'main' into unet-2d-model-grad-checkpointing3

8cdaf5a

yiyixuxu merged commit 648d968 into huggingface:main Dec 20, 2024
11 of 12 checks passed

yiyixuxu removed the close-to-merge label Dec 20, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable Gradient Checkpointing for UNet2DModel (New)#7201

Enable Gradient Checkpointing for UNet2DModel (New)#7201
yiyixuxu merged 14 commits intohuggingface:mainfrom
dg845:unet-2d-model-grad-checkpointing3

dg845 commented Mar 4, 2024

Uh oh!

sayakpaul left a comment

Uh oh!

HuggingFaceDocBuilderDev commented Mar 6, 2024

Uh oh!

github-actions bot commented Apr 3, 2024

Uh oh!

yiyixuxu commented Dec 3, 2024

Uh oh!

dg845 commented Dec 15, 2024

Uh oh!

sayakpaul left a comment

Uh oh!

yiyixuxu commented Dec 18, 2024

Uh oh!

hlky left a comment

Uh oh!

hlky Dec 18, 2024

Uh oh!

hlky left a comment

Uh oh!

hlky Dec 18, 2024

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

dg845 commented Mar 4, 2024

What does this PR do?

Before submitting

Who can review?

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Mar 6, 2024

Uh oh!

github-actions bot commented Apr 3, 2024

Uh oh!

yiyixuxu commented Dec 3, 2024

Uh oh!

dg845 commented Dec 15, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

yiyixuxu commented Dec 18, 2024

Uh oh!

hlky left a comment

Choose a reason for hiding this comment

Uh oh!

hlky Dec 18, 2024

Choose a reason for hiding this comment

Uh oh!

hlky left a comment

Choose a reason for hiding this comment

Uh oh!

hlky Dec 18, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants