diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 33b9b78b836..57eb1e94478 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -187,11 +187,13 @@ def get_megatron_muon_optimizer(
 
     assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed."
 
-    # dist-optim is not supported due to strong coupling with how DDP init grad buffer
-    # in thoery we can put some weight to use non-dist-muon and rest to dist-adam
-    # but there are strong dependency and assumption in DDP that prevent it
+    # Dist-opt is not supported due to strong coupling with how DDP init grad buffer
+    # In theory we can change DDP to enable use muon and dist-opt-adam together
     if config.use_distributed_optimizer:
         raise Exception('muon with dist optimizer is not supported.')
+    # only support bf16 w/o loss scale now
+    if config.fp16:
+        raise Exception('muon with fp16 is not supported.')
 
     # before this function receive properly created collection
     if pg_collection is None:
@@ -199,11 +201,30 @@ def get_megatron_muon_optimizer(
 
     log_single_rank(logger, logging.INFO, f'Setting up emerging optimizer with config {config}')
 
+    # Needed for torch_dist ckpt_format, unlike torch ckpt_format
+    # For other emerging optimizers, need to implement init_state_fn as well
+    # TODO(boxiangw): Improve usability after optimizer refactor
+    # TODO(boxiangw): support precision aware optimizer
+    def muon_init_state_fn(opt, config=None):
+        for group in opt.param_groups:
+            for p in group['params']:
+                if len(opt.state[p]) == 0:
+                    opt.state[p]['momentum_buffer'] = torch.zeros_like(p.data)
+
+    def adam_init_state_fn(opt, config=None):
+        for group in opt.param_groups:
+            for p in group['params']:
+                if len(opt.state[p]) == 0:
+                    if config is None or not config.use_precision_aware_optimizer:
+                        opt.state[p]['exp_avg'] = torch.zeros_like(p.data)
+                        opt.state[p]['exp_avg_sq'] = torch.zeros_like(p.data)
+                    else:
+                        opt.initialize_state(p)
+
     optimizers = []
     # record list of non/linear params
     linear_params = []
     nonlinear_params = []
-
     for model_chunk in model_chunks:
         # use config to determine qkv split shapes.
         # no need to check tp since tp splits by head and this is per head(group) dimension
@@ -236,52 +257,36 @@ def get_megatron_muon_optimizer(
             else:
                 nonlinear_params.append(param)
 
+    muon_kwargs = {
+        "lr": config.lr,
+        "momentum_beta": config.muon_momentum,
+        "use_nesterov": config.muon_use_nesterov,
+        "weight_decay": config.weight_decay,
+        "fp32_matmul_prec": config.muon_fp32_matmul_prec,
+        "num_ns_steps": config.muon_num_ns_steps,
+        "scale_mode": config.muon_scale_mode,
+        "split_qkv": config.muon_split_qkv,
+        "is_qkv_fn": lambda p: getattr(p, "is_qkv", False),
+        "qkv_split_shapes": qkv_split_shapes,
+        "extra_scale_factor": config.muon_extra_scale_factor,
+        "pg_collection": pg_collection,
+        "mode": config.muon_tp_mode,
+    }
+
     # freezing nonlinear params and get param groups for muon
     for param in nonlinear_params:
         param.requires_grad = False
 
     linear_param_groups = _get_param_groups(model_chunks, config, config_overrides)
+    # if layerwise distributed optimizer is not used, need to handle ep params separately
+    expert_param_groups = []
+    if not layer_wise_distributed_optimizer:
+        for group in linear_param_groups:
+            if group['is_expert_parallel']:
+                expert_param_groups.append(group)
+                linear_param_groups.remove(group)
 
-    optimizer = TensorParallelMuon(
-        linear_param_groups,
-        lr=config.lr,
-        momentum_beta=config.muon_momentum,
-        use_nesterov=config.muon_use_nesterov,
-        weight_decay=config.weight_decay,
-        fp32_matmul_prec=config.muon_fp32_matmul_prec,
-        num_ns_steps=config.muon_num_ns_steps,
-        scale_mode=config.muon_scale_mode,
-        split_qkv=config.muon_split_qkv,
-        is_qkv_fn=lambda p: getattr(p, 'is_qkv', False),
-        qkv_split_shapes=qkv_split_shapes,
-        extra_scale_factor=config.muon_extra_scale_factor,
-        pg_collection=pg_collection,
-        mode=config.muon_tp_mode,
-    )
-
-    # Needed for torch_dist ckpt_format, unlike torch ckpt_format
-    # For other emerging optimizers, need to implement init_state_fn as well
-    # TODO(boxiangw): Improve usability after optimizer refactor
-    # TODO(boxiangw): support precision aware optimizer
-    def muon_init_state_fn(opt, config=None):
-        for group in opt.param_groups:
-            for p in group['params']:
-                if len(opt.state[p]) == 0:
-                    opt.state[p]['momentum_buffer'] = torch.zeros_like(p.data)
-
-    def adam_init_state_fn(opt, config=None):
-        for group in opt.param_groups:
-            for p in group['params']:
-                if len(opt.state[p]) == 0:
-                    if config is None or not config.use_precision_aware_optimizer:
-                        opt.state[p]['exp_avg'] = torch.zeros_like(p.data)
-                        opt.state[p]['exp_avg_sq'] = torch.zeros_like(p.data)
-                    else:
-                        opt.initialize_state(p)
-
-    # need to wrap into megatron mix precision optimizer. (only support bf16 w/o loss scale now)
-    if config.fp16:
-        raise Exception('muon with fp16 is not supported.')
+    optimizer = TensorParallelMuon(linear_param_groups, **muon_kwargs)
 
     reset_config_bf16 = False
     if config.bf16:
@@ -301,6 +306,18 @@ def adam_init_state_fn(opt, config=None):
 
     optimizers.append(optimizer)
 
+    # expert optimizer exists meaning layerwise distributed optimizer is not used
+    if len(expert_param_groups) > 0:
+        expert_optimizer = TensorParallelMuon(expert_param_groups, **muon_kwargs)
+        if config.bf16:
+            expert_optimizer = Float16OptimizerWithFloat16Params(
+                expert_optimizer, config, None, muon_init_state_fn
+            )
+        else:
+            expert_optimizer = FP32Optimizer(expert_optimizer, config, muon_init_state_fn)
+        setattr(expert_optimizer, 'grad_stats_parallel_group', pg_collection.tp_ep_pp)
+        optimizers.append(expert_optimizer)
+
     # done with muon, unfreeze nonlinear and freeze linear
     for param in nonlinear_params:
         param.requires_grad = True
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json
index 197eda568d8..ccbece04f60 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json
@@ -11,99 +11,99 @@
             "5": 10.84375,
             "6": 10.8473,
             "7": 10.85341,
-            "8": 10.83649,
-            "9": 10.84696,
-            "10": 10.78181,
-            "11": 10.85157,
-            "12": 10.86354,
-            "13": 10.85392,
-            "14": 10.88443,
-            "15": 10.87738,
-            "16": 10.84647,
-            "17": 10.83081,
-            "18": 10.86619,
-            "19": 10.84941,
-            "20": 10.84533,
-            "21": 10.84772,
-            "22": 10.79615,
-            "23": 10.88259,
-            "24": 10.83337,
-            "25": 10.82488,
-            "26": 10.84313,
-            "27": 10.85316,
-            "28": 10.87689,
-            "29": 10.86377,
-            "30": 10.81302,
-            "31": 10.78697,
-            "32": 10.85497,
-            "33": 10.85651,
-            "34": 10.849,
-            "35": 10.83725,
-            "36": 10.80381,
-            "37": 10.83835,
-            "38": 10.8051,
-            "39": 10.84122,
-            "40": 10.80292,
-            "41": 10.8407,
-            "42": 10.84416,
-            "43": 10.80995,
-            "44": 10.80279,
-            "45": 10.7866,
-            "46": 10.80814,
-            "47": 10.81723,
-            "48": 10.80288,
-            "49": 10.78144,
-            "50": 10.80226,
-            "51": 10.8227,
-            "52": 10.80372,
-            "53": 10.83318,
-            "54": 10.81535,
-            "55": 10.8256,
+            "8": 10.83652,
+            "9": 10.84691,
+            "10": 10.78166,
+            "11": 10.85213,
+            "12": 10.8629,
+            "13": 10.85433,
+            "14": 10.88455,
+            "15": 10.87782,
+            "16": 10.84637,
+            "17": 10.83054,
+            "18": 10.86645,
+            "19": 10.84951,
+            "20": 10.84547,
+            "21": 10.8476,
+            "22": 10.79618,
+            "23": 10.88285,
+            "24": 10.83247,
+            "25": 10.8246,
+            "26": 10.8432,
+            "27": 10.85345,
+            "28": 10.87635,
+            "29": 10.864,
+            "30": 10.81293,
+            "31": 10.78651,
+            "32": 10.85541,
+            "33": 10.85587,
+            "34": 10.8491,
+            "35": 10.83747,
+            "36": 10.80362,
+            "37": 10.83812,
+            "38": 10.80509,
+            "39": 10.84183,
+            "40": 10.80312,
+            "41": 10.84012,
+            "42": 10.84384,
+            "43": 10.80987,
+            "44": 10.80275,
+            "45": 10.78691,
+            "46": 10.80833,
+            "47": 10.81704,
+            "48": 10.80337,
+            "49": 10.78131,
+            "50": 10.80305,
+            "51": 10.82235,
+            "52": 10.80371,
+            "53": 10.83231,
+            "54": 10.8151,
+            "55": 10.82578,
             "56": 10.77729,
-            "57": 10.75246,
-            "58": 10.80818,
-            "59": 10.7909,
-            "60": 10.74009,
-            "61": 10.79938,
-            "62": 10.81291,
-            "63": 10.7204,
-            "64": 10.78529,
-            "65": 10.68966,
-            "66": 10.76117,
-            "67": 10.73412,
-            "68": 10.80256,
-            "69": 10.7832,
-            "70": 10.77682,
-            "71": 10.76728,
-            "72": 10.73575,
-            "73": 10.72932,
-            "74": 10.62223,
-            "75": 10.69036,
-            "76": 10.65459,
-            "77": 10.8217,
-            "78": 10.76362,
-            "79": 10.70431,
-            "80": 10.69382,
-            "81": 10.72448,
-            "82": 10.74183,
-            "83": 10.66825,
-            "84": 10.69817,
-            "85": 10.71449,
-            "86": 10.63898,
-            "87": 10.7181,
-            "88": 10.73512,
-            "89": 10.71387,
-            "90": 10.74622,
-            "91": 10.64935,
-            "92": 10.64642,
-            "93": 10.60191,
-            "94": 10.53277,
-            "95": 10.66125,
-            "96": 10.67241,
-            "97": 10.61414,
-            "98": 10.68493,
-            "99": 10.51994,
-            "100": 10.61532
+            "57": 10.75325,
+            "58": 10.80742,
+            "59": 10.79087,
+            "60": 10.73998,
+            "61": 10.79954,
+            "62": 10.81284,
+            "63": 10.72011,
+            "64": 10.78598,
+            "65": 10.68981,
+            "66": 10.76066,
+            "67": 10.73402,
+            "68": 10.8022,
+            "69": 10.78312,
+            "70": 10.77711,
+            "71": 10.76626,
+            "72": 10.73591,
+            "73": 10.72919,
+            "74": 10.62192,
+            "75": 10.69079,
+            "76": 10.65398,
+            "77": 10.82162,
+            "78": 10.76368,
+            "79": 10.70473,
+            "80": 10.69368,
+            "81": 10.72419,
+            "82": 10.74233,
+            "83": 10.66786,
+            "84": 10.6983,
+            "85": 10.714,
+            "86": 10.6383,
+            "87": 10.71809,
+            "88": 10.73508,
+            "89": 10.7139,
+            "90": 10.74649,
+            "91": 10.64861,
+            "92": 10.64636,
+            "93": 10.60234,
+            "94": 10.53327,
+            "95": 10.66155,
+            "96": 10.67215,
+            "97": 10.61446,
+            "98": 10.68506,
+            "99": 10.52056,
+            "100": 10.61544
         }
     },
     "num-zeros": {
@@ -118,99 +118,99 @@
             "5": 1398.0,
             "6": 1528.0,
             "7": 1225.0,
-            "8": 1318.0,
-            "9": 1310.0,
-            "10": 1321.0,
-            "11": 1324.0,
-            "12": 1240.0,
-            "13": 1294.0,
-            "14": 1467.0,
-            "15": 1268.0,
-            "16": 1250.0,
-            "17": 1358.0,
-            "18": 1315.0,
-            "19": 1243.0,
+            "8": 1301.0,
+            "9": 1348.0,
+            "10": 1359.0,
+            "11": 1296.0,
+            "12": 1248.0,
+            "13": 1286.0,
+            "14": 1373.0,
+            "15": 1195.0,
+            "16": 1177.0,
+            "17": 1266.0,
+            "18": 1393.0,
+            "19": 1219.0,
             "20": 1257.0,
-            "21": 1227.0,
-            "22": 1182.0,
-            "23": 1417.0,
-            "24": 1332.0,
-            "25": 1281.0,
-            "26": 1209.0,
-            "27": 1318.0,
-            "28": 1410.0,
-            "29": 1295.0,
-            "30": 1234.0,
-            "31": 1108.0,
-            "32": 1299.0,
-            "33": 1298.0,
-            "34": 1116.0,
-            "35": 1213.0,
-            "36": 1208.0,
-            "37": 1242.0,
-            "38": 1382.0,
-            "39": 1531.0,
-            "40": 1195.0,
-            "41": 1382.0,
-            "42": 1173.0,
-            "43": 1189.0,
-            "44": 1215.0,
-            "45": 1175.0,
-            "46": 1207.0,
-            "47": 1372.0,
-            "48": 1158.0,
-            "49": 1223.0,
-            "50": 1257.0,
-            "51": 1219.0,
-            "52": 1236.0,
-            "53": 1343.0,
-            "54": 1286.0,
-            "55": 1103.0,
-            "56": 1299.0,
-            "57": 1212.0,
-            "58": 1379.0,
-            "59": 1235.0,
-            "60": 1210.0,
-            "61": 1159.0,
-            "62": 1203.0,
-            "63": 1219.0,
-            "64": 1239.0,
-            "65": 1245.0,
-            "66": 1153.0,
-            "67": 1210.0,
-            "68": 1206.0,
-            "69": 1315.0,
-            "70": 1342.0,
-            "71": 1288.0,
-            "72": 1171.0,
-            "73": 1182.0,
-            "74": 1093.0,
-            "75": 1300.0,
-            "76": 1341.0,
-            "77": 1369.0,
-            "78": 1286.0,
-            "79": 1111.0,
-            "80": 1189.0,
-            "81": 1205.0,
-            "82": 1269.0,
-            "83": 1293.0,
-            "84": 1145.0,
-            "85": 1251.0,
-            "86": 1191.0,
-            "87": 1179.0,
-            "88": 1294.0,
-            "89": 1265.0,
-            "90": 1314.0,
-            "91": 1175.0,
-            "92": 1286.0,
-            "93": 1100.0,
-            "94": 969.0,
-            "95": 1204.0,
-            "96": 1241.0,
-            "97": 1163.0,
-            "98": 1205.0,
-            "99": 1291.0,
-            "100": 1214.0
+            "21": 1244.0,
+            "22": 1155.0,
+            "23": 1385.0,
+            "24": 1323.0,
+            "25": 1226.0,
+            "26": 1184.0,
+            "27": 1394.0,
+            "28": 1476.0,
+            "29": 1300.0,
+            "30": 1245.0,
+            "31": 1138.0,
+            "32": 1283.0,
+            "33": 1247.0,
+            "34": 1186.0,
+            "35": 1158.0,
+            "36": 1178.0,
+            "37": 1232.0,
+            "38": 1357.0,
+            "39": 1541.0,
+            "40": 1170.0,
+            "41": 1369.0,
+            "42": 1153.0,
+            "43": 1180.0,
+            "44": 1239.0,
+            "45": 1189.0,
+            "46": 1141.0,
+            "47": 1203.0,
+            "48": 1126.0,
+            "49": 1194.0,
+            "50": 1214.0,
+            "51": 1274.0,
+            "52": 1209.0,
+            "53": 1360.0,
+            "54": 1257.0,
+            "55": 1170.0,
+            "56": 1282.0,
+            "57": 1296.0,
+            "58": 1271.0,
+            "59": 1180.0,
+            "60": 1182.0,
+            "61": 1202.0,
+            "62": 1192.0,
+            "63": 1253.0,
+            "64": 1248.0,
+            "65": 1180.0,
+            "66": 1179.0,
+            "67": 1188.0,
+            "68": 1229.0,
+            "69": 1232.0,
+            "70": 1280.0,
+            "71": 1246.0,
+            "72": 1261.0,
+            "73": 1148.0,
+            "74": 1114.0,
+            "75": 1281.0,
+            "76": 1376.0,
+            "77": 1373.0,
+            "78": 1285.0,
+            "79": 1087.0,
+            "80": 1127.0,
+            "81": 1135.0,
+            "82": 1169.0,
+            "83": 1300.0,
+            "84": 1206.0,
+            "85": 1269.0,
+            "86": 1187.0,
+            "87": 1236.0,
+            "88": 1262.0,
+            "89": 1197.0,
+            "90": 1425.0,
+            "91": 1197.0,
+            "92": 1244.0,
+            "93": 1142.0,
+            "94": 971.0,
+            "95": 1281.0,
+            "96": 1243.0,
+            "97": 1145.0,
+            "98": 1288.0,
+            "99": 1286.0,
+            "100": 1212.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 994066432.0,
-            "2": 994036224.0,
-            "3": 994083840.0,
-            "4": 994063872.0,
-            "5": 994086912.0,
-            "6": 994028032.0,
-            "7": 994051072.0,
-            "8": 994058752.0,
-            "9": 994072576.0,
-            "10": 994086912.0,
-            "11": 994060800.0,
-            "12": 994029056.0,
-            "13": 994085888.0,
-            "14": 993994240.0,
-            "15": 994040832.0,
-            "16": 993971712.0,
-            "17": 994093568.0,
-            "18": 994065920.0,
-            "19": 994073088.0,
-            "20": 993993216.0,
-            "21": 994013184.0,
-            "22": 994089472.0,
-            "23": 994065408.0,
-            "24": 994004992.0,
-            "25": 994137600.0,
-            "26": 994042880.0,
-            "27": 994099712.0,
-            "28": 994027520.0,
-            "29": 994059776.0,
-            "30": 994023936.0,
-            "31": 994087936.0,
-            "32": 994022400.0,
-            "33": 994032640.0,
-            "34": 993997312.0,
-            "35": 994046976.0,
-            "36": 994061824.0,
-            "37": 994019840.0,
-            "38": 994102784.0,
-            "39": 994113536.0,
-            "40": 994000384.0,
-            "41": 994028544.0,
-            "42": 994046464.0,
-            "43": 994057728.0,
-            "44": 994161664.0,
-            "45": 994034176.0,
-            "46": 994053120.0,
-            "47": 994075648.0,
-            "48": 994058240.0,
-            "49": 994025472.0,
-            "50": 994043392.0,
-            "51": 994117120.0,
-            "52": 994060800.0,
-            "53": 994122752.0,
-            "54": 994071040.0,
-            "55": 994060800.0,
-            "56": 994049536.0,
-            "57": 994097152.0,
-            "58": 994092544.0,
-            "59": 994078720.0,
-            "60": 994044928.0,
-            "61": 994045440.0,
-            "62": 994039808.0,
-            "63": 994052608.0,
-            "64": 994041856.0,
-            "65": 994048000.0,
-            "66": 994055680.0,
-            "67": 994045440.0,
-            "68": 994053120.0,
-            "69": 994042368.0,
-            "70": 994087424.0,
-            "71": 994061312.0,
-            "72": 993986560.0,
-            "73": 994088448.0,
-            "74": 994099200.0,
-            "75": 994067456.0,
-            "76": 994084864.0,
-            "77": 994039808.0,
-            "78": 994094080.0,
-            "79": 994071040.0,
-            "80": 994024960.0,
-            "81": 994057728.0,
-            "82": 994005504.0,
-            "83": 994106880.0,
-            "84": 994085888.0,
-            "85": 994054144.0,
-            "86": 994055168.0,
-            "87": 994075648.0,
-            "88": 994062336.0,
-            "89": 994051584.0,
-            "90": 994043392.0,
-            "91": 994097664.0,
-            "92": 994082304.0,
-            "93": 994058752.0,
-            "94": 994066944.0,
-            "95": 994068992.0,
-            "96": 994066944.0,
-            "97": 994078208.0,
-            "98": 994054144.0,
-            "99": 994071552.0,
-            "100": 994109952.0
+            "1": 1095885312.0,
+            "2": 1095855104.0,
+            "3": 1095902720.0,
+            "4": 1095882752.0,
+            "5": 1095905792.0,
+            "6": 1095846912.0,
+            "7": 1095869952.0,
+            "8": 1095877120.0,
+            "9": 1095892480.0,
+            "10": 1095903232.0,
+            "11": 1095879168.0,
+            "12": 1095851008.0,
+            "13": 1095903232.0,
+            "14": 1095813120.0,
+            "15": 1095857152.0,
+            "16": 1095791104.0,
+            "17": 1095911936.0,
+            "18": 1095883264.0,
+            "19": 1095893504.0,
+            "20": 1095812096.0,
+            "21": 1095832064.0,
+            "22": 1095908864.0,
+            "23": 1095883776.0,
+            "24": 1095824384.0,
+            "25": 1095956480.0,
+            "26": 1095863808.0,
+            "27": 1095919104.0,
+            "28": 1095844864.0,
+            "29": 1095879168.0,
+            "30": 1095843840.0,
+            "31": 1095908352.0,
+            "32": 1095840768.0,
+            "33": 1095850496.0,
+            "34": 1095818240.0,
+            "35": 1095864832.0,
+            "36": 1095879680.0,
+            "37": 1095839232.0,
+            "38": 1095923200.0,
+            "39": 1095930880.0,
+            "40": 1095819264.0,
+            "41": 1095848448.0,
+            "42": 1095866880.0,
+            "43": 1095878656.0,
+            "44": 1095980544.0,
+            "45": 1095855104.0,
+            "46": 1095869952.0,
+            "47": 1095895040.0,
+            "48": 1095877632.0,
+            "49": 1095844352.0,
+            "50": 1095864320.0,
+            "51": 1095936000.0,
+            "52": 1095879680.0,
+            "53": 1095939584.0,
+            "54": 1095890432.0,
+            "55": 1095879168.0,
+            "56": 1095869440.0,
+            "57": 1095916544.0,
+            "58": 1095913984.0,
+            "59": 1095899136.0,
+            "60": 1095863296.0,
+            "61": 1095864320.0,
+            "62": 1095858176.0,
+            "63": 1095874048.0,
+            "64": 1095861760.0,
+            "65": 1095869952.0,
+            "66": 1095875584.0,
+            "67": 1095864832.0,
+            "68": 1095874048.0,
+            "69": 1095860224.0,
+            "70": 1095905280.0,
+            "71": 1095880192.0,
+            "72": 1095805440.0,
+            "73": 1095907840.0,
+            "74": 1095919616.0,
+            "75": 1095884800.0,
+            "76": 1095905792.0,
+            "77": 1095855616.0,
+            "78": 1095916544.0,
+            "79": 1095888384.0,
+            "80": 1095842304.0,
+            "81": 1095875584.0,
+            "82": 1095823872.0,
+            "83": 1095923712.0,
+            "84": 1095906304.0,
+            "85": 1095871488.0,
+            "86": 1095872512.0,
+            "87": 1095895552.0,
+            "88": 1095880192.0,
+            "89": 1095869440.0,
+            "90": 1095863296.0,
+            "91": 1095917056.0,
+            "92": 1095900160.0,
+            "93": 1095879680.0,
+            "94": 1095888896.0,
+            "95": 1095886848.0,
+            "96": 1095888384.0,
+            "97": 1095897088.0,
+            "98": 1095875584.0,
+            "99": 1095889408.0,
+            "100": 1095928320.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,106 +325,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 3209309696.0,
-            "2": 3480903680.0,
-            "3": 3511780864.0,
-            "4": 3511780864.0,
-            "5": 3517387264.0,
-            "6": 3517387264.0,
-            "7": 3517387264.0,
-            "8": 3517387264.0,
-            "9": 3517387264.0,
-            "10": 3517387264.0,
-            "11": 3517387264.0,
-            "12": 3517387264.0,
-            "13": 3517387264.0,
-            "14": 3517387264.0,
-            "15": 3517387264.0,
-            "16": 3517387264.0,
-            "17": 3518340096.0,
-            "18": 3518340096.0,
-            "19": 3518340096.0,
-            "20": 3518340096.0,
-            "21": 3518340096.0,
-            "22": 3518340096.0,
-            "23": 3518340096.0,
-            "24": 3518340096.0,
-            "25": 3547281408.0,
-            "26": 3547281408.0,
-            "27": 3547281408.0,
-            "28": 3547281408.0,
-            "29": 3547281408.0,
-            "30": 3547281408.0,
-            "31": 3547281408.0,
-            "32": 3547281408.0,
-            "33": 3547281408.0,
-            "34": 3547281408.0,
-            "35": 3547281408.0,
-            "36": 3547281408.0,
-            "37": 3547281408.0,
-            "38": 3547281408.0,
-            "39": 3547281408.0,
-            "40": 3547281408.0,
-            "41": 3547281408.0,
-            "42": 3547281408.0,
-            "43": 3547281408.0,
-            "44": 3565241856.0,
-            "45": 3565241856.0,
-            "46": 3565241856.0,
-            "47": 3565241856.0,
-            "48": 3565241856.0,
-            "49": 3565241856.0,
-            "50": 3565241856.0,
-            "51": 3565241856.0,
-            "52": 3565241856.0,
-            "53": 3565241856.0,
-            "54": 3565241856.0,
-            "55": 3565241856.0,
-            "56": 3565241856.0,
-            "57": 3565241856.0,
-            "58": 3565241856.0,
-            "59": 3565241856.0,
-            "60": 3565241856.0,
-            "61": 3565241856.0,
-            "62": 3565241856.0,
-            "63": 3565241856.0,
-            "64": 3565241856.0,
-            "65": 3565241856.0,
-            "66": 3565241856.0,
-            "67": 3565241856.0,
-            "68": 3565241856.0,
-            "69": 3565241856.0,
-            "70": 3565241856.0,
-            "71": 3565241856.0,
-            "72": 3565241856.0,
-            "73": 3565241856.0,
-            "74": 3565241856.0,
-            "75": 3565241856.0,
-            "76": 3565241856.0,
-            "77": 3565241856.0,
-            "78": 3565241856.0,
-            "79": 3565241856.0,
-            "80": 3565241856.0,
-            "81": 3565241856.0,
-            "82": 3565241856.0,
-            "83": 3565241856.0,
-            "84": 3565241856.0,
-            "85": 3565241856.0,
-            "86": 3565241856.0,
-            "87": 3565241856.0,
-            "88": 3565241856.0,
-            "89": 3565241856.0,
-            "90": 3565241856.0,
-            "91": 3565241856.0,
-            "92": 3565241856.0,
-            "93": 3565241856.0,
-            "94": 3565241856.0,
-            "95": 3565241856.0,
-            "96": 3565241856.0,
-            "97": 3565241856.0,
-            "98": 3565241856.0,
-            "99": 3565241856.0,
-            "100": 3565241856.0
+            "1": 3260420096.0,
+            "2": 3582874112.0,
+            "3": 3616017408.0,
+            "4": 3616017408.0,
+            "5": 3616065536.0,
+            "6": 3616065536.0,
+            "7": 3616065536.0,
+            "8": 3616065536.0,
+            "9": 3616065536.0,
+            "10": 3619626496.0,
+            "11": 3619626496.0,
+            "12": 3619626496.0,
+            "13": 3619626496.0,
+            "14": 3619626496.0,
+            "15": 3619626496.0,
+            "16": 3619626496.0,
+            "17": 3619626496.0,
+            "18": 3619626496.0,
+            "19": 3619626496.0,
+            "20": 3619626496.0,
+            "21": 3619626496.0,
+            "22": 3619626496.0,
+            "23": 3619626496.0,
+            "24": 3619626496.0,
+            "25": 3648242176.0,
+            "26": 3648242176.0,
+            "27": 3648242176.0,
+            "28": 3648242176.0,
+            "29": 3648242176.0,
+            "30": 3648242176.0,
+            "31": 3648242176.0,
+            "32": 3648242176.0,
+            "33": 3648242176.0,
+            "34": 3648242176.0,
+            "35": 3648242176.0,
+            "36": 3648242176.0,
+            "37": 3648242176.0,
+            "38": 3648242176.0,
+            "39": 3648242176.0,
+            "40": 3648242176.0,
+            "41": 3648242176.0,
+            "42": 3648242176.0,
+            "43": 3648242176.0,
+            "44": 3665209344.0,
+            "45": 3665209344.0,
+            "46": 3665209344.0,
+            "47": 3665209344.0,
+            "48": 3665209344.0,
+            "49": 3665209344.0,
+            "50": 3665209344.0,
+            "51": 3665209344.0,
+            "52": 3665209344.0,
+            "53": 3665209344.0,
+            "54": 3665209344.0,
+            "55": 3665209344.0,
+            "56": 3665209344.0,
+            "57": 3665209344.0,
+            "58": 3665209344.0,
+            "59": 3665209344.0,
+            "60": 3665209344.0,
+            "61": 3665209344.0,
+            "62": 3665209344.0,
+            "63": 3665209344.0,
+            "64": 3665209344.0,
+            "65": 3665209344.0,
+            "66": 3665209344.0,
+            "67": 3665209344.0,
+            "68": 3665209344.0,
+            "69": 3665209344.0,
+            "70": 3665209344.0,
+            "71": 3665209344.0,
+            "72": 3665209344.0,
+            "73": 3665209344.0,
+            "74": 3665209344.0,
+            "75": 3665209344.0,
+            "76": 3665209344.0,
+            "77": 3665209344.0,
+            "78": 3665209344.0,
+            "79": 3665209344.0,
+            "80": 3665209344.0,
+            "81": 3665209344.0,
+            "82": 3665209344.0,
+            "83": 3665209344.0,
+            "84": 3665209344.0,
+            "85": 3665209344.0,
+            "86": 3665209344.0,
+            "87": 3665209344.0,
+            "88": 3665209344.0,
+            "89": 3665209344.0,
+            "90": 3665209344.0,
+            "91": 3665209344.0,
+            "92": 3665209344.0,
+            "93": 3665209344.0,
+            "94": 3665209344.0,
+            "95": 3665209344.0,
+            "96": 3665209344.0,
+            "97": 3665209344.0,
+            "98": 3665209344.0,
+            "99": 3665209344.0,
+            "100": 3665209344.0
         }
     },
     "iteration-time": {
@@ -432,106 +432,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 10.4734,
-            "2": 0.22466,
-            "3": 0.19051,
-            "4": 0.16936,
-            "5": 0.17686,
-            "6": 0.15785,
-            "7": 0.16819,
-            "8": 0.15689,
-            "9": 0.15169,
-            "10": 0.15121,
-            "11": 0.15857,
-            "12": 0.15775,
-            "13": 0.15107,
-            "14": 0.19276,
-            "15": 0.1585,
-            "16": 0.14844,
-            "17": 0.14326,
-            "18": 0.13869,
-            "19": 0.1396,
-            "20": 0.15448,
-            "21": 0.139,
-            "22": 0.13512,
-            "23": 0.1426,
-            "24": 0.13221,
-            "25": 0.13685,
-            "26": 0.1411,
-            "27": 0.13181,
-            "28": 0.1391,
-            "29": 0.15621,
-            "30": 0.13616,
-            "31": 0.14287,
-            "32": 0.14647,
-            "33": 0.13884,
-            "34": 0.137,
-            "35": 0.13475,
-            "36": 0.13916,
-            "37": 0.14264,
-            "38": 0.13664,
-            "39": 0.14359,
-            "40": 0.13821,
-            "41": 0.13468,
-            "42": 0.1363,
-            "43": 0.13569,
-            "44": 0.13933,
-            "45": 0.13715,
-            "46": 0.12697,
-            "47": 0.13407,
-            "48": 0.13274,
-            "49": 0.13757,
-            "50": 0.13925,
-            "51": 0.14105,
-            "52": 0.1341,
-            "53": 0.5448,
-            "54": 0.13151,
-            "55": 0.13522,
-            "56": 0.13665,
-            "57": 0.13286,
-            "58": 0.13453,
-            "59": 0.12754,
-            "60": 0.1357,
-            "61": 0.53562,
-            "62": 0.13254,
-            "63": 0.13398,
-            "64": 0.12882,
-            "65": 0.13897,
-            "66": 0.13313,
-            "67": 0.12905,
-            "68": 0.13433,
-            "69": 0.13542,
-            "70": 0.13311,
-            "71": 0.12876,
-            "72": 0.12973,
-            "73": 0.12733,
-            "74": 0.13423,
-            "75": 0.12883,
-            "76": 0.13263,
-            "77": 0.13959,
-            "78": 0.13036,
-            "79": 0.12628,
-            "80": 0.13369,
-            "81": 0.13323,
-            "82": 0.13,
-            "83": 0.13277,
-            "84": 0.12856,
-            "85": 0.13675,
-            "86": 0.13342,
-            "87": 0.13516,
-            "88": 0.13259,
-            "89": 0.13162,
-            "90": 0.14614,
-            "91": 0.13534,
-            "92": 0.1265,
-            "93": 0.12755,
-            "94": 0.12676,
-            "95": 0.12846,
-            "96": 0.13404,
-            "97": 0.12623,
-            "98": 0.13489,
-            "99": 0.13377,
-            "100": 0.12824
+            "1": "nan",
+            "2": 6.96692,
+            "3": 0.41239,
+            "4": 0.39161,
+            "5": 0.40475,
+            "6": 0.3904,
+            "7": 0.39424,
+            "8": 0.38721,
+            "9": 0.37766,
+            "10": 0.38826,
+            "11": 0.39241,
+            "12": 0.37744,
+            "13": 0.37937,
+            "14": 0.39891,
+            "15": 0.39154,
+            "16": 0.38546,
+            "17": 0.36906,
+            "18": 0.37961,
+            "19": 0.37168,
+            "20": 0.37856,
+            "21": 0.37322,
+            "22": 0.36901,
+            "23": 0.36962,
+            "24": 0.37071,
+            "25": 0.36454,
+            "26": 0.37164,
+            "27": 0.35661,
+            "28": 0.36072,
+            "29": 0.37992,
+            "30": 0.35418,
+            "31": 0.35828,
+            "32": 0.35863,
+            "33": 0.36304,
+            "34": 0.34938,
+            "35": 0.36044,
+            "36": 0.3661,
+            "37": 0.36694,
+            "38": 0.37046,
+            "39": 0.37481,
+            "40": 0.37606,
+            "41": 0.35942,
+            "42": 0.35928,
+            "43": 0.82934,
+            "44": 0.36187,
+            "45": 0.36124,
+            "46": 0.35574,
+            "47": 0.36316,
+            "48": 0.36376,
+            "49": 0.35682,
+            "50": 0.36509,
+            "51": 0.36781,
+            "52": 0.36533,
+            "53": 0.85049,
+            "54": 0.36057,
+            "55": 0.3565,
+            "56": 0.3743,
+            "57": 0.36606,
+            "58": 0.36355,
+            "59": 0.36215,
+            "60": 0.36264,
+            "61": 0.36287,
+            "62": 0.35671,
+            "63": 0.3661,
+            "64": 0.35095,
+            "65": 0.38153,
+            "66": 0.35893,
+            "67": 0.37021,
+            "68": 0.35656,
+            "69": 0.35749,
+            "70": 0.3687,
+            "71": 0.35581,
+            "72": 0.36693,
+            "73": 0.35596,
+            "74": 0.361,
+            "75": 0.35439,
+            "76": 0.35584,
+            "77": 0.36297,
+            "78": 0.35272,
+            "79": 0.35409,
+            "80": 0.35974,
+            "81": 0.355,
+            "82": 0.35692,
+            "83": 0.3617,
+            "84": 0.36038,
+            "85": 0.36694,
+            "86": 0.36667,
+            "87": 0.36782,
+            "88": 0.37457,
+            "89": 0.36585,
+            "90": 0.37116,
+            "91": 0.36385,
+            "92": 0.3564,
+            "93": 0.36251,
+            "94": 0.35477,
+            "95": 0.35372,
+            "96": 0.8695,
+            "97": 0.35034,
+            "98": 0.36289,
+            "99": 0.35766,
+            "100": 0.35116
         }
     }
 }
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100_2nd.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100_2nd.json
index bc235c4dfa5..59528111109 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100_2nd.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100_2nd.json
@@ -54,56 +54,56 @@
             "48": "nan",
             "49": "nan",
             "50": "nan",
