fix(optimizer): checkpoint native FP32 parameters in mixed precision

[FDecaYed]

What

Fix sharded optimizer-state mapping when mixed-precision parameter groups interleave BF16/FP16 parameters with parameters that remain natively FP32. State ids now follow the inner optimizer order, while FP32 master copies are paired only with converted low-precision parameters.

This is a generic checkpointing correctness fix and is detached from the DSv4 stack. It does not depend on the parameter-marking PR; DSv4 call sites can consume both independently.

Provenance

• Extracted from the high-precision DSv4 follow-up in [Tracking] DeepSeek-v4 HybridModel merge plan #5795 (commit 081eebc).
• The eventual precision-sensitive consumers originate in the CSA work from [dev] [DeepSeek-v4] Part 1: Hybrid Attention with CSA and HCA #4458 and [dev] [DeepSeek-v4] Part 4: Fusion Kernels for DSv4 Hybrid Attention #4894.
• Original DSv4 integration by @hxbai and the source authors.
• Reconstructed against current main; feature cutoff: 2026-07-21.
• GPTModel integration is intentionally excluded.

Testing

• uv run isort --check-only on both changed Python files
• uv run ruff check on both changed Python files
• uv run python -m compileall -q on both changed Python files
• Focused CUDA checkpoint test is included; local execution requires the project CUDA/Torch environment and is left to CI.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[FDecaYed]

/ok to test 851d75d

[FDecaYed]

Closing this standalone boundary because it is absorbed into P1. P1 combines selective-FP32 runtime and native-FP32 optimizer checkpoint correctness and will be published by updating #5929. Tracking PR: #5795.