diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py
index 61374a4ae08..c86dc30939e 100644
--- a/megatron/training/initialize.py
+++ b/megatron/training/initialize.py
@@ -14,6 +14,7 @@
 
 from megatron.core import mpu, tensor_parallel
 from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train
+from megatron.core.fusions.fused_bias_geglu import bias_geglu
 from megatron.core.fusions.fused_bias_gelu import bias_gelu
 from megatron.core.fusions.fused_bias_swiglu import bias_swiglu
 from megatron.core.parallel_state import create_group
@@ -526,57 +527,126 @@ def _warmup_jit_function(tp_size=None):
     else:
         dtype = torch.float32
 
-    # Warmup fused bias+gelu
-    bias = torch.rand(
-        args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda"
+    # Check if TE activation function is used (in which case, no need to warmup custom fusions)
+    use_te_activation_func = getattr(args, 'use_te_activation_func', False)
+    gated_linear_unit = getattr(args, 'gated_linear_unit', False)
+
+    # Warmup bias_swiglu: swiglu activation (F.silu + GLU)
+    warmup_bias_swiglu = (
+        not use_te_activation_func
+        and args.swiglu
+        and args.bias_swiglu_fusion
     )
-    input = torch.rand(
-        (
-            args.seq_length // args.context_parallel_size,
-            args.micro_batch_size,
-            args.ffn_hidden_size // args.tensor_model_parallel_size,
-        ),
-        dtype=dtype,
-        device="cuda",
+
+    warmup_bias_gelu = (
+        not use_te_activation_func
+        and not args.swiglu
+        and not getattr(args, 'quick_geglu', False)
+        and not gated_linear_unit
+        and args.bias_gelu_fusion
+    )
+
+    warmup_bias_geglu = (
+        not use_te_activation_func
+        and not args.swiglu
+        and not getattr(args, 'quick_geglu', False)
+        and gated_linear_unit
+        and args.bias_gelu_fusion
     )
-    # Warmup JIT fusions with the input grad_enable state of both forward
-    # prop and recomputation
-    for bias_grad, input_grad in zip([True, True], [False, True]):
-        bias.requires_grad, input.requires_grad = bias_grad, input_grad
-        for _ in range(5):
-            if args.swiglu:
+
+    if warmup_bias_swiglu:
+        bias = torch.rand(
+            args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda"
+        )
+        input = torch.rand(
+            (
+                args.seq_length // args.context_parallel_size,
+                args.micro_batch_size,
+                args.ffn_hidden_size // args.tensor_model_parallel_size,
+            ),
+            dtype=dtype,
+            device="cuda",
+        )
+        # Warmup JIT fusions with the input grad_enable state of both forward
+        # prop and recomputation
+        for bias_grad, input_grad in zip([True, True], [False, True]):
+            bias.requires_grad, input.requires_grad = bias_grad, input_grad
+            for _ in range(5):
                 output = bias_swiglu(input, bias)
-            else:
+        del bias, input, output
+
+    if warmup_bias_gelu:
+        bias = torch.rand(
+            args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda"
+        )
+        input = torch.rand(
+            (
+                args.seq_length // args.context_parallel_size,
+                args.micro_batch_size,
+                args.ffn_hidden_size // args.tensor_model_parallel_size,
+            ),
+            dtype=dtype,
+            device="cuda",
+        )
+        # Warmup JIT fusions with the input grad_enable state of both forward
+        # prop and recomputation
+        for bias_grad, input_grad in zip([True, True], [False, True]):
+            bias.requires_grad, input.requires_grad = bias_grad, input_grad
+            for _ in range(5):
                 output = bias_gelu(bias, input)
-    del bias, input, output
+        del bias, input, output
+
+    # Warmup fused bias+geglu (gated gelu)
+    if warmup_bias_geglu:
+        # For geglu, input size is 2x ffn_hidden_size (will be split into two halves)
+        bias = torch.rand(
+            (args.ffn_hidden_size // args.tensor_model_parallel_size) * 2, dtype=dtype, device="cuda"
+        )
+        input = torch.rand(
+            (
+                args.seq_length // args.context_parallel_size,
+                args.micro_batch_size,
+                (args.ffn_hidden_size // args.tensor_model_parallel_size) * 2,
+            ),
+            dtype=dtype,
+            device="cuda",
+        )
+        # Warmup JIT fusions with the input grad_enable state of both forward
+        # prop and recomputation
+        for bias_grad, input_grad in zip([True, True], [False, True]):
+            bias.requires_grad, input.requires_grad = bias_grad, input_grad
+            for _ in range(5):
+                output = bias_geglu(bias, input)
+        del bias, input, output
 
     # Warmup fused bias+dropout+add
-    if args.sequence_parallel:
-        # tp_size threaded by the caller (hetero MIMO language PGC); None -> mpu.
-        seq_length = args.seq_length // (tp_size or mpu.get_tensor_model_parallel_world_size())
-    else:
-        seq_length = args.seq_length
-    input = torch.rand(
-        (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size),
-        dtype=dtype,
-        device="cuda",
-    )
-    residual = torch.rand(
-        (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size),
-        dtype=dtype,
-        device="cuda",
-    )
-    bias = torch.rand((args.hidden_size), dtype=dtype, device="cuda").expand_as(residual)
-    dropout_rate = 0.1
-    # Warmup JIT fusions with the input grad_enable state of both forward
-    # prop and recomputation
-    for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]):
-        input.requires_grad = input_grad
-        bias.requires_grad = bias_grad
-        residual.requires_grad = residual_grad
-        for _ in range(5):
-            output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate)
-    del bias, input, residual, output
+    if args.bias_dropout_fusion:
+        if args.sequence_parallel:
+            # tp_size threaded by the caller (hetero MIMO language PGC); None -> mpu.
+            seq_length = args.seq_length // (tp_size or mpu.get_tensor_model_parallel_world_size())
+        else:
+            seq_length = args.seq_length
+        input = torch.rand(
+            (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size),
+            dtype=dtype,
+            device="cuda",
+        )
+        residual = torch.rand(
+            (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size),
+            dtype=dtype,
+            device="cuda",
+        )
+        bias = torch.rand((args.hidden_size), dtype=dtype, device="cuda").expand_as(residual)
+        dropout_rate = 0.1
+        # Warmup JIT fusions with the input grad_enable state of both forward
+        # prop and recomputation
+        for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]):
+            input.requires_grad = input_grad
+            bias.requires_grad = bias_grad
+            residual.requires_grad = residual_grad
+            for _ in range(5):
+                output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate)
+        del bias, input, residual, output
     torch.cuda.empty_cache()
 
 
diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_disable_bias_linear/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_disable_bias_linear/golden_values_dev_dgx_h100.json
index a93909a5a3e..3e60347bcd6 100644
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_disable_bias_linear/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_disable_bias_linear/golden_values_dev_dgx_h100.json
@@ -4,106 +4,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 10.95945,
-            "2": 10.94558,
-            "3": 10.94922,
-            "4": 10.94704,
-            "5": 10.95589,
-            "6": 10.95218,
-            "7": 10.95016,
-            "8": 10.95689,
-            "9": 10.94261,
-            "10": 10.94724,
-            "11": 10.93546,
-            "12": 10.94741,
-            "13": 10.91988,
-            "14": 10.92077,
-            "15": 10.90467,
-            "16": 10.8924,
-            "17": 10.89279,
-            "18": 10.89248,
-            "19": 10.87973,
-            "20": 10.80736,
-            "21": 10.78846,
-            "22": 10.77833,
-            "23": 10.78262,
-            "24": 10.75945,
-            "25": 10.75778,
-            "26": 10.73468,
-            "27": 10.69867,
-            "28": 10.61066,
-            "29": 10.57956,
-            "30": 10.56539,
-            "31": 10.56486,
-            "32": 10.55161,
-            "33": 10.52219,
-            "34": 10.47136,
-            "35": 10.48355,
-            "36": 10.46232,
-            "37": 10.42496,
-            "38": 10.42078,
-            "39": 10.39103,
-            "40": 10.37295,
-            "41": 10.35282,
-            "42": 10.35034,
-            "43": 10.31696,
-            "44": 10.30514,
-            "45": 10.29693,
-            "46": 10.26616,
-            "47": 10.2552,
-            "48": 10.21043,
-            "49": 10.2023,
-            "50": 10.21715,
-            "51": 10.21037,
-            "52": 10.1583,
-            "53": 10.15998,
-            "54": 10.13463,
-            "55": 10.11084,
-            "56": 10.12869,
-            "57": 10.11006,
-            "58": 10.12183,
-            "59": 10.06883,
-            "60": 10.09,
-            "61": 10.0453,
-            "62": 10.01564,
-            "63": 10.08031,
-            "64": 10.03238,
-            "65": 10.00175,
-            "66": 10.04593,
-            "67": 10.01595,
-            "68": 9.98437,
-            "69": 9.99683,
-            "70": 9.98003,
-            "71": 10.00129,
-            "72": 9.98841,
-            "73": 9.97265,
-            "74": 9.96398,
-            "75": 9.92912,
-            "76": 9.95785,
-            "77": 9.95872,
-            "78": 9.91029,
-            "79": 9.90855,
-            "80": 9.92441,
-            "81": 9.94842,
-            "82": 9.88868,
-            "83": 9.85308,
-            "84": 9.79712,
-            "85": 9.77773,
-            "86": 9.88383,
-            "87": 9.90616,
-            "88": 9.87921,
-            "89": 9.8183,
-            "90": 9.81557,
-            "91": 9.81758,
-            "92": 9.81601,
-            "93": 9.74139,
-            "94": 9.82066,
-            "95": 9.80836,
-            "96": 9.79432,
-            "97": 9.73962,
-            "98": 9.76311,
-            "99": 9.81735,
-            "100": 9.70263
+            "1": 10.95936,
+            "2": 10.94977,
+            "3": 10.94747,
+            "4": 10.94712,
+            "5": 10.95723,
+            "6": 10.9502,
+            "7": 10.95021,
+            "8": 10.95367,
+            "9": 10.9446,
+            "10": 10.95154,
+            "11": 10.93386,
+            "12": 10.94603,
+            "13": 10.92327,
+            "14": 10.92309,
+            "15": 10.90119,
+            "16": 10.8969,
+            "17": 10.89052,
+            "18": 10.89608,
+            "19": 10.88286,
+            "20": 10.80497,
+            "21": 10.79195,
+            "22": 10.77939,
+            "23": 10.78411,
+            "24": 10.76047,
+            "25": 10.75806,
+            "26": 10.73441,
+            "27": 10.6978,
+            "28": 10.61418,
+            "29": 10.58081,
+            "30": 10.56574,
+            "31": 10.56554,
+            "32": 10.55299,
+            "33": 10.5197,
+            "34": 10.47359,
+            "35": 10.48528,
+            "36": 10.46148,
+            "37": 10.42264,
+            "38": 10.42219,
+            "39": 10.38888,
+            "40": 10.37047,
+            "41": 10.35242,
+            "42": 10.34888,
+            "43": 10.31667,
+            "44": 10.30525,
+            "45": 10.2977,
+            "46": 10.26829,
+            "47": 10.2535,
+            "48": 10.21113,
+            "49": 10.20128,
+            "50": 10.21746,
+            "51": 10.20821,
+            "52": 10.1564,
+            "53": 10.16158,
+            "54": 10.13459,
+            "55": 10.10982,
+            "56": 10.12859,
+            "57": 10.1118,
+            "58": 10.12143,
+            "59": 10.06785,
+            "60": 10.09061,
+            "61": 10.04499,
+            "62": 10.01656,
+            "63": 10.08005,
+            "64": 10.03135,
+            "65": 10.00238,
+            "66": 10.04605,
+            "67": 10.01486,
+            "68": 9.98422,
+            "69": 9.99663,
+            "70": 9.97922,
+            "71": 10.00131,
+            "72": 9.98786,
+            "73": 9.97376,
+            "74": 9.96469,
+            "75": 9.93111,
+            "76": 9.95904,
+            "77": 9.95783,
+            "78": 9.91063,
+            "79": 9.90915,
+            "80": 9.9249,
+            "81": 9.94881,
+            "82": 9.8881,
+            "83": 9.85484,
+            "84": 9.79857,
+            "85": 9.77855,
+            "86": 9.88669,
+            "87": 9.90709,
+            "88": 9.87985,
+            "89": 9.82084,
+            "90": 9.81725,
+            "91": 9.8175,
+            "92": 9.81545,
+            "93": 9.74382,
+            "94": 9.82121,
+            "95": 9.80871,
+            "96": 9.793,
+            "97": 9.74031,
+            "98": 9.76505,
+            "99": 9.81886,
+            "100": 9.70162
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1757.0,
-            "2": 1634.0,
-            "3": 1778.0,
-            "4": 1689.0,
-            "5": 1718.0,
-            "6": 1665.0,
-            "7": 1932.0,
-            "8": 1688.0,
-            "9": 1691.0,
-            "10": 1675.0,
-            "11": 1611.0,
-            "12": 1641.0,
-            "13": 1752.0,
-            "14": 1772.0,
-            "15": 1704.0,
-            "16": 1707.0,
-            "17": 1774.0,
-            "18": 1653.0,
-            "19": 1654.0,
-            "20": 1674.0,
-            "21": 1884.0,
-            "22": 1717.0,
-            "23": 1682.0,
-            "24": 1842.0,
-            "25": 1672.0,
-            "26": 1827.0,
-            "27": 1781.0,
-            "28": 1787.0,
-            "29": 1874.0,
-            "30": 1756.0,
-            "31": 2047.0,
-            "32": 2000.0,
-            "33": 2016.0,
-            "34": 1927.0,
-            "35": 2059.0,
-            "36": 1953.0,
-            "37": 2219.0,
-            "38": 2079.0,
-            "39": 2287.0,
-            "40": 2267.0,
-            "41": 2353.0,
-            "42": 1984.0,
-            "43": 2366.0,
-            "44": 2237.0,
-            "45": 2406.0,
-            "46": 2437.0,
-            "47": 2490.0,
-            "48": 2610.0,
-            "49": 2824.0,
-            "50": 2462.0,
-            "51": 2550.0,
-            "52": 2650.0,
-            "53": 2604.0,
-            "54": 2782.0,
-            "55": 2379.0,
-            "56": 2607.0,
-            "57": 2254.0,
-            "58": 3543.0,
-            "59": 2766.0,
-            "60": 2903.0,
-            "61": 2703.0,
-            "62": 3189.0,
-            "63": 3372.0,
-            "64": 3554.0,
-            "65": 2665.0,
-            "66": 2967.0,
-            "67": 3732.0,
-            "68": 3289.0,
-            "69": 2882.0,
-            "70": 3230.0,
-            "71": 3124.0,
-            "72": 2888.0,
-            "73": 3497.0,
-            "74": 3070.0,
-            "75": 2934.0,
-            "76": 3294.0,
-            "77": 3651.0,
-            "78": 3172.0,
-            "79": 3231.0,
-            "80": 3072.0,
-            "81": 3623.0,
-            "82": 3038.0,
-            "83": 3112.0,
-            "84": 2955.0,
-            "85": 2647.0,
-            "86": 2960.0,
-            "87": 2632.0,
-            "88": 2986.0,
-            "89": 3002.0,
-            "90": 3706.0,
-            "91": 2999.0,
-            "92": 2941.0,
-            "93": 2999.0,
-            "94": 3127.0,
-            "95": 3168.0,
-            "96": 3607.0,
-            "97": 3401.0,
-            "98": 3405.0,
-            "99": 2933.0,
-            "100": 3025.0
+            "1": 1779.0,
+            "2": 1668.0,
+            "3": 1719.0,
+            "4": 1654.0,
+            "5": 1739.0,
+            "6": 1625.0,
+            "7": 1871.0,
+            "8": 1645.0,
+            "9": 1732.0,
+            "10": 1654.0,
+            "11": 1697.0,
+            "12": 1653.0,
+            "13": 1718.0,
+            "14": 1825.0,
+            "15": 1683.0,
+            "16": 1694.0,
+            "17": 1886.0,
+            "18": 1692.0,
+            "19": 1741.0,
+            "20": 1704.0,
+            "21": 1831.0,
+            "22": 1682.0,
+            "23": 1629.0,
+            "24": 1832.0,
+            "25": 1591.0,
+            "26": 1754.0,
+            "27": 1798.0,
+            "28": 1753.0,
+            "29": 1918.0,
+            "30": 1767.0,
+            "31": 2018.0,
+            "32": 1910.0,
+            "33": 2002.0,
+            "34": 2029.0,
+            "35": 2038.0,
+            "36": 1916.0,
+            "37": 2145.0,
+            "38": 2119.0,
+            "39": 2198.0,
+            "40": 2248.0,
+            "41": 2260.0,
+            "42": 1968.0,
+            "43": 2382.0,
+            "44": 2305.0,
+            "45": 2434.0,
+            "46": 2368.0,
+            "47": 2512.0,
+            "48": 2607.0,
+            "49": 2838.0,
+            "50": 2557.0,
+            "51": 2402.0,
+            "52": 2732.0,
+            "53": 2656.0,
+            "54": 2783.0,
+            "55": 2472.0,
+            "56": 2646.0,
+            "57": 2252.0,
+            "58": 3497.0,
+            "59": 2854.0,
+            "60": 2888.0,
+            "61": 2604.0,
+            "62": 3158.0,
+            "63": 3183.0,
+            "64": 3523.0,
+            "65": 2699.0,
+            "66": 3039.0,
+            "67": 3633.0,
+            "68": 3306.0,
+            "69": 2957.0,
+            "70": 3161.0,
+            "71": 3219.0,
+            "72": 2908.0,
+            "73": 3396.0,
+            "74": 3264.0,
+            "75": 2916.0,
+            "76": 3230.0,
+            "77": 3594.0,
+            "78": 3303.0,
+            "79": 3248.0,
+            "80": 2964.0,
+            "81": 3511.0,
+            "82": 3041.0,
+            "83": 3208.0,
+            "84": 2930.0,
+            "85": 2776.0,
+            "86": 2952.0,
+            "87": 2636.0,
+            "88": 3042.0,
+            "89": 3049.0,
+            "90": 3692.0,
+            "91": 3095.0,
+            "92": 2939.0,
+            "93": 3070.0,
+            "94": 3150.0,
+            "95": 3379.0,
+            "96": 3524.0,
+            "97": 3558.0,
+            "98": 3252.0,
+            "99": 3121.0,
+            "100": 3043.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 747244032.0,
-            "2": 747244032.0,
-            "3": 747244032.0,
-            "4": 747244032.0,
-            "5": 747244032.0,
-            "6": 747244032.0,
-            "7": 747244032.0,
-            "8": 747244032.0,
-            "9": 747244032.0,
-            "10": 747244032.0,
-            "11": 747244032.0,
-            "12": 747244032.0,
-            "13": 747244032.0,
-            "14": 747244032.0,
-            "15": 747244032.0,
-            "16": 747244032.0,
-            "17": 747244032.0,
-            "18": 747244032.0,
-            "19": 747244032.0,
-            "20": 747244032.0,
-            "21": 747244032.0,
-            "22": 747244032.0,
-            "23": 747244032.0,
-            "24": 747244032.0,
-            "25": 747244032.0,
-            "26": 747244032.0,
-            "27": 747244032.0,
-            "28": 747244032.0,
-            "29": 747244032.0,
-            "30": 747244032.0,
-            "31": 747244032.0,
-            "32": 747244032.0,
-            "33": 747244032.0,
-            "34": 747244032.0,
-            "35": 747244032.0,
-            "36": 747244032.0,
-            "37": 747244032.0,
-            "38": 747244032.0,
-            "39": 747244032.0,
-            "40": 747244032.0,
-            "41": 747244032.0,
-            "42": 747244032.0,
-            "43": 747244032.0,
-            "44": 747244032.0,
-            "45": 747244032.0,
-            "46": 747244032.0,
-            "47": 747244032.0,
-            "48": 747244032.0,
-            "49": 747244032.0,
-            "50": 747244032.0,
-            "51": 747244032.0,
-            "52": 747244032.0,
-            "53": 747244032.0,
-            "54": 747244032.0,
-            "55": 747244032.0,
-            "56": 747244032.0,
-            "57": 747244032.0,
-            "58": 747244032.0,
-            "59": 747244032.0,
-            "60": 747244032.0,
-            "61": 747244032.0,
-            "62": 747244032.0,
-            "63": 747244032.0,
-            "64": 747244032.0,
-            "65": 747244032.0,
-            "66": 747244032.0,
-            "67": 747244032.0,
-            "68": 747244032.0,
-            "69": 747244032.0,
-            "70": 747244032.0,
-            "71": 747244032.0,
-            "72": 747244032.0,
-            "73": 747244032.0,
-            "74": 747244032.0,
-            "75": 747244032.0,
-            "76": 747244032.0,
-            "77": 747244032.0,
-            "78": 747244032.0,
-            "79": 747244032.0,
-            "80": 747244032.0,
-            "81": 747244032.0,
-            "82": 747244032.0,
-            "83": 747244032.0,
-            "84": 747244032.0,
-            "85": 747244032.0,
-            "86": 747244032.0,
-            "87": 747244032.0,
-            "88": 747244032.0,
-            "89": 747244032.0,
-            "90": 747244032.0,
-            "91": 747244032.0,
-            "92": 747244032.0,
-            "93": 747244032.0,
-            "94": 747244032.0,
-            "95": 747244032.0,
-            "96": 747244032.0,
-            "97": 747244032.0,
-            "98": 747244032.0,
-            "99": 747244032.0,
-            "100": 747244032.0
+            "1": 747243520.0,
+            "2": 747243520.0,
+            "3": 747243520.0,
+            "4": 747243520.0,
+            "5": 747243520.0,
+            "6": 747243520.0,
+            "7": 747243520.0,
+            "8": 747243520.0,
+            "9": 747243520.0,
+            "10": 747243520.0,
+            "11": 747243520.0,
+            "12": 747243520.0,
+            "13": 747243520.0,
+            "14": 747243520.0,
+            "15": 747243520.0,
+            "16": 747243520.0,
+            "17": 747243520.0,
+            "18": 747243520.0,
+            "19": 747243520.0,
+            "20": 747243520.0,
+            "21": 747243520.0,
+            "22": 747243520.0,
+            "23": 747243520.0,
+            "24": 747243520.0,
+            "25": 747243520.0,
+            "26": 747243520.0,
+            "27": 747243520.0,
+            "28": 747243520.0,
+            "29": 747243520.0,
+            "30": 747243520.0,
+            "31": 747243520.0,
+            "32": 747243520.0,
+            "33": 747243520.0,
+            "34": 747243520.0,
+            "35": 747243520.0,
+            "36": 747243520.0,
+            "37": 747243520.0,
+            "38": 747243520.0,
+            "39": 747243520.0,
+            "40": 747243520.0,
+            "41": 747243520.0,
+            "42": 747243520.0,
+            "43": 747243520.0,
+            "44": 747243520.0,
+            "45": 747243520.0,
+            "46": 747243520.0,
+            "47": 747243520.0,
+            "48": 747243520.0,
+            "49": 747243520.0,
+            "50": 747243520.0,
+            "51": 747243520.0,
+            "52": 747243520.0,
+            "53": 747243520.0,
+            "54": 747243520.0,
+            "55": 747243520.0,
+            "56": 747243520.0,
+            "57": 747243520.0,
+            "58": 747243520.0,
+            "59": 747243520.0,
+            "60": 747243520.0,
+            "61": 747243520.0,
+            "62": 747243520.0,
+            "63": 747243520.0,
+            "64": 747243520.0,
+            "65": 747243520.0,
+            "66": 747243520.0,
+            "67": 747243520.0,
+            "68": 747243520.0,
+            "69": 747243520.0,
+            "70": 747243520.0,
+            "71": 747243520.0,
+            "72": 747243520.0,
+            "73": 747243520.0,
+            "74": 747243520.0,
+            "75": 747243520.0,
+            "76": 747243520.0,
+            "77": 747243520.0,
+            "78": 747243520.0,
+            "79": 747243520.0,
+            "80": 747243520.0,
+            "81": 747243520.0,
+            "82": 747243520.0,
+            "83": 747243520.0,
+            "84": 747243520.0,
+            "85": 747243520.0,
+            "86": 747243520.0,
+            "87": 747243520.0,
+            "88": 747243520.0,
+            "89": 747243520.0,
+            "90": 747243520.0,
+            "91": 747243520.0,
+            "92": 747243520.0,
+            "93": 747243520.0,
+            "94": 747243520.0,
+            "95": 747243520.0,
+            "96": 747243520.0,
+            "97": 747243520.0,
+            "98": 747243520.0,
+            "99": 747243520.0,
+            "100": 747243520.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,212 +325,213 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1927202816.0,
-            "2": 2211948544.0,
-            "3": 2211948544.0,
-            "4": 2211948544.0,
-            "5": 2211948544.0,
-            "6": 2211948544.0,
-            "7": 2211948544.0,
-            "8": 2211948544.0,
-            "9": 2211948544.0,
-            "10": 2211948544.0,
-            "11": 2211948544.0,
-            "12": 2211948544.0,
-            "13": 2211948544.0,
-            "14": 2211948544.0,
-            "15": 2211948544.0,
-            "16": 2211948544.0,
-            "17": 2211948544.0,
-            "18": 2211948544.0,
-            "19": 2211948544.0,
-            "20": 2211948544.0,
-            "21": 2211948544.0,
-            "22": 2211948544.0,
-            "23": 2211948544.0,
-            "24": 2211948544.0,
-            "25": 2211948544.0,
-            "26": 2211948544.0,
-            "27": 2211948544.0,
-            "28": 2211948544.0,
-            "29": 2211948544.0,
-            "30": 2211948544.0,
-            "31": 2211948544.0,
-            "32": 2211948544.0,
-            "33": 2211948544.0,
-            "34": 2211948544.0,
-            "35": 2211948544.0,
-            "36": 2211948544.0,
-            "37": 2211948544.0,
-            "38": 2211948544.0,
-            "39": 2211948544.0,
-            "40": 2211948544.0,
-            "41": 2211948544.0,
-            "42": 2211948544.0,
-            "43": 2211948544.0,
-            "44": 2211948544.0,
-            "45": 2211948544.0,
-            "46": 2211948544.0,
-            "47": 2211948544.0,
-            "48": 2211948544.0,
-            "49": 2211948544.0,
-            "50": 2211948544.0,
-            "51": 2211948544.0,
-            "52": 2211948544.0,
-            "53": 2211948544.0,
-            "54": 2211948544.0,
-            "55": 2211948544.0,
-            "56": 2211948544.0,
-            "57": 2211948544.0,
-            "58": 2211948544.0,
-            "59": 2211948544.0,
-            "60": 2211948544.0,
-            "61": 2211948544.0,
-            "62": 2211948544.0,
-            "63": 2211948544.0,
-            "64": 2211948544.0,
-            "65": 2211948544.0,
-            "66": 2211948544.0,
-            "67": 2211948544.0,
-            "68": 2211948544.0,
-            "69": 2211948544.0,
-            "70": 2211948544.0,
-            "71": 2211948544.0,
-            "72": 2211948544.0,
-            "73": 2211948544.0,
-            "74": 2211948544.0,
-            "75": 2211948544.0,
-            "76": 2211948544.0,
-            "77": 2211948544.0,
-            "78": 2211948544.0,
-            "79": 2211948544.0,
-            "80": 2211948544.0,
-            "81": 2211948544.0,
-            "82": 2211948544.0,
-            "83": 2211948544.0,
-            "84": 2211948544.0,
-            "85": 2211948544.0,
-            "86": 2211948544.0,
-            "87": 2211948544.0,
-            "88": 2211948544.0,
-            "89": 2211948544.0,
-            "90": 2211948544.0,
-            "91": 2211948544.0,
-            "92": 2211948544.0,
-            "93": 2211948544.0,
-            "94": 2211948544.0,
-            "95": 2211948544.0,
-            "96": 2211948544.0,
-            "97": 2211948544.0,
-            "98": 2211948544.0,
-            "99": 2211948544.0,
-            "100": 2211948544.0
+            "1": 1927136768.0,
+            "2": 2211882496.0,
+            "3": 2211882496.0,
+            "4": 2211882496.0,
+            "5": 2211882496.0,
+            "6": 2211882496.0,
+            "7": 2211882496.0,
+            "8": 2211882496.0,
+            "9": 2211882496.0,
+            "10": 2211882496.0,
+            "11": 2211882496.0,
+            "12": 2211882496.0,
+            "13": 2211882496.0,
+            "14": 2211882496.0,
+            "15": 2211882496.0,
+            "16": 2211882496.0,
+            "17": 2211882496.0,
+            "18": 2211882496.0,
+            "19": 2211882496.0,
+            "20": 2211882496.0,
+            "21": 2211882496.0,
+            "22": 2211882496.0,
+            "23": 2211882496.0,
+            "24": 2211882496.0,
+            "25": 2211882496.0,
+            "26": 2211882496.0,
+            "27": 2211882496.0,
+            "28": 2211882496.0,
+            "29": 2211882496.0,
+            "30": 2211882496.0,
+            "31": 2211882496.0,
+            "32": 2211882496.0,
+            "33": 2211882496.0,
+            "34": 2211882496.0,
+            "35": 2211882496.0,
+            "36": 2211882496.0,
+            "37": 2211882496.0,
+            "38": 2211882496.0,
+            "39": 2211882496.0,
+            "40": 2211882496.0,
+            "41": 2211882496.0,
+            "42": 2211882496.0,
+            "43": 2211882496.0,
+            "44": 2211882496.0,
+            "45": 2211882496.0,
+            "46": 2211882496.0,
+            "47": 2211882496.0,
+            "48": 2211882496.0,
+            "49": 2211882496.0,
+            "50": 2211882496.0,
+            "51": 2211882496.0,
+            "52": 2211882496.0,
+            "53": 2211882496.0,
+            "54": 2211882496.0,
+            "55": 2211882496.0,
+            "56": 2211882496.0,
+            "57": 2211882496.0,
+            "58": 2211882496.0,
+            "59": 2211882496.0,
+            "60": 2211882496.0,
+            "61": 2211882496.0,
+            "62": 2211882496.0,
+            "63": 2211882496.0,
+            "64": 2211882496.0,
+            "65": 2211882496.0,
+            "66": 2211882496.0,
+            "67": 2211882496.0,
+            "68": 2211882496.0,
+            "69": 2211882496.0,
+            "70": 2211882496.0,
+            "71": 2211882496.0,
+            "72": 2211882496.0,
+            "73": 2211882496.0,
+            "74": 2211882496.0,
+            "75": 2211882496.0,
+            "76": 2211882496.0,
+            "77": 2211882496.0,
+            "78": 2211882496.0,
+            "79": 2211882496.0,
+            "80": 2211882496.0,
+            "81": 2211882496.0,
+            "82": 2211882496.0,
+            "83": 2211882496.0,
+            "84": 2211882496.0,
+            "85": 2211882496.0,
+            "86": 2211882496.0,
+            "87": 2211882496.0,
+            "88": 2211882496.0,
+            "89": 2211882496.0,
+            "90": 2211882496.0,
+            "91": 2211882496.0,
+            "92": 2211882496.0,
+            "93": 2211882496.0,
+            "94": 2211882496.0,
+            "95": 2211882496.0,
+            "96": 2211882496.0,
+            "97": 2211882496.0,
+            "98": 2211882496.0,
+            "99": 2211882496.0,
+            "100": 2211882496.0
         }
     },
     "iteration-time": {
-        "start_step": 2,
+        "start_step": 1,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "2": 7.47187,
-            "3": 0.12583,
-            "4": 0.12513,
-            "5": 0.12312,
-            "6": 0.12002,
-            "7": 0.12266,
-            "8": 0.12083,
-            "9": 0.12109,
-            "10": 0.12778,
-            "11": 0.12,
-            "12": 0.11977,
-            "13": 0.1209,
-            "14": 0.1222,
-            "15": 0.11966,
-            "16": 0.1203,
-            "17": 0.12112,
-            "18": 0.12089,
-            "19": 0.11949,
-            "20": 0.12319,
-            "21": 0.11956,
-            "22": 0.12175,
-            "23": 0.11982,
-            "24": 0.11915,
-            "25": 0.11883,
-            "26": 0.1211,
-            "27": 0.11897,
-            "28": 0.11868,
-            "29": 0.1187,
-            "30": 0.1185,
-            "31": 0.11891,
-            "32": 0.12011,
-            "33": 0.11819,
-            "34": 0.11956,
-            "35": 0.11916,
-            "36": 0.12349,
-            "37": 0.12261,
-            "38": 0.12218,
-            "39": 0.12343,
-            "40": 0.12294,
-            "41": 0.12401,
-            "42": 0.13039,
-            "43": 0.12135,
-            "44": 0.12159,
-            "45": 0.12059,
-            "46": 0.12055,
-            "47": 0.12061,
-            "48": 0.12205,
-            "49": 0.12274,
-            "50": 0.12152,
-            "51": 0.15983,
-            "52": 0.13886,
-            "53": 0.12789,
-            "54": 0.12438,
-            "55": 0.12402,
-            "56": 0.12066,
-            "57": 0.12431,
-            "58": 0.12428,
-            "59": 0.1225,
-            "60": 0.12216,
-            "61": 0.12287,
-            "62": 0.1233,
-            "63": 0.12246,
-            "64": 0.12156,
-            "65": 0.12249,
-            "66": 0.12311,
-            "67": 0.12081,
-            "68": 0.12265,
-            "69": 0.12004,
-            "70": 0.12046,
-            "71": 0.12475,
-            "72": 0.12229,
-            "73": 0.1218,
-            "74": 0.12074,
-            "75": 0.12139,
-            "76": 0.1207,
-            "77": 0.12323,
-            "78": 0.12329,
-            "79": 0.12234,
-            "80": 0.12214,
-            "81": 0.12205,
-            "82": 0.12476,
-            "83": 0.12372,
-            "84": 0.12074,
-            "85": 0.1208,
-            "86": 0.12117,
-            "87": 0.11997,
-            "88": 0.1188,
-            "89": 0.12144,
-            "90": 0.1191,
-            "91": 0.12011,
-            "92": 0.12118,
-            "93": 0.1193,
-            "94": 0.11984,
-            "95": 0.11981,
-            "96": 0.1197,
-            "97": 0.11967,
-            "98": 0.1204,
-            "99": 0.12098,
-            "100": 0.12041
+            "1": "nan",
+            "2": 6.28856,
+            "3": 2.87408,
+            "4": 1.95928,
+            "5": 1.5832,
+            "6": 1.80267,
+            "7": 2.17865,
+            "8": 2.56834,
+            "9": 2.73831,
+            "10": 1.5937,
+            "11": 1.72001,
+            "12": 1.96787,
+            "13": 1.84992,
+            "14": 2.44652,
+            "15": 1.40887,
+            "16": 1.79479,
+            "17": 1.93826,
+            "18": 2.18185,
+            "19": 1.86202,
+            "20": 1.7889,
+            "21": 1.70263,
+            "22": 1.98469,
+            "23": 2.17113,
+            "24": 2.69296,
+            "25": 1.11094,
+            "26": 1.25066,
+            "27": 2.01477,
+            "28": 1.69372,
+            "29": 2.18649,
+            "30": 1.07977,
+            "31": 1.91029,
+            "32": 3.14707,
+            "33": 1.75155,
+            "34": 0.84322,
+            "35": 2.32734,
+            "36": 1.91784,
+            "37": 1.44452,
+            "38": 1.49868,
+            "39": 1.31399,
+            "40": 2.20354,
+            "41": 1.48409,
+            "42": 2.12089,
+            "43": 1.23836,
+            "44": 2.4032,
+            "45": 1.79923,
+            "46": 1.60394,
+            "47": 1.77471,
+            "48": 1.58995,
+            "49": 1.32347,
+            "50": 1.54706,
+            "51": 0.14886,
+            "52": 2.22136,
+            "53": 2.01456,
+            "54": 1.83129,
+            "55": 2.43884,
+            "56": 1.4282,
+            "57": 1.85358,
+            "58": 1.1592,
+            "59": 1.54137,
+            "60": 1.78592,
+            "61": 1.80887,
+            "62": 1.94828,
+            "63": 1.24083,
+            "64": 2.131,
+            "65": 1.1498,
+            "66": 1.88453,
+            "67": 1.86127,
+            "68": 1.87149,
+            "69": 1.80235,
+            "70": 2.09491,
+            "71": 2.06033,
+            "72": 1.32738,
+            "73": 1.8966,
+            "74": 0.96839,
+            "75": 1.45872,
+            "76": 1.85607,
+            "77": 1.79916,
+            "78": 2.08855,
+            "79": 2.44964,
+            "80": 2.41197,
+            "81": 0.61018,
+            "82": 1.88508,
+            "83": 1.20543,
+            "84": 1.68087,
+            "85": 0.95404,
+            "86": 1.82326,
+            "87": 2.0686,
+            "88": 1.53418,
+            "89": 1.84138,
+            "90": 1.85917,
+            "91": 1.53667,
+            "92": 2.66045,
+            "93": 0.92927,
+            "94": 1.63442,
+            "95": 1.86353,
+            "96": 1.14845,
+            "97": 1.63437,
+            "98": 1.18653,
+            "99": 2.43032,
+            "100": 1.43562
         }
     }
 }
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_persistent_disable_bias_linear/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_persistent_disable_bias_linear/golden_values_dev_dgx_h100.json
index 1eb20ba8556..2d247e92917 100644
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_persistent_disable_bias_linear/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_resume_torch_dist_persistent_disable_bias_linear/golden_values_dev_dgx_h100.json
@@ -4,106 +4,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 10.95945,
-            "2": 10.94558,
-            "3": 10.94922,
-            "4": 10.94704,
-            "5": 10.95589,
-            "6": 10.95218,
-            "7": 10.95016,
-            "8": 10.95689,
-            "9": 10.94261,
-            "10": 10.94724,
-            "11": 10.93546,
-            "12": 10.94741,
-            "13": 10.91988,
-            "14": 10.92077,
-            "15": 10.90467,
-            "16": 10.8924,
-            "17": 10.89279,
-            "18": 10.89248,
-            "19": 10.87973,
-            "20": 10.80736,
-            "21": 10.78846,
-            "22": 10.77833,
-            "23": 10.78262,
-            "24": 10.75945,
-            "25": 10.75778,
-            "26": 10.73468,
-            "27": 10.69867,
-            "28": 10.61066,
-            "29": 10.57956,
-            "30": 10.56539,
-            "31": 10.56486,
-            "32": 10.55161,
-            "33": 10.52219,
-            "34": 10.47136,
-            "35": 10.48355,
-            "36": 10.46232,
-            "37": 10.42496,
-            "38": 10.42078,
-            "39": 10.39103,
-            "40": 10.37295,
-            "41": 10.35282,
-            "42": 10.35034,
-            "43": 10.31696,
-            "44": 10.30514,
-            "45": 10.29693,
-            "46": 10.26616,
-            "47": 10.2552,
-            "48": 10.21043,
-            "49": 10.2023,
-            "50": 10.21715,
-            "51": 10.21037,
-            "52": 10.1583,
-            "53": 10.15998,
-            "54": 10.13463,
-            "55": 10.11084,
-            "56": 10.12869,
-            "57": 10.11006,
-            "58": 10.12183,
-            "59": 10.06883,
-            "60": 10.09,
-            "61": 10.0453,
-            "62": 10.01564,
-            "63": 10.08031,
-            "64": 10.03238,
-            "65": 10.00175,
-            "66": 10.04593,
-            "67": 10.01595,
-            "68": 9.98437,
-            "69": 9.99683,
-            "70": 9.98003,
-            "71": 10.00129,
-            "72": 9.98841,
-            "73": 9.97265,
-            "74": 9.96398,
-            "75": 9.92912,
-            "76": 9.95785,
-            "77": 9.95872,
-            "78": 9.91029,
-            "79": 9.90855,
-            "80": 9.92441,
-            "81": 9.94842,
-            "82": 9.88868,
-            "83": 9.85308,
-            "84": 9.79712,
-            "85": 9.77773,
-            "86": 9.88383,
-            "87": 9.90616,
-            "88": 9.87921,
-            "89": 9.8183,
-            "90": 9.81557,
-            "91": 9.81758,
-            "92": 9.81601,
-            "93": 9.74139,
-            "94": 9.82066,
-            "95": 9.80836,
-            "96": 9.79432,
-            "97": 9.73962,
-            "98": 9.76311,
-            "99": 9.81735,
-            "100": 9.70263
+            "1": 10.95936,
+            "2": 10.94977,
+            "3": 10.94747,
+            "4": 10.94712,
+            "5": 10.95723,
+            "6": 10.9502,
+            "7": 10.95021,
+            "8": 10.95367,
+            "9": 10.9446,
+            "10": 10.95154,
+            "11": 10.93386,
+            "12": 10.94603,
+            "13": 10.92327,
+            "14": 10.92309,
+            "15": 10.90119,
+            "16": 10.8969,
+            "17": 10.89052,
+            "18": 10.89608,
+            "19": 10.88286,
+            "20": 10.80497,
+            "21": 10.79195,
+            "22": 10.77939,
+            "23": 10.78411,
+            "24": 10.76047,
+            "25": 10.75806,
+            "26": 10.73441,
+            "27": 10.6978,
+            "28": 10.61418,
+            "29": 10.58081,
+            "30": 10.56574,
+            "31": 10.56554,
+            "32": 10.55299,
+            "33": 10.5197,
+            "34": 10.47359,
+            "35": 10.48528,
+            "36": 10.46148,
+            "37": 10.42264,
+            "38": 10.42219,
+            "39": 10.38888,
+            "40": 10.37047,
+            "41": 10.35242,
+            "42": 10.34888,
+            "43": 10.31667,
+            "44": 10.30525,
+            "45": 10.2977,
+            "46": 10.26829,
+            "47": 10.2535,
+            "48": 10.21113,
+            "49": 10.20128,
+            "50": 10.21746,
+            "51": 10.20821,
+            "52": 10.1564,
+            "53": 10.16158,
+            "54": 10.13459,
+            "55": 10.10982,
+            "56": 10.12859,
+            "57": 10.1118,
+            "58": 10.12143,
+            "59": 10.06785,
+            "60": 10.09061,
+            "61": 10.04499,
+            "62": 10.01656,
+            "63": 10.08005,
+            "64": 10.03135,
+            "65": 10.00238,
+            "66": 10.04605,
+            "67": 10.01486,
+            "68": 9.98422,
+            "69": 9.99663,
+            "70": 9.97922,
+            "71": 10.00131,
+            "72": 9.98786,
+            "73": 9.97376,
+            "74": 9.96469,
+            "75": 9.93111,
+            "76": 9.95904,
+            "77": 9.95783,
+            "78": 9.91063,
+            "79": 9.90915,
+            "80": 9.9249,
+            "81": 9.94881,
+            "82": 9.8881,
+            "83": 9.85484,
+            "84": 9.79857,
+            "85": 9.77855,
+            "86": 9.88669,
+            "87": 9.90709,
+            "88": 9.87985,
+            "89": 9.82084,
+            "90": 9.81725,
+            "91": 9.8175,
+            "92": 9.81545,
+            "93": 9.74382,
+            "94": 9.82121,
+            "95": 9.80871,
+            "96": 9.793,
+            "97": 9.74031,
+            "98": 9.76505,
+            "99": 9.81886,
+            "100": 9.70162
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1757.0,
-            "2": 1634.0,
-            "3": 1778.0,
-            "4": 1689.0,
-            "5": 1718.0,
-            "6": 1665.0,
-            "7": 1932.0,
-            "8": 1688.0,
-            "9": 1691.0,
-            "10": 1675.0,
-            "11": 1611.0,
-            "12": 1641.0,
-            "13": 1752.0,
-            "14": 1772.0,
-            "15": 1704.0,
-            "16": 1707.0,
-            "17": 1774.0,
-            "18": 1653.0,
-            "19": 1654.0,
-            "20": 1674.0,
-            "21": 1884.0,
-            "22": 1717.0,
-            "23": 1682.0,
-            "24": 1842.0,
-            "25": 1672.0,
-            "26": 1827.0,
-            "27": 1781.0,
-            "28": 1787.0,
-            "29": 1874.0,
-            "30": 1756.0,
-            "31": 2047.0,
-            "32": 2000.0,
-            "33": 2016.0,
-            "34": 1927.0,
-            "35": 2059.0,
-            "36": 1953.0,
-            "37": 2219.0,
-            "38": 2079.0,
-            "39": 2287.0,
-            "40": 2267.0,
-            "41": 2353.0,
-            "42": 1984.0,
-            "43": 2366.0,
-            "44": 2237.0,
-            "45": 2406.0,
-            "46": 2437.0,
-            "47": 2490.0,
-            "48": 2610.0,
-            "49": 2824.0,
-            "50": 2462.0,
-            "51": 2550.0,
-            "52": 2650.0,
-            "53": 2604.0,
-            "54": 2782.0,
-            "55": 2379.0,
-            "56": 2607.0,
-            "57": 2254.0,
-            "58": 3543.0,
-            "59": 2766.0,
-            "60": 2903.0,
-            "61": 2703.0,
-            "62": 3189.0,
-            "63": 3372.0,
-            "64": 3554.0,
-            "65": 2665.0,
-            "66": 2967.0,
-            "67": 3732.0,
-            "68": 3289.0,
-            "69": 2882.0,
-            "70": 3230.0,
-            "71": 3124.0,
-            "72": 2888.0,
-            "73": 3497.0,
-            "74": 3070.0,
-            "75": 2934.0,
-            "76": 3294.0,
-            "77": 3651.0,
-            "78": 3172.0,
-            "79": 3231.0,
-            "80": 3072.0,
-            "81": 3623.0,
-            "82": 3038.0,
-            "83": 3112.0,
-            "84": 2955.0,
-            "85": 2647.0,
-            "86": 2960.0,
-            "87": 2632.0,
-            "88": 2986.0,
-            "89": 3002.0,
-            "90": 3706.0,
-            "91": 2999.0,
-            "92": 2941.0,
-            "93": 2999.0,
-            "94": 3127.0,
-            "95": 3168.0,
-            "96": 3607.0,
-            "97": 3401.0,
-            "98": 3405.0,
-            "99": 2933.0,
-            "100": 3025.0
+            "1": 1779.0,
+            "2": 1668.0,
+            "3": 1719.0,
+            "4": 1654.0,
+            "5": 1739.0,
+            "6": 1625.0,
+            "7": 1871.0,
+            "8": 1645.0,
+            "9": 1732.0,
+            "10": 1654.0,
+            "11": 1697.0,
+            "12": 1653.0,
+            "13": 1718.0,
+            "14": 1825.0,
+            "15": 1683.0,
+            "16": 1694.0,
+            "17": 1886.0,
+            "18": 1692.0,
+            "19": 1741.0,
+            "20": 1704.0,
+            "21": 1831.0,
+            "22": 1682.0,
+            "23": 1629.0,
+            "24": 1832.0,
+            "25": 1591.0,
+            "26": 1754.0,
+            "27": 1798.0,
+            "28": 1753.0,
+            "29": 1918.0,
+            "30": 1767.0,
+            "31": 2018.0,
+            "32": 1910.0,
+            "33": 2002.0,
+            "34": 2029.0,
+            "35": 2038.0,
+            "36": 1916.0,
+            "37": 2145.0,
+            "38": 2119.0,
+            "39": 2198.0,
+            "40": 2248.0,
+            "41": 2260.0,
+            "42": 1968.0,
+            "43": 2382.0,
+            "44": 2305.0,
+            "45": 2434.0,
+            "46": 2368.0,
+            "47": 2512.0,
+            "48": 2607.0,
+            "49": 2838.0,
+            "50": 2557.0,
+            "51": 2402.0,
+            "52": 2732.0,
+            "53": 2656.0,
+            "54": 2783.0,
+            "55": 2472.0,
+            "56": 2646.0,
+            "57": 2252.0,
+            "58": 3497.0,
+            "59": 2854.0,
+            "60": 2888.0,
+            "61": 2604.0,
+            "62": 3158.0,
+            "63": 3183.0,
+            "64": 3523.0,
+            "65": 2699.0,
+            "66": 3039.0,
+            "67": 3633.0,
+            "68": 3306.0,
+            "69": 2957.0,
+            "70": 3161.0,
+            "71": 3219.0,
+            "72": 2908.0,
+            "73": 3396.0,
+            "74": 3264.0,
+            "75": 2916.0,
+            "76": 3230.0,
+            "77": 3594.0,
+            "78": 3303.0,
+            "79": 3248.0,
+            "80": 2964.0,
+            "81": 3511.0,
+            "82": 3041.0,
+            "83": 3208.0,
+            "84": 2930.0,
+            "85": 2776.0,
+            "86": 2952.0,
+            "87": 2636.0,
+            "88": 3042.0,
+            "89": 3049.0,
+            "90": 3692.0,
+            "91": 3095.0,
+            "92": 2939.0,
+            "93": 3070.0,
+            "94": 3150.0,
+            "95": 3379.0,
+            "96": 3524.0,
+            "97": 3558.0,
+            "98": 3252.0,
+            "99": 3121.0,
+            "100": 3043.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 747244032.0,
-            "2": 747244032.0,
-            "3": 747244032.0,
-            "4": 747244032.0,
-            "5": 747244032.0,
-            "6": 747244032.0,
-            "7": 747244032.0,
-            "8": 747244032.0,
-            "9": 747244032.0,
-            "10": 747244032.0,
-            "11": 747244032.0,
-            "12": 747244032.0,
-            "13": 747244032.0,
-            "14": 747244032.0,
-            "15": 747244032.0,
-            "16": 747244032.0,
-            "17": 747244032.0,
-            "18": 747244032.0,
-            "19": 747244032.0,
-            "20": 747244032.0,
-            "21": 747244032.0,
-            "22": 747244032.0,
-            "23": 747244032.0,
-            "24": 747244032.0,
-            "25": 747244032.0,
-            "26": 747244032.0,
-            "27": 747244032.0,
-            "28": 747244032.0,
-            "29": 747244032.0,
-            "30": 747244032.0,
-            "31": 747244032.0,
-            "32": 747244032.0,
-            "33": 747244032.0,
-            "34": 747244032.0,
-            "35": 747244032.0,
-            "36": 747244032.0,
-            "37": 747244032.0,
-            "38": 747244032.0,
-            "39": 747244032.0,
-            "40": 747244032.0,
-            "41": 747244032.0,
-            "42": 747244032.0,
-            "43": 747244032.0,
-            "44": 747244032.0,
-            "45": 747244032.0,
-            "46": 747244032.0,
-            "47": 747244032.0,
-            "48": 747244032.0,
-            "49": 747244032.0,
-            "50": 747244032.0,
-            "51": 747244032.0,
-            "52": 747244032.0,
-            "53": 747244032.0,
-            "54": 747244032.0,
-            "55": 747244032.0,
-            "56": 747244032.0,
-            "57": 747244032.0,
-            "58": 747244032.0,
-            "59": 747244032.0,
-            "60": 747244032.0,
-            "61": 747244032.0,
-            "62": 747244032.0,
-            "63": 747244032.0,
-            "64": 747244032.0,
-            "65": 747244032.0,
-            "66": 747244032.0,
-            "67": 747244032.0,
-            "68": 747244032.0,
-            "69": 747244032.0,
-            "70": 747244032.0,
-            "71": 747244032.0,
-            "72": 747244032.0,
-            "73": 747244032.0,
-            "74": 747244032.0,
-            "75": 747244032.0,
-            "76": 747244032.0,
-            "77": 747244032.0,
-            "78": 747244032.0,
-            "79": 747244032.0,
-            "80": 747244032.0,
-            "81": 747244032.0,
-            "82": 747244032.0,
-            "83": 747244032.0,
-            "84": 747244032.0,
-            "85": 747244032.0,
-            "86": 747244032.0,
-            "87": 747244032.0,
-            "88": 747244032.0,
-            "89": 747244032.0,
-            "90": 747244032.0,
-            "91": 747244032.0,
-            "92": 747244032.0,
-            "93": 747244032.0,
-            "94": 747244032.0,
-            "95": 747244032.0,
-            "96": 747244032.0,
-            "97": 747244032.0,
-            "98": 747244032.0,
-            "99": 747244032.0,
-            "100": 747244032.0
+            "1": 747243520.0,
+            "2": 747243520.0,
+            "3": 747243520.0,
+            "4": 747243520.0,
+            "5": 747243520.0,
+            "6": 747243520.0,
+            "7": 747243520.0,
+            "8": 747243520.0,
+            "9": 747243520.0,
+            "10": 747243520.0,
+            "11": 747243520.0,
+            "12": 747243520.0,
+            "13": 747243520.0,
+            "14": 747243520.0,
+            "15": 747243520.0,
+            "16": 747243520.0,
+            "17": 747243520.0,
+            "18": 747243520.0,
+            "19": 747243520.0,
+            "20": 747243520.0,
+            "21": 747243520.0,
+            "22": 747243520.0,
+            "23": 747243520.0,
+            "24": 747243520.0,
+            "25": 747243520.0,
+            "26": 747243520.0,
+            "27": 747243520.0,
+            "28": 747243520.0,
+            "29": 747243520.0,
+            "30": 747243520.0,
+            "31": 747243520.0,
+            "32": 747243520.0,
+            "33": 747243520.0,
+            "34": 747243520.0,
+            "35": 747243520.0,
+            "36": 747243520.0,
+            "37": 747243520.0,
+            "38": 747243520.0,
+            "39": 747243520.0,
+            "40": 747243520.0,
+            "41": 747243520.0,
+            "42": 747243520.0,
+            "43": 747243520.0,
+            "44": 747243520.0,
+            "45": 747243520.0,
+            "46": 747243520.0,
+            "47": 747243520.0,
+            "48": 747243520.0,
+            "49": 747243520.0,
+            "50": 747243520.0,
+            "51": 747243520.0,
+            "52": 747243520.0,
+            "53": 747243520.0,
+            "54": 747243520.0,
+            "55": 747243520.0,
+            "56": 747243520.0,
+            "57": 747243520.0,
+            "58": 747243520.0,
+            "59": 747243520.0,
+            "60": 747243520.0,
+            "61": 747243520.0,
+            "62": 747243520.0,
+            "63": 747243520.0,
+            "64": 747243520.0,
+            "65": 747243520.0,
+            "66": 747243520.0,
+            "67": 747243520.0,
+            "68": 747243520.0,
+            "69": 747243520.0,
+            "70": 747243520.0,
+            "71": 747243520.0,
+            "72": 747243520.0,
+            "73": 747243520.0,
+            "74": 747243520.0,
+            "75": 747243520.0,
+            "76": 747243520.0,
+            "77": 747243520.0,
+            "78": 747243520.0,
+            "79": 747243520.0,
+            "80": 747243520.0,
+            "81": 747243520.0,
+            "82": 747243520.0,
+            "83": 747243520.0,
+            "84": 747243520.0,
+            "85": 747243520.0,
+            "86": 747243520.0,
+            "87": 747243520.0,
+            "88": 747243520.0,
+            "89": 747243520.0,
+            "90": 747243520.0,
+            "91": 747243520.0,
+            "92": 747243520.0,
+            "93": 747243520.0,
+            "94": 747243520.0,
+            "95": 747243520.0,
+            "96": 747243520.0,
+            "97": 747243520.0,
+            "98": 747243520.0,
+            "99": 747243520.0,
+            "100": 747243520.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,212 +325,213 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1927202816.0,
-            "2": 2211948544.0,
-            "3": 2211948544.0,
-            "4": 2211948544.0,
-            "5": 2211948544.0,
-            "6": 2211948544.0,
-            "7": 2211948544.0,
-            "8": 2211948544.0,
-            "9": 2211948544.0,
-            "10": 2211948544.0,
-            "11": 2211948544.0,
-            "12": 2211948544.0,
-            "13": 2211948544.0,
-            "14": 2211948544.0,
-            "15": 2211948544.0,
-            "16": 2211948544.0,
-            "17": 2211948544.0,
-            "18": 2211948544.0,
-            "19": 2211948544.0,
-            "20": 2211948544.0,
-            "21": 2211948544.0,
-            "22": 2211948544.0,
-            "23": 2211948544.0,
-            "24": 2211948544.0,
-            "25": 2211948544.0,
-            "26": 2211948544.0,
-            "27": 2211948544.0,
-            "28": 2211948544.0,
-            "29": 2211948544.0,
-            "30": 2211948544.0,
-            "31": 2211948544.0,
-            "32": 2211948544.0,
-            "33": 2211948544.0,
-            "34": 2211948544.0,
-            "35": 2211948544.0,
-            "36": 2211948544.0,
-            "37": 2211948544.0,
-            "38": 2211948544.0,
-            "39": 2211948544.0,
-            "40": 2211948544.0,
-            "41": 2211948544.0,
-            "42": 2211948544.0,
-            "43": 2211948544.0,
-            "44": 2211948544.0,
-            "45": 2211948544.0,
-            "46": 2211948544.0,
-            "47": 2211948544.0,
-            "48": 2211948544.0,
-            "49": 2211948544.0,
-            "50": 2211948544.0,
-            "51": 2211948544.0,
-            "52": 2211948544.0,
-            "53": 2211948544.0,
-            "54": 2211948544.0,
-            "55": 2211948544.0,
-            "56": 2211948544.0,
-            "57": 2211948544.0,
-            "58": 2211948544.0,
-            "59": 2211948544.0,
-            "60": 2211948544.0,
-            "61": 2211948544.0,
-            "62": 2211948544.0,
-            "63": 2211948544.0,
-            "64": 2211948544.0,
-            "65": 2211948544.0,
-            "66": 2211948544.0,
-            "67": 2211948544.0,
-            "68": 2211948544.0,
-            "69": 2211948544.0,
-            "70": 2211948544.0,
-            "71": 2211948544.0,
-            "72": 2211948544.0,
-            "73": 2211948544.0,
-            "74": 2211948544.0,
-            "75": 2211948544.0,
-            "76": 2211948544.0,
-            "77": 2211948544.0,
-            "78": 2211948544.0,
-            "79": 2211948544.0,
-            "80": 2211948544.0,
-            "81": 2211948544.0,
-            "82": 2211948544.0,
-            "83": 2211948544.0,
-            "84": 2211948544.0,
-            "85": 2211948544.0,
-            "86": 2211948544.0,
-            "87": 2211948544.0,
-            "88": 2211948544.0,
-            "89": 2211948544.0,
-            "90": 2211948544.0,
-            "91": 2211948544.0,
-            "92": 2211948544.0,
-            "93": 2211948544.0,
-            "94": 2211948544.0,
-            "95": 2211948544.0,
-            "96": 2211948544.0,
-            "97": 2211948544.0,
-            "98": 2211948544.0,
-            "99": 2211948544.0,
-            "100": 2211948544.0
+            "1": 1927136768.0,
+            "2": 2211882496.0,
+            "3": 2211882496.0,
+            "4": 2211882496.0,
+            "5": 2211882496.0,
+            "6": 2211882496.0,
+            "7": 2211882496.0,
+            "8": 2211882496.0,
+            "9": 2211882496.0,
+            "10": 2211882496.0,
+            "11": 2211882496.0,
+            "12": 2211882496.0,
+            "13": 2211882496.0,
+            "14": 2211882496.0,
+            "15": 2211882496.0,
+            "16": 2211882496.0,
+            "17": 2211882496.0,
+            "18": 2211882496.0,
+            "19": 2211882496.0,
+            "20": 2211882496.0,
+            "21": 2211882496.0,
+            "22": 2211882496.0,
+            "23": 2211882496.0,
+            "24": 2211882496.0,
+            "25": 2211882496.0,
+            "26": 2211882496.0,
+            "27": 2211882496.0,
+            "28": 2211882496.0,
+            "29": 2211882496.0,
+            "30": 2211882496.0,
+            "31": 2211882496.0,
+            "32": 2211882496.0,
+            "33": 2211882496.0,
+            "34": 2211882496.0,
+            "35": 2211882496.0,
+            "36": 2211882496.0,
+            "37": 2211882496.0,
+            "38": 2211882496.0,
+            "39": 2211882496.0,
+            "40": 2211882496.0,
+            "41": 2211882496.0,
+            "42": 2211882496.0,
+            "43": 2211882496.0,
+            "44": 2211882496.0,
+            "45": 2211882496.0,
+            "46": 2211882496.0,
+            "47": 2211882496.0,
+            "48": 2211882496.0,
+            "49": 2211882496.0,
+            "50": 2211882496.0,
+            "51": 2211882496.0,
+            "52": 2211882496.0,
+            "53": 2211882496.0,
+            "54": 2211882496.0,
+            "55": 2211882496.0,
+            "56": 2211882496.0,
+            "57": 2211882496.0,
+            "58": 2211882496.0,
+            "59": 2211882496.0,
+            "60": 2211882496.0,
+            "61": 2211882496.0,
+            "62": 2211882496.0,
+            "63": 2211882496.0,
+            "64": 2211882496.0,
+            "65": 2211882496.0,
+            "66": 2211882496.0,
+            "67": 2211882496.0,
+            "68": 2211882496.0,
+            "69": 2211882496.0,
+            "70": 2211882496.0,
+            "71": 2211882496.0,
+            "72": 2211882496.0,
+            "73": 2211882496.0,
+            "74": 2211882496.0,
+            "75": 2211882496.0,
+            "76": 2211882496.0,
+            "77": 2211882496.0,
+            "78": 2211882496.0,
+            "79": 2211882496.0,
+            "80": 2211882496.0,
+            "81": 2211882496.0,
+            "82": 2211882496.0,
+            "83": 2211882496.0,
+            "84": 2211882496.0,
+            "85": 2211882496.0,
+            "86": 2211882496.0,
+            "87": 2211882496.0,
+            "88": 2211882496.0,
+            "89": 2211882496.0,
+            "90": 2211882496.0,
+            "91": 2211882496.0,
+            "92": 2211882496.0,
+            "93": 2211882496.0,
+            "94": 2211882496.0,
+            "95": 2211882496.0,
+            "96": 2211882496.0,
+            "97": 2211882496.0,
+            "98": 2211882496.0,
+            "99": 2211882496.0,
+            "100": 2211882496.0
         }
     },
     "iteration-time": {
-        "start_step": 2,
+        "start_step": 1,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "2": 5.72562,
-            "3": 0.12302,
-            "4": 0.12421,
-            "5": 0.1218,
-            "6": 0.12123,
-            "7": 0.12147,
-            "8": 0.12092,
-            "9": 0.12258,
-            "10": 0.12294,
-            "11": 0.12181,
-            "12": 0.1218,
-            "13": 0.12055,
-            "14": 0.12088,
-            "15": 0.12079,
-            "16": 0.12174,
-            "17": 0.12112,
-            "18": 0.12066,
-            "19": 0.12115,
-            "20": 0.12151,
-            "21": 0.12106,
-            "22": 0.12054,
-            "23": 0.12192,
-            "24": 0.1245,
-            "25": 0.12154,
-            "26": 0.12156,
-            "27": 0.12672,
-            "28": 0.12192,
-            "29": 0.12134,
-            "30": 0.12205,
-            "31": 0.1235,
-            "32": 0.12222,
-            "33": 0.12142,
-            "34": 0.12365,
-            "35": 0.1205,
-            "36": 0.1213,
-            "37": 0.12315,
-            "38": 0.1225,
-            "39": 0.12386,
-            "40": 0.12056,
-            "41": 0.1215,
-            "42": 0.1209,
-            "43": 0.12054,
-            "44": 0.1213,
-            "45": 0.12264,
-            "46": 0.12291,
-            "47": 0.1217,
-            "48": 0.1222,
-            "49": 0.12142,
-            "50": 0.1219,
-            "51": 0.26227,
-            "52": 0.15233,
-            "53": 0.12815,
-            "54": 0.12062,
-            "55": 0.1213,
-            "56": 0.12077,
-            "57": 0.12272,
-            "58": 0.12076,
-            "59": 0.12239,
-            "60": 0.12145,
-            "61": 0.11983,
-            "62": 0.11999,
-            "63": 0.12021,
-            "64": 0.11885,
-            "65": 0.11967,
-            "66": 0.11923,
-            "67": 0.12028,
-            "68": 0.12012,
-            "69": 0.11894,
-            "70": 0.12065,
-            "71": 0.12021,
-            "72": 0.11949,
-            "73": 0.11998,
-            "74": 0.11954,
-            "75": 0.12081,
-            "76": 0.11959,
-            "77": 0.12003,
-            "78": 0.11974,
-            "79": 0.11965,
-            "80": 0.11945,
-            "81": 0.12092,
-            "82": 0.12172,
-            "83": 0.12076,
-            "84": 0.12018,
-            "85": 0.12039,
-            "86": 0.12127,
-            "87": 0.1214,
-            "88": 0.12047,
-            "89": 0.11985,
-            "90": 0.12319,
-            "91": 0.12061,
-            "92": 0.12047,
-            "93": 0.12169,
-            "94": 0.12001,
-            "95": 0.12027,
-            "96": 0.11824,
-            "97": 0.11873,
-            "98": 0.11964,
-            "99": 0.11983,
-            "100": 0.11983
+            "1": "nan",
+            "2": 6.24618,
+            "3": 3.151,
+            "4": 2.31037,
+            "5": 1.79881,
+            "6": 2.29447,
+            "7": 2.74417,
+            "8": 2.52002,
+            "9": 2.61034,
+            "10": 1.40703,
+            "11": 2.23217,
+            "12": 2.20007,
+            "13": 2.3719,
+            "14": 2.61168,
+            "15": 1.90981,
+            "16": 2.30443,
+            "17": 1.99544,
+            "18": 2.44703,
+            "19": 1.98571,
+            "20": 1.84536,
+            "21": 2.1074,
+            "22": 2.00219,
+            "23": 1.74302,
+            "24": 2.1646,
+            "25": 1.44549,
+            "26": 2.40416,
+            "27": 1.43645,
+            "28": 2.20222,
+            "29": 1.74314,
+            "30": 1.21234,
+            "31": 1.35078,
+            "32": 2.70861,
+            "33": 2.42266,
+            "34": 1.51672,
+            "35": 2.53748,
+            "36": 1.83503,
+            "37": 1.92569,
+            "38": 2.06617,
+            "39": 2.0465,
+            "40": 1.92352,
+            "41": 2.77767,
+            "42": 2.04428,
+            "43": 1.42723,
+            "44": 1.91761,
+            "45": 2.17055,
+            "46": 1.75077,
+            "47": 1.81998,
+            "48": 1.81357,
+            "49": 1.69197,
+            "50": 1.51838,
+            "51": 0.15192,
+            "52": 2.02186,
+            "53": 2.16204,
+            "54": 1.92557,
+            "55": 2.22308,
+            "56": 1.99594,
+            "57": 1.83874,
+            "58": 1.46614,
+            "59": 1.97892,
+            "60": 1.16187,
+            "61": 2.02262,
+            "62": 1.95485,
+            "63": 0.9961,
+            "64": 1.88464,
+            "65": 1.22757,
+            "66": 2.15242,
+            "67": 1.35726,
+            "68": 1.77912,
+            "69": 1.79285,
+            "70": 2.02694,
+            "71": 2.01251,
+            "72": 1.02892,
+            "73": 1.68641,
+            "74": 0.85011,
+            "75": 1.43952,
+            "76": 1.72294,
+            "77": 2.05619,
+            "78": 2.05968,
+            "79": 2.46187,
+            "80": 2.11551,
+            "81": 1.19437,
+            "82": 1.85063,
+            "83": 1.26521,
+            "84": 1.70074,
+            "85": 1.40838,
+            "86": 1.32604,
+            "87": 2.8051,
+            "88": 1.67246,
+            "89": 1.42332,
+            "90": 1.95826,
+            "91": 1.48537,
+            "92": 2.47226,
+            "93": 1.88724,
+            "94": 1.7977,
+            "95": 1.22446,
+            "96": 1.50973,
+            "97": 1.25559,
+            "98": 1.24681,
+            "99": 1.80845,
+            "100": 1.85168
         }
     }
 }
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_h100.json
index 8bcf15522c7..18ef0f6881f 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_a2a_ovlp_8experts_etp1_ep4/golden_values_dev_dgx_h100.json
@@ -4,56 +4,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 10.90791,
-            "2": 10.90713,
-            "3": 10.91656,
-            "4": 10.90856,
-            "5": 10.91486,
-            "6": 10.8955,
-            "7": 10.90682,
-            "8": 10.90938,
-            "9": 10.90906,
-            "10": 10.91044,
-            "11": 10.90161,
-            "12": 10.9002,
-            "13": 10.88758,
-            "14": 10.88178,
-            "15": 10.87374,
-            "16": 10.85236,
-            "17": 10.85613,
-            "18": 10.84761,
-            "19": 10.85533,
-            "20": 10.77576,
-            "21": 10.76185,
-            "22": 10.75979,
-            "23": 10.74357,
-            "24": 10.7085,
-            "25": 10.70954,
-            "26": 10.69236,
-            "27": 10.66744,
-            "28": 10.60538,
-            "29": 10.57197,
-            "30": 10.541,
-            "31": 10.53121,
-            "32": 10.51574,
-            "33": 10.48005,
-            "34": 10.44884,
-            "35": 10.44593,
-            "36": 10.41946,
-            "37": 10.39946,
-            "38": 10.40313,
-            "39": 10.36823,
-            "40": 10.35085,
-            "41": 10.3285,
-            "42": 10.30918,
-            "43": 10.29587,
-            "44": 10.26912,
-            "45": 10.28118,
-            "46": 10.2386,
-            "47": 10.23166,
-            "48": 10.18908,
-            "49": 10.19252,
-            "50": 10.1881
+            "1": 10.91,
+            "2": 10.90764,
+            "3": 10.91566,
+            "4": 10.91381,
+            "5": 10.91535,
+            "6": 10.90118,
+            "7": 10.90956,
+            "8": 10.90801,
+            "9": 10.90729,
+            "10": 10.90956,
+            "11": 10.89984,
+            "12": 10.90167,
+            "13": 10.88926,
+            "14": 10.88509,
+            "15": 10.87107,
+            "16": 10.85095,
+            "17": 10.86067,
+            "18": 10.84812,
+            "19": 10.85627,
+            "20": 10.77262,
+            "21": 10.76683,
+            "22": 10.75357,
+            "23": 10.74507,
+            "24": 10.70772,
+            "25": 10.70928,
+            "26": 10.69118,
+            "27": 10.67155,
+            "28": 10.60299,
+            "29": 10.57531,
+            "30": 10.53828,
+            "31": 10.5313,
+            "32": 10.51519,
+            "33": 10.48188,
+            "34": 10.45201,
+            "35": 10.44937,
+            "36": 10.42228,
+            "37": 10.39953,
+            "38": 10.40635,
+            "39": 10.36886,
+            "40": 10.35046,
+            "41": 10.33025,
+            "42": 10.31136,
+            "43": 10.29898,
+            "44": 10.27125,
+            "45": 10.28474,
+            "46": 10.24045,
+            "47": 10.23523,
+            "48": 10.19278,
+            "49": 10.19476,
+            "50": 10.19184
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 32481.0,
-            "2": 32205.0,
-            "3": 31782.0,
-            "4": 32082.0,
-            "5": 31672.0,
-            "6": 30901.0,
-            "7": 32296.0,
-            "8": 30851.0,
-            "9": 32347.0,
-            "10": 32485.0,
-            "11": 31812.0,
-            "12": 31039.0,
-            "13": 32298.0,
-            "14": 32795.0,
-            "15": 31592.0,
-            "16": 30976.0,
-            "17": 32064.0,
-            "18": 32220.0,
-            "19": 32630.0,
-            "20": 32593.0,
-            "21": 31945.0,
-            "22": 32124.0,
-            "23": 32052.0,
-            "24": 33318.0,
-            "25": 31411.0,
-            "26": 32629.0,
-            "27": 32773.0,
-            "28": 32484.0,
-            "29": 32771.0,
-            "30": 32994.0,
-            "31": 34132.0,
-            "32": 34806.0,
-            "33": 33924.0,
-            "34": 34187.0,
-            "35": 35432.0,
-            "36": 35117.0,
-            "37": 35331.0,
-            "38": 35038.0,
-            "39": 36823.0,
-            "40": 38166.0,
-            "41": 36109.0,
-            "42": 35997.0,
-            "43": 37458.0,
-            "44": 37012.0,
-            "45": 40824.0,
-            "46": 38797.0,
-            "47": 38754.0,
-            "48": 39579.0,
-            "49": 43844.0,
-            "50": 39384.0
+            "1": 16628.0,
+            "2": 16624.0,
+            "3": 16082.0,
+            "4": 16379.0,
+            "5": 16275.0,
+            "6": 15872.0,
+            "7": 16861.0,
+            "8": 16026.0,
+            "9": 16555.0,
+            "10": 16899.0,
+            "11": 16553.0,
+            "12": 15990.0,
+            "13": 16930.0,
+            "14": 16753.0,
+            "15": 16015.0,
+            "16": 16170.0,
+            "17": 16381.0,
+            "18": 16670.0,
+            "19": 16342.0,
+            "20": 16460.0,
+            "21": 16426.0,
+            "22": 16453.0,
+            "23": 16537.0,
+            "24": 16866.0,
+            "25": 15501.0,
+            "26": 16689.0,
+            "27": 16669.0,
+            "28": 16329.0,
+            "29": 16299.0,
+            "30": 16485.0,
+            "31": 17244.0,
+            "32": 17360.0,
+            "33": 17071.0,
+            "34": 16996.0,
+            "35": 17753.0,
+            "36": 17364.0,
+            "37": 17823.0,
+            "38": 17542.0,
+            "39": 18334.0,
+            "40": 19250.0,
+            "41": 18258.0,
+            "42": 17843.0,
+            "43": 18802.0,
+            "44": 18562.0,
+            "45": 20762.0,
+            "46": 19921.0,
+            "47": 19428.0,
+            "48": 20192.0,
+            "49": 22112.0,
+            "50": 20108.0
         }
     },
     "mem-allocated-bytes": {
@@ -118,56 +118,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 1564053504.0,
-            "2": 1563238912.0,
-            "3": 1562403328.0,
-            "4": 1562403328.0,
-            "5": 1562403328.0,
-            "6": 1562403328.0,
-            "7": 1562403328.0,
-            "8": 1562513920.0,
-            "9": 1563245056.0,
-            "10": 1562403328.0,
-            "11": 1562484224.0,
-            "12": 1562403328.0,
-            "13": 1562433024.0,
-            "14": 1562863104.0,
-            "15": 1563126272.0,
-            "16": 1563629056.0,
-            "17": 1563629056.0,
-            "18": 1564133888.0,
-            "19": 1562616320.0,
-            "20": 1562403328.0,
-            "21": 1562403328.0,
-            "22": 1562895872.0,
-            "23": 1563697664.0,
-            "24": 1562403328.0,
-            "25": 1562640896.0,
-            "26": 1562403328.0,
-            "27": 1562403328.0,
-            "28": 1562403328.0,
-            "29": 1562691072.0,
-            "30": 1563299328.0,
-            "31": 1562403328.0,
-            "32": 1562403328.0,
-            "33": 1562915328.0,
-            "34": 1563009536.0,
-            "35": 1562771968.0,
-            "36": 1562403328.0,
-            "37": 1563407872.0,
-            "38": 1563094528.0,
-            "39": 1562403328.0,
-            "40": 1562403328.0,
-            "41": 1563132416.0,
-            "42": 1562603008.0,
-            "43": 1562563072.0,
-            "44": 1562528256.0,
-            "45": 1562403328.0,
-            "46": 1562589696.0,
-            "47": 1562403328.0,
-            "48": 1562632704.0,
-            "49": 1562403328.0,
-            "50": 1563123200.0
+            "1": 1561356288.0,
+            "2": 1560986112.0,
+            "3": 1560986112.0,
+            "4": 1560986112.0,
+            "5": 1561978368.0,
+            "6": 1560986112.0,
+            "7": 1561500160.0,
+            "8": 1560986112.0,
+            "9": 1561209344.0,
+            "10": 1561555456.0,
+            "11": 1561686528.0,
+            "12": 1561893376.0,
+            "13": 1562735104.0,
+            "14": 1560986112.0,
+            "15": 1560986112.0,
+            "16": 1562570240.0,
+            "17": 1562019328.0,
+            "18": 1562019328.0,
+            "19": 1560986112.0,
+            "20": 1561616896.0,
+            "21": 1560986112.0,
+            "22": 1560986112.0,
+            "23": 1560986112.0,
+            "24": 1561467392.0,
+            "25": 1560986112.0,
+            "26": 1560986112.0,
+            "27": 1560986112.0,
+            "28": 1560986112.0,
+            "29": 1560986112.0,
+            "30": 1560986112.0,
+            "31": 1560986112.0,
+            "32": 1561809408.0,
+            "33": 1561820672.0,
+            "34": 1562026496.0,
+            "35": 1561317888.0,
+            "36": 1561453056.0,
+            "37": 1561453056.0,
+            "38": 1561215488.0,
+            "39": 1561379328.0,
+            "40": 1561657856.0,
+            "41": 1560986112.0,
+            "42": 1561094656.0,
+            "43": 1560986112.0,
+            "44": 1560986112.0,
+            "45": 1561507328.0,
+            "46": 1561619968.0,
+            "47": 1561155072.0,
+            "48": 1561155072.0,
+            "49": 1561050624.0,
+            "50": 1560986112.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -175,112 +175,113 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 3481845248.0,
-            "2": 4045780992.0,
-            "3": 4052392960.0,
-            "4": 4057906176.0,
-            "5": 4057906176.0,
-            "6": 4057906176.0,
-            "7": 4057906176.0,
-            "8": 4057906176.0,
-            "9": 4060597248.0,
-            "10": 4065854464.0,
-            "11": 4065854464.0,
-            "12": 4065854464.0,
-            "13": 4065854464.0,
-            "14": 4065854464.0,
-            "15": 4065854464.0,
-            "16": 4065854464.0,
-            "17": 4065854464.0,
-            "18": 4065854464.0,
-            "19": 4065854464.0,
-            "20": 4065854464.0,
-            "21": 4065854464.0,
-            "22": 4065854464.0,
-            "23": 4065854464.0,
-            "24": 4065854464.0,
-            "25": 4065854464.0,
-            "26": 4065854464.0,
-            "27": 4065854464.0,
-            "28": 4065854464.0,
-            "29": 4065854464.0,
-            "30": 4065854464.0,
-            "31": 4065854464.0,
-            "32": 4065854464.0,
-            "33": 4065854464.0,
-            "34": 4065854464.0,
-            "35": 4065854464.0,
-            "36": 4065854464.0,
-            "37": 4065854464.0,
-            "38": 4065854464.0,
-            "39": 4065854464.0,
-            "40": 4065854464.0,
-            "41": 4065854464.0,
-            "42": 4065854464.0,
-            "43": 4065854464.0,
-            "44": 4065854464.0,
-            "45": 4065854464.0,
-            "46": 4065854464.0,
-            "47": 4065854464.0,
-            "48": 4065854464.0,
-            "49": 4065854464.0,
-            "50": 4065854464.0
+            "1": 3480969728.0,
+            "2": 4040542208.0,
+            "3": 4040542208.0,
+            "4": 4055748608.0,
+            "5": 4055748608.0,
+            "6": 4055748608.0,
+            "7": 4059894784.0,
+            "8": 4063235072.0,
+            "9": 4063235072.0,
+            "10": 4063235072.0,
+            "11": 4063235072.0,
+            "12": 4063235072.0,
+            "13": 4063235072.0,
+            "14": 4063235072.0,
+            "15": 4063235072.0,
+            "16": 4063235072.0,
+            "17": 4063235072.0,
+            "18": 4063235072.0,
+            "19": 4063235072.0,
+            "20": 4063235072.0,
+            "21": 4063235072.0,
+            "22": 4063235072.0,
+            "23": 4063235072.0,
+            "24": 4063235072.0,
+            "25": 4063235072.0,
+            "26": 4063235072.0,
+            "27": 4063235072.0,
+            "28": 4063235072.0,
+            "29": 4063235072.0,
+            "30": 4063235072.0,
+            "31": 4063235072.0,
+            "32": 4063235072.0,
+            "33": 4063235072.0,
+            "34": 4063235072.0,
+            "35": 4063235072.0,
+            "36": 4063235072.0,
+            "37": 4063235072.0,
+            "38": 4063235072.0,
+            "39": 4063235072.0,
+            "40": 4063235072.0,
+            "41": 4063235072.0,
+            "42": 4063235072.0,
+            "43": 4063235072.0,
+            "44": 4063235072.0,
+            "45": 4063235072.0,
+            "46": 4063235072.0,
+            "47": 4063235072.0,
+            "48": 4063235072.0,
+            "49": 4063235072.0,
+            "50": 4063235072.0
         }
     },
     "iteration-time": {
-        "start_step": 2,
+        "start_step": 1,
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "2": 7.65129,
-            "3": 0.49113,
-            "4": 0.46732,
-            "5": 0.45015,
-            "6": 0.44098,
-            "7": 0.43279,
-            "8": 0.43532,
-            "9": 0.41644,
-            "10": 0.41447,
-            "11": 0.417,
-            "12": 1.27039,
-            "13": 0.42857,
-            "14": 0.42043,
-            "15": 0.43429,
-            "16": 0.42646,
-            "17": 0.41839,
-            "18": 0.41875,
-            "19": 0.42078,
-            "20": 0.4152,
-            "21": 0.41839,
-            "22": 0.42007,
-            "23": 0.40978,
-            "24": 0.4028,
-            "25": 0.40842,
-            "26": 0.41505,
-            "27": 1.15357,
-            "28": 0.43831,
-            "29": 1.14215,
-            "30": 0.42585,
-            "31": 0.42393,
-            "32": 0.42386,
-            "33": 0.41322,
-            "34": 0.42071,
-            "35": 0.4168,
-            "36": 1.17628,
-            "37": 0.42372,
-            "38": 0.42557,
-            "39": 1.14457,
-            "40": 0.4147,
-            "41": 0.41313,
-            "42": 0.41232,
-            "43": 0.41219,
-            "44": 0.41084,
-            "45": 0.40381,
-            "46": 0.40844,
-            "47": 0.40717,
-            "48": 1.20976,
-            "49": 0.42107,
-            "50": 0.42035
+            "1": "nan",
+            "2": 8.15585,
+            "3": 0.46238,
+            "4": 0.94242,
+            "5": 0.44862,
+            "6": 1.05421,
+            "7": 0.86359,
+            "8": 1.54702,
+            "9": 1.00781,
+            "10": 0.428,
+            "11": 1.28207,
+            "12": 0.49529,
+            "13": 0.7824,
+            "14": 1.00029,
+            "15": 0.47713,
+            "16": 1.35821,
+            "17": 0.75064,
+            "18": 1.06931,
+            "19": 0.40765,
+            "20": 0.75622,
+            "21": 1.276,
+            "22": 0.60548,
+            "23": 0.67732,
+            "24": 0.64344,
+            "25": 0.47881,
+            "26": 1.32412,
+            "27": 0.46865,
+            "28": 0.89189,
+            "29": 0.52184,
+            "30": 0.42359,
+            "31": 0.87728,
+            "32": 1.26034,
+            "33": 0.60968,
+            "34": 1.46667,
+            "35": 0.4158,
+            "36": 0.88038,
+            "37": 0.91173,
+            "38": 0.42661,
+            "39": 0.77983,
+            "40": 1.25508,
+            "41": 0.77841,
+            "42": 0.91989,
+            "43": 0.59046,
+            "44": 0.75354,
+            "45": 0.60038,
+            "46": 0.97126,
+            "47": 0.54442,
+            "48": 0.91883,
+            "49": 0.52464,
+            "50": 0.40366
         }
     }
