[BUG]: Sharded optimizer loading retains old momentum absent from the checkpoint #6549

Is there an existing issue for this bug?

• I have searched existing issues and pull requests.
• I reproduced this on current main, commit 297c8cf6d91cb0bff795cba961d3aea1e66fa370.

Describe the bug

GeneralCheckpointIO.load_optimizer retains optimizer state that is absent from a sharded checkpoint. The loader restores parameter groups and adds each shard with optimizer.state.update(...), but never removes the destination optimizer's earlier state. This differs from native Optimizer.load_state_dict and unsharded loading.

A checkpoint saved before the first optimizer step has an empty state. Loading it into an optimizer that has since accumulated momentum leaves the old momentum active and changes the next training update. In the native CPU reproduction below, the loaded state still contains momentum_buffer=tensor([1.]); the next parameter is approximately 0.81, whereas restoring the empty checkpoint should produce 0.9. No accelerator synchronization or GPU operation is needed for this empty-shard checkpoint.

Minimal reproduction

import tempfile
import torch
from colossalai.checkpoint_io import GeneralCheckpointIO

p = torch.nn.Parameter(torch.tensor([1.0]))
optimizer = torch.optim.SGD([p], lr=0.1, momentum=0.9)
io = GeneralCheckpointIO()
with tempfile.TemporaryDirectory() as checkpoint:
    io.save_optimizer(optimizer, checkpoint, shard=True)
    p.grad = torch.ones_like(p)
    optimizer.step()
    optimizer.zero_grad()
    io.load_optimizer(optimizer, checkpoint)
    print(optimizer.state_dict()["state"])
    with torch.no_grad():
        p.fill_(1.0)
    p.grad = torch.ones_like(p)
    optimizer.step()
    print(p.item())

Expected behavior

Restoring a complete sharded optimizer checkpoint should replace the destination state, including removing parameter state absent from the checkpoint. Compatibility validation of parameter groups should happen before clearing the existing state.

Environment

Linux, Python 3.12, PyTorch 2.9.1+cpu, native ColossalAI main at the commit above. This report does not claim GPU or distributed validation.

Related work checked

All-state PR/issue searches for stale/reset optimizer state, empty optimizer checkpoints, and load_param_groups_into_optimizer found no equivalent fix. PR #6433 concerns migrating eager optimizer state to ZeRO master shards; #6490/#6520 concern master-weight precision on resume. Neither addresses replacing the destination state in this loader. The original sharded checkpoint implementation (#3984) already used additive state updates without first resetting the destination.