[BUG] The order of the param groups matters when creating optimizer

**Describe the bug**
A clear and concise description of what the bug is.

According to the code https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/optimizer/__init__.py#L131, the order of the param groups is not deterministic, and might be an issue when pipeline parallelism is enabled

**To Reproduce**
Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.

Run NeMo deepseek_v2_lite training with pp>1, warning like "There is difference in the common state dict in different ranks" appears and it turns out the order of the param groups in different pp stages is different, while only rank0 will save the param groups in ckpt.

**Expected behavior**
A clear and concise description of what you expected to happen.
A key sort might needed as follows

**Stack trace/logs**
If applicable, add the stack trace or logs from the time of the error.
I ran deepseek_v2_lite in NeMo with tp=2, pp=4 on 8 gpus (changed the layer_num so that pp could be 4). The warning related to ckpt saving is as follows:

There is difference in the common state dict in different ranks. The differences are {2: ([], [], [(('optimizer', 0, 'optimizer', 'param_groups', 0, 'wd_mult'), , ), (('optimizer', 0, 'optimizer', 'param_groups', 1, 'wd_mult'), , )]), 3: ([], [], [(('optimizer', 0, 'optimizer', 'param_groups', 0, 'wd_mult'), , ), (('optimizer', 0, 'optimizer', 'param_groups', 1, 'wd_mult'), , )]), 4: ([], [], [(('optimizer', 0, 'optimizer', 'param_groups', 0, 'wd_mult'), , ), (('optimizer', 0, 'optimizer', 'param_groups', 1, 'wd_mult'), , )]), 5: ([], [], [(('optimizer', 0, 'optimizer', 'param_groups', 0, 'wd_mult'), , ), (('optimizer', 0, 'optimizer', 'param_groups', 1, 'wd_mult'), , )]), 6: ([], [], [(('optimizer', 0, 'optimizer', 'param_groups', 0, 'wd_mult'), , ), (('optimizer', 0, 'optimizer', 'param_groups', 1, 'wd_mult'), , )]), 7: ([], [], [(('optimizer', 0, 'optimizer', 'param_groups', 0, 'wd_mult'), , ), (('optimizer', 0, 'optimizer', 'param_groups', 1, 'wd_mult'), , )])}

**Environment (please complete the following information):**
- Megatron-LM 56f8a96e7e328f1c9979b03b359fd284432050fc

Start at megatron/core/optimizer/__init__.py around line 131 and inspect how parameter groups are assembled. Reproduce the NeMo DeepSeek-V2-Lite training case with tp=2 and pp=4, then verify that parameter-group ordering is consistent across pipeline stages and the checkpoint common-state-dict warning no longer appears.

Written by the indexing model from the issue text.