Feat (brevitas_examples/llm): correct scale init with CPU offloading

[Giuseppe5]

Reason for this PR

Offload might cause models to end up on CPU + GPU.
If that is the case, the initialization of scale factor is not correct, since the scales that are on CPU are forgotten because of the offload.
This allows to initialize and fix the scales correctly.

Changes Made in this PR

post_forward has now a hook that stores the new parameters, if these have been updated, like the scale post initialization.

Testing Summary

NA

Risk Highlight

NA

• This PR includes code from another work (please detail).
• This PR contains API-breaking changes.
• This PR depends on work in another PR (please provide links/details).
• This PR introduces new dependencies (please detail).
• There are coverage gaps not covered by tests.
• Documentation updates required in subsequent PR.

Checklist

• Code comments added to any hard-to-understand areas, if applicable.
• Changes generate no new warnings.
• Updated any relevant tests, if applicable.
• No conflicts with destination dev branch.
• I reviewed my own code changes.
• Initial CI/CD passing.
• 1+ reviews given, and any review issues addressed and approved.
• Post-review full CI/CD passing.

[nickfraser]

Please add comments to the changes in main.py. Looks like internal accelerate hackery that deserves a bit of context.

[nickfraser]

Nice solution BTW, but me in 6 months won't remember how this works