diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index e2e60adfadf..15247c2be53 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -887,8 +887,9 @@ def group_params_for_buffers(
     Each distinct buffer is identified by a BufferKey with three dimensions:
     - param_dtype: storage dtype (torch.uint8 for FP8/NVFP4 parameters, else param.dtype).
     - grad_dtype: gradient reduction dtype (torch.float if grad_reduce_in_fp32, else param.dtype).
-    - is_expert_parallel: whether the parameter is expert-parallel (param.allreduce == False),
-      which requires a separate buffer with a different data-parallel group.
+    - is_expert_parallel: whether the parameter uses the expert topology (param.allreduce == False),
+      which requires a separate buffer for the expert data-parallel group.  This is true for experts
+      when expert-parallelism > 1 or expert-tensor-parallelism != tensor-parallelism.
 
     The param_indices track each parameter's position among same-dtype params (using
     the "fake" high-precision dtype for FP8/NVFP4 params), needed for loading non-native-fp8
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 51a9457883a..8f0117c68ec 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -8,6 +8,7 @@
 import io
 import os
 import pickle
+import re
 import warnings
 from contextlib import contextmanager, nullcontext
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, cast
@@ -84,6 +85,42 @@
         HAVE_TE = False
 
 _TE_CONFIG_TYPE_KEY = "transformer_engine_config_type"
+_EXPERT_PARAMETER_NAME_PATTERN = re.compile(r"(weight|bias)\d*")
+
+
+def _set_expert_parameter_attributes(
+    module: torch.nn.Module, parallel_mode: Optional[str], use_expert_pgs: bool
+) -> None:
+    """Set process-group and tensor-partition metadata on an expert TE module.
+
+    ``allreduce=False`` selects EDP for gradient reduction.
+
+    Weights and biases, including TEGroupedLinear's numbered parameters, are also marked as
+    TP-partitioned according to ``parallel_mode``; row-parallel biases remain replicated.
+
+    Any parameter which is partitioned along TP or ETP is marked with ``tensor_model_parallel``,
+    which ensures that all shards contribute to the gradient norm.
+
+    Args:
+        module: Transformer Engine module whose direct parameters should be marked.
+        parallel_mode: Tensor-parallel mode used by the module (``"column"``, ``"row"``, or None).
+        use_expert_pgs: Whether to use EP/ETP/EDP process groups instead of TP/CP/DP.
+    """
+    for name, param in module.named_parameters(recurse=False):
+        param.allreduce = not use_expert_pgs
+
+        name_match = _EXPERT_PARAMETER_NAME_PATTERN.fullmatch(name)
+        parameter_kind = name_match.group(1) if name_match else None
+        is_weight = parameter_kind == "weight"
+        is_bias = parameter_kind == "bias"
+        is_partitioned = parallel_mode in ("column", "row") and (
+            is_weight or (parallel_mode == "column" and is_bias)
+        )
+        if is_weight or is_bias:
+            param.tensor_model_parallel = is_partitioned
+        if is_partitioned:
+            param.partition_dim = 1 if parallel_mode == "row" else 0
+            param.partition_stride = 1
 
 
 class TransformerEngineConfigType(enum.Enum):
