Repository: pytorch
Issue: [DCP]Distributed checkpoint `set_optimizer_state_dict` cause optimizer step error when optimizer contains empty param group
Labels: oncall: distributed,module: optimizer,triaged

### 🐛 Describe the bug

DCP `set_optimizer_state_dict` introduce wrong param group and cause `optim.step` raise error when original state dict contains param group that doesn't have any parameters.

* Error Message

```
[rank1]: Traceback (most recent call last):
[rank1]:   File "/path/to/pytorch_bug/dcp_bug.py", line 45, in <module>
[rank1]:     optim_new.step()
[rank1]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/optim/optimizer.py", line 493, in wrapper
[rank1]:     out = func(*args, **kwargs)
[rank1]:           ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/optim/optimizer.py", line 91, in _use_grad
[rank1]:     ret = func(self, *args, **kwargs)
[rank1]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/optim/adamw.py", line 230, in step
[rank1]:     beta1, beta2 = cast(Tuple[float, float], group["betas"])
[rank1]:                                              ~~~~~^^^^^^^^^
[rank1]: KeyError: 'betas'
```

* Code

`torchrun --nnodes=1 --nproc-per-node=4 --standalone /path/to/pytorch_bug/dcp_bug.py`

```Python
import torch 
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict, set_optimizer_state_dict
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
)
from torch.distributed.tensor import Shard, DTensor, Replicate
import os
_world_size = int(os.environ["WORLD_SIZE"])
device_mesh = init_device_mesh(device_type="cuda", mesh_shape=(_world_size,))
class TestMod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(64, 64)
    
    def forward(self, x):
        return self.fc(x)
mod = TestMod().cuda()
parallelize_module(mod, device_mesh, {
    "fc": ColwiseParallel(use_local_output=False)
})
optim = torch.optim.AdamW([
    {"params": mod.parameters()},
    {"params": [], "lr": 0.2}, # empty pg group here
], lr=0.1)
optim_new = torch.optim.AdamW([
    {"params": mod.parameters()},
    {"params": [], "lr": 0.2}, # empty pg group here
], lr=0.1)

# init optimizer state
sample_inp = torch.randn(2, 128, 64).cuda()
sample_target = torch.randn(2, 128, 64).cuda()
loss_cls = torch.nn.MSELoss()
optim.zero_grad()
output = mod(sample_inp).redistribute(device_mesh, [Replicate()]).to_local()
loss = loss_cls(output, sample_target)
loss.backward()
optim.step()
# bug
optim_state_dict = get_optimizer_state_dict(mod, optim)
set_optimizer_state_dict(mod, optim_new, optim_state_dict)

optim_new.step()
```


### Versions

```
PyTorch version: 2.6.0.dev20241222+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.10 (main, Oct  3 2024, 07:29:13) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.210-4-velinux1-amd64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
Nvidia driver version: 535.86.10
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==2.1.3
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] pytorch-triton==3.2.0+git0d4682f0
[pip3] torch==2.6.0.dev20241222+cu124
[pip3] torchaudio==2.6.0.dev20241222+cu124
[pip3] torchdata==0.9.0
[pip3] torchpippy==0.2.0+1bcb2bf
[pip3] torchtitan==0.0.2
[pip3] torchvision==0.22.0.dev20241222+cu124
[pip3] triton==3.1.0
[conda] blas                      1.0                         mkl  
[conda] cuda-cudart               12.1.105                      0    nvidia
[conda] cuda-cupti                12.1.105                      0    nvidia
[conda] cuda-libraries            12.1.0                        0    nvidia
[conda] cuda-nvrtc                12.1.105                      0    nvidia
[conda] cuda-nvtx                 12.1.105                      0    nvidia
[conda] cuda-opencl               12.4.127                      0    nvidia
[conda] cuda-runtime              12.1.0                        0    nvidia
[conda] libcublas                 12.1.0.26                     0    nvidia
[conda] libcufft                  11.0.2.4                      0    nvidia
[conda] libcurand                 10.3.5.147                    0    nvidia
[conda] libcusolver               11.4.4.55                     0    nvidia
[conda] libcusparse               12.0.2.55                     0    nvidia
[conda] libnvjitlink              12.1.105                      0    nvidia
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py311h5eee18b_1  
[conda] mkl_fft                   1.3.11          py311h5eee18b_0  
[conda] mkl_random                1.2.8           py311ha02d727_0  
[conda] numpy                     2.1.3           py311h08b1b3b_0  
[conda] numpy-base                2.1.3           py311hf175353_0  
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.2                    pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] pytorch-cuda              12.1                 ha16c6d3_6    pytorch
[conda] pytorch-triton            3.2.0+git0d4682f0          pypi_0    pypi
[conda] torch                     2.6.0.dev20241222+cu124          pypi_0    pypi
[conda] torchaudio                2.6.0.dev20241222+cu124          pypi_0    pypi
[conda] torchdata                 0.9.0                    pypi_0    pypi
[conda] torchpippy                0.2.0+1bcb2bf            pypi_0    pypi
[conda] torchtitan                0.0.2                    pypi_0    pypi
[conda] torchvision               0.22.0.dev20241222+cu124          pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi
```

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar @LucasLLC @pradeepfn