cpu optimizer offloading can raise assertion error on checkpoint saving

**Describe the bug**

Git commit: Repro on 44af130cc4568d324860646996b6a5bfd6c5e3e6

Crashes when check-pointing when saving OPTIMIZER_OFFLOAD=1, PR=fp8, when it hits for the save interval from the DeepSeek configs from Megatron Model Zoo.

A clear and concise description of what the bug is

**Steps/Code to reproduce bug**
```python
24: [rank24]: Traceback (most recent call last):
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/pretrain_gpt.py", line 233, in
24: [rank24]: pretrain(
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 737, in pretrain
24: [rank24]: iteration, num_floating_point_operations_so_far = train(
24: [rank24]: ^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 2405, in train
24: [rank24]: should_exit = checkpoint_and_decide_exit(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1850, in checkpoint_and_decide_exit
24: [rank24]: save_checkpoint_and_time(
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1721, in save_checkpoint_and_time
24: [rank24]: save_checkpoint(
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 509, in save_checkpoint
24: [rank24]: state_dict = generate_state_dict(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 817, in generate_state_dict
24: [rank24]: optimizer_sd = optimizer.sharded_state_dict(state_dict, **(optim_sd_kwargs or {}))
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/optimizer/optimizer.py", line 1204, in sharded_state_dict
24: [rank24]: optim_state_dict = optimizer.sharded_state_dict(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 1252, in sharded_state_dict
24: [rank24]: param_state = self.sharded_param_state_dp_reshardable(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 1646, in sharded_param_state_dp_reshardable
24: [rank24]: assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), (
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: AssertionError: (torch.Size([]), 0, 1050820)
25: [rank25]: Traceback (most recent call last):
```

Not sure if this interacts with the singleton save. Turning OPTIMIZER_OFFLOAD=0 seems to mitigate this but need to double check.

A clear and concise description of what you expected to happen.

Checkpointing saving works properly when optimizer offload is enabled.
**Additional context**

Add any other context about the problem here.