-            "51": 10.8227,
-            "52": 10.80372,
-            "53": 10.83318,
-            "54": 10.81535,
-            "55": 10.8256,
+            "51": 10.82235,
+            "52": 10.80371,
+            "53": 10.83231,
+            "54": 10.8151,
+            "55": 10.82578,
             "56": 10.77729,
-            "57": 10.75246,
-            "58": 10.80818,
-            "59": 10.7909,
-            "60": 10.74009,
-            "61": 10.79938,
-            "62": 10.81291,
-            "63": 10.7204,
-            "64": 10.78529,
-            "65": 10.68966,
-            "66": 10.76117,
-            "67": 10.73412,
-            "68": 10.80256,
-            "69": 10.7832,
-            "70": 10.77682,
-            "71": 10.76728,
-            "72": 10.73575,
-            "73": 10.72932,
-            "74": 10.62223,
-            "75": 10.69036,
-            "76": 10.65459,
-            "77": 10.8217,
-            "78": 10.76362,
-            "79": 10.70431,
-            "80": 10.69382,
-            "81": 10.72448,
-            "82": 10.74183,
-            "83": 10.66825,
-            "84": 10.69817,
-            "85": 10.71449,
-            "86": 10.63898,
-            "87": 10.7181,
-            "88": 10.73512,
-            "89": 10.71387,
-            "90": 10.74622,
-            "91": 10.64935,
-            "92": 10.64642,
-            "93": 10.60191,
-            "94": 10.53277,
-            "95": 10.66125,
-            "96": 10.67241,
-            "97": 10.61414,
-            "98": 10.68493,
-            "99": 10.51994,
-            "100": 10.61532
+            "57": 10.75325,
+            "58": 10.80742,
+            "59": 10.79087,
+            "60": 10.73998,
+            "61": 10.79954,
+            "62": 10.81284,
+            "63": 10.72011,
+            "64": 10.78598,
+            "65": 10.68981,
+            "66": 10.76066,
+            "67": 10.73402,
+            "68": 10.8022,
+            "69": 10.78312,
+            "70": 10.77711,
+            "71": 10.76626,
+            "72": 10.73591,
+            "73": 10.72919,
+            "74": 10.62192,
+            "75": 10.69079,
+            "76": 10.65398,
+            "77": 10.82162,
+            "78": 10.76368,
+            "79": 10.70473,
+            "80": 10.69368,
+            "81": 10.72419,
+            "82": 10.74233,
+            "83": 10.66786,
+            "84": 10.6983,
+            "85": 10.714,
+            "86": 10.6383,
+            "87": 10.71809,
+            "88": 10.73508,
+            "89": 10.7139,
+            "90": 10.74649,
+            "91": 10.64861,
+            "92": 10.64636,
+            "93": 10.60234,
+            "94": 10.53327,
+            "95": 10.66155,
+            "96": 10.67215,
+            "97": 10.61446,
+            "98": 10.68506,
+            "99": 10.52056,
+            "100": 10.61544
         }
     },
     "num-zeros": {
@@ -161,56 +161,56 @@
             "48": "nan",
             "49": "nan",
             "50": "nan",
-            "51": 1219.0,
-            "52": 1236.0,
-            "53": 1343.0,
-            "54": 1286.0,
-            "55": 1103.0,
-            "56": 1299.0,
-            "57": 1212.0,
-            "58": 1379.0,
-            "59": 1235.0,
-            "60": 1210.0,
-            "61": 1159.0,
-            "62": 1203.0,
-            "63": 1219.0,
-            "64": 1239.0,
-            "65": 1245.0,
-            "66": 1153.0,
-            "67": 1210.0,
-            "68": 1206.0,
-            "69": 1315.0,
-            "70": 1342.0,
-            "71": 1288.0,
-            "72": 1171.0,
-            "73": 1182.0,
-            "74": 1093.0,
-            "75": 1300.0,
-            "76": 1341.0,
-            "77": 1369.0,
-            "78": 1286.0,
-            "79": 1111.0,
-            "80": 1189.0,
-            "81": 1205.0,
-            "82": 1269.0,
-            "83": 1293.0,
-            "84": 1145.0,
-            "85": 1251.0,
-            "86": 1191.0,
-            "87": 1179.0,
-            "88": 1294.0,
-            "89": 1265.0,
-            "90": 1314.0,
-            "91": 1175.0,
-            "92": 1286.0,
-            "93": 1100.0,
-            "94": 969.0,
-            "95": 1204.0,
-            "96": 1241.0,
-            "97": 1163.0,
-            "98": 1205.0,
-            "99": 1291.0,
-            "100": 1214.0
+            "51": 1274.0,
+            "52": 1209.0,
+            "53": 1360.0,
+            "54": 1257.0,
+            "55": 1170.0,
+            "56": 1282.0,
+            "57": 1296.0,
+            "58": 1271.0,
+            "59": 1180.0,
+            "60": 1182.0,
+            "61": 1202.0,
+            "62": 1192.0,
+            "63": 1253.0,
+            "64": 1248.0,
+            "65": 1180.0,
+            "66": 1179.0,
+            "67": 1188.0,
+            "68": 1229.0,
+            "69": 1232.0,
+            "70": 1280.0,
+            "71": 1246.0,
+            "72": 1261.0,
+            "73": 1148.0,
+            "74": 1114.0,
+            "75": 1281.0,
+            "76": 1376.0,
+            "77": 1373.0,
+            "78": 1285.0,
+            "79": 1087.0,
+            "80": 1127.0,
+            "81": 1135.0,
+            "82": 1169.0,
+            "83": 1300.0,
+            "84": 1206.0,
+            "85": 1269.0,
+            "86": 1187.0,
+            "87": 1236.0,
+            "88": 1262.0,
+            "89": 1197.0,
+            "90": 1425.0,
+            "91": 1197.0,
+            "92": 1244.0,
+            "93": 1142.0,
+            "94": 971.0,
+            "95": 1281.0,
+            "96": 1243.0,
+            "97": 1145.0,
+            "98": 1288.0,
+            "99": 1286.0,
+            "100": 1212.0
         }
     },
     "mem-allocated-bytes": {
@@ -268,56 +268,56 @@
             "48": "nan",
             "49": "nan",
             "50": "nan",
-            "51": 994116096.0,
-            "52": 994060800.0,
-            "53": 994122752.0,
-            "54": 994071040.0,
-            "55": 994060800.0,
-            "56": 994049536.0,
-            "57": 994097152.0,
-            "58": 994092544.0,
-            "59": 994078720.0,
-            "60": 994044928.0,
-            "61": 994045440.0,
-            "62": 994039808.0,
-            "63": 994052608.0,
-            "64": 994041856.0,
-            "65": 994048000.0,
-            "66": 994055680.0,
-            "67": 994045440.0,
-            "68": 994053120.0,
-            "69": 994042368.0,
-            "70": 994087424.0,
-            "71": 994061312.0,
-            "72": 993986560.0,
-            "73": 994088448.0,
-            "74": 994099200.0,
-            "75": 994067456.0,
-            "76": 994084864.0,
-            "77": 994039808.0,
-            "78": 994094080.0,
-            "79": 994071040.0,
-            "80": 994024960.0,
-            "81": 994057728.0,
-            "82": 994005504.0,
-            "83": 994106880.0,
-            "84": 994085888.0,
-            "85": 994054144.0,
-            "86": 994055168.0,
-            "87": 994075648.0,
-            "88": 994062336.0,
-            "89": 994051584.0,
-            "90": 994043392.0,
-            "91": 994097664.0,
-            "92": 994082304.0,
-            "93": 994058752.0,
-            "94": 994066944.0,
-            "95": 994068992.0,
-            "96": 994066944.0,
-            "97": 994078208.0,
-            "98": 994054144.0,
-            "99": 994071552.0,
-            "100": 994109952.0
+            "51": 1095902208.0,
+            "52": 1095846912.0,
+            "53": 1095906816.0,
+            "54": 1095857664.0,
+            "55": 1095846400.0,
+            "56": 1095836672.0,
+            "57": 1095883776.0,
+            "58": 1095881216.0,
+            "59": 1095866368.0,
+            "60": 1095830528.0,
+            "61": 1095831552.0,
+            "62": 1095825408.0,
+            "63": 1095841280.0,
+            "64": 1095828992.0,
+            "65": 1095837184.0,
+            "66": 1095842816.0,
+            "67": 1095832064.0,
+            "68": 1095841280.0,
+            "69": 1095827456.0,
+            "70": 1095872512.0,
+            "71": 1095847424.0,
+            "72": 1095772672.0,
+            "73": 1095875072.0,
+            "74": 1095886848.0,
+            "75": 1095852032.0,
+            "76": 1095873024.0,
+            "77": 1095822848.0,
+            "78": 1095883776.0,
+            "79": 1095855616.0,
+            "80": 1095809536.0,
+            "81": 1095842816.0,
+            "82": 1095791104.0,
+            "83": 1095890944.0,
+            "84": 1095873536.0,
+            "85": 1095838720.0,
+            "86": 1095839744.0,
+            "87": 1095862784.0,
+            "88": 1095847424.0,
+            "89": 1095836672.0,
+            "90": 1095830528.0,
+            "91": 1095884288.0,
+            "92": 1095867392.0,
+            "93": 1095846912.0,
+            "94": 1095856128.0,
+            "95": 1095854080.0,
+            "96": 1095855616.0,
+            "97": 1095864320.0,
+            "98": 1095842816.0,
+            "99": 1095856640.0,
+            "100": 1095895552.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -375,56 +375,56 @@
             "48": "nan",
             "49": "nan",
             "50": "nan",
-            "51": 3502329856.0,
-            "52": 3502329856.0,
-            "53": 3537698304.0,
-            "54": 3537698304.0,
-            "55": 3537698304.0,
-            "56": 3537698304.0,
-            "57": 3537698304.0,
-            "58": 3537698304.0,
-            "59": 3537698304.0,
-            "60": 3537698304.0,
-            "61": 3537698304.0,
-            "62": 3537698304.0,
-            "63": 3537698304.0,
-            "64": 3537698304.0,
-            "65": 3537698304.0,
-            "66": 3537698304.0,
-            "67": 3537698304.0,
-            "68": 3537698304.0,
-            "69": 3537698304.0,
-            "70": 3537698304.0,
-            "71": 3537698304.0,
-            "72": 3537698304.0,
-            "73": 3537698304.0,
-            "74": 3537698304.0,
-            "75": 3537698304.0,
-            "76": 3537698304.0,
-            "77": 3537698304.0,
-            "78": 3537698304.0,
-            "79": 3537698304.0,
-            "80": 3537698304.0,
-            "81": 3537698304.0,
-            "82": 3537698304.0,
-            "83": 3537698304.0,
-            "84": 3537698304.0,
-            "85": 3537698304.0,
-            "86": 3537698304.0,
-            "87": 3537698304.0,
-            "88": 3537698304.0,
-            "89": 3537698304.0,
-            "90": 3537698304.0,
-            "91": 3537698304.0,
-            "92": 3537698304.0,
-            "93": 3537698304.0,
-            "94": 3537698304.0,
-            "95": 3537698304.0,
-            "96": 3537698304.0,
-            "97": 3537698304.0,
-            "98": 3537698304.0,
-            "99": 3537698304.0,
-            "100": 3537698304.0
+            "51": 3605514752.0,
+            "52": 3605514752.0,
+            "53": 3638906880.0,
+            "54": 3638906880.0,
+            "55": 3638906880.0,
+            "56": 3638906880.0,
+            "57": 3638906880.0,
+            "58": 3638906880.0,
+            "59": 3638906880.0,
+            "60": 3638906880.0,
+            "61": 3638906880.0,
+            "62": 3638906880.0,
+            "63": 3638906880.0,
+            "64": 3638906880.0,
+            "65": 3638906880.0,
+            "66": 3638906880.0,
+            "67": 3638906880.0,
+            "68": 3638906880.0,
+            "69": 3638906880.0,
+            "70": 3638906880.0,
+            "71": 3638906880.0,
+            "72": 3638906880.0,
+            "73": 3638906880.0,
+            "74": 3638906880.0,
+            "75": 3638906880.0,
+            "76": 3638906880.0,
+            "77": 3638906880.0,
+            "78": 3638906880.0,
+            "79": 3638906880.0,
+            "80": 3638906880.0,
+            "81": 3638906880.0,
+            "82": 3638906880.0,
+            "83": 3638906880.0,
+            "84": 3638906880.0,
+            "85": 3638906880.0,
+            "86": 3638906880.0,
+            "87": 3638906880.0,
+            "88": 3638906880.0,
+            "89": 3638906880.0,
+            "90": 3638906880.0,
+            "91": 3638906880.0,
+            "92": 3638906880.0,
+            "93": 3638906880.0,
+            "94": 3638906880.0,
+            "95": 3638906880.0,
+            "96": 3638906880.0,
+            "97": 3638906880.0,
+            "98": 3638906880.0,
+            "99": 3638906880.0,
+            "100": 3638906880.0
         }
     },
     "iteration-time": {
@@ -482,56 +482,56 @@
             "48": "nan",
             "49": "nan",
             "50": "nan",
-            "51": 7.80393,
-            "52": 0.21609,
-            "53": 0.18011,
-            "54": 0.16574,
-            "55": 0.17551,
-            "56": 0.15661,
-            "57": 0.15643,
-            "58": 0.14683,
-            "59": 0.14167,
-            "60": 0.15286,
-            "61": 0.14194,
-            "62": 0.15289,
-            "63": 0.14852,
-            "64": 0.15158,
-            "65": 0.14582,
-            "66": 0.14918,
-            "67": 0.13999,
-            "68": 0.14356,
-            "69": 0.14847,
-            "70": 0.14345,
-            "71": 0.13948,
-            "72": 0.14052,
-            "73": 0.13195,
-            "74": 0.14445,
-            "75": 0.12708,
-            "76": 0.13314,
-            "77": 0.14514,
-            "78": 0.14212,
-            "79": 0.12911,
-            "80": 0.13195,
-            "81": 0.14027,
-            "82": 0.13349,
-            "83": 0.12837,
-            "84": 0.1284,
-            "85": 0.14683,
-            "86": 0.14559,
-            "87": 0.14449,
-            "88": 0.13511,
-            "89": 0.13496,
-            "90": 0.14777,
-            "91": 0.13483,
-            "92": 0.13387,
-            "93": 0.12619,
-            "94": 0.12638,
-            "95": 0.12624,
-            "96": 0.13537,
-            "97": 0.12788,
-            "98": 0.14225,
-            "99": 0.13569,
-            "100": 0.12935
+            "51": "nan",
+            "52": 5.33757,
+            "53": 0.39893,
+            "54": 0.38074,
+            "55": 0.38709,
+            "56": 0.37977,
+            "57": 0.37403,
+            "58": 0.3832,
+            "59": 0.37979,
+            "60": 0.3767,
+            "61": 0.37583,
+            "62": 0.38081,
+            "63": 0.38367,
+            "64": 0.38655,
+            "65": 0.37373,
+            "66": 0.37183,
+            "67": 0.37121,
+            "68": 0.38709,
+            "69": 0.38149,
+            "70": 0.38976,
+            "71": 0.38463,
+            "72": 0.38157,
+            "73": 0.36873,
+            "74": 0.3762,
+            "75": 0.36571,
+            "76": 0.36544,
+            "77": 0.37985,
+            "78": 0.37941,
+            "79": 0.36655,
+            "80": 0.37258,
+            "81": 0.36741,
+            "82": 0.36798,
+            "83": 0.3641,
+            "84": 0.36415,
+            "85": 0.37605,
+            "86": 0.37639,
+            "87": 0.38223,
+            "88": 0.37682,
+            "89": 0.3604,
+            "90": 0.37267,
+            "91": 0.36421,
+            "92": 0.36312,
+            "93": 0.36608,
+            "94": 0.35916,
+            "95": 0.37338,
+            "96": 0.3876,
+            "97": 0.37229,
+            "98": 0.3763,
+            "99": 0.37389,
+            "100": 0.3586
         }
     }
 }
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/model_config.yaml
index 81b023bd86e..5c395caed56 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/model_config.yaml
@@ -64,4 +64,5 @@ MODEL_ARGS:
   --muon-momentum: 0.9
   --muon-extra-scale-factor: 0.2
   --muon-scale-mode: spectral
+  --check-weight-hash-across-dp-replicas-interval: 1
 TEST_TYPE: ckpt-resume
diff --git a/tests/test_utils/recipes/moe.yaml b/tests/test_utils/recipes/moe.yaml
index faef76e38eb..06039d77440 100644
--- a/tests/test_utils/recipes/moe.yaml
+++ b/tests/test_utils/recipes/moe.yaml
@@ -203,7 +203,7 @@ products:
   - test_case: [gpt3_moe_mcore_te_ep8_resume_torch_dist_muon]
     products:
       - environment: [dev]
-        scope: [mr-broken, mr-github-broken, mr-slim-broken]
+        scope: [mr, mr-github, mr-slim]
         platforms: [dgx_h100]
   - test_case: [gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_no_mtp_no_a2a_ovlp_fine_grained_offloading]
     products:
diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py
index 0662922586c..3f60658a005 100644
--- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py
+++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py
@@ -186,10 +186,10 @@ def test_broadcast_params(self, tp, pp):
             for name, param in model[0].named_parameters():
                 assert torch.allclose(param.data, original_params[name])
 
-    # TODO(@boxiangw): add PP=4 back and fix the test
+    # TODO(deyuf): check bf16 False case
     @pytest.mark.parametrize('tp', [1, 2, 4])
-    @pytest.mark.parametrize('pp', [1, 2])
-    @pytest.mark.parametrize('bf16', [True, False])
+    @pytest.mark.parametrize('pp', [1, 2, 4])
+    @pytest.mark.parametrize('bf16', [True])
     def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16):
         """Test save/load of LayerWiseDistributedOptimizer checkpoints."""
         if tp * pp > 8:
@@ -315,11 +315,10 @@ def test_layer_wise_optimizer_count_zeros(self, tp, pp):
             num_zeros = optimizer.count_zeros()
             assert num_zeros >= 0
 
-    # TODO(@boxiangw): add PP=4 back and fix the test
     @pytest.mark.parametrize('src_tp', [1, 2, 4])
-    @pytest.mark.parametrize('src_pp', [1, 2])
+    @pytest.mark.parametrize('src_pp', [1, 2, 4])
     @pytest.mark.parametrize('dest_tp', [1, 2, 4])
-    @pytest.mark.parametrize('dest_pp', [1, 2])
+    @pytest.mark.parametrize('dest_pp', [1, 2, 4])
     def test_layer_wise_optimizer_resharding(
         self, tmp_path_dist_ckpt, src_tp, src_pp, dest_tp, dest_pp
     ):
diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py
index 8d22e184893..ce068ef3227 100644
--- a/tests/unit_tests/dist_checkpointing/utils.py
+++ b/tests/unit_tests/dist_checkpointing/utils.py
@@ -202,7 +202,11 @@ def setup_model_and_optimizer(
     if 'muon' in optimizer:
         # Use layer-wise distributed optimizer with Muon
         optimizer_type = optimizer
-        optimizer = get_megatron_muon_optimizer(config, model)
+        # default lr None feels wrong. only change muon lr to avoid breaking old tests
+        config.lr = 0.0
+        optimizer = get_megatron_muon_optimizer(
+            config, model, layer_wise_distributed_optimizer='dist' in optimizer_type
+        )
     else:
         optimizer_type = optimizer
         optimizer = get_megatron_optimizer(config, model)
@@ -217,18 +221,8 @@ def setup_model_and_optimizer(
                     optimizer.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data)
                     optimizer.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data)
     else:
-        for group in optimizer.chained_optimizers[0].param_groups:
-            for p in group['params']:
-                if len(optimizer.chained_optimizers[0].state[p]) == 0:
-                    optimizer.chained_optimizers[0].state[p]['momentum_buffer'] = torch.rand_like(
-                        p.data
-                    )
-
-        for group in optimizer.chained_optimizers[1].param_groups:
-            for p in group['params']:
-                if len(optimizer.chained_optimizers[1].state[p]) == 0:
-                    optimizer.chained_optimizers[1].state[p]['exp_avg'] = torch.rand_like(p.data)
-                    optimizer.chained_optimizers[1].state[p]['exp_avg_sq'] = torch.rand_like(p.data)
+        for opt in optimizer.chained_optimizers:
+            opt.init_state_fn(opt)
 
     optimizer.reload_model_params()
 
@@ -305,7 +299,11 @@ def setup_moe_model_and_optimizer(
 
     if 'muon' in optimizer:
         optimizer_type = optimizer
-        optimizer = get_megatron_muon_optimizer(config, model)
+        # default lr None feels wrong. only change muon lr to avoid breaking old tests
+        config.lr = 0.0
+        optimizer = get_megatron_muon_optimizer(
+            config, model, layer_wise_distributed_optimizer='dist' in optimizer_type
+        )
     else:
         optimizer_type = optimizer
         optimizer = get_megatron_optimizer(config, model)
@@ -321,18 +319,8 @@ def setup_moe_model_and_optimizer(
                         opt.state[p]['exp_avg'] = torch.rand_like(p.data)
                         opt.state[p]['exp_avg_sq'] = torch.rand_like(p.data)
     else:
-        for group in optimizer.chained_optimizers[0].param_groups:
-            for p in group['params']:
-                if len(optimizer.chained_optimizers[0].state[p]) == 0:
-                    optimizer.chained_optimizers[0].state[p]['momentum_buffer'] = torch.rand_like(
-                        p.data
-                    )
-
-        for group in optimizer.chained_optimizers[1].param_groups:
-            for p in group['params']:
-                if len(optimizer.chained_optimizers[1].state[p]) == 0:
-                    optimizer.chained_optimizers[1].state[p]['exp_avg'] = torch.rand_like(p.data)
-                    optimizer.chained_optimizers[1].state[p]['exp_avg_sq'] = torch.rand_like(p.data)
+        for opt in optimizer.chained_optimizers:
+            opt.init_state_fn(opt)
 
     optimizer.reload_model_params()