[ZenAlexa]

Description

Loading an FP32 optimizer checkpoint into FP16 states scales the checkpoint tensors in place. A second load then treats those already scaled values as unscaled moments or master weights, so training resumes with different optimizer state.

Write the scaled FP32 product directly to its FP16 destination. Keep the per-tensor scale scalar so loading a zero-dimensional parameter preserves its shape. The checkpoint format and optimizer API stay unchanged.

Verification

• Extended the existing save/load training scenario to restore the same state twice. The FP16 master weight, first moment and second moment cases all failed on main and passed with the change. The full fused optimizer test file had 33 passed and 6 skipped.
• Compared three resumed steps with uninterrupted training across 24 combinations of state dtype and parameter layout, including scalars and noncontiguous views. Parameters and optimizer states matched bit for bit.
• On one H200 with PyTorch 2.9.0+cu128, synchronized FusedAdam steps with FP32 parameters and both moments in FP16 took 0.183 → 0.176 ms for one 1024×1024 parameter, 4.249 → 4.015 ms for 32 parameters of 512×512, and 0.610 → 0.585 ms for one 4096×4096 parameter (medians of 11 interleaved measurements per case). These cover optimizer steps, not full training iterations. Both revisions ran against the same native extension built from a7aec214; the change is in the PyTorch optimizer.