fix: model weight updates with automatic_optimization=False in mixed precision training

[iamarunbrahma]

What does this PR do?

Fixes an issue where model weights were not being properly updated when using automatic_optimization=False with mixed precision training (16-mixed precision). The core issues were:

1. Improper gradient scaling/unscaling in the AMP plugin
2. Missing validation of gradient values before optimizer steps
3. Lack of proper error handling for repeated unscale operations

The changes include:

• Enhanced AMP plugin implementation with proper gradient scaling workflow
• Added explicit gradient unscaling before optimizer steps
• Improved error handling for edge cases
• Added warning when gradients become NaN/inf during training

Fixes #20215

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--20460.org.readthedocs.build/en/20460/

[lantiga]

Thank you for the contribution, will review today

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79%. Comparing base (be608fa) to head (4e6faf5).

❗ There is a different number of reports uploaded between BASE (be608fa) and HEAD (4e6faf5). Click for more details.

HEAD has 542 uploads less than BASE

Flag	BASE (be608fa)	HEAD (4e6faf5)
cpu	146	24
lightning_fabric	21	0
pytest	76	0
python3.9	36	6
lightning	109	18
python3.11	36	6
python3.10	19	3
gpu	2	0
python3.12	55	9
pytorch2.1	27	9
pytest-full	72	24
pytorch_lightning	18	6
pytorch2.2.2	9	3
pytorch2.3	9	3
pytorch2.5.1	18	6
pytorch2.4.1	9	3

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #20460     +/-   ##
=========================================
- Coverage      88%      79%     -9%     
=========================================
  Files         267      265      -2     
  Lines       23276    23288     +12     
=========================================
- Hits        20383    18322   -2061     
- Misses       2893     4966   +2073     

[lantiga]

something must have happened when submitting the PR, changes need to go in here
https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/plugins/precision/amp.py#L67

but they appear as a stand-alone file https://github.com/Lightning-AI/pytorch-lightning/pull/20460/files#diff-77ef63a4f3cecc394dba6f10386a36ba11362b5e4086b71fca4a4561caf8b0c8, with the method outside the class

[lantiga]

@iamarunbrahma looking at the signature for optimizer_step from, are you working off a pre-2.0 version of PyTorch Lightning? optimizer_idx was removed before 2.0.0 in #16539

[lantiga]

ok the more I look at this the more puzzled I am, can you clarify the code in the PR @iamarunbrahma, how it fits with the plugin optimizer_step and how it solves the original issue?

[lantiga]

BTW the title is contradicting the issue: the issue describes something happening with automatic_optimization=True, not False. I'm closing this one.