[Afloat16]

Fixes #6549.

Problem

General sharded optimizer loading restores parameter groups and accumulates state with optimizer.state.update(...), but it never removes the destination optimizer's earlier state. Parameters absent from a checkpoint retain stale momentum or Adam moments. In a native reproduction, saving a checkpoint before the first SGD step, accumulating momentum, then restoring that checkpoint makes the next update produce 0.81 instead of 0.9.

Changes

After the saved parameter groups pass compatibility validation and the parameter mapping is built, replace the destination state with a fresh defaultdict(dict). This happens once before loading the shards, so successive shards continue to accumulate normally. Incompatible groups still raise before changing existing groups or state.

Validation

Base: 297c8cf6d91cb0bff795cba961d3aea1e66fa370.

• Identical 25 new regression cases on unchanged main: 16 failed, 9 passed. After the fix: 25 passed.
• Including four existing related checkpoint cases: 29 passed.
• Public GeneralCheckpointIO.save_optimizer/load_optimizer tests restore real on-disk empty-state sharded checkpoints into already-trained SGD, Adam and AdamW optimizers, covering one/two parameter groups and float32/float64. Three subsequent updates match native Optimizer.load_state_dict exactly for parameters and every state tensor.
• Additional tests cover safetensors and legacy group files, saved parameter-ID mapping, removal of old non-parameter state, invalid group count/size/type without mutation, unsharded controls and loading into fresh optimizers.
• All applicable repository pre-commit hooks and git diff --check passed.

Command:

python -m pytest tests/test_checkpoint_io/test_sharded_optimizer_replacement.py tests/test_checkpoint_io/test_optimizer_safetensors_v2.py::test_shards_groups_and_prefetch tests/test_checkpoint_io/test_optimizer_safetensors_v2.py::test_optimizer_index_preserves_internal_link_names tests/test_checkpoint_io/test_optimizer_safetensors_v2.py::test_general_sync_checkpoint_next_step -q

Validated on Linux/Python 3.12 and PyTorch 2.15.0+cpu. Legacy-format cases follow the repository's PyTorch >=2.10 requirement. All 25 new tests run without replacing accelerator methods. Two existing nonempty-shard cases use their existing cpu_sync fixture because the CPU accelerator does not implement synchronize(); this is not a claim of complete native GPU or distributed checkpoint validation. The full repository suite was not run.