[PyTorch] Avoid temporary state casts when loading FusedAdam checkpoints

[bzantium]

Description

FusedAdam.load_state_dict() lets PyTorch cast parameter states to the parameter dtype before restoring them in the optimizer's configured dtypes. With BF16 parameters and FP32 moments, this creates temporary GPU allocations that are immediately discarded.

Load parameter-group metadata and non-parameter state through PyTorch, then restore parameter states through the existing TE path. Keep the original path for subclasses and registered load hooks, which can observe intermediate state. Restored states continue to own their storage; checkpoint tensors are copied.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Avoid redundant parameter-state casts during checkpoint loading.
• Preserve parameter-group validation, storage ownership, and hook/subclass behavior.
• Cover restored states, the next optimizer update, CPU/GPU checkpoints, supported floating-point state dtypes, DTensor placement, and peak allocation.

Testing

All 23 targeted GPU tests pass. On the base revision, the allocation regression fails and the other 22 pass. Reloading initialized state for 4M BF16 parameters reduces extra peak allocation from 24 MiB to within the test's 1 MiB bound.

Validated on B200 with PyTorch 2.11 / CUDA 13.1, using the changed optimizer Python source with TE 2.14.1 backend libraries. A fresh current-main TE build has not been tested. Black, Pylint, and the repository license check pass.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes