Make dynamo wrapped modules work with save_pretrained #2726

pcuenca · 2023-03-17T20:48:42Z

Experimental, not sure if there's a better way to do it. We'd also need to do something similar in modeling_utils.py.

Addresses #2709.

HuggingFaceDocBuilderDev · 2023-03-17T20:53:23Z

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

sayakpaul · 2023-03-18T03:38:36Z

I don't see other way around either.

patrickvonplaten · 2023-03-18T17:33:44Z

src/diffusers/pipelines/pipeline_utils.py


+            # Dynamo wraps the original mode and changes the class.
+            # Is there a principled way to obtain the original class?
+            if "_dynamo.eval_frame" in str(model_cls):


Think that's a bit cleaner

Suggested change

if "_dynamo.eval_frame" in str(model_cls):

if is_torch_version(">=", "2.0.0") and isinstance(sub_model, torch._dynamo.eval_frame.OptimizedModule):

@pcuenca quick question - why don't we go for this API? Think isinstance is a better check here than the string checking

Because torch._dynamo is not available in all platforms, just on cuda builds.

How about this then:

Suggested change

if "_dynamo.eval_frame" in str(model_cls):

if is_torch_version(">=", "2.0.0") and hasattr(torch, "_dynamo") and isinstance(sub_model, torch._dynamo.eval_frame.OptimizedModule):

It's a bit too verbose and distracting for my taste, specially if we are going to use it in multiple places. We could create a helper function somewhere. The other issue is that using a private symbol makes me uneasy about the class hierarchy changing in future revisions, making the code crash. Which is not necessarily bad, because if we use a string then we'd silently fail.

I see! Tbh, both ways are fine by me - I do prefer to avoid silent errors and generally don't like string "in" testing for classes as this is a very brittle test in my opinion.

Would prefer a helper function here, but also ok to leave as is as it's an edge case.

src/diffusers/pipelines/pipeline_utils.py

patrickvonplaten

Looks nice - thanks for the fix! Let's maybe do the same in modeling_utils.py. Would just recommend to use the isinstance(...) check instead of string matching check - think that's a bit less brittle.

The way you retrieve the original class here is good IMO

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Instead of the dynamo class.

pcuenca · 2023-03-20T15:15:41Z

I created a simple test that unsurfaced additional problems in the configuration (the wrong class was saved to config.json).

PyTorch CPU does not have _dynamo.

sayakpaul

Thanks!

patrickvonplaten · 2023-03-21T11:07:16Z

src/diffusers/pipelines/pipeline_utils.py

                register_dict = {name: (None, None)}
            else:
+                # register the original module, not the dynamo compiled one
+                if "_dynamo.eval_frame" in str(module.__class__):


Think this check is better:

if is_torch_version(">=", "2.0.0") and isinstance(sub_model, torch._dynamo.eval_frame.OptimizedModule):

As commented above, torch._dynamo does not exist in all platforms.

Ah sorry, now I get it

patrickvonplaten · 2023-03-21T11:07:20Z

src/diffusers/pipelines/pipeline_utils.py

+        # I didn't find a public API to get the original class.
+        sub_model = passed_class_obj[name]
+        model_cls = sub_model.__class__
+        if "_dynamo.eval_frame" in str(model_cls):


Think this check is better:

if is_torch_version(">=", "2.0.0") and isinstance(sub_model, torch._dynamo.eval_frame.OptimizedModule):

patrickvonplaten

Can we apply the same changes to the model save_pretrained function?

patrickvonplaten · 2023-03-21T11:27:15Z

Thanks for clarifying the design here. Could we then maybe go with: #2726 (comment) or do you prefer how it is now?

pcuenca · 2023-03-21T19:38:29Z

Can we apply the same changes to the model save_pretrained function?

It works, actually; at first I thought it wouldn't but I was wrong. When save_pretrained is invoked on a compiled model, self is the appropriate class and saving succeeds. I verified that the json file is correct and it refers to the right class. I also added a very small test case.