-}
+}
\ 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.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 c4aaa6c46c1..a307d679d46 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
@@ -4,106 +4,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 10.90111,
-            "2": 10.8957,
-            "3": 10.90636,
-            "4": 10.90606,
-            "5": 10.91349,
-            "6": 10.89747,
-            "7": 10.89834,
-            "8": 10.91366,
-            "9": 10.89559,
-            "10": 10.90375,
-            "11": 10.90014,
-            "12": 10.91887,
-            "13": 10.88393,
-            "14": 10.87584,
-            "15": 10.91035,
-            "16": 10.90325,
-            "17": 10.88717,
-            "18": 10.89416,
-            "19": 10.89487,
-            "20": 10.89092,
-            "21": 10.88636,
-            "22": 10.90341,
-            "23": 10.91256,
-            "24": 10.88547,
-            "25": 10.90164,
-            "26": 10.88848,
-            "27": 10.89705,
-            "28": 10.88168,
-            "29": 10.88997,
-            "30": 10.91097,
-            "31": 10.89128,
-            "32": 10.88933,
-            "33": 10.89772,
-            "34": 10.87339,
-            "35": 10.89595,
-            "36": 10.90751,
-            "37": 10.87112,
-            "38": 10.87839,
-            "39": 10.88579,
-            "40": 10.89247,
-            "41": 10.88364,
-            "42": 10.89496,
-            "43": 10.88058,
-            "44": 10.88492,
-            "45": 10.88361,
-            "46": 10.88582,
-            "47": 10.88565,
-            "48": 10.86424,
-            "49": 10.8759,
-            "50": 10.88259,
-            "51": 10.89305,
-            "52": 10.87123,
-            "53": 10.85905,
-            "54": 10.8702,
-            "55": 10.86807,
-            "56": 10.87324,
-            "57": 10.84465,
-            "58": 10.8563,
-            "59": 10.84628,
-            "60": 10.84343,
-            "61": 10.8604,
-            "62": 10.85675,
-            "63": 10.86117,
-            "64": 10.83852,
-            "65": 10.82804,
-            "66": 10.84584,
-            "67": 10.82846,
-            "68": 10.83007,
-            "69": 10.81894,
-            "70": 10.82797,
-            "71": 10.82677,
-            "72": 10.8088,
-            "73": 10.80753,
-            "74": 10.80659,
-            "75": 10.81472,
-            "76": 10.81086,
-            "77": 10.80817,
-            "78": 10.79255,
-            "79": 10.80218,
-            "80": 10.7879,
-            "81": 10.79487,
-            "82": 10.79556,
-            "83": 10.78629,
-            "84": 10.75658,
-            "85": 10.76241,
-            "86": 10.77699,
-            "87": 10.79624,
-            "88": 10.77534,
-            "89": 10.77559,
-            "90": 10.76414,
-            "91": 10.74033,
-            "92": 10.76012,
-            "93": 10.74683,
-            "94": 10.73435,
-            "95": 10.75231,
-            "96": 10.72344,
-            "97": 10.71428,
-            "98": 10.72631,
-            "99": 10.74637,
-            "100": 10.69592
+            "1": 10.89993,
+            "2": 10.89521,
+            "3": 10.91766,
+            "4": 10.90803,
+            "5": 10.90842,
+            "6": 10.89655,
+            "7": 10.90177,
+            "8": 10.90892,
+            "9": 10.89906,
+            "10": 10.90528,
+            "11": 10.90163,
+            "12": 10.92056,
+            "13": 10.88881,
+            "14": 10.87794,
+            "15": 10.90166,
+            "16": 10.90092,
+            "17": 10.88659,
+            "18": 10.89124,
+            "19": 10.89595,
+            "20": 10.89006,
+            "21": 10.88404,
+            "22": 10.89665,
+            "23": 10.91331,
+            "24": 10.89125,
+            "25": 10.90628,
+            "26": 10.89515,
+            "27": 10.89244,
+            "28": 10.88088,
+            "29": 10.89309,
+            "30": 10.90833,
+            "31": 10.8921,
+            "32": 10.88871,
+            "33": 10.89608,
+            "34": 10.87464,
+            "35": 10.89097,
+            "36": 10.90886,
+            "37": 10.88013,
+            "38": 10.87495,
+            "39": 10.87764,
+            "40": 10.8967,
+            "41": 10.88274,
+            "42": 10.89148,
+            "43": 10.87904,
+            "44": 10.8877,
+            "45": 10.89149,
+            "46": 10.88009,
+            "47": 10.88249,
+            "48": 10.86425,
+            "49": 10.87229,
+            "50": 10.88948,
+            "51": 10.88676,
+            "52": 10.87193,
+            "53": 10.86027,
+            "54": 10.86311,
+            "55": 10.86882,
+            "56": 10.87537,
+            "57": 10.84319,
+            "58": 10.85656,
+            "59": 10.84599,
+            "60": 10.84173,
+            "61": 10.86236,
+            "62": 10.85456,
+            "63": 10.86271,
+            "64": 10.83785,
+            "65": 10.82429,
+            "66": 10.84478,
+            "67": 10.82847,
+            "68": 10.82859,
+            "69": 10.81623,
+            "70": 10.82685,
+            "71": 10.82907,
+            "72": 10.81157,
+            "73": 10.8073,
+            "74": 10.80688,
+            "75": 10.8117,
+            "76": 10.80937,
+            "77": 10.80521,
+            "78": 10.7976,
+            "79": 10.80533,
+            "80": 10.78589,
+            "81": 10.79293,
+            "82": 10.79276,
+            "83": 10.78493,
+            "84": 10.75347,
+            "85": 10.76248,
+            "86": 10.7731,
+            "87": 10.79677,
+            "88": 10.77462,
+            "89": 10.76937,
+            "90": 10.76432,
+            "91": 10.74231,
+            "92": 10.75842,
+            "93": 10.74898,
+            "94": 10.73439,
+            "95": 10.74831,
+            "96": 10.71976,
+            "97": 10.71354,
+            "98": 10.72427,
+            "99": 10.74505,
+            "100": 10.69992
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1133.0,
-            "2": 1120.0,
-            "3": 1371.0,
-            "4": 1297.0,
-            "5": 1158.0,
-            "6": 1283.0,
-            "7": 1334.0,
-            "8": 1122.0,
-            "9": 1269.0,
-            "10": 1134.0,
-            "11": 1184.0,
-            "12": 1223.0,
-            "13": 1340.0,
-            "14": 1232.0,
-            "15": 1234.0,
-            "16": 1179.0,
-            "17": 1212.0,
-            "18": 1083.0,
-            "19": 1187.0,
-            "20": 1174.0,
-            "21": 1196.0,
-            "22": 1227.0,
-            "23": 1318.0,
-            "24": 1113.0,
-            "25": 1191.0,
-            "26": 1243.0,
-            "27": 1343.0,
-            "28": 1251.0,
-            "29": 1186.0,
-            "30": 1221.0,
-            "31": 1205.0,
-            "32": 1269.0,
-            "33": 1254.0,
-            "34": 1196.0,
-            "35": 1308.0,
-            "36": 1277.0,
-            "37": 1196.0,
-            "38": 1321.0,
-            "39": 1407.0,
-            "40": 1299.0,
-            "41": 1284.0,
-            "42": 1170.0,
-            "43": 1170.0,
-            "44": 1252.0,
-            "45": 1154.0,
-            "46": 1403.0,
-            "47": 1148.0,
-            "48": 1176.0,
-            "49": 1389.0,
-            "50": 1329.0,
-            "51": 1209.0,
-            "52": 1224.0,
-            "53": 1276.0,
-            "54": 1198.0,
-            "55": 1179.0,
-            "56": 1166.0,
-            "57": 1168.0,
-            "58": 1318.0,
-            "59": 1047.0,
-            "60": 1118.0,
-            "61": 1192.0,
-            "62": 1108.0,
-            "63": 1225.0,
-            "64": 1256.0,
-            "65": 1221.0,
-            "66": 1228.0,
-            "67": 1193.0,
-            "68": 1279.0,
-            "69": 1141.0,
-            "70": 1248.0,
-            "71": 1165.0,
-            "72": 1230.0,
-            "73": 1236.0,
-            "74": 1336.0,
-            "75": 1346.0,
-            "76": 1259.0,
-            "77": 1205.0,
-            "78": 1217.0,
-            "79": 1298.0,
-            "80": 1228.0,
-            "81": 1335.0,
-            "82": 1198.0,
-            "83": 1246.0,
-            "84": 1152.0,
-            "85": 1250.0,
-            "86": 1287.0,
-            "87": 1111.0,
-            "88": 1169.0,
-            "89": 1119.0,
-            "90": 1402.0,
-            "91": 1297.0,
-            "92": 1194.0,
-            "93": 1271.0,
-            "94": 1409.0,
-            "95": 980.0,
-            "96": 1170.0,
-            "97": 1212.0,
-            "98": 1325.0,
-            "99": 1125.0,
-            "100": 1206.0
+            "1": 1240.0,
+            "2": 1269.0,
+            "3": 1319.0,
+            "4": 1193.0,
+            "5": 1257.0,
+            "6": 1294.0,
+            "7": 1261.0,
+            "8": 1300.0,
+            "9": 1277.0,
+            "10": 1160.0,
+            "11": 1271.0,
+            "12": 1299.0,
+            "13": 1215.0,
+            "14": 1219.0,
+            "15": 1077.0,
+            "16": 1178.0,
+            "17": 1204.0,
+            "18": 1160.0,
+            "19": 1250.0,
+            "20": 1139.0,
+            "21": 1234.0,
+            "22": 1308.0,
+            "23": 1192.0,
+            "24": 1267.0,
+            "25": 1241.0,
+            "26": 1304.0,
+            "27": 1319.0,
+            "28": 1255.0,
+            "29": 1195.0,
+            "30": 1198.0,
+            "31": 1237.0,
+            "32": 1204.0,
+            "33": 1214.0,
+            "34": 1207.0,
+            "35": 1127.0,
+            "36": 1217.0,
+            "37": 1349.0,
+            "38": 1344.0,
+            "39": 1351.0,
+            "40": 1294.0,
+            "41": 1117.0,
+            "42": 1210.0,
+            "43": 1266.0,
+            "44": 1183.0,
+            "45": 1270.0,
+            "46": 1324.0,
+            "47": 1078.0,
+            "48": 1238.0,
+            "49": 1256.0,
+            "50": 1088.0,
+            "51": 1247.0,
+            "52": 1085.0,
+            "53": 1203.0,
+            "54": 1255.0,
+            "55": 1237.0,
+            "56": 1101.0,
+            "57": 1218.0,
+            "58": 1319.0,
+            "59": 1211.0,
+            "60": 1165.0,
+            "61": 1128.0,
+            "62": 1261.0,
+            "63": 1164.0,
+            "64": 1358.0,
+            "65": 1176.0,
+            "66": 1116.0,
+            "67": 1184.0,
+            "68": 1374.0,
+            "69": 1347.0,
+            "70": 1218.0,
+            "71": 1251.0,
+            "72": 1137.0,
+            "73": 1259.0,
+            "74": 1290.0,
+            "75": 1431.0,
+            "76": 1308.0,
+            "77": 1342.0,
+            "78": 1220.0,
+            "79": 1212.0,
+            "80": 1393.0,
+            "81": 1219.0,
+            "82": 1225.0,
+            "83": 1248.0,
+            "84": 1231.0,
+            "85": 1238.0,
+            "86": 1238.0,
+            "87": 1046.0,
+            "88": 1284.0,
+            "89": 1353.0,
+            "90": 1369.0,
+            "91": 1202.0,
+            "92": 1349.0,
+            "93": 1137.0,
+            "94": 1288.0,
+            "95": 1227.0,
+            "96": 1137.0,
+            "97": 1172.0,
+            "98": 1257.0,
+            "99": 1301.0,
+            "100": 1193.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 994063360.0,
-            "2": 994054144.0,
-            "3": 994029056.0,
-            "4": 994010624.0,
-            "5": 994059776.0,
-            "6": 994042368.0,
-            "7": 993970176.0,
-            "8": 994006528.0,
-            "9": 994008064.0,
-            "10": 994039808.0,
-            "11": 994015232.0,
-            "12": 994024960.0,
-            "13": 994029056.0,
-            "14": 994030592.0,
-            "15": 994007040.0,
-            "16": 993992704.0,
-            "17": 994001408.0,
-            "18": 994003968.0,
-            "19": 994032640.0,
-            "20": 993997312.0,
-            "21": 994052608.0,
-            "22": 994075648.0,
-            "23": 994049024.0,
-            "24": 993964032.0,
-            "25": 994045440.0,
-            "26": 994076672.0,
-            "27": 994028544.0,
-            "28": 993984512.0,
-            "29": 994056704.0,
-            "30": 993990656.0,
-            "31": 993996800.0,
-            "32": 994003456.0,
-            "33": 994011136.0,
-            "34": 993989120.0,
-            "35": 993974784.0,
-            "36": 994040832.0,
-            "37": 993988608.0,
-            "38": 994046976.0,
-            "39": 994054144.0,
-            "40": 994001408.0,
-            "41": 994021888.0,
-            "42": 993964544.0,
-            "43": 994032640.0,
-            "44": 994014208.0,
-            "45": 993989120.0,
-            "46": 994052096.0,
-            "47": 994011136.0,
-            "48": 994040320.0,
-            "49": 993998848.0,
-            "50": 994044928.0,
-            "51": 994030592.0,
-            "52": 994003968.0,
-            "53": 993953280.0,
-            "54": 994051072.0,
-            "55": 994039808.0,
-            "56": 994038784.0,
-            "57": 994098688.0,
-            "58": 994035200.0,
-            "59": 994040832.0,
-            "60": 994019328.0,
-            "61": 993988608.0,
-            "62": 993980928.0,
-            "63": 993987584.0,
-            "64": 994015744.0,
-            "65": 994023424.0,
-            "66": 994037248.0,
-            "67": 994018816.0,
-            "68": 994020352.0,
-            "69": 993985024.0,
-            "70": 994067456.0,
-            "71": 994047488.0,
-            "72": 994044416.0,
-            "73": 993973760.0,
-            "74": 994015744.0,
-            "75": 994060800.0,
-            "76": 994017792.0,
-            "77": 994060288.0,
-            "78": 994006016.0,
-            "79": 994072064.0,
-            "80": 993993728.0,
-            "81": 994046464.0,
-            "82": 993994752.0,
-            "83": 994054656.0,
-            "84": 994024448.0,
-            "85": 994032128.0,
-            "86": 994016256.0,
-            "87": 994030080.0,
-            "88": 994010112.0,
-            "89": 994001408.0,
-            "90": 994024960.0,
-            "91": 994011136.0,
-            "92": 994009600.0,
-            "93": 994040320.0,
-            "94": 994015232.0,
-            "95": 994036736.0,
-            "96": 994028544.0,
-            "97": 993997312.0,
-            "98": 994035712.0,
-            "99": 994051584.0,
-            "100": 994083328.0
+            "1": 994539520.0,
+            "2": 994600448.0,
+            "3": 994529280.0,
+            "4": 994469376.0,
+            "5": 994539008.0,
+            "6": 994576896.0,
+            "7": 994478080.0,
+            "8": 994490368.0,
+            "9": 994520576.0,
+            "10": 994538496.0,
+            "11": 994491904.0,
+            "12": 994523648.0,
+            "13": 994546176.0,
+            "14": 994506752.0,
+            "15": 994555392.0,
+            "16": 994532864.0,
+            "17": 994450432.0,
+            "18": 994511360.0,
+            "19": 994550272.0,
+            "20": 994529280.0,
+            "21": 994562560.0,
+            "22": 994536960.0,
+            "23": 994513408.0,
+            "24": 994521600.0,
+            "25": 994524672.0,
+            "26": 994595328.0,
+            "27": 994556928.0,
+            "28": 994487808.0,
+            "29": 994530304.0,
+            "30": 994516992.0,
+            "31": 994503680.0,
+            "32": 994458112.0,
+            "33": 994515968.0,
+            "34": 994542080.0,
+            "35": 994507264.0,
+            "36": 994511360.0,
+            "37": 994478080.0,
+            "38": 994489344.0,
+            "39": 994570752.0,
+            "40": 994491904.0,
+            "41": 994545152.0,
+            "42": 994475520.0,
+            "43": 994553856.0,
+            "44": 994544640.0,
+            "45": 994519040.0,
+            "46": 994569728.0,
+            "47": 994522112.0,
+            "48": 994481664.0,
+            "49": 994493440.0,
+            "50": 994513408.0,
+            "51": 994516992.0,
+            "52": 994521088.0,
+            "53": 994452480.0,
+            "54": 994532352.0,
+            "55": 994541056.0,
+            "56": 994513408.0,
+            "57": 994566656.0,
+            "58": 994548736.0,
+            "59": 994479104.0,
+            "60": 994532864.0,
+            "61": 994526720.0,
+            "62": 994499584.0,
+            "63": 994492928.0,
+            "64": 994521088.0,
+            "65": 994519552.0,
+            "66": 994519040.0,
+            "67": 994522112.0,
+            "68": 994549248.0,
+            "69": 994506752.0,
+            "70": 994511872.0,
+            "71": 994523136.0,
+            "72": 994551808.0,
+            "73": 994503680.0,
+            "74": 994517504.0,
+            "75": 994522624.0,
+            "76": 994521600.0,
+            "77": 994529792.0,
+            "78": 994544640.0,
+            "79": 994601472.0,
+            "80": 994529792.0,
+            "81": 994481152.0,
+            "82": 994483712.0,
+            "83": 994542592.0,
+            "84": 994544128.0,
+            "85": 994533888.0,
+            "86": 994503168.0,
+            "87": 994492416.0,
+            "88": 994517504.0,
+            "89": 994507264.0,
+            "90": 994519040.0,
+            "91": 994515456.0,
+            "92": 994498560.0,
+            "93": 994508288.0,
+            "94": 994493440.0,
+            "95": 994522112.0,
+            "96": 994510848.0,
+            "97": 994505216.0,
+            "98": 994539008.0,
+            "99": 994555904.0,
+            "100": 994522624.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,212 +325,213 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 3205999104.0,
-            "2": 3490994688.0,
-            "3": 3490994688.0,
-            "4": 3490994688.0,
-            "5": 3498078208.0,
-            "6": 3498078208.0,
-            "7": 3498078208.0,
-            "8": 3498078208.0,
-            "9": 3498078208.0,
-            "10": 3498078208.0,
-            "11": 3498078208.0,
-            "12": 3498078208.0,
-            "13": 3498078208.0,
-            "14": 3498078208.0,
-            "15": 3498078208.0,
-            "16": 3498078208.0,
-            "17": 3498078208.0,
-            "18": 3498078208.0,
-            "19": 3498078208.0,
-            "20": 3498078208.0,
-            "21": 3498078208.0,
-            "22": 3507077632.0,
-            "23": 3507077632.0,
-            "24": 3507077632.0,
-            "25": 3507077632.0,
-            "26": 3510859264.0,
-            "27": 3510859264.0,
-            "28": 3510859264.0,
-            "29": 3510859264.0,
-            "30": 3510859264.0,
-            "31": 3510859264.0,
-            "32": 3510859264.0,
-            "33": 3510859264.0,
-            "34": 3510859264.0,
-            "35": 3510859264.0,
-            "36": 3510859264.0,
-            "37": 3510859264.0,
-            "38": 3510859264.0,
-            "39": 3510859264.0,
-            "40": 3510859264.0,
-            "41": 3510859264.0,
-            "42": 3510859264.0,
-            "43": 3510859264.0,
-            "44": 3510859264.0,
-            "45": 3510859264.0,
-            "46": 3510859264.0,
-            "47": 3510859264.0,
-            "48": 3510859264.0,
-            "49": 3510859264.0,
-            "50": 3510859264.0,
-            "51": 3510859264.0,
-            "52": 3510859264.0,
-            "53": 3510859264.0,
-            "54": 3510859264.0,
-            "55": 3510859264.0,
-            "56": 3510859264.0,
-            "57": 3526288384.0,
-            "58": 3526288384.0,
-            "59": 3526288384.0,
-            "60": 3526288384.0,
-            "61": 3526288384.0,
-            "62": 3526288384.0,
-            "63": 3526288384.0,
-            "64": 3526288384.0,
-            "65": 3526288384.0,
-            "66": 3526288384.0,
-            "67": 3526288384.0,
-            "68": 3526288384.0,
-            "69": 3526288384.0,
-            "70": 3526288384.0,
-            "71": 3526288384.0,
-            "72": 3526288384.0,
-            "73": 3526288384.0,
-            "74": 3526288384.0,
-            "75": 3526288384.0,
-            "76": 3526288384.0,
-            "77": 3526288384.0,
-            "78": 3526288384.0,
-            "79": 3526288384.0,
-            "80": 3526288384.0,
-            "81": 3526288384.0,
-            "82": 3526288384.0,
-            "83": 3526288384.0,
-            "84": 3526288384.0,
-            "85": 3526288384.0,
-            "86": 3526288384.0,
-            "87": 3526288384.0,
-            "88": 3526288384.0,
-            "89": 3526288384.0,
-            "90": 3526288384.0,
-            "91": 3526288384.0,
-            "92": 3526288384.0,
-            "93": 3526288384.0,
-            "94": 3526288384.0,
-            "95": 3526288384.0,
-            "96": 3526288384.0,
-            "97": 3526288384.0,
-            "98": 3526288384.0,
-            "99": 3526288384.0,
-            "100": 3526288384.0
+            "1": 3190371840.0,
+            "2": 3531113472.0,
+            "3": 3531113472.0,
+            "4": 3531113472.0,
+            "5": 3531113472.0,
+            "6": 3531113472.0,
+            "7": 3531113472.0,
+            "8": 3531113472.0,
+            "9": 3531113472.0,
+            "10": 3531113472.0,
+            "11": 3531113472.0,
+            "12": 3531113472.0,
+            "13": 3531113472.0,
+            "14": 3531113472.0,
+            "15": 3531113472.0,
+            "16": 3531113472.0,
+            "17": 3531113472.0,
+            "18": 3531113472.0,
+            "19": 3531113472.0,
+            "20": 3531113472.0,
+            "21": 3531113472.0,
+            "22": 3531113472.0,
+            "23": 3531113472.0,
+            "24": 3531113472.0,
+            "25": 3531113472.0,
+            "26": 3531113472.0,
+            "27": 3531113472.0,
+            "28": 3531113472.0,
+            "29": 3531113472.0,
+            "30": 3531113472.0,
+            "31": 3531113472.0,
+            "32": 3531113472.0,
+            "33": 3531113472.0,
+            "34": 3531113472.0,
+            "35": 3531113472.0,
+            "36": 3531113472.0,
+            "37": 3531113472.0,
+            "38": 3531113472.0,
+            "39": 3531113472.0,
+            "40": 3531113472.0,
+            "41": 3531113472.0,
+            "42": 3531113472.0,
+            "43": 3531113472.0,
+            "44": 3531113472.0,
+            "45": 3531113472.0,
+            "46": 3531113472.0,
+            "47": 3531113472.0,
+            "48": 3531113472.0,
+            "49": 3531113472.0,
+            "50": 3531113472.0,
+            "51": 3531113472.0,
+            "52": 3531113472.0,
+            "53": 3531113472.0,
+            "54": 3531113472.0,
+            "55": 3531113472.0,
+            "56": 3531113472.0,
+            "57": 3531113472.0,
+            "58": 3531113472.0,
+            "59": 3531113472.0,
+            "60": 3531113472.0,
+            "61": 3531113472.0,
+            "62": 3531113472.0,
+            "63": 3531113472.0,
+            "64": 3531113472.0,
+            "65": 3531113472.0,
+            "66": 3531113472.0,
+            "67": 3531113472.0,
+            "68": 3531113472.0,
+            "69": 3531113472.0,
+            "70": 3531113472.0,
+            "71": 3531113472.0,
+            "72": 3531113472.0,
+            "73": 3531113472.0,
+            "74": 3531113472.0,
+            "75": 3531113472.0,
+            "76": 3531113472.0,
+            "77": 3531113472.0,
+            "78": 3531113472.0,
+            "79": 3531113472.0,
+            "80": 3531113472.0,
+            "81": 3531113472.0,
+            "82": 3531113472.0,
+            "83": 3531113472.0,
+            "84": 3531113472.0,
+            "85": 3531113472.0,
+            "86": 3531113472.0,
+            "87": 3531113472.0,
+            "88": 3531113472.0,
+            "89": 3531113472.0,
+            "90": 3531113472.0,
+            "91": 3531113472.0,
+            "92": 3531113472.0,
+            "93": 3531113472.0,
+            "94": 3531113472.0,
+            "95": 3531113472.0,
+            "96": 3531113472.0,
+            "97": 3531113472.0,
+            "98": 3531113472.0,
+            "99": 3531113472.0,
+            "100": 3531113472.0
         }
     },
     "iteration-time": {
-        "start_step": 2,
+        "start_step": 1,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "2": 4.80343,
-            "3": 0.24461,
-            "4": 0.21905,
-            "5": 0.23456,
-            "6": 0.22151,
-            "7": 0.22611,
-            "8": 0.2404,
-            "9": 0.21788,
-            "10": 0.22319,
-            "11": 0.22849,
-            "12": 0.21807,
-            "13": 0.2241,
-            "14": 0.21328,
-            "15": 0.21924,
-            "16": 0.19766,
-            "17": 0.2119,
-            "18": 0.20845,
-            "19": 0.2102,
-            "20": 0.21676,
-            "21": 0.2071,
-            "22": 0.21066,
-            "23": 0.20145,
-            "24": 0.20861,
-            "25": 0.22351,
-            "26": 0.20212,
-            "27": 0.19701,
-            "28": 0.19726,
-            "29": 0.21428,
-            "30": 0.20896,
-            "31": 0.20412,
-            "32": 0.21173,
-            "33": 0.20281,
-            "34": 0.20432,
-            "35": 0.20437,
-            "36": 0.20527,
-            "37": 0.20735,
-            "38": 0.19734,
-            "39": 0.19343,
-            "40": 0.20184,
-            "41": 0.19924,
-            "42": 0.19429,
-            "43": 0.20453,
-            "44": 0.19924,
-            "45": 0.20236,
-            "46": 0.19519,
-            "47": 0.1998,
-            "48": 0.20296,
-            "49": 0.20318,
-            "50": 0.19638,
-            "51": 0.25423,
-            "52": 0.22224,
-            "53": 0.20083,
-            "54": 0.19378,
-            "55": 0.19638,
-            "56": 0.19433,
-            "57": 0.21583,
-            "58": 0.20227,
-            "59": 0.19823,
-            "60": 0.2012,
-            "61": 0.20156,
-            "62": 0.21045,
-            "63": 0.19778,
-            "64": 0.19188,
-            "65": 0.19114,
-            "66": 0.19429,
-            "67": 0.21388,
-            "68": 0.19833,
-            "69": 0.20586,
-            "70": 0.19598,
-            "71": 0.1923,
-            "72": 0.19411,
-            "73": 0.20635,
-            "74": 0.20087,
-            "75": 0.19463,
-            "76": 0.19501,
-            "77": 0.20497,
-            "78": 0.19948,
-            "79": 0.19854,
-            "80": 0.19801,
-            "81": 0.21063,
-            "82": 0.20059,
-            "83": 0.19918,
-            "84": 0.20511,
-            "85": 0.19701,
-            "86": 0.20162,
-            "87": 0.20013,
-            "88": 0.20349,
-            "89": 0.19494,
-            "90": 0.19192,
-            "91": 0.20582,
-            "92": 0.19576,
-            "93": 0.20142,
-            "94": 0.20002,
-            "95": 0.1982,
-            "96": 0.20164,
-            "97": 0.20778,
-            "98": 0.19874,
-            "99": 0.19422,
-            "100": 0.20224
+            "1": "nan",
+            "2": 5.90335,
+            "3": 0.22209,
+            "4": 0.21435,
+            "5": 0.20363,
+            "6": 1.22751,
+            "7": 0.40895,
+            "8": 0.60458,
+            "9": 0.85046,
+            "10": 1.07523,
+            "11": 0.18485,
+            "12": 0.18656,
+            "13": 0.18633,
+            "14": 0.18854,
+            "15": 0.86533,
+            "16": 0.47264,
+            "17": 1.25739,
+            "18": 0.18467,
+            "19": 0.32408,
+            "20": 0.43829,
+            "21": 0.27053,
+            "22": 0.6587,
+            "23": 0.29606,
+            "24": 0.58257,
+            "25": 0.37525,
+            "26": 0.61472,
+            "27": 0.70636,
+            "28": 0.17638,
+            "29": 0.29873,
+            "30": 0.17349,
+            "31": 0.50698,
+            "32": 1.16714,
+            "33": 0.17948,
+            "34": 0.50932,
+            "35": 0.1738,
+            "36": 1.26524,
+            "37": 0.17053,
+            "38": 0.17677,
+            "39": 0.70732,
+            "40": 0.17105,
+            "41": 0.81031,
+            "42": 0.18223,
+            "43": 0.48259,
+            "44": 0.46028,
+            "45": 0.16318,
+            "46": 0.43782,
+            "47": 0.45902,
+            "48": 0.92326,
+            "49": 0.16733,
+            "50": 0.27526,
+            "51": 0.23105,
+            "52": 0.21306,
+            "53": 0.17912,
+            "54": 0.17977,
+            "55": 0.7502,
+            "56": 0.75115,
+            "57": 0.16553,
+            "58": 0.64842,
+            "59": 0.85205,
+            "60": 0.16501,
+            "61": 0.56475,
+            "62": 0.17289,
+            "63": 0.54266,
+            "64": 0.17548,
+            "65": 0.2394,
+            "66": 1.56348,
+            "67": 0.17286,
+            "68": 0.25695,
+            "69": 0.18,
+            "70": 0.74997,
+            "71": 0.3352,
+            "72": 0.26506,
+            "73": 0.57101,
+            "74": 0.18125,
+            "75": 0.74336,
+            "76": 0.17154,
+            "77": 1.5061,
+            "78": 0.16082,
+            "79": 0.28202,
+            "80": 0.41137,
+            "81": 0.18084,
+            "82": 0.24534,
+            "83": 0.41481,
+            "84": 0.60668,
+            "85": 0.86174,
+            "86": 0.17345,
+            "87": 0.32596,
+            "88": 0.17444,
+            "89": 0.62208,
+            "90": 0.39801,
+            "91": 0.38313,
+            "92": 1.18821,
+            "93": 0.16854,
+            "94": 0.36966,
+            "95": 0.17944,
+            "96": 0.18117,
+            "97": 0.26269,
+            "98": 1.21278,
+            "99": 0.23525,
+            "100": 0.40205
         }
     }
 }
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
index 03b84ec3a88..8fc13ff802e 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
@@ -4,106 +4,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 10.90111,
-            "2": 10.8957,
-            "3": 10.90626,
-            "4": 10.90562,
-            "5": 10.91338,
-            "6": 10.89684,
-            "7": 10.89792,
-            "8": 10.91285,
-            "9": 10.89279,
-            "10": 10.90027,
-            "11": 10.89132,
-            "12": 10.90672,
-            "13": 10.87062,
-            "14": 10.86054,
-            "15": 10.87324,
-            "16": 10.85911,
-            "17": 10.84015,
-            "18": 10.83726,
-            "19": 10.83943,
-            "20": 10.79178,
-            "21": 10.75607,
-            "22": 10.75557,
-            "23": 10.75782,
-            "24": 10.72,
-            "25": 10.72107,
-            "26": 10.69764,
-            "27": 10.68234,
-            "28": 10.61163,
-            "29": 10.5913,
-            "30": 10.57438,
-            "31": 10.56065,
-            "32": 10.54263,
-            "33": 10.51706,
-            "34": 10.47723,
-            "35": 10.48572,
-            "36": 10.47878,
+            "1": 10.89993,
+            "2": 10.89521,
+            "3": 10.91741,
+            "4": 10.90771,
+            "5": 10.90843,
+            "6": 10.89602,
+            "7": 10.9013,
+            "8": 10.90712,
+            "9": 10.89564,
+            "10": 10.90213,
+            "11": 10.89284,
+            "12": 10.90839,
+            "13": 10.87585,
+            "14": 10.86373,
+            "15": 10.86587,
+            "16": 10.85724,
+            "17": 10.84217,
+            "18": 10.8334,
+            "19": 10.84191,
+            "20": 10.79367,
+            "21": 10.75583,
+            "22": 10.75242,
+            "23": 10.75941,
+            "24": 10.72283,
+            "25": 10.72766,
+            "26": 10.70191,
+            "27": 10.67962,
+            "28": 10.61078,
+            "29": 10.59488,
+            "30": 10.57553,
+            "31": 10.55887,
+            "32": 10.54327,
+            "33": 10.51446,
+            "34": 10.47754,
+            "35": 10.48312,
+            "36": 10.47431,
             "37": 10.42879,
