Handle step key correctly in dp_reshardable checkpoint save with --optimizer-cpu-offload

[artkorenev]

Grafts NVIDIA#2874 (f4502eb1c9).

Wraps the optimizer step in LocalNonpersistentObject in
sharded_param_state_dp_reshardable so it is not persisted into the rank-0
common state.

🤖 Generated with Claude Code

[yushengsu-thu]

The graft matches upstream exactly, but I traced the load path with --optimizer-cpu-offload and hit two problems:

1. Old checkpoints won't load anymore. Pre-PR saves leaked the raw step into common.pt (since [2/8] feat: support partial checkpoint loading #3), and the load request now puts a LocalNonpersistentObject at the same path, so merge() in serialization.py:126 throws on the duplicate leaf. No strictness flag gets around that.

2. step gets reset to 1 on resume. The load skeleton is built after HDO's dummy_step(), LNO keeps that local value, and _set_main_param_and_optimizer_states then copies it over the step that load_state_dict just restored from param_groups (:863-873 vs :913). So Adam bias correction restarts on every resume. The old leak was actually what made resume work before. Upstream has the same bug, the code is identical there.

For 2 the easy fix is skipping step in _set_main_param_and_optimizer_states, same as load_parameter_state_from_fs_model_space already does (:1852). For 1, drop the key from the request when is_loading instead of wrapping it. If nobody has live cpu-offload + dp_reshardable checkpoints then 1 is just a release note, but I'd fix 2 before this lands since it breaks the exact resume flow the PR is for.

(Only cpu-offload / torch AdamW runs are affected — FusedAdam keeps step in param_groups, which is also why no existing test catches this.)