[BUG] dist_muon checkpoint load error

I save a chekpoint with **dist_muon**,

```
--muon-momentum 0.95 \
--optimizer dist_muon \
--muon-use-nesterov \
--muon-scale-mode unit_rms_norm \
--muon-fp32-matmul-prec high \
--muon-num-ns-steps 5 \
--muon-extra-scale-factor 1.0
--muon-tp-mode blockwise
--use-persistent-ckpt-worker
--async-save
```

**then load the ckpt, got error:**

```
192.168.11.3: [rank22]: Traceback (most recent call last):
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/examples/../pretrain_gpt.py", line 263, in
192.168.11.3: [rank22]: pretrain(
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/training/training.py", line 762, in pretrain
192.168.11.3: [rank22]: model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/training/training.py", line 1288, in setup_model_and_optimizer
192.168.11.3: [rank22]: args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/training/checkpointing.py", line 1717, in load_checkpoint
192.168.11.3: [rank22]: optimizer.load_state_dict(state_dict['optimizer'])
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/core/optimizer/layer_wise_optimizer.py", line 246, in load_state_dict
192.168.11.3: [rank22]: super().load_state_dict(state_dict)
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/core/optimizer/optimizer.py", line 1225, in load_state_dict
192.168.11.3: [rank22]: optimizer.load_state_dict(state)
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/core/optimizer/optimizer.py", line 856, in load_state_dict
192.168.11.3: [rank22]: state_dict[optimizer_key]['param_groups'] = self._filter_and_reorder_param_groups(
192.168.11.3: [rank22]: File "/kanliu/Megatron-LM/megatron/core/optimizer/optimizer.py", line 419, in _filter_and_reorder_param_groups
192.168.11.3: [rank22]: raise ValueError(
192.168.11.3: [rank22]: ValueError: Could not find parameter group with key (0.0, 1.0, False, False) in loaded checkpoint.
192.168.11.3: [rank22]: Available keys:
192.168.11.3: [rank22]: (1.0, 1.0, False, False)
192.168.11.3: [rank22]: Parameter group key definition: ('wd_mult', 'lr_mult', 'is_expert_parallel', 'is_decoupled_lr')
```

I have identified the cause as the use of “**--pipeline-model-parallel-layout**”， but I have no idea to solove it.
If no_wd, I print param info in _get_param_groups:

```
Parameter with wd_mult=0.0: module.module.decoder.final_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.final_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernormINFO:megatron.core.optimizer:Setting up optimizer with config OptimizerConfig(optimizer='adam', lr=0.0002, min_lr=0.0, decoupled_lr=None, decoupled_min_lr=None, weight_decay=0.1, fp8_recipe='delayed', fp16=False, bf16=False, reuse_grad_buf_for_mxfp8_param_ag=False, params_dtype=torch.bfloat16, use_precision_aware_optimizer=False, store_param_remainders=True, main_grads_dtype=torch.float32, main_params_dtype=torch.float32, exp_avg_dtype=torch.float32, exp_avg_sq_dtype=torch.float32, loss_scale=None, initial_loss_scale=4294967296, min_loss_scale=1.0, loss_scale_window=1000, hysteresis=2, adam_beta1=0.9, adam_beta2=0.95, adam_eps=1e-08, decoupled_weight_decay=True, sgd_momentum=0.9, muon_momentum=0.95, muon_split_qkv=True, muon_use_nesterov=True, muon_scale_mode='unit_rms_norm', muon_fp32_matmul_prec='high', muon_num_ns_steps=5, muon_tp_mode='blockwise', muon_extra_scale_factor=1.0, use_distributed_optimizer=False, overlap_param_gather=False, overlap_param_gather_with_optimizer_step=False, optimizer_cpu_offload=False, optimizer_offload_fraction=1.0, use_torch_optimizer_for_cpu_offload=False, overlap_cpu_optimizer_d2h_h2d=False, pin_cpu_grads=True, pin_cpu_params=True, clip_grad=1.0, log_num_zeros_in_grad=False, barrier_with_L1_time=True, timers=, config_logger_dir='')
Parameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernormParameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_projParameter with wd_mult=0.0: module.module.decoder.final_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.final_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_projParameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_projParameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_projParameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernormParameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernorm
Parameter with wd_mult=0.0: module.module.decoder.final_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.final_layernormParameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernormParameter with wd_mult=0.0: module.module.decoder.final_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.final_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.input_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.input_layernormParameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_projParameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernormParameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernorm
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight, shape: torch.Size([1536]), module: module.module.decoder.layers.0.self_attention.linear_q_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight, shape: torch.Size([512]), module: module.module.decoder.layers.0.self_attention.linear_kv_up_proj
Parameter with wd_mult=0.0: module.module.decoder.layers.0.pre_mlp_layernorm.weight, shape: torch.Size([7168]), module: module.module.decoder.layers.0.pre_mlp_layernorm
```
ALL layer get error~~， I caused by **layerwiseoptimizer**?