-            "38": 10.43527,
-            "39": 10.40593,
-            "40": 10.39405,
-            "41": 10.3711,
-            "42": 10.35553,
-            "43": 10.32405,
-            "44": 10.3079,
-            "45": 10.3161,
-            "46": 10.27786,
-            "47": 10.26864,
-            "48": 10.22395,
-            "49": 10.21645,
-            "50": 10.21926,
-            "51": 10.22194,
-            "52": 10.17303,
-            "53": 10.17725,
-            "54": 10.14489,
-            "55": 10.1139,
-            "56": 10.14246,
-            "57": 10.12615,
-            "58": 10.13392,
-            "59": 10.07842,
-            "60": 10.10018,
-            "61": 10.05693,
-            "62": 10.02102,
-            "63": 10.09347,
-            "64": 10.03541,
-            "65": 10.00572,
-            "66": 10.04666,
-            "67": 10.01884,
-            "68": 9.9791,
-            "69": 9.99387,
-            "70": 9.9811,
-            "71": 10.00411,
-            "72": 9.9866,
-            "73": 9.96195,
-            "74": 9.96004,
-            "75": 9.91415,
-            "76": 9.95265,
-            "77": 9.9512,
-            "78": 9.89997,
-            "79": 9.90165,
-            "80": 9.91122,
-            "81": 9.94213,
-            "82": 9.89797,
-            "83": 9.85138,
-            "84": 9.79603,
-            "85": 9.77409,
-            "86": 9.89333,
-            "87": 9.90529,
-            "88": 9.88326,
-            "89": 9.82622,
-            "90": 9.82545,
-            "91": 9.82891,
-            "92": 9.8238,
-            "93": 9.7521,
-            "94": 9.83174,
-            "95": 9.82115,
-            "96": 9.80774,
-            "97": 9.75361,
-            "98": 9.78431,
-            "99": 9.83399,
-            "100": 9.72316
+            "38": 10.43097,
+            "39": 10.40457,
+            "40": 10.39635,
+            "41": 10.36984,
+            "42": 10.35198,
+            "43": 10.32968,
+            "44": 10.30925,
+            "45": 10.31767,
+            "46": 10.27652,
+            "47": 10.26477,
+            "48": 10.2211,
+            "49": 10.2152,
+            "50": 10.21935,
+            "51": 10.22236,
+            "52": 10.17224,
+            "53": 10.17675,
+            "54": 10.14671,
+            "55": 10.11066,
+            "56": 10.1419,
+            "57": 10.12762,
+            "58": 10.13321,
+            "59": 10.07601,
+            "60": 10.10006,
+            "61": 10.05561,
+            "62": 10.02017,
+            "63": 10.09317,
+            "64": 10.03615,
+            "65": 10.00259,
+            "66": 10.04794,
+            "67": 10.01982,
+            "68": 9.98101,
+            "69": 9.99378,
+            "70": 9.98233,
+            "71": 10.00423,
+            "72": 9.98547,
+            "73": 9.96169,
+            "74": 9.95871,
+            "75": 9.91183,
+            "76": 9.95466,
+            "77": 9.95139,
+            "78": 9.89838,
+            "79": 9.89999,
+            "80": 9.9105,
+            "81": 9.94342,
+            "82": 9.89566,
+            "83": 9.85237,
+            "84": 9.79699,
+            "85": 9.77733,
+            "86": 9.89461,
+            "87": 9.90436,
+            "88": 9.88181,
+            "89": 9.8275,
+            "90": 9.82616,
+            "91": 9.8276,
+            "92": 9.8225,
+            "93": 9.75318,
+            "94": 9.83183,
+            "95": 9.8225,
+            "96": 9.80735,
+            "97": 9.75318,
+            "98": 9.78258,
+            "99": 9.83178,
+            "100": 9.72219
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1133.0,
-            "2": 1120.0,
-            "3": 1357.0,
-            "4": 1328.0,
-            "5": 1243.0,
-            "6": 1194.0,
-            "7": 1359.0,
-            "8": 1220.0,
-            "9": 1207.0,
-            "10": 1187.0,
-            "11": 1241.0,
-            "12": 1150.0,
-            "13": 1293.0,
-            "14": 1250.0,
-            "15": 1165.0,
-            "16": 1281.0,
-            "17": 1240.0,
-            "18": 1254.0,
-            "19": 1207.0,
-            "20": 1123.0,
-            "21": 1158.0,
-            "22": 1122.0,
-            "23": 1212.0,
-            "24": 1223.0,
-            "25": 1195.0,
-            "26": 1229.0,
-            "27": 1201.0,
-            "28": 1173.0,
-            "29": 1129.0,
-            "30": 1179.0,
-            "31": 1078.0,
-            "32": 1148.0,
-            "33": 1241.0,
-            "34": 1120.0,
-            "35": 1181.0,
-            "36": 1113.0,
-            "37": 1150.0,
-            "38": 1171.0,
-            "39": 1185.0,
-            "40": 1211.0,
-            "41": 1313.0,
-            "42": 1187.0,
-            "43": 1181.0,
-            "44": 1146.0,
-            "45": 1186.0,
-            "46": 1210.0,
-            "47": 1153.0,
-            "48": 1245.0,
-            "49": 1302.0,
-            "50": 1326.0,
-            "51": 1447.0,
-            "52": 1342.0,
-            "53": 1269.0,
-            "54": 1381.0,
-            "55": 1228.0,
-            "56": 1273.0,
-            "57": 1214.0,
-            "58": 50777.0,
-            "59": 1361.0,
-            "60": 1406.0,
-            "61": 42322.0,
-            "62": 49580.0,
-            "63": 1577.0,
-            "64": 43545.0,
-            "65": 1216.0,
-            "66": 1050479.0,
-            "67": 51707.0,
-            "68": 1528.0,
-            "69": 5467.0,
-            "70": 6590.0,
-            "71": 2418.0,
-            "72": 57795.0,
-            "73": 7777.0,
-            "74": 11706.0,
-            "75": 13832.0,
-            "76": 70208.0,
-            "77": 1681.0,
-            "78": 1563.0,
-            "79": 6703.0,
-            "80": 5581.0,
-            "81": 2703.0,
-            "82": 2750.0,
-            "83": 1486.0,
-            "84": 1050539.0,
-            "85": 1435.0,
-            "86": 1495.0,
-            "87": 2151913.0,
-            "88": 2377.0,
-            "89": 10766.0,
-            "90": 1060091.0,
-            "91": 1258.0,
-            "92": 2099760.0,
-            "93": 11822.0,
-            "94": 1776.0,
-            "95": 1050670.0,
-            "96": 116325.0,
-            "97": 53923.0,
-            "98": 53892.0,
-            "99": 5442.0,
-            "100": 1050643.0
+            "1": 1240.0,
+            "2": 1269.0,
+            "3": 1337.0,
+            "4": 1177.0,
+            "5": 1264.0,
+            "6": 1256.0,
+            "7": 1387.0,
+            "8": 1288.0,
+            "9": 1248.0,
+            "10": 1245.0,
+            "11": 1221.0,
+            "12": 1330.0,
+            "13": 1294.0,
+            "14": 1113.0,
+            "15": 1156.0,
+            "16": 1137.0,
+            "17": 1245.0,
+            "18": 1239.0,
+            "19": 1266.0,
+            "20": 1102.0,
+            "21": 1166.0,
+            "22": 1043.0,
+            "23": 1102.0,
+            "24": 1041.0,
+            "25": 1160.0,
+            "26": 1005.0,
+            "27": 1222.0,
+            "28": 1087.0,
+            "29": 1127.0,
+            "30": 1030.0,
+            "31": 1109.0,
+            "32": 1102.0,
+            "33": 1176.0,
+            "34": 1154.0,
+            "35": 1137.0,
+            "36": 1155.0,
+            "37": 1234.0,
+            "38": 1191.0,
+            "39": 1126.0,
+            "40": 1181.0,
+            "41": 1175.0,
+            "42": 1193.0,
+            "43": 1227.0,
+            "44": 1176.0,
+            "45": 1198.0,
+            "46": 1239.0,
+            "47": 1291.0,
+            "48": 1305.0,
+            "49": 1204.0,
+            "50": 1258.0,
+            "51": 1352.0,
+            "52": 53542.0,
+            "53": 1378.0,
+            "54": 1380.0,
+            "55": 1242.0,
+            "56": 2362.0,
+            "57": 63620.0,
+            "58": 1645.0,
+            "59": 1450.0,
+            "60": 1425.0,
+            "61": 1363.0,
+            "62": 1520.0,
+            "63": 48673.0,
+            "64": 1536.0,
+            "65": 1162.0,
+            "66": 9633.0,
+            "67": 1050634.0,
+            "68": 9678.0,
+            "69": 1106725.0,
+            "70": 1050586.0,
+            "71": 5618.0,
+            "72": 6508.0,
+            "73": 8688.0,
+            "74": 4522.0,
+            "75": 8597.0,
+            "76": 1562.0,
+            "77": 7800.0,
+            "78": 1654.0,
+            "79": 1566.0,
+            "80": 1479.0,
+            "81": 1655.0,
+            "82": 1646.0,
+            "83": 1464.0,
+            "84": 1502.0,
+            "85": 1386.0,
+            "86": 1542.0,
+            "87": 1157139.0,
+            "88": 6676.0,
+            "89": 46522.0,
+            "90": 1109010.0,
+            "91": 2355.0,
+            "92": 1604.0,
+            "93": 1388.0,
+            "94": 6804.0,
+            "95": 1060856.0,
+            "96": 107239.0,
+            "97": 59055.0,
+            "98": 49897.0,
+            "99": 1534.0,
+            "100": 1050768.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 788519936.0,
-            "2": 788510720.0,
-            "3": 788485120.0,
-            "4": 788465664.0,
-            "5": 788514816.0,
-            "6": 788501504.0,
-            "7": 788426240.0,
-            "8": 788464128.0,
-            "9": 788466176.0,
-            "10": 788499456.0,
-            "11": 788475392.0,
-            "12": 788483584.0,
-            "13": 788490240.0,
-            "14": 788493312.0,
-            "15": 788476416.0,
-            "16": 788467712.0,
-            "17": 788476416.0,
-            "18": 788483584.0,
-            "19": 788512256.0,
-            "20": 788488704.0,
-            "21": 788551680.0,
-            "22": 788594176.0,
-            "23": 788579328.0,
-            "24": 788489728.0,
-            "25": 788578816.0,
-            "26": 788620288.0,
-            "27": 788563456.0,
-            "28": 788538368.0,
-            "29": 788597760.0,
-            "30": 788569600.0,
-            "31": 788587008.0,
-            "32": 788580864.0,
-            "33": 788574208.0,
-            "34": 788563456.0,
-            "35": 788553216.0,
-            "36": 788605440.0,
-            "37": 788550144.0,
-            "38": 788575744.0,
-            "39": 788595200.0,
-            "40": 788552704.0,
-            "41": 788545536.0,
-            "42": 788503552.0,
-            "43": 788534272.0,
-            "44": 788486144.0,
-            "45": 788453888.0,
-            "46": 788497408.0,
-            "47": 788449792.0,
-            "48": 788434944.0,
-            "49": 788413952.0,
-            "50": 788384256.0,
-            "51": 788398592.0,
-            "52": 788370944.0,
-            "53": 788320256.0,
-            "54": 788355584.0,
-            "55": 788348416.0,
-            "56": 788313088.0,
-            "57": 788386816.0,
-            "58": 788362752.0,
-            "59": 788348416.0,
-            "60": 788368896.0,
-            "61": 788305920.0,
-            "62": 788297728.0,
-            "63": 788329472.0,
-            "64": 788335104.0,
-            "65": 788324352.0,
-            "66": 788331008.0,
-            "67": 788335104.0,
-            "68": 788327424.0,
-            "69": 788331520.0,
-            "70": 788349952.0,
-            "71": 788344832.0,
-            "72": 788392960.0,
-            "73": 788349952.0,
-            "74": 788381184.0,
-            "75": 788319744.0,
-            "76": 788291584.0,
-            "77": 788393472.0,
-            "78": 788356096.0,
-            "79": 788316160.0,
-            "80": 788275712.0,
-            "81": 788283904.0,
-            "82": 788240896.0,
-            "83": 788277248.0,
-            "84": 788317696.0,
-            "85": 788314624.0,
-            "86": 788337664.0,
-            "87": 788147200.0,
-            "88": 788256256.0,
-            "89": 788264448.0,
-            "90": 788247552.0,
-            "91": 788319232.0,
-            "92": 788239872.0,
-            "93": 788245504.0,
-            "94": 788387840.0,
-            "95": 788252672.0,
-            "96": 788294656.0,
-            "97": 788360192.0,
-            "98": 788326912.0,
-            "99": 788267008.0,
-            "100": 788357120.0
+            "1": 790002176.0,
+            "2": 790063104.0,
+            "3": 789992960.0,
+            "4": 789929984.0,
+            "5": 790002176.0,
+            "6": 790040576.0,
+            "7": 789941760.0,
+            "8": 789952512.0,
+            "9": 789984768.0,
+            "10": 790000640.0,
+            "11": 789958656.0,
+            "12": 789986304.0,
+            "13": 790013440.0,
+            "14": 789977600.0,
+            "15": 790033408.0,
+            "16": 790013952.0,
+            "17": 789932544.0,
+            "18": 789990912.0,
+            "19": 790029824.0,
+            "20": 790024192.0,
+            "21": 790074880.0,
+            "22": 790069760.0,
+            "23": 790042624.0,
+            "24": 790045184.0,
+            "25": 790058496.0,
+            "26": 790133760.0,
+            "27": 790103040.0,
+            "28": 790033408.0,
+            "29": 790097920.0,
+            "30": 790106112.0,
+            "31": 790086656.0,
+            "32": 790054400.0,
+            "33": 790088192.0,
+            "34": 790109696.0,
+            "35": 790074880.0,
+            "36": 790070784.0,
+            "37": 790050816.0,
+            "38": 790038016.0,
+            "39": 790094848.0,
+            "40": 790046208.0,
+            "41": 790056448.0,
+            "42": 789992960.0,
+            "43": 790049792.0,
+            "44": 790007808.0,
+            "45": 789982720.0,
+            "46": 790012928.0,
+            "47": 789961728.0,
+            "48": 789908480.0,
+            "49": 789893632.0,
+            "50": 789886464.0,
+            "51": 789872128.0,
+            "52": 789878784.0,
+            "53": 789821952.0,
+            "54": 789839360.0,
+            "55": 789838848.0,
+            "56": 789814272.0,
+            "57": 789883904.0,
+            "58": 789877760.0,
+            "59": 789819904.0,
+            "60": 789857280.0,
+            "61": 789826560.0,
+            "62": 789813760.0,
+            "63": 789826048.0,
+            "64": 789858304.0,
+            "65": 789836800.0,
+            "66": 789824000.0,
+            "67": 789825024.0,
+            "68": 789836800.0,
+            "69": 789865472.0,
+            "70": 789832704.0,
+            "71": 789809152.0,
+            "72": 789897728.0,
+            "73": 789852672.0,
+            "74": 789878272.0,
+            "75": 789811712.0,
+            "76": 789782528.0,
+            "77": 789867520.0,
+            "78": 789883392.0,
+            "79": 789828608.0,
+            "80": 789794304.0,
+            "81": 789789696.0,
+            "82": 789745664.0,
+            "83": 789785088.0,
+            "84": 789807104.0,
+            "85": 789825536.0,
+            "86": 789842432.0,
+            "87": 789649920.0,
+            "88": 789756416.0,
+            "89": 789772288.0,
+            "90": 789752832.0,
+            "91": 789829120.0,
+            "92": 789762048.0,
+            "93": 789776384.0,
+            "94": 789902336.0,
+            "95": 789748736.0,
+            "96": 789780480.0,
+            "97": 789849088.0,
+            "98": 789834752.0,
+            "99": 789754368.0,
+            "100": 789845504.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,212 +325,213 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 3121236480.0,
-            "2": 3284394496.0,
-            "3": 3284394496.0,
-            "4": 3284394496.0,
-            "5": 3289592320.0,
-            "6": 3289592320.0,
-            "7": 3289592320.0,
-            "8": 3289592320.0,
-            "9": 3289592320.0,
-            "10": 3289592320.0,
-            "11": 3289592320.0,
-            "12": 3289592320.0,
-            "13": 3289592320.0,
-            "14": 3289592320.0,
-            "15": 3289592320.0,
-            "16": 3289592320.0,
-            "17": 3289592320.0,
-            "18": 3289592320.0,
-            "19": 3289592320.0,
-            "20": 3289592320.0,
-            "21": 3315035648.0,
-            "22": 3342148608.0,
-            "23": 3342148608.0,
-            "24": 3342148608.0,
-            "25": 3342148608.0,
-            "26": 3359425024.0,
-            "27": 3359425024.0,
-            "28": 3359425024.0,
-            "29": 3359425024.0,
-            "30": 3359425024.0,
-            "31": 3359425024.0,
-            "32": 3359425024.0,
-            "33": 3359425024.0,
-            "34": 3359425024.0,
-            "35": 3359425024.0,
-            "36": 3359425024.0,
-            "37": 3359425024.0,
-            "38": 3359425024.0,
-            "39": 3359425024.0,
-            "40": 3359425024.0,
-            "41": 3359425024.0,
-            "42": 3359425024.0,
-            "43": 3359425024.0,
-            "44": 3359425024.0,
-            "45": 3359425024.0,
-            "46": 3359425024.0,
-            "47": 3359425024.0,
-            "48": 3359425024.0,
-            "49": 3359425024.0,
-            "50": 3359425024.0,
-            "51": 3359425024.0,
-            "52": 3359425024.0,
-            "53": 3359425024.0,
-            "54": 3359425024.0,
-            "55": 3359425024.0,
-            "56": 3359425024.0,
-            "57": 3359425024.0,
-            "58": 3359425024.0,
-            "59": 3359425024.0,
-            "60": 3359425024.0,
-            "61": 3359425024.0,
-            "62": 3359425024.0,
-            "63": 3359425024.0,
-            "64": 3359425024.0,
-            "65": 3359425024.0,
-            "66": 3359425024.0,
-            "67": 3359425024.0,
-            "68": 3359425024.0,
-            "69": 3359425024.0,
-            "70": 3359425024.0,
-            "71": 3359425024.0,
-            "72": 3359425024.0,
-            "73": 3359425024.0,
-            "74": 3359425024.0,
-            "75": 3359425024.0,
-            "76": 3359425024.0,
-            "77": 3359425024.0,
-            "78": 3359425024.0,
-            "79": 3359425024.0,
-            "80": 3359425024.0,
-            "81": 3359425024.0,
-            "82": 3359425024.0,
-            "83": 3359425024.0,
-            "84": 3359425024.0,
-            "85": 3359425024.0,
-            "86": 3359425024.0,
-            "87": 3359425024.0,
-            "88": 3359425024.0,
-            "89": 3359425024.0,
-            "90": 3359425024.0,
-            "91": 3359425024.0,
-            "92": 3359425024.0,
-            "93": 3359425024.0,
-            "94": 3359425024.0,
-            "95": 3359425024.0,
-            "96": 3359425024.0,
-            "97": 3359425024.0,
-            "98": 3359425024.0,
-            "99": 3359425024.0,
-            "100": 3359425024.0
+            "1": 3107662336.0,
+            "2": 3321466368.0,
+            "3": 3321466368.0,
+            "4": 3321466368.0,
+            "5": 3321466368.0,
+            "6": 3321466368.0,
+            "7": 3321466368.0,
+            "8": 3321466368.0,
+            "9": 3321466368.0,
+            "10": 3321466368.0,
+            "11": 3321466368.0,
+            "12": 3321466368.0,
+            "13": 3321466368.0,
+            "14": 3321466368.0,
+            "15": 3321466368.0,
+            "16": 3321466368.0,
+            "17": 3321466368.0,
+            "18": 3321466368.0,
+            "19": 3321466368.0,
+            "20": 3321466368.0,
+            "21": 3330454528.0,
+            "22": 3330454528.0,
+            "23": 3330454528.0,
+            "24": 3330454528.0,
+            "25": 3330454528.0,
+            "26": 3368265216.0,
+            "27": 3368265216.0,
+            "28": 3368265216.0,
+            "29": 3368265216.0,
+            "30": 3368265216.0,
+            "31": 3368265216.0,
+            "32": 3368265216.0,
+            "33": 3368265216.0,
+            "34": 3368265216.0,
+            "35": 3368265216.0,
+            "36": 3368265216.0,
+            "37": 3368265216.0,
+            "38": 3368265216.0,
+            "39": 3368265216.0,
+            "40": 3368265216.0,
+            "41": 3368265216.0,
+            "42": 3368265216.0,
+            "43": 3368265216.0,
+            "44": 3368265216.0,
+            "45": 3368265216.0,
+            "46": 3368265216.0,
+            "47": 3368265216.0,
+            "48": 3368265216.0,
+            "49": 3368265216.0,
+            "50": 3368265216.0,
+            "51": 3368265216.0,
+            "52": 3368265216.0,
+            "53": 3368265216.0,
+            "54": 3368265216.0,
+            "55": 3368265216.0,
+            "56": 3368265216.0,
+            "57": 3368265216.0,
+            "58": 3368265216.0,
+            "59": 3368265216.0,
+            "60": 3368265216.0,
+            "61": 3368265216.0,
+            "62": 3368265216.0,
+            "63": 3368265216.0,
+            "64": 3368265216.0,
+            "65": 3368265216.0,
+            "66": 3368265216.0,
+            "67": 3368265216.0,
+            "68": 3368265216.0,
+            "69": 3368265216.0,
+            "70": 3368265216.0,
+            "71": 3368265216.0,
+            "72": 3368265216.0,
+            "73": 3368265216.0,
+            "74": 3368265216.0,
+            "75": 3368265216.0,
+            "76": 3368265216.0,
+            "77": 3368265216.0,
+            "78": 3368265216.0,
+            "79": 3368265216.0,
+            "80": 3368265216.0,
+            "81": 3368265216.0,
+            "82": 3368265216.0,
+            "83": 3368265216.0,
+            "84": 3368265216.0,
+            "85": 3368265216.0,
+            "86": 3368265216.0,
+            "87": 3368265216.0,
+            "88": 3368265216.0,
+            "89": 3368265216.0,
+            "90": 3368265216.0,
+            "91": 3368265216.0,
+            "92": 3368265216.0,
+            "93": 3368265216.0,
+            "94": 3368265216.0,
+            "95": 3368265216.0,
+            "96": 3368265216.0,
+            "97": 3368265216.0,
+            "98": 3368265216.0,
+            "99": 3368265216.0,
+            "100": 3368265216.0
         }
     },
     "iteration-time": {
-        "start_step": 2,
+        "start_step": 1,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "2": 4.52323,
-            "3": 0.19663,
-            "4": 0.17543,
-            "5": 0.19069,
-            "6": 0.17051,
-            "7": 0.17616,
-            "8": 0.18284,
-            "9": 0.17765,
-            "10": 0.17291,
-            "11": 0.18329,
-            "12": 0.15973,
-            "13": 0.17233,
-            "14": 0.16664,
-            "15": 0.16783,
-            "16": 0.16783,
-            "17": 0.17292,
-            "18": 0.15603,
-            "19": 0.15638,
-            "20": 0.16092,
-            "21": 0.17656,
-            "22": 0.1677,
-            "23": 0.15908,
-            "24": 0.16291,
-            "25": 0.184,
-            "26": 0.15922,
-            "27": 0.17381,
-            "28": 0.18227,
-            "29": 0.1705,
-            "30": 0.17401,
-            "31": 0.1729,
-            "32": 0.18045,
-            "33": 0.17736,
-            "34": 0.18122,
-            "35": 0.15835,
-            "36": 0.17971,
-            "37": 0.17994,
-            "38": 0.18053,
-            "39": 0.18529,
-            "40": 0.16901,
-            "41": 0.16177,
-            "42": 0.17084,
-            "43": 0.16807,
-            "44": 0.16785,
-            "45": 0.16278,
-            "46": 0.17404,
-            "47": 0.16869,
-            "48": 0.15239,
-            "49": 0.15444,
-            "50": 0.16228,
-            "51": 1.71136,
-            "52": 0.19743,
-            "53": 0.22424,
-            "54": 0.15535,
-            "55": 0.15493,
-            "56": 0.14309,
-            "57": 0.16908,
-            "58": 0.21236,
-            "59": 0.15781,
-            "60": 0.21784,
-            "61": 0.14226,
-            "62": 0.17439,
-            "63": 0.22118,
-            "64": 0.15047,
-            "65": 0.16162,
-            "66": 0.17192,
-            "67": 0.15749,
-            "68": 0.14332,
-            "69": 0.16492,
-            "70": 0.15042,
-            "71": 0.17996,
-            "72": 0.15294,
-            "73": 0.21192,
-            "74": 0.15977,
-            "75": 0.15858,
-            "76": 0.15006,
-            "77": 0.16115,
-            "78": 0.15397,
-            "79": 0.15516,
-            "80": 0.14733,
-            "81": 0.15321,
-            "82": 0.15755,
-            "83": 0.15836,
-            "84": 0.15737,
-            "85": 0.1536,
-            "86": 0.15088,
-            "87": 0.147,
-            "88": 0.15501,
-            "89": 0.15595,
-            "90": 0.14272,
-            "91": 0.15399,
-            "92": 0.1428,
-            "93": 0.15036,
-            "94": 0.14851,
-            "95": 0.14892,
-            "96": 0.14683,
-            "97": 0.15235,
-            "98": 0.14893,
-            "99": 0.14814,
-            "100": 0.14019
+            "1": "nan",
+            "2": 6.72002,
+            "3": 0.19153,
+            "4": 0.1869,
+            "5": 0.18093,
+            "6": 1.71816,
+            "7": 0.43956,
+            "8": 1.22852,
+            "9": 0.50797,
+            "10": 1.77116,
+            "11": 0.15385,
+            "12": 0.1694,
+            "13": 0.16278,
+            "14": 0.73069,
+            "15": 0.74662,
+            "16": 0.6977,
+            "17": 1.59318,
+            "18": 0.30546,
+            "19": 0.16535,
+            "20": 0.81437,
+            "21": 0.1781,
+            "22": 0.85825,
+            "23": 0.24586,
+            "24": 0.59259,
+            "25": 0.6168,
+            "26": 0.8908,
+            "27": 0.58788,
+            "28": 0.31692,
+            "29": 0.3525,
+            "30": 0.17486,
+            "31": 1.71875,
+            "32": 0.78013,
+            "33": 0.17702,
+            "34": 0.17049,
+            "35": 0.59279,
+            "36": 0.5626,
+            "37": 0.47382,
+            "38": 0.38311,
+            "39": 0.94987,
+            "40": 0.16562,
+            "41": 1.44828,
+            "42": 0.16307,
+            "43": 0.52293,
+            "44": 0.15854,
+            "45": 0.97664,
+            "46": 0.15531,
+            "47": 1.14048,
+            "48": 0.25038,
+            "49": 0.91494,
+            "50": 0.15665,
+            "51": 0.21462,
+            "52": 0.20466,
+            "53": 0.16067,
+            "54": 0.16037,
+            "55": 0.93216,
+            "56": 1.23158,
+            "57": 0.1597,
+            "58": 0.32695,
+            "59": 0.70656,
+            "60": 0.15041,
+            "61": 0.7218,
+            "62": 0.18484,
+            "63": 0.66162,
+            "64": 0.15003,
+            "65": 0.78038,
+            "66": 0.86286,
+            "67": 0.15286,
+            "68": 0.41516,
+            "69": 0.80928,
+            "70": 0.29023,
+            "71": 0.4049,
+            "72": 0.1661,
+            "73": 0.86413,
+            "74": 0.4669,
+            "75": 0.37553,
+            "76": 0.28759,
+            "77": 1.21568,
+            "78": 0.15014,
+            "79": 0.29435,
+            "80": 0.62782,
+            "81": 0.15723,
+            "82": 0.30613,
+            "83": 0.33307,
+            "84": 0.40998,
+            "85": 0.86772,
+            "86": 0.21281,
+            "87": 0.77287,
+            "88": 0.1498,
+            "89": 0.70194,
+            "90": 0.39284,
+            "91": 0.43514,
+            "92": 1.35257,
+            "93": 0.37058,
+            "94": 0.15217,
+            "95": 0.14,
+            "96": 0.13888,
+            "97": 0.59838,
+            "98": 1.14766,
+            "99": 0.36092,
+            "100": 0.33679
         }
     }
 }
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/golden_values_dev_dgx_h100.json
index 4405a6078dd..aa854eb5860 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_muon/golden_values_dev_dgx_h100.json
@@ -4,106 +4,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 10.90111,
-            "2": 10.8957,
-            "3": 10.90636,
-            "4": 10.90618,
-            "5": 10.91363,
-            "6": 10.89761,
-            "7": 10.89776,
-            "8": 10.91415,
-            "9": 10.89561,
-            "10": 10.9037,
-            "11": 10.90078,
-            "12": 10.91896,
-            "13": 10.88436,
-            "14": 10.8755,
-            "15": 10.91016,
-            "16": 10.90356,
-            "17": 10.88647,
-            "18": 10.89373,
-            "19": 10.89488,
-            "20": 10.89145,
-            "21": 10.8867,
-            "22": 10.90339,
-            "23": 10.9122,
-            "24": 10.88545,
-            "25": 10.90147,
-            "26": 10.88807,
-            "27": 10.89667,
-            "28": 10.88181,
-            "29": 10.88999,
-            "30": 10.91088,
-            "31": 10.89164,
-            "32": 10.88874,
-            "33": 10.89699,
-            "34": 10.87341,
-            "35": 10.89671,
-            "36": 10.90754,
-            "37": 10.87082,
-            "38": 10.87845,
-            "39": 10.8861,
-            "40": 10.89256,
-            "41": 10.88453,
-            "42": 10.89553,
-            "43": 10.88019,
-            "44": 10.88447,
-            "45": 10.88361,
-            "46": 10.88675,
-            "47": 10.88597,
-            "48": 10.8638,
-            "49": 10.8766,
-            "50": 10.88241,
-            "51": 10.89271,
-            "52": 10.87137,
-            "53": 10.85905,
-            "54": 10.87009,
-            "55": 10.86786,
-            "56": 10.87291,
-            "57": 10.84458,
-            "58": 10.85672,
-            "59": 10.84682,
-            "60": 10.84354,
-            "61": 10.86079,
-            "62": 10.85703,
-            "63": 10.86049,
-            "64": 10.83806,
-            "65": 10.82753,
-            "66": 10.84631,
-            "67": 10.82788,
-            "68": 10.83033,
-            "69": 10.81867,
-            "70": 10.82807,
-            "71": 10.82687,
-            "72": 10.80846,
-            "73": 10.80745,
-            "74": 10.8061,
-            "75": 10.81421,
-            "76": 10.81095,
-            "77": 10.80817,
-            "78": 10.79232,
-            "79": 10.80249,
-            "80": 10.78847,
-            "81": 10.79536,
-            "82": 10.79574,
-            "83": 10.78549,
-            "84": 10.75707,
-            "85": 10.76214,
-            "86": 10.77729,
-            "87": 10.79586,
-            "88": 10.77534,
-            "89": 10.77602,
-            "90": 10.76492,
-            "91": 10.74145,
-            "92": 10.75946,
-            "93": 10.74599,
-            "94": 10.73392,
-            "95": 10.75202,
-            "96": 10.72408,
-            "97": 10.71554,
-            "98": 10.7261,
-            "99": 10.74663,
-            "100": 10.69609
+            "1": 10.89993,
+            "2": 10.89521,
+            "3": 10.91767,
+            "4": 10.90821,
+            "5": 10.90822,
+            "6": 10.89634,
+            "7": 10.90136,
+            "8": 10.90825,
+            "9": 10.89914,
+            "10": 10.9051,
+            "11": 10.90169,
+            "12": 10.92025,
+            "13": 10.88855,
+            "14": 10.87854,
+            "15": 10.9023,
+            "16": 10.90045,
+            "17": 10.88668,
+            "18": 10.89099,
+            "19": 10.89552,
+            "20": 10.88995,
+            "21": 10.88425,
+            "22": 10.89603,
+            "23": 10.91293,
+            "24": 10.89092,
+            "25": 10.9068,
+            "26": 10.89456,
+            "27": 10.89204,
+            "28": 10.88144,
+            "29": 10.89325,
+            "30": 10.90832,
+            "31": 10.89209,
+            "32": 10.88811,
+            "33": 10.89532,
+            "34": 10.87436,
+            "35": 10.89132,
+            "36": 10.90873,
+            "37": 10.87951,
+            "38": 10.87438,
+            "39": 10.87785,
+            "40": 10.8965,
+            "41": 10.88187,
+            "42": 10.89192,
+            "43": 10.87902,
+            "44": 10.88738,
+            "45": 10.89155,
+            "46": 10.87988,
+            "47": 10.88225,
+            "48": 10.86425,
+            "49": 10.87226,
+            "50": 10.8889,
+            "51": 10.88691,
+            "52": 10.87174,
+            "53": 10.86032,
+            "54": 10.86321,
+            "55": 10.86846,
+            "56": 10.87538,
+            "57": 10.84295,
+            "58": 10.85647,
+            "59": 10.84623,
+            "60": 10.8423,
+            "61": 10.8623,
+            "62": 10.85422,
+            "63": 10.86208,
+            "64": 10.83791,
+            "65": 10.82375,
+            "66": 10.8455,
+            "67": 10.82865,
+            "68": 10.8278,
+            "69": 10.81675,
+            "70": 10.82691,
+            "71": 10.82914,
+            "72": 10.81178,
+            "73": 10.80744,
+            "74": 10.80684,
+            "75": 10.81184,
+            "76": 10.80973,
+            "77": 10.80504,
+            "78": 10.79723,
+            "79": 10.80512,
+            "80": 10.78664,
+            "81": 10.79319,
+            "82": 10.79338,
+            "83": 10.78459,
+            "84": 10.75321,
+            "85": 10.76227,
+            "86": 10.77348,
+            "87": 10.79634,
+            "88": 10.77437,
+            "89": 10.76868,
+            "90": 10.76446,
+            "91": 10.74212,
+            "92": 10.75828,
+            "93": 10.74918,
+            "94": 10.73459,
+            "95": 10.74833,
+            "96": 10.71952,
+            "97": 10.71332,
+            "98": 10.72383,
+            "99": 10.74522,
+            "100": 10.69981
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1133.0,
-            "2": 1120.0,
-            "3": 1360.0,
-            "4": 1289.0,
-            "5": 1206.0,
-            "6": 1218.0,
-            "7": 1344.0,
-            "8": 1181.0,
-            "9": 1206.0,
-            "10": 1139.0,
-            "11": 1252.0,
-            "12": 1218.0,
-            "13": 1288.0,
-            "14": 1162.0,
-            "15": 1248.0,
-            "16": 1265.0,
-            "17": 1188.0,
-            "18": 1115.0,
-            "19": 1172.0,
-            "20": 1089.0,
-            "21": 1213.0,
-            "22": 1244.0,
-            "23": 1290.0,
-            "24": 1071.0,
-            "25": 1243.0,
-            "26": 1216.0,
-            "27": 1393.0,
-            "28": 1179.0,
-            "29": 1194.0,
-            "30": 1353.0,
-            "31": 1240.0,
-            "32": 1343.0,
-            "33": 1190.0,
-            "34": 1196.0,
-            "35": 1270.0,
-            "36": 1169.0,
-            "37": 1196.0,
-            "38": 1212.0,
-            "39": 1344.0,
-            "40": 1272.0,
-            "41": 1350.0,
-            "42": 1153.0,
-            "43": 1231.0,
-            "44": 1234.0,
-            "45": 1094.0,
-            "46": 1458.0,
-            "47": 1192.0,
-            "48": 1140.0,
-            "49": 1412.0,
-            "50": 1316.0,
-            "51": 1261.0,
-            "52": 1161.0,
-            "53": 1242.0,
-            "54": 1183.0,
-            "55": 1246.0,
-            "56": 1243.0,
-            "57": 1208.0,
-            "58": 1283.0,
-            "59": 1108.0,
-            "60": 1224.0,
-            "61": 1277.0,
-            "62": 1192.0,
-            "63": 1234.0,
-            "64": 1285.0,
-            "65": 1166.0,
-            "66": 1238.0,
-            "67": 1128.0,
-            "68": 1281.0,
-            "69": 1115.0,
-            "70": 1255.0,
-            "71": 1199.0,
-            "72": 1239.0,
-            "73": 1264.0,
-            "74": 1318.0,
-            "75": 1344.0,
-            "76": 1250.0,
-            "77": 1197.0,
-            "78": 1241.0,
-            "79": 1270.0,
-            "80": 1227.0,
-            "81": 1238.0,
-            "82": 1078.0,
-            "83": 1220.0,
-            "84": 1190.0,
-            "85": 1187.0,
-            "86": 1238.0,
-            "87": 1074.0,
-            "88": 1182.0,
-            "89": 1184.0,
-            "90": 1343.0,
-            "91": 1308.0,
-            "92": 1173.0,
-            "93": 1342.0,
-            "94": 1400.0,
-            "95": 1059.0,
-            "96": 1129.0,
-            "97": 1263.0,
-            "98": 1344.0,
-            "99": 1065.0,
-            "100": 1185.0
+            "1": 1240.0,
+            "2": 1269.0,
+            "3": 1308.0,
+            "4": 1111.0,
+            "5": 1207.0,
+            "6": 1300.0,
+            "7": 1364.0,
+            "8": 1332.0,
+            "9": 1270.0,
+            "10": 1270.0,
+            "11": 1199.0,
+            "12": 1358.0,
+            "13": 1266.0,
+            "14": 1242.0,
+            "15": 1067.0,
+            "16": 1155.0,
+            "17": 1211.0,
+            "18": 1120.0,
+            "19": 1209.0,
+            "20": 1115.0,
+            "21": 1205.0,
+            "22": 1281.0,
+            "23": 1272.0,
+            "24": 1349.0,
+            "25": 1210.0,
+            "26": 1223.0,
+            "27": 1327.0,
+            "28": 1196.0,
+            "29": 1188.0,
+            "30": 1229.0,
+            "31": 1305.0,
+            "32": 1205.0,
+            "33": 1160.0,
+            "34": 1161.0,
+            "35": 1080.0,
+            "36": 1219.0,
+            "37": 1273.0,
+            "38": 1390.0,
+            "39": 1392.0,
+            "40": 1424.0,
+            "41": 1084.0,
+            "42": 1303.0,
+            "43": 1238.0,
+            "44": 1116.0,
+            "45": 1316.0,
+            "46": 1361.0,
+            "47": 1140.0,
+            "48": 1253.0,
+            "49": 1281.0,
+            "50": 1067.0,
+            "51": 1180.0,
+            "52": 1101.0,
+            "53": 1224.0,
+            "54": 1201.0,
+            "55": 1139.0,
+            "56": 1086.0,
+            "57": 1221.0,
+            "58": 1375.0,
+            "59": 1154.0,
+            "60": 1161.0,
+            "61": 1167.0,
+            "62": 1257.0,
+            "63": 1153.0,
+            "64": 1322.0,
+            "65": 1122.0,
+            "66": 1154.0,
+            "67": 1140.0,
+            "68": 1372.0,
+            "69": 1324.0,
+            "70": 1235.0,
+            "71": 1255.0,
+            "72": 1111.0,
+            "73": 1251.0,
+            "74": 1301.0,
+            "75": 1363.0,
+            "76": 1339.0,
+            "77": 1210.0,
+            "78": 1283.0,
+            "79": 1189.0,
+            "80": 1309.0,
+            "81": 1281.0,
+            "82": 1230.0,
+            "83": 1270.0,
+            "84": 1288.0,
+            "85": 1296.0,
+            "86": 1236.0,
+            "87": 1023.0,
+            "88": 1258.0,
+            "89": 1285.0,
+            "90": 1407.0,
+            "91": 1244.0,
+            "92": 1307.0,
+            "93": 1216.0,
+            "94": 1303.0,
+            "95": 1237.0,
+            "96": 1134.0,
+            "97": 1196.0,
+            "98": 1317.0,
+            "99": 1349.0,
+            "100": 1146.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1096719872.0,
-            "2": 1096710656.0,
-            "3": 1096685568.0,
-            "4": 1096666112.0,
-            "5": 1096716288.0,
-            "6": 1096698368.0,
-            "7": 1096625664.0,
-            "8": 1096664064.0,
-            "9": 1096666624.0,
-            "10": 1096697344.0,
-            "11": 1096671744.0,
-            "12": 1096680448.0,
-            "13": 1096684032.0,
-            "14": 1096687104.0,
-            "15": 1096663040.0,
-            "16": 1096648192.0,
-            "17": 1096657920.0,
-            "18": 1096658944.0,
-            "19": 1096689152.0,
-            "20": 1096654336.0,
-            "21": 1096707584.0,
-            "22": 1096734208.0,
-            "23": 1096705024.0,
-            "24": 1096621568.0,
-            "25": 1096701952.0,
-            "26": 1096734208.0,
-            "27": 1096685056.0,
-            "28": 1096639488.0,
-            "29": 1096713216.0,
-            "30": 1096646656.0,
-            "31": 1096652800.0,
-            "32": 1096659456.0,
-            "33": 1096667648.0,
-            "34": 1096645120.0,
-            "35": 1096631296.0,
-            "36": 1096694272.0,
-            "37": 1096644096.0,
-            "38": 1096703488.0,
-            "39": 1096710144.0,
-            "40": 1096657408.0,
-            "41": 1096678912.0,
-            "42": 1096621568.0,
-            "43": 1096688128.0,
-            "44": 1096673792.0,
-            "45": 1096646656.0,
-            "46": 1096709632.0,
-            "47": 1096666624.0,
-            "48": 1096695808.0,
-            "49": 1096656896.0,
-            "50": 1096704000.0,
-            "51": 1096687616.0,
-            "52": 1096661504.0,
-            "53": 1096610816.0,
-            "54": 1096705024.0,
-            "55": 1096696320.0,
-            "56": 1096693248.0,
-            "57": 1096753664.0,
-            "58": 1096691712.0,
-            "59": 1096695808.0,
-            "60": 1096675840.0,
-            "61": 1096647168.0,
-            "62": 1096636416.0,
-            "63": 1096645632.0,
-            "64": 1096672256.0,
-            "65": 1096679936.0,
-            "66": 1096693248.0,
-            "67": 1096674816.0,
-            "68": 1096677376.0,
-            "69": 1096642560.0,
-            "70": 1096725504.0,
-            "71": 1096706560.0,
-            "72": 1096703488.0,
-            "73": 1096632320.0,
-            "74": 1096673792.0,
-            "75": 1096717312.0,
-            "76": 1096673280.0,
-            "77": 1096714752.0,
-            "78": 1096664064.0,
-            "79": 1096729600.0,
-            "80": 1096649216.0,
-            "81": 1096703488.0,
-            "82": 1096649728.0,
-            "83": 1096710144.0,
-            "84": 1096683008.0,
-            "85": 1096688640.0,
-            "86": 1096672256.0,
-            "87": 1096687104.0,
-            "88": 1096668672.0,
-            "89": 1096658944.0,
-            "90": 1096683008.0,
-            "91": 1096666112.0,
-            "92": 1096664576.0,
-            "93": 1096695808.0,
-            "94": 1096671744.0,
-            "95": 1096692736.0,
-            "96": 1096684544.0,
-            "97": 1096652800.0,
-            "98": 1096691200.0,
-            "99": 1096708608.0,
-            "100": 1096738304.0
+            "1": 1096325632.0,
+            "2": 1096386560.0,
+            "3": 1096315392.0,
+            "4": 1096254976.0,
+            "5": 1096325120.0,
+            "6": 1096363520.0,
+            "7": 1096264192.0,
+            "8": 1096275968.0,
+            "9": 1096305664.0,
+            "10": 1096323584.0,
+            "11": 1096276480.0,
+            "12": 1096307712.0,
+            "13": 1096329728.0,
+            "14": 1096292864.0,
+            "15": 1096339968.0,
+            "16": 1096317440.0,
+            "17": 1096236032.0,
+            "18": 1096297472.0,
+            "19": 1096334336.0,
+            "20": 1096314368.0,
+            "21": 1096350208.0,
+            "22": 1096322560.0,
+            "23": 1096301568.0,
+            "24": 1096308224.0,
+            "25": 1096311808.0,
+            "26": 1096379904.0,
+            "27": 1096344064.0,
+            "28": 1096274944.0,
+            "29": 1096314368.0,
+            "30": 1096304128.0,
+            "31": 1096288256.0,
+            "32": 1096247808.0,
+            "33": 1096303616.0,
+            "34": 1096329216.0,
+            "35": 1096293376.0,
+            "36": 1096297472.0,
+            "37": 1096264704.0,
+            "38": 1096275968.0,
+            "39": 1096357376.0,
+            "40": 1096279552.0,
+            "41": 1096329216.0,
+            "42": 1096260608.0,
+            "43": 1096339968.0,
+            "44": 1096332800.0,
+            "45": 1096305152.0,
+            "46": 1096355840.0,
+            "47": 1096307712.0,
+            "48": 1096266240.0,
+            "49": 1096279040.0,
+            "50": 1096301056.0,
+            "51": 1096302592.0,
+            "52": 1096305664.0,
+            "53": 1096240640.0,
+            "54": 1096317440.0,
+            "55": 1096325120.0,
+            "56": 1096299520.0,
+            "57": 1096352256.0,
+            "58": 1096334848.0,
+            "59": 1096264192.0,
+            "60": 1096317440.0,
+            "61": 1096312832.0,
+            "62": 1096285696.0,
+            "63": 1096279040.0,
+            "64": 1096306176.0,
+            "65": 1096305664.0,
+            "66": 1096304640.0,
+            "67": 1096308736.0,
+            "68": 1096334336.0,
+            "69": 1096293376.0,
+            "70": 1096297472.0,
+            "71": 1096308736.0,
+            "72": 1096337408.0,
+            "73": 1096290304.0,
+            "74": 1096303616.0,
+            "75": 1096306688.0,
+            "76": 1096308736.0,
+            "77": 1096314880.0,
+            "78": 1096329728.0,
+            "79": 1096388608.0,
+            "80": 1096315392.0,
+            "81": 1096268800.0,
+            "82": 1096268288.0,
+            "83": 1096327680.0,
+            "84": 1096331264.0,
+            "85": 1096320000.0,
+            "86": 1096287744.0,
+            "87": 1096280576.0,
+            "88": 1096303616.0,
+            "89": 1096294912.0,
+            "90": 1096305664.0,
+            "91": 1096297472.0,
+            "92": 1096288768.0,
+            "93": 1096294400.0,
+            "94": 1096279552.0,
+            "95": 1096308736.0,
+            "96": 1096294912.0,
+            "97": 1096288768.0,
+            "98": 1096327680.0,
+            "99": 1096341504.0,
+            "100": 1096311296.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,212 +325,213 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 3257109504.0,
-            "2": 3590822912.0,
-            "3": 3590822912.0,
-            "4": 3590822912.0,
-            "5": 3596996096.0,
-            "6": 3596996096.0,
-            "7": 3596996096.0,
-            "8": 3596996096.0,
-            "9": 3596996096.0,
-            "10": 3596996096.0,
-            "11": 3596996096.0,
-            "12": 3596996096.0,
-            "13": 3596996096.0,
-            "14": 3596996096.0,
-            "15": 3596996096.0,
-            "16": 3596996096.0,
-            "17": 3596996096.0,
-            "18": 3596996096.0,
-            "19": 3596996096.0,
-            "20": 3596996096.0,
-            "21": 3596996096.0,
-            "22": 3610441728.0,
-            "23": 3610441728.0,
-            "24": 3610441728.0,
-            "25": 3610441728.0,
-            "26": 3611658240.0,
-            "27": 3611658240.0,
-            "28": 3611658240.0,
-            "29": 3611658240.0,
-            "30": 3611658240.0,
-            "31": 3611658240.0,
-            "32": 3611658240.0,
-            "33": 3611658240.0,
-            "34": 3611658240.0,
-            "35": 3611658240.0,
-            "36": 3611658240.0,
-            "37": 3611658240.0,
-            "38": 3611658240.0,
-            "39": 3611658240.0,
-            "40": 3611658240.0,
-            "41": 3611658240.0,
-            "42": 3611658240.0,
-            "43": 3611658240.0,
-            "44": 3611658240.0,
-            "45": 3611658240.0,
-            "46": 3611658240.0,
-            "47": 3611658240.0,
-            "48": 3611658240.0,
-            "49": 3611658240.0,
-            "50": 3611658240.0,
-            "51": 3611658240.0,
-            "52": 3611658240.0,
-            "53": 3611658240.0,
-            "54": 3611658240.0,
-            "55": 3611658240.0,
-            "56": 3611658240.0,
-            "57": 3625479168.0,
-            "58": 3625479168.0,
-            "59": 3625479168.0,
-            "60": 3625479168.0,
-            "61": 3625479168.0,
-            "62": 3625479168.0,
-            "63": 3625479168.0,
-            "64": 3625479168.0,
-            "65": 3625479168.0,
-            "66": 3625479168.0,
-            "67": 3625479168.0,
-            "68": 3625479168.0,
-            "69": 3625479168.0,
-            "70": 3625479168.0,
-            "71": 3625479168.0,
-            "72": 3625479168.0,
-            "73": 3625479168.0,
-            "74": 3625479168.0,
-            "75": 3625479168.0,
-            "76": 3625479168.0,
-            "77": 3625479168.0,
-            "78": 3625479168.0,
-            "79": 3625479168.0,
-            "80": 3625479168.0,
-            "81": 3625479168.0,
-            "82": 3625479168.0,
-            "83": 3625479168.0,
-            "84": 3625479168.0,
-            "85": 3625479168.0,
-            "86": 3625479168.0,
-            "87": 3625479168.0,
-            "88": 3625479168.0,
-            "89": 3625479168.0,
-            "90": 3625479168.0,
-            "91": 3625479168.0,
-            "92": 3625479168.0,
-            "93": 3625479168.0,
-            "94": 3625479168.0,
-            "95": 3625479168.0,
-            "96": 3625479168.0,
-            "97": 3625479168.0,
-            "98": 3625479168.0,
-            "99": 3625479168.0,
-            "100": 3625479168.0
+            "1": 3241482240.0,
+            "2": 3627405824.0,
+            "3": 3627405824.0,
+            "4": 3627405824.0,
+            "5": 3627405824.0,
+            "6": 3627405824.0,
+            "7": 3627405824.0,
+            "8": 3627405824.0,
+            "9": 3627405824.0,
+            "10": 3627405824.0,
+            "11": 3627405824.0,
+            "12": 3627405824.0,
+            "13": 3627405824.0,
+            "14": 3627405824.0,
+            "15": 3627405824.0,
+            "16": 3627405824.0,
+            "17": 3627405824.0,
+            "18": 3627405824.0,
+            "19": 3627405824.0,
+            "20": 3627405824.0,
+            "21": 3627405824.0,
+            "22": 3627405824.0,
+            "23": 3627405824.0,
+            "24": 3627405824.0,
+            "25": 3627405824.0,
+            "26": 3627405824.0,
+            "27": 3627405824.0,
+            "28": 3627405824.0,
+            "29": 3627405824.0,
+            "30": 3627405824.0,
+            "31": 3627405824.0,
+            "32": 3627405824.0,
+            "33": 3627405824.0,
+            "34": 3627405824.0,
+            "35": 3627405824.0,
+            "36": 3627405824.0,
+            "37": 3627405824.0,
+            "38": 3627405824.0,
+            "39": 3627405824.0,
+            "40": 3627405824.0,
+            "41": 3627405824.0,
+            "42": 3627405824.0,
+            "43": 3627405824.0,
+            "44": 3627405824.0,
+            "45": 3627405824.0,
+            "46": 3627405824.0,
+            "47": 3627405824.0,
+            "48": 3627405824.0,
+            "49": 3627405824.0,
+            "50": 3627405824.0,
+            "51": 3627405824.0,
+            "52": 3627405824.0,
+            "53": 3627405824.0,
+            "54": 3627405824.0,
+            "55": 3627405824.0,
+            "56": 3627405824.0,
+            "57": 3627405824.0,
+            "58": 3627405824.0,
+            "59": 3627405824.0,
+            "60": 3627405824.0,
+            "61": 3627405824.0,
+            "62": 3627405824.0,
+            "63": 3627405824.0,
+            "64": 3627405824.0,
+            "65": 3627405824.0,
+            "66": 3627405824.0,
+            "67": 3627405824.0,
+            "68": 3627405824.0,
+            "69": 3627405824.0,
+            "70": 3627405824.0,
+            "71": 3627405824.0,
+            "72": 3627405824.0,
+            "73": 3627405824.0,
+            "74": 3627405824.0,
+            "75": 3627405824.0,
+            "76": 3627405824.0,
+            "77": 3627405824.0,
+            "78": 3627405824.0,
+            "79": 3630189056.0,
+            "80": 3630189056.0,
+            "81": 3630189056.0,
+            "82": 3630189056.0,
+            "83": 3630189056.0,
+            "84": 3630189056.0,
+            "85": 3630189056.0,
+            "86": 3630189056.0,
+            "87": 3630189056.0,
+            "88": 3630189056.0,
+            "89": 3630189056.0,
+            "90": 3630189056.0,
+            "91": 3630189056.0,
+            "92": 3630189056.0,
+            "93": 3630189056.0,
+            "94": 3630189056.0,
+            "95": 3630189056.0,
+            "96": 3630189056.0,
+            "97": 3630189056.0,
+            "98": 3630189056.0,
+            "99": 3630189056.0,
+            "100": 3630189056.0
         }
     },
     "iteration-time": {
-        "start_step": 2,
+        "start_step": 1,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "2": 7.98917,
-            "3": 0.61774,
-            "4": 0.60327,
-            "5": 0.61118,
-            "6": 0.60341,
-            "7": 0.60198,
-            "8": 0.60804,
-            "9": 0.59243,
-            "10": 0.60275,
-            "11": 0.61409,
-            "12": 0.59478,
-            "13": 0.59416,
-            "14": 0.61157,
-            "15": 0.59528,
-            "16": 0.58201,
-            "17": 0.59897,
-            "18": 0.5747,
-            "19": 0.5852,
-            "20": 0.58745,
-            "21": 0.58032,
-            "22": 0.5811,
-            "23": 0.57991,
-            "24": 0.58907,
-            "25": 0.58791,
-            "26": 0.59379,
-            "27": 0.58561,
-            "28": 0.58974,
-            "29": 0.5884,
-            "30": 0.60742,
-            "31": 0.58261,
-            "32": 0.58956,
-            "33": 0.58555,
-            "34": 0.58583,
-            "35": 0.59496,
-            "36": 0.6018,
-            "37": 0.59609,
-            "38": 0.58356,
-            "39": 0.57802,
-            "40": 0.57925,
-            "41": 0.57195,
-            "42": 0.57595,
-            "43": 0.57774,
-            "44": 0.56745,
-            "45": 0.57916,
-            "46": 0.56772,
-            "47": 0.57299,
-            "48": 0.56576,
-            "49": 0.56807,
-            "50": 0.56685,
-            "51": 0.63208,
-            "52": 0.59958,
-            "53": 0.57872,
-            "54": 0.56971,
-            "55": 0.57134,
-            "56": 0.56955,
-            "57": 0.57154,
-            "58": 0.57832,
-            "59": 0.57163,
-            "60": 0.57009,
-            "61": 0.57588,
-            "62": 0.58149,
-            "63": 0.57169,
-            "64": 0.56388,
-            "65": 0.56972,
-            "66": 0.56979,
-            "67": 0.56465,
-            "68": 0.56683,
-            "69": 0.57178,
-            "70": 0.57356,
-            "71": 0.5728,
-            "72": 0.57218,
-            "73": 0.58671,
-            "74": 0.57786,
-            "75": 0.57202,
-            "76": 0.56452,
-            "77": 0.57511,
-            "78": 0.57141,
-            "79": 0.56707,
-            "80": 0.56945,
-            "81": 0.57006,
-            "82": 0.58157,
-            "83": 0.5684,
-            "84": 0.5733,
-            "85": 0.5679,
-            "86": 0.57571,
-            "87": 0.57676,
-            "88": 0.57866,
-            "89": 0.57356,
-            "90": 0.57899,
-            "91": 0.57243,
-            "92": 0.6007,
-            "93": 0.59639,
-            "94": 0.58713,
-            "95": 0.59701,
-            "96": 0.60178,
-            "97": 0.57955,
-            "98": 0.5712,
-            "99": 0.57566,
-            "100": 0.58236
+            "1": "nan",
+            "2": 5.82645,
+            "3": 0.56935,
+            "4": 0.53226,
+            "5": 0.5234,
+            "6": 0.51861,
+            "7": 0.52442,
+            "8": 1.1669,
+            "9": 0.61585,
+            "10": 1.28981,
+            "11": 0.50078,
+            "12": 0.51545,
+            "13": 0.51278,
+            "14": 0.50805,
+            "15": 0.50172,
+            "16": 0.50244,
+            "17": 0.49784,
+            "18": 0.61979,
+            "19": 0.56119,
+            "20": 0.51081,
+            "21": 0.50024,
+            "22": 1.5289,
+            "23": 0.50151,
+            "24": 0.50503,
+            "25": 0.48873,
+            "26": 0.49486,
+            "27": 0.49958,
+            "28": 0.50892,
+            "29": 0.49877,
+            "30": 0.49887,
+            "31": 0.49276,
+            "32": 0.74532,
+            "33": 0.49725,
+            "34": 0.49,
+            "35": 0.49566,
+            "36": 0.56904,
+            "37": 0.50053,
+            "38": 0.48543,
+            "39": 0.4915,
+            "40": 0.48911,
+            "41": 0.49237,
+            "42": 0.50502,
+            "43": 0.6971,
+            "44": 0.48466,
+            "45": 0.48657,
+            "46": 0.48912,
+            "47": 0.49872,
+            "48": 0.49649,
+            "49": 0.50023,
+            "50": 0.49167,
+            "51": 0.54286,
+            "52": 0.537,
+            "53": 0.50613,
+            "54": 0.49945,
+            "55": 0.48685,
+            "56": 0.4947,
+            "57": 0.49311,
+            "58": 0.4867,
+            "59": 0.55079,
+            "60": 0.48821,
+            "61": 0.88185,
+            "62": 0.49817,
+            "63": 0.49948,
+            "64": 0.49443,
+            "65": 0.49186,
+            "66": 0.68473,
+            "67": 0.49826,
+            "68": 0.48954,
+            "69": 0.52156,
+            "70": 0.49207,
+            "71": 0.49505,
+            "72": 0.49817,
+            "73": 0.50732,
+            "74": 0.4906,
+            "75": 0.49748,
+            "76": 0.48422,
+            "77": 0.48902,
+            "78": 0.48466,
+            "79": 0.49314,
+            "80": 0.48808,
+            "81": 0.49806,
+            "82": 0.49986,
+            "83": 0.49598,
+            "84": 0.49257,
+            "85": 0.49093,
+            "86": 0.49165,
+            "87": 0.49788,
+            "88": 0.50477,
+            "89": 0.50357,
+            "90": 0.48442,
+            "91": 0.49447,
+            "92": 0.50294,
+            "93": 0.50096,
+            "94": 0.49868,
+            "95": 0.49472,
+            "96": 0.51062,
+            "97": 0.49231,
+            "98": 0.716,
+            "99": 0.49303,
+            "100": 0.50087
         }
     }
 }
