Fix model copy after dispatch_model

[austinapatel]

What does this PR do?

Fixes #1835 where a model dispatched with dispatch_model was unable to be properly copied with copy.deepcopy. Specifically, the copied model would incorrectly still reference the align device hook forward function of the original model, rather than its own forward function. This PR includes a fix to this issue as well as a test to validate the fix.

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,
  Pull Request section?
• 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.

@SunMarc @sgugger

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[SunMarc]

Hi @austinapatel, thanks a lot for fixing this ! Just a few nits :) Can you have a second look @BenjaminBossan ?

[austinapatel]

Thanks for the review @SunMarc! I have pushed changes per your feedback!

[BenjaminBossan]

In general, this LGTM. I only have one question, apart from that this can be merged.

[BenjaminBossan]

I wonder about this line versus:

module.forward = MethodType(functools.update_wrapper(new_forward, old_forward), module)

@austinapatel Do you think that would be preferable? I thought it might affect pickling, but unfortunately, both variants don't allow it, but the error messages are different.

[austinapatel]

Thanks for the suggestion and review! I just tried your code snippet, and I run into the same issue as this comment: #1971 (comment) since functools.update_wrapper needs to be called last (test_add_and_remove_hooks in test_hooks.py fails). If I try module.forward = functools.update_wrapper(MethodType(new_forward, module), old_forward) I get the following error: AttributeError: 'method' object has no attribute '__module__' when calling update_wrapper. My understanding is that update_wrapper is now being called on a method rather than on a function, which doesn't play as nicely with update_wrapper.

[BenjaminBossan]

Thanks for explaining and testing, this makes sense. I'm a bit wary about binding the module via partial to basically imitate a bound method, but it seems that all other solutions don't work, so I guess it's the best we can do.

[BenjaminBossan]

Thanks so much @austinapatel for all the work you put into this.

[austinapatel]

Thanks @SunMarc and @BenjaminBossan for the reviews and feedback!