Restore generation unwrap state if generate() raises

[YeonwooSung]

What does this PR do?

_unwrap_model_for_generation disables gradient checkpointing and, on ZeRO-3, removes DeepSpeed hooks, then restores them only after a successful yield. If generate() raises, checkpointing stays off and the hooks stay removed.

This wraps the yield in try/finally, matching _override_model_generation_config in the same file. ZeRO-3 add_hooks is also in an inner finally so it runs even when generation fails inside GatheredParameters.

Adds test_restores_gradient_checkpointing_on_error.

Fixes #6659

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? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

---

Note

Low Risk
Narrow bugfix to generation context cleanup; behavior on success is unchanged and failures now restore training state instead of leaving the model misconfigured.

Overview
Fixes a cleanup gap in _unwrap_model_for_generation: when generate() (or other code inside the context) raises, gradient checkpointing could stay disabled and, on DeepSpeed ZeRO-3 with gathered params, optimizer hooks might not be reattached.

The yield path is now wrapped in try/finally so gradient checkpointing is re-enabled on exit, mirroring _override_model_generation_config. For ZeRO-3 with GatheredParameters, add_hooks runs in an inner finally so hooks are restored even if generation fails mid-context.

Adds test_restores_gradient_checkpointing_on_error to assert checkpointing is turned back on after a RuntimeError inside the context.

^{Reviewed by Cursor Bugbot for commit 0139109. Bugbot is set up for automated code reviews on this repo. Configure here.}