\ No newline at end of file
diff --git a/tests/unit_tests/test_initialize_jit_warmup.py b/tests/unit_tests/test_initialize_jit_warmup.py
new file mode 100644
index 00000000000..826f59950ee
--- /dev/null
+++ b/tests/unit_tests/test_initialize_jit_warmup.py
@@ -0,0 +1,141 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+"""Tests that JIT warmup skips fused kernels when disabled via arguments.
+
+Uses mocks only (no CUDA required). Patches use
+``megatron.training.initialize.<name>`` because ``_warmup_jit_function`` closes
+over that module's globals (imported ``get_args``, ``torch``, fusion ops).
+"""
+
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from megatron.training.initialize import _warmup_jit_function
+
+
+def _args(**overrides):
+    base = dict(
+        bf16=False,
+        fp16=False,
+        use_te_activation_func=False,
+        gated_linear_unit=False,
+        quick_geglu=False,
+        swiglu=False,
+        bias_swiglu_fusion=True,
+        bias_gelu_fusion=True,
+        bias_dropout_fusion=False,
+        sequence_parallel=False,
+        ffn_hidden_size=64,
+        tensor_model_parallel_size=1,
+        seq_length=8,
+        context_parallel_size=1,
+        micro_batch_size=1,
+        hidden_size=32,
+    )
+    base.update(overrides)
+    return SimpleNamespace(**base)
+
+
+def _fake_rand(*size, **kwargs):
+    """Mirror ``torch.rand`` shapes used in ``_warmup_jit_function`` but on CPU."""
+    dtype = kwargs.get("dtype", torch.float32)
+    if len(size) == 1 and isinstance(size[0], tuple):
+        shape = size[0]
+    elif size:
+        shape = size if len(size) > 1 else (size[0],)
+    else:
+        shape = ()
+    return torch.zeros(shape, dtype=dtype, device="cpu")
+
+
+@mock.patch("megatron.training.initialize.torch.cuda.empty_cache")
+@mock.patch("megatron.training.initialize.get_args")
+class TestJitWarmupSkippedWhenFusionDisabled:
+    @mock.patch("megatron.training.initialize.bias_swiglu")
+    def test_skips_bias_swiglu(self, mock_swiglu, mock_get_args, _empty_cache):
+        mock_get_args.return_value = _args(swiglu=True, bias_swiglu_fusion=False)
+        _warmup_jit_function()
+        mock_swiglu.assert_not_called()
+
+    @mock.patch("megatron.training.initialize.bias_gelu")
+    def test_skips_bias_gelu(self, mock_gelu, mock_get_args, _empty_cache):
+        mock_get_args.return_value = _args(
+            swiglu=False, bias_gelu_fusion=False, gated_linear_unit=False
+        )
+        _warmup_jit_function()
+        mock_gelu.assert_not_called()
+
+    @mock.patch("megatron.training.initialize.bias_geglu")
+    def test_skips_bias_geglu(self, mock_geglu, mock_get_args, _empty_cache):
+        mock_get_args.return_value = _args(
+            swiglu=False, bias_gelu_fusion=False, gated_linear_unit=True
+        )
+        _warmup_jit_function()
+        mock_geglu.assert_not_called()
+
+    @mock.patch("megatron.training.initialize.bias_geglu")
+    @mock.patch("megatron.training.initialize.bias_gelu")
+    def test_skips_gelu_and_geglu_when_quick_geglu(
+        self, mock_gelu, mock_geglu, mock_get_args, _empty_cache
+    ):
+        # Even with bias_gelu_fusion True, quick_geglu skips gelu and geglu warmups.
+        mock_get_args.return_value = _args(
+            swiglu=False, quick_geglu=True, bias_gelu_fusion=True, gated_linear_unit=True
+        )
+        _warmup_jit_function()
+        mock_gelu.assert_not_called()
+        mock_geglu.assert_not_called()
+
+    @mock.patch("megatron.training.initialize.bias_dropout_add_fused_train")
+    def test_skips_bias_dropout_fusion(self, mock_dropout, mock_get_args, _empty_cache):
+        # With swiglu=True, gelu/geglu warmups are off; disable swiglu fusion so no CUDA path runs.
+        mock_get_args.return_value = _args(
+            swiglu=True, bias_swiglu_fusion=False, bias_dropout_fusion=False
+        )
+        _warmup_jit_function()
+        mock_dropout.assert_not_called()
+
+    @mock.patch("megatron.training.initialize.bias_dropout_add_fused_train")
+    @mock.patch("megatron.training.initialize.bias_geglu")
+    @mock.patch("megatron.training.initialize.bias_gelu")
+    @mock.patch("megatron.training.initialize.bias_swiglu")
+    def test_te_activation_skips_activation_warmups(
+        self, mock_swiglu, mock_gelu, mock_geglu, mock_dropout, mock_get_args, _empty_cache
+    ):
+        # bias_dropout_fusion is independent of use_te_activation_func; keep False
+        # so this test stays CPU-only without mocking the dropout path + mpu.
+        mock_get_args.return_value = _args(
+            use_te_activation_func=True,
+            swiglu=True,
+            bias_swiglu_fusion=True,
+            gated_linear_unit=True,
+            bias_gelu_fusion=True,
+            bias_dropout_fusion=False,
+        )
+        _warmup_jit_function()
+        mock_swiglu.assert_not_called()
+        mock_gelu.assert_not_called()
+        mock_geglu.assert_not_called()
+        mock_dropout.assert_not_called()
+
+
+@mock.patch("megatron.training.initialize.bias_dropout_add_fused_train")
+@mock.patch("megatron.training.initialize.bias_geglu")
+@mock.patch("megatron.training.initialize.bias_gelu")
+@mock.patch("megatron.training.initialize.torch.cuda.empty_cache")
+@mock.patch("megatron.training.initialize.torch.rand", side_effect=_fake_rand)
+@mock.patch("megatron.training.initialize.bias_swiglu")
+@mock.patch("megatron.training.initialize.get_args")
+def test_calls_bias_swiglu_when_fusion_enabled(
+    mock_get_args, mock_swiglu, mock_torch_rand, _empty_cache, mock_gelu, mock_geglu, mock_dropout
+):
+    mock_get_args.return_value = _args(swiglu=True, bias_swiglu_fusion=True)
+    mock_swiglu.return_value = torch.tensor(0.0)
+    _warmup_jit_function()
+    assert mock_swiglu.call_count == 10  # zip([True,True],[False,True]) × range(5)
+    mock_gelu.assert_not_called()
+    mock_geglu.assert_not_called()
+    mock_dropout.assert_not_called()
+    assert mock_torch_rand.call_count == 2