pcuenca · 2023-03-21T19:41:02Z

@patrickvonplaten so the only thing missing here would be to decide whether to use a helper function to verify the class. I'm undecided because I think that testing for a private class that might change is not very much better than testing for a name, but we might want to make it fail so we can fix it when it changes. What do you think?

patrickvonplaten

Thanks! Summarized my thoughts here: https://github.com/huggingface/diffusers/pull/2726/files#r1146130475

Good to go whenever you want :-)

pcuenca · 2023-03-27T14:08:37Z

I used the helper function and I think it's ok, can one of you give it a final look before merge? @sayakpaul @patrickvonplaten

sayakpaul · 2023-03-28T06:11:29Z

Nice! Looks clean to me :)

* Workaround for saving dynamo-wrapped models. * Accept suggestion from code review Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * Apply workaround when overriding pipeline components. * Ensure the correct config.json is saved to disk. Instead of the dynamo class. * Save correct module (not compiled one) * Add test * style * fix docstrings * Go back to using string comparisons. PyTorch CPU does not have _dynamo. * Simple test for save_pretrained of compiled models. * Helper function to test whether module is compiled. --------- Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Workaround for saving dynamo-wrapped models.

72f1b13

pcuenca requested a review from patrickvonplaten March 17, 2023 20:48

patrickvonplaten reviewed Mar 18, 2023

View reviewed changes

src/diffusers/pipelines/pipeline_utils.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed Mar 18, 2023

View reviewed changes

pcuenca and others added 6 commits March 20, 2023 14:09

Accept suggestion from code review

9351bc8

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Apply workaround when overriding pipeline components.

351b47b

Ensure the correct config.json is saved to disk.

cb76bde

Instead of the dynamo class.

Save correct module (not compiled one)

e990b53

Add test

f803679

style

ed7a625

pcuenca requested review from patrickvonplaten and sayakpaul March 20, 2023 15:15

pcuenca added 2 commits March 20, 2023 18:11

fix docstrings

0f3a8f0

Go back to using string comparisons.

2c4ba90

PyTorch CPU does not have _dynamo.

sayakpaul approved these changes Mar 21, 2023

View reviewed changes

patrickvonplaten reviewed Mar 21, 2023

View reviewed changes

Simple test for save_pretrained of compiled models.

4680faf

Merge branch 'main' into dynamo-save-pretrained

f00d4d5

patrickvonplaten approved these changes Mar 23, 2023

View reviewed changes

Helper function to test whether module is compiled.

e41331b

pcuenca merged commit 81125d8 into main Mar 28, 2023

pcuenca deleted the dynamo-save-pretrained branch March 28, 2023 07:03

	if "_dynamo.eval_frame" in str(model_cls):
	if is_torch_version(">=", "2.0.0") and isinstance(sub_model, torch._dynamo.eval_frame.OptimizedModule):

	if "_dynamo.eval_frame" in str(model_cls):
	if is_torch_version(">=", "2.0.0") and hasattr(torch, "_dynamo") and isinstance(sub_model, torch._dynamo.eval_frame.OptimizedModule):

Make dynamo wrapped modules work with save_pretrained #2726

Make dynamo wrapped modules work with save_pretrained #2726

Uh oh!

Conversation

pcuenca commented Mar 17, 2023

Uh oh!

HuggingFaceDocBuilderDev commented Mar 17, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sayakpaul commented Mar 18, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten Mar 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca commented Mar 20, 2023

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Mar 21, 2023

Uh oh!

pcuenca commented Mar 21, 2023

Uh oh!

pcuenca commented Mar 21, 2023

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca commented Mar 27, 2023

Uh oh!

sayakpaul commented Mar 28, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

HuggingFaceDocBuilderDev commented Mar 17, 2023 •

edited

Loading

patrickvonplaten Mar 21, 2023 •

edited

Loading