@@ -942,6 +979,10 @@ def __init__(
             tp_size = get_pg_size(tp_group)
 
         self.expert_parallel = self.config.expert_model_parallel_size > 1
+        use_expert_pgs = is_expert and (
+            self.expert_parallel
+            or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+        )
         if is_expert:
             rng_tracker_name = get_expert_parallel_rng_tracker_name()
         else:
@@ -1007,11 +1048,10 @@ def __init__(
                 **extra_kwargs,
             )
 
-        for param in self.parameters():
-            if is_expert:
-                # Reduce the gradient on the expert_data_parallel group for expert linear layers
-                setattr(param, "allreduce", not self.expert_parallel)
-            else:
+        if is_expert:
+            _set_expert_parameter_attributes(self, parallel_mode, use_expert_pgs)
+        else:
+            for param in self.parameters():
                 # Reduce the gradient on DP group
                 setattr(param, "allreduce", True)
                 if parallel_mode == "duplicated":
@@ -1431,6 +1471,13 @@ def __init__(
                     self.bias.zero_()
                 setattr(self.bias, "allreduce", True)
 
+        if is_expert:
+            use_expert_pgs = (
+                config.expert_model_parallel_size > 1
+                or config.expert_tensor_parallel_size != config.tensor_model_parallel_size
+            )
+            _set_expert_parameter_attributes(self, "column", use_expert_pgs)
+
     def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
         """Sharding along axis 0, bias sharded"""
         state_dict = self.state_dict(prefix="", keep_vars=True)
@@ -1680,6 +1727,13 @@ def __init__(
                 setattr(self.bias, "allreduce", True)
                 setattr(self.bias, "sequence_parallel", config.sequence_parallel)
 
+        if is_expert:
+            use_expert_pgs = (
+                config.expert_model_parallel_size > 1
+                or config.expert_tensor_parallel_size != config.tensor_model_parallel_size
+            )
+            _set_expert_parameter_attributes(self, "row", use_expert_pgs)
+
     def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
         """Sharding along axis 1, bias not sharded"""
         state_dict = self.state_dict(prefix="", keep_vars=True)
@@ -2106,6 +2160,10 @@ def __init__(
             extra_kwargs["ub_name"] = tp_comm_buffer_name
 
             self.expert_parallel = self.config.expert_model_parallel_size > 1
+            use_expert_pgs = is_expert and (
+                self.expert_parallel
+                or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+            )
             if is_expert:
                 extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name()
 
@@ -2182,23 +2240,7 @@ def __init__(
                     **extra_kwargs,
                 )
 
-            for param in self.parameters():
-                setattr(param, "allreduce", not (is_expert and self.expert_parallel))
-
-            # Explicitly stamp partition_dim and partition_stride on expert weight
-            # tensors when explicit_expert_comm cleared parallel_mode.  TE ≤2.12
-            # set these internally; TE ≥2.13 no longer does (parallel_mode=None
-            # is passed due to explicit_expert_comm).  The resharding/refit planner
-            # relies on partition_dim to correctly plan TP gather/scatter operations.
-            # NOTE: we intentionally do NOT stamp tensor_model_parallel here —
-            # doing so would change num-zeros gradient counting.
-            if self.explicit_expert_comm and original_parallel_mode in ("column", "row"):
-                part_dim = 0 if original_parallel_mode == "column" else 1
-                for i in range(num_gemms):
-                    weight = getattr(self, f"weight{i}", None)
-                    if weight is not None:
-                        setattr(weight, "partition_dim", part_dim)
-                        setattr(weight, "partition_stride", 1)
+            _set_expert_parameter_attributes(self, original_parallel_mode, use_expert_pgs)
 
             self._register_load_state_dict_pre_hook(
                 type(self)._normalize_grouped_parameter_keys, with_module=True
diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index b7ea0bbe069..70f757f2889 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -47,7 +47,6 @@
 if HAVE_EMERGING_OPTIMIZERS:
     from emerging_optimizers.scalar_optimizers import Lion
 
-from megatron.core import parallel_state
 from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer
 from megatron.core.optimizer_param_scheduler import (
     ParamGroupOverride,
@@ -686,11 +685,12 @@ def init_state_fn(opt, config=None):
         setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group)
 
     if pg_collection is None or not hasattr(pg_collection, 'tp'):
-        tp_group = parallel_state.get_tensor_model_parallel_group()
-    else:
-        tp_group = pg_collection.tp
-    # TODO(M4): plumb tp_group through optimizer constructors so this setattr disappears.
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+    tp_group = pg_collection.tp
+    expert_tp_group = getattr(pg_collection, 'expt_tp', tp_group)
+    # TODO(M4): plumb TP groups through optimizer constructors so these setattrs disappear.
     setattr(optimizer, 'tp_group', tp_group)
+    setattr(optimizer, 'expert_tp_group', expert_tp_group)
 
     return optimizer
 
@@ -898,11 +898,10 @@ def _get_megatron_emerging_optimizer(
             else:
                 optimizer = FP32Optimizer(optimizer, config, init_state_fn)
             setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group)
-            if pg_collection is None or not hasattr(pg_collection, 'tp'):
-                tp_group = parallel_state.get_tensor_model_parallel_group()
-            else:
-                tp_group = pg_collection.tp
+            tp_group = pg_collection.tp
+            expert_tp_group = getattr(pg_collection, 'expt_tp', tp_group)
             setattr(optimizer, 'tp_group', tp_group)
+            setattr(optimizer, 'expert_tp_group', expert_tp_group)
             results.append(optimizer)
             continue
         else:
diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py
index 55848e104ae..ad2064ede36 100644
--- a/megatron/core/optimizer/clip_grads.py
+++ b/megatron/core/optimizer/clip_grads.py
@@ -197,6 +197,7 @@ def count_zeros_fp32(
     grad_stats_parallel_group: torch.distributed.ProcessGroup,
     use_decoupled_grad: bool = False,
     tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    expert_tp_group: Optional[torch.distributed.ProcessGroup] = None,
 ) -> float:
     """Counts the number of zero values in the gradients of the given parameters.
 
@@ -242,7 +243,9 @@ def count_zeros_fp32(
             total_num_zeros += num_zeros
             continue
         is_not_shared = param_is_not_shared(param)
-        is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param, tp_group=tp_group)
+        is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(
+            param, tp_group=tp_group, expert_tp_group=expert_tp_group
+        )
         is_not_gtp_duplicate = param_is_not_gtp_duplicate(param)
         if grad_not_none and is_not_shared and is_not_tp_duplicate and is_not_gtp_duplicate:
             grad_obj = getattr(param, grad_attr)
diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py
index 77cd20de10e..dac16f4a2ee 100644
--- a/megatron/core/optimizer/optimizer.py
+++ b/megatron/core/optimizer/optimizer.py
@@ -176,6 +176,8 @@ def _is_separate_grad_norm_group(grad_norm_group: Optional[str]) -> bool:
 
 def copy_optimizer_param_metadata(destination: torch.Tensor, source: torch.Tensor) -> None:
     """Copy optimizer-relevant metadata when creating param views/copies."""
+    if hasattr(source, 'allreduce'):
+        destination.allreduce = source.allreduce
     if hasattr(source, 'shared'):
         destination.shared = source.shared
     if hasattr(source, GRAD_NORM_GROUP_ATTR):
@@ -266,7 +268,9 @@ def _filter_grads_for_norm(
             grad_not_none = grad is not None
             is_not_shared = param_is_not_shared(param)
             is_not_tp_duplicate = tensor_parallel.param_is_not_tensor_parallel_duplicate(
-                param, getattr(self, 'tp_group', None)
+                param,
+                tp_group=getattr(self, 'tp_group', None),
+                expert_tp_group=getattr(self, 'expert_tp_group', None),
             )
             is_not_gtp_duplicate = tensor_parallel.param_is_not_gtp_duplicate(param)
             if grad_not_none and is_not_shared and is_not_tp_duplicate and is_not_gtp_duplicate:
@@ -434,6 +438,7 @@ def count_zeros(self) -> float:
                 and getattr(params[0], "__fsdp_param__", False)
             ),
             tp_group=getattr(self, 'tp_group', None),
+            expert_tp_group=getattr(self, 'expert_tp_group', None),
         )
 
     @abstractmethod
@@ -1785,6 +1790,8 @@ def count_zeros(self):
                     self.config.use_precision_aware_optimizer
                     and getattr(params[0], "__fsdp_param__", False)
                 ),
+                tp_group=getattr(self.chained_optimizers[0], 'tp_group', None),
+                expert_tp_group=getattr(self.chained_optimizers[0], 'expert_tp_group', None),
             )
         else:
             num_zeros_in_grad = 0
diff --git a/megatron/core/post_training/modelopt/layers.py b/megatron/core/post_training/modelopt/layers.py
index 04e03a36458..5f1a746e95b 100644
--- a/megatron/core/post_training/modelopt/layers.py
+++ b/megatron/core/post_training/modelopt/layers.py
@@ -159,7 +159,12 @@ def __init__(
         for param in self.parameters():
             if is_expert:
                 # Reduce the gradient on the expert_data_parallel group for expert linear layers
-                setattr(param, "allreduce", self.config.expert_model_parallel_size == 1)
+                use_expert_groups = (
+                    self.config.expert_model_parallel_size > 1
+                    or self.config.expert_tensor_parallel_size
+                    != self.config.tensor_model_parallel_size
+                )
+                setattr(param, "allreduce", not use_expert_groups)
             else:
                 # Reduce the gradient on DP group
                 setattr(param, "allreduce", True)
diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py
index e6e55a96a75..3ff635b362b 100644
--- a/megatron/core/tensor_parallel/layers.py
+++ b/megatron/core/tensor_parallel/layers.py
@@ -92,11 +92,17 @@
     dist_reduce_scatter_func = torch.distributed._reduce_scatter_base
 
 
-def param_is_not_tensor_parallel_duplicate(param, tp_group=None):
-    """Returns true if the passed-in parameter is not a duplicate parameter
-    on another TP rank."""
+def param_is_not_tensor_parallel_duplicate(param, tp_group=None, expert_tp_group=None):
+    """Return whether a parameter contributes to a unique model-parallel shard.
+
+    Parameters reduced over expert data parallel groups use the expert tensor-parallel
+    group for duplicate filtering. Other parameters use the regular tensor-parallel group.
+    """
     if hasattr(param, "tensor_model_parallel") and param.tensor_model_parallel:
         return True
+    # allreduce=False marks parameters reduced over expert DP, so filter their duplicates over ETP.
+    if not getattr(param, "allreduce", True) and expert_tp_group is not None:
+        tp_group = expert_tp_group
     # Prefer provided tp_group when available (new explicit path).
     if tp_group is not None:
         return tp_group.rank() == 0
@@ -952,6 +958,10 @@ def __init__(
         world_size = get_pg_size(self.tp_group)
         rank = get_pg_rank(self.tp_group)
         self.explicit_expert_comm = self.is_expert and (world_size > 1 or self.expert_parallel)
+        use_expert_pgs = self.is_expert and (
+            self.expert_parallel
+            or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+        )
         self.output_size_per_partition = divide(output_size, world_size)
 
         # Parameters.
@@ -1004,7 +1014,7 @@ def __init__(
                         tensor=self.weight, is_parallel=True, dim=0, stride=stride
                     )
 
-            setattr(self.weight, "allreduce", not (self.is_expert and self.expert_parallel))
+            setattr(self.weight, "allreduce", not use_expert_pgs)
         else:
             self.weight = None
 
@@ -1037,7 +1047,7 @@ def __init__(
                 # Always initialize bias to zero.
                 with torch.no_grad():
                     self.bias.zero_()
-            setattr(self.bias, "allreduce", not (self.is_expert and self.expert_parallel))
+            setattr(self.bias, "allreduce", not use_expert_pgs)
         else:
             self.register_parameter("bias", None)
 
@@ -1366,7 +1376,11 @@ def __init__(
                 set_tensor_model_parallel_attributes(
                     tensor=self.weight, is_parallel=True, dim=1, stride=stride
                 )
-        setattr(self.weight, "allreduce", not (self.is_expert and self.expert_parallel))
+        use_expert_pgs = self.is_expert and (
+            self.expert_parallel
+            or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+        )
+        setattr(self.weight, "allreduce", not use_expert_pgs)
 
         self.gtp_remat_size = 1
         _pg = ProcessGroupCollection.use_mpu_process_groups(
@@ -1395,7 +1409,7 @@ def __init__(
                 # Always initialize bias to zero.
                 with torch.no_grad():
                     self.bias.zero_()
-            setattr(self.bias, "allreduce", not (self.is_expert and self.expert_parallel))
+            setattr(self.bias, "allreduce", not use_expert_pgs)
             setattr(self.bias, "sequence_parallel", self.sequence_parallel)
         else:
             self.register_parameter("bias", None)
diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py
index 4e7d55a4577..30617ef9b4c 100644
--- a/megatron/training/utils/common_utils.py
+++ b/megatron/training/utils/common_utils.py
@@ -112,6 +112,8 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
 
     gtp_rank = mpu.get_gtp_weight_remat_rank()
     egtp_rank = mpu.get_expert_gtp_weight_remat_rank()
+    tp_group = mpu.get_tensor_model_parallel_group()
+    expert_tp_group = mpu.get_expert_tensor_parallel_group()
 
     for model_chunk in model:
         for param in model_chunk.parameters():
@@ -119,7 +121,9 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
 
             # Filter TP duplicates. GTP_remat params are always unique across TP ranks
             # so skip this check for them.
-            if not is_gtp and not param_is_not_tensor_parallel_duplicate(param):
+            if not is_gtp and not param_is_not_tensor_parallel_duplicate(
+                param, tp_group=tp_group, expert_tp_group=expert_tp_group
+            ):
                 continue
             is_expert = not getattr(param, 'allreduce', True)
 
diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_h100.json
index 8a7452aa143..3eefe1b2bb3 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective/golden_values_dev_dgx_h100.json
@@ -6,54 +6,54 @@
         "values": {
             "1": 10.92563,
             "2": 10.91638,
-            "3": 10.92433,
-            "4": 10.93217,
-            "5": 10.92999,
-            "6": 10.92662,
-            "7": 10.92571,
-            "8": 10.92333,
-            "9": 10.92825,
-            "10": 10.91605,
-            "11": 10.91854,
-            "12": 10.92399,
-            "13": 10.91037,
-            "14": 10.90685,
-            "15": 10.90136,
-            "16": 10.88661,
-            "17": 10.88849,
-            "18": 10.88662,
-            "19": 10.8857,
-            "20": 10.83765,
-            "21": 10.82761,
-            "22": 10.81538,
-            "23": 10.8078,
-            "24": 10.78018,
-            "25": 10.778,
-            "26": 10.76109,
-            "27": 10.74912,
-            "28": 10.69195,
-            "29": 10.66617,
-            "30": 10.63122,
-            "31": 10.62222,
-            "32": 10.61543,
-            "33": 10.57901,
-            "34": 10.54677,
-            "35": 10.54608,
-            "36": 10.53419,
-            "37": 10.50598,
-            "38": 10.50458,
-            "39": 10.47274,
-            "40": 10.45062,
-            "41": 10.42727,
-            "42": 10.41444,
-            "43": 10.40126,
-            "44": 10.3705,
-            "45": 10.38167,
-            "46": 10.33539,
-            "47": 10.32458,
-            "48": 10.28718,
-            "49": 10.28599,
-            "50": 10.27739
+            "3": 10.92459,
+            "4": 10.93245,
+            "5": 10.93019,
+            "6": 10.92651,
+            "7": 10.92566,
+            "8": 10.92291,
+            "9": 10.92842,
+            "10": 10.91709,
+            "11": 10.9186,
+            "12": 10.92364,
+            "13": 10.91032,
+            "14": 10.906,
+            "15": 10.90106,
+            "16": 10.88602,
+            "17": 10.88857,
+            "18": 10.8863,
+            "19": 10.88563,
+            "20": 10.83831,
+            "21": 10.82801,
+            "22": 10.81606,
+            "23": 10.80814,
+            "24": 10.77971,
+            "25": 10.77746,
+            "26": 10.76123,
+            "27": 10.74966,
+            "28": 10.69186,
+            "29": 10.66597,
+            "30": 10.63047,
+            "31": 10.62182,
+            "32": 10.61491,
+            "33": 10.5784,
+            "34": 10.54596,
+            "35": 10.54573,
+            "36": 10.53409,
+            "37": 10.5049,
+            "38": 10.50267,
+            "39": 10.47204,
+            "40": 10.44914,
+            "41": 10.426,
+            "42": 10.41331,
+            "43": 10.39913,
+            "44": 10.36866,
+            "45": 10.37972,
+            "46": 10.3331,
+            "47": 10.32171,
+            "48": 10.28474,
+            "49": 10.28349,
+            "50": 10.27386
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 18949.0,
-            "2": 19099.0,
-            "3": 19050.0,
-            "4": 19118.0,
-            "5": 19083.0,
-            "6": 18713.0,
-            "7": 18983.0,
-            "8": 18573.0,
-            "9": 19074.0,
-            "10": 19633.0,
-            "11": 18864.0,
-            "12": 18680.0,
-            "13": 18983.0,
-            "14": 19635.0,
-            "15": 18534.0,
-            "16": 18776.0,
-            "17": 19027.0,
-            "18": 19179.0,
-            "19": 19256.0,
-            "20": 19131.0,
-            "21": 18678.0,
-            "22": 19277.0,
-            "23": 19003.0,
-            "24": 19177.0,
-            "25": 18328.0,
-            "26": 18877.0,
-            "27": 19140.0,
-            "28": 18575.0,
-            "29": 18512.0,
-            "30": 18566.0,
-            "31": 18581.0,
-            "32": 19224.0,
-            "33": 18078.0,
-            "34": 18612.0,
-            "35": 18599.0,
-            "36": 18325.0,
-            "37": 18458.0,
-            "38": 18488.0,
-            "39": 18900.0,
-            "40": 19303.0,
-            "41": 18755.0,
-            "42": 18613.0,
-            "43": 19044.0,
-            "44": 18919.0,
-            "45": 20364.0,
-            "46": 19848.0,
-            "47": 20037.0,
-            "48": 19837.0,
-            "49": 21694.0,
-            "50": 20114.0
+            "1": 36671.0,
+            "2": 37173.0,
+            "3": 36899.0,
+            "4": 37068.0,
+            "5": 36679.0,
+            "6": 35978.0,
+            "7": 37031.0,
+            "8": 36229.0,
+            "9": 37090.0,
+            "10": 37967.0,
+            "11": 36941.0,
+            "12": 36081.0,
+            "13": 36942.0,
+            "14": 37392.0,
+            "15": 36186.0,
+            "16": 36099.0,
+            "17": 36904.0,
+            "18": 36981.0,
+            "19": 37455.0,
+            "20": 36609.0,
+            "21": 36447.0,
+            "22": 36641.0,
+            "23": 36893.0,
+            "24": 37094.0,
+            "25": 35824.0,
+            "26": 36997.0,
+            "27": 36498.0,
+            "28": 35858.0,
+            "29": 36103.0,
+            "30": 35720.0,
+            "31": 36188.0,
+            "32": 36884.0,
+            "33": 36092.0,
+            "34": 36008.0,
+            "35": 36767.0,
+            "36": 35759.0,
+            "37": 36357.0,
+            "38": 35639.0,
+            "39": 36859.0,
+            "40": 37400.0,
+            "41": 36473.0,
+            "42": 36460.0,
+            "43": 37019.0,
+            "44": 36370.0,
+            "45": 39844.0,
+            "46": 38478.0,
+            "47": 38793.0,
+            "48": 38437.0,
+            "49": 42666.0,
+            "50": 39545.0
         }
     },
     "mem-allocated-bytes": {
@@ -118,56 +118,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 1027089408.0,
-            "2": 1027091456.0,
-            "3": 1027087360.0,
-            "4": 1027088384.0,
-            "5": 1027091456.0,
-            "6": 1027091456.0,
-            "7": 1027088896.0,
-            "8": 1027092480.0,
+            "1": 1027088896.0,
+            "2": 1027090944.0,
+            "3": 1027086848.0,
+            "4": 1027087360.0,
+            "5": 1027091968.0,
+            "6": 1027089920.0,
+            "7": 1027088384.0,
+            "8": 1027091968.0,
             "9": 1027091968.0,
-            "10": 1027089408.0,
+            "10": 1027089920.0,
             "11": 1027089920.0,
-            "12": 1027092480.0,
+            "12": 1027091456.0,
             "13": 1027090944.0,
-            "14": 1027092480.0,
-            "15": 1027090432.0,
-            "16": 1027088384.0,
-            "17": 1027089408.0,
-            "18": 1027090944.0,
-            "19": 1027088384.0,
-            "20": 1027090432.0,
-            "21": 1027092480.0,
+            "14": 1027091456.0,
+            "15": 1027089920.0,
+            "16": 1027088896.0,
+            "17": 1027087872.0,
+            "18": 1027090432.0,
+            "19": 1027088896.0,
+            "20": 1027089408.0,
+            "21": 1027091456.0,
             "22": 1027089920.0,
-            "23": 1027093504.0,
+            "23": 1027094016.0,
             "24": 1027092480.0,
             "25": 1027089408.0,
-            "26": 1027090944.0,
-            "27": 1027087360.0,
-            "28": 1027090432.0,
-            "29": 1027090432.0,
-            "30": 1027089920.0,
-            "31": 1027089408.0,
-            "32": 1027093504.0,
-            "33": 1027094016.0,
-            "34": 1027093504.0,
-            "35": 1027085824.0,
-            "36": 1027087872.0,
+            "26": 1027089408.0,
+            "27": 1027086848.0,
+            "28": 1027090944.0,
+            "29": 1027089920.0,
+            "30": 1027090432.0,
+            "31": 1027090432.0,
+            "32": 1027092992.0,
+            "33": 1027092480.0,
+            "34": 1027091968.0,
+            "35": 1027085312.0,
+            "36": 1027086336.0,
             "37": 1027088896.0,
             "38": 1027089920.0,
-            "39": 1027088384.0,
-            "40": 1027091968.0,
+            "39": 1027087360.0,
+            "40": 1027090944.0,
             "41": 1027088384.0,
-            "42": 1027089408.0,
+            "42": 1027087872.0,
             "43": 1027087872.0,
-            "44": 1027091456.0,
+            "44": 1027088384.0,
             "45": 1027090432.0,
             "46": 1027086336.0,
-            "47": 1027088384.0,
-            "48": 1027087360.0,
-            "49": 1027087360.0,
-            "50": 1027089920.0
+            "47": 1027086848.0,
+            "48": 1027086848.0,
+            "49": 1027086336.0,
+            "50": 1027089408.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -175,113 +175,112 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 3058326528.0,
-            "2": 3298517504.0,
-            "3": 3298517504.0,
-            "4": 3298517504.0,
-            "5": 3300747776.0,
-            "6": 3300747776.0,
-            "7": 3300747776.0,
-            "8": 3300747776.0,
-            "9": 3300747776.0,
-            "10": 3300747776.0,
-            "11": 3300747776.0,
-            "12": 3300747776.0,
-            "13": 3300747776.0,
-            "14": 3300747776.0,
-            "15": 3300747776.0,
-            "16": 3300747776.0,
-            "17": 3300747776.0,
-            "18": 3300747776.0,
-            "19": 3300747776.0,
-            "20": 3300747776.0,
-            "21": 3300747776.0,
-            "22": 3300747776.0,
-            "23": 3300747776.0,
-            "24": 3300747776.0,
-            "25": 3300747776.0,
-            "26": 3300747776.0,
-            "27": 3300747776.0,
-            "28": 3300747776.0,
-            "29": 3300747776.0,
-            "30": 3300747776.0,
-            "31": 3300747776.0,
-            "32": 3300747776.0,
-            "33": 3300747776.0,
-            "34": 3300872192.0,
-            "35": 3300872192.0,
-            "36": 3300872192.0,
-            "37": 3300872192.0,
-            "38": 3300872192.0,
-            "39": 3300872192.0,
-            "40": 3300872192.0,
-            "41": 3300872192.0,
-            "42": 3300872192.0,
-            "43": 3300872192.0,
-            "44": 3300872192.0,
-            "45": 3300872192.0,
-            "46": 3300872192.0,
-            "47": 3300872192.0,
-            "48": 3300872192.0,
-            "49": 3300872192.0,
-            "50": 3300872192.0
+            "1": 3059096576.0,
+            "2": 3298776064.0,
+            "3": 3298776064.0,
+            "4": 3298776064.0,
+            "5": 3298776064.0,
+            "6": 3298776064.0,
+            "7": 3298776064.0,
+            "8": 3299136512.0,
+            "9": 3299136512.0,
+            "10": 3299136512.0,
+            "11": 3299136512.0,
+            "12": 3299136512.0,
+            "13": 3299320320.0,
+            "14": 3299397120.0,
+            "15": 3299397120.0,
+            "16": 3299397120.0,
+            "17": 3299397120.0,
+            "18": 3299397120.0,
+            "19": 3299397120.0,
+            "20": 3299397120.0,
+            "21": 3299397120.0,
+            "22": 3299397120.0,
+            "23": 3300247552.0,
+            "24": 3300247552.0,
+            "25": 3300247552.0,
+            "26": 3300247552.0,
+            "27": 3300247552.0,
+            "28": 3300247552.0,
+            "29": 3300247552.0,
+            "30": 3300247552.0,
+            "31": 3300247552.0,
+            "32": 3300247552.0,
+            "33": 3300247552.0,
+            "34": 3300554752.0,
+            "35": 3300554752.0,
+            "36": 3300554752.0,
+            "37": 3300554752.0,
+            "38": 3300554752.0,
+            "39": 3300554752.0,
+            "40": 3300554752.0,
+            "41": 3300554752.0,
+            "42": 3300554752.0,
+            "43": 3300554752.0,
+            "44": 3300554752.0,
+            "45": 3300554752.0,
+            "46": 3300554752.0,
+            "47": 3300554752.0,
+            "48": 3300554752.0,
+            "49": 3300554752.0,
+            "50": 3300554752.0
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 7.37375,
-            "3": 0.25401,
-            "4": 0.23,
-            "5": 0.23156,
-            "6": 0.22618,
-            "7": 0.22033,
-            "8": 0.2124,
-            "9": 0.21458,
-            "10": 0.2112,
-            "11": 0.22058,
-            "12": 0.21214,
-            "13": 0.20964,
-            "14": 0.21773,
-            "15": 0.21046,
-            "16": 0.21558,
-            "17": 0.21724,
-            "18": 0.21042,
-            "19": 0.2121,
-            "20": 0.21156,
-            "21": 0.2121,
-            "22": 0.20983,
-            "23": 0.22142,
-            "24": 0.21088,
-            "25": 0.21096,
-            "26": 0.2105,
-            "27": 0.21223,
-            "28": 0.21432,
-            "29": 0.20728,
-            "30": 0.20861,
-            "31": 0.20793,
-            "32": 0.20812,
-            "33": 0.20817,
-            "34": 0.20922,
-            "35": 0.20912,
-            "36": 0.21051,
-            "37": 0.21278,
-            "38": 0.21391,
-            "39": 0.2131,
-            "40": 0.21335,
-            "41": 0.21205,
-            "42": 0.20975,
-            "43": 0.2117,
-            "44": 0.21456,
-            "45": 0.21588,
-            "46": 0.21062,
-            "47": 0.21618,
-            "48": 0.21235,
-            "49": 0.21609,
-            "50": 0.21536
+            "2": 5.12545,
+            "3": 0.33231,
+            "4": 0.34056,
+            "5": 0.32602,
+            "6": 0.3278,
+            "7": 0.32667,
+            "8": 0.30797,
+            "9": 0.30763,
+            "10": 0.30676,
+            "11": 0.31256,
+            "12": 0.3065,
+            "13": 0.3017,
+            "14": 0.29823,
+            "15": 0.30407,
+            "16": 0.3043,
+            "17": 0.29941,
+            "18": 1.04497,
+            "19": 0.30398,
+            "20": 0.30088,
+            "21": 0.31032,
+            "22": 0.30474,
+            "23": 0.30372,
+            "24": 0.30325,
+            "25": 0.31021,
+            "26": 0.2994,
+            "27": 0.30871,
+            "28": 0.29653,
+            "29": 0.29361,
+            "30": 0.29596,
+            "31": 0.2957,
+            "32": 0.29902,
+            "33": 0.2982,
+            "34": 0.29551,
+            "35": 0.29378,
+            "36": 0.31272,
+            "37": 0.30903,
+            "38": 0.30702,
+            "39": 0.30131,
+            "40": 0.3076,
+            "41": 0.30128,
+            "42": 0.29866,
+            "43": 0.30149,
+            "44": 0.30224,
+            "45": 0.29648,
+            "46": 0.30138,
+            "47": 0.29993,
+            "48": 0.29673,
+            "49": 0.29749,
+            "50": 0.30226
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective_1node/golden_values_dev_dgx_gb200.json
index 7c9dfc9c650..2ce3b9d4103 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective_1node/golden_values_dev_dgx_gb200.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective_1node/golden_values_dev_dgx_gb200.json
@@ -6,54 +6,54 @@
         "values": {
             "1": 10.92337,
             "2": 10.91184,
-            "3": 10.92435,
-            "4": 10.91391,
-            "5": 10.91783,
-            "6": 10.90822,
-            "7": 10.90214,
-            "8": 10.90968,
-            "9": 10.92327,
-            "10": 10.91324,
-            "11": 10.90152,
-            "12": 10.91168,
-            "13": 10.88952,
-            "14": 10.88978,
-            "15": 10.88462,
-            "16": 10.87656,
-            "17": 10.87413,
-            "18": 10.8627,
-            "19": 10.87155,
-            "20": 10.8213,
-            "21": 10.8075,
-            "22": 10.79454,
-            "23": 10.79087,
-            "24": 10.75698,
-            "25": 10.76203,
-            "26": 10.75144,
-            "27": 10.73619,
-            "28": 10.67458,
-            "29": 10.63986,
-            "30": 10.62269,
-            "31": 10.6127,
-            "32": 10.59666,
-            "33": 10.57595,
-            "34": 10.53962,
-            "35": 10.53801,
-            "36": 10.52144,
-            "37": 10.48719,
-            "38": 10.49901,
-            "39": 10.45863,
-            "40": 10.43984,
-            "41": 10.42707,
-            "42": 10.41355,
-            "43": 10.38977,
-            "44": 10.35254,
-            "45": 10.36992,
-            "46": 10.32982,
-            "47": 10.31281,
-            "48": 10.27777,
-            "49": 10.27075,
-            "50": 10.27166
+            "3": 10.92433,
+            "4": 10.91442,
+            "5": 10.91813,
+            "6": 10.90832,
+            "7": 10.90215,
+            "8": 10.90903,
+            "9": 10.92402,
+            "10": 10.91272,
+            "11": 10.902,
+            "12": 10.91244,
+            "13": 10.88903,
+            "14": 10.88983,
+            "15": 10.88458,
+            "16": 10.8759,
+            "17": 10.87389,
+            "18": 10.86392,
+            "19": 10.87172,
+            "20": 10.82085,
+            "21": 10.80728,
+            "22": 10.79583,
+            "23": 10.79119,
+            "24": 10.75658,
+            "25": 10.76154,
+            "26": 10.75136,
+            "27": 10.73586,
+            "28": 10.67557,
+            "29": 10.64032,
+            "30": 10.62211,
+            "31": 10.61219,
+            "32": 10.59656,
+            "33": 10.57524,
+            "34": 10.53944,
+            "35": 10.53691,
+            "36": 10.52053,
+            "37": 10.48602,
+            "38": 10.49787,
+            "39": 10.45741,
+            "40": 10.43859,
+            "41": 10.42607,
+            "42": 10.41228,
+            "43": 10.38805,
+            "44": 10.35099,
+            "45": 10.36785,
+            "46": 10.32756,
+            "47": 10.31064,
+            "48": 10.27588,
+            "49": 10.26883,
+            "50": 10.26942
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 18852.0,
-            "2": 19201.0,
-            "3": 18829.0,
-            "4": 18799.0,
-            "5": 18661.0,
-            "6": 18606.0,
-            "7": 19036.0,
-            "8": 18992.0,
-            "9": 19205.0,
-            "10": 19409.0,
-            "11": 18725.0,
-            "12": 18412.0,
-            "13": 19021.0,
-            "14": 19235.0,
-            "15": 18796.0,
-            "16": 18661.0,
-            "17": 18978.0,
-            "18": 18747.0,
-            "19": 18689.0,
-            "20": 18999.0,
-            "21": 18460.0,
-            "22": 18943.0,
-            "23": 19110.0,
-            "24": 18724.0,
-            "25": 18209.0,
-            "26": 18944.0,
-            "27": 18769.0,
-            "28": 18225.0,
-            "29": 18553.0,
-            "30": 18154.0,
-            "31": 18660.0,
-            "32": 19011.0,
-            "33": 18551.0,
-            "34": 18477.0,
-            "35": 18870.0,
-            "36": 18407.0,
-            "37": 18996.0,
-            "38": 18270.0,
-            "39": 19069.0,
-            "40": 19467.0,
-            "41": 18520.0,
-            "42": 18295.0,
-            "43": 19123.0,
-            "44": 18827.0,
-            "45": 20268.0,
-            "46": 19722.0,
-            "47": 19576.0,
-            "48": 19950.0,
-            "49": 21813.0,
-            "50": 20205.0
+            "1": 36489.0,
+            "2": 36983.0,
+            "3": 36910.0,
+            "4": 36780.0,
+            "5": 36558.0,
+            "6": 35770.0,
+            "7": 36868.0,
+            "8": 36266.0,
+            "9": 36754.0,
+            "10": 37570.0,
+            "11": 36863.0,
+            "12": 35673.0,
+            "13": 36548.0,
+            "14": 36969.0,
+            "15": 36370.0,
+            "16": 35940.0,
+            "17": 36259.0,
+            "18": 36835.0,
+            "19": 36756.0,
+            "20": 36499.0,
+            "21": 36348.0,
+            "22": 37016.0,
+            "23": 36564.0,
+            "24": 36769.0,
+            "25": 35732.0,
+            "26": 36536.0,
+            "27": 36730.0,
+            "28": 35823.0,
+            "29": 35621.0,
+            "30": 35794.0,
+            "31": 36536.0,
+            "32": 36948.0,
+            "33": 35774.0,
+            "34": 35891.0,
+            "35": 36873.0,
+            "36": 36000.0,
+            "37": 36496.0,
+            "38": 35759.0,
+            "39": 37107.0,
+            "40": 37664.0,
+            "41": 36067.0,
+            "42": 35632.0,
+            "43": 37356.0,
+            "44": 36500.0,
+            "45": 39206.0,
+            "46": 38653.0,
+            "47": 38503.0,
+            "48": 38931.0,
+            "49": 41941.0,
+            "50": 39489.0
         }
     },
     "mem-allocated-bytes": {
@@ -121,53 +121,53 @@
             "1": 1388729344.0,
             "2": 1388729856.0,
             "3": 1388729344.0,
-            "4": 1388727808.0,
-            "5": 1388728832.0,
-            "6": 1388729856.0,
+            "4": 1388728320.0,
+            "5": 1388728320.0,
+            "6": 1388729344.0,
             "7": 1388731392.0,
-            "8": 1388730880.0,
-            "9": 1388731392.0,
-            "10": 1388731392.0,
-            "11": 1388729344.0,
-            "12": 1388728320.0,
-            "13": 1388732928.0,
-            "14": 1388727808.0,
-            "15": 1388730880.0,
+            "8": 1388730368.0,
+            "9": 1388731904.0,
+            "10": 1388730368.0,
+            "11": 1388728832.0,
+            "12": 1388727808.0,
+            "13": 1388733952.0,
+            "14": 1388727296.0,
+            "15": 1388730368.0,
             "16": 1388728832.0,
-            "17": 1388729856.0,
+            "17": 1388729344.0,
             "18": 1388729856.0,
-            "19": 1388730880.0,
-            "20": 1388731904.0,
-            "21": 1388729856.0,
-            "22": 1388730368.0,
-            "23": 1388731392.0,
+            "19": 1388731392.0,
+            "20": 1388730368.0,
+            "21": 1388731904.0,
+            "22": 1388730880.0,
+            "23": 1388731904.0,
             "24": 1388730880.0,
             "25": 1388732416.0,
-            "26": 1388732416.0,
-            "27": 1388731904.0,
-            "28": 1388728832.0,
-            "29": 1388730880.0,
-            "30": 1388732928.0,
-            "31": 1388731392.0,
-            "32": 1388728832.0,
+            "26": 1388732928.0,
+            "27": 1388729856.0,
+            "28": 1388729856.0,
+            "29": 1388731392.0,
+            "30": 1388732416.0,
+            "31": 1388731904.0,
+            "32": 1388728320.0,
             "33": 1388731904.0,
-            "34": 1388733952.0,
+            "34": 1388733440.0,
             "35": 1388731392.0,
-            "36": 1388730368.0,
+            "36": 1388730880.0,
             "37": 1388731392.0,
-            "38": 1388730368.0,
-            "39": 1388731904.0,
-            "40": 1388731904.0,
+            "38": 1388729856.0,
+            "39": 1388733440.0,
+            "40": 1388732416.0,
             "41": 1388729344.0,
             "42": 1388732416.0,
-            "43": 1388731904.0,
-            "44": 1388729344.0,
-            "45": 1388731904.0,
-            "46": 1388732416.0,
-            "47": 1388733440.0,
-            "48": 1388731392.0,
-            "49": 1388735488.0,
-            "50": 1388730368.0
+            "43": 1388730880.0,
+            "44": 1388730368.0,
+            "45": 1388732416.0,
+            "46": 1388731904.0,
+            "47": 1388733952.0,
+            "48": 1388731904.0,
+            "49": 1388734976.0,
+            "50": 1388731392.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -176,112 +176,111 @@
         "step_interval": 1,
         "values": {
             "1": 3179891200.0,
-            "2": 3662152192.0,
-            "3": 3662164480.0,
-            "4": 3663250944.0,
-            "5": 3663250944.0,
-            "6": 3663250944.0,
-            "7": 3663250944.0,
-            "8": 3663250944.0,
-            "9": 3663250944.0,
-            "10": 3663250944.0,
-            "11": 3663250944.0,
-            "12": 3663250944.0,
-            "13": 3663250944.0,
-            "14": 3663250944.0,
-            "15": 3663250944.0,
-            "16": 3663250944.0,
-            "17": 3663250944.0,
-            "18": 3663250944.0,
-            "19": 3663250944.0,
-            "20": 3663250944.0,
-            "21": 3663250944.0,
-            "22": 3663250944.0,
-            "23": 3663250944.0,
-            "24": 3663250944.0,
-            "25": 3663250944.0,
-            "26": 3663250944.0,
-            "27": 3663250944.0,
-            "28": 3663250944.0,
-            "29": 3663250944.0,
-            "30": 3663250944.0,
-            "31": 3663250944.0,
-            "32": 3663250944.0,
-            "33": 3663250944.0,
-            "34": 3663250944.0,
-            "35": 3663897600.0,
-            "36": 3663897600.0,
-            "37": 3663897600.0,
-            "38": 3663897600.0,
-            "39": 3663897600.0,
-            "40": 3663897600.0,
-            "41": 3663897600.0,
-            "42": 3664121344.0,
-            "43": 3664121344.0,
-            "44": 3664593408.0,
-            "45": 3664593408.0,
-            "46": 3664593408.0,
-            "47": 3664593408.0,
-            "48": 3664593408.0,
-            "49": 3664994816.0,
-            "50": 3664994816.0
+            "2": 3661529600.0,
+            "3": 3661897728.0,
+            "4": 3663147008.0,
+            "5": 3663147008.0,
+            "6": 3663147008.0,
+            "7": 3663147008.0,
+            "8": 3663147008.0,
+            "9": 3663147008.0,
+            "10": 3663147008.0,
+            "11": 3663147008.0,
+            "12": 3663147008.0,
+            "13": 3663147008.0,
+            "14": 3663147008.0,
+            "15": 3663147008.0,
+            "16": 3663147008.0,
+            "17": 3663147008.0,
+            "18": 3663147008.0,
+            "19": 3663147008.0,
+            "20": 3663147008.0,
+            "21": 3663147008.0,
+            "22": 3663147008.0,
+            "23": 3663147008.0,
+            "24": 3663147008.0,
+            "25": 3663147008.0,
+            "26": 3663147008.0,
+            "27": 3663147008.0,
+            "28": 3663147008.0,
+            "29": 3663147008.0,
+            "30": 3663147008.0,
+            "31": 3663147008.0,
+            "32": 3663147008.0,
+            "33": 3663147008.0,
+            "34": 3663147008.0,
+            "35": 3663633408.0,
+            "36": 3663633408.0,
+            "37": 3663633408.0,
+            "38": 3663633408.0,
+            "39": 3663633408.0,
+            "40": 3664276992.0,
+            "41": 3664276992.0,
+            "42": 3664276992.0,
+            "43": 3664276992.0,
+            "44": 3664562176.0,
+            "45": 3664562176.0,
+            "46": 3664562176.0,
+            "47": 3664562176.0,
+            "48": 3664941056.0,
+            "49": 3664941056.0,
+            "50": 3664941056.0
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 6.85618,
-            "3": 2.07992,
-            "4": 1.87215,
-            "5": 1.48011,
-            "6": 2.02399,
-            "7": 2.13514,
-            "8": 2.16634,
-            "9": 1.77359,
-            "10": 1.31073,
-            "11": 1.64395,
-            "12": 1.59434,
-            "13": 1.52469,
-            "14": 1.92437,
-            "15": 1.76445,
-            "16": 1.88589,
-            "17": 2.05598,
-            "18": 1.86971,
-            "19": 2.0144,
-            "20": 1.80886,
-            "21": 1.87488,
-            "22": 1.82297,
-            "23": 2.50818,
-            "24": 2.75346,
-            "25": 0.74091,
-            "26": 2.16027,
-            "27": 2.04848,
-            "28": 1.31746,
-            "29": 1.64889,
-            "30": 1.25634,
-            "31": 1.91444,
-            "32": 3.0018,
-            "33": 2.02046,
-            "34": 0.90031,
-            "35": 2.58547,
-            "36": 2.4573,
-            "37": 2.09483,
-            "38": 0.8511,
-            "39": 1.81657,
-            "40": 2.84934,
-            "41": 1.28469,
-            "42": 2.64455,
-            "43": 1.75379,
-            "44": 2.06723,
-            "45": 1.55995,
-            "46": 1.02072,
-            "47": 2.05419,
-            "48": 2.10853,
-            "49": 1.68931,
-            "50": 1.50254
+            "2": 5.15666,
+            "3": 1.89051,
+            "4": 1.70427,
+            "5": 1.36821,
+            "6": 2.22402,
+            "7": 2.0206,
+            "8": 2.64192,
+            "9": 2.49779,
+            "10": 1.43263,
+            "11": 1.40881,
+            "12": 1.54998,
+            "13": 1.38339,
+            "14": 1.66132,
+            "15": 1.27345,
+            "16": 1.64809,
+            "17": 1.39588,
+            "18": 1.42002,
+            "19": 1.40276,
+            "20": 1.33212,
+            "21": 1.40375,
+            "22": 1.69832,
+            "23": 1.54171,
+            "24": 2.83335,
+            "25": 1.15593,
+            "26": 1.50513,
+            "27": 1.49099,
+            "28": 11.01099,
+            "29": 3.57548,
+            "30": 1.01686,
+            "31": 7.90208,
+            "32": 9.8648,
+            "33": 4.94174,
+            "34": 1.21046,
+            "35": 1.99545,
+            "36": 1.50634,
+            "37": 1.39215,
+            "38": 1.49176,
+            "39": 1.71133,
+            "40": 2.30027,
+            "41": 1.40087,
+            "42": 2.15518,
+            "43": 1.49416,
+            "44": 1.23264,
+            "45": 1.42777,
+            "46": 0.83186,
+            "47": 3.49767,
+            "48": 10.11677,
+            "49": 3.23631,
+            "50": 2.68346
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_h100.json
index 83ebb282949..8267b6a3f65 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM/golden_values_dev_dgx_h100.json
@@ -6,54 +6,54 @@
         "values": {
             "1": 10.92563,
             "2": 10.91638,
-            "3": 10.92433,
-            "4": 10.93217,
-            "5": 10.92999,
-            "6": 10.92662,
-            "7": 10.92571,
-            "8": 10.92333,
-            "9": 10.92825,
-            "10": 10.91605,
-            "11": 10.91854,
-            "12": 10.92399,
-            "13": 10.91037,
-            "14": 10.90685,
-            "15": 10.90136,
-            "16": 10.88661,
-            "17": 10.88849,
-            "18": 10.88662,
-            "19": 10.8857,
-            "20": 10.83765,
-            "21": 10.82761,
-            "22": 10.81538,
-            "23": 10.8078,
-            "24": 10.78018,
-            "25": 10.778,
-            "26": 10.76109,
-            "27": 10.74912,
-            "28": 10.69195,
-            "29": 10.66617,
-            "30": 10.63122,
-            "31": 10.62222,
-            "32": 10.61543,
-            "33": 10.57901,
-            "34": 10.54677,
-            "35": 10.54608,
-            "36": 10.53419,
-            "37": 10.50598,
-            "38": 10.50458,
-            "39": 10.47274,
-            "40": 10.45062,
-            "41": 10.42727,
-            "42": 10.41444,
-            "43": 10.40126,
-            "44": 10.3705,
-            "45": 10.38167,
-            "46": 10.33539,
-            "47": 10.32458,
-            "48": 10.28718,
-            "49": 10.28599,
-            "50": 10.27739
+            "3": 10.92459,
+            "4": 10.93245,
+            "5": 10.93019,
+            "6": 10.92651,
+            "7": 10.92566,
+            "8": 10.92291,
+            "9": 10.92842,
+            "10": 10.91709,
+            "11": 10.9186,
+            "12": 10.92364,
+            "13": 10.91032,
+            "14": 10.906,
+            "15": 10.90106,
+            "16": 10.88602,
+            "17": 10.88857,
+            "18": 10.8863,
+            "19": 10.88563,
+            "20": 10.83831,
+            "21": 10.82801,
+            "22": 10.81606,
+            "23": 10.80814,
+            "24": 10.77971,
+            "25": 10.77746,
+            "26": 10.76123,
+            "27": 10.74966,
+            "28": 10.69186,
+            "29": 10.66597,
+            "30": 10.63047,
+            "31": 10.62182,
+            "32": 10.61491,
+            "33": 10.5784,
+            "34": 10.54596,
+            "35": 10.54573,
+            "36": 10.53409,
+            "37": 10.5049,
+            "38": 10.50267,
+            "39": 10.47204,
+            "40": 10.44914,
+            "41": 10.426,
+            "42": 10.41331,
+            "43": 10.39913,
+            "44": 10.36866,
+            "45": 10.37972,
+            "46": 10.3331,
+            "47": 10.32171,
+            "48": 10.28474,
+            "49": 10.28349,
+            "50": 10.27386
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 18949.0,
-            "2": 19099.0,
-            "3": 19050.0,
-            "4": 19118.0,
-            "5": 19083.0,
-            "6": 18713.0,
-            "7": 18983.0,
-            "8": 18573.0,
-            "9": 19074.0,
-            "10": 19633.0,
-            "11": 18864.0,
-            "12": 18680.0,
-            "13": 18983.0,
-            "14": 19635.0,
-            "15": 18534.0,
-            "16": 18776.0,
-            "17": 19027.0,
-            "18": 19179.0,
-            "19": 19256.0,
-            "20": 19131.0,
-            "21": 18678.0,
-            "22": 19277.0,
-            "23": 19003.0,
-            "24": 19177.0,
-            "25": 18328.0,
-            "26": 18877.0,
-            "27": 19140.0,
-            "28": 18575.0,
-            "29": 18512.0,
-            "30": 18566.0,
-            "31": 18581.0,
-            "32": 19224.0,
-            "33": 18078.0,
-            "34": 18612.0,
-            "35": 18599.0,
-            "36": 18325.0,
-            "37": 18458.0,
-            "38": 18488.0,
-            "39": 18900.0,
-            "40": 19303.0,
-            "41": 18755.0,
-            "42": 18613.0,
-            "43": 19044.0,
-            "44": 18919.0,
-            "45": 20364.0,
-            "46": 19848.0,
-            "47": 20037.0,
-            "48": 19837.0,
-            "49": 21694.0,
-            "50": 20114.0
+            "1": 36671.0,
+            "2": 37173.0,
+            "3": 36899.0,
+            "4": 37068.0,
+            "5": 36679.0,
+            "6": 35978.0,
+            "7": 37031.0,
+            "8": 36229.0,
+            "9": 37090.0,
+            "10": 37967.0,
+            "11": 36941.0,
+            "12": 36081.0,
+            "13": 36942.0,
+            "14": 37392.0,
+            "15": 36186.0,
+            "16": 36099.0,
+            "17": 36904.0,
+            "18": 36981.0,
+            "19": 37455.0,
+            "20": 36609.0,
+            "21": 36447.0,
+            "22": 36641.0,
+            "23": 36893.0,
+            "24": 37094.0,
+            "25": 35824.0,
+            "26": 36997.0,
+            "27": 36498.0,
+            "28": 35858.0,
+            "29": 36103.0,
+            "30": 35720.0,
+            "31": 36188.0,
+            "32": 36884.0,
+            "33": 36092.0,
+            "34": 36008.0,
+            "35": 36767.0,
+            "36": 35759.0,
+            "37": 36357.0,
+            "38": 35639.0,
+            "39": 36859.0,
+            "40": 37400.0,
+            "41": 36473.0,
+            "42": 36460.0,
+            "43": 37019.0,
+            "44": 36370.0,
+            "45": 39844.0,
+            "46": 38478.0,
+            "47": 38793.0,
+            "48": 38437.0,
+            "49": 42666.0,
+            "50": 39545.0
         }
     },
     "mem-allocated-bytes": {
@@ -118,56 +118,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 1027089408.0,
-            "2": 1027091456.0,
-            "3": 1027087360.0,
-            "4": 1027088384.0,
-            "5": 1027091456.0,
-            "6": 1027091456.0,
-            "7": 1027088896.0,
-            "8": 1027092480.0,
+            "1": 1027088896.0,
+            "2": 1027090944.0,
+            "3": 1027086848.0,
+            "4": 1027087360.0,
+            "5": 1027091968.0,
+            "6": 1027089920.0,
+            "7": 1027088384.0,
+            "8": 1027091968.0,
             "9": 1027091968.0,
-            "10": 1027089408.0,
+            "10": 1027089920.0,
             "11": 1027089920.0,
-            "12": 1027092480.0,
+            "12": 1027091456.0,
             "13": 1027090944.0,
-            "14": 1027092480.0,
-            "15": 1027090432.0,
-            "16": 1027088384.0,
-            "17": 1027089408.0,
-            "18": 1027090944.0,
-            "19": 1027088384.0,
-            "20": 1027090432.0,
-            "21": 1027092480.0,
+            "14": 1027091456.0,
+            "15": 1027089920.0,
+            "16": 1027088896.0,
+            "17": 1027087872.0,
+            "18": 1027090432.0,
+            "19": 1027088896.0,
+            "20": 1027089408.0,
+            "21": 1027091456.0,
             "22": 1027089920.0,
-            "23": 1027093504.0,
+            "23": 1027094016.0,
             "24": 1027092480.0,
             "25": 1027089408.0,
-            "26": 1027090944.0,
-            "27": 1027087360.0,
-            "28": 1027090432.0,
-            "29": 1027090432.0,
-            "30": 1027089920.0,
-            "31": 1027089408.0,
-            "32": 1027093504.0,
-            "33": 1027094016.0,
-            "34": 1027093504.0,
-            "35": 1027085824.0,
-            "36": 1027087872.0,
+            "26": 1027089408.0,
+            "27": 1027086848.0,
+            "28": 1027090944.0,
+            "29": 1027089920.0,
+            "30": 1027090432.0,
+            "31": 1027090432.0,
+            "32": 1027092992.0,
+            "33": 1027092480.0,
+            "34": 1027091968.0,
+            "35": 1027085312.0,
+            "36": 1027086336.0,
             "37": 1027088896.0,
             "38": 1027089920.0,
-            "39": 1027088384.0,
-            "40": 1027091968.0,
+            "39": 1027087360.0,
+            "40": 1027090944.0,
             "41": 1027088384.0,
-            "42": 1027089408.0,
+            "42": 1027087872.0,
             "43": 1027087872.0,
-            "44": 1027091456.0,
+            "44": 1027088384.0,
             "45": 1027090432.0,
             "46": 1027086336.0,
-            "47": 1027088384.0,
-            "48": 1027087360.0,
-            "49": 1027087360.0,
-            "50": 1027089920.0
+            "47": 1027086848.0,
+            "48": 1027086848.0,
+            "49": 1027086336.0,
+            "50": 1027089408.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -175,113 +175,112 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 3058326528.0,
-            "2": 3298517504.0,
-            "3": 3298517504.0,
-            "4": 3298517504.0,
-            "5": 3300747776.0,
-            "6": 3300747776.0,
-            "7": 3300747776.0,
-            "8": 3300747776.0,
-            "9": 3300747776.0,
-            "10": 3300747776.0,
-            "11": 3300747776.0,
-            "12": 3300747776.0,
-            "13": 3300747776.0,
-            "14": 3300747776.0,
-            "15": 3300747776.0,
-            "16": 3300747776.0,
-            "17": 3300747776.0,
-            "18": 3300747776.0,
-            "19": 3300747776.0,
-            "20": 3300747776.0,
-            "21": 3300747776.0,
-            "22": 3300747776.0,
-            "23": 3300747776.0,
-            "24": 3300747776.0,
-            "25": 3300747776.0,
-            "26": 3300747776.0,
-            "27": 3300747776.0,
-            "28": 3300747776.0,
-            "29": 3300747776.0,
-            "30": 3300747776.0,
-            "31": 3300747776.0,
-            "32": 3300747776.0,
-            "33": 3300747776.0,
-            "34": 3300872192.0,
-            "35": 3300872192.0,
-            "36": 3300872192.0,
-            "37": 3300872192.0,
-            "38": 3300872192.0,
-            "39": 3300872192.0,
-            "40": 3300872192.0,
-            "41": 3300872192.0,
-            "42": 3300872192.0,
-            "43": 3300872192.0,
-            "44": 3300872192.0,
-            "45": 3300872192.0,
-            "46": 3300872192.0,
-            "47": 3300872192.0,
-            "48": 3300872192.0,
-            "49": 3300872192.0,
-            "50": 3300872192.0
+            "1": 3059096576.0,
+            "2": 3298776064.0,
+            "3": 3298776064.0,
+            "4": 3298776064.0,
+            "5": 3298776064.0,
+            "6": 3298776064.0,
+            "7": 3298776064.0,
+            "8": 3299136512.0,
+            "9": 3299136512.0,
+            "10": 3299136512.0,
+            "11": 3299136512.0,
+            "12": 3299136512.0,
+            "13": 3299320320.0,
+            "14": 3299397120.0,
+            "15": 3299397120.0,
+            "16": 3299397120.0,
+            "17": 3299397120.0,
+            "18": 3299397120.0,
+            "19": 3299397120.0,
+            "20": 3299397120.0,
+            "21": 3299397120.0,
+            "22": 3299397120.0,
+            "23": 3300247552.0,
+            "24": 3300247552.0,
+            "25": 3300247552.0,
+            "26": 3300247552.0,
+            "27": 3300247552.0,
+            "28": 3300247552.0,
+            "29": 3300247552.0,
+            "30": 3300247552.0,
+            "31": 3300247552.0,
+            "32": 3300247552.0,
+            "33": 3300247552.0,
+            "34": 3300554752.0,
+            "35": 3300554752.0,
+            "36": 3300554752.0,
+            "37": 3300554752.0,
+            "38": 3300554752.0,
+            "39": 3300554752.0,
+            "40": 3300554752.0,
+            "41": 3300554752.0,
+            "42": 3300554752.0,
+            "43": 3300554752.0,
+            "44": 3300554752.0,
+            "45": 3300554752.0,
+            "46": 3300554752.0,
+            "47": 3300554752.0,
+            "48": 3300554752.0,
+            "49": 3300554752.0,
+            "50": 3300554752.0
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 7.27219,
-            "3": 0.24679,
-            "4": 0.22635,
-            "5": 0.23051,
-            "6": 0.22263,
-            "7": 0.21818,
-            "8": 0.21355,
-            "9": 0.21356,
-            "10": 0.21043,
-            "11": 0.21544,
-            "12": 0.21111,
-            "13": 0.21015,
-            "14": 0.21431,
-            "15": 0.21165,
-            "16": 0.21367,
-            "17": 0.21668,
-            "18": 0.2084,
-            "19": 0.20834,
-            "20": 0.20701,
-            "21": 0.21147,
-            "22": 0.20775,
-            "23": 0.2219,
-            "24": 0.21061,
-            "25": 0.20661,
-            "26": 0.21028,
-            "27": 0.2129,
-            "28": 0.20786,
-            "29": 0.20797,
-            "30": 0.20789,
-            "31": 0.20896,
-            "32": 0.20624,
-            "33": 0.20688,
-            "34": 0.20637,
-            "35": 0.20779,
-            "36": 0.20898,
-            "37": 0.20801,
-            "38": 0.2083,
-            "39": 0.20824,
-            "40": 0.20749,
-            "41": 0.20582,
-            "42": 0.20712,
-            "43": 0.21062,
-            "44": 0.21109,
-            "45": 0.21193,
-            "46": 0.20663,
-            "47": 0.21151,
-            "48": 0.20703,
-            "49": 0.21392,
-            "50": 0.21062
+            "2": 5.16303,
+            "3": 0.31975,
+            "4": 0.32114,
+            "5": 0.32174,
+            "6": 0.31434,
+            "7": 0.30635,
+            "8": 0.30462,
+            "9": 0.30412,
+            "10": 0.30505,
+            "11": 0.31288,
+            "12": 0.29834,
+            "13": 0.30086,
+            "14": 0.29552,
+            "15": 0.28913,
+            "16": 0.29571,
+            "17": 0.29148,
+            "18": 0.2864,
+            "19": 0.28774,
+            "20": 0.28853,
+            "21": 0.2911,
+            "22": 0.29212,
+            "23": 0.2979,
+            "24": 0.29998,
+            "25": 0.29745,
+            "26": 0.29211,
+            "27": 0.30024,
+            "28": 0.29179,
+            "29": 0.29311,
+            "30": 0.29515,
+            "31": 0.29371,
+            "32": 0.29669,
+            "33": 0.29283,
+            "34": 0.29194,
+            "35": 0.29361,
+            "36": 0.2973,
+            "37": 0.29273,
+            "38": 0.29215,
+            "39": 0.29439,
+            "40": 0.295,
+            "41": 0.28702,
+            "42": 0.29175,
+            "43": 0.28749,
+            "44": 0.29187,
+            "45": 1.10566,
+            "46": 0.28901,
+            "47": 0.2914,
+            "48": 0.30221,
+            "49": 0.30073,
+            "50": 0.29095
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM_1node/golden_values_dev_dgx_gb200.json
index 4e1261dee6a..594a99c55b6 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM_1node/golden_values_dev_dgx_gb200.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM_1node/golden_values_dev_dgx_gb200.json
@@ -6,54 +6,54 @@
         "values": {
             "1": 10.92337,
             "2": 10.91184,
-            "3": 10.92435,
-            "4": 10.91391,
-            "5": 10.91783,
-            "6": 10.90822,
-            "7": 10.90214,
-            "8": 10.90968,
-            "9": 10.92327,
-            "10": 10.91324,
-            "11": 10.90152,
-            "12": 10.91168,
-            "13": 10.88952,
-            "14": 10.88978,
-            "15": 10.88462,
-            "16": 10.87656,
-            "17": 10.87413,
-            "18": 10.8627,
-            "19": 10.87155,
-            "20": 10.8213,
-            "21": 10.8075,
-            "22": 10.79454,
-            "23": 10.79087,
-            "24": 10.75698,
-            "25": 10.76203,
-            "26": 10.75144,
-            "27": 10.73619,
-            "28": 10.67458,
-            "29": 10.63986,
-            "30": 10.62269,
-            "31": 10.6127,
-            "32": 10.59666,
-            "33": 10.57595,
-            "34": 10.53962,
-            "35": 10.53801,
-            "36": 10.52144,
-            "37": 10.48719,
-            "38": 10.49901,
-            "39": 10.45863,
-            "40": 10.43984,
-            "41": 10.42707,
-            "42": 10.41355,
-            "43": 10.38977,
-            "44": 10.35254,
-            "45": 10.36992,
-            "46": 10.32982,
-            "47": 10.31281,
-            "48": 10.27777,
-            "49": 10.27075,
-            "50": 10.27166
+            "3": 10.92433,
+            "4": 10.91442,
+            "5": 10.91813,
+            "6": 10.90832,
+            "7": 10.90215,
+            "8": 10.90903,
+            "9": 10.92402,
+            "10": 10.91272,
+            "11": 10.902,
+            "12": 10.91244,
+            "13": 10.88903,
+            "14": 10.88983,
+            "15": 10.88458,
+            "16": 10.8759,
+            "17": 10.87389,
+            "18": 10.86392,
+            "19": 10.87172,
+            "20": 10.82085,
+            "21": 10.80728,
+            "22": 10.79583,
+            "23": 10.79119,
+            "24": 10.75658,
+            "25": 10.76154,
+            "26": 10.75136,
+            "27": 10.73586,
+            "28": 10.67557,
+            "29": 10.64032,
+            "30": 10.62211,
+            "31": 10.61219,
+            "32": 10.59656,
+            "33": 10.57524,
+            "34": 10.53944,
+            "35": 10.53691,
+            "36": 10.52053,
+            "37": 10.48602,
+            "38": 10.49787,
+            "39": 10.45741,
+            "40": 10.43859,
+            "41": 10.42607,
+            "42": 10.41228,
+            "43": 10.38805,
+            "44": 10.35099,
+            "45": 10.36785,
+            "46": 10.32756,
+            "47": 10.31064,
+            "48": 10.27588,
+            "49": 10.26883,
+            "50": 10.26942
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 18852.0,
-            "2": 19201.0,
-            "3": 18829.0,
-            "4": 18799.0,
-            "5": 18661.0,
-            "6": 18606.0,
-            "7": 19036.0,
-            "8": 18992.0,
-            "9": 19205.0,
-            "10": 19409.0,
-            "11": 18725.0,
-            "12": 18412.0,
-            "13": 19021.0,
-            "14": 19235.0,
-            "15": 18796.0,
-            "16": 18661.0,
-            "17": 18978.0,
-            "18": 18747.0,
-            "19": 18689.0,
-            "20": 18999.0,
-            "21": 18460.0,
-            "22": 18943.0,
-            "23": 19110.0,
-            "24": 18724.0,
-            "25": 18209.0,
-            "26": 18944.0,
-            "27": 18769.0,
-            "28": 18225.0,
-            "29": 18553.0,
-            "30": 18154.0,
-            "31": 18660.0,
-            "32": 19011.0,
-            "33": 18551.0,
-            "34": 18477.0,
-            "35": 18870.0,
-            "36": 18407.0,
-            "37": 18996.0,
-            "38": 18270.0,
-            "39": 19069.0,
-            "40": 19467.0,
-            "41": 18520.0,
-            "42": 18295.0,
-            "43": 19123.0,
-            "44": 18827.0,
-            "45": 20268.0,
-            "46": 19722.0,
-            "47": 19576.0,
-            "48": 19950.0,
-            "49": 21813.0,
-            "50": 20205.0
+            "1": 36489.0,
+            "2": 36983.0,
+            "3": 36910.0,
+            "4": 36780.0,
+            "5": 36558.0,
+            "6": 35770.0,
+            "7": 36868.0,
+            "8": 36266.0,
+            "9": 36754.0,
+            "10": 37570.0,
+            "11": 36863.0,
+            "12": 35673.0,
+            "13": 36548.0,
+            "14": 36969.0,
+            "15": 36370.0,
+            "16": 35940.0,
+            "17": 36259.0,
+            "18": 36835.0,
+            "19": 36756.0,
+            "20": 36499.0,
+            "21": 36348.0,
+            "22": 37016.0,
+            "23": 36564.0,
+            "24": 36769.0,
+            "25": 35732.0,
+            "26": 36536.0,
+            "27": 36730.0,
+            "28": 35823.0,
+            "29": 35621.0,
+            "30": 35794.0,
+            "31": 36536.0,
+            "32": 36948.0,
+            "33": 35774.0,
+            "34": 35891.0,
+            "35": 36873.0,
+            "36": 36000.0,
+            "37": 36496.0,
+            "38": 35759.0,
+            "39": 37107.0,
+            "40": 37664.0,
+            "41": 36067.0,
+            "42": 35632.0,
+            "43": 37356.0,
+            "44": 36500.0,
+            "45": 39206.0,
+            "46": 38653.0,
+            "47": 38503.0,
+            "48": 38931.0,
+            "49": 41941.0,
+            "50": 39489.0
         }
     },
     "mem-allocated-bytes": {
@@ -121,53 +121,53 @@
             "1": 1388729344.0,
             "2": 1388729856.0,
             "3": 1388729344.0,
-            "4": 1388727808.0,
-            "5": 1388728832.0,
-            "6": 1388729856.0,
+            "4": 1388728320.0,
+            "5": 1388728320.0,
+            "6": 1388729344.0,
             "7": 1388731392.0,
-            "8": 1388730880.0,
-            "9": 1388731392.0,
-            "10": 1388731392.0,
-            "11": 1388729344.0,
-            "12": 1388728320.0,
-            "13": 1388732928.0,
-            "14": 1388727808.0,
-            "15": 1388730880.0,
+            "8": 1388730368.0,
+            "9": 1388731904.0,
+            "10": 1388730368.0,
+            "11": 1388728832.0,
+            "12": 1388727808.0,
+            "13": 1388733952.0,
+            "14": 1388727296.0,
+            "15": 1388730368.0,
             "16": 1388728832.0,
-            "17": 1388729856.0,
+            "17": 1388729344.0,
             "18": 1388729856.0,
-            "19": 1388730880.0,
-            "20": 1388731904.0,
-            "21": 1388729856.0,
-            "22": 1388730368.0,
-            "23": 1388731392.0,
+            "19": 1388731392.0,
+            "20": 1388730368.0,
+            "21": 1388731904.0,
+            "22": 1388730880.0,
+            "23": 1388731904.0,
             "24": 1388730880.0,
             "25": 1388732416.0,
-            "26": 1388732416.0,
-            "27": 1388731904.0,
-            "28": 1388728832.0,
-            "29": 1388730880.0,
-            "30": 1388732928.0,
-            "31": 1388731392.0,
-            "32": 1388728832.0,
+            "26": 1388732928.0,
+            "27": 1388729856.0,
+            "28": 1388729856.0,
+            "29": 1388731392.0,
+            "30": 1388732416.0,
+            "31": 1388731904.0,
+            "32": 1388728320.0,
             "33": 1388731904.0,
-            "34": 1388733952.0,
+            "34": 1388733440.0,
             "35": 1388731392.0,
-            "36": 1388730368.0,
+            "36": 1388730880.0,
             "37": 1388731392.0,
-            "38": 1388730368.0,
-            "39": 1388731904.0,
-            "40": 1388731904.0,
+            "38": 1388729856.0,
+            "39": 1388733440.0,
+            "40": 1388732416.0,
             "41": 1388729344.0,
             "42": 1388732416.0,
-            "43": 1388731904.0,
-            "44": 1388729344.0,
-            "45": 1388731904.0,
-            "46": 1388732416.0,
-            "47": 1388733440.0,
-            "48": 1388731392.0,
-            "49": 1388735488.0,
-            "50": 1388730368.0
+            "43": 1388730880.0,
+            "44": 1388730368.0,
+            "45": 1388732416.0,
+            "46": 1388731904.0,
+            "47": 1388733952.0,
+            "48": 1388731904.0,
+            "49": 1388734976.0,
+            "50": 1388731392.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -175,113 +175,112 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 3179891200.0,
-            "2": 3662227968.0,
-            "3": 3662227968.0,
-            "4": 3663250944.0,
-            "5": 3663250944.0,
-            "6": 3663250944.0,
-            "7": 3663250944.0,
-            "8": 3663250944.0,
-            "9": 3663250944.0,
-            "10": 3663250944.0,
-            "11": 3663250944.0,
-            "12": 3663250944.0,
-            "13": 3663250944.0,
-            "14": 3663250944.0,
-            "15": 3663250944.0,
-            "16": 3663250944.0,
-            "17": 3663250944.0,
-            "18": 3663250944.0,
-            "19": 3663250944.0,
-            "20": 3663250944.0,
-            "21": 3663250944.0,
-            "22": 3663250944.0,
-            "23": 3663250944.0,
-            "24": 3663250944.0,
-            "25": 3663250944.0,
-            "26": 3663250944.0,
-            "27": 3663250944.0,
-            "28": 3663250944.0,
-            "29": 3663250944.0,
-            "30": 3663250944.0,
-            "31": 3663250944.0,
-            "32": 3663250944.0,
-            "33": 3663250944.0,
-            "34": 3664421888.0,
-            "35": 3664421888.0,
-            "36": 3664421888.0,
-            "37": 3664421888.0,
-            "38": 3664421888.0,
-            "39": 3664421888.0,
-            "40": 3664421888.0,
-            "41": 3664421888.0,
-            "42": 3664421888.0,
-            "43": 3664421888.0,
-            "44": 3664535552.0,
-            "45": 3664535552.0,
-            "46": 3664535552.0,
-            "47": 3664535552.0,
-            "48": 3664535552.0,
-            "49": 3664612864.0,
-            "50": 3664612864.0
+            "1": 3179466240.0,
+            "2": 3662152192.0,
+            "3": 3662152192.0,
+            "4": 3663147008.0,
+            "5": 3663147008.0,
+            "6": 3663147008.0,
+            "7": 3663147008.0,
+            "8": 3663147008.0,
+            "9": 3663147008.0,
+            "10": 3663147008.0,
+            "11": 3663147008.0,
+            "12": 3663147008.0,
+            "13": 3663147008.0,
+            "14": 3663147008.0,
+            "15": 3663147008.0,
+            "16": 3663147008.0,
+            "17": 3663147008.0,
+            "18": 3663147008.0,
+            "19": 3663147008.0,
+            "20": 3663147008.0,
+            "21": 3663147008.0,
+            "22": 3663147008.0,
+            "23": 3663147008.0,
+            "24": 3663147008.0,
+            "25": 3663147008.0,
+            "26": 3663147008.0,
+            "27": 3663147008.0,
+            "28": 3663147008.0,
+            "29": 3663147008.0,
+            "30": 3663147008.0,
+            "31": 3663147008.0,
+            "32": 3663147008.0,
+            "33": 3663147008.0,
+            "34": 3663147008.0,
+            "35": 3663633408.0,
+            "36": 3663633408.0,
+            "37": 3663633408.0,
+            "38": 3663633408.0,
+            "39": 3663633408.0,
+            "40": 3664276992.0,
+            "41": 3664276992.0,
+            "42": 3664276992.0,
+            "43": 3664276992.0,
+            "44": 3664562176.0,
+            "45": 3664562176.0,
+            "46": 3664562176.0,
+            "47": 3664562176.0,
+            "48": 3664941056.0,
+            "49": 3664941056.0,
+            "50": 3664941056.0
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 6.76257,
-            "3": 2.41181,
-            "4": 2.14671,
-            "5": 1.68613,
-            "6": 2.35596,
-            "7": 1.85689,
-            "8": 3.55826,
-            "9": 2.10313,
-            "10": 1.69174,
-            "11": 2.07594,
-            "12": 2.11237,
-            "13": 2.0346,
-            "14": 2.34804,
-            "15": 1.56257,
-            "16": 1.96304,
-            "17": 2.0887,
-            "18": 2.24792,
-            "19": 1.71276,
-            "20": 1.64151,
-            "21": 2.03443,
-            "22": 1.86292,
-            "23": 1.95656,
-            "24": 2.17465,
-            "25": 2.28198,
-            "26": 1.89479,
-            "27": 1.24118,
-            "28": 2.07544,
-            "29": 1.97422,
-            "30": 1.20941,
-            "31": 1.60534,
-            "32": 2.43502,
-            "33": 2.08637,
-            "34": 1.84735,
-            "35": 2.40386,
-            "36": 1.40304,
-            "37": 2.14664,
-            "38": 0.89545,
-            "39": 2.42476,
-            "40": 2.25422,
-            "41": 1.29899,
-            "42": 2.6352,
-            "43": 1.42406,
-            "44": 2.03708,
-            "45": 2.10233,
-            "46": 0.7566,
-            "47": 1.72618,
-            "48": 1.51965,
-            "49": 1.43953,
-            "50": 1.66225
+            "2": 5.39007,
+            "3": 2.22863,
+            "4": 1.54551,
+            "5": 1.86051,
+            "6": 1.81825,
+            "7": 1.83536,
+            "8": 2.64367,
+            "9": 2.04074,
+            "10": 1.39125,
+            "11": 1.58824,
+            "12": 1.51442,
+            "13": 1.55316,
+            "14": 2.01252,
+            "15": 1.07894,
+            "16": 1.60711,
+            "17": 1.8311,
+            "18": 1.79223,
+            "19": 1.69767,
+            "20": 2.25425,
+            "21": 1.10313,
+            "22": 1.1802,
+            "23": 2.17979,
+            "24": 1.83343,
+            "25": 1.4927,
+            "26": 1.46116,
+            "27": 1.40226,
+            "28": 2.35976,
+            "29": 1.44806,
+            "30": 1.08009,
+            "31": 1.20756,
+            "32": 2.55991,
+            "33": 1.71293,
+            "34": 2.84147,
+            "35": 2.40341,
+            "36": 1.99923,
+            "37": 2.07605,
+            "38": 1.16269,
+            "39": 1.45685,
+            "40": 2.61788,
+            "41": 1.76194,
+            "42": 2.30587,
+            "43": 1.5739,
+            "44": 1.95111,
+            "45": 1.65086,
+            "46": 0.91469,
+            "47": 1.7889,
+            "48": 1.52903,
+            "49": 1.16314,
+            "50": 1.17655
         }
     }
-}
\ 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 06a1d993f8b..8bcf15522c7 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
@@ -6,54 +6,54 @@
         "values": {
             "1": 10.90791,
             "2": 10.90713,
-            "3": 10.91668,
-            "4": 10.90899,
-            "5": 10.91483,
-            "6": 10.89524,
-            "7": 10.90676,
-            "8": 10.90896,
-            "9": 10.90939,
-            "10": 10.91077,
-            "11": 10.901,
-            "12": 10.89922,
-            "13": 10.88807,
-            "14": 10.88197,
-            "15": 10.87251,
-            "16": 10.85287,
-            "17": 10.85704,
-            "18": 10.84826,
-            "19": 10.85507,
-            "20": 10.77687,
-            "21": 10.76073,
-            "22": 10.7605,
-            "23": 10.74325,
-            "24": 10.70838,
-            "25": 10.70981,
-            "26": 10.69235,
-            "27": 10.66868,
-            "28": 10.60599,
-            "29": 10.57223,
-            "30": 10.54151,
-            "31": 10.53199,
-            "32": 10.51634,
-            "33": 10.481,
-            "34": 10.44913,
-            "35": 10.44632,
-            "36": 10.42066,
-            "37": 10.40067,
-            "38": 10.40454,
-            "39": 10.36981,
-            "40": 10.35244,
-            "41": 10.33008,
-            "42": 10.31128,
-            "43": 10.29795,
-            "44": 10.27171,
-            "45": 10.28363,
-            "46": 10.24114,
-            "47": 10.23434,
-            "48": 10.19197,
-            "49": 10.19498,
-            "50": 10.19073
+            "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
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 16592.0,
-            "2": 16504.0,
-            "3": 16575.0,
-            "4": 16376.0,
-            "5": 16183.0,
-            "6": 16014.0,
-            "7": 16818.0,
-            "8": 15883.0,
-            "9": 16693.0,
-            "10": 16580.0,
-            "11": 16233.0,
-            "12": 16030.0,
-            "13": 16714.0,
-            "14": 16720.0,
-            "15": 15895.0,
-            "16": 16103.0,
-            "17": 16623.0,
-            "18": 16575.0,
-            "19": 16738.0,
-            "20": 16598.0,
-            "21": 16235.0,
-            "22": 16514.0,
-            "23": 16299.0,
-            "24": 16586.0,
-            "25": 15770.0,
-            "26": 16540.0,
-            "27": 16550.0,
-            "28": 16182.0,
-            "29": 16694.0,
-            "30": 16454.0,
-            "31": 17100.0,
-            "32": 17390.0,
-            "33": 17062.0,
-            "34": 17233.0,
-            "35": 17703.0,
-            "36": 17351.0,
-            "37": 18011.0,
-            "38": 17621.0,
-            "39": 18243.0,
-            "40": 18971.0,
-            "41": 18220.0,
-            "42": 17966.0,
-            "43": 18752.0,
-            "44": 18809.0,
-            "45": 20890.0,
-            "46": 19846.0,
-            "47": 19418.0,
-            "48": 20136.0,
-            "49": 22380.0,
-            "50": 20145.0
+            "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
         }
     },
     "mem-allocated-bytes": {
@@ -118,56 +118,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 1562452480.0,
-            "2": 1560886272.0,
-            "3": 1560936960.0,
-            "4": 1560987648.0,
-            "5": 1560990208.0,
-            "6": 1560987648.0,
-            "7": 1561697792.0,
-            "8": 1561797632.0,
-            "9": 1560936960.0,
-            "10": 1560987648.0,
-            "11": 1562030592.0,
-            "12": 1561056256.0,
-            "13": 1561038336.0,
-            "14": 1561974784.0,
-            "15": 1561152000.0,
-            "16": 1561089024.0,
-            "17": 1561038336.0,
-            "18": 1562183680.0,
-            "19": 1562019328.0,
-            "20": 1561089024.0,
-            "21": 1561960960.0,
-            "22": 1561631744.0,
-            "23": 1561836032.0,
-            "24": 1561089024.0,
-            "25": 1561139712.0,
-            "26": 1561089024.0,
-            "27": 1561139712.0,
-            "28": 1561190400.0,
-            "29": 1561772544.0,
-            "30": 1561604096.0,
-            "31": 1562013184.0,
-            "32": 1561406464.0,
-            "33": 1561139712.0,
-            "34": 1561755648.0,
-            "35": 1561589248.0,
-            "36": 1561190400.0,
-            "37": 1561662976.0,
-            "38": 1561190400.0,
-            "39": 1561241088.0,
-            "40": 1563022336.0,
-            "41": 1564137984.0,
-            "42": 1561291776.0,
-            "43": 1561952768.0,
-            "44": 1561291776.0,
-            "45": 1561715200.0,
-            "46": 1561291776.0,
-            "47": 1561342464.0,
-            "48": 1561291776.0,
-            "49": 1561342464.0,
-            "50": 1561291776.0
+            "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
         }
     },
     "mem-max-allocated-bytes": {
@@ -175,113 +175,112 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 3479876608.0,
-            "2": 4042335744.0,
-            "3": 4047595520.0,
-            "4": 4053274112.0,
-            "5": 4053274112.0,
-            "6": 4053274112.0,
-            "7": 4057433088.0,
-            "8": 4061642240.0,
-            "9": 4061642240.0,
-            "10": 4065529344.0,
-            "11": 4065529344.0,
-            "12": 4065529344.0,
-            "13": 4065529344.0,
-            "14": 4065529344.0,
-            "15": 4065529344.0,
-            "16": 4065529344.0,
-            "17": 4065529344.0,
-            "18": 4065529344.0,
-            "19": 4065529344.0,
-            "20": 4065529344.0,
-            "21": 4065529344.0,
-            "22": 4065529344.0,
-            "23": 4065529344.0,
-            "24": 4065529344.0,
-            "25": 4065529344.0,
-            "26": 4065529344.0,
-            "27": 4065529344.0,
-            "28": 4065529344.0,
-            "29": 4065529344.0,
-            "30": 4065529344.0,
-            "31": 4065529344.0,
-            "32": 4065529344.0,
-            "33": 4065529344.0,
-            "34": 4065529344.0,
-            "35": 4065529344.0,
-            "36": 4065529344.0,
-            "37": 4065529344.0,
-            "38": 4065529344.0,
-            "39": 4065529344.0,
-            "40": 4065529344.0,
-            "41": 4065529344.0,
-            "42": 4065529344.0,
-            "43": 4065529344.0,
-            "44": 4065529344.0,
-            "45": 4065529344.0,
-            "46": 4065529344.0,
-            "47": 4065529344.0,
-            "48": 4065529344.0,
-            "49": 4065529344.0,
-            "50": 4065529344.0
+            "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
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 11.44286,
-            "3": 0.39137,
-            "4": 0.33071,
-            "5": 0.32257,
-            "6": 0.32404,
-            "7": 0.30595,
-            "8": 0.29297,
-            "9": 0.29395,
-            "10": 0.27912,
-            "11": 0.30251,
-            "12": 0.28669,
-            "13": 0.28455,
-            "14": 0.28124,
-            "15": 0.2876,
-            "16": 0.27705,
-            "17": 0.28277,
-            "18": 0.28818,
-            "19": 0.29518,
-            "20": 0.28783,
-            "21": 0.28453,
-            "22": 0.28955,
-            "23": 0.27766,
-            "24": 0.278,
-            "25": 0.28149,
-            "26": 0.29603,
-            "27": 0.27934,
-            "28": 0.29048,
-            "29": 0.29607,
-            "30": 0.28981,
-            "31": 0.32857,
-            "32": 0.29071,
-            "33": 0.29613,
-            "34": 0.2968,
-            "35": 0.30616,
-            "36": 0.30069,
-            "37": 0.29431,
-            "38": 0.29876,
-            "39": 0.30582,
-            "40": 0.28349,
-            "41": 0.28535,
-            "42": 0.28254,
-            "43": 0.2788,
-            "44": 0.27508,
-            "45": 0.27863,
-            "46": 0.27541,
-            "47": 0.27561,
-            "48": 0.27969,
-            "49": 0.27721,
-            "50": 0.27313
+            "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
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
index 635fcea4a97..455d3cbe1c4 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
@@ -7,103 +7,103 @@
             "1": 10.92486,
             "2": 10.91069,
             "3": 10.91839,
-            "4": 10.91715,
-            "5": 10.90503,
-            "6": 10.90196,
-            "7": 10.89732,
-            "8": 10.91344,
-            "9": 10.91625,
-            "10": 10.91028,
-            "11": 10.90163,
-            "12": 10.89726,
-            "13": 10.8879,
-            "14": 10.89483,
-            "15": 10.87506,
-            "16": 10.87056,
-            "17": 10.86912,
-            "18": 10.85168,
-            "19": 10.87022,
-            "20": 10.78801,
-            "21": 10.77233,
-            "22": 10.76715,
-            "23": 10.75842,
-            "24": 10.71926,
-            "25": 10.71997,
-            "26": 10.71229,
-            "27": 10.68551,
-            "28": 10.61305,
-            "29": 10.58637,
-            "30": 10.56568,
-            "31": 10.55773,
-            "32": 10.54888,
-            "33": 10.50977,
-            "34": 10.48172,
-            "35": 10.47011,
-            "36": 10.45293,
-            "37": 10.42772,
-            "38": 10.43271,
-            "39": 10.40299,
-            "40": 10.3775,
-            "41": 10.36875,
-            "42": 10.33113,
-            "43": 10.31542,
-            "44": 10.29012,
-            "45": 10.30282,
-            "46": 10.2657,
-            "47": 10.25567,
-            "48": 10.20709,
-            "49": 10.21058,
-            "50": 10.21068,
+            "4": 10.9172,
+            "5": 10.90494,
+            "6": 10.90213,
+            "7": 10.8971,
+            "8": 10.91319,
+            "9": 10.9167,
+            "10": 10.9103,
+            "11": 10.90137,
+            "12": 10.8968,
+            "13": 10.88802,
+            "14": 10.89532,
+            "15": 10.87553,
+            "16": 10.87001,
+            "17": 10.86926,
+            "18": 10.8521,
+            "19": 10.86958,
+            "20": 10.78798,
+            "21": 10.77234,
+            "22": 10.7674,
+            "23": 10.75859,
+            "24": 10.7193,
+            "25": 10.72025,
+            "26": 10.71224,
+            "27": 10.68505,
+            "28": 10.61329,
+            "29": 10.5867,
+            "30": 10.56566,
+            "31": 10.55778,
+            "32": 10.54894,
+            "33": 10.5097,
+            "34": 10.48124,
+            "35": 10.46985,
+            "36": 10.45295,
+            "37": 10.42776,
+            "38": 10.43249,
+            "39": 10.40289,
+            "40": 10.37722,
+            "41": 10.36871,
+            "42": 10.33138,
+            "43": 10.31516,
+            "44": 10.29024,
+            "45": 10.3027,
+            "46": 10.26552,
+            "47": 10.25565,
+            "48": 10.20692,
+            "49": 10.21051,
+            "50": 10.21041,
             "51": 10.21195,
-            "52": 10.16251,
-            "53": 10.16325,
-            "54": 10.13402,
-            "55": 10.10872,
-            "56": 10.13453,
-            "57": 10.13277,
-            "58": 10.12413,
-            "59": 10.06521,
-            "60": 10.09524,
-            "61": 10.04762,
-            "62": 10.01553,
-            "63": 10.08297,
-            "64": 10.03279,
-            "65": 9.99846,
-            "66": 10.03919,
-            "67": 10.01295,
-            "68": 9.97759,
-            "69": 9.99341,
-            "70": 9.97104,
-            "71": 9.9982,
-            "72": 9.97568,
-            "73": 9.95991,
-            "74": 9.95299,
-            "75": 9.9144,
-            "76": 9.95006,
-            "77": 9.94205,
-            "78": 9.89904,
-            "79": 9.89709,
-            "80": 9.91042,
-            "81": 9.93365,
-            "82": 9.88344,
-            "83": 9.83983,
-            "84": 9.7821,
-            "85": 9.76275,
-            "86": 9.87789,
-            "87": 9.9007,
-            "88": 9.87402,
-            "89": 9.82463,
-            "90": 9.81377,
-            "91": 9.8198,
-            "92": 9.81606,
-            "93": 9.74349,
-            "94": 9.82172,
-            "95": 9.81227,
-            "96": 9.79491,
-            "97": 9.74649,
-            "98": 9.7688,
-            "99": 9.81824,
-            "100": 9.70741
+            "52": 10.16255,
+            "53": 10.1631,
+            "54": 10.13394,
+            "55": 10.10862,
+            "56": 10.13462,
+            "57": 10.13253,
+            "58": 10.12393,
+            "59": 10.06506,
+            "60": 10.09504,
+            "61": 10.04746,
+            "62": 10.01501,
+            "63": 10.08262,
+            "64": 10.03252,
+            "65": 9.99817,
+            "66": 10.03863,
+            "67": 10.01262,
+            "68": 9.97709,
+            "69": 9.99273,
+            "70": 9.97043,
+            "71": 9.99761,
+            "72": 9.97478,
+            "73": 9.95903,
+            "74": 9.95198,
+            "75": 9.91332,
+            "76": 9.9493,
+            "77": 9.94113,
+            "78": 9.89802,
+            "79": 9.89619,
+            "80": 9.90937,
+            "81": 9.93263,
+            "82": 9.88265,
+            "83": 9.83871,
+            "84": 9.78108,
+            "85": 9.76154,
+            "86": 9.87689,
+            "87": 9.89981,
+            "88": 9.87312,
+            "89": 9.82362,
+            "90": 9.81273,
+            "91": 9.81873,
+            "92": 9.81493,
+            "93": 9.74223,
+            "94": 9.82036,
+            "95": 9.81099,
+            "96": 9.79374,
+            "97": 9.74486,
+            "98": 9.76728,
+            "99": 9.81697,
+            "100": 9.70593
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 2429.0,
-            "2": 2591.0,
-            "3": 2604.0,
-            "4": 2665.0,
-            "5": 2610.0,
-            "6": 2487.0,
-            "7": 2589.0,
-            "8": 2535.0,
-            "9": 2677.0,
-            "10": 2474.0,
-            "11": 2585.0,
-            "12": 2605.0,
-            "13": 2530.0,
-            "14": 2651.0,
-            "15": 2509.0,
-            "16": 2506.0,
-            "17": 2646.0,
-            "18": 2601.0,
-            "19": 2552.0,
-            "20": 2518.0,
-            "21": 2539.0,
-            "22": 2594.0,
-            "23": 2531.0,
-            "24": 2604.0,
-            "25": 2474.0,
-            "26": 2505.0,
-            "27": 2647.0,
-            "28": 2551.0,
-            "29": 2735.0,
-            "30": 2709.0,
-            "31": 2746.0,
-            "32": 2729.0,
-            "33": 2672.0,
-            "34": 2741.0,
-            "35": 2722.0,
-            "36": 2761.0,
-            "37": 2860.0,
-            "38": 2827.0,
-            "39": 3030.0,
-            "40": 3060.0,
-            "41": 3129.0,
-            "42": 2813.0,
-            "43": 3059.0,
-            "44": 3088.0,
-            "45": 3301.0,
-            "46": 3239.0,
-            "47": 3241.0,
-            "48": 3278.0,
-            "49": 3483.0,
-            "50": 3340.0,
-            "51": 3328.0,
-            "52": 3384.0,
-            "53": 3253.0,
-            "54": 3558.0,
-            "55": 3290.0,
-            "56": 3548.0,
-            "57": 2934.0,
-            "58": 3989.0,
-            "59": 3538.0,
-            "60": 3638.0,
-            "61": 3440.0,
-            "62": 3763.0,
-            "63": 3857.0,
-            "64": 4201.0,
-            "65": 3318.0,
-            "66": 3743.0,
-            "67": 4019.0,
-            "68": 3853.0,
-            "69": 3501.0,
-            "70": 3812.0,
-            "71": 3749.0,
-            "72": 3597.0,
-            "73": 4178.0,
-            "74": 3676.0,
-            "75": 3677.0,
-            "76": 4080.0,
-            "77": 4149.0,
-            "78": 4236.0,
-            "79": 7606.0,
-            "80": 33298.0,
-            "81": 8170.0,
-            "82": 528771.0,
-            "83": 3458.0,
-            "84": 31908.0,
-            "85": 528749.0,
-            "86": 529194.0,
-            "87": 61926.0,
-            "88": 529038.0,
-            "89": 529251.0,
-            "90": 28836.0,
-            "91": 528719.0,
-            "92": 529072.0,
-            "93": 1053703.0,
-            "94": 529234.0,
-            "95": 553148.0,
-            "96": 560606.0,
-            "97": 529810.0,
-            "98": 529332.0,
-            "99": 529265.0,
-            "100": 529071.0
+            "1": 6308.0,
+            "2": 6526.0,
+            "3": 6460.0,
+            "4": 6587.0,
+            "5": 6499.0,
+            "6": 6370.0,
+            "7": 6708.0,
+            "8": 6525.0,
+            "9": 6677.0,
+            "10": 6625.0,
+            "11": 6727.0,
+            "12": 6349.0,
+            "13": 6508.0,
+            "14": 6824.0,
+            "15": 6084.0,
+            "16": 6481.0,
+            "17": 6610.0,
+            "18": 6402.0,
+            "19": 6412.0,
+            "20": 6120.0,
+            "21": 6510.0,
+            "22": 6553.0,
+            "23": 6598.0,
+            "24": 6752.0,
+            "25": 6592.0,
+            "26": 6412.0,
+            "27": 6775.0,
+            "28": 6714.0,
+            "29": 7049.0,
+            "30": 6871.0,
+            "31": 7154.0,
+            "32": 7296.0,
+            "33": 6998.0,
+            "34": 7308.0,
+            "35": 7361.0,
+            "36": 7195.0,
+            "37": 7698.0,
+            "38": 7541.0,
+            "39": 7777.0,
+            "40": 7986.0,
+            "41": 8348.0,
+            "42": 7583.0,
+            "43": 8268.0,
+            "44": 7990.0,
+            "45": 8716.0,
+            "46": 8372.0,
+            "47": 8571.0,
+            "48": 8629.0,
+            "49": 8993.0,
+            "50": 8812.0,
+            "51": 8718.0,
+            "52": 9112.0,
+            "53": 8324.0,
+            "54": 9142.0,
+            "55": 8346.0,
+            "56": 9464.0,
+            "57": 7897.0,
+            "58": 10393.0,
+            "59": 9474.0,
+            "60": 9199.0,
+            "61": 8898.0,
+            "62": 9492.0,
+            "63": 10239.0,
+            "64": 10243.0,
+            "65": 8621.0,
+            "66": 9522.0,
+            "67": 10217.0,
+            "68": 9767.0,
+            "69": 8996.0,
+            "70": 9595.0,
+            "71": 9956.0,
+            "72": 9400.0,
+            "73": 10200.0,
+            "74": 14258.0,
+            "75": 8955.0,
+            "76": 10156.0,
+            "77": 10151.0,
+            "78": 10834.0,
+            "79": 76326.0,
+            "80": 132648.0,
+            "81": 80347.0,
+            "82": 1118255.0,
+            "83": 66887.0,
+            "84": 1113524.0,
+            "85": 2106481.0,
+            "86": 2107956.0,
+            "87": 290100.0,
+            "88": 2107621.0,
+            "89": 2107900.0,
+            "90": 58643.0,
+            "91": 2106770.0,
+            "92": 2107002.0,
+            "93": 3155613.0,
+            "94": 2107794.0,
+            "95": 2155564.0,
+            "96": 2170246.0,
+            "97": 2108416.0,
+            "98": 2107390.0,
+            "99": 2107633.0,
+            "100": 2107960.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 628064256.0,
-            "2": 628065280.0,
-            "3": 628065280.0,
-            "4": 628065280.0,
-            "5": 628065280.0,
-            "6": 628065280.0,
-            "7": 628065280.0,
-            "8": 628065280.0,
-            "9": 628065280.0,
-            "10": 628065280.0,
-            "11": 628065280.0,
-            "12": 628065280.0,
-            "13": 628065280.0,
-            "14": 628065280.0,
-            "15": 628065280.0,
-            "16": 628065280.0,
-            "17": 628065280.0,
-            "18": 628065280.0,
-            "19": 628065280.0,
-            "20": 628065280.0,
-            "21": 628065280.0,
-            "22": 628065280.0,
-            "23": 628065280.0,
-            "24": 628065280.0,
-            "25": 628065280.0,
-            "26": 628065280.0,
-            "27": 628065280.0,
-            "28": 628065280.0,
-            "29": 628065280.0,
-            "30": 628065280.0,
-            "31": 628065280.0,
-            "32": 628065280.0,
-            "33": 628065280.0,
-            "34": 628065280.0,
-            "35": 628065280.0,
-            "36": 628065280.0,
-            "37": 628065280.0,
-            "38": 628065280.0,
-            "39": 628065280.0,
-            "40": 628065280.0,
-            "41": 628065280.0,
-            "42": 628065280.0,
-            "43": 628065280.0,
-            "44": 628065280.0,
-            "45": 628065280.0,
-            "46": 628065280.0,
-            "47": 628065280.0,
-            "48": 628065280.0,
-            "49": 628065280.0,
-            "50": 628065280.0,
-            "51": 628065280.0,
-            "52": 628065280.0,
-            "53": 628065280.0,
-            "54": 628065280.0,
-            "55": 628065280.0,
-            "56": 628065280.0,
-            "57": 628065280.0,
-            "58": 628065280.0,
-            "59": 628065280.0,
-            "60": 628065280.0,
-            "61": 628065280.0,
-            "62": 628065280.0,
-            "63": 628065280.0,
-            "64": 628065280.0,
-            "65": 628065280.0,
-            "66": 628065280.0,
-            "67": 628065280.0,
-            "68": 628065280.0,
-            "69": 628065280.0,
-            "70": 628065280.0,
-            "71": 628065280.0,
-            "72": 628065280.0,
-            "73": 628065280.0,
-            "74": 628065280.0,
-            "75": 628065280.0,
-            "76": 628065280.0,
-            "77": 628065280.0,
-            "78": 628065280.0,
-            "79": 628065280.0,
-            "80": 628065280.0,
-            "81": 628065280.0,
-            "82": 628065280.0,
-            "83": 628065280.0,
-            "84": 628065280.0,
-            "85": 628065280.0,
-            "86": 628065280.0,
-            "87": 628065280.0,
-            "88": 628065280.0,
-            "89": 628065280.0,
-            "90": 628065280.0,
-            "91": 628065280.0,
-            "92": 628065280.0,
-            "93": 628065280.0,
-            "94": 628065280.0,
-            "95": 628065280.0,
-            "96": 628065280.0,
-            "97": 628065280.0,
-            "98": 628065280.0,
-            "99": 628065280.0,
-            "100": 628065280.0
+            "1": 628063744.0,
+            "2": 628064768.0,
+            "3": 628064768.0,
+            "4": 628064768.0,
+            "5": 628064768.0,
+            "6": 628064768.0,
+            "7": 628064768.0,
+            "8": 628064768.0,
+            "9": 628064768.0,
+            "10": 628064768.0,
+            "11": 628064768.0,
+            "12": 628064768.0,
+            "13": 628064768.0,
+            "14": 628064768.0,
+            "15": 628064768.0,
+            "16": 628064768.0,
+            "17": 628064768.0,
+            "18": 628064768.0,
+            "19": 628064768.0,
+            "20": 628064768.0,
+            "21": 628064768.0,
+            "22": 628064768.0,
+            "23": 628064768.0,
+            "24": 628064768.0,
+            "25": 628064768.0,
+            "26": 628064768.0,
+            "27": 628064768.0,
+            "28": 628064768.0,
+            "29": 628064768.0,
+            "30": 628064768.0,
+            "31": 628064768.0,
+            "32": 628064768.0,
+            "33": 628064768.0,
+            "34": 628064768.0,
+            "35": 628064768.0,
+            "36": 628064768.0,
+            "37": 628064768.0,
+            "38": 628064768.0,
+            "39": 628064768.0,
+            "40": 628064768.0,
+            "41": 628064768.0,
+            "42": 628064768.0,
+            "43": 628064768.0,
+            "44": 628064768.0,
+            "45": 628064768.0,
+            "46": 628064768.0,
+            "47": 628064768.0,
+            "48": 628064768.0,
+            "49": 628064768.0,
+            "50": 628064768.0,
+            "51": 628064768.0,
+            "52": 628064768.0,
+            "53": 628064768.0,
+            "54": 628064768.0,
+            "55": 628064768.0,
+            "56": 628064768.0,
+            "57": 628064768.0,
+            "58": 628064768.0,
+            "59": 628064768.0,
+            "60": 628064768.0,
+            "61": 628064768.0,
+            "62": 628064768.0,
+            "63": 628064768.0,
+            "64": 628064768.0,
+            "65": 628064768.0,
+            "66": 628064768.0,
+            "67": 628064768.0,
+            "68": 628064768.0,
+            "69": 628064768.0,
+            "70": 628064768.0,
+            "71": 628064768.0,
+            "72": 628064768.0,
+            "73": 628064768.0,
+            "74": 628064768.0,
+            "75": 628064768.0,
+            "76": 628064768.0,
+            "77": 628064768.0,
+            "78": 628064768.0,
+            "79": 628064768.0,
+            "80": 628064768.0,
+            "81": 628064768.0,
+            "82": 628064768.0,
+            "83": 628064768.0,
+            "84": 628064768.0,
+            "85": 628064768.0,
+            "86": 628064768.0,
+            "87": 628064768.0,
+            "88": 628064768.0,
+            "89": 628064768.0,
+            "90": 628064768.0,
+            "91": 628064768.0,
+            "92": 628064768.0,
+            "93": 628064768.0,
+            "94": 628064768.0,
+            "95": 628064768.0,
+            "96": 628064768.0,
+            "97": 628064768.0,
+            "98": 628064768.0,
+            "99": 628064768.0,
+            "100": 628064768.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,213 +325,212 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 974434304.0,
-            "2": 1143423488.0,
-            "3": 1143897600.0,
-            "4": 1146645504.0,
-            "5": 1147860992.0,
-            "6": 1147860992.0,
-            "7": 1148556800.0,
-            "8": 1148556800.0,
-            "9": 1148556800.0,
-            "10": 1148556800.0,
-            "11": 1148556800.0,
-            "12": 1148556800.0,
-            "13": 1148556800.0,
-            "14": 1148556800.0,
-            "15": 1148556800.0,
-            "16": 1148556800.0,
-            "17": 1148556800.0,
-            "18": 1148556800.0,
-            "19": 1148556800.0,
-            "20": 1148556800.0,
-            "21": 1148556800.0,
-            "22": 1148556800.0,
-            "23": 1148556800.0,
-            "24": 1148556800.0,
-            "25": 1148556800.0,
-            "26": 1149693440.0,
-            "27": 1149693440.0,
-            "28": 1149693440.0,
-            "29": 1149693440.0,
-            "30": 1149693440.0,
-            "31": 1149693440.0,
-            "32": 1149693440.0,
-            "33": 1149693440.0,
-            "34": 1149693440.0,
-            "35": 1149693440.0,
-            "36": 1149693440.0,
-            "37": 1149693440.0,
-            "38": 1149693440.0,
-            "39": 1149693440.0,
-            "40": 1149693440.0,
-            "41": 1149693440.0,
-            "42": 1149693440.0,
-            "43": 1149693440.0,
-            "44": 1149693440.0,
-            "45": 1149693440.0,
-            "46": 1149693440.0,
-            "47": 1149693440.0,
-            "48": 1149693440.0,
-            "49": 1149693440.0,
-            "50": 1149693440.0,
-            "51": 1149693440.0,
-            "52": 1149693440.0,
-            "53": 1149693440.0,
-            "54": 1149693440.0,
-            "55": 1149693440.0,
-            "56": 1149693440.0,
-            "57": 1149693440.0,
-            "58": 1149693440.0,
-            "59": 1149693440.0,
-            "60": 1149693440.0,
-            "61": 1149693440.0,
-            "62": 1149693440.0,
-            "63": 1149693440.0,
-            "64": 1149693440.0,
-            "65": 1149693440.0,
-            "66": 1149693440.0,
-            "67": 1149693440.0,
-            "68": 1149693440.0,
-            "69": 1149693440.0,
-            "70": 1149693440.0,
-            "71": 1149693440.0,
-            "72": 1149693440.0,
-            "73": 1149693440.0,
-            "74": 1149693440.0,
-            "75": 1149693440.0,
-            "76": 1149693440.0,
-            "77": 1149693440.0,
-            "78": 1149693440.0,
-            "79": 1149693440.0,
-            "80": 1149693440.0,
-            "81": 1149693440.0,
-            "82": 1149693440.0,
-            "83": 1149693440.0,
-            "84": 1149693440.0,
-            "85": 1149693440.0,
-            "86": 1149693440.0,
-            "87": 1149693440.0,
-            "88": 1149693440.0,
-            "89": 1149693440.0,
-            "90": 1149693440.0,
-            "91": 1149693440.0,
-            "92": 1149693440.0,
-            "93": 1149693440.0,
-            "94": 1149693440.0,
-            "95": 1149693440.0,
-            "96": 1149693440.0,
-            "97": 1149693440.0,
-            "98": 1149693440.0,
-            "99": 1149693440.0,
-            "100": 1149693440.0
+            "1": 974433792.0,
+            "2": 1143422976.0,
+            "3": 1143894016.0,
+            "4": 1147345408.0,
+            "5": 1147816448.0,
+            "6": 1147816448.0,
+            "7": 1148524544.0,
+            "8": 1148524544.0,
+            "9": 1148524544.0,
+            "10": 1148524544.0,
+            "11": 1148524544.0,
+            "12": 1148524544.0,
+            "13": 1148524544.0,
+            "14": 1148524544.0,
+            "15": 1148524544.0,
+            "16": 1148524544.0,
+            "17": 1148524544.0,
+            "18": 1148524544.0,
+            "19": 1148524544.0,
+            "20": 1148524544.0,
+            "21": 1148524544.0,
+            "22": 1148524544.0,
+            "23": 1148524544.0,
+            "24": 1148524544.0,
+            "25": 1148524544.0,
+            "26": 1149626368.0,
+            "27": 1149626368.0,
+            "28": 1149626368.0,
+            "29": 1149626368.0,
+            "30": 1149626368.0,
+            "31": 1149626368.0,
+            "32": 1149626368.0,
+            "33": 1149626368.0,
+            "34": 1149626368.0,
+            "35": 1149626368.0,
+            "36": 1149626368.0,
+            "37": 1149626368.0,
+            "38": 1149626368.0,
+            "39": 1149626368.0,
+            "40": 1149626368.0,
+            "41": 1149626368.0,
+            "42": 1149626368.0,
+            "43": 1149626368.0,
+            "44": 1149626368.0,
+            "45": 1149626368.0,
+            "46": 1149626368.0,
+            "47": 1149626368.0,
+            "48": 1149626368.0,
+            "49": 1149626368.0,
+            "50": 1149626368.0,
+            "51": 1149626368.0,
+            "52": 1149626368.0,
+            "53": 1149626368.0,
+            "54": 1149626368.0,
+            "55": 1149626368.0,
+            "56": 1149626368.0,
+            "57": 1149626368.0,
+            "58": 1149626368.0,
+            "59": 1149626368.0,
+            "60": 1149626368.0,
+            "61": 1149626368.0,
+            "62": 1149626368.0,
+            "63": 1149626368.0,
+            "64": 1149626368.0,
+            "65": 1149626368.0,
+            "66": 1149626368.0,
+            "67": 1149626368.0,
+            "68": 1149626368.0,
+            "69": 1149626368.0,
+            "70": 1149626368.0,
+            "71": 1149626368.0,
+            "72": 1149626368.0,
+            "73": 1149626368.0,
+            "74": 1149626368.0,
+            "75": 1149626368.0,
+            "76": 1149626368.0,
+            "77": 1149626368.0,
+            "78": 1149626368.0,
+            "79": 1149626368.0,
+            "80": 1149626368.0,
+            "81": 1149626368.0,
+            "82": 1149626368.0,
+            "83": 1149626368.0,
+            "84": 1149626368.0,
+            "85": 1149626368.0,
+            "86": 1149626368.0,
+            "87": 1149626368.0,
+            "88": 1149626368.0,
+            "89": 1149626368.0,
+            "90": 1149626368.0,
+            "91": 1149626368.0,
+            "92": 1149626368.0,
+            "93": 1149626368.0,
+            "94": 1149626368.0,
+            "95": 1149626368.0,
+            "96": 1149626368.0,
+            "97": 1149626368.0,
+            "98": 1149626368.0,
+            "99": 1149626368.0,
+            "100": 1149626368.0
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 8.40786,
-            "3": 0.89447,
-            "4": 0.87487,
-            "5": 0.85695,
-            "6": 0.85548,
-            "7": 0.86384,
-            "8": 0.8398,
-            "9": 0.83442,
-            "10": 0.83457,
-            "11": 0.83165,
-            "12": 0.82049,
-            "13": 0.81938,
-            "14": 0.8372,
-            "15": 0.81635,
-            "16": 0.82269,
-            "17": 0.81755,
-            "18": 0.82139,
-            "19": 0.81834,
-            "20": 0.81571,
-            "21": 0.82027,
-            "22": 0.81783,
-            "23": 0.82434,
-            "24": 0.8179,
-            "25": 0.81779,
-            "26": 0.80609,
-            "27": 0.81441,
-            "28": 0.83081,
-            "29": 0.82504,
-            "30": 0.81873,
-            "31": 0.82454,
-            "32": 0.81663,
-            "33": 0.80909,
-            "34": 0.82198,
-            "35": 0.81846,
-            "36": 0.81614,
-            "37": 0.81026,
-            "38": 0.84604,
-            "39": 0.82085,
-            "40": 0.8318,
-            "41": 0.82267,
-            "42": 0.81837,
-            "43": 0.87684,
-            "44": 0.81896,
-            "45": 0.82655,
-            "46": 0.8241,
-            "47": 0.82308,
-            "48": 0.81433,
-            "49": 0.83989,
-            "50": 0.82395,
-            "51": 0.87417,
-            "52": 0.8737,
-            "53": 0.81483,
-            "54": 0.82825,
-            "55": 0.83667,
-            "56": 0.83546,
-            "57": 0.83562,
-            "58": 0.83505,
-            "59": 0.83375,
-            "60": 0.83021,
-            "61": 0.82875,
-            "62": 0.83214,
-            "63": 0.83746,
-            "64": 0.83687,
-            "65": 0.8281,
-            "66": 0.8317,
-            "67": 0.82752,
-            "68": 0.82693,
-            "69": 0.83293,
-            "70": 0.83375,
-            "71": 0.8272,
-            "72": 0.82716,
-            "73": 0.83134,
-            "74": 1.39559,
-            "75": 1.46874,
-            "76": 0.83059,
-            "77": 0.83236,
-            "78": 0.83428,
-            "79": 0.835,
-            "80": 0.83444,
-            "81": 0.83542,
-            "82": 0.84117,
-            "83": 0.83432,
-            "84": 0.82381,
-            "85": 0.831,
-            "86": 1.48456,
-            "87": 1.37924,
-            "88": 0.82167,
-            "89": 0.82408,
-            "90": 0.81692,
-            "91": 0.81059,
-            "92": 0.81301,
-            "93": 0.8096,
-            "94": 0.8091,
-            "95": 0.80549,
-            "96": 0.80731,
-            "97": 0.81231,
-            "98": 0.8007,
-            "99": 0.80887,
-            "100": 0.81095
+            "2": 8.99025,
+            "3": 0.90833,
+            "4": 0.86431,
+            "5": 0.87482,
+            "6": 0.8573,
+            "7": 0.85553,
+            "8": 0.83513,
+            "9": 0.83483,
+            "10": 0.83557,
+            "11": 0.83528,
+            "12": 1.58117,
+            "13": 0.82194,
+            "14": 0.81823,
+            "15": 0.81808,
+            "16": 0.81282,
+            "17": 0.82548,
+            "18": 0.81502,
+            "19": 0.81167,
+            "20": 0.81094,
+            "21": 0.83617,
+            "22": 0.828,
+            "23": 0.82514,
+            "24": 0.85341,
+            "25": 0.81784,
+            "26": 0.81255,
+            "27": 0.81988,
+            "28": 0.84249,
+            "29": 1.54481,
+            "30": 0.82635,
+            "31": 0.81779,
+            "32": 0.83092,
+            "33": 0.82788,
+            "34": 0.8237,
+            "35": 0.83024,
+            "36": 0.81681,
+            "37": 0.81326,
+            "38": 1.61795,
+            "39": 0.84407,
+            "40": 0.85127,
+            "41": 0.82922,
+            "42": 0.83611,
+            "43": 0.81901,
+            "44": 2.40554,
+            "45": 0.81924,
+            "46": 0.84478,
+            "47": 0.8247,
+            "48": 1.5457,
+            "49": 0.81497,
+            "50": 0.80868,
+            "51": 0.95984,
+            "52": 0.95192,
+            "53": 0.8168,
+            "54": 0.83474,
+            "55": 0.82917,
+            "56": 0.81693,
+            "57": 0.81814,
+            "58": 0.80755,
+            "59": 0.80597,
+            "60": 0.81208,
+            "61": 0.81909,
+            "62": 0.80757,
+            "63": 0.8268,
+            "64": 0.81229,
+            "65": 0.80745,
+            "66": 0.81789,
+            "67": 0.79693,
+            "68": 0.81825,
+            "69": 0.82426,
+            "70": 0.8204,
+            "71": 0.80984,
+            "72": 0.8026,
+            "73": 0.80831,
+            "74": 0.86362,
+            "75": 0.82144,
+            "76": 0.87744,
+            "77": 0.81254,
+            "78": 0.82353,
+            "79": 0.85476,
+            "80": 0.8258,
+            "81": 0.81312,
+            "82": 0.81209,
+            "83": 0.80825,
+            "84": 0.81139,
+            "85": 0.80849,
+            "86": 0.81299,
+            "87": 0.81396,
+            "88": 0.80619,
+            "89": 0.79982,
+            "90": 0.80843,
+            "91": 0.81705,
+            "92": 0.8122,
+            "93": 0.8039,
+            "94": 0.80977,
+            "95": 0.81732,
+            "96": 0.80769,
+            "97": 0.81238,
+            "98": 0.80923,
+            "99": 0.80613,
+            "100": 0.80743
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node/golden_values_dev_dgx_gb200.json
index 03672a09acd..d1508989d13 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node/golden_values_dev_dgx_gb200.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node/golden_values_dev_dgx_gb200.json
@@ -7,103 +7,103 @@
             "1": 10.92169,
             "2": 10.90762,
             "3": 10.90914,
-            "4": 10.91959,
-            "5": 10.90569,
-            "6": 10.91601,
-            "7": 10.92268,
-            "8": 10.90935,
-            "9": 10.91155,
-            "10": 10.90713,
-            "11": 10.88867,
-            "12": 10.90462,
-            "13": 10.89519,
-            "14": 10.89552,
-            "15": 10.8746,
-            "16": 10.86175,
-            "17": 10.86443,
-            "18": 10.85034,
-            "19": 10.85454,
-            "20": 10.78432,
-            "21": 10.77769,
-            "22": 10.75615,
-            "23": 10.74474,
-            "24": 10.71765,
-            "25": 10.71594,
-            "26": 10.70996,
-            "27": 10.67222,
-            "28": 10.60454,
-            "29": 10.57083,
-            "30": 10.54221,
-            "31": 10.54476,
-            "32": 10.53155,
-            "33": 10.49746,
-            "34": 10.4747,
-            "35": 10.47049,
-            "36": 10.44085,
-            "37": 10.41646,
-            "38": 10.41713,
-            "39": 10.38512,
-            "40": 10.36837,
-            "41": 10.34662,
-            "42": 10.3344,
-            "43": 10.31362,
-            "44": 10.28901,
-            "45": 10.29313,
-            "46": 10.26547,
-            "47": 10.24518,
-            "48": 10.1985,
-            "49": 10.19468,
-            "50": 10.19581,
-            "51": 10.20229,
-            "52": 10.15471,
-            "53": 10.16268,
-            "54": 10.12733,
-            "55": 10.09567,
-            "56": 10.12483,
-            "57": 10.11356,
-            "58": 10.12404,
-            "59": 10.07408,
-            "60": 10.09371,
-            "61": 10.04305,
-            "62": 10.01618,
-            "63": 10.08733,
-            "64": 10.03798,
-            "65": 10.01758,
-            "66": 10.04082,
-            "67": 10.0182,
-            "68": 9.98006,
-            "69": 10.00358,
-            "70": 9.98391,
-            "71": 10.0103,
-            "72": 9.99161,
-            "73": 9.98074,
-            "74": 9.96226,
-            "75": 9.94239,
-            "76": 9.97139,
-            "77": 9.96076,
-            "78": 9.91543,
-            "79": 9.92224,
-            "80": 9.93093,
-            "81": 9.95886,
-            "82": 9.8954,
-            "83": 9.86306,
-            "84": 9.79599,
-            "85": 9.78812,
-            "86": 9.88186,
-            "87": 9.90774,
-            "88": 9.88268,
-            "89": 9.82997,
-            "90": 9.8221,
-            "91": 9.82939,
-            "92": 9.82008,
-            "93": 9.76022,
-            "94": 9.83152,
-            "95": 9.82801,
-            "96": 9.80754,
-            "97": 9.75369,
-            "98": 9.78408,
-            "99": 9.826,
-            "100": 9.71837
+            "4": 10.91973,
+            "5": 10.90565,
+            "6": 10.91579,
+            "7": 10.92266,
+            "8": 10.90944,
+            "9": 10.9121,
+            "10": 10.90717,
+            "11": 10.88827,
+            "12": 10.90454,
+            "13": 10.89547,
+            "14": 10.89568,
+            "15": 10.87449,
+            "16": 10.86135,
+            "17": 10.86449,
+            "18": 10.85059,
+            "19": 10.85449,
+            "20": 10.78456,
+            "21": 10.77782,
+            "22": 10.75571,
+            "23": 10.7451,
+            "24": 10.71744,
+            "25": 10.71583,
+            "26": 10.70999,
+            "27": 10.6725,
+            "28": 10.60441,
+            "29": 10.57075,
+            "30": 10.54191,
+            "31": 10.54486,
+            "32": 10.53129,
+            "33": 10.49749,
+            "34": 10.47481,
+            "35": 10.47098,
+            "36": 10.4409,
+            "37": 10.41647,
+            "38": 10.41736,
+            "39": 10.3848,
+            "40": 10.36814,
+            "41": 10.34667,
+            "42": 10.33413,
+            "43": 10.31323,
+            "44": 10.28878,
+            "45": 10.29329,
+            "46": 10.26533,
+            "47": 10.24502,
+            "48": 10.1982,
+            "49": 10.19453,
+            "50": 10.19587,
+            "51": 10.20209,
+            "52": 10.15439,
+            "53": 10.1624,
+            "54": 10.12711,
+            "55": 10.09528,
+            "56": 10.12434,
+            "57": 10.11303,
+            "58": 10.12378,
+            "59": 10.07361,
+            "60": 10.09334,
+            "61": 10.04273,
+            "62": 10.01574,
+            "63": 10.08695,
+            "64": 10.03735,
+            "65": 10.01698,
+            "66": 10.0401,
+            "67": 10.01756,
+            "68": 9.97942,
+            "69": 10.00274,
+            "70": 9.98318,
+            "71": 10.00977,
+            "72": 9.99083,
+            "73": 9.97972,
+            "74": 9.96126,
+            "75": 9.94138,
+            "76": 9.97037,
+            "77": 9.95962,
+            "78": 9.91416,
+            "79": 9.92096,
+            "80": 9.92938,
+            "81": 9.95731,
+            "82": 9.89416,
+            "83": 9.86158,
+            "84": 9.79459,
+            "85": 9.78662,
+            "86": 9.8805,
+            "87": 9.90664,
+            "88": 9.88117,
+            "89": 9.82889,
+            "90": 9.82091,
+            "91": 9.82792,
+            "92": 9.81871,
+            "93": 9.75871,
+            "94": 9.82985,
+            "95": 9.82669,
+            "96": 9.80575,
+            "97": 9.75194,
+            "98": 9.78246,
+            "99": 9.82444,
+            "100": 9.71677
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 2485.0,
-            "2": 2410.0,
-            "3": 2558.0,
-            "4": 2521.0,
-            "5": 2379.0,
-            "6": 2355.0,
-            "7": 2651.0,
-            "8": 2390.0,
-            "9": 2608.0,
-            "10": 2524.0,
-            "11": 2482.0,
-            "12": 2420.0,
-            "13": 2711.0,
-            "14": 2610.0,
-            "15": 2447.0,
-            "16": 2348.0,
-            "17": 2554.0,
-            "18": 2448.0,
-            "19": 2569.0,
-            "20": 2425.0,
-            "21": 2494.0,
-            "22": 2534.0,
-            "23": 2440.0,
-            "24": 2576.0,
-            "25": 2394.0,
-            "26": 2550.0,
-            "27": 2531.0,
-            "28": 2635.0,
-            "29": 2658.0,
-            "30": 2646.0,
-            "31": 2811.0,
-            "32": 2768.0,
-            "33": 32343.0,
-            "34": 2854.0,
-            "35": 527663.0,
-            "36": 527633.0,
-            "37": 25982.0,
-            "38": 527719.0,
-            "39": 527760.0,
-            "40": 527781.0,
-            "41": 31784.0,
-            "42": 527721.0,
-            "43": 528039.0,
-            "44": 527836.0,
-            "45": 30434.0,
-            "46": 528026.0,
-            "47": 528121.0,
-            "48": 528129.0,
-            "49": 528389.0,
-            "50": 550759.0,
-            "51": 531404.0,
-            "52": 533197.0,
-            "53": 1059399.0,
-            "54": 554922.0,
-            "55": 553732.0,
-            "56": 1053105.0,
-            "57": 1052632.0,
-            "58": 1078718.0,
-            "59": 557322.0,
-            "60": 582238.0,
-            "61": 1085352.0,
-            "62": 1578412.0,
-            "63": 557697.0,
-            "64": 554190.0,
-            "65": 1578355.0,
-            "66": 1059050.0,
-            "67": 1086225.0,
-            "68": 1053764.0,
-            "69": 561544.0,
-            "70": 531044.0,
-            "71": 1054678.0,
-            "72": 1053885.0,
-            "73": 1578126.0,
-            "74": 1055279.0,
-            "75": 556894.0,
-            "76": 1081506.0,
-            "77": 552287.0,
-            "78": 556331.0,
-            "79": 1077913.0,
-            "80": 578071.0,
-            "81": 1084863.0,
-            "82": 1080774.0,
-            "83": 1577846.0,
-            "84": 575407.0,
-            "85": 558901.0,
-            "86": 554895.0,
-            "87": 575490.0,
-            "88": 533778.0,
-            "89": 1053554.0,
-            "90": 1578611.0,
-            "91": 1053399.0,
-            "92": 1053583.0,
-            "93": 529118.0,
-            "94": 1053702.0,
-            "95": 558281.0,
-            "96": 1055792.0,
-            "97": 1054543.0,
-            "98": 1054324.0,
-            "99": 554749.0,
-            "100": 1578781.0
+            "1": 6375.0,
+            "2": 6416.0,
+            "3": 6525.0,
+            "4": 6457.0,
+            "5": 6273.0,
+            "6": 6109.0,
+            "7": 6669.0,
+            "8": 6430.0,
+            "9": 6708.0,
+            "10": 6443.0,
+            "11": 6458.0,
+            "12": 6151.0,
+            "13": 6572.0,
+            "14": 6699.0,
+            "15": 6360.0,
+            "16": 6322.0,
+            "17": 6591.0,
+            "18": 6232.0,
+            "19": 6393.0,
+            "20": 6142.0,
+            "21": 6576.0,
+            "22": 6560.0,
+            "23": 6296.0,
+            "24": 6476.0,
+            "25": 6394.0,
+            "26": 6513.0,
+            "27": 6523.0,
+            "28": 6752.0,
+            "29": 7074.0,
+            "30": 6676.0,
+            "31": 7133.0,
+            "32": 7218.0,
+            "33": 14417.0,
+            "34": 7490.0,
+            "35": 1056500.0,
+            "36": 1056500.0,
+            "37": 53811.0,
+            "38": 1056688.0,
+            "39": 1056914.0,
+            "40": 1057239.0,
+            "41": 65654.0,
+            "42": 1056693.0,
+            "43": 1057447.0,
+            "44": 1057197.0,
+            "45": 62972.0,
+            "46": 1057458.0,
+            "47": 1057702.0,
+            "48": 1057863.0,
+            "49": 1058586.0,
+            "50": 1061397.0,
+            "51": 1064572.0,
+            "52": 1068311.0,
+            "53": 2171476.0,
+            "54": 1111879.0,
+            "55": 1109170.0,
+            "56": 2108974.0,
+            "57": 2106394.0,
+            "58": 2159236.0,
+            "59": 1117431.0,
+            "60": 2166475.0,
+            "61": 2224014.0,
+            "62": 3161907.0,
+            "63": 1127153.0,
+            "64": 1119365.0,
+            "65": 4205396.0,
+            "66": 2119442.0,
+            "67": 2175193.0,
+            "68": 1107351.0,
+            "69": 2169556.0,
+            "70": 2152217.0,
+            "71": 3157239.0,
+            "72": 3156914.0,
+            "73": 3206334.0,
+            "74": 3199858.0,
+            "75": 2165206.0,
+            "76": 3159262.0,
+            "77": 1163799.0,
+            "78": 2162285.0,
+            "79": 2167416.0,
+            "80": 1167052.0,
+            "81": 2212175.0,
+            "82": 2175237.0,
+            "83": 2170605.0,
+            "84": 1156708.0,
+            "85": 2167112.0,
+            "86": 2159256.0,
+            "87": 1153212.0,
+            "88": 1137179.0,
+            "89": 2155808.0,
+            "90": 4206118.0,
+            "91": 3155614.0,
+            "92": 3155708.0,
+            "93": 2107178.0,
+            "94": 3155959.0,
+            "95": 2109446.0,
+            "96": 2161322.0,
+            "97": 3157406.0,
+            "98": 3156865.0,
+            "99": 2203211.0,
+            "100": 4205763.0
         }
     },
     "mem-allocated-bytes": {
@@ -325,213 +325,212 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 1322636288.0,
-            "2": 1613951488.0,
-            "3": 1617380864.0,
-            "4": 1620309504.0,
-            "5": 1620309504.0,
-            "6": 1620309504.0,
-            "7": 1620309504.0,
-            "8": 1620309504.0,
-            "9": 1620309504.0,
-            "10": 1620309504.0,
-            "11": 1620309504.0,
-            "12": 1620309504.0,
-            "13": 1620309504.0,
-            "14": 1620309504.0,
-            "15": 1620309504.0,
-            "16": 1620309504.0,
-            "17": 1620309504.0,
-            "18": 1620309504.0,
-            "19": 1620309504.0,
-            "20": 1620309504.0,
-            "21": 1620309504.0,
-            "22": 1620309504.0,
-            "23": 1620309504.0,
-            "24": 1620309504.0,
-            "25": 1620309504.0,
-            "26": 1620309504.0,
-            "27": 1620309504.0,
-            "28": 1620309504.0,
-            "29": 1620309504.0,
-            "30": 1620309504.0,
-            "31": 1620309504.0,
-            "32": 1620309504.0,
-            "33": 1620309504.0,
-            "34": 1620309504.0,
-            "35": 1620309504.0,
-            "36": 1620309504.0,
-            "37": 1620309504.0,
-            "38": 1620309504.0,
-            "39": 1620309504.0,
-            "40": 1620309504.0,
-            "41": 1620309504.0,
-            "42": 1620309504.0,
-            "43": 1620309504.0,
-            "44": 1620309504.0,
-            "45": 1620309504.0,
-            "46": 1620309504.0,
-            "47": 1620309504.0,
-            "48": 1620309504.0,
-            "49": 1620309504.0,
-            "50": 1620309504.0,
-            "51": 1620309504.0,
-            "52": 1620309504.0,
-            "53": 1620309504.0,
-            "54": 1620309504.0,
-            "55": 1620309504.0,
-            "56": 1620309504.0,
-            "57": 1620309504.0,
-            "58": 1620309504.0,
-            "59": 1620309504.0,
-            "60": 1620309504.0,
-            "61": 1620309504.0,
-            "62": 1620309504.0,
-            "63": 1620309504.0,
-            "64": 1620309504.0,
-            "65": 1620309504.0,
-            "66": 1620309504.0,
-            "67": 1620309504.0,
-            "68": 1620309504.0,
-            "69": 1620309504.0,
-            "70": 1620309504.0,
-            "71": 1620309504.0,
-            "72": 1620309504.0,
-            "73": 1620309504.0,
-            "74": 1620309504.0,
-            "75": 1620309504.0,
-            "76": 1620309504.0,
-            "77": 1620309504.0,
-            "78": 1620309504.0,
-            "79": 1620309504.0,
-            "80": 1620309504.0,
-            "81": 1620309504.0,
-            "82": 1620309504.0,
-            "83": 1620309504.0,
-            "84": 1620309504.0,
-            "85": 1620309504.0,
-            "86": 1620309504.0,
-            "87": 1620309504.0,
-            "88": 1620309504.0,
-            "89": 1620309504.0,
-            "90": 1620309504.0,
-            "91": 1620309504.0,
-            "92": 1620309504.0,
-            "93": 1620309504.0,
-            "94": 1620309504.0,
-            "95": 1620309504.0,
-            "96": 1620309504.0,
-            "97": 1620309504.0,
-            "98": 1620309504.0,
-            "99": 1620309504.0,
-            "100": 1620309504.0
+            "1": 1322703872.0,
+            "2": 1613727232.0,
+            "3": 1618076160.0,
+            "4": 1618448384.0,
+            "5": 1618448384.0,
+            "6": 1618448384.0,
+            "7": 1618448384.0,
+            "8": 1618448384.0,
+            "9": 1618448384.0,
+            "10": 1618448384.0,
+            "11": 1620705792.0,
+            "12": 1620705792.0,
+            "13": 1620705792.0,
+            "14": 1620705792.0,
+            "15": 1620705792.0,
+            "16": 1620705792.0,
+            "17": 1620705792.0,
+            "18": 1620705792.0,
+            "19": 1620705792.0,
+            "20": 1620705792.0,
+            "21": 1620705792.0,
+            "22": 1620705792.0,
+            "23": 1620705792.0,
+            "24": 1620705792.0,
+            "25": 1620705792.0,
+            "26": 1620705792.0,
+            "27": 1620705792.0,
+            "28": 1620705792.0,
+            "29": 1620705792.0,
+            "30": 1620705792.0,
+            "31": 1620705792.0,
+            "32": 1620705792.0,
+            "33": 1620705792.0,
+            "34": 1620705792.0,
+            "35": 1620705792.0,
+            "36": 1620705792.0,
+            "37": 1620705792.0,
+            "38": 1620705792.0,
+            "39": 1620705792.0,
+            "40": 1620705792.0,
+            "41": 1620705792.0,
+            "42": 1620705792.0,
+            "43": 1620705792.0,
+            "44": 1620705792.0,
+            "45": 1620705792.0,
+            "46": 1620705792.0,
+            "47": 1620705792.0,
+            "48": 1620705792.0,
+            "49": 1620705792.0,
+            "50": 1620705792.0,
+            "51": 1620705792.0,
+            "52": 1620705792.0,
+            "53": 1620705792.0,
+            "54": 1620705792.0,
+            "55": 1620705792.0,
+            "56": 1620705792.0,
+            "57": 1620705792.0,
+            "58": 1620705792.0,
+            "59": 1620705792.0,
+            "60": 1620705792.0,
+            "61": 1620705792.0,
+            "62": 1620705792.0,
+            "63": 1620705792.0,
+            "64": 1620705792.0,
+            "65": 1620705792.0,
+            "66": 1620705792.0,
+            "67": 1620705792.0,
+            "68": 1620705792.0,
+            "69": 1620705792.0,
+            "70": 1620705792.0,
+            "71": 1620705792.0,
+            "72": 1620705792.0,
+            "73": 1620705792.0,
+            "74": 1620705792.0,
+            "75": 1620705792.0,
+            "76": 1620705792.0,
+            "77": 1620705792.0,
+            "78": 1620705792.0,
+            "79": 1620705792.0,
+            "80": 1620705792.0,
+            "81": 1620705792.0,
+            "82": 1620705792.0,
+            "83": 1620705792.0,
+            "84": 1620705792.0,
+            "85": 1620705792.0,
+            "86": 1620705792.0,
+            "87": 1620705792.0,
+            "88": 1620705792.0,
+            "89": 1620705792.0,
+            "90": 1620705792.0,
+            "91": 1620705792.0,
+            "92": 1620705792.0,
+            "93": 1620705792.0,
+            "94": 1620705792.0,
+            "95": 1620705792.0,
+            "96": 1620705792.0,
+            "97": 1620705792.0,
+            "98": 1620705792.0,
+            "99": 1620705792.0,
+            "100": 1620705792.0
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 8.44616,
-            "3": 4.35946,
-            "4": 3.58536,
-            "5": 3.43258,
-            "6": 4.01618,
-            "7": 3.41532,
-            "8": 4.69585,
-            "9": 4.42146,
-            "10": 3.2698,
-            "11": 4.19498,
-            "12": 3.71608,
-            "13": 4.078,
-            "14": 3.7934,
-            "15": 2.75437,
-            "16": 2.90972,
-            "17": 3.42375,
-            "18": 3.12947,
-            "19": 3.74111,
-            "20": 2.77586,
-            "21": 3.59019,
-            "22": 3.15397,
-            "23": 4.71787,
-            "24": 3.61079,
-            "25": 3.91574,
-            "26": 4.22139,
-            "27": 3.87761,
-            "28": 3.93585,
-            "29": 3.34081,
-            "30": 3.02171,
-            "31": 3.39522,
-            "32": 5.21075,
-            "33": 2.81038,
-            "34": 3.41835,
-            "35": 4.45162,
-            "36": 4.20517,
-            "37": 3.9982,
-            "38": 3.03186,
-            "39": 6.15837,
-            "40": 3.06211,
-            "41": 3.45761,
-            "42": 4.42809,
-            "43": 3.47773,
-            "44": 3.66168,
-            "45": 3.60833,
-            "46": 2.92395,
-            "47": 2.64427,
-            "48": 2.86604,
-            "49": 2.50808,
-            "50": 3.04166,
-            "51": 2.18664,
-            "52": 3.25348,
-            "53": 3.69136,
-            "54": 3.55702,
-            "55": 3.0236,
-            "56": 4.7145,
-            "57": 2.71755,
-            "58": 3.34597,
-            "59": 3.09264,
-            "60": 3.14001,
-            "61": 3.06867,
-            "62": 4.22075,
-            "63": 2.45818,
-            "64": 3.80461,
-            "65": 2.8975,
-            "66": 3.83227,
-            "67": 3.41683,
-            "68": 3.28576,
-            "69": 3.16606,
-            "70": 3.69457,
-            "71": 3.84527,
-            "72": 2.94589,
-            "73": 2.78343,
-            "74": 3.33954,
-            "75": 3.40883,
-            "76": 3.35809,
-            "77": 3.09918,
-            "78": 3.74403,
-            "79": 3.86757,
-            "80": 3.56784,
-            "81": 3.62657,
-            "82": 2.52733,
-            "83": 3.21777,
-            "84": 2.77215,
-            "85": 3.51637,
-            "86": 3.14139,
-            "87": 4.25648,
-            "88": 3.21899,
-            "89": 3.79927,
-            "90": 3.25394,
-            "91": 3.48754,
-            "92": 3.57614,
-            "93": 2.08265,
-            "94": 3.45887,
-            "95": 3.83381,
-            "96": 3.66511,
-            "97": 3.09335,
-            "98": 2.62447,
-            "99": 4.07289,
-            "100": 3.1867
+            "2": 8.11788,
+            "3": 4.18228,
+            "4": 2.72856,
+            "5": 2.74768,
+            "6": 3.58337,
+            "7": 3.05344,
+            "8": 3.63101,
+            "9": 3.55617,
+            "10": 3.39428,
+            "11": 3.27369,
+            "12": 3.05256,
+            "13": 4.19382,
+            "14": 3.51376,
+            "15": 3.33174,
+            "16": 3.34286,
+            "17": 3.28752,
+            "18": 3.36605,
+            "19": 3.34594,
+            "20": 3.06058,
+            "21": 3.73314,
+            "22": 2.06924,
+            "23": 3.19849,
+            "24": 2.73243,
+            "25": 3.48441,
+            "26": 5.64283,
+            "27": 9.3865,
+            "28": 2.82874,
+            "29": 2.66418,
+            "30": 3.68009,
+            "31": 3.99777,
+            "32": 4.22254,
+            "33": 2.35028,
+            "34": 2.60248,
+            "35": 4.32557,
+            "36": 2.99871,
+            "37": 3.06038,
+            "38": 2.56642,
+            "39": 3.5962,
+            "40": 2.97409,
+            "41": 3.18763,
+            "42": 3.96569,
+            "43": 2.75426,
+            "44": 3.76212,
+            "45": 2.68676,
+            "46": 2.25774,
+            "47": 2.65821,
+            "48": 3.58569,
+            "49": 3.10933,
+            "50": 2.72659,
+            "51": 3.10832,
+            "52": 2.54657,
+            "53": 3.5435,
+            "54": 3.48014,
+            "55": 3.02386,
+            "56": 4.34098,
+            "57": 2.76378,
+            "58": 3.52668,
+            "59": 2.98688,
+            "60": 3.29326,
+            "61": 2.79132,
+            "62": 20.76609,
+            "63": 10.52649,
+            "64": 4.66917,
+            "65": 4.40657,
+            "66": 4.07985,
+            "67": 2.93174,
+            "68": 2.54979,
+            "69": 3.25024,
+            "70": 2.92684,
+            "71": 13.6022,
+            "72": 6.85737,
+            "73": 3.00302,
+            "74": 2.27875,
+            "75": 2.94009,
+            "76": 3.77221,
+            "77": 2.98374,
+            "78": 3.97115,
+            "79": 20.85088,
+            "80": 2.66986,
+            "81": 3.20797,
+            "82": 3.31175,
+            "83": 2.6616,
+            "84": 2.451,
+            "85": 2.63993,
+            "86": 2.98263,
+            "87": 3.18667,
+            "88": 2.23769,
+            "89": 2.79162,
+            "90": 2.86908,
+            "91": 3.40123,
+            "92": 3.10172,
+            "93": 2.49279,
+            "94": 2.03585,
+            "95": 2.71749,
+            "96": 2.79808,
+            "97": 2.9601,
+            "98": 2.37092,
+            "99": 3.48981,
+            "100": 2.32584
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json
index 1a8739a347a..745ed36d8c9 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json
@@ -6,104 +6,104 @@
         "values": {
             "1": 10.95659,
             "2": 10.95273,
-            "3": 10.97293,
-            "4": 10.95528,
-            "5": 10.95882,
-            "6": 10.96034,
-            "7": 10.94954,
-            "8": 10.95612,
-            "9": 10.96238,
-            "10": 10.95507,
-            "11": 10.94969,
-            "12": 10.94911,
-            "13": 10.94368,
-            "14": 10.9398,
-            "15": 10.91315,
-            "16": 10.89286,
-            "17": 10.89445,
-            "18": 10.88436,
-            "19": 10.88801,
-            "20": 10.81677,
-            "21": 10.77929,
-            "22": 10.77943,
-            "23": 10.75103,
-            "24": 10.73693,
-            "25": 10.70864,
-            "26": 10.70217,
-            "27": 10.66685,
-            "28": 10.59001,
-            "29": 10.57517,
-            "30": 10.53988,
-            "31": 10.54916,
-            "32": 10.49405,
-            "33": 10.45568,
-            "34": 10.45441,
-            "35": 10.41606,
-            "36": 10.40471,
-            "37": 10.37388,
-            "38": 10.38042,
-            "39": 10.33633,
-            "40": 10.3374,
-            "41": 10.29239,
-            "42": 10.24602,
-            "43": 10.23775,
-            "44": 10.20504,
-            "45": 10.23936,
-            "46": 10.16799,
-            "47": 10.16516,
-            "48": 10.11263,
-            "49": 10.11908,
-            "50": 10.0989,
-            "51": 10.11379,
-            "52": 10.07064,
-            "53": 10.03526,
-            "54": 10.01886,
-            "55": 9.97027,
-            "56": 10.01809,
-            "57": 10.00228,
-            "58": 10.00891,
-            "59": 9.93698,
-            "60": 9.97704,
-            "61": 9.92067,
-            "62": 9.86061,
-            "63": 9.97239,
-            "64": 9.91741,
-            "65": 9.88037,
-            "66": 9.90572,
-            "67": 9.88892,
-            "68": 9.81415,
-            "69": 9.83834,
-            "70": 9.82934,
-            "71": 9.85471,
-            "72": 9.8455,
-            "73": 9.79839,
-            "74": 9.79383,
-            "75": 9.74221,
-            "76": 9.81225,
-            "77": 9.80864,
-            "78": 9.76159,
-            "79": 9.73676,
-            "80": 9.76097,
-            "81": 9.801,
-            "82": 9.72391,
-            "83": 9.66545,
-            "84": 9.62619,
-            "85": 9.59104,
-            "86": 9.73785,
-            "87": 9.72688,
-            "88": 9.73448,
-            "89": 9.63538,
-            "90": 9.62968,
-            "91": 9.67388,
-            "92": 9.63789,
-            "93": 9.53753,
-            "94": 9.65601,
-            "95": 9.62946,
-            "96": 9.63424,
-            "97": 9.54678,
-            "98": 9.59626,
-            "99": 9.64175,
-            "100": 9.53549
+            "3": 10.97295,
+            "4": 10.95542,
+            "5": 10.95849,
+            "6": 10.96054,
+            "7": 10.94934,
+            "8": 10.95616,
+            "9": 10.96262,
+            "10": 10.95503,
+            "11": 10.9491,
+            "12": 10.9493,
+            "13": 10.94313,
+            "14": 10.93907,
+            "15": 10.91333,
+            "16": 10.89331,
+            "17": 10.8942,
+            "18": 10.8846,
+            "19": 10.88746,
+            "20": 10.8169,
+            "21": 10.77918,
+            "22": 10.78015,
+            "23": 10.75133,
+            "24": 10.73664,
+            "25": 10.70889,
+            "26": 10.7017,
+            "27": 10.66701,
+            "28": 10.5899,
+            "29": 10.57557,
+            "30": 10.5395,
+            "31": 10.54921,
+            "32": 10.49385,
+            "33": 10.4557,
+            "34": 10.45468,
+            "35": 10.41609,
+            "36": 10.40507,
+            "37": 10.37438,
+            "38": 10.38081,
+            "39": 10.33703,
+            "40": 10.33794,
+            "41": 10.29217,
+            "42": 10.2462,
+            "43": 10.23797,
+            "44": 10.20525,
+            "45": 10.23977,
+            "46": 10.16827,
+            "47": 10.16534,
+            "48": 10.11329,
+            "49": 10.1194,
+            "50": 10.09951,
+            "51": 10.11426,
+            "52": 10.07085,
+            "53": 10.03538,
+            "54": 10.01908,
+            "55": 9.97069,
+            "56": 10.0184,
+            "57": 10.0028,
+            "58": 10.00966,
+            "59": 9.93792,
+            "60": 9.97762,
+            "61": 9.92129,
+            "62": 9.86095,
+            "63": 9.97263,
+            "64": 9.9177,
+            "65": 9.88092,
+            "66": 9.90582,
+            "67": 9.88943,
+            "68": 9.81505,
+            "69": 9.8386,
+            "70": 9.82966,
+            "71": 9.85498,
+            "72": 9.84522,
+            "73": 9.79852,
+            "74": 9.79437,
+            "75": 9.74264,
+            "76": 9.81212,
+            "77": 9.80875,
+            "78": 9.76202,
+            "79": 9.73761,
+            "80": 9.7607,
+            "81": 9.80103,
+            "82": 9.72386,
+            "83": 9.66588,
+            "84": 9.6266,
+            "85": 9.59092,
+            "86": 9.73755,
+            "87": 9.72718,
+            "88": 9.73452,
+            "89": 9.63555,
+            "90": 9.62955,
+            "91": 9.67368,
+            "92": 9.63796,
+            "93": 9.53744,
+            "94": 9.65635,
+            "95": 9.62912,
+            "96": 9.63405,
+            "97": 9.54616,
+            "98": 9.59568,
+            "99": 9.64142,
+            "100": 9.53537
         }
     },
     "num-zeros": {
@@ -111,106 +111,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 22985512.0,
-            "2": 22866852.0,
-            "3": 22718796.0,
-            "4": 22793114.0,
-            "5": 22800332.0,
-            "6": 22758732.0,
-            "7": 22889360.0,
-            "8": 22616950.0,
-            "9": 22770544.0,
-            "10": 22482352.0,
-            "11": 22768012.0,
-            "12": 22646638.0,
-            "13": 23376220.0,
-            "14": 23020932.0,
-            "15": 22728648.0,
-            "16": 22844252.0,
-            "17": 22956366.0,
-            "18": 23025410.0,
-            "19": 23121762.0,
-            "20": 22737712.0,
-            "21": 22939090.0,
-            "22": 22975288.0,
-            "23": 22636468.0,
-            "24": 22885640.0,
-            "25": 22646650.0,
-            "26": 23036326.0,
-            "27": 22820200.0,
-            "28": 23031704.0,
-            "29": 23007740.0,
-            "30": 22978056.0,
-            "31": 22931440.0,
-            "32": 22671854.0,
-            "33": 22753936.0,
-            "34": 23115324.0,
-            "35": 22764088.0,
-            "36": 22708302.0,
-            "37": 23140388.0,
-            "38": 22991024.0,
-            "39": 23018000.0,
-            "40": 22766666.0,
-            "41": 23101500.0,
-            "42": 22700204.0,
-            "43": 23019132.0,
-            "44": 22716592.0,
-            "45": 22868230.0,
-            "46": 22743320.0,
-            "47": 22871964.0,
-            "48": 22852496.0,
-            "49": 22908240.0,
-            "50": 22654376.0,
-            "51": 22713842.0,
-            "52": 22833092.0,
-            "53": 22987628.0,
-            "54": 22806964.0,
-            "55": 22950768.0,
-            "56": 22669906.0,
-            "57": 23234446.0,
-            "58": 22699600.0,
-            "59": 22862406.0,
-            "60": 23046738.0,
-            "61": 22688380.0,
-            "62": 22743124.0,
-            "63": 22644394.0,
-            "64": 23031802.0,
-            "65": 23243748.0,
-            "66": 22705416.0,
-            "67": 22986284.0,
-            "68": 22949944.0,
-            "69": 23193570.0,
-            "70": 22838384.0,
-            "71": 22750404.0,
-            "72": 23155280.0,
-            "73": 23168580.0,
-            "74": 22970434.0,
-            "75": 22903816.0,
-            "76": 22714100.0,
-            "77": 23011992.0,
-            "78": 23010412.0,
-            "79": 22845512.0,
-            "80": 22958324.0,
-            "81": 22850196.0,
-            "82": 22746284.0,
-            "83": 22741626.0,
-            "84": 23135720.0,
-            "85": 22945772.0,
-            "86": 23108174.0,
-            "87": 22369184.0,
-            "88": 22565120.0,
-            "89": 22738524.0,
-            "90": 22782056.0,
-            "91": 22941044.0,
-            "92": 22680736.0,
-            "93": 22647836.0,
-            "94": 23168884.0,
-            "95": 22702208.0,
-            "96": 22867360.0,
-            "97": 22852572.0,
-            "98": 22897152.0,
-            "99": 22645740.0,
-            "100": 23029680.0
+            "1": 22986408.0,
+            "2": 22867856.0,
+            "3": 22719884.0,
+            "4": 22794008.0,
+            "5": 22801154.0,
+            "6": 22759694.0,
+            "7": 22890364.0,
+            "8": 22617896.0,
+            "9": 22771528.0,
+            "10": 22483364.0,
+            "11": 22768948.0,
+            "12": 22647494.0,
+            "13": 23377220.0,
+            "14": 23021868.0,
+            "15": 22729498.0,
+            "16": 22845248.0,
+            "17": 22957354.0,
+            "18": 23026330.0,
+            "19": 23122828.0,
+            "20": 22738700.0,
+            "21": 22939944.0,
+            "22": 22976226.0,
+            "23": 22637450.0,
+            "24": 22886616.0,
+            "25": 22647556.0,
+            "26": 23037216.0,
+            "27": 22821184.0,
+            "28": 23032704.0,
+            "29": 23008756.0,
+            "30": 22979208.0,
+            "31": 22932472.0,
+            "32": 22672852.0,
+            "33": 22755016.0,
+            "34": 23116560.0,
+            "35": 22765664.0,
+            "36": 22709672.0,
+            "37": 23141764.0,
+            "38": 22992714.0,
+            "39": 23022566.0,
+            "40": 22767830.0,
+            "41": 23106848.0,
+            "42": 23749706.0,
+            "43": 24068728.0,
+            "44": 23766336.0,
+            "45": 22873326.0,
+            "46": 23792938.0,
+            "47": 23922180.0,
+            "48": 23902652.0,
+            "49": 23957900.0,
+            "50": 23704272.0,
+            "51": 23763660.0,
+            "52": 23883720.0,
+            "53": 24037434.0,
+            "54": 23856996.0,
+            "55": 24001224.0,
+            "56": 23720200.0,
+            "57": 24284806.0,
+            "58": 23749452.0,
+            "59": 23914064.0,
+            "60": 24098812.0,
+            "61": 23742456.0,
+            "62": 22745048.0,
+            "63": 23693956.0,
+            "64": 24081672.0,
+            "65": 24297776.0,
+            "66": 23755356.0,
+            "67": 24036968.0,
+            "68": 25048580.0,
+            "69": 24243380.0,
+            "70": 23891872.0,
+            "71": 24848952.0,
+            "72": 24205908.0,
+            "73": 24221136.0,
+            "74": 25068696.0,
+            "75": 23952946.0,
+            "76": 23764776.0,
+            "77": 25110290.0,
+            "78": 24060988.0,
+            "79": 23897268.0,
+            "80": 24012984.0,
+            "81": 23901444.0,
+            "82": 23796884.0,
+            "83": 22742696.0,
+            "84": 24186118.0,
+            "85": 23995964.0,
+            "86": 24178450.0,
+            "87": 23419336.0,
+            "88": 23615672.0,
+            "89": 23792564.0,
+            "90": 23832136.0,
+            "91": 23991288.0,
+            "92": 23731182.0,
+            "93": 22649058.0,
+            "94": 24219030.0,
+            "95": 22703490.0,
+            "96": 23918500.0,
+            "97": 23902440.0,
+            "98": 22898616.0,
+            "99": 23695624.0,
+            "100": 24079520.0
         }
     },
     "mem-allocated-bytes": {
@@ -218,106 +218,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 815727104.0,
-            "2": 793283584.0,
-            "3": 835202048.0,
-            "4": 807708672.0,
-            "5": 808458240.0,
-            "6": 804466688.0,
-            "7": 800862208.0,
-            "8": 808080384.0,
-            "9": 807653376.0,
-            "10": 800862208.0,
-            "11": 808208384.0,
-            "12": 807350272.0,
-            "13": 800862208.0,
-            "14": 808185856.0,
-            "15": 807186432.0,
-            "16": 800862208.0,
-            "17": 807448576.0,
-            "18": 807286784.0,
-            "19": 800862208.0,
-            "20": 807684096.0,
-            "21": 807890944.0,
-            "22": 800862208.0,
-            "23": 808510464.0,
-            "24": 808387584.0,
-            "25": 804035584.0,
-            "26": 800862208.0,
-            "27": 807825408.0,
-            "28": 806744064.0,
-            "29": 800862208.0,
-            "30": 807996416.0,
-            "31": 807682048.0,
-            "32": 803694592.0,
-            "33": 800862208.0,
-            "34": 807350272.0,
-            "35": 806928384.0,
-            "36": 800862208.0,
-            "37": 807727104.0,
-            "38": 807112704.0,
-            "39": 800862208.0,
-            "40": 807997440.0,
-            "41": 807677952.0,
-            "42": 803500032.0,
-            "43": 800862208.0,
-            "44": 807511040.0,
-            "45": 807274496.0,
-            "46": 800862208.0,
-            "47": 807894016.0,
-            "48": 807567360.0,
-            "49": 803500032.0,
-            "50": 800862208.0,
-            "51": 808185856.0,
-            "52": 804275200.0,
-            "53": 800862208.0,
-            "54": 807925760.0,
-            "55": 807542784.0,
-            "56": 803713024.0,
-            "57": 800862208.0,
-            "58": 807350272.0,
-            "59": 807393280.0,
-            "60": 800862208.0,
-            "61": 807858176.0,
-            "62": 807350272.0,
-            "63": 803500032.0,
-            "64": 800862208.0,
-            "65": 807711744.0,
-            "66": 807874560.0,
-            "67": 803500032.0,
-            "68": 807350272.0,
-            "69": 807612416.0,
-            "70": 804241408.0,
-            "71": 800862208.0,
-            "72": 807661568.0,
-            "73": 806744064.0,
-            "74": 800862208.0,
-            "75": 807350272.0,
-            "76": 807350272.0,
-            "77": 800862208.0,
-            "78": 800862208.0,
-            "79": 807976960.0,
-            "80": 808120320.0,
-            "81": 804230144.0,
-            "82": 800862208.0,
-            "83": 808173568.0,
-            "84": 807569408.0,
-            "85": 803500032.0,
-            "86": 800862208.0,
-            "87": 807350272.0,
-            "88": 807421952.0,
-            "89": 803729408.0,
-            "90": 800862208.0,
-            "91": 808009728.0,
-            "92": 808318976.0,
-            "93": 806807552.0,
-            "94": 800862208.0,
-            "95": 800862208.0,
-            "96": 808192000.0,
-            "97": 807350272.0,
-            "98": 806744064.0,
-            "99": 800862208.0,
-            "100": 807636992.0
+            "1": 783287296.0,
+            "2": 788025856.0,
+            "3": 780923392.0,
+            "4": 807817216.0,
+            "5": 808268288.0,
+            "6": 807809536.0,
+            "7": 805578240.0,
+            "8": 802091520.0,
+            "9": 805330432.0,
+            "10": 802091520.0,
+            "11": 802091520.0,
+            "12": 805496320.0,
+            "13": 802091520.0,
+            "14": 805298688.0,
+            "15": 807817216.0,
+            "16": 802091520.0,
+            "17": 804969984.0,
+            "18": 802091520.0,
+            "19": 804729344.0,
+            "20": 804729344.0,
+            "21": 802091520.0,
+            "22": 807724032.0,
+            "23": 802091520.0,
+            "24": 802091520.0,
+            "25": 808419840.0,
+            "26": 804729344.0,
+            "27": 802091520.0,
+            "28": 808036864.0,
+            "29": 808786944.0,
+            "30": 805179392.0,
+            "31": 804972032.0,
+            "32": 802091520.0,
+            "33": 807858688.0,
+            "34": 804729344.0,
+            "35": 802091520.0,
+            "36": 808122880.0,
+            "37": 804729344.0,
+            "38": 805076480.0,
+            "39": 802091520.0,
+            "40": 807704064.0,
+            "41": 804729344.0,
+            "42": 802091520.0,
+            "43": 807645696.0,
+            "44": 802091520.0,
+            "45": 807986688.0,
+            "46": 804811264.0,
+            "47": 802091520.0,
+            "48": 802091520.0,
+            "49": 808298496.0,
+            "50": 802091520.0,
+            "51": 804729344.0,
+            "52": 807973376.0,
+            "53": 807367168.0,
+            "54": 804729344.0,
+            "55": 802091520.0,
+            "56": 802091520.0,
+            "57": 807793152.0,
+            "58": 805179392.0,
+            "59": 805179392.0,
+            "60": 804938240.0,
+            "61": 805179392.0,
+            "62": 805579264.0,
+            "63": 808347136.0,
+            "64": 804729344.0,
+            "65": 804729344.0,
+            "66": 802091520.0,
+            "67": 802091520.0,
+            "68": 802091520.0,
+            "69": 802091520.0,
+            "70": 807817216.0,
+            "71": 805188096.0,
+            "72": 802091520.0,
+            "73": 802091520.0,
+            "74": 802091520.0,
+            "75": 807825920.0,
+            "76": 805512704.0,
+            "77": 802091520.0,
+            "78": 802091520.0,
+            "79": 808472576.0,
+            "80": 802091520.0,
+            "81": 802091520.0,
+            "82": 807986688.0,
+            "83": 805179392.0,
+            "84": 802091520.0,
+            "85": 802091520.0,
+            "86": 807367168.0,
+            "87": 804729344.0,
+            "88": 802091520.0,
+            "89": 808003072.0,
+            "90": 802091520.0,
+            "91": 804729344.0,
+            "92": 807367168.0,
+            "93": 802091520.0,
+            "94": 807367168.0,
+            "95": 802091520.0,
+            "96": 802091520.0,
+            "97": 808018944.0,
+            "98": 802091520.0,
+            "99": 807722496.0,
+            "100": 805179392.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -325,106 +325,106 @@
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": 964055552.0,
-            "2": 1186166272.0,
-            "3": 1196239360.0,
-            "4": 1196239360.0,
-            "5": 1196239360.0,
-            "6": 1196239360.0,
-            "7": 1196239360.0,
-            "8": 1196239360.0,
-            "9": 1196239360.0,
-            "10": 1196239360.0,
-            "11": 1196239360.0,
-            "12": 1196239360.0,
-            "13": 1196239360.0,
-            "14": 1196239360.0,
-            "15": 1196239360.0,
-            "16": 1196239360.0,
-            "17": 1196239360.0,
-            "18": 1196239360.0,
-            "19": 1196239360.0,
-            "20": 1196239360.0,
-            "21": 1196239360.0,
-            "22": 1196239360.0,
-            "23": 1196239360.0,
-            "24": 1196239360.0,
-            "25": 1196239360.0,
-            "26": 1196239360.0,
-            "27": 1196239360.0,
-            "28": 1196239360.0,
-            "29": 1196239360.0,
-            "30": 1196239360.0,
-            "31": 1196239360.0,
-            "32": 1196239360.0,
-            "33": 1196239360.0,
-            "34": 1196239360.0,
-            "35": 1196239360.0,
-            "36": 1196239360.0,
-            "37": 1196239360.0,
-            "38": 1196239360.0,
-            "39": 1196239360.0,
-            "40": 1196239360.0,
-            "41": 1196239360.0,
-            "42": 1196239360.0,
-            "43": 1196239360.0,
-            "44": 1196239360.0,
-            "45": 1196239360.0,
-            "46": 1196239360.0,
-            "47": 1196239360.0,
-            "48": 1196239360.0,
-            "49": 1196239360.0,
-            "50": 1196239360.0,
-            "51": 1196239360.0,
-            "52": 1196239360.0,
-            "53": 1196239360.0,
-            "54": 1196239360.0,
-            "55": 1196239360.0,
-            "56": 1196239360.0,
-            "57": 1196239360.0,
-            "58": 1196239360.0,
-            "59": 1196239360.0,
-            "60": 1196239360.0,
-            "61": 1196239360.0,
-            "62": 1196239360.0,
-            "63": 1196239360.0,
-            "64": 1196239360.0,
-            "65": 1196239360.0,
-            "66": 1196239360.0,
-            "67": 1196239360.0,
-            "68": 1196239360.0,
-            "69": 1196239360.0,
-            "70": 1196239360.0,
-            "71": 1196239360.0,
-            "72": 1196239360.0,
-            "73": 1196239360.0,
-            "74": 1196239360.0,
-            "75": 1196239360.0,
-            "76": 1196239360.0,
-            "77": 1196239360.0,
-            "78": 1196239360.0,
-            "79": 1196239360.0,
-            "80": 1196239360.0,
-            "81": 1196239360.0,
-            "82": 1196239360.0,
-            "83": 1196239360.0,
-            "84": 1196239360.0,
-            "85": 1196239360.0,
-            "86": 1196239360.0,
-            "87": 1196239360.0,
-            "88": 1196239360.0,
-            "89": 1196239360.0,
-            "90": 1196239360.0,
-            "91": 1196239360.0,
-            "92": 1196239360.0,
-            "93": 1196239360.0,
-            "94": 1196239360.0,
-            "95": 1196239360.0,
-            "96": 1196239360.0,
-            "97": 1196239360.0,
-            "98": 1196239360.0,
-            "99": 1196239360.0,
-            "100": 1196239360.0
+            "1": 931044352.0,
+            "2": 1150597632.0,
+            "3": 1150597632.0,
+            "4": 1150597632.0,
+            "5": 1150597632.0,
+            "6": 1150597632.0,
+            "7": 1150597632.0,
+            "8": 1150597632.0,
+            "9": 1150597632.0,
+            "10": 1150597632.0,
+            "11": 1150597632.0,
+            "12": 1150597632.0,
+            "13": 1150597632.0,
+            "14": 1150597632.0,
+            "15": 1150597632.0,
+            "16": 1150597632.0,
+            "17": 1150597632.0,
+            "18": 1150597632.0,
+            "19": 1150597632.0,
+            "20": 1150597632.0,
+            "21": 1150597632.0,
+            "22": 1150597632.0,
+            "23": 1150597632.0,
+            "24": 1150597632.0,
+            "25": 1150597632.0,
+            "26": 1150597632.0,
+            "27": 1150597632.0,
+            "28": 1150597632.0,
+            "29": 1150597632.0,
+            "30": 1150597632.0,
+            "31": 1150597632.0,
+            "32": 1150597632.0,
+            "33": 1150597632.0,
+            "34": 1150597632.0,
+            "35": 1150597632.0,
+            "36": 1150597632.0,
+            "37": 1150597632.0,
+            "38": 1150597632.0,
+            "39": 1150597632.0,
+            "40": 1150597632.0,
+            "41": 1150597632.0,
+            "42": 1150597632.0,
+            "43": 1150597632.0,
+            "44": 1150597632.0,
+            "45": 1150597632.0,
+            "46": 1150597632.0,
+            "47": 1150597632.0,
+            "48": 1150597632.0,
+            "49": 1150597632.0,
+            "50": 1150597632.0,
+            "51": 1150597632.0,
+            "52": 1150597632.0,
+            "53": 1150597632.0,
+            "54": 1150597632.0,
+            "55": 1150597632.0,
+            "56": 1150597632.0,
+            "57": 1150597632.0,
+            "58": 1150597632.0,
+            "59": 1150597632.0,
+            "60": 1150597632.0,
+            "61": 1150597632.0,
+            "62": 1150597632.0,
+            "63": 1150597632.0,
+            "64": 1150597632.0,
+            "65": 1150597632.0,
+            "66": 1150597632.0,
+            "67": 1150597632.0,
+            "68": 1150597632.0,
+            "69": 1150597632.0,
+            "70": 1150597632.0,
+            "71": 1150597632.0,
+            "72": 1150597632.0,
+            "73": 1150597632.0,
+            "74": 1150597632.0,
+            "75": 1150597632.0,
+            "76": 1150597632.0,
+            "77": 1150597632.0,
+            "78": 1150597632.0,
+            "79": 1150597632.0,
+            "80": 1150597632.0,
+            "81": 1150597632.0,
+            "82": 1150597632.0,
+            "83": 1150597632.0,
+            "84": 1150597632.0,
+            "85": 1150597632.0,
+            "86": 1150597632.0,
+            "87": 1150597632.0,
+            "88": 1150597632.0,
+            "89": 1150597632.0,
+            "90": 1150597632.0,
+            "91": 1150597632.0,
+            "92": 1150597632.0,
+            "93": 1150597632.0,
+            "94": 1150597632.0,
+            "95": 1150597632.0,
+            "96": 1150597632.0,
+            "97": 1150597632.0,
+            "98": 1150597632.0,
+            "99": 1150597632.0,
+            "100": 1150597632.0
         }
     },
     "mtp_1 loss": {
@@ -434,211 +434,210 @@
         "values": {
             "1": 10.91546,
             "2": 10.92323,
-            "3": 10.93384,
-            "4": 10.92739,
-            "5": 10.90724,
-            "6": 10.91817,
-            "7": 10.92486,
-            "8": 10.92528,
-            "9": 10.93457,
+            "3": 10.93381,
+            "4": 10.92692,
+            "5": 10.90737,
+            "6": 10.91868,
+            "7": 10.92452,
+            "8": 10.92547,
+            "9": 10.93431,
             "10": 10.9265,
-            "11": 10.91896,
-            "12": 10.91863,
-            "13": 10.92814,
-            "14": 10.91203,
-            "15": 10.92041,
-            "16": 10.92467,
-            "17": 10.92235,
-            "18": 10.90719,
-            "19": 10.91438,
-            "20": 10.90506,
-            "21": 10.91161,
-            "22": 10.89778,
-            "23": 10.90483,
-            "24": 10.88964,
-            "25": 10.89765,
-            "26": 10.88453,
-            "27": 10.89849,
-            "28": 10.89069,
+            "11": 10.91849,
+            "12": 10.91841,
+            "13": 10.92772,
+            "14": 10.91224,
+            "15": 10.92052,
+            "16": 10.92447,
+            "17": 10.92184,
+            "18": 10.90722,
+            "19": 10.91489,
+            "20": 10.90515,
+            "21": 10.91178,
+            "22": 10.89783,
+            "23": 10.90484,
+            "24": 10.89011,
+            "25": 10.89839,
+            "26": 10.88455,
+            "27": 10.89857,
+            "28": 10.89064,
             "29": 10.87558,
-            "30": 10.88029,
-            "31": 10.87277,
-            "32": 10.87869,
-            "33": 10.87024,
-            "34": 10.8682,
-            "35": 10.85932,
-            "36": 10.86187,
-            "37": 10.85499,
-            "38": 10.85713,
-            "39": 10.84881,
-            "40": 10.86288,
-            "41": 10.85314,
-            "42": 10.8473,
-            "43": 10.84576,
-            "44": 10.83821,
-            "45": 10.84904,
+            "30": 10.88023,
+            "31": 10.87305,
+            "32": 10.87883,
+            "33": 10.86963,
+            "34": 10.86778,
+            "35": 10.85942,
+            "36": 10.86134,
+            "37": 10.85492,
+            "38": 10.85701,
+            "39": 10.8491,
+            "40": 10.86294,
+            "41": 10.85349,
+            "42": 10.84748,
+            "43": 10.84529,
+            "44": 10.83797,
+            "45": 10.84927,
             "46": 10.8381,
-            "47": 10.83803,
-            "48": 10.831,
-            "49": 10.82926,
-            "50": 10.82229,
-            "51": 10.82153,
-            "52": 10.82122,
-            "53": 10.80664,
-            "54": 10.81103,
-            "55": 10.79443,
-            "56": 10.79963,
-            "57": 10.78961,
-            "58": 10.79824,
-            "59": 10.78095,
-            "60": 10.77503,
-            "61": 10.77627,
-            "62": 10.7614,
-            "63": 10.78392,
-            "64": 10.75466,
-            "65": 10.75002,
-            "66": 10.75702,
-            "67": 10.73504,
-            "68": 10.72878,
-            "69": 10.72595,
-            "70": 10.72543,
-            "71": 10.72482,
-            "72": 10.71955,
-            "73": 10.71178,
-            "74": 10.70369,
-            "75": 10.68547,
-            "76": 10.69478,
-            "77": 10.69055,
-            "78": 10.68188,
-            "79": 10.66968,
-            "80": 10.67688,
-            "81": 10.66904,
-            "82": 10.65016,
-            "83": 10.6267,
-            "84": 10.61015,
-            "85": 10.60262,
-            "86": 10.6432,
-            "87": 10.63641,
-            "88": 10.63101,
-            "89": 10.59551,
-            "90": 10.58424,
-            "91": 10.60768,
-            "92": 10.58305,
-            "93": 10.56222,
-            "94": 10.59342,
-            "95": 10.57615,
-            "96": 10.57208,
-            "97": 10.55416,
-            "98": 10.55921,
-            "99": 10.55818,
-            "100": 10.5284
+            "47": 10.83839,
+            "48": 10.83123,
+            "49": 10.82966,
+            "50": 10.82273,
+            "51": 10.82188,
+            "52": 10.82116,
+            "53": 10.80672,
+            "54": 10.81085,
+            "55": 10.79459,
+            "56": 10.79964,
+            "57": 10.78951,
+            "58": 10.79809,
+            "59": 10.78098,
+            "60": 10.77499,
+            "61": 10.77669,
+            "62": 10.76134,
+            "63": 10.78372,
+            "64": 10.75495,
+            "65": 10.75047,
+            "66": 10.75725,
+            "67": 10.73518,
+            "68": 10.72884,
+            "69": 10.72584,
+            "70": 10.72542,
+            "71": 10.72504,
+            "72": 10.71966,
+            "73": 10.7118,
+            "74": 10.70416,
+            "75": 10.68562,
+            "76": 10.69517,
+            "77": 10.69068,
+            "78": 10.682,
+            "79": 10.66958,
+            "80": 10.67703,
+            "81": 10.66911,
+            "82": 10.65004,
+            "83": 10.62667,
+            "84": 10.61036,
+            "85": 10.60259,
+            "86": 10.64333,
+            "87": 10.63617,
+            "88": 10.63091,
+            "89": 10.59516,
+            "90": 10.58423,
+            "91": 10.60752,
+            "92": 10.58284,
+            "93": 10.56206,
+            "94": 10.59328,
+            "95": 10.57611,
+            "96": 10.57209,
+            "97": 10.55386,
+            "98": 10.55913,
+            "99": 10.55814,
+            "100": 10.52809
         }
     },
     "iteration-time": {
-        "start_step": 1,
+        "start_step": 2,
         "end_step": 100,
         "step_interval": 1,
         "values": {
-            "1": "nan",
-            "2": 18.24385,
-            "3": 1.24992,
-            "4": 3.44822,
-            "5": 0.67115,
-            "6": 0.67992,
-            "7": 0.67235,
-            "8": 0.67169,
-            "9": 0.67242,
-            "10": 0.66924,
-            "11": 0.67194,
-            "12": 0.66696,
-            "13": 0.66783,
-            "14": 0.66806,
-            "15": 0.66792,
-            "16": 0.66722,
-            "17": 0.66884,
-            "18": 0.66891,
-            "19": 0.67057,
-            "20": 0.67024,
-            "21": 0.67476,
-            "22": 0.6704,
-            "23": 0.66892,
-            "24": 0.67043,
-            "25": 0.67258,
-            "26": 0.67099,
-            "27": 0.67203,
-            "28": 0.67141,
-            "29": 0.67162,
-            "30": 0.67618,
-            "31": 0.67022,
-            "32": 0.68537,
-            "33": 0.67019,
-            "34": 0.66964,
-            "35": 0.67288,
-            "36": 0.66938,
-            "37": 0.67603,
-            "38": 0.66977,
-            "39": 0.67445,
-            "40": 0.67455,
-            "41": 0.6717,
-            "42": 0.67202,
-            "43": 0.67216,
-            "44": 0.67433,
-            "45": 0.67073,
-            "46": 0.6702,
-            "47": 0.67187,
-            "48": 0.67865,
-            "49": 0.67059,
-            "50": 0.67768,
-            "51": 0.7684,
-            "52": 0.67476,
-            "53": 0.67246,
-            "54": 0.67426,
-            "55": 0.67523,
-            "56": 0.67431,
-            "57": 0.67379,
-            "58": 0.67592,
-            "59": 0.67389,
-            "60": 0.67679,
-            "61": 0.67409,
-            "62": 0.67265,
-            "63": 0.67543,
-            "64": 0.67577,
-            "65": 0.6745,
-            "66": 0.67687,
-            "67": 0.67327,
-            "68": 0.67244,
-            "69": 0.67241,
-            "70": 0.67191,
-            "71": 0.67044,
-            "72": 0.67049,
-            "73": 0.67597,
-            "74": 0.67288,
-            "75": 0.67123,
-            "76": 0.67032,
-            "77": 0.66955,
-            "78": 0.68133,
-            "79": 0.67997,
-            "80": 0.68011,
-            "81": 0.68168,
-            "82": 0.68012,
-            "83": 0.68054,
-            "84": 0.67091,
-            "85": 0.67421,
-            "86": 0.67093,
-            "87": 0.68073,
-            "88": 0.67264,
-            "89": 0.67707,
-            "90": 0.6819,
-            "91": 0.67945,
-            "92": 0.6829,
-            "93": 0.68217,
-            "94": 0.68005,
-            "95": 0.68097,
-            "96": 0.68088,
-            "97": 0.68296,
-            "98": 0.68201,
-            "99": 0.67603,
-            "100": 0.67638
+            "2": 29.93393,
+            "3": 1.98973,
+            "4": 5.35658,
+            "5": 1.20635,
+            "6": 1.19444,
+            "7": 1.18934,
+            "8": 1.16351,
+            "9": 1.13996,
+            "10": 1.15149,
+            "11": 1.13288,
+            "12": 1.13898,
+            "13": 1.14688,
+            "14": 1.14899,
+            "15": 1.14352,
+            "16": 1.14048,
+            "17": 1.13989,
+            "18": 1.15969,
+            "19": 1.15356,
+            "20": 1.1122,
+            "21": 1.1575,
+            "22": 1.15642,
+            "23": 1.17142,
+            "24": 1.20778,
+            "25": 1.24095,
+            "26": 1.17461,
+            "27": 1.16007,
+            "28": 1.19187,
+            "29": 1.1912,
+            "30": 1.19282,
+            "31": 1.19842,
+            "32": 1.17999,
+            "33": 1.16481,
+            "34": 1.17327,
+            "35": 1.17844,
+            "36": 1.16523,
+            "37": 1.17743,
+            "38": 1.20861,
+            "39": 1.17587,
+            "40": 1.16604,
+            "41": 1.17477,
+            "42": 1.16262,
+            "43": 1.16073,
+            "44": 1.15779,
+            "45": 1.18144,
+            "46": 1.16745,
+            "47": 1.16008,
+            "48": 1.20221,
+            "49": 1.18658,
+            "50": 1.16544,
+            "51": 1.26104,
+            "52": 1.21112,
+            "53": 1.20935,
+            "54": 1.2683,
+            "55": 1.07576,
+            "56": 1.08833,
+            "57": 1.08595,
+            "58": 1.07682,
+            "59": 1.10548,
+            "60": 1.10443,
+            "61": 1.12092,
+            "62": 1.11365,
+            "63": 1.10675,
+            "64": 1.09502,
+            "65": 1.09559,
+            "66": 1.09977,
+            "67": 1.09642,
+            "68": 1.08417,
+            "69": 1.09903,
+            "70": 1.08224,
+            "71": 1.08075,
+            "72": 1.08284,
+            "73": 1.08425,
+            "74": 1.09211,
+            "75": 1.0851,
+            "76": 1.0861,
+            "77": 1.07301,
+            "78": 1.08354,
+            "79": 1.07668,
+            "80": 1.07698,
+            "81": 1.07855,
+            "82": 1.07964,
+            "83": 1.07692,
+            "84": 1.08014,
+            "85": 1.08469,
+            "86": 1.08449,
+            "87": 1.07474,
+            "88": 1.06482,
+            "89": 1.05622,
+            "90": 1.04827,
+            "91": 1.04974,
+            "92": 1.04658,
+            "93": 1.04466,
+            "94": 1.03782,
+            "95": 1.04608,
+            "96": 1.03914,
+            "97": 1.04721,
+            "98": 1.05712,
+            "99": 1.05243,
+            "100": 1.08418
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py b/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py
index 6461dbee751..fe56317f39a 100644
--- a/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py
+++ b/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py
@@ -8,9 +8,15 @@
 
 from megatron.core import parallel_state
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.extensions.transformer_engine import (
+    TEColumnParallelGroupedLinear,
+    TEColumnParallelLinear,
+)
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_layer_with_transformer_engine_submodules,
 )
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
 from megatron.core.transformer.spec_utils import get_submodules
@@ -104,6 +110,121 @@ def get_moe_model_and_buffers(
     )
 
 
+def _build_expert_linear(implementation: str, config: TransformerConfig) -> torch.nn.Module:
+    common_kwargs = {
+        "input_size": config.hidden_size,
+        "output_size": config.ffn_hidden_size,
+        "config": config,
+        "init_method": config.init_method,
+        "bias": False,
+        "skip_bias_add": False,
+        "is_expert": True,
+    }
+    if implementation == "native":
+        return ColumnParallelLinear(
+            **common_kwargs,
+            gather_output=False,
+            tp_group=parallel_state.get_expert_tensor_parallel_group(),
+        )
+    if implementation == "transformer_engine":
+        return TEColumnParallelLinear(
+            **common_kwargs,
+            gather_output=False,
+            tp_group=parallel_state.get_expert_tensor_parallel_group(),
+        )
+    if implementation == "transformer_engine_grouped":
+        return TEColumnParallelGroupedLinear(
+            num_gemms=config.num_moe_experts,
+            **common_kwargs,
+            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+        )
+    raise AssertionError(f"Unsupported implementation: {implementation}")
+
+
+@pytest.mark.parametrize(
+    ("tensor_model_parallel_size", "expert_tensor_parallel_size"), [(2, 1), (1, 2)]
+)
+@pytest.mark.parametrize(
+    "implementation", ["native", "transformer_engine", "transformer_engine_grouped"]
+)
+def test_expert_grad_sync_uses_expert_data_parallel_group(
+    implementation: str, tensor_model_parallel_size: int, expert_tensor_parallel_size: int
+):
+    """Expert gradients must not be reduced over ordinary DP when ETP differs from TP."""
+    if Utils.world_size < 4 or Utils.world_size % 4 != 0:
+        pytest.skip("Test requires a world size divisible by four")
+    if Utils.world_size > 16:
+        pytest.skip("Rank-encoded gradients are intended for small unit-test world sizes")
+
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=tensor_model_parallel_size,
+        expert_model_parallel_size=1,
+        expert_tensor_parallel_size=expert_tensor_parallel_size,
+    )
+    try:
+        # Per-token loss leaves DDP's pre-collective gradient scaling at one.
+        config = TransformerConfig(
+            num_layers=1,
+            hidden_size=8,
+            num_attention_heads=4,
+            ffn_hidden_size=16,
+            num_moe_experts=2,
+            moe_ffn_hidden_size=16,
+            moe_router_topk=2,
+            tensor_model_parallel_size=tensor_model_parallel_size,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=expert_tensor_parallel_size,
+            calculate_per_token_loss=True,
+            gradient_accumulation_fusion=False,
+            perform_initialization=False,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            add_bias_linear=False,
+        )
+        module = _build_expert_linear(implementation, config).cuda()
+        model = DistributedDataParallel(
+            config,
+            ddp_config=DistributedDataParallelConfig(
+                grad_reduce_in_fp32=True,
+                overlap_grad_reduce=False,
+                use_distributed_optimizer=False,
+                average_in_collective=False,
+            ),
+            module=module,
+        )
+
+        expert_dp_group = parallel_state.get_expert_data_parallel_group(
+            partial_expert_data_parallel=True
+        )
+        ordinary_dp_group = parallel_state.get_data_parallel_group(
+            with_context_parallel=True, partial_data_parallel=True
+        )
+        expert_dp_ranks = torch.distributed.get_process_group_ranks(expert_dp_group)
+        ordinary_dp_ranks = torch.distributed.get_process_group_ranks(ordinary_dp_group)
+        assert expert_dp_ranks != ordinary_dp_ranks
+
+        # Powers of two give every rank set a distinct sum, exposing the wrong collective group.
+        rank_value = float(2 ** torch.distributed.get_rank())
+        expected_value = float(sum(2**rank for rank in expert_dp_ranks))
+        ordinary_dp_value = float(sum(2**rank for rank in ordinary_dp_ranks))
+        assert expected_value != ordinary_dp_value
+
+        for param in model.parameters():
+            param.main_grad.fill_(rank_value)
+        model.finish_grad_sync()
+
+        for param in model.parameters():
+            torch.testing.assert_close(
+                param.main_grad, torch.full_like(param.main_grad, expected_value), rtol=0, atol=0
+            )
+
+        assert not model.buffers
+        assert len(model.expert_parallel_buffers) == 1
+        assert all(param.allreduce is False for param in model.parameters())
+    finally:
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.parametrize("use_distributed_optimizer", [False, True])
 @pytest.mark.parametrize("overlap_grad_reduce", [False, True])
 @pytest.mark.parametrize("average_in_collective", [False, True])
diff --git a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py
index 44d6fa21178..a76746d7674 100644
--- a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py
+++ b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py
@@ -8,6 +8,7 @@
     RowParallelLinear,
     VocabParallelEmbedding,
     copy_tensor_model_parallel_attributes,
+    param_is_not_tensor_parallel_duplicate,
 )
 from megatron.core.transformer.transformer_config import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
@@ -100,3 +101,60 @@ def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes():
 
     assert destination.is_qkv is True
     assert destination.qkv_split_shapes == source.qkv_split_shapes
+
+
+def test_non_allreduce_param_uses_expert_tp_group_for_duplicate_filter():
+    class RankGroup:
+        def __init__(self, rank):
+            self._rank = rank
+
+        def rank(self):
+            return self._rank
+
+    param = torch.empty(1)
+    regular_tp_group = RankGroup(rank=1)
+    expert_tp_group = RankGroup(rank=0)
+    assert not param_is_not_tensor_parallel_duplicate(param, regular_tp_group)
+
+    param.allreduce = False
+    assert param_is_not_tensor_parallel_duplicate(
+        param, tp_group=regular_tp_group, expert_tp_group=expert_tp_group
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_expert_linear_parameters_use_expert_topology_metadata():
+    Utils.initialize_model_parallel(tensor_model_parallel_size=2, expert_tensor_parallel_size=1)
+    cfg = TransformerConfig(
+        num_layers=1,
+        hidden_size=8,
+        num_attention_heads=4,
+        tensor_model_parallel_size=2,
+        expert_tensor_parallel_size=1,
+        use_cpu_initialization=True,
+        perform_initialization=False,
+    )
+
+    column = ColumnParallelLinear(
+        input_size=8,
+        output_size=8,
+        init_method=cfg.init_method,
+        bias=True,
+        config=cfg,
+        gather_output=False,
+        skip_bias_add=False,
+        is_expert=True,
+    )
+    row = RowParallelLinear(
+        input_size=8,
+        output_size=8,
+        init_method=cfg.init_method,
+        bias=True,
+        input_is_parallel=True,
+        config=cfg,
+        skip_bias_add=False,
+        is_expert=True,
+    )
+
+    for param in (*column.parameters(), *row.parameters()):
+        assert param.allreduce is False
diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py
index 94613b7096c..d822f9bbf3f 100644
--- a/tests/unit_tests/test_optimizer.py
+++ b/tests/unit_tests/test_optimizer.py
@@ -27,6 +27,7 @@
     get_standard_config_overrides,
 )
 from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
+from megatron.core.optimizer.optimizer import copy_optimizer_param_metadata
 from megatron.core.optimizer_param_scheduler import ParamGroupOverride
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
@@ -80,6 +81,16 @@ def forward(self, x):
         return x
 
 
+def test_copy_optimizer_param_metadata_preserves_allreduce():
+    source = torch.empty(1)
+    destination = torch.empty_like(source)
+    source.allreduce = False
+
+    copy_optimizer_param_metadata(destination, source)
+
+    assert destination.allreduce is False
+
+
 @patch('torch.distributed.get_world_size', return_value=1)
 @patch(
     'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj)
diff --git a/tests/unit_tests/training/test_param_norm.py b/tests/unit_tests/training/test_param_norm.py
new file mode 100644
index 00000000000..d58eef225d3
--- /dev/null
+++ b/tests/unit_tests/training/test_param_norm.py
@@ -0,0 +1,202 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import math
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training.utils import common_utils
+from tests.unit_tests.test_utilities import Utils
+
+
+def _build_tiny_moe_gpt(
+    tensor_parallel_size: int,
+    expert_parallel_size: int,
+    expert_tensor_parallel_size: int,
+    bf16: bool = False,
+) -> GPTModel:
+    config = TransformerConfig(
+        num_layers=1,
+        hidden_size=8,
+        num_attention_heads=4,
+        ffn_hidden_size=16,
+        num_moe_experts=2,
+        moe_ffn_hidden_size=16,
+        moe_shared_expert_intermediate_size=16,
+        moe_router_topk=1,
+        moe_router_pre_softmax=True,
+        tensor_model_parallel_size=tensor_parallel_size,
+        expert_model_parallel_size=expert_parallel_size,
+        expert_tensor_parallel_size=expert_tensor_parallel_size,
+        sequence_parallel=tensor_parallel_size > 1,
+        use_cpu_initialization=True,
+        add_bias_linear=False,
+        normalization="RMSNorm",
+        moe_grouped_gemm=True,
+        bf16=bf16,
+        params_dtype=torch.bfloat16 if bf16 else torch.float32,
+    )
+    model = GPTModel(
+        config=config,
+        transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(
+            num_experts=config.num_moe_experts, moe_grouped_gemm=True
+        ),
+        vocab_size=16,
+        max_sequence_length=8,
+        position_embedding_type="rope",
+    )
+    assert any(".shared_experts." in name for name, _ in model.named_parameters())
+    return model.cuda()
+
+
+def _fill_parameters_with_ones(model: GPTModel) -> None:
+    with torch.no_grad():
+        for param in model.parameters():
+            param.fill_(1.0)
+
+
+@pytest.mark.parametrize(
+    ("tensor_parallel_size", "expert_parallel_size", "expert_tensor_parallel_size"),
+    ((2, 2, 1), (2, 1, 2), (4, 1, 2), (2, 1, 4)),
+    ids=("expert-parallel", "expert-tensor-parallel", "tp-larger-than-etp", "etp-larger-than-tp"),
+)
+def test_moe_param_norm_counts_each_logical_parameter_once(
+    monkeypatch,
+    tensor_parallel_size: int,
+    expert_parallel_size: int,
+    expert_tensor_parallel_size: int,
+):
+    """Parameter norm should be invariant to expert and expert-tensor parallelism."""
+    if Utils.world_size < 4 or Utils.world_size % 4 != 0:
+        pytest.skip("test requires a world size divisible by four")
+
+    monkeypatch.setattr(
+        common_utils, "get_args", lambda: SimpleNamespace(use_megatron_fsdp=False, bf16=False)
+    )
+
+    try:
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=1,
+        )
+        reference_model = _build_tiny_moe_gpt(
+            tensor_parallel_size=1, expert_parallel_size=1, expert_tensor_parallel_size=1
+        )
+        _fill_parameters_with_ones(reference_model)
+        expected_numel = sum(param.numel() for param in reference_model.parameters())
+        expected_norm = math.sqrt(expected_numel)
+        reference_norm = common_utils.calc_params_l2_norm(reference_model)
+
+        assert reference_norm == pytest.approx(expected_norm)
+        del reference_model
+
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tensor_parallel_size,
+            expert_model_parallel_size=expert_parallel_size,
+            expert_tensor_parallel_size=expert_tensor_parallel_size,
+        )
+        distributed_model = _build_tiny_moe_gpt(
+            tensor_parallel_size=tensor_parallel_size,
+            expert_parallel_size=expert_parallel_size,
+            expert_tensor_parallel_size=expert_tensor_parallel_size,
+        )
+        _fill_parameters_with_ones(distributed_model)
+
+        actual_norm = common_utils.calc_params_l2_norm(distributed_model)
+
+        assert actual_norm == pytest.approx(expected_norm)
+    finally:
+        Utils.destroy_model_parallel()
+
+
+@pytest.mark.parametrize("use_distributed_optimizer", (False, True), ids=("optimizer", "distopt"))
+@pytest.mark.parametrize(
+    ("tensor_parallel_size", "expert_parallel_size", "expert_tensor_parallel_size"),
+    ((2, 2, 1), (2, 1, 2), (4, 1, 2), (2, 1, 4)),
+    ids=("expert-parallel", "expert-tensor-parallel", "tp-larger-than-etp", "etp-larger-than-tp"),
+)
+def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once(
+    tensor_parallel_size: int,
+    expert_parallel_size: int,
+    expert_tensor_parallel_size: int,
+    use_distributed_optimizer: bool,
+):
+    """Gradient clipping should use each logical parameter's gradient exactly once."""
+    if Utils.world_size < 4 or Utils.world_size % 4 != 0:
+        pytest.skip("test requires a world size divisible by four")
+
+    try:
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=1,
+        )
+        reference_model = _build_tiny_moe_gpt(
+            tensor_parallel_size=1, expert_parallel_size=1, expert_tensor_parallel_size=1, bf16=True
+        )
+        expected_numel = sum(param.numel() for param in reference_model.parameters())
+        expected_norm = math.sqrt(expected_numel)
+        del reference_model
+
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tensor_parallel_size,
+            expert_model_parallel_size=expert_parallel_size,
+            expert_tensor_parallel_size=expert_tensor_parallel_size,
+        )
+        model = _build_tiny_moe_gpt(
+            tensor_parallel_size=tensor_parallel_size,
+            expert_parallel_size=expert_parallel_size,
+            expert_tensor_parallel_size=expert_tensor_parallel_size,
+            bf16=True,
+        )
+        ddp_config = DistributedDataParallelConfig(
+            grad_reduce_in_fp32=True, use_distributed_optimizer=use_distributed_optimizer
+        )
+        model = DistributedDataParallel(model.config, ddp_config, model)
+
+        max_norm = expected_norm / 2.0
+        optimizer = get_megatron_optimizer(
+            OptimizerConfig(
+                optimizer="adam",
+                lr=0.0,
+                bf16=True,
+                clip_grad=max_norm,
+                use_distributed_optimizer=use_distributed_optimizer,
+            ),
+            [model],
+        )
+
+        for param in model.parameters():
+            assert hasattr(param, "main_grad")
+            param.main_grad.fill_(1.0)
+
+        update_successful, actual_norm, _ = optimizer.step()
+
+        assert update_successful
+        actual_norm_value = (
+            actual_norm.item() if isinstance(actual_norm, torch.Tensor) else actual_norm
+        )
+        assert actual_norm_value == pytest.approx(expected_norm)
+
+        expected_clip_coefficient = max_norm / (expected_norm + 1.0e-6)
+        grads_checked = 0
+        for param in optimizer.get_parameters():
+            if param.grad is None:
+                continue
+            torch.testing.assert_close(
+                param.grad,
+                torch.full_like(param.grad, expected_clip_coefficient),
+                rtol=1.0e-5,
+                atol=1.0e-6,
+            )
+            grads_checked += 1
+        assert grads_checked > 0
+    finally:
+        Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py b/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py
index d0a411091f7..84ca109a15e 100644
--- a/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py
+++ b/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py
@@ -36,6 +36,44 @@ def _empty_load_args():
     return {}, True, [], [], []
 
 
+@pytest.mark.parametrize(("parallel_mode", "partition_dim"), (("column", 0), ("row", 1)))
+def test_expert_parameter_attributes_use_expert_topology(parallel_mode, partition_dim):
+    module = torch.nn.Module()
+    module.register_parameter("weight0", torch.nn.Parameter(torch.empty(4, 4)))
+    module.register_parameter("bias0", torch.nn.Parameter(torch.empty(4)))
+
+    te_ext._set_expert_parameter_attributes(
+        module, parallel_mode=parallel_mode, use_expert_pgs=True
+    )
+
+    assert module.weight0.allreduce is False
+    assert module.weight0.tensor_model_parallel is True
+    assert module.weight0.partition_dim == partition_dim
+    assert module.bias0.allreduce is False
+    assert module.bias0.tensor_model_parallel is (parallel_mode == "column")
+
+
+@pytest.mark.parametrize(
+    ("name", "is_partitioned"),
+    (
+        ("weight", True),
+        ("weight12", True),
+        ("bias", True),
+        ("bias12", True),
+        ("weight_scale", False),
+        ("bias_extra", False),
+    ),
+)
+def test_expert_parameter_attributes_match_parameter_names(name, is_partitioned):
+    module = torch.nn.Module()
+    module.register_parameter(name, torch.nn.Parameter(torch.empty(4)))
+
+    te_ext._set_expert_parameter_attributes(module, parallel_mode="column", use_expert_pgs=True)
+
+    param = module.get_parameter(name)
+    assert getattr(param, "tensor_model_parallel", False) is is_partitioned
+
+
 def test_split_empty_extra_state_for_stateless_recipe():
     module = _grouped_linear_stub(num_gemms=2)
     module.fp8_meta = {"fp8_checkpoint": True}