Remove explicit copy of model.forward() and instead check/store the id to unwrap

[muellerzr]

What does this PR do?

This PR solves another bug with fp8 where having a link to the original forward in the model would cause yet another memory spike:

(From the NVIDIA team):

we still experience OOM in fp8, but in model function call. In fp8, the forward step takes 15 GB of memory, while in bf16, it is 3GB.

The solution is to get rid of _original_forward_func (which was shown to be the issue) and to instead rely on the id of the function and don't make copies of it/copies of references.

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.

@BenjaminBossan @SunMarc

[HuggingFaceDocBuilderDev]

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

[BenjaminBossan]

Nice work identifying and fixing the issue.

Regarding the solution, it isn't really pretty, but there is probably no pretty solution.

What confuses me is why the previous code would lead to a memory spike but not this one. Clearly, it can't be the case that, by holding a reference to the original module's forward method, accelerate somehow inhibits garbage collection. Since we can still retrieve the original method through unwrapping, a reference to the original module has to exist.

It could be some kind of circular reference issue, but that should be detected by the GC. I wonder if this points to some kind of more fundamental bug. Also, it's really unclear to me how it related to fp8.

Regarding id, I very rarely make us of it. Can we be sure that it will do the right thing? I think it should work, but we should consider all possible edge cases.