From 01368760764adc65efb3d4c94f47cf125f0e9af8 Mon Sep 17 00:00:00 2001
From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
Date: Thu, 22 Jan 2026 18:50:03 -0800
Subject: [PATCH 01/79] Supporting inference when called within an asyncio loop
 (#2816)

---
 .../inference/contexts/dynamic_context.py     |  7 ++++++
 .../core/inference/engines/dynamic_engine.py  | 23 +++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index b4e50ff6c8c..c6b30f47f78 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1616,6 +1616,13 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int]
         metadata_types = req.get_metadata_types()
         for m, m_type in zip(metadata, metadata_types):
             label, _, _ = m_type
+            if not isinstance(m, torch.Tensor):
+                m = torch.as_tensor(
+                    m,
+                    device=self.request_metadata[label].device,
+                    dtype=self.request_metadata[label].dtype,
+                )
+
             self.request_metadata[label][current_id] = m
 
         # Handle length and block assignments.
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 906b46efed5..c56f91bbbe9 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import asyncio
+import concurrent.futures
 import logging
 import multiprocessing
 import os
@@ -1368,11 +1369,29 @@ async def async_step(
         # Keep for compatibility with current test suite.
         return ret
 
+    def _run_coroutine_sync(self, coro):
+        """Run a coroutine synchronously, handling the case when already in an event loop.
+
+        This method safely runs an async coroutine from synchronous code, even when
+        called from within an already running event loop (e.g., when used with async
+        frameworks like pytriton).
+        """
+        try:
+            # Check if there's already a running event loop
+            asyncio.get_running_loop()
+            # We're inside a running loop - run in a separate thread
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(asyncio.run, coro)
+                return future.result()
+        except RuntimeError:
+            # No running loop - safe to use run_until_complete
+            return self._loop.run_until_complete(coro)
+
     def step_modern(
         self,
     ) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest], float]:
         """Synchronous wrapper for `self.async_step`."""
-        return self._loop.run_until_complete(self.async_step())
+        return self._run_coroutine_sync(self.async_step())
 
     def step_legacy(
         self, sampling_params: SamplingParams
@@ -1383,7 +1402,7 @@ def step_legacy(
             "0.16. Please use `step_modern()` going forward, which will eventually "
             "be renamed to `step()`."
         )
-        result = self._loop.run_until_complete(self.async_step())
+        result = self._run_coroutine_sync(self.async_step())
         active_requests = [self.get_request(i) for i in result["active_request_ids"]]
         finished_requests = [r.merge() for r in result["finished_request_records"]]
         return active_requests, finished_requests, result["step_time"]

From 03e09154eeffd6030bb846d6c778a80901db057b Mon Sep 17 00:00:00 2001
From: JavaZero <71128095+JavaZeroo@users.noreply.github.com>
Date: Fri, 23 Jan 2026 10:53:41 +0800
Subject: [PATCH 02/79] Update type hints and doc strings for moe_utils.py
 (#2821)

Co-authored-by: Xin Yao <yaox12@outlook.com>
Co-authored-by: Aaron Gokaslan <aaronGokaslan@gmail.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
---
 megatron/core/transformer/moe/moe_utils.py | 285 ++++++++++++++++-----
 1 file changed, 220 insertions(+), 65 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 3bef7d46924..5fdeda23dea 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -3,7 +3,7 @@
 import functools
 import math
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 
@@ -38,7 +38,7 @@
 
 
 # MOE logging
-_MOE_LAYER_WISE_LOGGING_TRACKER = {}
+_MOE_LAYER_WISE_LOGGING_TRACKER: dict = {}
 
 
 def switch_load_balancing_loss_func(
@@ -49,7 +49,7 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
-):
+) -> torch.Tensor:
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
     and Global Load Balancing Loss(https://arxiv.org/abs/2501.11873) for details.
@@ -99,6 +99,8 @@ def switch_load_balancing_loss_func(
         topk (int): The number of experts selected for each token.
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
+        fused (bool): Whether to use the fused version of the auxiliary loss.
+
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
@@ -121,12 +123,13 @@ def switch_load_balancing_loss_func(
     return aux_loss
 
 
-def z_loss_func(logits, z_loss_coeff):
+def z_loss_func(logits: torch.Tensor, z_loss_coeff: float) -> torch.Tensor:
     """Encourages the router's logits to remain small to enhance stability.
     Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
     Args:
         logits (torch.Tensor): The logits of the router.
+        z_loss_coeff (float): The coefficient for the z-loss.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
@@ -136,8 +139,16 @@ def z_loss_func(logits, z_loss_coeff):
     return z_loss
 
 
-def sinkhorn(cost: torch.Tensor, tol: float = 0.0001):
-    """Sinkhorn based MoE routing function"""
+def sinkhorn(cost: torch.Tensor, tol: float = 0.0001) -> torch.Tensor:
+    """Sinkhorn based MoE routing function.
+
+    Args:
+        cost (torch.Tensor): The cost tensor.
+        tol (float): The tolerance for the Sinkhorn algorithm.
+
+    Returns:
+        torch.Tensor: The routing probabilities.
+    """
     cost = torch.exp(cost)
     d0 = torch.ones(cost.size(0), device=cost.device, dtype=cost.dtype)
     d1 = torch.ones(cost.size(1), device=cost.device, dtype=cost.dtype)
@@ -153,7 +164,9 @@ def sinkhorn(cost: torch.Tensor, tol: float = 0.0001):
     return d1 * cost * d0.unsqueeze(1)
 
 
-def get_capacity(num_tokens: int, num_experts: int, capacity_factor: float, min_capacity=None):
+def get_capacity(
+    num_tokens: int, num_experts: int, capacity_factor: float, min_capacity: Optional[int] = None
+) -> int:
     """
     Calculate the capacity of each expert.
 
@@ -164,7 +177,7 @@ def get_capacity(num_tokens: int, num_experts: int, capacity_factor: float, min_
         min_capacity (int, optional): Minimum capacity. Defaults to None.
 
     Returns:
-        Tensor: Capacity of each expert.
+        int: Capacity of each expert.
     """
     capacity = math.ceil((num_tokens / num_experts) * capacity_factor)
     if min_capacity is not None and capacity < min_capacity:
@@ -178,7 +191,7 @@ class MoEAuxLossAutoScaler(torch.autograd.Function):
     main_loss_backward_scale: Optional[torch.Tensor] = None
 
     @staticmethod
-    def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor):
+    def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor:
         """Preserve the aux_loss by storing it in the context to avoid garbage collection.
 
         Args:
@@ -192,7 +205,7 @@ def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor):
         return output
 
     @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """Compute and scale the gradient for auxiliary loss..
 
         Args:
@@ -212,7 +225,7 @@ def backward(ctx, grad_output: torch.Tensor):
         return grad_output, scaled_aux_loss_grad
 
     @staticmethod
-    def set_loss_scale(scale: torch.Tensor):
+    def set_loss_scale(scale: torch.Tensor) -> None:
         """set the scale of the aux loss.
 
         Args:
@@ -226,13 +239,13 @@ def set_loss_scale(scale: torch.Tensor):
 
 
 def permute(
-    tokens,
-    routing_map,
+    tokens: torch.Tensor,
+    routing_map: torch.Tensor,
     probs: Optional[torch.Tensor] = None,
     num_out_tokens: Optional[int] = None,
     fused: bool = False,
     drop_and_pad: bool = False,
-):
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
     """Permute the tokens and probs based on the mask.
     Tokens with the same designated expert will be grouped together.
     The shape of mask is [tokens, num_experts], it indicates which experts were selected
@@ -252,6 +265,10 @@ def permute(
                                        and pads the number of tokens to the expert capacity.
                                        If set to true, routing_map has a fixed number of non-zeros
                                        in each column.
+
+    Returns:
+        Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+            The permuted tokens, permuted probs, and sorted indices.
     """
     if fused and probs is None:
         if not HAVE_TE or fused_permute is None:
@@ -320,7 +337,7 @@ def unpermute(
     routing_map: Optional[torch.Tensor] = None,
     fused: bool = False,
     drop_and_pad: bool = False,
-):
+) -> torch.Tensor:
     """
     Restore the original order of tokens after permutation. If probs are provided, it
     will also apply them to the tokens before restoring the order.
@@ -408,8 +425,20 @@ def sort_chunks_by_idxs(
     sorted_idxs: torch.Tensor,
     probs: Optional[torch.Tensor] = None,
     fused: bool = False,
-):
-    """Split and sort the input tensor based on the split_sizes and sorted indices."""
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Split and sort the input tensor based on the split_sizes and sorted indices.
+
+    Args:
+        input (torch.Tensor): The input tensor.
+        split_sizes (torch.Tensor): The split sizes.
+        sorted_idxs (torch.Tensor): The sorted indices.
+        probs (torch.Tensor, optional): The probs tensor. Defaults to None.
+        fused (bool, optional): Whether to use the fused version of the sort_chunks_by_idxs
+                                function. Defaults to False.
+
+    Returns:
+        Tuple[torch.Tensor, Optional[torch.Tensor]]: The sorted output tensor and permuted probs.
+    """
     if fused and probs is None:
         if not HAVE_TE or fused_sort_chunks_by_index is None:
             raise ValueError(
@@ -442,7 +471,7 @@ def group_limited_topk(
     num_experts: int,
     num_groups: int,
     group_topk: int,
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """Perform top-k routing on a subset of expert groups.
 
     When using group-limited routing:
@@ -538,19 +567,26 @@ def topk_routing_with_score_function(
     score_function: str = "softmax",
     expert_bias: Optional[torch.Tensor] = None,
     fused: bool = False,
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute the routing probabilities and map for top-k selection with score function.
+
     Args:
         logits (torch.Tensor): Logits tensor.
         topk (int): The number of experts to select for each token.
-        use_pre_softmax (bool): Whether to apply softmax or sigmoid before top-k selection.
-        num_groups (int): Number of groups for routed experts.
-        group_topk (int): Number of selected groups for each token.
-        scaling_factor (float): Scaling factor of routing score in top-k selection.
-        score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
-        expert_bias (torch.Tensor): The bias added to logits for expert routing.
+        use_pre_softmax (bool, optional): Whether to apply softmax or sigmoid before top-k
+                                          selection. Defaults to False.
+        num_groups (int, optional): Number of groups for routed experts. Defaults to None.
+        group_topk (int, optional): Number of selected groups for each token. Defaults to None.
+        scaling_factor (float, optional): Scaling factor of routing score in top-k selection.
+                                         Defaults to None.
+        score_function (str, optional): The score function to use. Can be either "softmax" or
+                                        "sigmoid". Defaults to "softmax".
+        expert_bias (torch.Tensor, optional): The bias added to logits for expert routing.
+                                              Defaults to None.
+        fused (bool, optional): Whether to use the fused version. Defaults to False.
+
     Returns:
-        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        Tuple[torch.Tensor, torch.Tensor]:
             - routing_probs (torch.Tensor): A tensor of shape [num_tokens, num_experts] containing
               the routing probabilities for each token to each expert.
             - routing_map (torch.Tensor): A mask tensor of shape [num_tokens, num_experts]
@@ -575,7 +611,25 @@ def topk_routing_with_score_function(
             expert_bias=expert_bias,
         )
 
-    def compute_topk(scores, topk, num_groups=None, group_topk=None):
+    def compute_topk(
+        scores: torch.Tensor,
+        topk: int,
+        num_groups: Optional[int] = None,
+        group_topk: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute the top-k indices for the given scores.
+
+        Args:
+            scores (torch.Tensor): The scores tensor.
+            topk (int): The number of top-k indices to compute.
+            num_groups (int, optional): The number of groups to compute the top-k indices for.
+                                        Defaults to None.
+            group_topk (int, optional): The number of top-k indices to compute for each group.
+                                        Defaults to None.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: The top-k indices and the top-k scores.
+        """
         if group_topk:
             return group_limited_topk(
                 scores=scores,
@@ -631,14 +685,17 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
 
 def compute_routing_scores_for_aux_loss(
     logits: torch.Tensor, topk: int, score_function: str, fused: bool = False
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute routing scores based on the score function.
 
     Args:
         logits (torch.Tensor): The logits tensor after gating, shape: [num_tokens, num_experts].
+        topk (int): The number of top-k indices to compute.
+        score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
+        fused (bool, optional): Whether to use the fused version. Defaults to False.
 
     Returns:
-        torch.Tensor: The normalized routing scores.
+        Tuple[torch.Tensor, torch.Tensor]: The routing map and the normalized routing scores.
     """
     if fused:
         if not HAVE_TE or fused_compute_score_for_moe_aux_loss is None:
@@ -669,7 +726,7 @@ def apply_router_token_dropping(
     capacity_factor: float,
     drop_policy: str = "probs",
     pad_to_capacity: bool = False,
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """Apply token dropping to top-k expert selection.
 
     This function enforces expert capacity limits by dropping tokens that exceed
@@ -682,8 +739,9 @@ def apply_router_token_dropping(
             indicating which experts were selected for each token.
         router_topk (int): Number of experts selected per token.
         capacity_factor (float): The capacity factor of each expert.
-        drop_policy (str): Policy to drop tokens - "probs" or "position".
-        pad_to_capacity (bool): Whether to pad to capacity.
+        drop_policy (str, optional): Policy to drop tokens - "probs" or "position".
+                                     Defaults to "probs".
+        pad_to_capacity (bool, optional): Whether to pad to capacity. Defaults to False.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]:
@@ -735,18 +793,20 @@ def save_to_aux_losses_tracker(
     reduce_group: Optional[torch.distributed.ProcessGroup] = None,
     avg_group: Optional[torch.distributed.ProcessGroup] = None,
     reduce_group_has_dp: bool = False,
-):
+) -> None:
     """Save the auxiliary loss for logging.
     Args:
         name (str): The name of the loss.
         loss (torch.Tensor): The loss tensor.
         layer_number (int): Layer index of the loss.
         num_layers (int): The number of total layers.
-        reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss.
-        avg_group (torch.distributed.ProcessGroup): The group for averaging the loss.
-        reduce_group_has_dp (bool): Whether the reduce group has data parallel ranks.
+        reduce_group (torch.distributed.ProcessGroup, optional): The group for reducing the loss.
+                                                                 Defaults to None.
+        avg_group (torch.distributed.ProcessGroup, optional): The group for averaging the loss.
+                                                              Defaults to None.
+        reduce_group_has_dp (bool, optional): Whether the reduce group has data parallel ranks.
             Set this to True if the reduce group has data parallel ranks. This flag is used to
-            ensure the correct reduction in aux loss tracking.
+            ensure the correct reduction in aux loss tracking. Defaults to False.
     """
     # Skip aux loss logging if layer_number is None.
     if layer_number is None:
@@ -762,7 +822,7 @@ def save_to_aux_losses_tracker(
     tracker[name]["reduce_group_has_dp"] = reduce_group_has_dp
 
 
-def clear_aux_losses_tracker():
+def clear_aux_losses_tracker() -> None:
     """Clear the auxiliary losses."""
     tracker = get_moe_layer_wise_logging_tracker()
     for name in tracker:
@@ -771,8 +831,15 @@ def clear_aux_losses_tracker():
 
 def reduce_aux_losses_tracker_across_ranks(
     track_names: Optional[List[str]] = None, pg_collection: Optional[ProcessGroupCollection] = None
-):
-    """Collect and reduce the auxiliary losses across ranks."""
+) -> None:
+    """Collect and reduce the auxiliary losses across ranks.
+
+    Args:
+        track_names (Optional[List[str]], optional):
+            The names of the losses to track. Defaults to None.
+        pg_collection (Optional[ProcessGroupCollection], optional):
+            The process group collection. Defaults to None.
+    """
     tracker = get_moe_layer_wise_logging_tracker()
     if track_names is None:
         track_names = tracker.keys()
@@ -810,18 +877,38 @@ def reduce_aux_losses_tracker_across_ranks(
 def track_moe_metrics(
     loss_scale: float,
     iteration: int,
-    writer,
-    wandb_writer=None,
-    total_loss_dict=None,
-    per_layer_logging=False,
+    writer: Optional["SummaryWriter"] = None,
+    wandb_writer: Optional["wandb.Run"] = None,
+    total_loss_dict: Optional[dict[str, torch.Tensor]] = None,
+    per_layer_logging: bool = False,
     force_initialize: bool = False,
     track_names: Optional[List[str]] = None,
     num_layers: Optional[int] = None,
     moe_layer_freq: Optional[Union[int, List[int]]] = None,
     mtp_num_layers: Optional[int] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
-):
-    """Track the MoE metrics for logging."""
+) -> None:
+    """Track the MoE metrics for logging.
+
+    Args:
+        loss_scale (float): The loss scale.
+        iteration (int): The iteration.
+        writer (SummaryWriter, optional): The tensorboard writer. Defaults to None.
+        wandb_writer (wandb.Run, optional): The wandb writer. Defaults to None.
+        total_loss_dict (dict[str, torch.Tensor], optional): The total loss dictionary.
+                                                             Defaults to None.
+        per_layer_logging (bool, optional): Whether to log per layer. Defaults to False.
+        force_initialize (bool, optional): Whether to force initialize the tracker.
+                                           Defaults to False.
+        track_names (List[str], optional): The names of the losses to track. Defaults to None.
+        num_layers (int, optional): The number of layers. Defaults to None.
+        moe_layer_freq (Union[int, List[int]], optional): The frequency of the MoE layers.
+                                                          Defaults to None.
+        mtp_num_layers (int, optional): The number of layers in the model parallel group.
+                                        Defaults to None.
+        pg_collection (ProcessGroupCollection, optional): The process group collection.
+                                                          Defaults to None.
+    """
     # Aux loss logging
     tracker = get_moe_layer_wise_logging_tracker()
     # Initialize the tracker if force_initialize is True
@@ -884,13 +971,18 @@ def track_moe_metrics(
     clear_aux_losses_tracker()
 
 
-def get_updated_expert_bias(tokens_per_expert, expert_bias, expert_bias_update_rate):
+def get_updated_expert_bias(
+    tokens_per_expert: torch.Tensor, expert_bias: torch.Tensor, expert_bias_update_rate: float
+) -> torch.Tensor:
     """Update expert bias for biased expert routing. See https://arxiv.org/abs/2408.15664v1#
 
     Args:
         tokens_per_expert (torch.Tensor): The number of tokens assigned to each expert.
         expert_bias (torch.Tensor): The bias for each expert.
         expert_bias_udpate_rate (float): The update rate for the expert bias.
+
+    Returns:
+        torch.Tensor: The updated expert bias.
     """
     with torch.no_grad():
         # All Reduce Across TPxCPxDP group
@@ -905,13 +997,20 @@ def get_updated_expert_bias(tokens_per_expert, expert_bias, expert_bias_update_r
         return updated_expert_bias
 
 
-def maybe_move_tensor_to_cpu(tensor, as_numpy=False, record_stream=False):
+def maybe_move_tensor_to_cpu(
+    tensor: torch.Tensor, as_numpy: bool = False, record_stream: bool = False
+) -> torch.Tensor:
     """Move a tensor to CPU if it is on GPU.
     Args:
-        tensor (torch.Tensor or None): The tensor to move to CPU.
-        as_numpy (bool): Whether to convert the tensor to a numpy array.
-        record_stream (bool): Whether to record the stream of the tensor, to prevent memory leak
-                              when the DtoH data transfer is on a side stream.
+        tensor (torch.Tensor): The tensor to move to CPU.
+        as_numpy (bool, optional): Whether to convert the tensor to a numpy array.
+                                   Defaults to False.
+        record_stream (bool, optional): Whether to record the stream of the tensor, to prevent
+                                        memory leak when the DtoH data transfer is on a side
+                                        stream. Defaults to False.
+
+    Returns:
+        torch.Tensor: The tensor moved to CPU.
     """
     if torch.is_tensor(tensor) and tensor.is_cuda:
         cpu_tensor = tensor.to(torch.device("cpu"), non_blocking=True)
@@ -923,7 +1022,7 @@ def maybe_move_tensor_to_cpu(tensor, as_numpy=False, record_stream=False):
     return tensor
 
 
-def get_moe_layer_wise_logging_tracker():
+def get_moe_layer_wise_logging_tracker() -> dict:
     """Return the moe layer wise tracker."""
     global _MOE_LAYER_WISE_LOGGING_TRACKER
     return _MOE_LAYER_WISE_LOGGING_TRACKER
@@ -939,25 +1038,43 @@ class RandomSTE(torch.autograd.Function):
     """
 
     @staticmethod
-    def forward(ctx, logits):
+    def forward(ctx, logits: torch.Tensor) -> torch.Tensor:
         """
         Forward pass returns random logits with rank-specific seed.
+
+        Args:
+            logits (torch.Tensor): The logits.
+
+        Returns:
+            torch.Tensor: The random logits.
         """
         with get_cuda_rng_tracker().fork(get_expert_parallel_rng_tracker_name()):
             random_logits = logits.clone().normal_()
         return random_logits
 
     @staticmethod
-    def backward(ctx, grad_output):
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
         """
         Backward pass propagates the gradient for logits.
+
+        Args:
+            grad_output (torch.Tensor): The gradient output.
+
+        Returns:
+            torch.Tensor: The gradient input.
         """
         return grad_output
 
 
-def apply_random_logits(logits):
+def apply_random_logits(logits: torch.Tensor) -> torch.Tensor:
     """
     Apply the RandomSTE function to the logits.
+
+    Args:
+        logits (torch.Tensor): The logits.
+
+    Returns:
+        torch.Tensor: The random logits.
     """
     return RandomSTE.apply(logits)
 
@@ -969,10 +1086,23 @@ class RouterGatingLinearFunction(torch.autograd.Function):
 
     @staticmethod
     def forward(
-        ctx, inp: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, router_dtype: torch.dtype
-    ):
+        ctx,
+        inp: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        router_dtype: torch.dtype,
+    ) -> torch.Tensor:
         """
         Forward pass of the RouterGatingLinearFunction function.
+
+        Args:
+            inp (torch.Tensor): The input tensor.
+            weight (torch.Tensor): The weight tensor.
+            bias (torch.Tensor): The bias tensor. Could be None.
+            router_dtype (torch.dtype): The router dtype.
+
+        Returns:
+            torch.Tensor: The output tensor.
         """
         ctx.save_for_backward(inp, weight, bias)
         ctx.router_dtype = router_dtype
@@ -995,9 +1125,18 @@ def forward(
         return output
 
     @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):
+    def backward(
+        ctx, grad_output: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], None]:
         """
         Backward pass of the RouterGatingLinearFunction function.
+
+        Args:
+            grad_output (torch.Tensor): The gradient output.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], None]:
+                The gradient input, gradient weight, gradient bias, and None.
         """
         inp, weight, bias = ctx.saved_tensors
         inp_shape = inp.shape
@@ -1024,18 +1163,34 @@ def backward(ctx, grad_output: torch.Tensor):
 
 
 def router_gating_linear(
-    inp: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, router_dtype: torch.dtype
-):
+    inp: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], router_dtype: torch.dtype
+) -> torch.Tensor:
     """
     Customized linear layer for router gating.
     This linear layer accepts bfloat16 input and weight, and can return output with router_dtype.
     It can reduce the memory usage by avoiding saving the intermediate high precision tensors.
+
+    Args:
+        inp (torch.Tensor): The input tensor.
+        weight (torch.Tensor): The weight tensor.
+        bias (torch.Tensor): The bias tensor. Could be None.
+        router_dtype (torch.dtype): The router dtype.
+
+    Returns:
+        torch.Tensor: The output tensor.
     """
     return RouterGatingLinearFunction.apply(inp, weight, bias, router_dtype)
 
 
-def get_align_size_for_quantization(config: TransformerConfig):
-    """Get the alignment size for quantization."""
+def get_align_size_for_quantization(config: TransformerConfig) -> int:
+    """Get the alignment size for quantization.
+
+    Args:
+        config (TransformerConfig): The configuration.
+
+    Returns:
+        int: The alignment size for quantization.
+    """
     if config.fp8:
         return get_fp8_align_size(config.fp8_recipe)
     elif config.fp4:
@@ -1045,7 +1200,7 @@ def get_align_size_for_quantization(config: TransformerConfig):
 
 # TODO(Hepteract): delete the usage of the global parallel_state.
 # Initialize process groups with the global parallel_state.
-def get_default_pg_collection():
+def get_default_pg_collection() -> ProcessGroupCollection:
     """Get the default process groups for MoE.
 
     Returns:
@@ -1080,7 +1235,7 @@ def __init__(self, moe_layer, return_step: str, **kwargs):
 
     def get_early_return_outputs(
         self, hidden_states: torch.Tensor, shared_expert_output: torch.Tensor
-    ):
+    ) -> List[torch.Tensor]:
         """
         Get the CUDA graph early return outputs for the MoE layer, including the intermediate
         tensors and the intermediate attributes of the token dispatcher.

From 10c6f010e902f704626c14cdcbe5e7d9e266e357 Mon Sep 17 00:00:00 2001
From: HaochenYuan <106647990+HaochenYuan@users.noreply.github.com>
Date: Fri, 23 Jan 2026 14:37:40 +0800
Subject: [PATCH 03/79] Remove calculation of padding token in moe routing loss
 (#2142)

Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
---
 .../core/extensions/transformer_engine.py     |   2 +-
 .../common/model_chunk_schedule_plan.py       |   2 +
 .../core/models/gpt/fine_grained_callables.py |  21 +-
 megatron/core/models/gpt/gpt_model.py         |  37 +++-
 megatron/core/models/mamba/mamba_model.py     |   2 +
 megatron/core/ssm/mamba_block.py              |   2 +
 megatron/core/transformer/mlp.py              |   2 +-
 megatron/core/transformer/moe/moe_layer.py    |  21 +-
 megatron/core/transformer/moe/moe_utils.py    |  91 +++++++--
 megatron/core/transformer/moe/router.py       | 160 ++++++++++-----
 .../core/transformer/transformer_block.py     |  14 +-
 .../core/transformer/transformer_layer.py     |  24 ++-
 .../a2a_overlap/test_schedule_chunk_1f1b.py   | 116 ++++++++++-
 .../a2a_overlap/test_schedule_layer_1f1b.py   |   4 +-
 .../transformer/moe/test_aux_loss.py          | 182 ++++++++++++++++++
 .../transformer/moe/test_routers.py           |  47 +++++
 16 files changed, 639 insertions(+), 88 deletions(-)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 63694fc1172..ef8527e9e5e 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -2161,7 +2161,7 @@ def forward_post_hook(module, *_) -> None:
                     "TEFusedMLP module does not support submodules with post-backward hooks"
                 )
 
-        def forward(self, hidden_states: torch.Tensor) -> Tuple[Tensor, Optional[Tensor]]:
+        def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Optional[Tensor]]:
             """Forward."""
 
             # Construct fused impl if needed
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 71aa1ab97f0..033e8e808f9 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -281,6 +281,7 @@ def __init__(
         extra_block_kwargs=None,
         runtime_gather_output: Optional[bool] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask=None,
     ):
         """Initialize the schedule plan of all Transformer layers' sub-modules.
 
@@ -323,6 +324,7 @@ def __init__(
         self._model_chunk_state.mtp_hidden_states = None
         self._model_chunk_state.loss_mask = loss_mask
         self._model_chunk_state.packed_seq_params = packed_seq_params
+        self._model_chunk_state.padding_mask = padding_mask
         self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
         self._model_chunk_state.runtime_gather_output = runtime_gather_output
         self._model_chunk_state.model = model
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 9234e142c6c..6f2f6b1cb80 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -131,13 +131,19 @@ def forward_impl(self):
         if not self.gpt_model.pre_process:
             self.chunk_state.decoder_input = self.gpt_model.decoder.input_tensor
         # Run GPTModel._preprocess
-        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
-            self.gpt_model._preprocess(
-                input_ids=self.chunk_state.input_ids,
-                position_ids=self.chunk_state.position_ids,
-                decoder_input=self.chunk_state.decoder_input,
-                packed_seq_params=self.chunk_state.packed_seq_params,
-            )
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = self.gpt_model._preprocess(
+            input_ids=self.chunk_state.input_ids,
+            position_ids=self.chunk_state.position_ids,
+            decoder_input=self.chunk_state.decoder_input,
+            packed_seq_params=self.chunk_state.packed_seq_params,
+            padding_mask=self.chunk_state.padding_mask,
         )
 
         # Saved for later use
@@ -146,6 +152,7 @@ def forward_impl(self):
         self.chunk_state.rotary_pos_cos = rotary_pos_cos
         self.chunk_state.rotary_pos_sin = rotary_pos_sin
         self.chunk_state.sequence_len_offset = sequence_len_offset
+        self.chunk_state.padding_mask = padding_mask
         return decoder_input
 
 
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index e70221d2cfa..e287344c13d 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -288,6 +288,7 @@ def _preprocess(
         decoder_input: Tensor = None,
         inference_context: BaseInferenceContext = None,
         packed_seq_params: PackedSeqParams = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Preprocesses inputs for the transformer decoder.
 
@@ -304,7 +305,20 @@ def _preprocess(
         if decoder_input is not None:
             pass
         elif self.pre_process:
+            if padding_mask is not None:
+                assert padding_mask.shape == input_ids.shape, (
+                    f"padding_mask shape {padding_mask.shape} does not match "
+                    f"input_ids shape {input_ids.shape}"
+                )
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+            if padding_mask is not None and self.config.sequence_parallel:
+                padding_mask = (
+                    tensor_parallel.scatter_to_sequence_parallel_region(
+                        padding_mask.transpose(0, 1).contiguous()
+                    )
+                    .transpose(0, 1)
+                    .contiguous()
+                )
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
@@ -423,6 +437,7 @@ def _preprocess(
             rotary_pos_cos,
             rotary_pos_sin,
             sequence_len_offset,
+            padding_mask,
         )
         if rotary_pos_cos_sin is not None:
             # only in the case of flashinfer fused rope will we
@@ -466,6 +481,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -476,6 +492,9 @@ def forward(
         Args:
             runtime_gather_output (bool): Gather output at runtime. Default None means
                 `parallel_output` arg in the constructor will be used.
+            padding_mask (Tensor, optional): Padding mask for MoE routing.
+                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Only used for MoE layers to exclude padding tokens from routing computations.
         """
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()
@@ -488,13 +507,19 @@ def forward(
             decoder_input=decoder_input,
             inference_context=inference_context,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
-        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset) = (
-            preproc_output[:5]
-        )
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = preproc_output[:6]
 
-        rotary_pos_cos_sin = preproc_output[5] if len(preproc_output) == 6 else None
+        rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
 
         # Run decoder.
         hidden_states = self.decoder(
@@ -507,6 +532,7 @@ def forward(
             rotary_pos_cos_sin=rotary_pos_cos_sin,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
+            padding_mask=padding_mask,
             **(extra_block_kwargs or {}),
         )
 
@@ -723,6 +749,7 @@ def build_schedule_plan(
         runtime_gather_output: Optional[bool] = None,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Builds a computation schedule plan for the model.
 
@@ -748,6 +775,7 @@ def build_schedule_plan(
             inference_params (InferenceParams, optional):
                 Parameters for inference. Defaults to None.
             loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
+            padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
 
         Returns:
             TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -769,6 +797,7 @@ def build_schedule_plan(
             extra_block_kwargs,
             runtime_gather_output,
             loss_mask,
+            padding_mask,
         )
 
     def sharded_state_dict(
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 0d71ead4b0f..8d45e1d0147 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -185,6 +185,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
+        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the Mamba model. This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -254,6 +255,7 @@ def forward(
             inference_context=inference_context,
             rotary_pos_emb=rotary_pos_emb,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
         if not self.post_process:
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 9e41aca8253..ef41faae143 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -211,6 +211,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
+        padding_mask=None,
     ):
         """
         Forward function of the MambaStack class.
@@ -293,6 +294,7 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
                         )
                     else:  # MambaLayer
                         hidden_states = layer(
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index 2eae0178eea..2bc3949a421 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -148,7 +148,7 @@ def __init__(
             tp_group=tp_group,
         )
 
-    def forward(self, hidden_states, per_token_scale=None):
+    def forward(self, hidden_states, per_token_scale=None, **kwargs):
         """Perform the forward pass through the MLP block."""
         # [s, b, 4 * h/p]
         nvtx_range_push(suffix="linear_fc1")
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 2fad0f8e5b7..08c68ae3aed 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -239,13 +239,13 @@ def __init__(
         self.cudagraph_tensor_store = MoECudaGraphTensorStore()
 
     @maybe_skip_or_early_return_by_cudagraph("route")
-    def route(self, hidden_states: torch.Tensor):
+    def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Compute token routing for preprocessing.
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
         """
-        probs, routing_map = apply_module(self.router)(hidden_states)
+        probs, routing_map = apply_module(self.router)(hidden_states, padding_mask)
         return probs, routing_map
 
     @maybe_skip_or_early_return_by_cudagraph("preprocess")
@@ -346,7 +346,7 @@ def router_and_preprocess(self, hidden_states: torch.Tensor):
         hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map)
         return hidden_states, probs, residual
 
-    def forward(self, hidden_states: torch.Tensor):
+    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Forward pass for the MoE layer.
 
         The forward pass comprises four main steps:
@@ -356,8 +356,10 @@ def forward(self, hidden_states: torch.Tensor):
         4. Combine: The outputs from the experts are combined and returned.
 
         Args:
-            hidden_states (torch.Tensor): The input tensor to the MoE layer.
-
+            hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size].
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
         Returns:
             A tuple containing the output tensor and the MLP bias, if any.
         """
@@ -366,12 +368,15 @@ def forward(self, hidden_states: torch.Tensor):
                 "During training, performance may degrade if MoE and tensor parallelism"
                 "are enabled without also enabling sequence parallelism."
             )
+        # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states
+        if padding_mask is not None:
+            padding_mask = padding_mask.transpose(0, 1).bool()
 
         # MoE forward: route -> dispatch -> compute -> combine
         def custom_forward(hidden_states):
             try:
                 shared_expert_output = self.shared_experts_compute(hidden_states)
-                probs, routing_map = self.route(hidden_states)
+                probs, routing_map = self.route(hidden_states, padding_mask)
                 hidden_states, probs = self.preprocess(hidden_states, probs, routing_map)
             except MoECudaGraphPartialCaptureSignal as e:
                 # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator.
@@ -398,7 +403,9 @@ def custom_forward(hidden_states):
                     hidden_states,
                 )
             else:
-                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+                outputs = tensor_parallel.checkpoint(
+                    custom_forward, False, hidden_states, padding_mask
+                )
         else:
             outputs = custom_forward(hidden_states)
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 5fdeda23dea..28c486545a2 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -12,6 +12,7 @@
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -49,6 +50,7 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
@@ -100,10 +102,19 @@ def switch_load_balancing_loss_func(
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
         fused (bool): Whether to use the fused version of the auxiliary loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
 
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
+    # Apply padding mask to probs if provided
+    if padding_mask is not None:
+        # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
+        mask_expanded = padding_mask.unsqueeze(-1)
+        probs = probs * mask_expanded
+
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -123,19 +134,35 @@ def switch_load_balancing_loss_func(
     return aux_loss
 
 
-def z_loss_func(logits: torch.Tensor, z_loss_coeff: float) -> torch.Tensor:
+def z_loss_func(
+    logits: torch.Tensor, z_loss_coeff: float, padding_mask: Optional[torch.Tensor] = None
+) -> torch.Tensor:
     """Encourages the router's logits to remain small to enhance stability.
     Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
     Args:
         logits (torch.Tensor): The logits of the router.
         z_loss_coeff (float): The coefficient for the z-loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                               Shape [num_tokens]. True = padding (exclude),
+                                               False = valid (include). Defaults to None.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
     """
-
-    z_loss = torch.mean(torch.square(torch.logsumexp(logits, dim=-1))) * z_loss_coeff
+    logsum = torch.logsumexp(logits, dim=-1)
+    z_loss_values = torch.square(logsum)
+
+    if padding_mask is not None:
+        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
+        valid_mask = ~padding_mask
+        # Only compute z_loss for valid (non-padding) tokens
+        z_loss_values = z_loss_values * valid_mask
+        # Compute mean over valid tokens only
+        num_valid_tokens = valid_mask.sum()
+        z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff
+    else:
+        z_loss = torch.mean(z_loss_values) * z_loss_coeff
     return z_loss
 
 
@@ -185,6 +212,28 @@ def get_capacity(
     return capacity
 
 
+def get_tokens_per_expert_and_token_count(
+    routing_map: torch.Tensor,
+    reduce_group: torch.distributed.ProcessGroup,
+    topk: int = None,
+    with_padding_mask: bool = False,
+) -> torch.Tensor:
+    """
+    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with padding mask.
+    """
+    local_tokens_per_expert = routing_map.sum(dim=0)
+    global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
+        local_tokens_per_expert, reduce_group
+    )
+    if with_padding_mask:
+        local_num_tokens = local_tokens_per_expert.sum() / topk
+        total_num_tokens = global_tokens_per_expert.sum() / topk
+    else:
+        local_num_tokens = routing_map.shape[0]
+        total_num_tokens = local_num_tokens * reduce_group.size()
+    return global_tokens_per_expert, local_num_tokens, total_num_tokens
+
+
 class MoEAuxLossAutoScaler(torch.autograd.Function):
     """An AutoScaler that triggers the backward pass and scales the grad for auxiliary loss."""
 
@@ -684,7 +733,11 @@ def compute_topk(
 
 
 def compute_routing_scores_for_aux_loss(
-    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False
+    logits: torch.Tensor,
+    topk: int,
+    score_function: str,
+    fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute routing scores based on the score function.
 
@@ -693,6 +746,9 @@ def compute_routing_scores_for_aux_loss(
         topk (int): The number of top-k indices to compute.
         score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
         fused (bool, optional): Whether to use the fused version. Defaults to False.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: The routing map and the normalized routing scores.
@@ -702,20 +758,27 @@ def compute_routing_scores_for_aux_loss(
             raise ValueError(
                 "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
             )
-        return fused_compute_score_for_moe_aux_loss(
+        routing_map, scores = fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
-
-    if score_function == "softmax":
-        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
-    elif score_function == "sigmoid":
-        scores = torch.sigmoid(logits)
-        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
     else:
-        raise ValueError(f"Invalid score_function: {score_function}")
+        if score_function == "softmax":
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        elif score_function == "sigmoid":
+            scores = torch.sigmoid(logits)
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        else:
+            raise ValueError(f"Invalid score_function: {score_function}")
+
+        _, top_indices = torch.topk(scores, k=topk, dim=1)
+        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
 
-    _, top_indices = torch.topk(scores, k=topk, dim=1)
-    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    # Apply padding mask to scores if provided
+    if padding_mask is not None:
+        # Invert padding_mask and make True indicates valid tokens
+        valid_mask = (~padding_mask).unsqueeze(-1)
+        routing_map = routing_map * valid_mask
+        scores = scores * valid_mask
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index c22ca4e8446..8f94a7312ce 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -1,12 +1,11 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
 from megatron.core.jit import jit_fuser
-from megatron.core.tensor_parallel import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,
@@ -14,6 +13,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
+    get_tokens_per_expert_and_token_count,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -268,22 +268,29 @@ def is_aux_loss_enabled(self) -> bool:
         return False
 
     def _apply_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        with_padding_mask: bool = False,
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
-        tokens_per_expert = routing_map.sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                topk=self.topk,
+                with_padding_mask=with_padding_mask,
+            )
         )
-        num_tokens = routing_map.shape[0]
-        total_num_tokens = num_tokens * self.tp_cp_group.size()
 
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
-            tokens_per_expert=tokens_per_expert,
+            tokens_per_expert=global_tokens_per_expert,
             total_num_tokens=total_num_tokens,
             topk=self.topk,
             num_experts=self.config.num_moe_experts,
@@ -291,7 +298,12 @@ def _apply_aux_loss(
             fused=self.config.moe_router_fusion,
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group
+            probs,
+            aux_loss_coeff,
+            aux_loss,
+            "load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -302,6 +314,7 @@ def _apply_seq_aux_loss(
         routing_map: torch.Tensor,
         seq_length: int,
         bsz: int,
+        with_padding_mask: bool = False,
     ):
         """Apply the sequence-level auxiliary loss for the given scores and routing map.
 
@@ -315,17 +328,21 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
-        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+        routing_map = routing_map.reshape(seq_length, -1)
+
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk * bsz,
+            )
         )
 
-        total_num_tokens = seq_length * self.tp_cp_group.size()
-
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
-                tokens_per_expert=tokens_per_expert,
+                tokens_per_expert=global_tokens_per_expert,
                 total_num_tokens=total_num_tokens,
                 topk=self.topk,
                 num_experts=self.config.num_moe_experts,
@@ -334,31 +351,43 @@ def _apply_seq_aux_loss(
             )
             / bsz
         )
+
         probs = self.attach_and_log_load_balancing_loss(
-            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group
+            probs,
+            seq_aux_loss_coeff,
+            aux_loss,
+            "seq_load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
     def _apply_global_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        with_padding_mask: bool = False,
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
         if global_aux_loss_coeff == 0:
             return probs
 
-        tokens_per_expert = routing_map.sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_dp_cp_group
+        # Use unified function to compute tokens_per_expert and num_tokens
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_dp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk,
+            )
         )
 
-        self.global_tokens_per_expert += tokens_per_expert
+        self.global_tokens_per_expert += global_tokens_per_expert
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        num_tokens = scores_for_aux_loss.shape[0]
-        total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
-
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
             tokens_per_expert=averated_tokens_per_expert,
@@ -375,6 +404,7 @@ def _apply_global_aux_loss(
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
             reduce_group_has_dp=True,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -386,18 +416,22 @@ def attach_and_log_load_balancing_loss(
         aux_loss_name: str,
         reduce_group: torch.distributed.ProcessGroup,
         reduce_group_has_dp: bool = False,
+        valid_token_count: Optional[Union[int, torch.Tensor]] = None,
     ):
         """Attach aux loss function to activation and add to logging.
 
         Args:
-            activation (torch.Tensor): The activation tensor to attach the loss to.
-            aux_loss_coeff (float): The coefficient for the auxiliary loss.
-            aux_loss (torch.Tensor): The auxiliary loss tensor.
-            aux_loss_name (str): The name of the auxiliary loss for logging.
-            reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss.
+            activation (torch.Tensor): Activation tensor to attach the aux loss to.
+            aux_loss_coeff (float): Coefficient for the aux loss.
+            aux_loss (torch.Tensor): Computed aux loss.
+            aux_loss_name (str): Name of the aux loss for logging.
+            reduce_group (torch.distributed.ProcessGroup): Process group for reduction.
             reduce_group_has_dp (bool): Whether the reduce group has data parallel ranks.
                 Set this to True if the reduce group has data parallel ranks. This flag is used to
                 ensure the correct reduction in aux loss tracking.
+            valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding
+                padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor).
+                If None, uses activation.shape[0]. Defaults to None.
         """
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
         # add the aux loss logging value to other layer's since it is difficult to get the
@@ -422,17 +456,22 @@ def attach_and_log_load_balancing_loss(
             # which scales both the main_loss gradient and aux_loss gradient by
             # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads function.
             # To correct this scaling, we need to scale the aux_loss by num_local_tokens here.
-            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0])
+            # Use valid_token_count (excluding padding) if provided, otherwise use total tokens.
+            num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0]
+            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens)
         else:
             activation = MoEAuxLossAutoScaler.apply(activation, aux_loss)
         return activation
 
-    def apply_z_loss(self, logits):
+    def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
         """Encourages the router's logits to remain small to enhance stability.
         Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
         Args:
             logits (torch.Tensor): The logits of the router.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [num_tokens]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             torch.Tensor: The logits after applying the z-loss.
@@ -440,7 +479,7 @@ def apply_z_loss(self, logits):
         if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled():
             # Skip Z loss calculations when using torch.no_grad() or checkpointing.
             moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()
-            z_loss = z_loss_func(logits, moe_z_loss_coeff)
+            z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask)
             if self.calculate_per_token_loss:
                 # The expected final scaling for z_loss gradients is
                 # 1/(num_micro_batches * dp_size).
@@ -449,7 +488,9 @@ def apply_z_loss(self, logits):
                 # which scales both the main_loss gradient and z_loss gradient by
                 # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads().
                 # To correct this scaling, we need to scale the z_loss by num_local_tokens here.
-                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * logits.shape[0])
+                # Count valid tokens: sum of inverted mask (False -> True = valid)
+                num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0]
+                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens)
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
@@ -483,20 +524,27 @@ def apply_input_jitter(self, input: torch.Tensor):
             return input
 
     @jit_fuser
-    def _apply_expert_bias(self, routing_map: torch.Tensor):
+    def _apply_expert_bias(
+        self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
+    ):
         """
         Update expert bias and tokens_per_expert
         Prevent extra local tokens accumulation on evaluation or activation recomputation
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
+                if padding_mask is not None:
+                    routing_map = routing_map & (~padding_mask)
                 self.local_tokens_per_expert += routing_map.sum(dim=0)
 
-    def routing(self, logits: torch.Tensor):
+    def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Top-k routing function
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
@@ -506,8 +554,12 @@ def routing(self, logits: torch.Tensor):
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
 
+        # Flatten padding_mask to [num_tokens] if provided
+        if padding_mask is not None:
+            padding_mask = padding_mask.reshape(-1)
+
         # Apply Z-Loss
-        logits = self.apply_z_loss(logits)
+        logits = self.apply_z_loss(logits, padding_mask=padding_mask)
 
         # Calculate probs and routing_map for token dispatching
         if self.routing_type == "sinkhorn":
@@ -540,18 +592,35 @@ def routing(self, logits: torch.Tensor):
         if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
-                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion
+                logits,
+                self.topk,
+                self.score_function,
+                fused=self.config.moe_router_fusion,
+                padding_mask=padding_mask,
+            )
+            probs = self._apply_aux_loss(
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
-            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)
             probs = self._apply_seq_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                seq_length,
+                bsz,
+                with_padding_mask=padding_mask is not None,
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
 
         # Optionally apply expert bias
-        self._apply_expert_bias(routing_map)
+        self._apply_expert_bias(routing_map, padding_mask=padding_mask)
 
         return probs, routing_map
 
@@ -561,12 +630,15 @@ def reset_global_aux_loss_tracker(self):
             self.global_tokens_per_expert.zero_()
             self.ga_steps.zero_()
 
-    def forward(self, input: torch.Tensor):
+    def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """
         Forward pass of the router.
 
         Args:
             input (torch.Tensor): Input tensor.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
         """
         self._maintain_float32_expert_bias()
 
@@ -578,7 +650,7 @@ def forward(self, input: torch.Tensor):
             # Apply force load balancing with random logits for benchmark
             logits = apply_random_logits(logits)
 
-        probs, routing_map = self.routing(logits)
+        probs, routing_map = self.routing(logits, padding_mask=padding_mask)
 
         return probs, routing_map
 
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index ea4464b4784..831b5546d53 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -453,12 +453,18 @@ def _checkpointed_forward(
         attention_bias: Tensor,
         packed_seq_params: PackedSeqParams,
         use_inner_quantization_context: bool,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Forward method with activation checkpointing."""
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+                hidden_states,
+                attention_mask,
+                context,
+                context_mask,
+                rotary_pos_emb,
+                padding_mask=None,
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
@@ -489,6 +495,7 @@ def custom_forward(
                             attention_bias=attention_bias,
                             inference_context=None,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
                         )
                 return hidden_states, context
 
@@ -508,6 +515,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                 )
             else:
                 return tensor_parallel.checkpoint(
@@ -518,6 +526,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                 )
 
         if self.config.recompute_method == 'uniform':
@@ -623,6 +632,7 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         dynamic_inference_decode_only: Optional[bool] = None,
@@ -732,6 +742,7 @@ def forward(
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,
                     use_inner_quantization_context=use_inner_quantization_context,
+                    padding_mask=padding_mask,
                 )
             else:
                 for l_no, layer in enumerate(self.layers):
@@ -764,6 +775,7 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
+                            padding_mask=padding_mask,
                         )
 
                     if (
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 920c3b8fcba..97ca1ec222e 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import functools
 import logging
 import warnings
 from abc import ABC
@@ -488,7 +489,11 @@ def forward(self, *args, **kwargs):
         # runners in the cuda graph manager
         kwargs.pop("dynamic_inference_decode_only", None)
         hidden_states, context = self._forward_attention(*args, **kwargs)
-        output = self._forward_mlp(hidden_states, kwargs.get("inference_context", None))
+        output = self._forward_mlp(
+            hidden_states,
+            kwargs.get("inference_context", None),
+            padding_mask=kwargs.get("padding_mask", None),
+        )
         return output, context
 
     def _forward_attention(
@@ -505,6 +510,7 @@ def _forward_attention(
         inference_context: Optional[Any] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[Any] = None,
     ):
@@ -635,13 +641,18 @@ def _forward_attention(
 
         return hidden_states, context
 
-    def _forward_mlp(self, hidden_states, inference_context=None):
+    def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None):
         """
         Perform a forward pass through the feed-forward layer.
 
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
-
+                Shape [seq_length, batch_size, hidden_size].
+            inference_context: Inference context for optimizations.
+            padding_mask (Tensor, optional): Padding mask for MoE routing.
+                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Only used for MoE layers to exclude padding tokens from aux loss computations.
+                The MoELayer will internally transform this to [seq_length, bsz] format.
         Returns:
             output (Tensor): Transformed hidden states of shape [s, b, h].
         """
@@ -689,10 +700,13 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     self.pg_collection.tp,
                     pre_mlp_layernorm_output,
+                    padding_mask=padding_mask,
                 )
             else:
                 mlp_output_with_bias = tensor_parallel.checkpoint(
-                    self.mlp, False, pre_mlp_layernorm_output
+                    functools.partial(self.mlp, padding_mask=padding_mask),
+                    False,
+                    pre_mlp_layernorm_output,
                 )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
@@ -712,7 +726,7 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 # Set the residual for fused reduce-scatter + add + layer-norm + all-gather
                 # operation in MLP's fc2.
                 self._set_fc2_residual(residual)
-            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
 
         if self.recompute_pre_mlp_layernorm:
             # discard the output of the pre-mlp layernorm and register the recompute
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index 81e61a3404a..6c59dd3f9e3 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -23,7 +23,7 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-def build_model(config):
+def build_model(config, use_padding_mask=False):
     seq_len = 32
     max_seq_len = 300
     # ids = random.sample([i for i in range(max_seq_len)], seq_len)
@@ -39,6 +39,12 @@ def build_model(config):
         "attention_mask": torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda(),
     }
 
+    # Optionally add padding_mask with same shape as input_ids
+    if use_padding_mask:
+        padding_mask = torch.zeros((1, seq_len), dtype=torch.bool).cuda()
+        padding_mask[0, -8:] = True
+        data["padding_mask"] = padding_mask
+
     # build layer spec
     transformer_layer_spec = get_gpt_decoder_block_spec(config=config, use_transformer_engine=True)
     mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec.layer_specs[-1], True)
@@ -48,7 +54,7 @@ def build_model(config):
         config=config,
         transformer_layer_spec=transformer_layer_spec,
         mtp_block_spec=mtp_block_spec,
-        vocab_size=100,
+        vocab_size=128,
         pre_process=True,
         post_process=True,
         max_sequence_length=max_seq_len,
@@ -174,3 +180,109 @@ def test_1f1b_schedule_model_chunk(self, mtp_layers, dispatcher_type, fp8_flag,
                 gpt_models[i] = None
             gc.collect()
             torch.cuda.empty_cache()
+
+    @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0")
+    @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types())
+    @pytest.mark.parametrize("layers", [[2, 1], [1, 1]])
+    @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+    def test_1f1b_schedule_model_chunk_with_padding_mask(self, dispatcher_type, layers, tp_size):
+        """
+        Verifies all-to-all overlap optimization with padding_mask produces
+        the same results as the reference implementation with various TP/EP/CP combinations.
+        """
+        # Re-initialize model parallel with the specified configuration
+        Utils.destroy_model_parallel()
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=1,
+            expert_model_parallel_size=4,
+            expert_tensor_parallel_size=1,
+        )
+        set_streams()
+
+        microbatches = 1
+
+        gpt_models = []
+        schedule_plans = []
+        ref_captures = []
+        datas = []
+
+        # create TransformerConfig
+        extra_kwargs = {
+            "moe_token_dispatcher_type": dispatcher_type,
+            "tensor_model_parallel_size": tp_size,
+            "sequence_parallel": tp_size > 1,
+        }
+        if dispatcher_type == "flex":
+            extra_kwargs["moe_flex_dispatcher_backend"] = "deepep"
+            extra_kwargs["moe_router_dtype"] = "fp32"
+        with deterministic_mode():
+            for layer_num in layers:
+                output_tensors = []
+                # build config
+                config = get_test_config(num_layers=layer_num, extra_kwargs=extra_kwargs)
+                # build model with padding_mask
+                gpt_model, schedule_plan, data = build_model(config, use_padding_mask=True)
+                gpt_model.cuda()
+                gpt_models.append(gpt_model)
+                datas.append(data)
+                schedule_plans.append(schedule_plan)
+
+                # run reference
+                for _ in range(microbatches):
+                    loss = gpt_model.forward(**data)
+                    loss = float16_to_fp32(loss)
+                    loss.backward(torch.ones_like(loss))
+                    output_tensors.append(loss)
+
+                capture = {"outputs": output_tensors}
+                for name, param in gpt_model.named_parameters():
+                    capture[name] = param.grad
+                ref_captures.append(capture)
+                gpt_model.zero_grad()
+            assert gpt_models[0].embedding is not None
+            assert gpt_models[1].embedding is not None
+            # run a2a overlap
+            capture_0 = {"outputs": []}
+            capture_1 = {"outputs": []}
+            a2a_captures = [capture_0, capture_1]
+            for i in range(microbatches):
+                # 1st forward
+                if i > 0:
+                    assert (
+                        schedule_plans[0].pre_process is None
+                    ), "pre_process should be released after backward"
+                    schedule_plans[0] = gpt_models[0].build_schedule_plan(**datas[0])
+                    schedule_plans[1] = gpt_models[1].build_schedule_plan(**datas[1])
+                f_input_0 = TransformerModelChunkSchedulePlan.run(schedule_plans[0], None)
+                capture_0["outputs"].append(f_input_0)
+                # overlap
+                f_input_1 = TransformerModelChunkSchedulePlan.run(
+                    schedule_plans[1], schedule_plans[0], b_grad=torch.ones_like(f_input_0)
+                )
+                capture_1["outputs"].append(f_input_1)
+                # last backward
+                TransformerModelChunkSchedulePlan.run(
+                    None, schedule_plans[1], b_grad=torch.ones_like(f_input_1)
+                )
+            for i in range(len(gpt_models)):
+                for name, param in gpt_models[i].named_parameters():
+                    a2a_captures[i][name] = param.grad
+
+            # compare results
+            for i in range(len(ref_captures)):
+                comp_res = compare_captures(ref_captures[i], a2a_captures[i], True, True)
+                assert comp_res[0], f"[rank {torch.distributed.get_rank()}] {comp_res[1]}"
+
+            # release resources is necessary, otherwise later testcases will oom
+            for i in range(len(schedule_plans)):
+                schedule_plans[i] = None
+                ref_captures[i] = None
+                a2a_captures[i] = None
+                for k in datas[i]:
+                    datas[i][k] = None
+                datas[i] = None
+                gpt_models[i].zero_grad()
+                gpt_models[i] = None
+            gc.collect()
+            torch.cuda.empty_cache()
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
index 0fd2c445c9f..c6c4a75af99 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
@@ -502,8 +502,8 @@ def test_mtp_layer_overlap(self, dispatcher_type, fp8_flag):
             position_ids = torch.tensor(data, dtype=torch.int64).repeat((1, 1)).cuda()
             attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda()
             # get rotary pos emb
-            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _ = gpt_model._preprocess(
-                input_ids, position_ids
+            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _, _padding_mask = (
+                gpt_model._preprocess(input_ids, position_ids)
             )
             # reset model
             params = reset_model(gpt_model)
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 621e200c2cb..ccd11bf29af 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -577,3 +577,185 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
                 reduce_from_tensor_model_parallel_region(aux_loss, router.tp_cp_group)
             assert aux_loss.item() == 1, f"{aux_loss_type}: {aux_loss.item()}"
             clear_aux_losses_tracker()
+
+
+class TestPaddingMaskAuxLoss:
+    """Test padding mask support in various aux loss types."""
+
+    def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1, sequence_parallel=False):
+        """Initialize model parallel with given configuration.
+
+        Args:
+            tp_size: Tensor parallel size.
+            ep_size: Expert parallel size.
+            cp_size: Context parallel size.
+        """
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=1,
+            context_parallel_size=cp_size,
+            expert_model_parallel_size=ep_size,
+        )
+        _set_random_seed(seed_=123, data_parallel_random_init=False)
+
+        # Store parallel configuration
+        self.tp_size = tp_size
+        self.ep_size = ep_size
+        self.cp_size = cp_size
+
+        # Default configuration
+        self.default_transformer_config = TransformerConfig(
+            num_layers=1,
+            hidden_size=12,
+            num_attention_heads=8,
+            num_moe_experts=32,
+            use_cpu_initialization=True,
+            moe_router_load_balancing_type="aux_loss",
+            moe_router_topk=8,
+            moe_aux_loss_coeff=1.0,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            add_bias_linear=False,
+            tensor_model_parallel_size=tp_size,
+            expert_model_parallel_size=ep_size,
+            context_parallel_size=cp_size,
+            sequence_parallel=sequence_parallel and tp_size > 1,
+        )
+
+    def new_router(self, **kwargs):
+        """Create a new router with updated configuration."""
+        pg_collection = get_default_pg_collection()
+        new_transformer_config = dataclasses.replace(self.default_transformer_config, **kwargs)
+        router = TopKRouter(config=new_transformer_config, pg_collection=pg_collection)
+        router.set_layer_number(0)
+        return router
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size):
+        """Test that padding tokens are correctly excluded from aux loss calculation."""
+        # Initialize model parallel with given configuration
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+
+        try:
+            clear_aux_losses_tracker()
+
+            router = self.new_router(
+                moe_router_load_balancing_type=aux_loss_type,
+                moe_aux_loss_coeff=1.0,
+                moe_router_dtype="fp64",
+            ).cuda()
+
+            seq_len = 32
+            batch_size = 2
+            hidden_size = router.config.hidden_size
+
+            # Create input with padding
+            hidden_states_full = torch.randn(
+                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+            )
+
+            # Create padding mask: first half valid, second half padding
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = True
+
+            # Test with padding mask
+            router.weight.grad = None
+            scores_with_mask, routing_map_with_mask = router(
+                hidden_states_full, padding_mask=padding_mask
+            )
+            scores_with_mask.backward(torch.zeros_like(scores_with_mask))
+
+            loss_name = {
+                "aux_loss": "load_balancing_loss",
+                "seq_aux_loss": "seq_load_balancing_loss",
+                "global_aux_loss": "global_load_balancing_loss",
+            }[aux_loss_type]
+
+            tracker = get_moe_layer_wise_logging_tracker()
+            aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
+            grad_with_mask = router.weight.grad.clone()
+
+            # Test without padding (with only half of the tokens)
+            clear_aux_losses_tracker()
+            router.weight.grad = None
+            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
+            scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
+            scores_without_mask.backward(torch.zeros_like(scores_without_mask))
+
+            aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
+            grad_without_mask = router.weight.grad.clone()
+
+            # The aux loss with mask should be equal to the aux loss without mask
+            assert torch.equal(aux_loss_with_mask, aux_loss_without_mask)
+            assert torch.equal(grad_with_mask, grad_without_mask)
+
+            clear_aux_losses_tracker()
+        finally:
+            # Always cleanup model parallel
+            Utils.destroy_model_parallel()
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
+        """Test that padding mask works correctly with z_loss."""
+        # Initialize model parallel with given configuration
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+
+        try:
+            clear_aux_losses_tracker()
+
+            router = self.new_router(
+                moe_router_load_balancing_type="aux_loss",
+                moe_aux_loss_coeff=0.0,
+                moe_z_loss_coeff=1.0,
+                moe_router_dtype="fp32",
+            ).cuda()
+
+            seq_len = 32
+            batch_size = 2
+            hidden_size = router.config.hidden_size
+
+            # Create input
+            hidden_states_full = torch.randn(
+                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+            )
+
+            # Create padding mask: first half valid, second half padding
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = True
+
+            # Test with padding mask
+            router.weight.grad = None
+            scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
+            scores_with_mask.sum().backward()
+
+            tracker = get_moe_layer_wise_logging_tracker()
+            z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
+            grad_with_mask = router.weight.grad.clone()
+
+            # Test without padding (with only half of the tokens)
+            clear_aux_losses_tracker()
+            router.weight.grad = None
+            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
+            scores_without_mask, _ = router(hidden_states_valid)
+            scores_without_mask.sum().backward()
+
+            z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
+            grad_without_mask = router.weight.grad.clone()
+
+            # The z_loss with mask should be close to the z_loss without mask
+            assert torch.equal(z_loss_with_mask, z_loss_without_mask)
+            assert torch.equal(grad_with_mask, grad_without_mask)
+
+            clear_aux_losses_tracker()
+        finally:
+            # Always cleanup model parallel
+            Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index 904595928de..4d6b5ee2c3e 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -127,6 +127,53 @@ def test_aux_loss(self):
         out.sum().mul_(0).backward()
         assert self.sequential_mlp.router.weight.grad.abs().sum() > 0
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_router_with_padding_mask(self):
+        """Test that padding mask correctly excludes padding tokens from routing."""
+        self.router = self.router.cuda()
+        seq_len = 32
+        batch_size = 2
+        hidden_size = self.router.config.hidden_size
+
+        # Create input with shape [seq_len, batch_size, hidden_size]
+        hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16()
+
+        # Create padding mask: first half valid, second half padding
+        # padding_mask shape: [seq_len, batch_size]
+        # Convention: True = padding (exclude), False = valid (include)
+        padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+        padding_mask[seq_len // 2 :, :] = True  # Second half is padding
+
+        # Test forward pass with padding mask
+        with torch.no_grad():
+            probs_with_mask, routing_map_with_mask = self.router(
+                hidden_states, padding_mask=padding_mask
+            )
+
+            # Test forward pass without padding mask (only valid tokens)
+            hidden_states_valid = hidden_states[: seq_len // 2, :, :]
+            probs_without_mask, routing_map_without_mask = self.router(hidden_states_valid)
+
+            # The valid part of routing with mask should match routing without mask
+            probs_valid_part = probs_with_mask.reshape(seq_len, batch_size, -1)[
+                : seq_len // 2, :, :
+            ]
+            probs_valid_part = probs_valid_part.reshape(-1, probs_valid_part.shape[-1])
+
+            # Check that shapes are as expected
+            assert probs_with_mask.shape == (
+                seq_len * batch_size,
+                self.router.config.num_moe_experts,
+            )
+            assert routing_map_with_mask.shape == (
+                seq_len * batch_size,
+                self.router.config.num_moe_experts,
+            )
+
+            # Verify that probs for valid tokens are similar
+            assert torch.equal(probs_valid_part, probs_without_mask)
+
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_router_dtype(self):

From 029f48f2f545efdaef2e05d3e7c4a1b6dc7e7ccb Mon Sep 17 00:00:00 2001
From: Jon Barker <jbarker@nvidia.com>
Date: Fri, 23 Jan 2026 10:36:15 -0700
Subject: [PATCH 04/79] Bug fix with --no-use-tokenizer-from-checkpoint-args
 (#3049)

Co-authored-by: Jon Barker <jbarker@oci-hsg-cs-001-vscode-01.cm.cluster>
---
 megatron/training/checkpointing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index a4605121f82..e49d7f47767 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1413,10 +1413,10 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('moe_latent_size', force=True)
 
     # Tokenizer args.
-    _set_arg('tokenizer_type', force=True)
     # Using checkpoint version might not always be safe (e.g., if running on different cluster).
     if args.use_tokenizer_model_from_checkpoint_args:
         _set_arg('tokenizer_model', force=True)
+        _set_arg('tokenizer_type', force=True)
     _set_arg('tiktoken_pattern', force=True)
     _set_arg('padded_vocab_size')
 

From 06836790b6f532d381d2760207392b1d8e5b5bbb Mon Sep 17 00:00:00 2001
From: Dong Hyuk Chang <donghyukc@nvidia.com>
Date: Fri, 23 Jan 2026 16:04:47 -0500
Subject: [PATCH 05/79] Revert "Bug fix with
 --no-use-tokenizer-from-checkpoint-args (#3049)" (#3057)

Signed-off-by: Dong Hyuk Chang <donghyukc@nvidia.com>
---
 megatron/training/checkpointing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index e49d7f47767..a4605121f82 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1413,10 +1413,10 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('moe_latent_size', force=True)
 
     # Tokenizer args.
+    _set_arg('tokenizer_type', force=True)
     # Using checkpoint version might not always be safe (e.g., if running on different cluster).
     if args.use_tokenizer_model_from_checkpoint_args:
         _set_arg('tokenizer_model', force=True)
-        _set_arg('tokenizer_type', force=True)
     _set_arg('tiktoken_pattern', force=True)
     _set_arg('padded_vocab_size')
 

From 93567e80a49ea461635718e69a4581b4aa8fe5bb Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 23 Jan 2026 14:03:21 -0800
Subject: [PATCH 06/79] Add health endpoint to dynamic text gen server (#3009)

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../endpoints/__init__.py                     |  3 +-
 .../endpoints/health.py                       | 38 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/health.py

diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py
index 1945fd10dba..f2b0661dace 100644
--- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py
+++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py
@@ -4,7 +4,8 @@
 try:
     from .chat_completions import bp as ChatCompletions
     from .completions import bp as Completions
+    from .health import bp as Health
 
-    __all__ = [Completions, ChatCompletions]
+    __all__ = [Completions, ChatCompletions, Health]
 except ImportError:
     __all__ = []
diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/health.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/health.py
new file mode 100644
index 00000000000..a9d0a678b44
--- /dev/null
+++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/health.py
@@ -0,0 +1,38 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+try:
+    from flask import Blueprint, current_app, jsonify
+
+    bp = Blueprint('health_api', __name__)
+
+    @bp.route('/health', methods=['GET'])
+    @bp.route('/v1/health', methods=['GET'])
+    async def health():
+        """
+        Handles GET requests for service health.
+        Checks if the inference client is initialized and reachable.
+        """
+        status_response = {"status": "ok", "service": "Megatron Inference Server", "ready": False}
+
+        try:
+            client = current_app.config.get('client')
+
+            if client is not None:
+                status_response["ready"] = True
+                return jsonify(status_response), 200
+            else:
+                logger.error("Health check failed: Client not found in app config.")
+                status_response["status"] = "error"
+                status_response["details"] = "Inference client not initialized"
+                return jsonify(status_response), 503
+
+        except Exception as e:
+            logger.error(f"Health check failed with exception: {e}")
+            return jsonify({"status": "error", "details": str(e)}), 500
+
+except ImportError as e:
+    logger.warning(f"Could not import flask: {e}")

From 359330160da0ecf855e905a33ea781f24cf286e6 Mon Sep 17 00:00:00 2001
From: GitHub Actions <github-actions[bot]@users.noreply.github.com>
Date: Sat, 24 Jan 2026 00:11:33 +0000
Subject: [PATCH 07/79] Update copy-pr-bot.yaml [skip ci]

---
 .github/copy-pr-bot.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml
index 8e92aabe027..2ece09a999f 100644
--- a/.github/copy-pr-bot.yaml
+++ b/.github/copy-pr-bot.yaml
@@ -1,4 +1,4 @@
 enabled: true
 auto_sync_draft: false
 auto_sync_ready: true
-trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "pablo-garay", "parthmannan", "pthombre", "rogerwaleffe", "sanandaraj5597", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]
+trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "pablo-garay", "parthmannan", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]

From 30dea5d8f0213d2799911bafef801bbdfadd8f17 Mon Sep 17 00:00:00 2001
From: Dong Hyuk Chang <donghyukc@nvidia.com>
Date: Fri, 23 Jan 2026 18:32:12 -0500
Subject: [PATCH 08/79] ci: Skip test_precision_aware_optimizer (#3062)

Signed-off-by: Dong Hyuk Chang <donghyukc@nvidia.com>
---
 tests/unit_tests/test_optimizer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py
index 1f5bbc3f14c..6b1da8c4e3f 100644
--- a/tests/unit_tests/test_optimizer.py
+++ b/tests/unit_tests/test_optimizer.py
@@ -384,6 +384,7 @@ def test_precision_aware_fused_adam():
     "moment_dtype",
     [torch.float32, torch.float16, torch.bfloat16, torch.uint8],
 )
+@pytest.mark.skip(reason="inconsistent ci test runs resulting in NCCL errors")
 def test_precision_aware_optimizer(
     precision: str,
     main_params_dtype: torch.dtype,

From 485ed1817d18e25817cb26f0813e6bd1e608eeb7 Mon Sep 17 00:00:00 2001
From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com>
Date: Fri, 23 Jan 2026 16:25:46 -1000
Subject: [PATCH 09/79] Support multimodule communication (#2031)

Co-authored-by: shifangx <shifangx@nvidia.com>
Co-authored-by: Mcore Bot <mcore-bot@nvidia.com>
---
 .../pipeline_parallel/bridge_communicator.py  |   3 -
 .../multimodule_communicator.py               | 531 ++++++++++++
 .../test_multimodule_communicator.py          | 782 ++++++++++++++++++
 3 files changed, 1313 insertions(+), 3 deletions(-)
 create mode 100644 megatron/core/pipeline_parallel/multimodule_communicator.py
 create mode 100644 tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py

diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py
index a67ded6bf08..f1e74a2f16d 100644
--- a/megatron/core/pipeline_parallel/bridge_communicator.py
+++ b/megatron/core/pipeline_parallel/bridge_communicator.py
@@ -628,9 +628,6 @@ def send_forward_recv_backward(
                 dist.broadcast(
                     shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
                 )
-                dist.broadcast(
-                    shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
-                )
 
                 # Broadcast the tensors to all ranks in the group
                 dist.broadcast(
diff --git a/megatron/core/pipeline_parallel/multimodule_communicator.py b/megatron/core/pipeline_parallel/multimodule_communicator.py
new file mode 100644
index 00000000000..1e8da3468e2
--- /dev/null
+++ b/megatron/core/pipeline_parallel/multimodule_communicator.py
@@ -0,0 +1,531 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.distributed as dist
+
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.pipeline_parallel.bridge_communicator import BridgeCommunicator
+from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator
+
+# Types
+Shape = Union[List[int], torch.Size]
+
+
+@dataclass
+class RankModuleInfo:
+    """Information about a rank in a module.
+
+    Attributes:
+        pp_rank: The stage index of the current rank within the module's pipeline.
+        pp_size: The total number of pipeline stages (ranks) in the module.
+        p2p_communicator: Intra-module point-to-point communicator.
+        bridge_comms_as_src_module: Bridge communicators for outgoing connections
+            from this module to downstream modules. One module may have multiple
+            bridge communicators if it has multiple outgoing connections.
+        bridge_comms_as_dest_module: Bridge communicators for incoming connections
+            to this module from upstream modules. One module may have multiple
+            bridge communicators if it has multiple incoming connections.
+        is_source_stage: True if this rank is at the absolute first stage in the
+            overall model (no incoming connections).
+        is_terminal_stage: True if this rank is at the absolute last stage in the
+            overall model (no outgoing connections).
+    """
+
+    pp_rank: int
+    pp_size: int
+    p2p_communicator: Optional[P2PCommunicator]
+    bridge_comms_as_src_module: Optional[List[BridgeCommunicator]]
+    bridge_comms_as_dest_module: Optional[List[BridgeCommunicator]]
+    is_source_stage: Optional[bool] = True
+    is_terminal_stage: Optional[bool] = True
+
+
+class MultiModulePipelineCommunicator:
+    """Communicator for a multi-module pipeline."""
+
+    def __init__(
+        self,
+        module_to_grid_map: Dict[str, HyperCommGrid],
+        topology: Dict[str, List[str]],
+        config: ModelParallelConfig,
+        dim_mapping: Dict[str, List[int]] = None,
+    ):
+        """
+        Initialize the MultiModulePipelineCommunicator.
+
+        Args:
+            module_to_grid_map (dict): A dictionary mapping module names to HyperCommGrids.
+                Example:
+                    module_to_grid_map = {
+                        'image_encoder': image_encoder_grid,
+                        'audio_encoder': audio_encoder_grid,
+                        'llm': llm_grid,
+                        'generator': generator_grid
+                    }
+            topology (dict): A dictionary mapping module names to lists of outgoing modules.
+                Example:
+                    topology = {
+                        'image_encoder': ['llm'],
+                        'audio_encoder': ['llm'],
+                        'llm': ['generator'],
+                        'generator': []
+                    }
+            config (ModelParallelConfig): A ModelParallelConfig object.
+            dim_mapping (Dict[str, List[int]]): Dimension mapping for sequence, batch, hidden.
+                Example:
+                    dim_mapping = {'s': 0, 'h': 2, 'b': 1}
+                Default: None
+        """
+        self.module_to_grid_map = module_to_grid_map
+        self.topology = topology
+        self.config = config
+        self.dim_mapping = dim_mapping
+        self.current_rank = dist.get_rank()
+
+        # Build bridge communicators for all modules
+        self.bridge_comms = []
+        self._build_bridge_comms()
+
+        self.rank_module_map = {}
+        self._build_rank_module_info_map()
+
+    def _build_bridge_comms(self):
+        """Construct and store BridgeCommunicator objects that describe the outgoing
+        communication relationships for all of the modules.
+        """
+        for src_module_name, src_grid in self.module_to_grid_map.items():
+            for dest_module_name in self.topology[src_module_name]:
+                dest_grid = self.module_to_grid_map[dest_module_name]
+                bridge_comm = BridgeCommunicator(
+                    src_grid=src_grid,
+                    dest_grid=dest_grid,
+                    dim_mapping=self.dim_mapping,
+                    comm_dtype=self.config.pipeline_dtype,
+                    src_module_name=src_module_name,
+                    dest_module_name=dest_module_name,
+                )
+                self.bridge_comms.append(bridge_comm)
+
+    @property
+    def is_pp_first_stage(self):
+        """Return True if the current rank has the absolute first stage in the overall model.
+
+        The absolute first stage is defined as:
+        1. The current rank must be in the first PP stage (pp_rank == 0) of some module
+        2. That module must be a source module (no incoming connections in topology)
+        """
+        for module_name, rank_module_info in self.rank_module_map.items():
+            # Check if this rank is at the first PP stage of this module
+            if rank_module_info.pp_rank == 0:
+                # Check if this module is a source module (no incoming connections)
+                if self._is_source_module(module_name):
+                    return True
+        return False
+
+    @property
+    def is_pp_last_stage(self):
+        """Return True if the current rank has the absolute last stage in the overall model.
+
+        The absolute last stage is defined as:
+        1. The current rank must be in the last PP stage of some module
+        2. That module must be a sink module (no outgoing connections in topology)
+        """
+        for module_name, rank_module_info in self.rank_module_map.items():
+            # Check if this rank is at the last PP stage of this module
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # Check if this module is a sink module (no outgoing connections)
+                if self._is_sink_module(module_name):
+                    return True
+        return False
+
+    def _is_source_module(self, module_name: str) -> bool:
+        """Check if a module is a source module (has no incoming connections)."""
+        # A module is a source if no other module lists it as a destination
+        for src_module, dest_modules in self.topology.items():
+            if module_name in dest_modules:
+                return False
+        return True
+
+    def _is_sink_module(self, module_name: str) -> bool:
+        """Check if a module is a sink module (has no outgoing connections)."""
+        return len(self.topology.get(module_name, [])) == 0
+
+    def is_current_rank_in_grid(self, grid: HyperCommGrid) -> bool:
+        """Check if the current rank is in the grid."""
+        return grid.rank_offset <= self.current_rank < grid.rank_offset + grid.size
+
+    @property
+    def num_warmup_microbatches(self):
+        """Calculate the number of warmup microbatches for the current rank.
+
+        Uses the same simple logic as P2PCommunicator:
+        total_pipeline_stages - current_rank_stage - 1
+
+        Returns:
+            int: Number of warmup microbatches for this rank
+        """
+        # Get total pipeline depth across all modules
+        total_stages = self.compute_total_pipeline_stages(self.topology, self.module_to_grid_map)
+
+        # Get current rank's position in the overall pipeline (0-indexed)
+        # Use compute_total_pipeline_stages with current rank to get cumulative position
+        if self.rank_module_map:
+            # Take the first module this rank belongs to
+            # TODO: ykarnati - improve this logic.
+            module_name = next(iter(self.rank_module_map.keys()))
+            current_stage = (
+                self.compute_total_pipeline_stages(
+                    self.topology,
+                    self.module_to_grid_map,
+                    rank=self.current_rank,
+                    module_name=module_name,
+                )
+                - 1
+            )  # Convert from 1-indexed to 0-indexed
+        else:
+            current_stage = 0
+
+        assert (
+            current_stage <= total_stages
+        ), f"current_stage: {current_stage} is greater than total_stages: {total_stages}"
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"current_stage: {current_stage} total_stages: {total_stages} "
+            f"num_warmup_microbatches: {total_stages - current_stage - 1}"
+        )
+        return total_stages - current_stage - 1
+
+    def _build_rank_module_info_map(self):
+        """For each module in the current rank, initialize the P2P communicator
+        and build the bridge communicator info for the module.
+        Each rank may hold multiple modules when colocated.
+        """
+        for module_name, module_grid in self.module_to_grid_map.items():
+            if self.is_current_rank_in_grid(module_grid):
+                # Initialize P2P communicator
+                pp_group = module_grid.get_pg('pp')
+                p2p_comm = P2PCommunicator(pp_group, self.config)
+                pp_size = dist.get_world_size(pp_group)
+                rank_in_pp_group = dist.get_group_rank(pp_group, self.current_rank)
+                pp_rank = rank_in_pp_group % pp_size
+
+                bridge_comms_as_dest_module = []
+                bridge_comms_as_src_module = []
+                # If first stage, check if the module has any incoming modules
+                # If so, initialize bridge communicator
+                if pp_rank == 0:
+                    for bridge_comm in self.bridge_comms:
+                        if (
+                            bridge_comm.is_current_rank_in_grid(bridge_comm.dest_grid)
+                            and bridge_comm.dest_module_name == module_name
+                        ):
+                            bridge_comms_as_dest_module.append(bridge_comm)
+                # If last stage, check if the module has any outgoing modules
+                # If so, initialize bridge communicator
+                if pp_rank == pp_size - 1:
+                    for bridge_comm in self.bridge_comms:
+                        if (
+                            bridge_comm.is_current_rank_in_grid(bridge_comm.src_grid)
+                            and bridge_comm.src_module_name == module_name
+                        ):
+                            bridge_comms_as_src_module.append(bridge_comm)
+                # Build RankModuleInfo for the module
+                rank_module_info = RankModuleInfo(
+                    pp_rank=pp_rank,
+                    pp_size=pp_size,
+                    p2p_communicator=p2p_comm,
+                    bridge_comms_as_dest_module=bridge_comms_as_dest_module,
+                    bridge_comms_as_src_module=bridge_comms_as_src_module,
+                )
+                self.rank_module_map[module_name] = rank_module_info
+
+    def recv_forward(
+        self, tensor_shape: Optional[Shape] = None, is_first_stage: bool = False
+    ) -> Dict[str, torch.Tensor]:
+        """Receive forward activation tensor.
+
+        Args:
+            tensor_shape: Expected activation tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[receive_forward] tensors_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
+        )
+        input_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+
+            if rank_module_info.pp_rank == 0:
+                # If first stage, and has incoming modules, receive forward activation
+                # from incoming modules.
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    input_dict[bridge_comm.src_module_name] = bridge_comm.recv_forward()
+            else:
+                # If not first stage, receive forward activation tensor from P2P communicator.
+                input_dict[module_name] = rank_module_info.p2p_communicator.recv_forward(
+                    tensor_shapes=tensor_shape, is_first_stage=False
+                )
+        return input_dict
+
+    def send_forward(self, output_dict: Dict[str, torch.Tensor], is_last_stage: bool = False):
+        """Send forward activation tensor.
+
+        Args:
+            output_dict: A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_forward] output_dict keys: {output_dict.keys()}, is_last_stage: {is_last_stage}"
+        )
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has outgoing modules, send forward activation
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    bridge_comm.send_forward(output_dict[module_name])
+            else:
+                # If not last stage, send forward activation by using P2P communicator.
+                rank_module_info.p2p_communicator.send_forward(
+                    output_dict[module_name], is_last_stage=False
+                )
+
+    def send_forward_recv_backward(
+        self,
+        output_dict: Dict[str, torch.Tensor],
+        tensor_shape: Optional[Shape] = None,
+        is_last_stage: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Send forward activation tensor and receive backward activation tensor.
+
+        Args:
+            output_dict: A dictionary mapping module names to tensors.
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_forward_recv_backward] output_dict keys: {output_dict.keys()}, "
+            f"tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
+        )
+        grad_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has outgoing modules, send forward activation and
+                # receive backward gradient by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    grad_dict[bridge_comm.src_module_name] = bridge_comm.send_forward_recv_backward(
+                        output_dict[module_name]
+                    )
+            else:
+                # If not last stage, send forward activation and receive backward gradient
+                # by using P2P communicator.
+                grad_dict[module_name] = (
+                    rank_module_info.p2p_communicator.send_forward_recv_backward(
+                        output_dict[module_name], tensor_shapes=tensor_shape, is_last_stage=False
+                    )
+                )
+        return grad_dict
+
+    def send_backward_recv_forward(
+        self,
+        grad_dict: Dict[str, torch.Tensor],
+        tensor_shape: Optional[Shape] = None,
+        is_first_stage: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Send backward activation tensor and receive forward activation tensor.
+
+        Args:
+            grad_dict: A dictionary mapping module names to tensors.
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_backward_recv_forward] grad_dict keys: {grad_dict.keys()}, "
+            f"tensor_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
+        )
+        input_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == 0:
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    # If first stage, and has incoming modules, send backward gradient and
+                    # receive forward activation by using bridge communicator.
+                    input_dict[bridge_comm.src_module_name] = (
+                        bridge_comm.send_backward_recv_forward(
+                            grad_dict[bridge_comm.src_module_name]
+                        )
+                    )
+            else:
+                # If not first stage, send backward gradient and receive forward activation
+                # by using P2P communicator.
+                input_dict[module_name] = (
+                    rank_module_info.p2p_communicator.send_backward_recv_forward(
+                        grad_dict[module_name], tensor_shapes=tensor_shape, is_first_stage=False
+                    )
+                )
+        return input_dict
+
+    def recv_backward(
+        self, tensor_shape: Optional[Shape] = None, is_last_stage: bool = False
+    ) -> Dict[str, torch.Tensor]:
+        """Receive backward activation tensor.
+
+        Args:
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[recv_backward] tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
+        )
+        grad_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has incoming modules, receive backward gradient
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    grad_dict[bridge_comm.src_module_name] = bridge_comm.recv_backward()
+            else:
+                # If not last stage, receive backward gradient by using P2P communicator.
+                grad_dict[module_name] = rank_module_info.p2p_communicator.recv_backward(
+                    tensor_shapes=tensor_shape, is_last_stage=False
+                )
+        return grad_dict
+
+    def send_backward(self, grad_dict: Dict[str, torch.Tensor], is_first_stage: bool = False):
+        """Send backward activation tensor.
+
+        Args:
+            grad_dict: A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_backward] grad_dict keys: {grad_dict.keys()}, is_first_stage: {is_first_stage}"
+        )
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == 0:
+                # If first stage, and has incoming modules, send backward activation
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    bridge_comm.send_backward(grad_dict[bridge_comm.src_module_name])
+            else:
+                # If not first stage, send backward activation by using P2P communicator.
+                rank_module_info.p2p_communicator.send_backward(
+                    grad_dict[module_name], is_first_stage=False
+                )
+
+    @staticmethod
+    def compute_total_pipeline_stages(
+        topology: Dict[str, List[str]],
+        module_to_grid_map: Dict[str, HyperCommGrid],
+        rank: Optional[int] = None,
+        module_name: Optional[str] = None,
+    ) -> int:
+        """Compute the total number of pipeline stages across a multi-module chain.
+
+        Interprets ``topology`` as a directed acyclic graph (DAG) where nodes are modules
+        and edges indicate forward data flow from source to destination modules. Each node
+        is assigned a weight equal to its pipeline parallel size (number of PP stages).
+
+        The total number of stages is defined as the length of the longest path in this DAG
+        under node weights.
+
+        If ``rank`` is None (default), returns the maximum over all terminal (sink) modules of
+        the sum of PP sizes along a path ending at that terminal. For example, given:
+
+            image_encoder ->\
+                              -> llm -> generator
+            audio_encoder  ->/
+
+        the total is: max(pp(image_encoder), pp(audio_encoder)) + pp(llm) + pp(generator).
+
+        If ``rank`` is provided, the result is the total number of pipeline stages up to (and
+        including) the PP stage that ``rank`` occupies inside its module. In this case, the
+        weight of the target module equals (pp_rank_index(rank) + 1) instead of the module's
+        full PP size; other modules still contribute their full PP sizes. If the rank belongs to
+        multiple modules (colocation), pass ``module_name`` to disambiguate; otherwise the
+        maximum across all candidate modules containing the rank is returned.
+
+        Args:
+            topology: Mapping from a module to its list of outgoing modules.
+            module_to_grid_map: Mapping from module name to its ``HyperCommGrid``.
+
+        Returns:
+            The total number of pipeline stages along the longest path given the constraints.
+
+        Raises:
+            ValueError: If the topology contains cycles; or has no terminal nodes when
+                ``rank`` is None
+        """
+        nodes = set(module_to_grid_map.keys())
+        # Build adjacency and reverse-adjacency (predecessors).
+        adj: Dict[str, List[str]] = {node: list(topology.get(node, [])) for node in nodes}
+        preds: Dict[str, List[str]] = {node: [] for node in nodes}
+        for src, outs in adj.items():
+            for dst in outs:
+                preds[dst].append(src)
+
+        # Identify terminal nodes (no outgoing edges) for the rank=None case.
+        sinks = [node for node, outs in adj.items() if not outs]
+        if rank is None and not sinks:
+            raise ValueError(
+                "Topology must be a DAG with at least one terminal (no outgoing) module."
+            )
+
+        def pp_size(name: str) -> int:
+            grid = module_to_grid_map[name]
+            pp_dim_index = grid.dim_names.index('pp')
+            return grid.shape[pp_dim_index]
+
+        def partial_weight_for_target(target: str) -> Optional[int]:
+            if rank is None:
+                return None
+            grid = module_to_grid_map.get(target)
+            rank_groups = grid._gen_rank_enum(['pp'])
+            stage_index: Optional[int] = None
+            for group in rank_groups:
+                if rank in group:
+                    stage_index = group.index(rank)
+                    break
+            return stage_index + 1
+
+        def longest_path_to(target: str) -> int:
+            visiting = set()
+            partial = partial_weight_for_target(target)
+
+            def weight(name: str) -> int:
+                if partial is not None and name == target:
+                    return partial
+                return pp_size(name)
+
+            def dfs(node: str) -> int:
+                if node in visiting:
+                    raise ValueError("Topology contains cycles; expected a DAG.")
+                visiting.add(node)
+                best = 0
+                for p in preds.get(node, []):
+                    val = dfs(p)
+                    if val > best:
+                        best = val
+                visiting.remove(node)
+                return weight(node) + best
+
+            return dfs(target)
+
+        if rank is None:
+            return max(longest_path_to(sink) for sink in sinks)
+
+        return longest_path_to(module_name)
diff --git a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
new file mode 100644
index 00000000000..22f790cc0a9
--- /dev/null
+++ b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
@@ -0,0 +1,782 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import os
+import sys
+
+import pytest
+import torch
+import torch.distributed as dist
+from packaging import version
+
+from megatron.core import parallel_state
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.pipeline_parallel.multimodule_communicator import MultiModulePipelineCommunicator
+from tests.unit_tests.pipeline_parallel.test_bridge_communicator import (
+    _avg_params,
+    _create_transformer_block,
+    _get_pg_collection_from_grid,
+    create_hypercomm_grid,
+    get_transformer_block_and_grid,
+)
+from tests.unit_tests.test_utilities import Utils
+
+
+class TestMultiModulePipelineCommunicator:
+
+    @classmethod
+    def setup_class(cls):
+        """Set up distributed environment for the entire test class."""
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+        world_size = dist.get_world_size()
+        if world_size != 8:
+            pytest.skip(
+                f"These tests require 8 GPUs, but only {world_size} are available.",
+                allow_module_level=True,
+            )
+
+    def teardown_class(cls):
+        Utils.destroy_model_parallel()
+
+    def test_multimodule_communicator_init(self):
+        """Test MultiModulePipelineCommunicator initialization."""
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=2, cp=1, pp=1, dp=1)
+
+        # Define module-grid mapping
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        # Define module computation topology
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(bf16=True)
+        # Initialize communicator
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+        # Test attributes match expectations
+        assert mllm_comm.module_to_grid_map == module_to_grid_map
+        assert mllm_comm.topology == topology
+        assert mllm_comm.config == config
+        assert mllm_comm.current_rank == dist.get_rank()
+
+    def test_compute_total_pipeline_stages(self):
+        """Test compute_total_pipeline_stages for overall chain and until specific ranks."""
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Define module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+
+        # Overall total pipeline stages: max(1,1) + 2 + 1 = 4
+        total = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
+            topology, module_to_grid_map
+        )
+        assert total == 4
+
+        llm_pp_rank = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
+            topology, module_to_grid_map, rank=2, module_name='llm'
+        )
+        assert llm_pp_rank == 2
+
+    def test_send_forward_recv_forward(self):
+        """Test send_forward and recv_forward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate forward communication for each module
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder sends output forward
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder sends output forward
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                # LLM stage receives both image and audio outputs
+                input_dict = mllm_comm.recv_forward()
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+                mllm_comm.send_forward(output_dict)
+            else:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+                mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator module receives final LLM output
+            input_dict = mllm_comm.recv_forward()
+            assert input_dict['llm'].shape == (1, 32, 128)
+
+    def test_send_forward_recv_forward_with_different_pp_size(self):
+        """Test for the case when pp(image_encoder) != pp(audio_encoder)."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=2, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=4, tp=1, cp=1, pp=4, dp=1)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+        }
+        topology = {'image_encoder': ['llm'], 'audio_encoder': ['llm'], 'llm': []}
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate forward communication for each module
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            if dist.get_rank() == 0:
+                # Image encoder sends output forward
+                mllm_comm.send_forward(output_dict)
+            else:
+                # Image stage receives image outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 8, 128))
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder sends output forward
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+            if dist.get_rank() == 4:
+                # LLM stage receives both image and audio outputs
+                input_dict = mllm_comm.recv_forward()
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+                mllm_comm.send_forward(output_dict)
+            elif dist.get_rank() == 5 or dist.get_rank() == 6:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+                mllm_comm.send_forward(output_dict)
+            elif dist.get_rank() == 7:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+
+    def test_send_backward_recv_backward(self):
+        """Test send_backward and recv_backward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate backward communication for each module
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator sends gradient backward
+            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
+            mllm_comm.send_backward(grad_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                # LLM receives expanded gradient and sends backward
+                received_grad = mllm_comm.recv_backward()
+                assert received_grad['llm'].shape == (2, 32, 128)
+                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                mllm_comm.send_backward(grad_dict)
+            else:
+                # LLM receives gradient and sends backward to both image/audio encoders
+                received_grad = mllm_comm.recv_backward(tensor_shape=(2, 32, 128))
+                assert received_grad['llm'].shape == (2, 32, 128)
+                grad_dict = {
+                    'image_encoder': torch.randn(2, 8, 128).cuda(),
+                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
+                }
+                mllm_comm.send_backward(grad_dict)
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder receives its gradient
+            received_grad = mllm_comm.recv_backward()
+            assert received_grad['image_encoder'].shape == (2, 8, 128)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder receives its gradient
+            received_grad = mllm_comm.recv_backward()
+            assert received_grad['audio_encoder'].shape == (2, 16, 128)
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    def test_send_forward_recv_backward_send_backward_recv_forward(self):
+        """Test send_forward_recv_backward and send_backward_recv_forward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate bidirectional send/recv for forward and backward in pipeline
+
+        # Encoder stages send forward to the first stage of LLM, and receive backward from the first stage of LLM
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+            assert received_grad['image_encoder'].shape == (2, 8, 128)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+            assert received_grad['audio_encoder'].shape == (2, 16, 128)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                grad_dict = {
+                    'image_encoder': torch.randn(2, 8, 128).cuda(),
+                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
+                }
+                input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+
+        # First stage of LLM sends forward to the second stage of LLM, and receive backward from the second stage of LLM
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                received_grad = mllm_comm.send_forward_recv_backward(
+                    output_dict, tensor_shape=(2, 32, 128)
+                )
+                assert received_grad['llm'].shape == (2, 32, 128)
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                input_dict = mllm_comm.send_backward_recv_forward(
+                    grad_dict, tensor_shape=(2, 32, 128)
+                )
+                assert input_dict['llm'].shape == (2, 32, 128)
+
+        # Second stage of LLM sends forward to generator, and receive backward from generator
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+                assert received_grad['llm'].shape == (2, 32, 128)
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
+            input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
+            assert input_dict['llm'].shape == (1, 32, 128)
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    def test_send_forward_recv_forward_with_transformer_blocks(self):
+        """Test send_forward and recv_forward operations."""
+
+        # Set model/test dimensions for easier debugging and output comparison
+        hidden_size = 16
+        sequence_length = 2
+        micro_batch_size = 2
+
+        # For reproducibility, set a fixed seed
+        torch.manual_seed(12345)
+        dtype = torch.float32
+
+        # Create random input hidden states tensor
+        hidden_states = torch.randn(
+            (sequence_length, micro_batch_size, hidden_size), device="cuda"
+        ).to(dtype)
+        current_rank = dist.get_rank()
+
+        # ========== Initialize tensor model-parallel environment ==========
+        parallel_state_tp = 2
+        Utils.initialize_model_parallel(tensor_model_parallel_size=2)
+
+        # ========== Build reference 1D grid and transformer block for weight sharing ==========
+        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
+        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
+        ref_block = _create_transformer_block(
+            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
+        )
+        _avg_params(
+            ref_block, ref_grid.get_pg("dp")
+        )  # Ensure parameters are averaged across data parallel (DP)
+
+        # ========== Create different transformer blocks for each model stage ==========
+        # Image encoder
+        image_encoder_block, image_encoder_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=1,
+            grid_offset=0,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # Audio encoder
+        audio_encoder_block, audio_encoder_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=1,
+            grid_offset=1,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # LLM (Large Language Model) block with tensor & pipeline parallelism
+        llm_block, llm_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=2,
+            cp_size=1,
+            pp_size=2,
+            dp_size=1,
+            grid_offset=2,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # Generator block (final stage) with DP=2
+        generator_block, generator_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=2,
+            grid_offset=6,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        # ========== Define module-to-grid correspondence and pipeline topology ==========
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],  # image_encoder sends output to llm
+            'audio_encoder': ['llm'],  # audio_encoder sends output to llm
+            'llm': ['generator'],  # llm sends output to generator
+            'generator': [],  # generator is the final module
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        # Define dimension mapping for sequence, batch, hidden
+        dim_mapping = {'s': 0, 'h': 2, 'b': 1}
+        seq_dim = dim_mapping['s']
+
+        # Communication handler for multi-module pipeline (send/recv abstraction)
+        mllm_comm = MultiModulePipelineCommunicator(
+            module_to_grid_map, topology, config, dim_mapping=dim_mapping
+        )
+
+        # ========== Run actual distributed pipeline blocks (per process, depending on role) ==========
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder rank: run forward and send output
+            image_encoder_output = image_encoder_block(
+                hidden_states=hidden_states, attention_mask=None
+            )
+            output_dict = {'image_encoder': image_encoder_output}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder rank: run forward and send output
+            audio_encoder_output = audio_encoder_block(
+                hidden_states=hidden_states, attention_mask=None
+            )
+            output_dict = {'audio_encoder': audio_encoder_output}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                # LLM stage 0 (receives both image and audio, concatenates along seq_dim)
+                input_dict = mllm_comm.recv_forward()
+                llm_output = llm_block(
+                    hidden_states=torch.cat(
+                        [input_dict['image_encoder'], input_dict['audio_encoder']], dim=seq_dim
+                    ),
+                    attention_mask=None,
+                )
+                output_dict = {'llm': llm_output}
+                mllm_comm.send_forward(output_dict)
+            else:
+                # LLM stage 1 (receives output of previous LLM stage)
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(sequence_length * 2, micro_batch_size, hidden_size)
+                )
+                llm_output = llm_block(hidden_states=input_dict['llm'], attention_mask=None)
+                output_dict = {'llm': llm_output}
+                mllm_comm.send_forward(output_dict)
+
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator block: only receives from llm and runs forward
+            input_dict = mllm_comm.recv_forward()
+            generator_output = generator_block(hidden_states=input_dict['llm'], attention_mask=None)
+
+        # ========== Build a reference (serial/global) pipeline for correctness checking ==========
+        global_image_encoder_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_audio_encoder_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_llm_block_pp_rank_0, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_llm_block_pp_rank_1, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_generator_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        # Run each stage sequentially as a global pipeline (for truth)
+        global_image_encoder_output = global_image_encoder_block(
+            hidden_states=hidden_states, attention_mask=None
+        )
+        global_audio_encoder_output = global_audio_encoder_block(
+            hidden_states=hidden_states, attention_mask=None
+        )
+        # Compare output between global and distributed blocks for image/audio stage
+        if current_rank == 0:
+            torch.testing.assert_close(
+                global_image_encoder_output, image_encoder_output, rtol=1e-3, atol=1e-3
+            )
+        if current_rank == 1:
+            torch.testing.assert_close(
+                global_audio_encoder_output, audio_encoder_output, rtol=1e-3, atol=1e-3
+            )
+
+        # Feed outputs to LLM stages (emulate pipeline cut with concatenation)
+        global_llm_input = torch.cat(
+            [global_image_encoder_output, global_audio_encoder_output], dim=seq_dim
+        )
+        global_llm_pp_rank_0_output = global_llm_block_pp_rank_0(
+            hidden_states=global_llm_input, attention_mask=None
+        )
+        if current_rank == 2 or current_rank == 3:
+            torch.testing.assert_close(
+                global_llm_pp_rank_0_output, llm_output, rtol=1e-3, atol=1e-3
+            )
+        global_llm_pp_rank_1_output = global_llm_block_pp_rank_1(
+            hidden_states=global_llm_pp_rank_0_output, attention_mask=None
+        )
+        if current_rank == 4 or current_rank == 5:
+            torch.testing.assert_close(
+                global_llm_pp_rank_1_output, llm_output, rtol=1e-3, atol=1e-3
+            )
+
+        # Generator output and comparison to distributed output (for each DP chunk)
+        global_generator_block_output = global_generator_block(
+            hidden_states=global_llm_pp_rank_1_output, attention_mask=None
+        )
+        global_generator_block_chunks = torch.split(
+            global_generator_block_output, global_generator_block_output.shape[1] // 2, dim=1
+        )
+        if current_rank == 6:
+            torch.testing.assert_close(
+                global_generator_block_chunks[0], generator_output, rtol=1e-3, atol=1e-3
+            )
+        if current_rank == 7:
+            torch.testing.assert_close(
+                global_generator_block_chunks[1], generator_output, rtol=1e-3, atol=1e-3
+            )
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    @pytest.mark.parametrize(
+        "grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp",
+        [
+            (2, 1, 1, 2, 1, 1, 2),  # TP2PP1DP1 to TP2PP1DP1
+            (2, 1, 1, 2, 2, 1, 2),  # TP2PP1DP1 to TP2PP2DP1
+            (2, 2, 1, 2, 2, 1, 2),  # TP2PP2DP1 to TP2PP2DP1
+            (4, 1, 1, 4, 1, 1, 4),  # TP4DP1 to TP4DP1
+            (2, 1, 2, 4, 1, 1, 2),  # TP2DP2 to TP4DP1
+            (4, 1, 1, 2, 1, 2, 2),  # TP4DP1 to TP2DP2
+            (2, 1, 2, 1, 1, 4, 2),  # TP2DP2 to TP1DP4
+        ],
+    )
+    def test_send_forward_recv_forward_with_transformer_blocks_and_different_parallelisms(
+        self, grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp
+    ):
+        """Test bridge communicator with two transformer blocks having different process group configurations."""
+        # Model and input configuration
+        hidden_size = 16
+        sequence_length = 2
+        micro_batch_size = 8
+        torch.manual_seed(12345)
+        dtype = torch.float32
+
+        # Create random input tensor on CUDA
+        hidden_states = torch.randn(
+            (sequence_length, micro_batch_size, hidden_size), device="cuda"
+        ).to(dtype)
+        hidden_states_ref = hidden_states.clone()
+        current_rank = dist.get_rank()
+
+        # Initialize model parallel with desired TP
+        Utils.initialize_model_parallel(tensor_model_parallel_size=parallel_state_tp)
+
+        # Build a reference grid and block for parameter sharing & DP averaging
+        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
+        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
+        ref_block = _create_transformer_block(
+            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
+        )
+        _avg_params(
+            ref_block, ref_grid.get_pg("dp")
+        )  # Synchronize parameters across DP for reproducibility
+
+        # ====== Create two transformer block+grid pairs with different TP/DP settings ======
+        block_grid_1, grid_1 = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=grid1_tp,
+            pp_size=grid1_pp,
+            dp_size=grid1_dp,
+            grid_offset=0,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        block_grid_2, grid_2 = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=grid2_tp,
+            pp_size=grid2_pp,
+            dp_size=grid2_dp,
+            grid_offset=grid_1.size,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        dist.barrier()  # Synchronize ranks before communication
+
+        # Module-grid map and pipeline communication topology
+        module_to_grid_map = {'image_encoder': grid_1, 'llm': grid_2}
+        topology = {
+            'image_encoder': ['llm'],  # image_encoder sends forward results to llm
+            'llm': [],  # llm is the last stage here
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(
+            module_to_grid_map, topology, config, dim_mapping={'s': 0, 'h': 2, 'b': 1}
+        )
+
+        output_grid_2 = None
+        # If current rank is in the first grid, run first block and send output
+        if grid_1 is not None and mllm_comm.is_current_rank_in_grid(grid_1):
+            rank_module_info = mllm_comm.rank_module_map['image_encoder']
+            if rank_module_info.pp_rank == 0:
+                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'image_encoder': hidden_states})
+            else:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(sequence_length, micro_batch_size, hidden_size)
+                )
+                hidden_states = input_dict['image_encoder']
+                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'image_encoder': hidden_states})
+
+        # If current rank is in second grid, receive and run the second block
+        if grid_2 is not None and mllm_comm.is_current_rank_in_grid(grid_2):
+            rank_module_info = mllm_comm.rank_module_map['llm']
+            if rank_module_info.pp_rank == 0:
+                input_dict = mllm_comm.recv_forward()
+                hidden_states = input_dict['image_encoder']
+                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+                if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                    output_grid_2 = hidden_states
+                else:
+                    mllm_comm.send_forward({'llm': hidden_states})
+            elif rank_module_info.pp_rank < rank_module_info.pp_size - 1:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(
+                        sequence_length,
+                        (grid1_dp * micro_batch_size) // grid2_dp,
+                        hidden_size,
+                    )
+                )
+                hidden_states = input_dict['llm']
+                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'llm': hidden_states})
+            else:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(
+                        sequence_length,
+                        (grid1_dp * micro_batch_size) // grid2_dp,
+                        hidden_size,
+                    )
+                )
+                hidden_states = input_dict['llm']
+                output_grid_2 = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+
+                # Compute expected output shape based on change in DP size (chunk/expand batch dimension appropriately)
+                factor = max(grid1_dp, grid2_dp) // min(grid1_dp, grid2_dp)
+                expected_output_shape = (
+                    sequence_length,
+                    (
+                        micro_batch_size * factor
+                        if grid1_dp > grid2_dp
+                        else micro_batch_size // factor
+                    ),
+                    hidden_size,
+                )
+                assert (
+                    output_grid_2.shape == expected_output_shape
+                ), f"Output2 shape mismatch: {output_grid_2.shape}"
+
+        # ====== Reference: global (replicated) pipeline forward for correctness checking ======
+        global_block_1, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_block_2, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        for i in range(grid1_pp):
+            hidden_states_ref = global_block_1(hidden_states=hidden_states_ref, attention_mask=None)
+
+        for i in range(grid2_pp):
+            hidden_states_ref = global_block_2(hidden_states=hidden_states_ref, attention_mask=None)
+
+        # Output comparison under different DP compositions between grids
+        if (
+            grid_2 is not None
+            and mllm_comm.is_current_rank_in_grid(grid_2)
+            and rank_module_info.pp_rank == rank_module_info.pp_size - 1
+        ):
+            if grid1_dp == grid2_dp:
+                # DP size matches: all outputs directly compared
+                torch.testing.assert_close(hidden_states_ref, output_grid_2, rtol=1e-3, atol=1e-3)
+            elif grid1_dp < grid2_dp:
+                # If grid2 expands DP: each output_grid_2 chunk corresponds to a split of the reference output
+                grid2_dp_ranks = grid_2._gen_rank_enum([x for x in grid_2.dim_names if x != "dp"])
+                global_block_2_chunks = torch.split(
+                    hidden_states_ref, hidden_states_ref.shape[1] // (grid2_dp // grid1_dp), dim=1
+                )
+                relevant_chunk = None
+                for i, dp_ranks in enumerate(grid2_dp_ranks):
+                    if current_rank in dp_ranks:
+                        relevant_chunk = global_block_2_chunks[i % len(global_block_2_chunks)]
+                torch.testing.assert_close(relevant_chunk, output_grid_2, rtol=1e-3, atol=1e-3)
+            else:
+                # If DP shrinks (grid1_dp > grid2_dp): just compare the relevant first chunk
+                output_grid_2_first_chunk = torch.chunk(output_grid_2, grid1_dp // grid2_dp, dim=1)[
+                    0
+                ]
+                torch.testing.assert_close(
+                    hidden_states_ref, output_grid_2_first_chunk, rtol=1e-3, atol=1e-3
+                )

From 0972f020ddf36d0e7583343e664536cb95248f5f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Sat, 24 Jan 2026 09:28:35 +0100
Subject: [PATCH 10/79] Revert "Support multimodule communication (#2031)"
 (#3068)

---
 .../pipeline_parallel/bridge_communicator.py  |   3 +
 .../multimodule_communicator.py               | 531 ------------
 .../test_multimodule_communicator.py          | 782 ------------------
 3 files changed, 3 insertions(+), 1313 deletions(-)
 delete mode 100644 megatron/core/pipeline_parallel/multimodule_communicator.py
 delete mode 100644 tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py

diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py
index f1e74a2f16d..a67ded6bf08 100644
--- a/megatron/core/pipeline_parallel/bridge_communicator.py
+++ b/megatron/core/pipeline_parallel/bridge_communicator.py
@@ -628,6 +628,9 @@ def send_forward_recv_backward(
                 dist.broadcast(
                     shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
                 )
+                dist.broadcast(
+                    shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
+                )
 
                 # Broadcast the tensors to all ranks in the group
                 dist.broadcast(
diff --git a/megatron/core/pipeline_parallel/multimodule_communicator.py b/megatron/core/pipeline_parallel/multimodule_communicator.py
deleted file mode 100644
index 1e8da3468e2..00000000000
--- a/megatron/core/pipeline_parallel/multimodule_communicator.py
+++ /dev/null
@@ -1,531 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-import logging
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Union
-
-import torch
-import torch.distributed as dist
-
-from megatron.core.hyper_comm_grid import HyperCommGrid
-from megatron.core.model_parallel_config import ModelParallelConfig
-from megatron.core.pipeline_parallel.bridge_communicator import BridgeCommunicator
-from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator
-
-# Types
-Shape = Union[List[int], torch.Size]
-
-
-@dataclass
-class RankModuleInfo:
-    """Information about a rank in a module.
-
-    Attributes:
-        pp_rank: The stage index of the current rank within the module's pipeline.
-        pp_size: The total number of pipeline stages (ranks) in the module.
-        p2p_communicator: Intra-module point-to-point communicator.
-        bridge_comms_as_src_module: Bridge communicators for outgoing connections
-            from this module to downstream modules. One module may have multiple
-            bridge communicators if it has multiple outgoing connections.
-        bridge_comms_as_dest_module: Bridge communicators for incoming connections
-            to this module from upstream modules. One module may have multiple
-            bridge communicators if it has multiple incoming connections.
-        is_source_stage: True if this rank is at the absolute first stage in the
-            overall model (no incoming connections).
-        is_terminal_stage: True if this rank is at the absolute last stage in the
-            overall model (no outgoing connections).
-    """
-
-    pp_rank: int
-    pp_size: int
-    p2p_communicator: Optional[P2PCommunicator]
-    bridge_comms_as_src_module: Optional[List[BridgeCommunicator]]
-    bridge_comms_as_dest_module: Optional[List[BridgeCommunicator]]
-    is_source_stage: Optional[bool] = True
-    is_terminal_stage: Optional[bool] = True
-
-
-class MultiModulePipelineCommunicator:
-    """Communicator for a multi-module pipeline."""
-
-    def __init__(
-        self,
-        module_to_grid_map: Dict[str, HyperCommGrid],
-        topology: Dict[str, List[str]],
-        config: ModelParallelConfig,
-        dim_mapping: Dict[str, List[int]] = None,
-    ):
-        """
-        Initialize the MultiModulePipelineCommunicator.
-
-        Args:
-            module_to_grid_map (dict): A dictionary mapping module names to HyperCommGrids.
-                Example:
-                    module_to_grid_map = {
-                        'image_encoder': image_encoder_grid,
-                        'audio_encoder': audio_encoder_grid,
-                        'llm': llm_grid,
-                        'generator': generator_grid
-                    }
-            topology (dict): A dictionary mapping module names to lists of outgoing modules.
-                Example:
-                    topology = {
-                        'image_encoder': ['llm'],
-                        'audio_encoder': ['llm'],
-                        'llm': ['generator'],
-                        'generator': []
-                    }
-            config (ModelParallelConfig): A ModelParallelConfig object.
-            dim_mapping (Dict[str, List[int]]): Dimension mapping for sequence, batch, hidden.
-                Example:
-                    dim_mapping = {'s': 0, 'h': 2, 'b': 1}
-                Default: None
-        """
-        self.module_to_grid_map = module_to_grid_map
-        self.topology = topology
-        self.config = config
-        self.dim_mapping = dim_mapping
-        self.current_rank = dist.get_rank()
-
-        # Build bridge communicators for all modules
-        self.bridge_comms = []
-        self._build_bridge_comms()
-
-        self.rank_module_map = {}
-        self._build_rank_module_info_map()
-
-    def _build_bridge_comms(self):
-        """Construct and store BridgeCommunicator objects that describe the outgoing
-        communication relationships for all of the modules.
-        """
-        for src_module_name, src_grid in self.module_to_grid_map.items():
-            for dest_module_name in self.topology[src_module_name]:
-                dest_grid = self.module_to_grid_map[dest_module_name]
-                bridge_comm = BridgeCommunicator(
-                    src_grid=src_grid,
-                    dest_grid=dest_grid,
-                    dim_mapping=self.dim_mapping,
-                    comm_dtype=self.config.pipeline_dtype,
-                    src_module_name=src_module_name,
-                    dest_module_name=dest_module_name,
-                )
-                self.bridge_comms.append(bridge_comm)
-
-    @property
-    def is_pp_first_stage(self):
-        """Return True if the current rank has the absolute first stage in the overall model.
-
-        The absolute first stage is defined as:
-        1. The current rank must be in the first PP stage (pp_rank == 0) of some module
-        2. That module must be a source module (no incoming connections in topology)
-        """
-        for module_name, rank_module_info in self.rank_module_map.items():
-            # Check if this rank is at the first PP stage of this module
-            if rank_module_info.pp_rank == 0:
-                # Check if this module is a source module (no incoming connections)
-                if self._is_source_module(module_name):
-                    return True
-        return False
-
-    @property
-    def is_pp_last_stage(self):
-        """Return True if the current rank has the absolute last stage in the overall model.
-
-        The absolute last stage is defined as:
-        1. The current rank must be in the last PP stage of some module
-        2. That module must be a sink module (no outgoing connections in topology)
-        """
-        for module_name, rank_module_info in self.rank_module_map.items():
-            # Check if this rank is at the last PP stage of this module
-            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
-                # Check if this module is a sink module (no outgoing connections)
-                if self._is_sink_module(module_name):
-                    return True
-        return False
-
-    def _is_source_module(self, module_name: str) -> bool:
-        """Check if a module is a source module (has no incoming connections)."""
-        # A module is a source if no other module lists it as a destination
-        for src_module, dest_modules in self.topology.items():
-            if module_name in dest_modules:
-                return False
-        return True
-
-    def _is_sink_module(self, module_name: str) -> bool:
-        """Check if a module is a sink module (has no outgoing connections)."""
-        return len(self.topology.get(module_name, [])) == 0
-
-    def is_current_rank_in_grid(self, grid: HyperCommGrid) -> bool:
-        """Check if the current rank is in the grid."""
-        return grid.rank_offset <= self.current_rank < grid.rank_offset + grid.size
-
-    @property
-    def num_warmup_microbatches(self):
-        """Calculate the number of warmup microbatches for the current rank.
-
-        Uses the same simple logic as P2PCommunicator:
-        total_pipeline_stages - current_rank_stage - 1
-
-        Returns:
-            int: Number of warmup microbatches for this rank
-        """
-        # Get total pipeline depth across all modules
-        total_stages = self.compute_total_pipeline_stages(self.topology, self.module_to_grid_map)
-
-        # Get current rank's position in the overall pipeline (0-indexed)
-        # Use compute_total_pipeline_stages with current rank to get cumulative position
-        if self.rank_module_map:
-            # Take the first module this rank belongs to
-            # TODO: ykarnati - improve this logic.
-            module_name = next(iter(self.rank_module_map.keys()))
-            current_stage = (
-                self.compute_total_pipeline_stages(
-                    self.topology,
-                    self.module_to_grid_map,
-                    rank=self.current_rank,
-                    module_name=module_name,
-                )
-                - 1
-            )  # Convert from 1-indexed to 0-indexed
-        else:
-            current_stage = 0
-
-        assert (
-            current_stage <= total_stages
-        ), f"current_stage: {current_stage} is greater than total_stages: {total_stages}"
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"current_stage: {current_stage} total_stages: {total_stages} "
-            f"num_warmup_microbatches: {total_stages - current_stage - 1}"
-        )
-        return total_stages - current_stage - 1
-
-    def _build_rank_module_info_map(self):
-        """For each module in the current rank, initialize the P2P communicator
-        and build the bridge communicator info for the module.
-        Each rank may hold multiple modules when colocated.
-        """
-        for module_name, module_grid in self.module_to_grid_map.items():
-            if self.is_current_rank_in_grid(module_grid):
-                # Initialize P2P communicator
-                pp_group = module_grid.get_pg('pp')
-                p2p_comm = P2PCommunicator(pp_group, self.config)
-                pp_size = dist.get_world_size(pp_group)
-                rank_in_pp_group = dist.get_group_rank(pp_group, self.current_rank)
-                pp_rank = rank_in_pp_group % pp_size
-
-                bridge_comms_as_dest_module = []
-                bridge_comms_as_src_module = []
-                # If first stage, check if the module has any incoming modules
-                # If so, initialize bridge communicator
-                if pp_rank == 0:
-                    for bridge_comm in self.bridge_comms:
-                        if (
-                            bridge_comm.is_current_rank_in_grid(bridge_comm.dest_grid)
-                            and bridge_comm.dest_module_name == module_name
-                        ):
-                            bridge_comms_as_dest_module.append(bridge_comm)
-                # If last stage, check if the module has any outgoing modules
-                # If so, initialize bridge communicator
-                if pp_rank == pp_size - 1:
-                    for bridge_comm in self.bridge_comms:
-                        if (
-                            bridge_comm.is_current_rank_in_grid(bridge_comm.src_grid)
-                            and bridge_comm.src_module_name == module_name
-                        ):
-                            bridge_comms_as_src_module.append(bridge_comm)
-                # Build RankModuleInfo for the module
-                rank_module_info = RankModuleInfo(
-                    pp_rank=pp_rank,
-                    pp_size=pp_size,
-                    p2p_communicator=p2p_comm,
-                    bridge_comms_as_dest_module=bridge_comms_as_dest_module,
-                    bridge_comms_as_src_module=bridge_comms_as_src_module,
-                )
-                self.rank_module_map[module_name] = rank_module_info
-
-    def recv_forward(
-        self, tensor_shape: Optional[Shape] = None, is_first_stage: bool = False
-    ) -> Dict[str, torch.Tensor]:
-        """Receive forward activation tensor.
-
-        Args:
-            tensor_shape: Expected activation tensor shape
-
-        Returns:
-            A dictionary mapping module names to tensors.
-        """
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"[receive_forward] tensors_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
-        )
-        input_dict = {}
-        for module_name, rank_module_info in self.rank_module_map.items():
-
-            if rank_module_info.pp_rank == 0:
-                # If first stage, and has incoming modules, receive forward activation
-                # from incoming modules.
-                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
-                    input_dict[bridge_comm.src_module_name] = bridge_comm.recv_forward()
-            else:
-                # If not first stage, receive forward activation tensor from P2P communicator.
-                input_dict[module_name] = rank_module_info.p2p_communicator.recv_forward(
-                    tensor_shapes=tensor_shape, is_first_stage=False
-                )
-        return input_dict
-
-    def send_forward(self, output_dict: Dict[str, torch.Tensor], is_last_stage: bool = False):
-        """Send forward activation tensor.
-
-        Args:
-            output_dict: A dictionary mapping module names to tensors.
-        """
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"[send_forward] output_dict keys: {output_dict.keys()}, is_last_stage: {is_last_stage}"
-        )
-        for module_name, rank_module_info in self.rank_module_map.items():
-            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
-                # If last stage, and has outgoing modules, send forward activation
-                # by using bridge communicator.
-                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
-                    bridge_comm.send_forward(output_dict[module_name])
-            else:
-                # If not last stage, send forward activation by using P2P communicator.
-                rank_module_info.p2p_communicator.send_forward(
-                    output_dict[module_name], is_last_stage=False
-                )
-
-    def send_forward_recv_backward(
-        self,
-        output_dict: Dict[str, torch.Tensor],
-        tensor_shape: Optional[Shape] = None,
-        is_last_stage: bool = False,
-    ) -> Dict[str, torch.Tensor]:
-        """Send forward activation tensor and receive backward activation tensor.
-
-        Args:
-            output_dict: A dictionary mapping module names to tensors.
-            tensor_shape: Expected gradient tensor shape
-
-        Returns:
-            A dictionary mapping module names to tensors.
-        """
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"[send_forward_recv_backward] output_dict keys: {output_dict.keys()}, "
-            f"tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
-        )
-        grad_dict = {}
-        for module_name, rank_module_info in self.rank_module_map.items():
-            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
-                # If last stage, and has outgoing modules, send forward activation and
-                # receive backward gradient by using bridge communicator.
-                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
-                    grad_dict[bridge_comm.src_module_name] = bridge_comm.send_forward_recv_backward(
-                        output_dict[module_name]
-                    )
-            else:
-                # If not last stage, send forward activation and receive backward gradient
-                # by using P2P communicator.
-                grad_dict[module_name] = (
-                    rank_module_info.p2p_communicator.send_forward_recv_backward(
-                        output_dict[module_name], tensor_shapes=tensor_shape, is_last_stage=False
-                    )
-                )
-        return grad_dict
-
-    def send_backward_recv_forward(
-        self,
-        grad_dict: Dict[str, torch.Tensor],
-        tensor_shape: Optional[Shape] = None,
-        is_first_stage: bool = False,
-    ) -> Dict[str, torch.Tensor]:
-        """Send backward activation tensor and receive forward activation tensor.
-
-        Args:
-            grad_dict: A dictionary mapping module names to tensors.
-            tensor_shape: Expected gradient tensor shape
-
-        Returns:
-            A dictionary mapping module names to tensors.
-        """
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"[send_backward_recv_forward] grad_dict keys: {grad_dict.keys()}, "
-            f"tensor_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
-        )
-        input_dict = {}
-        for module_name, rank_module_info in self.rank_module_map.items():
-            if rank_module_info.pp_rank == 0:
-                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
-                    # If first stage, and has incoming modules, send backward gradient and
-                    # receive forward activation by using bridge communicator.
-                    input_dict[bridge_comm.src_module_name] = (
-                        bridge_comm.send_backward_recv_forward(
-                            grad_dict[bridge_comm.src_module_name]
-                        )
-                    )
-            else:
-                # If not first stage, send backward gradient and receive forward activation
-                # by using P2P communicator.
-                input_dict[module_name] = (
-                    rank_module_info.p2p_communicator.send_backward_recv_forward(
-                        grad_dict[module_name], tensor_shapes=tensor_shape, is_first_stage=False
-                    )
-                )
-        return input_dict
-
-    def recv_backward(
-        self, tensor_shape: Optional[Shape] = None, is_last_stage: bool = False
-    ) -> Dict[str, torch.Tensor]:
-        """Receive backward activation tensor.
-
-        Args:
-            tensor_shape: Expected gradient tensor shape
-
-        Returns:
-            A dictionary mapping module names to tensors.
-        """
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"[recv_backward] tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
-        )
-        grad_dict = {}
-        for module_name, rank_module_info in self.rank_module_map.items():
-            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
-                # If last stage, and has incoming modules, receive backward gradient
-                # by using bridge communicator.
-                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
-                    grad_dict[bridge_comm.src_module_name] = bridge_comm.recv_backward()
-            else:
-                # If not last stage, receive backward gradient by using P2P communicator.
-                grad_dict[module_name] = rank_module_info.p2p_communicator.recv_backward(
-                    tensor_shapes=tensor_shape, is_last_stage=False
-                )
-        return grad_dict
-
-    def send_backward(self, grad_dict: Dict[str, torch.Tensor], is_first_stage: bool = False):
-        """Send backward activation tensor.
-
-        Args:
-            grad_dict: A dictionary mapping module names to tensors.
-        """
-        logging.debug(
-            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
-            f"[send_backward] grad_dict keys: {grad_dict.keys()}, is_first_stage: {is_first_stage}"
-        )
-        for module_name, rank_module_info in self.rank_module_map.items():
-            if rank_module_info.pp_rank == 0:
-                # If first stage, and has incoming modules, send backward activation
-                # by using bridge communicator.
-                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
-                    bridge_comm.send_backward(grad_dict[bridge_comm.src_module_name])
-            else:
-                # If not first stage, send backward activation by using P2P communicator.
-                rank_module_info.p2p_communicator.send_backward(
-                    grad_dict[module_name], is_first_stage=False
-                )
-
-    @staticmethod
-    def compute_total_pipeline_stages(
-        topology: Dict[str, List[str]],
-        module_to_grid_map: Dict[str, HyperCommGrid],
-        rank: Optional[int] = None,
-        module_name: Optional[str] = None,
-    ) -> int:
-        """Compute the total number of pipeline stages across a multi-module chain.
-
-        Interprets ``topology`` as a directed acyclic graph (DAG) where nodes are modules
-        and edges indicate forward data flow from source to destination modules. Each node
-        is assigned a weight equal to its pipeline parallel size (number of PP stages).
-
-        The total number of stages is defined as the length of the longest path in this DAG
-        under node weights.
-
-        If ``rank`` is None (default), returns the maximum over all terminal (sink) modules of
-        the sum of PP sizes along a path ending at that terminal. For example, given:
-
-            image_encoder ->\
-                              -> llm -> generator
-            audio_encoder  ->/
-
-        the total is: max(pp(image_encoder), pp(audio_encoder)) + pp(llm) + pp(generator).
-
-        If ``rank`` is provided, the result is the total number of pipeline stages up to (and
-        including) the PP stage that ``rank`` occupies inside its module. In this case, the
-        weight of the target module equals (pp_rank_index(rank) + 1) instead of the module's
-        full PP size; other modules still contribute their full PP sizes. If the rank belongs to
-        multiple modules (colocation), pass ``module_name`` to disambiguate; otherwise the
-        maximum across all candidate modules containing the rank is returned.
-
-        Args:
-            topology: Mapping from a module to its list of outgoing modules.
-            module_to_grid_map: Mapping from module name to its ``HyperCommGrid``.
-
-        Returns:
-            The total number of pipeline stages along the longest path given the constraints.
-
-        Raises:
-            ValueError: If the topology contains cycles; or has no terminal nodes when
-                ``rank`` is None
-        """
-        nodes = set(module_to_grid_map.keys())
-        # Build adjacency and reverse-adjacency (predecessors).
-        adj: Dict[str, List[str]] = {node: list(topology.get(node, [])) for node in nodes}
-        preds: Dict[str, List[str]] = {node: [] for node in nodes}
-        for src, outs in adj.items():
-            for dst in outs:
-                preds[dst].append(src)
-
-        # Identify terminal nodes (no outgoing edges) for the rank=None case.
-        sinks = [node for node, outs in adj.items() if not outs]
-        if rank is None and not sinks:
-            raise ValueError(
-                "Topology must be a DAG with at least one terminal (no outgoing) module."
-            )
-
-        def pp_size(name: str) -> int:
-            grid = module_to_grid_map[name]
-            pp_dim_index = grid.dim_names.index('pp')
-            return grid.shape[pp_dim_index]
-
-        def partial_weight_for_target(target: str) -> Optional[int]:
-            if rank is None:
-                return None
-            grid = module_to_grid_map.get(target)
-            rank_groups = grid._gen_rank_enum(['pp'])
-            stage_index: Optional[int] = None
-            for group in rank_groups:
-                if rank in group:
-                    stage_index = group.index(rank)
-                    break
-            return stage_index + 1
-
-        def longest_path_to(target: str) -> int:
-            visiting = set()
-            partial = partial_weight_for_target(target)
-
-            def weight(name: str) -> int:
-                if partial is not None and name == target:
-                    return partial
-                return pp_size(name)
-
-            def dfs(node: str) -> int:
-                if node in visiting:
-                    raise ValueError("Topology contains cycles; expected a DAG.")
-                visiting.add(node)
-                best = 0
-                for p in preds.get(node, []):
-                    val = dfs(p)
-                    if val > best:
-                        best = val
-                visiting.remove(node)
-                return weight(node) + best
-
-            return dfs(target)
-
-        if rank is None:
-            return max(longest_path_to(sink) for sink in sinks)
-
-        return longest_path_to(module_name)
diff --git a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
deleted file mode 100644
index 22f790cc0a9..00000000000
--- a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
+++ /dev/null
@@ -1,782 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-import logging
-import os
-import sys
-
-import pytest
-import torch
-import torch.distributed as dist
-from packaging import version
-
-from megatron.core import parallel_state
-from megatron.core.hyper_comm_grid import HyperCommGrid
-from megatron.core.model_parallel_config import ModelParallelConfig
-from megatron.core.pipeline_parallel.multimodule_communicator import MultiModulePipelineCommunicator
-from tests.unit_tests.pipeline_parallel.test_bridge_communicator import (
-    _avg_params,
-    _create_transformer_block,
-    _get_pg_collection_from_grid,
-    create_hypercomm_grid,
-    get_transformer_block_and_grid,
-)
-from tests.unit_tests.test_utilities import Utils
-
-
-class TestMultiModulePipelineCommunicator:
-
-    @classmethod
-    def setup_class(cls):
-        """Set up distributed environment for the entire test class."""
-        if not dist.is_initialized():
-            dist.init_process_group(backend="nccl")
-        if torch.cuda.is_available():
-            torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
-
-        world_size = dist.get_world_size()
-        if world_size != 8:
-            pytest.skip(
-                f"These tests require 8 GPUs, but only {world_size} are available.",
-                allow_module_level=True,
-            )
-
-    def teardown_class(cls):
-        Utils.destroy_model_parallel()
-
-    def test_multimodule_communicator_init(self):
-        """Test MultiModulePipelineCommunicator initialization."""
-
-        # Create process group grids for each module
-        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
-        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
-        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
-        generator_grid = create_hypercomm_grid(offset=6, tp=2, cp=1, pp=1, dp=1)
-
-        # Define module-grid mapping
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-            'generator': generator_grid,
-        }
-        # Define module computation topology
-        topology = {
-            'image_encoder': ['llm'],
-            'audio_encoder': ['llm'],
-            'llm': ['generator'],
-            'generator': [],
-        }
-        config = ModelParallelConfig(bf16=True)
-        # Initialize communicator
-        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
-        # Test attributes match expectations
-        assert mllm_comm.module_to_grid_map == module_to_grid_map
-        assert mllm_comm.topology == topology
-        assert mllm_comm.config == config
-        assert mllm_comm.current_rank == dist.get_rank()
-
-    def test_compute_total_pipeline_stages(self):
-        """Test compute_total_pipeline_stages for overall chain and until specific ranks."""
-
-        # Create process group grids for each module
-        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
-        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
-        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
-        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
-
-        # Define module-grid mapping and topology
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-            'generator': generator_grid,
-        }
-        topology = {
-            'image_encoder': ['llm'],
-            'audio_encoder': ['llm'],
-            'llm': ['generator'],
-            'generator': [],
-        }
-
-        # Overall total pipeline stages: max(1,1) + 2 + 1 = 4
-        total = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
-            topology, module_to_grid_map
-        )
-        assert total == 4
-
-        llm_pp_rank = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
-            topology, module_to_grid_map, rank=2, module_name='llm'
-        )
-        assert llm_pp_rank == 2
-
-    def test_send_forward_recv_forward(self):
-        """Test send_forward and recv_forward operations."""
-        if not dist.is_initialized():
-            pytest.skip("Distributed not initialized")
-
-        # Create process group grids for each module
-        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
-        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
-        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
-        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
-
-        # Set up module-grid mapping and topology
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-            'generator': generator_grid,
-        }
-        topology = {
-            'image_encoder': ['llm'],
-            'audio_encoder': ['llm'],
-            'llm': ['generator'],
-            'generator': [],
-        }
-        config = ModelParallelConfig(pipeline_dtype=torch.float)
-        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
-
-        # Simulate forward communication for each module
-        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
-            # Image encoder sends output forward
-            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
-            mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
-            # Audio encoder sends output forward
-            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
-            mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
-            if dist.get_rank() == 2 or dist.get_rank() == 3:
-                # LLM stage receives both image and audio outputs
-                input_dict = mllm_comm.recv_forward()
-                assert input_dict['image_encoder'].shape == (2, 8, 128)
-                assert input_dict['audio_encoder'].shape == (2, 16, 128)
-                mllm_comm.send_forward(output_dict)
-            else:
-                # LLM stage receives concatenated LLM outputs
-                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
-                assert input_dict['llm'].shape == (2, 32, 128)
-                mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(generator_grid):
-            # Generator module receives final LLM output
-            input_dict = mllm_comm.recv_forward()
-            assert input_dict['llm'].shape == (1, 32, 128)
-
-    def test_send_forward_recv_forward_with_different_pp_size(self):
-        """Test for the case when pp(image_encoder) != pp(audio_encoder)."""
-        if not dist.is_initialized():
-            pytest.skip("Distributed not initialized")
-
-        # Create process group grids for each module
-        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=2, dp=1)
-        audio_encoder_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=1, dp=1)
-        llm_grid = create_hypercomm_grid(offset=4, tp=1, cp=1, pp=4, dp=1)
-
-        # Set up module-grid mapping and topology
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-        }
-        topology = {'image_encoder': ['llm'], 'audio_encoder': ['llm'], 'llm': []}
-        config = ModelParallelConfig(pipeline_dtype=torch.float)
-        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
-
-        # Simulate forward communication for each module
-        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
-            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
-            if dist.get_rank() == 0:
-                # Image encoder sends output forward
-                mllm_comm.send_forward(output_dict)
-            else:
-                # Image stage receives image outputs
-                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 8, 128))
-                assert input_dict['image_encoder'].shape == (2, 8, 128)
-                mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
-            # Audio encoder sends output forward
-            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
-            mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
-            if dist.get_rank() == 4:
-                # LLM stage receives both image and audio outputs
-                input_dict = mllm_comm.recv_forward()
-                assert input_dict['image_encoder'].shape == (2, 8, 128)
-                assert input_dict['audio_encoder'].shape == (2, 16, 128)
-                mllm_comm.send_forward(output_dict)
-            elif dist.get_rank() == 5 or dist.get_rank() == 6:
-                # LLM stage receives concatenated LLM outputs
-                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
-                assert input_dict['llm'].shape == (2, 32, 128)
-                mllm_comm.send_forward(output_dict)
-            elif dist.get_rank() == 7:
-                # LLM stage receives concatenated LLM outputs
-                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
-                assert input_dict['llm'].shape == (2, 32, 128)
-
-    def test_send_backward_recv_backward(self):
-        """Test send_backward and recv_backward operations."""
-        if not dist.is_initialized():
-            pytest.skip("Distributed not initialized")
-
-        # Create process group grids for each module
-        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
-        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
-        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
-        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
-
-        # Set up module-grid mapping and topology
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-            'generator': generator_grid,
-        }
-        topology = {
-            'image_encoder': ['llm'],
-            'audio_encoder': ['llm'],
-            'llm': ['generator'],
-            'generator': [],
-        }
-        config = ModelParallelConfig(pipeline_dtype=torch.float)
-        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
-
-        # Simulate backward communication for each module
-        if mllm_comm.is_current_rank_in_grid(generator_grid):
-            # Generator sends gradient backward
-            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
-            mllm_comm.send_backward(grad_dict)
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            if dist.get_rank() == 4 or dist.get_rank() == 5:
-                # LLM receives expanded gradient and sends backward
-                received_grad = mllm_comm.recv_backward()
-                assert received_grad['llm'].shape == (2, 32, 128)
-                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
-                mllm_comm.send_backward(grad_dict)
-            else:
-                # LLM receives gradient and sends backward to both image/audio encoders
-                received_grad = mllm_comm.recv_backward(tensor_shape=(2, 32, 128))
-                assert received_grad['llm'].shape == (2, 32, 128)
-                grad_dict = {
-                    'image_encoder': torch.randn(2, 8, 128).cuda(),
-                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
-                }
-                mllm_comm.send_backward(grad_dict)
-        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
-            # Image encoder receives its gradient
-            received_grad = mllm_comm.recv_backward()
-            assert received_grad['image_encoder'].shape == (2, 8, 128)
-        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
-            # Audio encoder receives its gradient
-            received_grad = mllm_comm.recv_backward()
-            assert received_grad['audio_encoder'].shape == (2, 16, 128)
-
-    @pytest.mark.skipif(
-        version.parse(torch.__version__) < version.parse('2.3.0'),
-        reason="Feature requires PyTorch 2.3 or later",
-    )
-    def test_send_forward_recv_backward_send_backward_recv_forward(self):
-        """Test send_forward_recv_backward and send_backward_recv_forward operations."""
-        if not dist.is_initialized():
-            pytest.skip("Distributed not initialized")
-
-        # Create process group grids for each module
-        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
-        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
-        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
-        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
-
-        # Set up module-grid mapping and topology
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-            'generator': generator_grid,
-        }
-        topology = {
-            'image_encoder': ['llm'],
-            'audio_encoder': ['llm'],
-            'llm': ['generator'],
-            'generator': [],
-        }
-        config = ModelParallelConfig(pipeline_dtype=torch.float)
-        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
-
-        # Simulate bidirectional send/recv for forward and backward in pipeline
-
-        # Encoder stages send forward to the first stage of LLM, and receive backward from the first stage of LLM
-        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
-            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
-            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
-            assert received_grad['image_encoder'].shape == (2, 8, 128)
-        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
-            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
-            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
-            assert received_grad['audio_encoder'].shape == (2, 16, 128)
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            if dist.get_rank() == 2 or dist.get_rank() == 3:
-                grad_dict = {
-                    'image_encoder': torch.randn(2, 8, 128).cuda(),
-                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
-                }
-                input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
-                assert input_dict['image_encoder'].shape == (2, 8, 128)
-                assert input_dict['audio_encoder'].shape == (2, 16, 128)
-
-        # First stage of LLM sends forward to the second stage of LLM, and receive backward from the second stage of LLM
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            if dist.get_rank() == 2 or dist.get_rank() == 3:
-                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
-                received_grad = mllm_comm.send_forward_recv_backward(
-                    output_dict, tensor_shape=(2, 32, 128)
-                )
-                assert received_grad['llm'].shape == (2, 32, 128)
-            if dist.get_rank() == 4 or dist.get_rank() == 5:
-                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
-                input_dict = mllm_comm.send_backward_recv_forward(
-                    grad_dict, tensor_shape=(2, 32, 128)
-                )
-                assert input_dict['llm'].shape == (2, 32, 128)
-
-        # Second stage of LLM sends forward to generator, and receive backward from generator
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            if dist.get_rank() == 4 or dist.get_rank() == 5:
-                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
-                received_grad = mllm_comm.send_forward_recv_backward(output_dict)
-                assert received_grad['llm'].shape == (2, 32, 128)
-        if mllm_comm.is_current_rank_in_grid(generator_grid):
-            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
-            input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
-            assert input_dict['llm'].shape == (1, 32, 128)
-
-    @pytest.mark.skipif(
-        version.parse(torch.__version__) < version.parse('2.3.0'),
-        reason="Feature requires PyTorch 2.3 or later",
-    )
-    def test_send_forward_recv_forward_with_transformer_blocks(self):
-        """Test send_forward and recv_forward operations."""
-
-        # Set model/test dimensions for easier debugging and output comparison
-        hidden_size = 16
-        sequence_length = 2
-        micro_batch_size = 2
-
-        # For reproducibility, set a fixed seed
-        torch.manual_seed(12345)
-        dtype = torch.float32
-
-        # Create random input hidden states tensor
-        hidden_states = torch.randn(
-            (sequence_length, micro_batch_size, hidden_size), device="cuda"
-        ).to(dtype)
-        current_rank = dist.get_rank()
-
-        # ========== Initialize tensor model-parallel environment ==========
-        parallel_state_tp = 2
-        Utils.initialize_model_parallel(tensor_model_parallel_size=2)
-
-        # ========== Build reference 1D grid and transformer block for weight sharing ==========
-        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
-        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
-        ref_block = _create_transformer_block(
-            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
-        )
-        _avg_params(
-            ref_block, ref_grid.get_pg("dp")
-        )  # Ensure parameters are averaged across data parallel (DP)
-
-        # ========== Create different transformer blocks for each model stage ==========
-        # Image encoder
-        image_encoder_block, image_encoder_grid = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=1,
-            cp_size=1,
-            pp_size=1,
-            dp_size=1,
-            grid_offset=0,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        # Audio encoder
-        audio_encoder_block, audio_encoder_grid = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=1,
-            cp_size=1,
-            pp_size=1,
-            dp_size=1,
-            grid_offset=1,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        # LLM (Large Language Model) block with tensor & pipeline parallelism
-        llm_block, llm_grid = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=2,
-            cp_size=1,
-            pp_size=2,
-            dp_size=1,
-            grid_offset=2,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        # Generator block (final stage) with DP=2
-        generator_block, generator_grid = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=1,
-            cp_size=1,
-            pp_size=1,
-            dp_size=2,
-            grid_offset=6,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-
-        # ========== Define module-to-grid correspondence and pipeline topology ==========
-        module_to_grid_map = {
-            'image_encoder': image_encoder_grid,
-            'audio_encoder': audio_encoder_grid,
-            'llm': llm_grid,
-            'generator': generator_grid,
-        }
-        topology = {
-            'image_encoder': ['llm'],  # image_encoder sends output to llm
-            'audio_encoder': ['llm'],  # audio_encoder sends output to llm
-            'llm': ['generator'],  # llm sends output to generator
-            'generator': [],  # generator is the final module
-        }
-        config = ModelParallelConfig(pipeline_dtype=torch.float)
-        # Define dimension mapping for sequence, batch, hidden
-        dim_mapping = {'s': 0, 'h': 2, 'b': 1}
-        seq_dim = dim_mapping['s']
-
-        # Communication handler for multi-module pipeline (send/recv abstraction)
-        mllm_comm = MultiModulePipelineCommunicator(
-            module_to_grid_map, topology, config, dim_mapping=dim_mapping
-        )
-
-        # ========== Run actual distributed pipeline blocks (per process, depending on role) ==========
-        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
-            # Image encoder rank: run forward and send output
-            image_encoder_output = image_encoder_block(
-                hidden_states=hidden_states, attention_mask=None
-            )
-            output_dict = {'image_encoder': image_encoder_output}
-            mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
-            # Audio encoder rank: run forward and send output
-            audio_encoder_output = audio_encoder_block(
-                hidden_states=hidden_states, attention_mask=None
-            )
-            output_dict = {'audio_encoder': audio_encoder_output}
-            mllm_comm.send_forward(output_dict)
-        if mllm_comm.is_current_rank_in_grid(llm_grid):
-            if dist.get_rank() == 2 or dist.get_rank() == 3:
-                # LLM stage 0 (receives both image and audio, concatenates along seq_dim)
-                input_dict = mllm_comm.recv_forward()
-                llm_output = llm_block(
-                    hidden_states=torch.cat(
-                        [input_dict['image_encoder'], input_dict['audio_encoder']], dim=seq_dim
-                    ),
-                    attention_mask=None,
-                )
-                output_dict = {'llm': llm_output}
-                mllm_comm.send_forward(output_dict)
-            else:
-                # LLM stage 1 (receives output of previous LLM stage)
-                input_dict = mllm_comm.recv_forward(
-                    tensor_shape=(sequence_length * 2, micro_batch_size, hidden_size)
-                )
-                llm_output = llm_block(hidden_states=input_dict['llm'], attention_mask=None)
-                output_dict = {'llm': llm_output}
-                mllm_comm.send_forward(output_dict)
-
-        if mllm_comm.is_current_rank_in_grid(generator_grid):
-            # Generator block: only receives from llm and runs forward
-            input_dict = mllm_comm.recv_forward()
-            generator_output = generator_block(hidden_states=input_dict['llm'], attention_mask=None)
-
-        # ========== Build a reference (serial/global) pipeline for correctness checking ==========
-        global_image_encoder_block, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        global_audio_encoder_block, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        global_llm_block_pp_rank_0, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        global_llm_block_pp_rank_1, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        global_generator_block, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-
-        # Run each stage sequentially as a global pipeline (for truth)
-        global_image_encoder_output = global_image_encoder_block(
-            hidden_states=hidden_states, attention_mask=None
-        )
-        global_audio_encoder_output = global_audio_encoder_block(
-            hidden_states=hidden_states, attention_mask=None
-        )
-        # Compare output between global and distributed blocks for image/audio stage
-        if current_rank == 0:
-            torch.testing.assert_close(
-                global_image_encoder_output, image_encoder_output, rtol=1e-3, atol=1e-3
-            )
-        if current_rank == 1:
-            torch.testing.assert_close(
-                global_audio_encoder_output, audio_encoder_output, rtol=1e-3, atol=1e-3
-            )
-
-        # Feed outputs to LLM stages (emulate pipeline cut with concatenation)
-        global_llm_input = torch.cat(
-            [global_image_encoder_output, global_audio_encoder_output], dim=seq_dim
-        )
-        global_llm_pp_rank_0_output = global_llm_block_pp_rank_0(
-            hidden_states=global_llm_input, attention_mask=None
-        )
-        if current_rank == 2 or current_rank == 3:
-            torch.testing.assert_close(
-                global_llm_pp_rank_0_output, llm_output, rtol=1e-3, atol=1e-3
-            )
-        global_llm_pp_rank_1_output = global_llm_block_pp_rank_1(
-            hidden_states=global_llm_pp_rank_0_output, attention_mask=None
-        )
-        if current_rank == 4 or current_rank == 5:
-            torch.testing.assert_close(
-                global_llm_pp_rank_1_output, llm_output, rtol=1e-3, atol=1e-3
-            )
-
-        # Generator output and comparison to distributed output (for each DP chunk)
-        global_generator_block_output = global_generator_block(
-            hidden_states=global_llm_pp_rank_1_output, attention_mask=None
-        )
-        global_generator_block_chunks = torch.split(
-            global_generator_block_output, global_generator_block_output.shape[1] // 2, dim=1
-        )
-        if current_rank == 6:
-            torch.testing.assert_close(
-                global_generator_block_chunks[0], generator_output, rtol=1e-3, atol=1e-3
-            )
-        if current_rank == 7:
-            torch.testing.assert_close(
-                global_generator_block_chunks[1], generator_output, rtol=1e-3, atol=1e-3
-            )
-
-    @pytest.mark.skipif(
-        version.parse(torch.__version__) < version.parse('2.3.0'),
-        reason="Feature requires PyTorch 2.3 or later",
-    )
-    @pytest.mark.parametrize(
-        "grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp",
-        [
-            (2, 1, 1, 2, 1, 1, 2),  # TP2PP1DP1 to TP2PP1DP1
-            (2, 1, 1, 2, 2, 1, 2),  # TP2PP1DP1 to TP2PP2DP1
-            (2, 2, 1, 2, 2, 1, 2),  # TP2PP2DP1 to TP2PP2DP1
-            (4, 1, 1, 4, 1, 1, 4),  # TP4DP1 to TP4DP1
-            (2, 1, 2, 4, 1, 1, 2),  # TP2DP2 to TP4DP1
-            (4, 1, 1, 2, 1, 2, 2),  # TP4DP1 to TP2DP2
-            (2, 1, 2, 1, 1, 4, 2),  # TP2DP2 to TP1DP4
-        ],
-    )
-    def test_send_forward_recv_forward_with_transformer_blocks_and_different_parallelisms(
-        self, grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp
-    ):
-        """Test bridge communicator with two transformer blocks having different process group configurations."""
-        # Model and input configuration
-        hidden_size = 16
-        sequence_length = 2
-        micro_batch_size = 8
-        torch.manual_seed(12345)
-        dtype = torch.float32
-
-        # Create random input tensor on CUDA
-        hidden_states = torch.randn(
-            (sequence_length, micro_batch_size, hidden_size), device="cuda"
-        ).to(dtype)
-        hidden_states_ref = hidden_states.clone()
-        current_rank = dist.get_rank()
-
-        # Initialize model parallel with desired TP
-        Utils.initialize_model_parallel(tensor_model_parallel_size=parallel_state_tp)
-
-        # Build a reference grid and block for parameter sharing & DP averaging
-        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
-        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
-        ref_block = _create_transformer_block(
-            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
-        )
-        _avg_params(
-            ref_block, ref_grid.get_pg("dp")
-        )  # Synchronize parameters across DP for reproducibility
-
-        # ====== Create two transformer block+grid pairs with different TP/DP settings ======
-        block_grid_1, grid_1 = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=grid1_tp,
-            pp_size=grid1_pp,
-            dp_size=grid1_dp,
-            grid_offset=0,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-
-        block_grid_2, grid_2 = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=grid2_tp,
-            pp_size=grid2_pp,
-            dp_size=grid2_dp,
-            grid_offset=grid_1.size,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-
-        dist.barrier()  # Synchronize ranks before communication
-
-        # Module-grid map and pipeline communication topology
-        module_to_grid_map = {'image_encoder': grid_1, 'llm': grid_2}
-        topology = {
-            'image_encoder': ['llm'],  # image_encoder sends forward results to llm
-            'llm': [],  # llm is the last stage here
-        }
-        config = ModelParallelConfig(pipeline_dtype=torch.float)
-        mllm_comm = MultiModulePipelineCommunicator(
-            module_to_grid_map, topology, config, dim_mapping={'s': 0, 'h': 2, 'b': 1}
-        )
-
-        output_grid_2 = None
-        # If current rank is in the first grid, run first block and send output
-        if grid_1 is not None and mllm_comm.is_current_rank_in_grid(grid_1):
-            rank_module_info = mllm_comm.rank_module_map['image_encoder']
-            if rank_module_info.pp_rank == 0:
-                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
-                mllm_comm.send_forward({'image_encoder': hidden_states})
-            else:
-                input_dict = mllm_comm.recv_forward(
-                    tensor_shape=(sequence_length, micro_batch_size, hidden_size)
-                )
-                hidden_states = input_dict['image_encoder']
-                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
-                mllm_comm.send_forward({'image_encoder': hidden_states})
-
-        # If current rank is in second grid, receive and run the second block
-        if grid_2 is not None and mllm_comm.is_current_rank_in_grid(grid_2):
-            rank_module_info = mllm_comm.rank_module_map['llm']
-            if rank_module_info.pp_rank == 0:
-                input_dict = mllm_comm.recv_forward()
-                hidden_states = input_dict['image_encoder']
-                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
-                if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
-                    output_grid_2 = hidden_states
-                else:
-                    mllm_comm.send_forward({'llm': hidden_states})
-            elif rank_module_info.pp_rank < rank_module_info.pp_size - 1:
-                input_dict = mllm_comm.recv_forward(
-                    tensor_shape=(
-                        sequence_length,
-                        (grid1_dp * micro_batch_size) // grid2_dp,
-                        hidden_size,
-                    )
-                )
-                hidden_states = input_dict['llm']
-                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
-                mllm_comm.send_forward({'llm': hidden_states})
-            else:
-                input_dict = mllm_comm.recv_forward(
-                    tensor_shape=(
-                        sequence_length,
-                        (grid1_dp * micro_batch_size) // grid2_dp,
-                        hidden_size,
-                    )
-                )
-                hidden_states = input_dict['llm']
-                output_grid_2 = block_grid_2(hidden_states=hidden_states, attention_mask=None)
-
-                # Compute expected output shape based on change in DP size (chunk/expand batch dimension appropriately)
-                factor = max(grid1_dp, grid2_dp) // min(grid1_dp, grid2_dp)
-                expected_output_shape = (
-                    sequence_length,
-                    (
-                        micro_batch_size * factor
-                        if grid1_dp > grid2_dp
-                        else micro_batch_size // factor
-                    ),
-                    hidden_size,
-                )
-                assert (
-                    output_grid_2.shape == expected_output_shape
-                ), f"Output2 shape mismatch: {output_grid_2.shape}"
-
-        # ====== Reference: global (replicated) pipeline forward for correctness checking ======
-        global_block_1, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-        global_block_2, _ = get_transformer_block_and_grid(
-            ref_block,
-            tp_size=parallel_state_tp,
-            use_global_parallel_state=True,
-            hidden_size=hidden_size,
-            dtype=dtype,
-        )
-
-        for i in range(grid1_pp):
-            hidden_states_ref = global_block_1(hidden_states=hidden_states_ref, attention_mask=None)
-
-        for i in range(grid2_pp):
-            hidden_states_ref = global_block_2(hidden_states=hidden_states_ref, attention_mask=None)
-
-        # Output comparison under different DP compositions between grids
-        if (
-            grid_2 is not None
-            and mllm_comm.is_current_rank_in_grid(grid_2)
-            and rank_module_info.pp_rank == rank_module_info.pp_size - 1
-        ):
-            if grid1_dp == grid2_dp:
-                # DP size matches: all outputs directly compared
-                torch.testing.assert_close(hidden_states_ref, output_grid_2, rtol=1e-3, atol=1e-3)
-            elif grid1_dp < grid2_dp:
-                # If grid2 expands DP: each output_grid_2 chunk corresponds to a split of the reference output
-                grid2_dp_ranks = grid_2._gen_rank_enum([x for x in grid_2.dim_names if x != "dp"])
-                global_block_2_chunks = torch.split(
-                    hidden_states_ref, hidden_states_ref.shape[1] // (grid2_dp // grid1_dp), dim=1
-                )
-                relevant_chunk = None
-                for i, dp_ranks in enumerate(grid2_dp_ranks):
-                    if current_rank in dp_ranks:
-                        relevant_chunk = global_block_2_chunks[i % len(global_block_2_chunks)]
-                torch.testing.assert_close(relevant_chunk, output_grid_2, rtol=1e-3, atol=1e-3)
-            else:
-                # If DP shrinks (grid1_dp > grid2_dp): just compare the relevant first chunk
-                output_grid_2_first_chunk = torch.chunk(output_grid_2, grid1_dp // grid2_dp, dim=1)[
-                    0
-                ]
-                torch.testing.assert_close(
-                    hidden_states_ref, output_grid_2_first_chunk, rtol=1e-3, atol=1e-3
-                )

From 4cfaa7d598b243f21ab0c18a3d3aef543682969f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Sat, 24 Jan 2026 17:45:37 +0100
Subject: [PATCH 11/79] Revert "Remove calculation of padding token in moe
 routing loss (#2142)" (#3069)

---
 .../core/extensions/transformer_engine.py     |   2 +-
 .../common/model_chunk_schedule_plan.py       |   2 -
 .../core/models/gpt/fine_grained_callables.py |  21 +-
 megatron/core/models/gpt/gpt_model.py         |  37 +---
 megatron/core/models/mamba/mamba_model.py     |   2 -
 megatron/core/ssm/mamba_block.py              |   2 -
 megatron/core/transformer/mlp.py              |   2 +-
 megatron/core/transformer/moe/moe_layer.py    |  21 +-
 megatron/core/transformer/moe/moe_utils.py    |  91 ++-------
 megatron/core/transformer/moe/router.py       | 160 +++++----------
 .../core/transformer/transformer_block.py     |  14 +-
 .../core/transformer/transformer_layer.py     |  24 +--
 .../a2a_overlap/test_schedule_chunk_1f1b.py   | 116 +----------
 .../a2a_overlap/test_schedule_layer_1f1b.py   |   4 +-
 .../transformer/moe/test_aux_loss.py          | 182 ------------------
 .../transformer/moe/test_routers.py           |  47 -----
 16 files changed, 88 insertions(+), 639 deletions(-)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index ef8527e9e5e..63694fc1172 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -2161,7 +2161,7 @@ def forward_post_hook(module, *_) -> None:
                     "TEFusedMLP module does not support submodules with post-backward hooks"
                 )
 
-        def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Optional[Tensor]]:
+        def forward(self, hidden_states: torch.Tensor) -> Tuple[Tensor, Optional[Tensor]]:
             """Forward."""
 
             # Construct fused impl if needed
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 033e8e808f9..71aa1ab97f0 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -281,7 +281,6 @@ def __init__(
         extra_block_kwargs=None,
         runtime_gather_output: Optional[bool] = None,
         loss_mask: Optional[Tensor] = None,
-        padding_mask=None,
     ):
         """Initialize the schedule plan of all Transformer layers' sub-modules.
 
@@ -324,7 +323,6 @@ def __init__(
         self._model_chunk_state.mtp_hidden_states = None
         self._model_chunk_state.loss_mask = loss_mask
         self._model_chunk_state.packed_seq_params = packed_seq_params
-        self._model_chunk_state.padding_mask = padding_mask
         self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
         self._model_chunk_state.runtime_gather_output = runtime_gather_output
         self._model_chunk_state.model = model
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 6f2f6b1cb80..9234e142c6c 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -131,19 +131,13 @@ def forward_impl(self):
         if not self.gpt_model.pre_process:
             self.chunk_state.decoder_input = self.gpt_model.decoder.input_tensor
         # Run GPTModel._preprocess
-        (
-            decoder_input,
-            rotary_pos_emb,
-            rotary_pos_cos,
-            rotary_pos_sin,
-            sequence_len_offset,
-            padding_mask,
-        ) = self.gpt_model._preprocess(
-            input_ids=self.chunk_state.input_ids,
-            position_ids=self.chunk_state.position_ids,
-            decoder_input=self.chunk_state.decoder_input,
-            packed_seq_params=self.chunk_state.packed_seq_params,
-            padding_mask=self.chunk_state.padding_mask,
+        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
+            self.gpt_model._preprocess(
+                input_ids=self.chunk_state.input_ids,
+                position_ids=self.chunk_state.position_ids,
+                decoder_input=self.chunk_state.decoder_input,
+                packed_seq_params=self.chunk_state.packed_seq_params,
+            )
         )
 
         # Saved for later use
@@ -152,7 +146,6 @@ def forward_impl(self):
         self.chunk_state.rotary_pos_cos = rotary_pos_cos
         self.chunk_state.rotary_pos_sin = rotary_pos_sin
         self.chunk_state.sequence_len_offset = sequence_len_offset
-        self.chunk_state.padding_mask = padding_mask
         return decoder_input
 
 
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index e287344c13d..e70221d2cfa 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -288,7 +288,6 @@ def _preprocess(
         decoder_input: Tensor = None,
         inference_context: BaseInferenceContext = None,
         packed_seq_params: PackedSeqParams = None,
-        padding_mask: Optional[Tensor] = None,
     ):
         """Preprocesses inputs for the transformer decoder.
 
@@ -305,20 +304,7 @@ def _preprocess(
         if decoder_input is not None:
             pass
         elif self.pre_process:
-            if padding_mask is not None:
-                assert padding_mask.shape == input_ids.shape, (
-                    f"padding_mask shape {padding_mask.shape} does not match "
-                    f"input_ids shape {input_ids.shape}"
-                )
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
-            if padding_mask is not None and self.config.sequence_parallel:
-                padding_mask = (
-                    tensor_parallel.scatter_to_sequence_parallel_region(
-                        padding_mask.transpose(0, 1).contiguous()
-                    )
-                    .transpose(0, 1)
-                    .contiguous()
-                )
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
@@ -437,7 +423,6 @@ def _preprocess(
             rotary_pos_cos,
             rotary_pos_sin,
             sequence_len_offset,
-            padding_mask,
         )
         if rotary_pos_cos_sin is not None:
             # only in the case of flashinfer fused rope will we
@@ -481,7 +466,6 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
-        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -492,9 +476,6 @@ def forward(
         Args:
             runtime_gather_output (bool): Gather output at runtime. Default None means
                 `parallel_output` arg in the constructor will be used.
-            padding_mask (Tensor, optional): Padding mask for MoE routing.
-                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
-                Only used for MoE layers to exclude padding tokens from routing computations.
         """
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()
@@ -507,19 +488,13 @@ def forward(
             decoder_input=decoder_input,
             inference_context=inference_context,
             packed_seq_params=packed_seq_params,
-            padding_mask=padding_mask,
         )
 
-        (
-            decoder_input,
-            rotary_pos_emb,
-            rotary_pos_cos,
-            rotary_pos_sin,
-            sequence_len_offset,
-            padding_mask,
-        ) = preproc_output[:6]
+        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset) = (
+            preproc_output[:5]
+        )
 
-        rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
+        rotary_pos_cos_sin = preproc_output[5] if len(preproc_output) == 6 else None
 
         # Run decoder.
         hidden_states = self.decoder(
@@ -532,7 +507,6 @@ def forward(
             rotary_pos_cos_sin=rotary_pos_cos_sin,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
-            padding_mask=padding_mask,
             **(extra_block_kwargs or {}),
         )
 
@@ -749,7 +723,6 @@ def build_schedule_plan(
         runtime_gather_output: Optional[bool] = None,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
-        padding_mask: Optional[Tensor] = None,
     ):
         """Builds a computation schedule plan for the model.
 
@@ -775,7 +748,6 @@ def build_schedule_plan(
             inference_params (InferenceParams, optional):
                 Parameters for inference. Defaults to None.
             loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
-            padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
 
         Returns:
             TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -797,7 +769,6 @@ def build_schedule_plan(
             extra_block_kwargs,
             runtime_gather_output,
             loss_mask,
-            padding_mask,
         )
 
     def sharded_state_dict(
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 8d45e1d0147..0d71ead4b0f 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -185,7 +185,6 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
-        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the Mamba model. This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -255,7 +254,6 @@ def forward(
             inference_context=inference_context,
             rotary_pos_emb=rotary_pos_emb,
             packed_seq_params=packed_seq_params,
-            padding_mask=padding_mask,
         )
 
         if not self.post_process:
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index ef41faae143..9e41aca8253 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -211,7 +211,6 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
-        padding_mask=None,
     ):
         """
         Forward function of the MambaStack class.
@@ -294,7 +293,6 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                             packed_seq_params=packed_seq_params,
-                            padding_mask=padding_mask,
                         )
                     else:  # MambaLayer
                         hidden_states = layer(
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index 2bc3949a421..2eae0178eea 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -148,7 +148,7 @@ def __init__(
             tp_group=tp_group,
         )
 
-    def forward(self, hidden_states, per_token_scale=None, **kwargs):
+    def forward(self, hidden_states, per_token_scale=None):
         """Perform the forward pass through the MLP block."""
         # [s, b, 4 * h/p]
         nvtx_range_push(suffix="linear_fc1")
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 08c68ae3aed..2fad0f8e5b7 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -239,13 +239,13 @@ def __init__(
         self.cudagraph_tensor_store = MoECudaGraphTensorStore()
 
     @maybe_skip_or_early_return_by_cudagraph("route")
-    def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+    def route(self, hidden_states: torch.Tensor):
         """Compute token routing for preprocessing.
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
         """
-        probs, routing_map = apply_module(self.router)(hidden_states, padding_mask)
+        probs, routing_map = apply_module(self.router)(hidden_states)
         return probs, routing_map
 
     @maybe_skip_or_early_return_by_cudagraph("preprocess")
@@ -346,7 +346,7 @@ def router_and_preprocess(self, hidden_states: torch.Tensor):
         hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map)
         return hidden_states, probs, residual
 
-    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+    def forward(self, hidden_states: torch.Tensor):
         """Forward pass for the MoE layer.
 
         The forward pass comprises four main steps:
@@ -356,10 +356,8 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
         4. Combine: The outputs from the experts are combined and returned.
 
         Args:
-            hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size].
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            hidden_states (torch.Tensor): The input tensor to the MoE layer.
+
         Returns:
             A tuple containing the output tensor and the MLP bias, if any.
         """
@@ -368,15 +366,12 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
                 "During training, performance may degrade if MoE and tensor parallelism"
                 "are enabled without also enabling sequence parallelism."
             )
-        # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states
-        if padding_mask is not None:
-            padding_mask = padding_mask.transpose(0, 1).bool()
 
         # MoE forward: route -> dispatch -> compute -> combine
         def custom_forward(hidden_states):
             try:
                 shared_expert_output = self.shared_experts_compute(hidden_states)
-                probs, routing_map = self.route(hidden_states, padding_mask)
+                probs, routing_map = self.route(hidden_states)
                 hidden_states, probs = self.preprocess(hidden_states, probs, routing_map)
             except MoECudaGraphPartialCaptureSignal as e:
                 # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator.
@@ -403,9 +398,7 @@ def custom_forward(hidden_states):
                     hidden_states,
                 )
             else:
-                outputs = tensor_parallel.checkpoint(
-                    custom_forward, False, hidden_states, padding_mask
-                )
+                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
         else:
             outputs = custom_forward(hidden_states)
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 28c486545a2..5fdeda23dea 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -12,7 +12,6 @@
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
-from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -50,7 +49,6 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
-    padding_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
@@ -102,19 +100,10 @@ def switch_load_balancing_loss_func(
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
         fused (bool): Whether to use the fused version of the auxiliary loss.
-        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                               Shape in [num_tokens]. True for valid tokens,
-                                               False for padding tokens. Defaults to None.
 
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
-    # Apply padding mask to probs if provided
-    if padding_mask is not None:
-        # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
-        mask_expanded = padding_mask.unsqueeze(-1)
-        probs = probs * mask_expanded
-
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -134,35 +123,19 @@ def switch_load_balancing_loss_func(
     return aux_loss
 
 
-def z_loss_func(
-    logits: torch.Tensor, z_loss_coeff: float, padding_mask: Optional[torch.Tensor] = None
-) -> torch.Tensor:
+def z_loss_func(logits: torch.Tensor, z_loss_coeff: float) -> torch.Tensor:
     """Encourages the router's logits to remain small to enhance stability.
     Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
     Args:
         logits (torch.Tensor): The logits of the router.
         z_loss_coeff (float): The coefficient for the z-loss.
-        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                                               Shape [num_tokens]. True = padding (exclude),
-                                               False = valid (include). Defaults to None.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
     """
-    logsum = torch.logsumexp(logits, dim=-1)
-    z_loss_values = torch.square(logsum)
-
-    if padding_mask is not None:
-        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
-        valid_mask = ~padding_mask
-        # Only compute z_loss for valid (non-padding) tokens
-        z_loss_values = z_loss_values * valid_mask
-        # Compute mean over valid tokens only
-        num_valid_tokens = valid_mask.sum()
-        z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff
-    else:
-        z_loss = torch.mean(z_loss_values) * z_loss_coeff
+
+    z_loss = torch.mean(torch.square(torch.logsumexp(logits, dim=-1))) * z_loss_coeff
     return z_loss
 
 
@@ -212,28 +185,6 @@ def get_capacity(
     return capacity
 
 
-def get_tokens_per_expert_and_token_count(
-    routing_map: torch.Tensor,
-    reduce_group: torch.distributed.ProcessGroup,
-    topk: int = None,
-    with_padding_mask: bool = False,
-) -> torch.Tensor:
-    """
-    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with padding mask.
-    """
-    local_tokens_per_expert = routing_map.sum(dim=0)
-    global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
-        local_tokens_per_expert, reduce_group
-    )
-    if with_padding_mask:
-        local_num_tokens = local_tokens_per_expert.sum() / topk
-        total_num_tokens = global_tokens_per_expert.sum() / topk
-    else:
-        local_num_tokens = routing_map.shape[0]
-        total_num_tokens = local_num_tokens * reduce_group.size()
-    return global_tokens_per_expert, local_num_tokens, total_num_tokens
-
-
 class MoEAuxLossAutoScaler(torch.autograd.Function):
     """An AutoScaler that triggers the backward pass and scales the grad for auxiliary loss."""
 
@@ -733,11 +684,7 @@ def compute_topk(
 
 
 def compute_routing_scores_for_aux_loss(
-    logits: torch.Tensor,
-    topk: int,
-    score_function: str,
-    fused: bool = False,
-    padding_mask: Optional[torch.Tensor] = None,
+    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute routing scores based on the score function.
 
@@ -746,9 +693,6 @@ def compute_routing_scores_for_aux_loss(
         topk (int): The number of top-k indices to compute.
         score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
         fused (bool, optional): Whether to use the fused version. Defaults to False.
-        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                               Shape in [num_tokens]. True for valid tokens,
-                                               False for padding tokens. Defaults to None.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: The routing map and the normalized routing scores.
@@ -758,27 +702,20 @@ def compute_routing_scores_for_aux_loss(
             raise ValueError(
                 "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
             )
-        routing_map, scores = fused_compute_score_for_moe_aux_loss(
+        return fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
-    else:
-        if score_function == "softmax":
-            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
-        elif score_function == "sigmoid":
-            scores = torch.sigmoid(logits)
-            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
-        else:
-            raise ValueError(f"Invalid score_function: {score_function}")
 
-        _, top_indices = torch.topk(scores, k=topk, dim=1)
-        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    if score_function == "softmax":
+        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+    elif score_function == "sigmoid":
+        scores = torch.sigmoid(logits)
+        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+    else:
+        raise ValueError(f"Invalid score_function: {score_function}")
 
-    # Apply padding mask to scores if provided
-    if padding_mask is not None:
-        # Invert padding_mask and make True indicates valid tokens
-        valid_mask = (~padding_mask).unsqueeze(-1)
-        routing_map = routing_map * valid_mask
-        scores = scores * valid_mask
+    _, top_indices = torch.topk(scores, k=topk, dim=1)
+    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 8f94a7312ce..c22ca4e8446 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -1,11 +1,12 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from abc import ABC, abstractmethod
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 
 from megatron.core.jit import jit_fuser
+from megatron.core.tensor_parallel import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,
@@ -13,7 +14,6 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
-    get_tokens_per_expert_and_token_count,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -268,29 +268,22 @@ def is_aux_loss_enabled(self) -> bool:
         return False
 
     def _apply_aux_loss(
-        self,
-        probs: torch.Tensor,
-        scores_for_aux_loss: torch.Tensor,
-        routing_map: torch.Tensor,
-        with_padding_mask: bool = False,
+        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
-
-        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
-            get_tokens_per_expert_and_token_count(
-                routing_map=routing_map,
-                reduce_group=self.tp_cp_group,
-                topk=self.topk,
-                with_padding_mask=with_padding_mask,
-            )
+        tokens_per_expert = routing_map.sum(dim=0)
+        tokens_per_expert = reduce_from_tensor_model_parallel_region(
+            tokens_per_expert, self.tp_cp_group
         )
+        num_tokens = routing_map.shape[0]
+        total_num_tokens = num_tokens * self.tp_cp_group.size()
 
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
-            tokens_per_expert=global_tokens_per_expert,
+            tokens_per_expert=tokens_per_expert,
             total_num_tokens=total_num_tokens,
             topk=self.topk,
             num_experts=self.config.num_moe_experts,
@@ -298,12 +291,7 @@ def _apply_aux_loss(
             fused=self.config.moe_router_fusion,
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs,
-            aux_loss_coeff,
-            aux_loss,
-            "load_balancing_loss",
-            self.tp_cp_group,
-            valid_token_count=local_num_tokens,
+            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group
         )
         return probs
 
@@ -314,7 +302,6 @@ def _apply_seq_aux_loss(
         routing_map: torch.Tensor,
         seq_length: int,
         bsz: int,
-        with_padding_mask: bool = False,
     ):
         """Apply the sequence-level auxiliary loss for the given scores and routing map.
 
@@ -328,21 +315,17 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
-        routing_map = routing_map.reshape(seq_length, -1)
-
-        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
-            get_tokens_per_expert_and_token_count(
-                routing_map=routing_map,
-                reduce_group=self.tp_cp_group,
-                with_padding_mask=with_padding_mask,
-                topk=self.topk * bsz,
-            )
+        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
+        tokens_per_expert = reduce_from_tensor_model_parallel_region(
+            tokens_per_expert, self.tp_cp_group
         )
 
+        total_num_tokens = seq_length * self.tp_cp_group.size()
+
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
-                tokens_per_expert=global_tokens_per_expert,
+                tokens_per_expert=tokens_per_expert,
                 total_num_tokens=total_num_tokens,
                 topk=self.topk,
                 num_experts=self.config.num_moe_experts,
@@ -351,43 +334,31 @@ def _apply_seq_aux_loss(
             )
             / bsz
         )
-
         probs = self.attach_and_log_load_balancing_loss(
-            probs,
-            seq_aux_loss_coeff,
-            aux_loss,
-            "seq_load_balancing_loss",
-            self.tp_cp_group,
-            valid_token_count=local_num_tokens,
+            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group
         )
         return probs
 
     def _apply_global_aux_loss(
-        self,
-        probs: torch.Tensor,
-        scores_for_aux_loss: torch.Tensor,
-        routing_map: torch.Tensor,
-        with_padding_mask: bool = False,
+        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
         if global_aux_loss_coeff == 0:
             return probs
 
-        # Use unified function to compute tokens_per_expert and num_tokens
-        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
-            get_tokens_per_expert_and_token_count(
-                routing_map=routing_map,
-                reduce_group=self.tp_dp_cp_group,
-                with_padding_mask=with_padding_mask,
-                topk=self.topk,
-            )
+        tokens_per_expert = routing_map.sum(dim=0)
+        tokens_per_expert = reduce_from_tensor_model_parallel_region(
+            tokens_per_expert, self.tp_dp_cp_group
         )
 
-        self.global_tokens_per_expert += global_tokens_per_expert
+        self.global_tokens_per_expert += tokens_per_expert
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
+        num_tokens = scores_for_aux_loss.shape[0]
+        total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
+
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
             tokens_per_expert=averated_tokens_per_expert,
@@ -404,7 +375,6 @@ def _apply_global_aux_loss(
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
             reduce_group_has_dp=True,
-            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -416,22 +386,18 @@ def attach_and_log_load_balancing_loss(
         aux_loss_name: str,
         reduce_group: torch.distributed.ProcessGroup,
         reduce_group_has_dp: bool = False,
-        valid_token_count: Optional[Union[int, torch.Tensor]] = None,
     ):
         """Attach aux loss function to activation and add to logging.
 
         Args:
-            activation (torch.Tensor): Activation tensor to attach the aux loss to.
-            aux_loss_coeff (float): Coefficient for the aux loss.
-            aux_loss (torch.Tensor): Computed aux loss.
-            aux_loss_name (str): Name of the aux loss for logging.
-            reduce_group (torch.distributed.ProcessGroup): Process group for reduction.
+            activation (torch.Tensor): The activation tensor to attach the loss to.
+            aux_loss_coeff (float): The coefficient for the auxiliary loss.
+            aux_loss (torch.Tensor): The auxiliary loss tensor.
+            aux_loss_name (str): The name of the auxiliary loss for logging.
+            reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss.
             reduce_group_has_dp (bool): Whether the reduce group has data parallel ranks.
                 Set this to True if the reduce group has data parallel ranks. This flag is used to
                 ensure the correct reduction in aux loss tracking.
-            valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding
-                padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor).
-                If None, uses activation.shape[0]. Defaults to None.
         """
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
         # add the aux loss logging value to other layer's since it is difficult to get the
@@ -456,22 +422,17 @@ def attach_and_log_load_balancing_loss(
             # which scales both the main_loss gradient and aux_loss gradient by
             # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads function.
             # To correct this scaling, we need to scale the aux_loss by num_local_tokens here.
-            # Use valid_token_count (excluding padding) if provided, otherwise use total tokens.
-            num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0]
-            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens)
+            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0])
         else:
             activation = MoEAuxLossAutoScaler.apply(activation, aux_loss)
         return activation
 
-    def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
+    def apply_z_loss(self, logits):
         """Encourages the router's logits to remain small to enhance stability.
         Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
         Args:
             logits (torch.Tensor): The logits of the router.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape in [num_tokens]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
 
         Returns:
             torch.Tensor: The logits after applying the z-loss.
@@ -479,7 +440,7 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
         if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled():
             # Skip Z loss calculations when using torch.no_grad() or checkpointing.
             moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()
-            z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask)
+            z_loss = z_loss_func(logits, moe_z_loss_coeff)
             if self.calculate_per_token_loss:
                 # The expected final scaling for z_loss gradients is
                 # 1/(num_micro_batches * dp_size).
@@ -488,9 +449,7 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
                 # which scales both the main_loss gradient and z_loss gradient by
                 # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads().
                 # To correct this scaling, we need to scale the z_loss by num_local_tokens here.
-                # Count valid tokens: sum of inverted mask (False -> True = valid)
-                num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0]
-                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens)
+                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * logits.shape[0])
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
@@ -524,27 +483,20 @@ def apply_input_jitter(self, input: torch.Tensor):
             return input
 
     @jit_fuser
-    def _apply_expert_bias(
-        self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
-    ):
+    def _apply_expert_bias(self, routing_map: torch.Tensor):
         """
         Update expert bias and tokens_per_expert
         Prevent extra local tokens accumulation on evaluation or activation recomputation
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
-                if padding_mask is not None:
-                    routing_map = routing_map & (~padding_mask)
                 self.local_tokens_per_expert += routing_map.sum(dim=0)
 
-    def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+    def routing(self, logits: torch.Tensor):
         """Top-k routing function
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
@@ -554,12 +506,8 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
 
-        # Flatten padding_mask to [num_tokens] if provided
-        if padding_mask is not None:
-            padding_mask = padding_mask.reshape(-1)
-
         # Apply Z-Loss
-        logits = self.apply_z_loss(logits, padding_mask=padding_mask)
+        logits = self.apply_z_loss(logits)
 
         # Calculate probs and routing_map for token dispatching
         if self.routing_type == "sinkhorn":
@@ -592,35 +540,18 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
-                logits,
-                self.topk,
-                self.score_function,
-                fused=self.config.moe_router_fusion,
-                padding_mask=padding_mask,
-            )
-            probs = self._apply_aux_loss(
-                probs,
-                scores_for_aux_loss,
-                routing_map_for_aux_loss,
-                with_padding_mask=padding_mask is not None,
+                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion
             )
+            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)
             probs = self._apply_seq_aux_loss(
-                probs,
-                scores_for_aux_loss,
-                routing_map_for_aux_loss,
-                seq_length,
-                bsz,
-                with_padding_mask=padding_mask is not None,
+                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz
             )
             probs = self._apply_global_aux_loss(
-                probs,
-                scores_for_aux_loss,
-                routing_map_for_aux_loss,
-                with_padding_mask=padding_mask is not None,
+                probs, scores_for_aux_loss, routing_map_for_aux_loss
             )
 
         # Optionally apply expert bias
-        self._apply_expert_bias(routing_map, padding_mask=padding_mask)
+        self._apply_expert_bias(routing_map)
 
         return probs, routing_map
 
@@ -630,15 +561,12 @@ def reset_global_aux_loss_tracker(self):
             self.global_tokens_per_expert.zero_()
             self.ga_steps.zero_()
 
-    def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+    def forward(self, input: torch.Tensor):
         """
         Forward pass of the router.
 
         Args:
             input (torch.Tensor): Input tensor.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
         """
         self._maintain_float32_expert_bias()
 
@@ -650,7 +578,7 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No
             # Apply force load balancing with random logits for benchmark
             logits = apply_random_logits(logits)
 
-        probs, routing_map = self.routing(logits, padding_mask=padding_mask)
+        probs, routing_map = self.routing(logits)
 
         return probs, routing_map
 
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index 831b5546d53..ea4464b4784 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -453,18 +453,12 @@ def _checkpointed_forward(
         attention_bias: Tensor,
         packed_seq_params: PackedSeqParams,
         use_inner_quantization_context: bool,
-        padding_mask: Optional[Tensor] = None,
     ):
         """Forward method with activation checkpointing."""
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states,
-                attention_mask,
-                context,
-                context_mask,
-                rotary_pos_emb,
-                padding_mask=None,
+                hidden_states, attention_mask, context, context_mask, rotary_pos_emb
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
@@ -495,7 +489,6 @@ def custom_forward(
                             attention_bias=attention_bias,
                             inference_context=None,
                             packed_seq_params=packed_seq_params,
-                            padding_mask=padding_mask,
                         )
                 return hidden_states, context
 
@@ -515,7 +508,6 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
-                    padding_mask,
                 )
             else:
                 return tensor_parallel.checkpoint(
@@ -526,7 +518,6 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
-                    padding_mask,
                 )
 
         if self.config.recompute_method == 'uniform':
@@ -632,7 +623,6 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
-        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         dynamic_inference_decode_only: Optional[bool] = None,
@@ -742,7 +732,6 @@ def forward(
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,
                     use_inner_quantization_context=use_inner_quantization_context,
-                    padding_mask=padding_mask,
                 )
             else:
                 for l_no, layer in enumerate(self.layers):
@@ -775,7 +764,6 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
-                            padding_mask=padding_mask,
                         )
 
                     if (
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 97ca1ec222e..920c3b8fcba 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -1,6 +1,5 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-import functools
 import logging
 import warnings
 from abc import ABC
@@ -489,11 +488,7 @@ def forward(self, *args, **kwargs):
         # runners in the cuda graph manager
         kwargs.pop("dynamic_inference_decode_only", None)
         hidden_states, context = self._forward_attention(*args, **kwargs)
-        output = self._forward_mlp(
-            hidden_states,
-            kwargs.get("inference_context", None),
-            padding_mask=kwargs.get("padding_mask", None),
-        )
+        output = self._forward_mlp(hidden_states, kwargs.get("inference_context", None))
         return output, context
 
     def _forward_attention(
@@ -510,7 +505,6 @@ def _forward_attention(
         inference_context: Optional[Any] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
-        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[Any] = None,
     ):
@@ -641,18 +635,13 @@ def _forward_attention(
 
         return hidden_states, context
 
-    def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None):
+    def _forward_mlp(self, hidden_states, inference_context=None):
         """
         Perform a forward pass through the feed-forward layer.
 
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
-                Shape [seq_length, batch_size, hidden_size].
-            inference_context: Inference context for optimizations.
-            padding_mask (Tensor, optional): Padding mask for MoE routing.
-                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
-                Only used for MoE layers to exclude padding tokens from aux loss computations.
-                The MoELayer will internally transform this to [seq_length, bsz] format.
+
         Returns:
             output (Tensor): Transformed hidden states of shape [s, b, h].
         """
@@ -700,13 +689,10 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
                     tensor_parallel.random.get_cuda_rng_tracker,
                     self.pg_collection.tp,
                     pre_mlp_layernorm_output,
-                    padding_mask=padding_mask,
                 )
             else:
                 mlp_output_with_bias = tensor_parallel.checkpoint(
-                    functools.partial(self.mlp, padding_mask=padding_mask),
-                    False,
-                    pre_mlp_layernorm_output,
+                    self.mlp, False, pre_mlp_layernorm_output
                 )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
@@ -726,7 +712,7 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
                 # Set the residual for fused reduce-scatter + add + layer-norm + all-gather
                 # operation in MLP's fc2.
                 self._set_fc2_residual(residual)
-            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
+            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
 
         if self.recompute_pre_mlp_layernorm:
             # discard the output of the pre-mlp layernorm and register the recompute
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index 6c59dd3f9e3..81e61a3404a 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -23,7 +23,7 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-def build_model(config, use_padding_mask=False):
+def build_model(config):
     seq_len = 32
     max_seq_len = 300
     # ids = random.sample([i for i in range(max_seq_len)], seq_len)
@@ -39,12 +39,6 @@ def build_model(config, use_padding_mask=False):
         "attention_mask": torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda(),
     }
 
-    # Optionally add padding_mask with same shape as input_ids
-    if use_padding_mask:
-        padding_mask = torch.zeros((1, seq_len), dtype=torch.bool).cuda()
-        padding_mask[0, -8:] = True
-        data["padding_mask"] = padding_mask
-
     # build layer spec
     transformer_layer_spec = get_gpt_decoder_block_spec(config=config, use_transformer_engine=True)
     mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec.layer_specs[-1], True)
@@ -54,7 +48,7 @@ def build_model(config, use_padding_mask=False):
         config=config,
         transformer_layer_spec=transformer_layer_spec,
         mtp_block_spec=mtp_block_spec,
-        vocab_size=128,
+        vocab_size=100,
         pre_process=True,
         post_process=True,
         max_sequence_length=max_seq_len,
@@ -180,109 +174,3 @@ def test_1f1b_schedule_model_chunk(self, mtp_layers, dispatcher_type, fp8_flag,
                 gpt_models[i] = None
             gc.collect()
             torch.cuda.empty_cache()
-
-    @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0")
-    @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types())
-    @pytest.mark.parametrize("layers", [[2, 1], [1, 1]])
-    @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
-    def test_1f1b_schedule_model_chunk_with_padding_mask(self, dispatcher_type, layers, tp_size):
-        """
-        Verifies all-to-all overlap optimization with padding_mask produces
-        the same results as the reference implementation with various TP/EP/CP combinations.
-        """
-        # Re-initialize model parallel with the specified configuration
-        Utils.destroy_model_parallel()
-        Utils.initialize_model_parallel(
-            tensor_model_parallel_size=tp_size,
-            pipeline_model_parallel_size=1,
-            expert_model_parallel_size=4,
-            expert_tensor_parallel_size=1,
-        )
-        set_streams()
-
-        microbatches = 1
-
-        gpt_models = []
-        schedule_plans = []
-        ref_captures = []
-        datas = []
-
-        # create TransformerConfig
-        extra_kwargs = {
-            "moe_token_dispatcher_type": dispatcher_type,
-            "tensor_model_parallel_size": tp_size,
-            "sequence_parallel": tp_size > 1,
-        }
-        if dispatcher_type == "flex":
-            extra_kwargs["moe_flex_dispatcher_backend"] = "deepep"
-            extra_kwargs["moe_router_dtype"] = "fp32"
-        with deterministic_mode():
-            for layer_num in layers:
-                output_tensors = []
-                # build config
-                config = get_test_config(num_layers=layer_num, extra_kwargs=extra_kwargs)
-                # build model with padding_mask
-                gpt_model, schedule_plan, data = build_model(config, use_padding_mask=True)
-                gpt_model.cuda()
-                gpt_models.append(gpt_model)
-                datas.append(data)
-                schedule_plans.append(schedule_plan)
-
-                # run reference
-                for _ in range(microbatches):
-                    loss = gpt_model.forward(**data)
-                    loss = float16_to_fp32(loss)
-                    loss.backward(torch.ones_like(loss))
-                    output_tensors.append(loss)
-
-                capture = {"outputs": output_tensors}
-                for name, param in gpt_model.named_parameters():
-                    capture[name] = param.grad
-                ref_captures.append(capture)
-                gpt_model.zero_grad()
-            assert gpt_models[0].embedding is not None
-            assert gpt_models[1].embedding is not None
-            # run a2a overlap
-            capture_0 = {"outputs": []}
-            capture_1 = {"outputs": []}
-            a2a_captures = [capture_0, capture_1]
-            for i in range(microbatches):
-                # 1st forward
-                if i > 0:
-                    assert (
-                        schedule_plans[0].pre_process is None
-                    ), "pre_process should be released after backward"
-                    schedule_plans[0] = gpt_models[0].build_schedule_plan(**datas[0])
-                    schedule_plans[1] = gpt_models[1].build_schedule_plan(**datas[1])
-                f_input_0 = TransformerModelChunkSchedulePlan.run(schedule_plans[0], None)
-                capture_0["outputs"].append(f_input_0)
-                # overlap
-                f_input_1 = TransformerModelChunkSchedulePlan.run(
-                    schedule_plans[1], schedule_plans[0], b_grad=torch.ones_like(f_input_0)
-                )
-                capture_1["outputs"].append(f_input_1)
-                # last backward
-                TransformerModelChunkSchedulePlan.run(
-                    None, schedule_plans[1], b_grad=torch.ones_like(f_input_1)
-                )
-            for i in range(len(gpt_models)):
-                for name, param in gpt_models[i].named_parameters():
-                    a2a_captures[i][name] = param.grad
-
-            # compare results
-            for i in range(len(ref_captures)):
-                comp_res = compare_captures(ref_captures[i], a2a_captures[i], True, True)
-                assert comp_res[0], f"[rank {torch.distributed.get_rank()}] {comp_res[1]}"
-
-            # release resources is necessary, otherwise later testcases will oom
-            for i in range(len(schedule_plans)):
-                schedule_plans[i] = None
-                ref_captures[i] = None
-                a2a_captures[i] = None
-                for k in datas[i]:
-                    datas[i][k] = None
-                datas[i] = None
-                gpt_models[i].zero_grad()
-                gpt_models[i] = None
-            gc.collect()
-            torch.cuda.empty_cache()
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
index c6c4a75af99..0fd2c445c9f 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
@@ -502,8 +502,8 @@ def test_mtp_layer_overlap(self, dispatcher_type, fp8_flag):
             position_ids = torch.tensor(data, dtype=torch.int64).repeat((1, 1)).cuda()
             attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda()
             # get rotary pos emb
-            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _, _padding_mask = (
-                gpt_model._preprocess(input_ids, position_ids)
+            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _ = gpt_model._preprocess(
+                input_ids, position_ids
             )
             # reset model
             params = reset_model(gpt_model)
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index ccd11bf29af..621e200c2cb 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -577,185 +577,3 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
                 reduce_from_tensor_model_parallel_region(aux_loss, router.tp_cp_group)
             assert aux_loss.item() == 1, f"{aux_loss_type}: {aux_loss.item()}"
             clear_aux_losses_tracker()
-
-
-class TestPaddingMaskAuxLoss:
-    """Test padding mask support in various aux loss types."""
-
-    def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1, sequence_parallel=False):
-        """Initialize model parallel with given configuration.
-
-        Args:
-            tp_size: Tensor parallel size.
-            ep_size: Expert parallel size.
-            cp_size: Context parallel size.
-        """
-        Utils.initialize_model_parallel(
-            tensor_model_parallel_size=tp_size,
-            pipeline_model_parallel_size=1,
-            context_parallel_size=cp_size,
-            expert_model_parallel_size=ep_size,
-        )
-        _set_random_seed(seed_=123, data_parallel_random_init=False)
-
-        # Store parallel configuration
-        self.tp_size = tp_size
-        self.ep_size = ep_size
-        self.cp_size = cp_size
-
-        # Default configuration
-        self.default_transformer_config = TransformerConfig(
-            num_layers=1,
-            hidden_size=12,
-            num_attention_heads=8,
-            num_moe_experts=32,
-            use_cpu_initialization=True,
-            moe_router_load_balancing_type="aux_loss",
-            moe_router_topk=8,
-            moe_aux_loss_coeff=1.0,
-            bf16=True,
-            params_dtype=torch.bfloat16,
-            add_bias_linear=False,
-            tensor_model_parallel_size=tp_size,
-            expert_model_parallel_size=ep_size,
-            context_parallel_size=cp_size,
-            sequence_parallel=sequence_parallel and tp_size > 1,
-        )
-
-    def new_router(self, **kwargs):
-        """Create a new router with updated configuration."""
-        pg_collection = get_default_pg_collection()
-        new_transformer_config = dataclasses.replace(self.default_transformer_config, **kwargs)
-        router = TopKRouter(config=new_transformer_config, pg_collection=pg_collection)
-        router.set_layer_number(0)
-        return router
-
-    @pytest.mark.internal
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
-    @pytest.mark.parametrize(
-        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
-    )
-    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size):
-        """Test that padding tokens are correctly excluded from aux loss calculation."""
-        # Initialize model parallel with given configuration
-        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
-
-        try:
-            clear_aux_losses_tracker()
-
-            router = self.new_router(
-                moe_router_load_balancing_type=aux_loss_type,
-                moe_aux_loss_coeff=1.0,
-                moe_router_dtype="fp64",
-            ).cuda()
-
-            seq_len = 32
-            batch_size = 2
-            hidden_size = router.config.hidden_size
-
-            # Create input with padding
-            hidden_states_full = torch.randn(
-                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
-            )
-
-            # Create padding mask: first half valid, second half padding
-            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
-            padding_mask[seq_len // 2 :, :] = True
-
-            # Test with padding mask
-            router.weight.grad = None
-            scores_with_mask, routing_map_with_mask = router(
-                hidden_states_full, padding_mask=padding_mask
-            )
-            scores_with_mask.backward(torch.zeros_like(scores_with_mask))
-
-            loss_name = {
-                "aux_loss": "load_balancing_loss",
-                "seq_aux_loss": "seq_load_balancing_loss",
-                "global_aux_loss": "global_load_balancing_loss",
-            }[aux_loss_type]
-
-            tracker = get_moe_layer_wise_logging_tracker()
-            aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
-            grad_with_mask = router.weight.grad.clone()
-
-            # Test without padding (with only half of the tokens)
-            clear_aux_losses_tracker()
-            router.weight.grad = None
-            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
-            scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
-            scores_without_mask.backward(torch.zeros_like(scores_without_mask))
-
-            aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
-            grad_without_mask = router.weight.grad.clone()
-
-            # The aux loss with mask should be equal to the aux loss without mask
-            assert torch.equal(aux_loss_with_mask, aux_loss_without_mask)
-            assert torch.equal(grad_with_mask, grad_without_mask)
-
-            clear_aux_losses_tracker()
-        finally:
-            # Always cleanup model parallel
-            Utils.destroy_model_parallel()
-
-    @pytest.mark.internal
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @pytest.mark.parametrize(
-        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
-    )
-    def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
-        """Test that padding mask works correctly with z_loss."""
-        # Initialize model parallel with given configuration
-        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
-
-        try:
-            clear_aux_losses_tracker()
-
-            router = self.new_router(
-                moe_router_load_balancing_type="aux_loss",
-                moe_aux_loss_coeff=0.0,
-                moe_z_loss_coeff=1.0,
-                moe_router_dtype="fp32",
-            ).cuda()
-
-            seq_len = 32
-            batch_size = 2
-            hidden_size = router.config.hidden_size
-
-            # Create input
-            hidden_states_full = torch.randn(
-                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
-            )
-
-            # Create padding mask: first half valid, second half padding
-            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
-            padding_mask[seq_len // 2 :, :] = True
-
-            # Test with padding mask
-            router.weight.grad = None
-            scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
-            scores_with_mask.sum().backward()
-
-            tracker = get_moe_layer_wise_logging_tracker()
-            z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
-            grad_with_mask = router.weight.grad.clone()
-
-            # Test without padding (with only half of the tokens)
-            clear_aux_losses_tracker()
-            router.weight.grad = None
-            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
-            scores_without_mask, _ = router(hidden_states_valid)
-            scores_without_mask.sum().backward()
-
-            z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
-            grad_without_mask = router.weight.grad.clone()
-
-            # The z_loss with mask should be close to the z_loss without mask
-            assert torch.equal(z_loss_with_mask, z_loss_without_mask)
-            assert torch.equal(grad_with_mask, grad_without_mask)
-
-            clear_aux_losses_tracker()
-        finally:
-            # Always cleanup model parallel
-            Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index 4d6b5ee2c3e..904595928de 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -127,53 +127,6 @@ def test_aux_loss(self):
         out.sum().mul_(0).backward()
         assert self.sequential_mlp.router.weight.grad.abs().sum() > 0
 
-    @pytest.mark.internal
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    def test_router_with_padding_mask(self):
-        """Test that padding mask correctly excludes padding tokens from routing."""
-        self.router = self.router.cuda()
-        seq_len = 32
-        batch_size = 2
-        hidden_size = self.router.config.hidden_size
-
-        # Create input with shape [seq_len, batch_size, hidden_size]
-        hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16()
-
-        # Create padding mask: first half valid, second half padding
-        # padding_mask shape: [seq_len, batch_size]
-        # Convention: True = padding (exclude), False = valid (include)
-        padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2 :, :] = True  # Second half is padding
-
-        # Test forward pass with padding mask
-        with torch.no_grad():
-            probs_with_mask, routing_map_with_mask = self.router(
-                hidden_states, padding_mask=padding_mask
-            )
-
-            # Test forward pass without padding mask (only valid tokens)
-            hidden_states_valid = hidden_states[: seq_len // 2, :, :]
-            probs_without_mask, routing_map_without_mask = self.router(hidden_states_valid)
-
-            # The valid part of routing with mask should match routing without mask
-            probs_valid_part = probs_with_mask.reshape(seq_len, batch_size, -1)[
-                : seq_len // 2, :, :
-            ]
-            probs_valid_part = probs_valid_part.reshape(-1, probs_valid_part.shape[-1])
-
-            # Check that shapes are as expected
-            assert probs_with_mask.shape == (
-                seq_len * batch_size,
-                self.router.config.num_moe_experts,
-            )
-            assert routing_map_with_mask.shape == (
-                seq_len * batch_size,
-                self.router.config.num_moe_experts,
-            )
-
-            # Verify that probs for valid tokens are similar
-            assert torch.equal(probs_valid_part, probs_without_mask)
-
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_router_dtype(self):

From dbde759da15b154e94f17a60e0a231d035c5ea84 Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Sat, 24 Jan 2026 09:10:11 -0800
Subject: [PATCH 12/79] Add ability to save wgrads and dgrads (#3032)

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 .../distributed/distributed_data_parallel.py  |   9 +-
 .../core/distributed/finalize_model_grads.py  |   3 +-
 .../core/distributed/param_and_grad_buffer.py |  22 ++--
 megatron/core/pipeline_parallel/schedules.py  |  16 ++-
 megatron/training/arguments.py                |   4 +
 megatron/training/checkpointing.py            |  36 +++++
 megatron/training/dgrad_logging.py            | 123 ++++++++++++++++++
 megatron/training/training.py                 |  40 +++++-
 .../distributed/test_param_and_grad_buffer.py |  55 ++++++++
 tests/unit_tests/test_training.py             |  59 +++++++++
 10 files changed, 350 insertions(+), 17 deletions(-)
 create mode 100644 megatron/training/dgrad_logging.py

diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index e831d7cf4ec..421b279b17d 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -93,6 +93,7 @@ def __init__(
         # disable_bucketing is True (e.g., we might not want to break up model parameters
         # into buckets for model chunks after the first in the interleaved schedule).
         self.bucket_size = self.ddp_config.bucket_size
+        self.force_all_reduce = False
         if isinstance(self.pp_group, list):
             pp_rank = self.pp_group[0].rank()
         else:
@@ -440,7 +441,9 @@ def hook(*unused):
                 param.grad = None
 
                 if self.ddp_config.overlap_grad_reduce:
-                    self.param_to_bucket_group[param].register_grad_ready(param)
+                    self.param_to_bucket_group[param].register_grad_ready(
+                        param, self.force_all_reduce
+                    )
 
         return hook
 
@@ -519,7 +522,7 @@ def start_grad_sync(self, *unused):
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
             bucket_group.start_grad_sync()
 
-    def finish_grad_sync(self):
+    def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Finishes grad sync (all-reduce or reduce-scatter) communication operations
         for all model gradients.
@@ -529,7 +532,7 @@ def finish_grad_sync(self):
         communication ops.
         """
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
-            bucket_group.finish_grad_sync()
+            bucket_group.finish_grad_sync(force_all_reduce=force_all_reduce)
 
     def scale_gradients(self, scaling_factor: float):
         """Scale all gradients inside the buffers by `scaling_factor`."""
diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py
index ddaeb7e8d84..a52592bb269 100644
--- a/megatron/core/distributed/finalize_model_grads.py
+++ b/megatron/core/distributed/finalize_model_grads.py
@@ -397,6 +397,7 @@ def finalize_model_grads(
     model: List[torch.nn.Module],
     num_tokens: Optional[torch.Tensor] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """
     All-reduce all model grads across DP replicas, layernorm grads for sequence parallelism,
@@ -439,7 +440,7 @@ def finalize_model_grads(
     if config.timers is not None:
         config.timers('all-grads-sync', log_level=1).start(barrier=config.barrier_with_L1_time)
     for model_chunk in model:
-        model_chunk.finish_grad_sync()
+        model_chunk.finish_grad_sync(force_all_reduce=force_all_reduce)
     if config.timers is not None:
         config.timers('all-grads-sync').stop()
 
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index 50cf3e0ea37..e7e32ddc081 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -355,7 +355,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
                 if len(fp8_params) > 0:
                     post_all_gather_processing(fp8_params)
 
-    def start_grad_sync(self):
+    def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Initiates grad sync (all-reduce or reduce-scatter) communication operations
         for all buckets in the bucket group.
@@ -426,7 +426,7 @@ def start_grad_sync(self):
         grad_reduce_handle = None
         with stream_context, _coalescing_manager(communication_group, async_ops=async_op) as cm:
             for idx, bucket in enumerate(self.buckets):
-                if self.ddp_config.use_distributed_optimizer:
+                if self.ddp_config.use_distributed_optimizer and not force_all_reduce:
                     if self.cached_grad_buffer_shard_list[idx] is None:
                         self.cached_grad_buffer_shard_list[idx] = shard_buffer(
                             bucket.grad_data, self.intra_distributed_optimizer_instance_size
@@ -442,6 +442,10 @@ def start_grad_sync(self):
                         async_op=async_op,
                     )
                 else:
+                    if torch.distributed.get_rank() == 0 and force_all_reduce:
+                        logger.info(
+                            f"Performing reduction using all_reduce because {force_all_reduce=}"
+                        )
                     torch.distributed.all_reduce(
                         bucket.grad_data, op=reduce_op, group=communication_group, async_op=async_op
                     )
@@ -476,7 +480,7 @@ def start_grad_sync(self):
                     )
 
         if async_op:
-            if self.ddp_config.reduce_scatter_with_fp32_accumulation:
+            if self.ddp_config.reduce_scatter_with_fp32_accumulation and not force_all_reduce:
                 assert (
                     len(self.buckets) == 1
                 ), "Only 1 bucket supported with reduce_scatter_with_fp32_accumulation=True"
@@ -494,7 +498,7 @@ def start_grad_sync(self):
             # None.
             self.grad_reduce_handle = None
 
-    def finish_grad_sync(self):
+    def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Finishes grad sync (all-reduce or reduce-scatter) communication operations
         for all buckets in the bucket group.
@@ -506,13 +510,13 @@ def finish_grad_sync(self):
         self.param_gather_dispatched = False
         # If overlap_grad_reduce is False, start (and finish) synchronous communication call here.
         if not self.ddp_config.overlap_grad_reduce:
-            self.start_grad_sync()
+            self.start_grad_sync(force_all_reduce=force_all_reduce)
             return
         # If first batch, start asynchronous communication here. register_grad_ready() launches
         # asynchronous communication only once self.golden_per_param_grad_ready_counts is
         # populated at the end of this first batch.
         if self.is_first_batch:
-            self.start_grad_sync()
+            self.start_grad_sync(force_all_reduce=force_all_reduce)
         # When using multiple DistOpt instances, we don't need to sync here as we launch
         # communications on a separate communication stream.
         if self.ddp_config.num_distributed_optimizer_instances > 1:
@@ -526,7 +530,9 @@ def finish_grad_sync(self):
         self.grad_reduce_handle.wait()
         self.grad_reduce_handle = None
 
-    def register_grad_ready(self, param: torch.nn.Parameter):
+    def register_grad_ready(
+        self, param: torch.nn.Parameter, force_all_reduce: Optional[bool] = False
+    ):
         """
         Registers grads for the passed-in param to be "ready" for grad sync.
 
@@ -546,7 +552,7 @@ def register_grad_ready(self, param: torch.nn.Parameter):
             if not self.is_first_batch:
                 if self.per_param_grad_ready_counts == self.golden_per_param_grad_ready_counts:
                     assert len(self.per_param_grad_ready_counts) == len(self.params)
-                    self.start_grad_sync()
+                    self.start_grad_sync(force_all_reduce=force_all_reduce)
 
 
 class _ParamAndGradBuffer:
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index f1bfc340faf..edca62be375 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -111,9 +111,9 @@ def forward_step(data_iterator, model):
     decoder_seq_length (int, optional): The sequence length for the decoder in a dual-stack
         transformer. This is ignored for a single-stack transformer.
 
-    forward_only (optional, default = False): Perform only the forward step
+    forward_only (optional, default = False): Perform only the forward step.
 
-    collect_non_loss_data (optional, bool, default=False): TODO
+    collect_non_loss_data (optional, bool, default=False): TODO.
 
     first_val_step (bool, optional): Is the first step of the validation phase. Used by
         Transformer Engine modules to only update their fp8 weights only on the first validation
@@ -125,11 +125,17 @@ def forward_step(data_iterator, model):
         respective list of shapes. Thus it is not used in the other forward-backward functions
         which have different shape handling.
 
+    force_all_reduce (bool, optional): If true, force use of all-reduce for gradient reduction
+        instead of reduce-scatter (if using distributed optimizer) in this iteration to ensure all
+        data-parallel ranks have fully reduced gradients. This is useful for easier wgrad saving
+        (can just inspect DP replica 0 to get full set of wgrads for entire model).
+
     Args:
         pp_size (Optional[int]): Pipeline model parallel size to use.
         vp_size (Optional[int]): Virtual pipeline model parallel size to use.
             If both pp_size and vp_size are None, both values fall back to parallel_state.
             Otherwise, provided values are used as-is and None is treated as an explicit input.
+
     """
     if pp_size is None and vp_size is None:
         pp_size = parallel_state.get_pipeline_model_parallel_world_size()
@@ -524,6 +530,7 @@ def forward_backward_no_pipelining(
     adjust_tensor_shapes_fn: Optional[Callable] = None,  # unused
     p2p_communicator: Optional[P2PCommunicator] = None,  # unused
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """Run forward and backward passes with no pipeline parallelism"""
 
@@ -674,6 +681,7 @@ def forward_backward_no_pipelining(
             [model],
             total_num_tokens if config.calculate_per_token_loss else None,
             pg_collection=pg_collection,
+            force_all_reduce=force_all_reduce,
         )
 
     if not forward_only and config.fine_grained_activation_offloading:
@@ -832,6 +840,7 @@ def forward_backward_pipelining_with_interleaving(
     adjust_tensor_shapes_fn: Optional[Callable] = None,  # unused
     p2p_communicator: Optional[P2PCommunicator] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """Run interleaved 1F1B schedule (model split into model chunks), with
     communication between pipeline stages as needed.
@@ -1911,6 +1920,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None):
             model,
             total_num_tokens if config.calculate_per_token_loss else None,
             pg_collection=pg_collection,
+            force_all_reduce=force_all_reduce,
         )
 
     if not forward_only and config.fine_grained_activation_offloading:
@@ -1974,6 +1984,7 @@ def forward_backward_pipelining_without_interleaving(
     adjust_tensor_shapes_fn: Optional[Callable] = None,
     p2p_communicator: Optional[P2PCommunicator] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """Run non-interleaved 1F1B schedule, with communication between pipeline
     stages. Returns dictionary with losses if the last stage, empty dict otherwise."""
@@ -2301,6 +2312,7 @@ def enable_grad_sync():
             [model],
             total_num_tokens if config.calculate_per_token_loss else None,
             pg_collection=pg_collection,
+            force_all_reduce=force_all_reduce,
         )
 
     if not forward_only and config.fine_grained_activation_offloading:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 9fadb3f9900..3c2a0e52f71 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2597,6 +2597,10 @@ def _add_checkpointing_args(parser):
                        help='Output directory to save checkpoints to.')
     group.add_argument('--save-interval', '--persistent-save-interval', type=int, default=None,
                        help='Number of iterations between persistent checkpoint saves.')
+    group.add_argument('--save-wgrads-interval', type=int, default=None,
+                       help='Number of iterations between wgrad (main_grad) saves.')
+    group.add_argument('--save-dgrads-interval', type=int, default=None,
+                       help='Number of iterations between dgrad saves.')
     group.add_argument('--save-retain-interval', type=int, default=None,
                        help='Number of iterations between retained checkpoints (other'
                        'checkpoints _except the last checkpoint_ are automatically deleted).')
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index a4605121f82..93e8a8acd6a 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -10,6 +10,7 @@
 import threading
 import types
 from argparse import Namespace
+from datetime import datetime
 from enum import Enum, auto
 from logging import getLogger
 from pathlib import Path
@@ -432,6 +433,41 @@ def _build_sharded_state_dict_metadata(args: Namespace, dp_cp_group: Optional[to
     metadata['dp_cp_group'] = dp_cp_group
     return metadata
 
+
+def save_grads(save_dir, state_dict, iteration, grad_label):
+    """Persist state_dict of grads onto disk. In case of wgrads, this collection should
+    be performed before the grads are cleared but after they are reduced.
+
+    NOTE: wgrads for non-expert layers will be duplicated if using expert parallelism, but
+    this can be handled in postprocessing."""
+
+    print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} "
+                 f"from iteration {iteration:7d}")
+
+    if mpu.get_expert_data_parallel_rank() == 0:
+        # Create saving directory.
+        ep_rank = mpu.get_expert_model_parallel_rank()
+        pp_rank = mpu.get_pipeline_model_parallel_rank()
+        tp_rank = mpu.get_tensor_model_parallel_rank()
+        assert save_dir is not None
+        assert iteration is not None
+        save_dir = os.path.join(save_dir, grad_label, f"iter_{iteration:07d}")
+        os.makedirs(save_dir, exist_ok=True)
+
+        # Save state_dict.
+        checkpoint_name = f"mp_rank_{tp_rank:02d}"
+        if mpu.get_pipeline_model_parallel_world_size() > 1:
+            checkpoint_name += f"_{pp_rank:03d}"
+        if mpu.get_expert_model_parallel_world_size() > 1:
+            checkpoint_name += f"_{ep_rank:03d}"
+        full_save_path = os.path.join(save_dir, f"{checkpoint_name}.pth")
+        # Convert back to dict (e.g., from collections.defaultdict) for easy loading later.
+        torch.save(dict(state_dict), full_save_path)
+
+    print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} "
+                 f"from iteration {iteration:7d}")
+
+
 def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far,
                     checkpointing_context=None, pipeline_rank=None, expert_rank=None, tensor_rank=None, pipeline_parallel=None, expert_parallel=None, non_persistent_ckpt=False,
                     train_data_iterator=None, preprocess_common_state_dict_fn = None, release=False, tp_group: Optional[torch.distributed.ProcessGroup] = None, pp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None):
diff --git a/megatron/training/dgrad_logging.py b/megatron/training/dgrad_logging.py
new file mode 100644
index 00000000000..c046b4709fb
--- /dev/null
+++ b/megatron/training/dgrad_logging.py
@@ -0,0 +1,123 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+
+"""dgrad logging using backward hooks."""
+
+from collections import defaultdict
+import torch
+import torch.nn as nn
+
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+
+from .checkpointing import save_grads
+from .utils import unwrap_model
+
+
+def _get_linear_types():
+    """Build tuple of linear layer types to capture gradients from."""
+    types = [nn.Linear, nn.Embedding, ColumnParallelLinear, RowParallelLinear]
+
+    # Add Transformer Engine layers if available.
+    try:
+        from megatron.core.extensions.transformer_engine import (
+            TELinear,
+            TEColumnParallelLinear,
+            TERowParallelLinear,
+            TELayerNormColumnParallelLinear,
+        )
+        types.extend([TELinear, TEColumnParallelLinear, TERowParallelLinear,
+                      TELayerNormColumnParallelLinear])
+    except ImportError:
+        pass
+
+    try:
+        from megatron.core.extensions.transformer_engine import (
+            TEGroupedLinear,
+            TEColumnParallelGroupedLinear,
+            TERowParallelGroupedLinear,
+        )
+        if TEGroupedLinear is not None:
+            types.extend([TEGroupedLinear, TEColumnParallelGroupedLinear,
+                          TERowParallelGroupedLinear])
+    except ImportError:
+        pass
+
+    return tuple(types)
+
+
+LINEAR_TYPES = _get_linear_types()
+
+
+class DataGradLogger:
+    """Captures and saves gradients from all linear layers using backward hooks.
+    
+    NOTE: Right now, we only save the dgrads for the last microbatch in a batch on DP replica 0.
+    The code below would need to be extended to save dgrads for all microbatches in a batch."""
+
+    def __init__(self, save_dir: str):
+        self._save_dir = save_dir
+        self._dgrads_state_dict = defaultdict(dict)
+        self._hooks = []
+
+    def _make_hook(self, model_chunk_name: str, module_name: str):
+        """Create a backward hook for a named module."""
+        def hook(_, grad_input, grad_output):
+            for idx, grad in enumerate(grad_output):
+                if grad is not None:
+                    grad_name = f"{module_name}/output{idx}"
+                    self._dgrads_state_dict[model_chunk_name][grad_name] = grad.detach().cpu()
+            for idx, grad in enumerate(grad_input):
+                if grad is not None:
+                    grad_name = f"{module_name}/input{idx}"
+                    self._dgrads_state_dict[model_chunk_name][grad_name] = grad.detach().cpu()
+        return hook
+
+    def save(self, iteration: int):
+        """Save captured gradients to disk and clear the buffer."""
+        if not self._dgrads_state_dict:
+            return
+        save_grads(self._save_dir, self._dgrads_state_dict, iteration, "dgrads")
+        self._dgrads_state_dict.clear()
+
+    def register_hooks(self, model: torch.nn.Module):
+        """Find and register hooks on all linear layers."""
+        assert len(self._hooks) == 0
+        for model_chunk_id, model_chunk in enumerate(model):
+            unwrapped_model_chunk = unwrap_model(model_chunk)
+            for module_name, module in unwrapped_model_chunk.named_modules():
+                if isinstance(module, LINEAR_TYPES):
+                    model_chunk_name = f"model_chunk{model_chunk_id}"
+                    handle = module.register_full_backward_hook(
+                        self._make_hook(model_chunk_name, module_name)
+                    )
+                    self._hooks.append(handle)
+
+    def remove_hooks(self):
+        """Remove all registered hooks."""
+        for handle in self._hooks:
+            handle.remove()
+        self._hooks.clear()
+
+
+_LOGGER = None
+
+
+def enable_dgrad_logging(model: torch.nn.Module, save_dir: str):
+    """Enable dgrad logging on a model."""
+    global _LOGGER
+    if _LOGGER is None:
+        _LOGGER = DataGradLogger(save_dir)
+    _LOGGER.register_hooks(model)
+
+
+def disable_dgrad_logging():
+    """Disable dgrad logging on a model."""
+    global _LOGGER
+    assert _LOGGER is not None
+    _LOGGER.remove_hooks()
+
+
+def save_dgrads(iteration: int):
+    """Save dgrads to disk."""
+    global _LOGGER
+    assert _LOGGER is not None
+    _LOGGER.save(iteration)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 08a655ba80d..b7040e7bbb9 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -31,6 +31,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
         _STARTUP_TIMESTAMPS['main_entry'] = main_entry
 
 
+from collections import defaultdict
 import copy
 import dataclasses
 from datetime import datetime, timedelta
@@ -100,7 +101,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 )
 from megatron.core.optimizer import get_standard_config_overrides
 from megatron.training.checkpointing import load_checkpoint
-from megatron.training.checkpointing import save_checkpoint
+from megatron.training.checkpointing import save_checkpoint, save_grads
 from megatron.training.checkpointing import checkpoint_exists
 from megatron.training.checkpointing import get_loaded_iteration
 from megatron.core.full_cuda_graph import FullCudaGraphWrapper
@@ -187,6 +188,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
     get_energy_monitor,
 )
 from . import one_logger_utils
+from .dgrad_logging import enable_dgrad_logging, disable_dgrad_logging, save_dgrads
 
 from . import ft_integration
 
@@ -1604,16 +1606,22 @@ def dummy_train_step(data_iterator):
             batch = get_batch_on_this_cp_rank(batch)
 
 
-def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func):
+def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None):
     """Single training step."""
     args = get_args()
     timers = get_timers()
 
     rerun_state_machine = get_rerun_state_machine()
+    save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and
+                                     (iteration + 1) % args.save_dgrads_interval == 0)
+    save_wgrads_in_this_iteration = (args.save_wgrads_interval is not None and
+                                     (iteration + 1) % args.save_wgrads_interval == 0)
     while rerun_state_machine.should_run_forward_backward(data_iterator):
         # Set grad to zero.
         for model_chunk in model:
             model_chunk.zero_grad_buffer()
+            # If saving main_grads in this iteration, then all-reduce instead of reduce-scatter.
+            model_chunk.force_all_reduce = save_wgrads_in_this_iteration
         optimizer.zero_grad()
 
         if has_nvidia_modelopt:
@@ -1636,6 +1644,8 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
                     optim_instance._copy_main_params_to_param_buffer()
 
         # Forward pass.
+        if save_dgrads_in_this_iteration:
+            enable_dgrad_logging(model, args.save)
         losses_reduced = forward_backward_func(
             forward_step_func=forward_step_func,
             data_iterator=data_iterator,
@@ -1646,7 +1656,31 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
             decoder_seq_length=args.decoder_seq_length,
             forward_only=False,
             adjust_tensor_shapes_fn=adjust_tensor_shapes_fn,
+            force_all_reduce=save_wgrads_in_this_iteration,
         )
+        if save_dgrads_in_this_iteration:
+            save_dgrads(iteration + 1)
+            disable_dgrad_logging()
+
+        # Reset force_all_reduce field.
+        for model_chunk in model:
+            model_chunk.force_all_reduce = False
+
+    # Checkpoint main_grads.
+    if save_wgrads_in_this_iteration:
+        # Collect state_dict of wgrads (each param's .main_grad field).
+        state_dict = defaultdict(dict)
+        for model_chunk_id, model_chunk in enumerate(model):
+            model_chunk_name = f"model_chunk{model_chunk_id}"
+            unwrapped_model_chunk = unwrap_model(model_chunk)
+            for param_name, param in unwrapped_model_chunk.named_parameters():
+                if getattr(param, "main_grad", None) is not None:
+                    main_grad_on_cpu = param.main_grad.cpu()
+                    state_dict[model_chunk_name][param_name] = main_grad_on_cpu
+
+        # iteration is 0-indexed, move to 1-indexed for checkpoint name and logging.
+        save_grads(args.save, state_dict, iteration + 1, "wgrads")
+
     should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit()
     if should_exit:
         return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0
@@ -2748,7 +2782,7 @@ def get_e2e_base_metrics():
             num_zeros_in_grad,
             max_attention_logit,
         ) = train_step(
-            forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func
+            forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration
         )
         ft_integration.on_training_step_end()
         if should_checkpoint:
diff --git a/tests/unit_tests/distributed/test_param_and_grad_buffer.py b/tests/unit_tests/distributed/test_param_and_grad_buffer.py
index ac0c6a6c422..609b2cc5a71 100644
--- a/tests/unit_tests/distributed/test_param_and_grad_buffer.py
+++ b/tests/unit_tests/distributed/test_param_and_grad_buffer.py
@@ -3,6 +3,7 @@
 import contextlib
 import math
 from typing import Optional
+from unittest import mock
 
 import pytest
 import torch
@@ -265,3 +266,57 @@ def test_grad_sync(
             param_and_grad_buffer.grad_data.data.fill_(1.0)
 
     Utils.destroy_model_parallel()
+
+
+@pytest.mark.parametrize("force_all_reduce", [False, True])
+def test_force_all_reduce_uses_correct_collective(force_all_reduce: bool):
+    """Test that force_all_reduce=True causes all-reduce to be used instead of reduce-scatter."""
+    Utils.initialize_model_parallel()
+
+    input_dim = 100
+    output_dim = 100
+    num_layers = 2
+    model, param_and_grad_buffer, _ = get_model_and_buffers(
+        input_dim=input_dim,
+        output_dim=output_dim,
+        num_layers=num_layers,
+        bias=True,
+        shared_embedding=False,
+        bucket_size=None,
+        use_distributed_optimizer=True,  # This normally uses reduce-scatter.
+        overlap_grad_reduce=False,
+        average_in_collective=False,
+    )
+
+    # Mock the collective operations to track which one is called.
+    with (
+        mock.patch('torch.distributed.all_reduce') as mock_all_reduce,
+        mock.patch(
+            'megatron.core.distributed.param_and_grad_buffer.dist_reduce_scatter_func'
+        ) as mock_reduce_scatter,
+    ):
+        # Set up the mocks to be no-ops.
+        mock_all_reduce.return_value = None
+        mock_reduce_scatter.return_value = None
+
+        # Trigger the grad sync via the DDP model's finish_grad_sync method.
+        model.finish_grad_sync(force_all_reduce=force_all_reduce)
+
+        if force_all_reduce:
+            # When force_all_reduce=True, all_reduce should be called.
+            assert (
+                mock_all_reduce.called
+            ), "Expected all_reduce to be called when force_all_reduce=True"
+            assert (
+                not mock_reduce_scatter.called
+            ), "Expected reduce_scatter NOT to be called when force_all_reduce=True"
+        else:
+            # When force_all_reduce=False with distributed optimizer, reduce_scatter should be called.
+            assert (
+                mock_reduce_scatter.called
+            ), "Expected reduce_scatter to be called when force_all_reduce=False"
+            assert (
+                not mock_all_reduce.called
+            ), "Expected all_reduce NOT to be called when force_all_reduce=False"
+
+    Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/test_training.py b/tests/unit_tests/test_training.py
index fef4bfbc5ef..2fd85724826 100644
--- a/tests/unit_tests/test_training.py
+++ b/tests/unit_tests/test_training.py
@@ -1,10 +1,16 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+from collections import defaultdict
+from pathlib import Path
 from types import SimpleNamespace
 
+import torch
+
+from megatron.training.checkpointing import save_grads
 from megatron.training.global_vars import set_args
 from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
 from megatron.training.training import build_train_valid_test_data_iterators
+from tests.unit_tests.dist_checkpointing import TempNamedDir
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -74,3 +80,56 @@ def old_round_impl(after, multiple):
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
+
+
+class TestSaveGrads:
+    """Tests for the save_grads function."""
+
+    def setup_method(self, method):
+        Utils.initialize_model_parallel(1, 1)
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    def test_save_grads(self, tmp_path_dist_ckpt):
+        """Test that save_grads creates the correct directory structure and saves
+        state_dict correctly.
+
+        With TP=1, PP=1 on 8 GPUs, we have 8 DP ranks. Only the rank with
+        expert_data_parallel_rank==0 should save. All ranks verify the result.
+        """
+        save_dir = str(tmp_path_dist_ckpt / "test_save_grads")
+
+        with TempNamedDir(save_dir, sync=True) as save_dir:
+            # Create a mock state_dict with gradients (use deterministic values for reproducibility).
+            state_dict = defaultdict(dict)
+            state_dict["model_chunk0"]["layer.weight"] = torch.arange(16).reshape(4, 4).float()
+            state_dict["model_chunk0"]["layer.bias"] = torch.arange(4).float()
+
+            iteration = 100
+            grad_label = "wgrads"
+
+            # All ranks call save_grads, but only expert_data_parallel_rank==0 actually saves.
+            save_grads(save_dir, dict(state_dict), iteration, grad_label)
+
+            # Synchronize before checking results since only rank 0 saves.
+            torch.distributed.barrier()
+
+            # All ranks verify the file was created by rank 0.
+            expected_dir = Path(save_dir) / grad_label / f"iter_{iteration:07d}"
+            assert expected_dir.exists(), f"Expected directory {expected_dir} to exist"
+
+            expected_file = expected_dir / "mp_rank_00.pth"
+            assert expected_file.exists(), f"Expected file {expected_file} to exist"
+
+            # Verify saved content.
+            loaded = torch.load(expected_file)
+            assert "model_chunk0" in loaded
+            assert "layer.weight" in loaded["model_chunk0"]
+            assert "layer.bias" in loaded["model_chunk0"]
+            assert torch.equal(
+                loaded["model_chunk0"]["layer.weight"], state_dict["model_chunk0"]["layer.weight"]
+            )
+            assert torch.equal(
+                loaded["model_chunk0"]["layer.bias"], state_dict["model_chunk0"]["layer.bias"]
+            )

From 55fe705062d1f6af6da7435ca605dc49073c3968 Mon Sep 17 00:00:00 2001
From: Charlie Truong <chtruong@nvidia.com>
Date: Sat, 24 Jan 2026 13:04:44 -0600
Subject: [PATCH 13/79] ci: Mark test_mode_partial_cudagraph unit tests as
 flaky (#3064)

Signed-off-by: Charlie Truong <chtruong@nvidia.com>
---
 .../a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py  | 2 ++
 tests/unit_tests/conftest.py                              | 8 ++++++++
 tests/unit_tests/transformer/test_cuda_graphs.py          | 2 ++
 3 files changed, 12 insertions(+)

diff --git a/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py
index 719bd5df18f..85586095bd7 100644
--- a/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py
@@ -333,6 +333,8 @@ def _run_test_helper(
 
         return loss_list
 
+    @pytest.mark.flaky
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(
         not (HAVE_TE and is_te_min_version("2.10.0")),
         reason="Partial CUDA graph support requires TransformerEngine version >= 2.10.0",
diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py
index 362d102200e..e251a3c1e7e 100644
--- a/tests/unit_tests/conftest.py
+++ b/tests/unit_tests/conftest.py
@@ -38,6 +38,14 @@ def pytest_sessionfinish(session, exitstatus):
         session.exitstatus = 0
 
 
+@pytest.fixture(scope="session", autouse=True)
+def cleanup():
+    yield
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
 @pytest.fixture(scope="function", autouse=True)
 def set_env():
     if is_te_min_version("1.3"):
diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py
index d4866b4839e..4696a3ed439 100644
--- a/tests/unit_tests/transformer/test_cuda_graphs.py
+++ b/tests/unit_tests/transformer/test_cuda_graphs.py
@@ -1013,6 +1013,8 @@ def _run_test_helper(
 
         return torch.tensor(loss_list)
 
+    @pytest.mark.flaky
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(
         not (HAVE_TE and is_te_min_version("2.10.0")),
         reason="Partial CUDA graph UT support requires TransformerEngine version >= 2.10.0",

From 3a7d74dacb977508d8c0e321ada26bb3db7b226a Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Sat, 24 Jan 2026 13:38:20 -0800
Subject: [PATCH 14/79] Keep FSDP's and DDP's finish_grad_sync API identical
 (#3070)

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 .../core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
index c1c11721f7e..f3708a35dd8 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
@@ -1105,7 +1105,7 @@ def attach_grad_to_optimizer_state(self):
         """
         self.param_and_grad_buffer.update_main_grads()
 
-    def finish_grad_sync(self):
+    def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Finishes grad sync (all-reduce or reduce-scatter) communication operations
         for all model gradients. Call prior to the optimization step to resolve
@@ -1114,6 +1114,9 @@ def finish_grad_sync(self):
         When overlap_grad_reduce is set to True, waits for asynchronous communication
         calls to complete. When overlap_grad_reduce is set to False, calls synchronous
         communication ops.
+
+        NOTE: force_all_reduce is included as an argument to maintain API compatibility
+        with DDP.force_grad_sync.
         """
         # Synchronize gradient reduce-scatter operations for all model gradients.
         self.synchronize_gradient_reduce()

From 389436b33e88ccdde4098851a60ff792c7ee8e48 Mon Sep 17 00:00:00 2001
From: Jon Barker <jbarker@nvidia.com>
Date: Sat, 24 Jan 2026 15:17:26 -0700
Subject: [PATCH 15/79] (REPLAY) Bug fix with
 --no-use-tokenizer-from-checkpoint-args (#3059)

Co-authored-by: Jon Barker <jbarker@oci-hsg-cs-001-vscode-01.cm.cluster>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
---
 megatron/training/checkpointing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 93e8a8acd6a..b73466ccfde 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1449,10 +1449,10 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('moe_latent_size', force=True)
 
     # Tokenizer args.
-    _set_arg('tokenizer_type', force=True)
     # Using checkpoint version might not always be safe (e.g., if running on different cluster).
     if args.use_tokenizer_model_from_checkpoint_args:
         _set_arg('tokenizer_model', force=True)
+        _set_arg('tokenizer_type', force=True)
     _set_arg('tiktoken_pattern', force=True)
     _set_arg('padded_vocab_size')
 

From 53a2b19a26f15e2edd363db6b5b29d2baf7420ac Mon Sep 17 00:00:00 2001
From: Siddharth Singh <136645615+sidsingh-nvidia@users.noreply.github.com>
Date: Sat, 24 Jan 2026 17:19:51 -0600
Subject: [PATCH 16/79] Optimizing post-processing of requests (#2920)

---
 .../data_parallel_inference_coordinator.py    | 24 ++++-
 .../core/inference/engines/dynamic_engine.py  | 35 ++++---
 megatron/core/inference/inference_request.py  | 96 +++++++++++--------
 3 files changed, 101 insertions(+), 54 deletions(-)

diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py
index 7f81065bcd4..3a1747facb2 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator.py
@@ -65,7 +65,7 @@ class DataParallelInferenceCoordinator:
         next_request_id (int): A counter for generating unique server-side request IDs.
     """
 
-    def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
+    def __init__(self, inference_coordinator_port: int, data_parallel_size: int, tokenizer):
         """
         Initializes the inference coordinator.
 
@@ -116,6 +116,7 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
         self.request_id_to_client_request_id = {}
 
         self.next_request_id = 0
+        self.tokenizer = tokenizer
 
     def get_next_data_parallel_rank(self):
         """
@@ -261,6 +262,7 @@ def start(self):
                 finished_request_records = deserialized_payload[1]
 
                 for finished_request_record in finished_request_records:
+                    self.detokenize(finished_request_record)
                     fid = finished_request_record["requests"][0]["request_id"]
                     client_identity = self.request_id_to_client_id[fid]
                     client_request_identity = self.request_id_to_client_request_id[fid]
@@ -280,9 +282,25 @@ def start(self):
             else:
                 raise UnknownHeaderError(header)
 
+    def detokenize(self, finished_request_record):
+        """
+        Detokenizes the generated tokens in the finished request record.
+
+        This method uses the coordinator's tokenizer to convert the list of
+        generated token IDs back into human-readable text.
+
+        Args:
+            finished_request_record (dict): The record containing the generated
+                tokens to be detokenized. It is modified in place.
+        """
+        for request in finished_request_record["requests"]:
+            if request["prompt"] is None:
+                request["prompt"] = self.tokenizer.detokenize(request["prompt_tokens"][1])
+            request["generated_text"] = self.tokenizer.detokenize(request["generated_tokens"])
+
     @classmethod
     def entrypoint(
-        cls, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int
+        cls, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int, tokenizer
     ):
         """
         Class method to instantiate and run the coordinator, for use in a separate process.
@@ -296,7 +314,7 @@ def entrypoint(
             inference_coordinator_port (int): The port to bind to.
             data_parallel_size (int): The number of expected TP-coordinators.
         """
-        coordinator = cls(inference_coordinator_port, data_parallel_size)
+        coordinator = cls(inference_coordinator_port, data_parallel_size, tokenizer)
         ready_event.set()
         try:
             coordinator.start()
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index c56f91bbbe9..c42246d5624 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -413,6 +413,7 @@ async def start_listening_to_data_parallel_coordinator(
                     coordinator_ready_event,
                     inference_coordinator_port,
                     get_pg_size(self.pg_collection.dp),
+                    self.controller.tokenizer,
                 ),
             )
             self.inference_coordinator_process.start()
@@ -1205,6 +1206,7 @@ async def async_bookkeep(
                 cuda_graph_request_count (int): The CUDA graph batch size matching this step.
         """
         # Increment finished_request_count.
+        range_push("bookkeeping")
         cuda_graph_request_count = None
 
         if step_result is not None:
@@ -1248,26 +1250,33 @@ async def async_bookkeep(
             finished_request_records.append(failed_entry.record)
             failed_entry.future.set_result(failed_entry.record)
         self.failed_request_ids.clear()
+        range_pop()
 
-        # Detokenize all finished requests (critical for InferenceClient, which
-        # doesn't necessarily have the tokenizer).
-        for record in finished_request_records:
-            for request in record.requests:
-                if request.prompt is None:
-                    request.prompt = self.controller.tokenizer.detokenize(
-                        request.prompt_tokens.tolist()
+        # Detokenize all finished requests if not using
+        # the coordinator. Otherwise, the coordinator will
+        # overlap detokenization with the engine.
+        if not self.use_coordinator:
+            range_push("detokenization")
+            for record in finished_request_records:
+                for request in record.requests:
+                    if request.prompt is None:
+                        request.prompt = self.controller.tokenizer.detokenize(
+                            request.prompt_tokens.tolist()
+                        )
+                    request.generated_text = self.controller.tokenizer.detokenize(
+                        request.generated_tokens
                     )
-                request.generated_text = self.controller.tokenizer.detokenize(
-                    request.generated_tokens
-                )
+            range_pop()
 
         # Handle necessary ZMQ DP coordinator communication.
         if self.use_coordinator and self.is_mp_coordinator and finished_request_records:
+            range_push("coordinator_communication")
             payload = msgpack.packb(
                 [Headers.ENGINE_REPLY.value, [r.serialize() for r in finished_request_records]],
                 use_bin_type=True,
             )
             self.socket_for_receiving_requests.send(payload)
+            range_pop()
 
         # Log KV cache utilization stats to W&B
         if context_state["kv_stats"] is not None:
@@ -1461,7 +1470,7 @@ def schedule_requests(self) -> int:
             int: The number of messages that were received and processed in this batch.
         """
 
-        torch.cuda.nvtx.range_push("drain_zmq_socket")
+        range_push("drain_zmq_socket")
         all_messages = []
         if self.is_mp_coordinator:
             while True:
@@ -1494,7 +1503,7 @@ def schedule_requests(self) -> int:
             else:
                 all_messages = []
 
-        torch.cuda.nvtx.range_pop()
+        range_pop()
         for message in all_messages:
             data = msgpack.unpackb(message, raw=False)
             header = Headers(data[0])
@@ -1507,7 +1516,9 @@ def schedule_requests(self) -> int:
             if header == Headers.SUBMIT_REQUEST:
                 request_id, prompt, sampling_params = data[1:]
                 sampling_params = SamplingParams.deserialize(sampling_params)
+                range_push("add_request")
                 self.add_request(request_id, prompt, sampling_params)
+                range_pop()
             elif header == Headers.PAUSE:
                 # Pause thyself.
                 self.received_pause = True
diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py
index 8bd0dd0aff4..6a7354220f9 100644
--- a/megatron/core/inference/inference_request.py
+++ b/megatron/core/inference/inference_request.py
@@ -1,7 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
-import io
 import time
 import warnings
 from dataclasses import asdict, dataclass, field
@@ -15,33 +14,34 @@
 from megatron.core.utils import experimental_api
 
 
-def serialize_tensor(tensor: torch.Tensor) -> bytes:
+def serialize_tensor(tensor: torch.Tensor) -> List:
     """Serialize tensor to bytes.
 
     Args:
         tensor (Tensor): Tensor.
 
     Returns:
-        (bytes) Byte representation of tensor.
+        (List) Tensor as a list
     """
-    buffer = io.BytesIO()
-    torch.save(tensor, buffer)
-    buffer.seek(0)
-    tensor_bytes = buffer.read()
-    return tensor_bytes
+    torch.cuda.nvtx.range_push("serialize_tensor")
 
+    # simply convert tensor into a list
+    tensor = tensor.cpu().tolist()
 
-def deserialize_tensor(tensor_bytes: bytes) -> torch.Tensor:
+    torch.cuda.nvtx.range_pop()
+    return tensor
+
+
+def deserialize_tensor(tensor_as_list: List) -> torch.Tensor:
     """Deserialize tensor from bytes.
 
     Args:
-        tensor_bytes (bytes): Byte representation of tensor.
+        tensor_as_list (List): List representation of tensor.
 
     Returns:
         (Tensor) Tensor.
     """
-    buffer = io.BytesIO(tensor_bytes)
-    tensor = torch.load(buffer)
+    tensor = torch.tensor(tensor_as_list)
     return tensor
 
 
@@ -99,17 +99,21 @@ def serialize(self) -> dict:
             (dict) A dictionary representation of the instance suitable for
                 serialization.
         """
-
         # Dataclass to dict.
-        obj = asdict(self)
+        # do not use asdict(self) - it has very high CPU overheads
+        # and if there are tensors, it will try to deepcopy them
+        obj = self.__dict__.copy()  # shallow dict copy
         obj["status"] = self.status.name if self.status else None
+        obj["sampling_params"] = self.sampling_params.serialize() if self.sampling_params else None
+        obj["inference_parameters"] = (
+            self.inference_parameters.serialize() if self.inference_parameters else None
+        )
 
         # Serialize tensors.
         obj = {
             k: (("tensor", serialize_tensor(v)) if isinstance(v, torch.Tensor) else v)
             for k, v in obj.items()
         }
-
         return obj
 
     @classmethod
@@ -125,14 +129,31 @@ def deserialize(cls, obj: dict) -> "InferenceRequest":
 
         # Initialize request.
         request = cls(**obj)
-        request.status = None if obj["status"] is None else Status[obj["status"]]
+        request._post_deserialize(obj)
+        return request
 
-        # Deserialize tensors.
+    def _post_deserialize(self, obj: dict):
+        """
+        This is called after the dataclass is initialized to handle any special
+        deserialization logic.
+        """
+        # Deserialize status.
+        self.status = None if obj["status"] is None else Status[obj["status"]]
+        self.sampling_params = (
+            None
+            if obj["sampling_params"] is None
+            else SamplingParams.deserialize(obj["sampling_params"])
+        )
+        self.inference_parameters = (
+            None
+            if obj["inference_parameters"] is None
+            else SamplingParams.deserialize(obj["inference_parameters"])
+        )
+
+        # Deserialize tensors and sampling params.
         for k, v in obj.items():
             if isinstance(v, list) and len(v) == 2 and v[0] == "tensor":
-                setattr(request, k, deserialize_tensor(v[1]))
-
-        return request
+                setattr(self, k, deserialize_tensor(v[1]))
 
 
 class DynamicInferenceEventType(Enum):
@@ -197,7 +218,10 @@ def serialize(self) -> dict:
         """
 
         # Dataclass to dict.
-        obj = asdict(self)
+        torch.cuda.nvtx.range_push("DynamicInferenceEvent.serialize")
+        # do not use asdict(self) - it has very high CPU overheads
+        # and if there are tensors, it will try to deepcopy them
+        obj = self.__dict__.copy()
         obj["type"] = self.type.name
 
         # Serialize payload.
@@ -205,7 +229,7 @@ def serialize(self) -> dict:
             from .contexts.dynamic_context import ContextErrorFactory  # avoid circular import.
 
             obj["payload"] = ContextErrorFactory.serialize(self.payload)
-
+        torch.cuda.nvtx.range_pop()
         return obj
 
     @classmethod
@@ -247,7 +271,7 @@ class DynamicInferenceRequest(InferenceRequest):
     # remaining prompt tokens are used for chunked prefill
     remaining_prompt_tokens: Optional[torch.Tensor] = None
     latency: Optional[float] = None
-    finished_chunk_token_count = 0
+    finished_chunk_token_count: int = 0
     stop_word_ids: Optional[List[List[int]]] = None  # Tokenized stop words (populated internally)
 
     def __post_init__(self):
@@ -275,30 +299,22 @@ def __str__(self):
             )
         )
 
-    def serialize(self):
+    def serialize(self) -> dict:
         """Converts the instance into a serializable dictionary.
 
         Returns:
             (dict) A dictionary representation of the instance suitable for
                 serialization.
         """
+        torch.cuda.nvtx.range_push("DynamicInferenceRequest.serialize")
         obj = super().serialize()
         obj["events"] = [e.serialize() for e in self.events]
+        torch.cuda.nvtx.range_pop()
         return obj
 
-    @classmethod
-    def deserialize(cls, obj: dict) -> "DynamicInferenceRequest":
-        """Deserialize request.
-
-        Args:
-            obj (dict): Serialized request data.
-
-        Returns:
-            (DynamicInferenceRequest) Deserialized request.
-        """
-        request = super().deserialize(obj)
-        request.events = [DynamicInferenceEvent.deserialize(e) for e in obj["events"]]
-        return request
+    def _post_deserialize(self, obj):
+        super()._post_deserialize(obj)
+        self.events = [DynamicInferenceEvent.deserialize(e) for e in obj["events"]]
 
     @property
     def tracked_metadata(self) -> List[Any]:
@@ -517,8 +533,10 @@ def serialize(self) -> dict:
             (dict) A dictionary representation of the instance suitable for
                 serialization.
         """
-        obj = asdict(self)
-        obj["requests"] = [r.serialize() for r in self.requests]
+        torch.cuda.nvtx.range_push("DynamicInferenceRequestRecord.serialize")
+        obj = self.__dict__.copy()  # shallow dict copy
+        obj["requests"] = [r.serialize() for r in obj["requests"]]
+        torch.cuda.nvtx.range_pop()
         return obj
 
     @classmethod

From 369e0eba792526d346390970ba53f2f6c40a2e69 Mon Sep 17 00:00:00 2001
From: Siddharth Singh <136645615+sidsingh-nvidia@users.noreply.github.com>
Date: Sun, 25 Jan 2026 15:07:18 -0600
Subject: [PATCH 17/79] Fix broken functional tests in #2920 (#3071)

---
 .../inference/gpt/gpt_dynamic_inference_with_coordinator.py     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
index f354b122a7e..18191fd38af 100644
--- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
+++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
@@ -152,7 +152,7 @@ async def main(
                     "generated_tokens": req.generated_tokens,
                     "latency": req.latency,  # InferenceClient populates this field in the returned future.
                 }
-                if req.sampling_params["return_log_probs"]:
+                if req.sampling_params.return_log_probs:
                     result_dict["logprobs"] = req.prompt_log_probs + req.generated_log_probs
                 throughput = len(req.generated_tokens) / req.latency
                 throughputs.append(throughput)

From b2e93909f09ce5a84d42beb844e905eb7edd7818 Mon Sep 17 00:00:00 2001
From: Deyu Fu <deyuf@nvidia.com>
Date: Mon, 26 Jan 2026 13:41:39 +0800
Subject: [PATCH 18/79] fix ep weight gradnorm/num_zero calculation error for
 muon (#3024)

---
 megatron/core/optimizer/muon.py               | 105 +-
 .../golden_values_dev_dgx_h100.json           | 968 +++++++++---------
 .../golden_values_dev_dgx_h100_2nd.json       | 498 ++++-----
 .../model_config.yaml                         |   1 +
 tests/test_utils/recipes/moe.yaml             |   2 +-
 .../test_layer_wise_optimizer.py              |  11 +-
 tests/unit_tests/dist_checkpointing/utils.py  |  40 +-
 7 files changed, 815 insertions(+), 810 deletions(-)

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

From 497d42de9862a446a08c4c2649f68d4aa5a7b7a7 Mon Sep 17 00:00:00 2001
From: Maanu Grover <109391026+maanug-nv@users.noreply.github.com>
Date: Sun, 25 Jan 2026 21:45:40 -0800
Subject: [PATCH 19/79] [training migration] Add LoggerConfig dataclass (#2414)

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/arguments.py       |  90 ++-----------------
 megatron/training/training_config.py | 127 +++++++++++++++++++++++++++
 2 files changed, 132 insertions(+), 85 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 3c2a0e52f71..ecf6b234d43 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2025,83 +2025,11 @@ def _add_config_logger_args(parser):
 
 
 def _add_logging_args(parser):
-    group = parser.add_argument_group(title='logging')
-
-    group.add_argument('--log-params-norm', action='store_true',
-                       help='If set, calculate and log parameters norm.')
-    group.add_argument('--log-num-zeros-in-grad', action='store_true',
-                       help='If set, calculate and log the number of zeros in gradient.')
-    group.add_argument('--log-throughput', action='store_true',
-                       help='If set, calculate and log throughput per GPU.')
-    group.add_argument('--log-progress', action='store_true',
-                       help='If set, log progress (in terms of number of processed tokens and '
-                       'number of floating-point operations) to progress.txt file in checkpoint '
-                       'directory.')
-    group.add_argument('--timing-log-level', type=int,
-                       default=0, choices=range(0,3),
-                       help='Granularity level to measure and report timing. '
-                       '   0: report only iteration time and make sure timing '
-                       '      does not introduce extra overhead.'
-                       '   1: report timing for operations that are executed '
-                       '      very limited times (basically once) during '
-                       '      each iteration (such as gradient all-reduce) '
-                       '   2: report timing for operations that migh be '
-                       '      executed numerous times during each iteration. '
-                       'Note that setting the level to 1 or 2 might '
-                       'cause increase in iteration time.')
-    group.add_argument('--log-energy', action='store_true',
-                       help='If set, log energy consumption (in Joules)')
-    group.add_argument('--no-barrier-with-level-1-timing', action='store_false',
-                       help='If not set, use barrier with level 1 time '
-                       'measurements. Note that this is up to the user '
-                       'to make sure calling barrier with their timers '
-                       'will not result in hangs. This can happen if for '
-                       'example the user adds a level 1 timer that is not '
-                       'called by all ranks.',
-                       dest='barrier_with_L1_time')
-    group.add_argument('--timing-log-option', type=str, default='minmax',
-                       choices=['max', 'minmax', 'all'],
-                       help='Options for logging timing:'
-                       '  max: report the max timing across all ranks'
-                       '  minmax: report min and max timings across all ranks'
-                       '  all: report timings of all ranks.')
-    group.add_argument('--tensorboard-log-interval', type=int, default=1,
-                       help='Report to tensorboard interval.')
-    group.add_argument('--tensorboard-queue-size', type=int, default=1000,
-                       help='Size of the tensorboard queue for pending events '
-                       'and summaries before one of the "add" calls forces a '
-                       'flush to disk.')
-    group.add_argument('--log-timers-to-tensorboard', action='store_true',
-                       help='If set, write timers to tensorboard.')
-    group.add_argument('--no-log-loss-scale-to-tensorboard',
-                       action='store_false',
-                       help='Disable loss-scale logging to tensorboard.',
-                       dest='log_loss_scale_to_tensorboard')
-    group.add_argument('--log-validation-ppl-to-tensorboard',
-                       action='store_true',
-                       help='If set, write validation perplexity to '
-                       'tensorboard.')
-    group.add_argument('--log-memory-to-tensorboard',
-                       action='store_true',
-                       help='Enable memory logging to tensorboard.')
-    group.add_argument('--log-world-size-to-tensorboard',
-                       action='store_true',
-                       help='Enable world size logging to tensorboard.')
-    group.add_argument('--log-max-attention-logit', action='store_true',
-                       help='Enable max attention logit logging to tensorboard.')
-    group.add_argument('--wandb-project', type=str, default='',
-                       help='The wandb project name. Ignore wandb by default.')
-    group.add_argument('--wandb-entity', type=str, default='',
-                       help='The wandb entity name. It is useful when '
-                       'there are multiple sub-projects in a project. '
-                       'https://community.wandb.ai/t/how-do-i-decide-which-account-private-or-team-to-upload-the-run-to/5704 '
-                       'Ignore wandb by default.')
-    group.add_argument('--wandb-exp-name', type=str, default='',
-                       help='The wandb experiment name.')
-    group.add_argument('--wandb-save-dir', type=str, default='',
-                       help='Path to save the wandb results locally.')
-    group.add_argument('--logging-level', type=int, default=None,
-                       help='Set default logging level')
+    from megatron.training.training_config import LoggerConfig
+
+    log_factory = ArgumentGroupFactory(LoggerConfig, exclude = ["log_throughput_to_tensorboard", "throughput_window_size", "memory_keys", "log_l2_norm_grad_to_tensorboard", "log_runtime_to_tensorboard", "runtime_time_unit", "filter_warnings", "modules_to_filter", "set_level_for_all_loggers", "save_config_filepath"])
+    group = log_factory.build_group(parser, title="logging")
+
     return parser
 
 
@@ -2400,14 +2328,6 @@ def _add_training_args(parser):
     group.add_argument('--checkpoint-activations', action='store_true',
                        help='Checkpoint activation to allow for training '
                        'with larger models, sequences, and batch sizes.')
-    group.add_argument('--log-interval', type=int, default=100,
-                       help='Report loss and timing interval.')
-    group.add_argument('--log-memory-interval', type=int, default=None,
-                       help='Report memory interval.')
-    group.add_argument('--log-device-memory-used', action='store_true',
-                       help='Log device memory used (as reported by nvidia-smi).')
-    group.add_argument('--tensorboard-dir', type=str, default=None,
-                       help='Write TensorBoard logs to this directory.')
     group.add_argument('--no-masked-softmax-fusion',
                        action='store_false',
                        help='Disable fusion of query_key_value scaling, '
diff --git a/megatron/training/training_config.py b/megatron/training/training_config.py
index d91972cf3c6..617c5cf5dfa 100644
--- a/megatron/training/training_config.py
+++ b/megatron/training/training_config.py
@@ -193,3 +193,130 @@ class SchedulerConfig:
 
     wsd_decay_steps: int | None = field(init=False, default=None)
     """Number of samples to decay WSD weight decay. Calculated at runtime."""
+
+
+@dataclass(kw_only=True)
+class LoggerConfig:
+    """Configuration settings for logging, including TensorBoard and WandB."""
+
+    log_interval: int = 100
+    """Report loss and timing interval."""
+
+    log_params_norm: bool = False
+    """If set, calculate and log parameters norm."""
+
+    log_throughput: bool = False
+    """If set, calculate and log throughput per GPU."""
+
+    log_throughput_to_tensorboard: bool = False
+    """Enable throughput logging to tensorboard."""
+
+    throughput_window_size: int = 100
+    """Number of batches to use for a rolling average of throughput."""
+
+    log_progress: bool = False
+    """If set, log progress (in terms of number of processed tokens and number of floating-point operations)
+    to progress.txt file in checkpoint directory.
+    """
+
+    timing_log_level: Literal[0, 1, 2] = 0
+    """Granularity level to measure and report timing.
+    0: report only iteration time and make sure timing does not introduce extra overhead.
+    1: report timing for operations that are executed very limited times (basically once) during each iteration
+        (such as gradient all-reduce)
+    2: report timing for operations that migh be executed numerous times during each iteration.
+    Note that setting the level to 1 or 2 might cause increase in iteration time.
+    """
+
+    timing_log_option: Literal["max", "minmax", "all"] = "minmax"
+    """Options for logging timing:
+    max: report the max timing across all ranks
+    minmax: report min and max timings across all ranks
+    all: report timings of all ranks.
+    """
+
+    tensorboard_dir: str | None = None
+    """Write TensorBoard logs to this directory."""
+
+    tensorboard_log_interval: int = 1
+    """Report to tensorboard interval."""
+
+    tensorboard_queue_size: int = 1000
+    """Size of the tensorboard queue for pending events and summaries
+    before one of the 'add' calls forces a flush to disk.
+    """
+
+    log_timers_to_tensorboard: bool = False
+    """If set, write timers to tensorboard."""
+
+    log_loss_scale_to_tensorboard: bool = True
+    """Disable loss-scale logging to tensorboard."""
+
+    log_validation_ppl_to_tensorboard: bool = False
+    """If set, write validation perplexity to tensorboard."""
+
+    log_memory_to_tensorboard: bool = False
+    """Enable memory logging to tensorboard."""
+
+    memory_keys: dict[str, str] | None = None
+    """Names of memory statistics to log from `torch.cuda.memory_stats()`"""
+
+    log_memory_interval: int | None = None
+    """Report memory interval."""
+
+    log_device_memory_used: bool = False
+    """Log device memory used (as reported by nvidia-smi)."""
+
+    log_l2_norm_grad_to_tensorboard: bool = False
+    """Enable gradients logging to tensorboard."""
+
+    log_num_zeros_in_grad: bool = False
+    """If set, calculate and log the number of zeros in gradient."""
+
+    log_max_attention_logit: bool = False
+    """Enable max attention logit logging to tensorboard."""
+
+    log_runtime_to_tensorboard: bool = False
+    """Enable runtime metrics logging to tensorboard."""
+
+    runtime_time_unit: str = "hours"
+    """Time unit to use for time logging. """
+
+    barrier_with_L1_time: bool = field(default=True, metadata={"argparse_meta": {"arg_names": ["--no-barrier-with-level-1-timing"]}})
+    """If not disabled, use barrier with level 1 time measurements. Note that this is up to the user to
+    make sure calling barrier with their timers will not result in hangs. This can happen if for
+    example the user adds a level 1 timer that is not called by all ranks.
+    """
+
+    log_world_size_to_tensorboard: bool = False
+    """Enable world size logging to tensorboard."""
+
+    wandb_project: str | None = None
+    """The wandb project name. Ignore wandb by default."""
+
+    wandb_exp_name: str | None = None
+    """The wandb experiment name."""
+
+    wandb_save_dir: str | None = None
+    """Path to save the wandb results locally."""
+
+    wandb_entity: str | None = None
+    """The wandb entity name. It is useful when there are multiple sub-projects in a project."""
+
+    logging_level: int | None = None
+    """Set default logging level"""
+
+    filter_warnings: bool = True
+    """Filter out warning messages"""
+
+    modules_to_filter: list[str] | None = None
+    """List of modules to filter out from the logs"""
+
+    set_level_for_all_loggers: bool = False
+    """Set the logging level for all loggers. If False, only level for NeMo loggers will be set."""
+
+    log_energy: bool = False
+    """If set, log energy consumption (in Joules)."""
+
+    save_config_filepath: str | None = None
+    """If set, save the task configuration (ConfigContainer) to this file."""

From 06d0f46c03ab1175290f238948e975e2366c4698 Mon Sep 17 00:00:00 2001
From: Hexin Wang <160587990+hexinw-nvidia@users.noreply.github.com>
Date: Mon, 26 Jan 2026 08:19:41 -0800
Subject: [PATCH 20/79] Added --ft-num-warmup-iters option. (#3052)

---
 megatron/training/arguments.py      | 5 +++++
 megatron/training/ft_integration.py | 7 +++++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index ecf6b234d43..47d0001d6e5 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2013,6 +2013,11 @@ def _add_ft_package_args(parser):
     group.add_argument('--calc-ft-timeouts', action='store_true',
                        help='If set, FT package will try to automatically compute the timeouts. '
                        'Note: This feature is for Nvidia internal use only.')
+    group.add_argument('--ft-num-warmup-iters', type=int, default=5,
+                       help='Number of warmup iterations before monitoring step section and '
+                       'out-of-section timeouts. The first N iterations are excluded from '
+                       'timeout monitoring as they can be significantly slower than steady-state. '
+                       'Default: 5. Note: This feature is for Nvidia internal use only.')
     return parser
 
 
diff --git a/megatron/training/ft_integration.py b/megatron/training/ft_integration.py
index 670cf492602..f3532e75639 100644
--- a/megatron/training/ft_integration.py
+++ b/megatron/training/ft_integration.py
@@ -60,7 +60,7 @@
 _seen_tr_iters_cnt = 0
 _curr_eval_iter_idx = 0
 
-_NUM_WARMUP_ITERS = 1
+_NUM_WARMUP_ITERS = 1  # Will be set by --ft-num-warmup-iters (default: 5)
 _MIN_ITERS_FOR_STEP_TIMEOUT_UPDATE = 16
 
 
@@ -105,7 +105,10 @@ def setup() -> None:
     global _is_calculating_timeouts
     _is_calculating_timeouts = args.calc_ft_timeouts
 
-    cli.init_workload_monitoring()
+    global _NUM_WARMUP_ITERS
+    _NUM_WARMUP_ITERS = args.ft_num_warmup_iters
+
+    cli.init_workload_monitoring(num_warmup_iters=_NUM_WARMUP_ITERS)
     _load_state_if_exists()
     if os.environ.get("RANK") == "0":
         print(f"FT: initialized. Timeouts={cli.section_timeouts}", flush=True)

From 642fdd9010320cb2438575a0c9eecd49cfdf5eaf Mon Sep 17 00:00:00 2001
From: Jimmy Zhang <133159885+jiemingz@users.noreply.github.com>
Date: Mon, 26 Jan 2026 11:42:16 -0500
Subject: [PATCH 21/79] Reapply "Various CUDA graph improvements on capture
 time, replay time, memory footprint (#2572)" (#3056)

Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
---
 .../core/models/gpt/fine_grained_callables.py |    4 +-
 .../core/models/mamba/mamba_layer_specs.py    |    9 +-
 megatron/core/ssm/mamba_layer.py              |    8 +
 megatron/core/tensor_parallel/random.py       |   51 +-
 megatron/core/transformer/cuda_graphs.py      | 1305 ++++++++++-------
 megatron/core/transformer/module.py           |   11 +-
 megatron/core/transformer/moe/moe_layer.py    |   63 +-
 megatron/core/transformer/moe/moe_utils.py    |    4 +
 .../core/transformer/moe/token_dispatcher.py  |    6 +-
 .../core/transformer/transformer_config.py    |   87 +-
 .../core/transformer/transformer_layer.py     |  226 ++-
 megatron/training/arguments.py                |   15 +-
 megatron/training/training.py                 |    4 +-
 .../golden_values_dev_dgx_h100.json           |   16 +-
 .../golden_values_dev_dgx_h100.json           |   16 +-
 15 files changed, 1166 insertions(+), 659 deletions(-)

diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 9234e142c6c..bbeee561110 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -545,7 +545,9 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor):
         """
         residual = node.layer_state.residual
         shared_expert_output = getattr(node.layer_state, 'shared_expert_output', None)
-        output = layer.mlp.combine(output, shared_expert_output)
+        output = layer.mlp.combine(output)
+        output = layer.mlp.postprocess(output, shared_expert_output)
+
         mlp_output_with_bias = (output, None)
         if hasattr(layer, 'cuda_graphs') and layer.cuda_graphs:
             layer.mlp.cudagraph_tensor_store.clear()
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index f83275ed9c6..b87124bab1d 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -20,7 +20,11 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
-from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+from megatron.core.transformer.transformer_layer import (
+    MoETransformerLayer,
+    TransformerLayer,
+    TransformerLayerSubmodules,
+)
 
 moe = get_moe_module_spec(
     use_te=True,
@@ -78,8 +82,7 @@
             ),
         ),
         moe_layer=ModuleSpec(
-            # TODO (rwaleffe): change this to be an "MoELayer" to work with CudaGraphs?
-            module=TransformerLayer,
+            module=MoETransformerLayer,
             submodules=TransformerLayerSubmodules(
                 pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
             ),
diff --git a/megatron/core/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py
index 48ea84566d5..ac6e8b5bf40 100644
--- a/megatron/core/ssm/mamba_layer.py
+++ b/megatron/core/ssm/mamba_layer.py
@@ -16,6 +16,7 @@
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.module import GraphableMegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
@@ -85,6 +86,13 @@ def __init__(
         self.mamba_bda = build_module(submodules.mamba_bda)
         self.bias_dropout_add_exec_handler = torch.enable_grad
 
+    def create_mcore_cudagraph_manager(self, config):
+        """Register the mamba layer for cudagraphs."""
+        from megatron.core.transformer.cuda_graphs import CudaGraphManager
+
+        if not self.config.cuda_graph_scope or CudaGraphScope.mamba in self.config.cuda_graph_scope:
+            self.cudagraph_manager = CudaGraphManager(config)
+
     def mamba_state_shapes_per_request(self) -> Tuple[Tuple[int], Tuple[int]]:
         """Returns the Mamba conv and ssm states shapes per request."""
         return self.mixer.mamba_state_shapes_per_request()
diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py
index 7f3a5ab0365..bf00717ab6c 100644
--- a/megatron/core/tensor_parallel/random.py
+++ b/megatron/core/tensor_parallel/random.py
@@ -472,6 +472,27 @@ def _fork_rng():
         _set_all_rng_states(*current_states)
 
 
+# Global flag that's toggled whenever inside a checkpointing context
+IS_CHECKPOINTING = False
+
+
+def _set_checkpointing():
+    """Set state to checkpointing enabled."""
+    global IS_CHECKPOINTING
+    IS_CHECKPOINTING = True
+
+
+def _unset_checkpointing():
+    """Unset state to checkpointing enabled."""
+    global IS_CHECKPOINTING
+    IS_CHECKPOINTING = False
+
+
+def is_checkpointing():
+    """Check if currently in a checkpoint context."""
+    return IS_CHECKPOINTING
+
+
 class CheckpointFunction(torch.autograd.Function):
     """Checkpoint Function
 
@@ -484,6 +505,8 @@ class CheckpointFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, run_function, distribute_saved_activations, *args):
         """Forward pass."""
+        _set_checkpointing()
+
         ctx.run_function = run_function
         ctx.distribute_saved_activations = distribute_saved_activations
 
@@ -504,6 +527,7 @@ def forward(ctx, run_function, distribute_saved_activations, *args):
         # Store everything.
         ctx.save_for_backward(*args)
 
+        _unset_checkpointing()
         return outputs
 
     # pylint: disable=missing-function-docstring
@@ -515,6 +539,8 @@ def backward(ctx, *args):
                 "Checkpointing is not compatible with .grad(), "
                 "please use .backward() if possible"
             )
+        _set_checkpointing()
+
         inputs = ctx.saved_tensors
         if ctx.distribute_saved_activations:
             safely_set_viewless_tensor_data(
@@ -539,6 +565,8 @@ def backward(ctx, *args):
         )
         torch.autograd.backward(outputs, args)
         grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)
+
+        _unset_checkpointing()
         return (None, None) + grads
 
 
@@ -615,6 +643,14 @@ def __init__(self, fp8=False):
 
     def checkpoint(self, run_function, *args):
         """Checkpoint function."""
+
+        # If in cuda graph warmup, disable checkpointing, as 'discard_output_and_register_recompute'
+        # may be called in a separate graph warmup.
+        from megatron.core.transformer.cuda_graphs import is_graph_warmup
+
+        if is_graph_warmup():
+            return run_function(*args)
+
         self.run_function = run_function
 
         self.rng_states = _get_all_rng_states()
@@ -628,11 +664,14 @@ def checkpoint(self, run_function, *args):
     def _recompute(self, _):
         """Used as a hook to recompute the output."""
 
-        if self.ctx is None:
-            # The recomputation has been triggered already. Just return.
+        from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup
+
+        # The recomputation has been triggered already. Just return.
+        # Handle cudagraphs, do nothing if currently in graph warmup
+        if self.ctx is None or is_graph_warmup():
             return
 
-        if not torch.autograd._is_checkpoint_valid():
+        if not torch.autograd._is_checkpoint_valid() and not is_graph_capturing():
             raise RuntimeError(
                 "Checkpointing is not compatible with .grad(), "
                 "please use .backward() if possible"
@@ -691,6 +730,12 @@ def discard_output_and_register_recompute(self, hook_tensor):
         in the forward pass and the gradient of the hook_tensor is computed before the recomputed
         tensors are used.
         """
+
+        from megatron.core.transformer.cuda_graphs import is_graph_warmup
+
+        if is_graph_warmup():
+            return
+
         # use resize to release the output tensor memory and still keep the metadata in the tensors.
         # the metadata is still needed for backward
         for output in self.outputs:
diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py
index 1e3e3edc558..3a0632d4ee7 100644
--- a/megatron/core/transformer/cuda_graphs.py
+++ b/megatron/core/transformer/cuda_graphs.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import dataclasses
 import gc
 import inspect
 import logging
@@ -8,15 +9,16 @@
 import time
 from collections import defaultdict
 from contextlib import nullcontext
-from dataclasses import fields, is_dataclass
+from copy import deepcopy
+from dataclasses import dataclass, is_dataclass
 from enum import Enum
 from functools import partial
-from itertools import zip_longest
+from itertools import chain, zip_longest
 from math import ceil
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 import torch
-from torch.utils._pytree import tree_flatten
+from torch.utils._pytree import tree_map as tree_map_pyt
 
 from megatron.core import parallel_state
 from megatron.core.num_microbatches_calculator import get_num_microbatches
@@ -24,9 +26,9 @@
     CudaRNGStatesTracker,
     get_all_rng_states,
     get_cuda_rng_tracker,
+    is_checkpointing,
 )
 from megatron.core.transformer.enums import CudaGraphScope
-from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import (
@@ -39,6 +41,7 @@
 
 try:
     import transformer_engine as te  # pylint: disable=unused-import
+    from transformer_engine.pytorch.distributed import is_fp8_activation_recompute_enabled
     from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
     from transformer_engine.pytorch.graph import (
         make_graphed_callables,
@@ -48,6 +51,7 @@
     from transformer_engine.pytorch.graph import set_capture_end as te_set_capture_end
     from transformer_engine.pytorch.graph import set_capture_start as te_set_capture_start
     from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
+    from transformer_engine.pytorch.utils import make_weak_ref
 
     HAVE_TE_GRAPHS = True
 except:
@@ -61,7 +65,7 @@
     HAVE_TQDM = False
 
 _IS_GRAPH_CAPTURING = False
-
+_IS_GRAPH_WARMUP = False
 logger = logging.getLogger(__name__)
 
 # Freeze GC during capture.
@@ -79,7 +83,6 @@
 
 def is_graph_capturing():
     """Query if currently capturing."""
-    global _IS_GRAPH_CAPTURING
     return _IS_GRAPH_CAPTURING
 
 
@@ -95,6 +98,39 @@ def _set_capture_end():
     _IS_GRAPH_CAPTURING = False
 
 
+def is_graph_warmup():
+    """Query if currently warming up for graph capture."""
+    return _IS_GRAPH_WARMUP
+
+
+def _set_warmup_start():
+    """Set graph warmup has started."""
+    global _IS_GRAPH_WARMUP
+    _IS_GRAPH_WARMUP = True
+
+
+def _set_warmup_end():
+    """Set graph warmup has ended."""
+    global _IS_GRAPH_WARMUP
+
+
+@dataclass
+class CudagraphBufferMetadata:
+    """
+    Metadata saved to tensors during cudagraph capture. This data will be used to determine
+    during graph captue when a cudagraph can reuse a buffer or directly write its output into
+    a subsequent's graph's input.
+    """
+
+    is_cudagraph_input: bool = False
+    is_cudagraph_output: bool = False
+    input_use_count: int = 0
+    cudagraph_reuse_ref_count: int = 0
+    capture_reuse_count: int = 0
+    fwd_cudagraph_buffer: torch.Tensor = None
+    bwd_cudagraph_buffer: torch.Tensor = None
+
+
 class ArgMetadata:
     """Arg meta."""
 
@@ -105,9 +141,83 @@ def __init__(self, arg):
             self.dtype = arg.dtype
             self.device = arg.device
             self.value = arg.data_ptr()
+            self.requires_grad = arg.requires_grad
+            if hasattr(arg, "cg_buffer_metadata"):
+                # Its important this is a reference copy
+                self.cg_buffer_metadata = arg.cg_buffer_metadata
         else:
             self.value = arg
 
+    def zeros_like(self):
+        """Reconstruct a tensor with the properties as the meta arg."""
+
+        assert self.type == torch.Tensor
+        return torch.zeros(
+            *self.shape, dtype=self.dtype, device=self.device, requires_grad=self.requires_grad
+        )
+
+
+class TensorReusePool:
+    """
+    A pool-like list of tensors that can be reused as input and output buffers during graph capture.
+    Also maintains strong references to all tensors created by this pool, so that they will never be
+    freed by the memory allocator.
+    """
+
+    """Record strong references to buffers created by the pool so they cannot be deallocated between
+    graph captures."""
+    tensor_strong_refs: list = []
+
+    """Record the data_ptrs of buffers created by the pool to check when a tensor came was 
+    allocated from this pool. """
+    tensor_strong_refs_dataptrs: set = set()
+
+    """Buffers that have been returned to the pool and are available for reuse. """
+    pool: list[torch.Tensor] = []
+
+    def insert(self, tensor: torch.Tensor):
+        """Return a tensor to the pool reuse."""
+        assert self.owns(tensor)
+        self.pool.append(tensor)
+
+    def owns(self, tensor: torch.Tensor):
+        """Check if a tensor was created from this pool."""
+        return tensor.data_ptr() in self.tensor_strong_refs_dataptrs
+
+    def get(self, meta: ArgMetadata):
+        """Try to get a buffer from the pool. If a matching tensor is already in the pool, its
+        assumed to be available and returned. Otherwise, allocate a new buffer."""
+
+        assert isinstance(meta, ArgMetadata)
+        # Find first matching buffer in pool
+        for i, buf in enumerate(self.pool):
+            if buf.shape == meta.shape and buf.dtype == meta.dtype and buf.device == meta.device:
+                return self.pool.pop(i)
+
+        out = meta.zeros_like()
+        self.tensor_strong_refs.append(out)
+        self.tensor_strong_refs_dataptrs.add(out.data_ptr())
+        return out
+
+
+def tree_map(func, tree):
+    """
+    Wrapper around pytorch's tree_map, but also recurses into dataclasses.
+    """
+
+    def wrapper(arg):
+        # If it's a dataclass, map over its fields
+        if is_dataclass(arg) and not isinstance(arg, type):
+            changes = {
+                f.name: tree_map_pyt(func, getattr(arg, f.name)) for f in dataclasses.fields(arg)
+            }
+            return dataclasses.replace(arg, **changes)
+
+        # Otherwise, apply the user function
+        return func(arg)
+
+    return tree_map_pyt(wrapper, tree)
+
 
 def _check_supported_type(meta):
     """Check if arg meta is a supported type for cudagraph input/outputs."""
@@ -125,35 +235,16 @@ def _check_supported_type(meta):
         int,
         str,
         float,
+        dataclass,
         StaticInferenceContext,
         DynamicInferenceContext,
+        ArgMetadata,
     }
     assert meta.type in _SUPPORTED_TYPES or is_dataclass(
         meta.value
     ), f"Cudagraphs recieved an arg of type {meta.type} which is not supported."
 
 
-def _determine_if_transformer_decoder_layer(base_module):
-    """Determine if the given module is a transformer decoder layer."""
-    # import modules here to avoid a circular import
-    from megatron.core.ssm.mamba_layer import MambaLayer
-    from megatron.core.transformer.transformer_layer import BaseTransformerLayer, TransformerLayer
-
-    is_potential_decoder_layer = isinstance(
-        base_module, (TransformerLayer, BaseTransformerLayer, MambaLayer)
-    )
-    if not is_potential_decoder_layer:
-        return False
-    if isinstance(base_module, TransformerLayer) and not isinstance(
-        base_module.cross_attention, IdentityOp
-    ):
-        # If the layer has a cross attention, it is not a decoder layer
-        return False
-    else:
-        # Otherwise it is a decoder layer
-        return True
-
-
 def _determine_if_first_last_layer_of_this_vp_chunk(base_module):
     """Determine if the given module is the first/last layer of the PP+VPP chunk it belongs to.
     Returns a tuple of two booleans indicating if the module is the first/last layer of the chunk.
@@ -163,6 +254,9 @@ def _determine_if_first_last_layer_of_this_vp_chunk(base_module):
     from megatron.core.transformer.transformer_block import get_num_layers_to_build
     from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
 
+    if not hasattr(base_module, "layer_number"):
+        return True, True
+
     # find all first/last layers of this PP stage
     first_layer_numbers = []
     last_layer_numbers = []
@@ -218,6 +312,10 @@ def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Gen
     return gen
 
 
+fwd_buffer_reuse_ref_count = 0
+bwd_buffer_reuse_ref_count = 0
+
+
 class _CudagraphGlobalRecord:
     """A global datastructure that records of the ordering of all _CudaGraphRunner's
     first fwd or bwd passes. 'create_cudagraphs' will use this to create
@@ -229,13 +327,16 @@ class _CudagraphGlobalRecord:
 
     """A record of fwd and bwd graph creation, populated with 'record_fwd_graph' and
     'record_bwd_graph."""
-    cudagraph_record = []
-    cudagraph_inference_record = []
+    cudagraph_record: list[tuple] = []
+    cudagraph_inference_record: list[tuple] = []
+
+    """A pool-like data structure to reuse input and output buffers across cudagraph."""
+    tensor_reuse_pool = TensorReusePool()
 
     @classmethod
-    def record_fwd_graph(cls, runner, args, kwargs):
+    def record_fwd_graph(cls, runner, args, kwargs, out):
         """Record a fwd graph to 'cudagraph_record"""
-        cls.cudagraph_record.append((runner, "fwd", args, kwargs))
+        cls.cudagraph_record.append((runner, "fwd", args, kwargs, out))
 
     @classmethod
     def record_bwd_graph(cls, runner):
@@ -246,7 +347,6 @@ def record_bwd_graph(cls, runner):
     def create_cudagraphs(cls):
         """Iterate through 'cudagraph_record' creating graphs in the order in which
         they were recorded."""
-
         # Cudagraphs have already been created, check that no cudagraphed modules ran in eager mode
         if cls.cudagraph_created:
             assert len(cls.cudagraph_record) == 0, (
@@ -260,8 +360,6 @@ def create_cudagraphs(cls):
             return
 
         # Otherwise, create all the recorded cudagraphs.
-        logging.getLogger(__name__).info(f"Creating {len(cls.cudagraph_record)} CUDA graphs")
-
         has_te_modules = False
         if HAVE_TE_GRAPHS:
             for g in cls.cudagraph_record:
@@ -270,21 +368,27 @@ def create_cudagraphs(cls):
                     [isinstance(m, TransformerEngineBaseModule) for m in base_module.modules()]
                 )
 
-        # If graphing only transformer layers with self attention, then apply the following
-        # transformer layer specific optimizations that reduce memory usage and tensor copies:
-        # These eventually will become unneccessary with:
-        # https://github.com/pytorch/pytorch/pull/137318
-        # 1. Some inputs to TransformerLayer (e.g. rotary_emb) are the same over all layers
-        #    and only need to be set once.
-        # 2. Because the next layer consumes the previous layer's hidden states, all fwd
-        #    cudagraphs can alternate reusing the same hidden_state input, output buffer.
-        #    Similarly, bwd graphs can alternate the same output, input grad buffers.
-        optimize_transformer_layer_graph_buffers = all(
-            [g[0].reuse_input_output_buffer for g in cls.cudagraph_record]
-        )
-        if optimize_transformer_layer_graph_buffers:
-            prev_fwd_hidden_state_output = None
-            prev_bwd_hidden_state_inputgrad = None
+        progress_bar = enumerate(cls.cudagraph_record)
+        time_start = time.time()
+        mem_stats_start = torch.cuda.memory_stats()
+
+        if torch.distributed.get_rank() == 0:
+            if HAVE_TQDM:
+                progress_bar = tqdm(
+                    progress_bar, "create cuda graphs", total=len(cls.cudagraph_record)
+                )
+
+            logger.info(f"Creating {len(cls.cudagraph_record)} CUDA graphs")
+            if not HAVE_TE_GRAPHS:
+                logger.warning(
+                    "Transformer Engine was not detected while capturing training cudagraphs."
+                    "As a result cudagraph memory overhead may significantly increase as "
+                    "Transformer Engine's weak reference feature is used on cudagraph input and "
+                    "output buffers. This allows the memory of input and output buffers to be "
+                    " reclaimed across graphs while remaining valid buffers for when the graph "
+                    "is replayed. For more information see: "
+                    "https://github.com/NVIDIA/TransformerEngine/blob/v2.10/transformer_engine/pytorch/utils.py#L759"  # pylint: disable=line-too-long
+                )
 
         gc.collect()
         torch.cuda.empty_cache()
@@ -293,6 +397,8 @@ def create_cudagraphs(cls):
         if has_te_modules:
             te_set_capture_start()
 
+        global bwd_buffer_reuse_ref_count, fwd_buffer_reuse_ref_count
+
         def format_mem_bytes(mem_bytes):
             for power, suffix in [(4, "tb"), (3, "gb"), (2, "mb"), (1, "kb"), (0, "bytes")]:
                 suffix_bytes = 1024**power
@@ -300,58 +406,25 @@ def format_mem_bytes(mem_bytes):
                     return "%.1f %s" % (mem_bytes / suffix_bytes, suffix)
             return "%d bytes" % mem_bytes
 
-        time_start = time.time()
-        mem_stats_start = torch.cuda.memory_stats()
-        progress_bar = enumerate(cls.cudagraph_record)
-        if HAVE_TQDM:
-            progress_bar = tqdm(progress_bar, "create cuda graphs", total=len(cls.cudagraph_record))
         for g_idx, g in progress_bar:
+            if torch.distributed.get_rank() == 0:
+                mem_stats = torch.cuda.memory_stats()
+                progress_str = "create cuda graphs | mem: alloc %s, res %s" % (
+                    format_mem_bytes(mem_stats["allocated_bytes.all.current"]),
+                    format_mem_bytes(mem_stats["reserved_bytes.all.current"]),
+                )
+                if HAVE_TQDM:
+                    progress_bar.set_description(progress_str)
+                elif g_idx % 100 == 0 or g_idx == len(cls.cudagraph_record) - 1:
+                    logger.info(f"{g_idx}/{len(cls.cudagraph_record)}. {progress_str}")
 
             runner, graph_type = g[0:2]
-
-            mem_stats = torch.cuda.memory_stats()
-            progress_str = "create cuda graphs | mem: alloc %s, res %s" % (
-                format_mem_bytes(mem_stats["allocated_bytes.all.current"]),
-                format_mem_bytes(mem_stats["reserved_bytes.all.current"]),
-            )
-            if HAVE_TQDM:
-                progress_bar.set_description(progress_str)
-            elif g_idx % 100 == 0 or g_idx == len(cls.cudagraph_record) - 1:
-                logger.info(f"{g_idx}/{len(cls.cudagraph_record)}. {progress_str}")
-
-            if optimize_transformer_layer_graph_buffers:
-                if graph_type == 'fwd':
-                    args, kwargs = g[2:]
-
-                    if not runner.is_first_layer:
-                        kwargs['hidden_states'] = prev_fwd_hidden_state_output
-                    runner.create_fwd_graph(args, kwargs, clone_inputs=False)
-
-                    # The output of TransformerLayer is: (hidden_states, None)
-                    # The output of MambaLayer is: (hidden_states,)
-                    # make sure to get the hidden states tensor from the tuple
-                    prev_fwd_hidden_state_output = runner.fwd_graph_outputs[0]
-
-                else:
-                    # In vision models, encoder and decoder transformers have different
-                    # hidden_states shapes. Each has its own first and last layers that
-                    # are noncontiguous. Reset prev_bwd_hidden_state_inputgrad to None at
-                    # each last layer to avoid shape mismatch when transitioning between
-                    # encoder and decoder.
-                    if runner.is_last_layer:
-                        prev_bwd_hidden_state_inputgrad = None
-
-                    runner.create_bwd_graph(prev_bwd_hidden_state_inputgrad)
-
-                    # The first input grad TransformerLayer is for 'hidden_states'
-                    prev_bwd_hidden_state_inputgrad = runner.static_grad_inputs[0]
+            if graph_type == 'fwd':
+                args, kwargs, out = g[2:]
+                runner.create_fwd_graph(args, kwargs, out, clone_inputs=True)
             else:
-                runner, graph_type = g[0:2]
-                if graph_type == 'fwd':
-                    args, kwargs = g[2:]
-                    runner.create_fwd_graph(args, kwargs)
-                else:
-                    runner.create_bwd_graph()
+                assert fwd_buffer_reuse_ref_count == 0
+                runner.create_bwd_graph()
 
         # Memory usage.
         time_end = time.time()
@@ -367,16 +440,18 @@ def format_mem_bytes(mem_bytes):
                 - mem_stats_start["reserved_bytes.all.current"]
             ),
         }
-        logger.info(
-            "> built %d cuda graph(s) in %.2f sec, with total memory usage: "
-            "allocated %s, reserved %s."
-            % (
-                len(cls.cudagraph_record),
-                capture_stats["time"],
-                format_mem_bytes(capture_stats["allocated_bytes"]),
-                format_mem_bytes(capture_stats["reserved_bytes"]),
+
+        if torch.distributed.get_rank() == 0:
+            logger.info(
+                "> built %d cuda graph(s) in %.2f sec, with total memory usage: "
+                "allocated %s, reserved %s."
+                % (
+                    len(cls.cudagraph_record),
+                    capture_stats["time"],
+                    format_mem_bytes(capture_stats["allocated_bytes"]),
+                    format_mem_bytes(capture_stats["reserved_bytes"]),
+                )
             )
-        )
 
         # Mark cuda graphs as created.
         for g in cls.cudagraph_record:
@@ -392,6 +467,8 @@ def format_mem_bytes(mem_bytes):
         if has_te_modules:
             te_set_capture_end()
 
+        torch.cuda.set_stream(torch.cuda.default_stream())
+
         # Return capture time and memory usage.
         return capture_stats
 
@@ -425,8 +502,7 @@ def delete_cuda_graphs():
         runner.bwd_graph_recorded = False
         runner.fwd_graph = None
         runner.bwd_graph = None
-        runner.fwd_mempool = None
-        runner.bwd_mempool = None
+        runner.mempool = None
 
     # Reset global tracking state
     _CudagraphGlobalRecord.cudagraph_created = False
@@ -438,8 +514,6 @@ def delete_cuda_graphs():
     torch.cuda.empty_cache()
 
     CudaGraphManager.global_mempool = None
-    CudaGraphManager.fwd_mempools = None
-    CudaGraphManager.bwd_mempool = None
 
 
 class _GraphStatus(Enum):
@@ -500,38 +574,42 @@ def forward(ctx, runner, is_first_microbatch, *inputs):
         ), "Fwd cudagraph received a different number of tensors than what it was graphed with!"
 
         # Copy new data into fwd graph input buffer
+        need_copy_inputs = []
         for user_input, cudagraph_input in zip(inputs, runner.fwd_graph_input_surface):
-            if user_input.data_ptr() != cudagraph_input.data_ptr():
+            if (
+                hasattr(cudagraph_input, "can_skip_replay_copy")
+                and cudagraph_input.can_skip_replay_copy
+            ):
+                need_copy_inputs.append(user_input)
+                assert user_input.data_ptr() == cudagraph_input.data_ptr()
+            else:
                 cudagraph_input.copy_(user_input)
 
         ctx.runner = runner
-        if runner.fp8_enabled or runner.fp4_enabled:
-            for m in runner.base_module.modules():
-                if isinstance(m, TransformerEngineBaseModule):
-                    m.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group()
-                    m.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
+        ctx.save_for_backward(*need_copy_inputs)
 
-                    if is_te_min_version("1.13.0"):
-                        FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(m.fp8_meta)
-                    else:
-                        FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(
-                            m.fp8_meta, fp8_weights=m._get_fp8_params()
-                        )
+        if runner.fp8_enabled or runner.fp4_enabled:
+            if isinstance(FP8GlobalStateManager.get_fp8_recipe(), te.common.recipe.DelayedScaling):
+                for m in runner.base_module.modules():
+                    if isinstance(m, TransformerEngineBaseModule):
+                        m.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group()
+                        m.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
+
+                        if is_te_min_version("1.13.0"):
+                            FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(m.fp8_meta)
+                        else:
+                            FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(
+                                m.fp8_meta, fp8_weights=m._get_fp8_params()
+                            )
 
-            is_first_fp8_module = FP8GlobalStateManager.is_first_fp8_module()
-            if is_first_fp8_module:
+            # Note that FP8GlobalStateManager.is_first_fp8_module() is inacccurate as each
+            # layer may be in its own fp8 context, when the fp8 recipe != delayed_scaling
+            if runner.is_first_layer and (runner.fp8_param_cache_updated != is_first_microbatch):
                 FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(not is_first_microbatch)
-            ctx.is_first_fp8_module = is_first_fp8_module
+                runner.fp8_param_cache_updated = is_first_microbatch
 
         runner.fwd_graph.replay()
-
-        # if last transformer layer, return a clone of the cudagraph output buffer, as releasing
-        # the cudagraph output buffer into the rest of the system may allow it to be corrupted
-        if runner.is_last_layer:
-            out = tuple(o.clone().detach() for o in runner.fwd_graph_output_surface)
-        else:
-            out = tuple(o.detach() for o in runner.fwd_graph_output_surface)
-        return out
+        return runner.fwd_graph_output_surface
 
     @staticmethod
     def backward(ctx, *grads):
@@ -548,16 +626,28 @@ def backward(ctx, *grads):
             runner.static_grad_outputs
         ), "Bwd cudagraph received a different number of tensors than what it was graphed with!"
 
+        need_copy_inputs = list(ctx.saved_tensors)
+        for cudagraph_input in runner.fwd_graph_input_surface:
+            if (
+                hasattr(cudagraph_input, "can_skip_replay_copy")
+                and cudagraph_input.can_skip_replay_copy
+            ):
+                cudagraph_input.copy_(need_copy_inputs.pop(0))
+
         # Copy new data into bwd graph input buffer
         for user_output_grad, cudagraph_output_grad in zip(grads, runner.static_grad_outputs):
+            if cudagraph_output_grad is None:
+                continue
             if user_output_grad.data_ptr() != cudagraph_output_grad.data_ptr():
                 cudagraph_output_grad.copy_(user_output_grad)
 
         runner.bwd_graph.replay()
         runner.status = _GraphStatus.FWD_READY
 
-        # Update FP8/FP4 scale factors if needed
-        if (runner.fp8_enabled or runner.fp4_enabled) and ctx.is_first_fp8_module:
+        # Update FP8 scale factors if needed
+        if runner.fp8_enabled and isinstance(
+            FP8GlobalStateManager.get_fp8_recipe(), te.common.recipe.DelayedScaling
+        ):
             FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
 
         # If using gradient_accumulation_fusion, whenever `main_grad` is calculated
@@ -566,18 +656,12 @@ def backward(ctx, *grads):
         for param, grad_added in runner.groundtruth_grad_added_to_main_grad.items():
             param.grad_added_to_main_grad = grad_added
 
-        grads, is_dummy_grad = runner.get_input_grads_with_dummy_flags()
-        if runner.is_first_layer:
-            output_grads = tuple(
-                b.clone().detach() if not (b is None or dummy) else b
-                for dummy, b in zip(is_dummy_grad, grads)
-            )
-        else:
-            output_grads = tuple(
-                b.detach() if not (b is None or dummy) else b
-                for dummy, b in zip(is_dummy_grad, grads)
-            )
-        return None, None, *output_grads
+        # Replaying the next bwd graph destroys the data held in static_grad_inputs, so clone
+        # wgrads as autograd may launch the next graph before wgrads are accumulated
+        dgrads = runner.static_grad_inputs[: runner.num_dgrads]
+        wgrads = (g.clone() for g in runner.static_grad_inputs[runner.num_dgrads :])
+
+        return None, None, *dgrads, *wgrads
 
 
 class _CudaGraphRunner(torch.nn.Module):
@@ -588,23 +672,20 @@ class _CudaGraphRunner(torch.nn.Module):
     def __init__(
         self,
         base_module: MegatronModule,
-        fwd_mempool: int,
-        bwd_mempool: int,
+        mempool: int,
         fwd_graph_input_args: List[Any],
         fwd_graph_input_kwargs: Dict[str, Any],
-        share_cudagraph_io_buffers=None,
+        func,
+        need_backward,
     ):
         """Creates a _CudaGraphRunner, which holds a single pair of fwd and bwd cudagraphs, which
         are not created until this runner records its graph creation into
-        '_CudagraphGlobalRecord', and 'create_cudagraphs()' is called. share_cudagraph_io_buffers
-        is a boolean flag to indicate whether to reuse the cudagraph input and output buffers for
-        transformer layer specific optimizations that reduce memory usage and tensor copies."""
+        '_CudagraphGlobalRecord', and 'create_cudagraphs()' is called."""
 
         super().__init__()
 
         self.base_module = base_module
-        self.fwd_mempool = fwd_mempool
-        self.bwd_mempool = bwd_mempool
+        self.mempool = mempool
 
         self.fwd_graph_input_arg_metas = [ArgMetadata(a) for a in fwd_graph_input_args]
         self.fwd_graph_input_kwarg_metas = {
@@ -624,14 +705,30 @@ def __init__(
         self.fp8_enabled = False
         self.fp4_enabled = False
         self.deallocate_pipeline_outputs = False
-        self.num_warmup_steps = 2
-        if isinstance(self.base_module.config, TransformerConfig):
+
+        self.grad_enabled = need_backward and torch.is_grad_enabled()
+        self.func = super(MegatronModule, self.base_module).__call__ if func is None else func
+        self.is_first_layer, self.is_last_layer = _determine_if_first_last_layer_of_this_vp_chunk(
+            base_module
+        )
+
+        # We use this attribute to record the value of 'is_first_microbatch' each fwd cudagraph
+        # replay so that way we only update the value of this flag in FP8GlobalStateManager when
+        # it changes which incurs an HtoD sync
+        if self.is_first_layer:
+            self.fp8_param_cache_updated = None
+
+        if hasattr(self.base_module, "config") and isinstance(
+            self.base_module.config, TransformerConfig
+        ):
             self.fuse_wgrad_accumulation = self.base_module.config.gradient_accumulation_fusion
             self.backward_retain_grad = self.base_module.config.cuda_graph_retain_backward_graph
-            self.fp8_enabled = self.base_module.config.fp8 is not None
-            self.fp4_enabled = self.base_module.config.fp4 is not None
             self.deallocate_pipeline_outputs = self.base_module.config.deallocate_pipeline_outputs
             self.num_warmup_steps = self.base_module.config.cuda_graph_warmup_steps
+            self.fp8_enabled = self.base_module.config.fp8 is not None
+            self.fp4_enabled = self.base_module.config.fp4 is not None
+            self.fp8_runtime_enabled = None
+            self.fp4_runtime_enabled = None
 
             if self.fp8_enabled:
                 self.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
@@ -643,84 +740,91 @@ def __init__(
                 self.fp4_recipe = get_fp4_recipe(self.base_module.config)
                 FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False)
 
-        # Decide whether to reuse the input and output buffer, and if so,
-        # whether this layer is the first layer (which needs an input buffer)
-        # or the last layer (which needs an output buffer)
-
-        self.is_transformer_decoder_layer = _determine_if_transformer_decoder_layer(base_module)
-        self.reuse_input_output_buffer = (
-            share_cudagraph_io_buffers and self.is_transformer_decoder_layer
-        )
-        if self.reuse_input_output_buffer:
-            self.is_first_layer, self.is_last_layer = (
-                _determine_if_first_last_layer_of_this_vp_chunk(base_module)
-            )
-        else:
-            self.is_first_layer, self.is_last_layer = True, True
-
     def __str__(self):
         return "%s; hid %s" % (
             self.base_module.__class__.__name__,
             tuple(self.fwd_graph_input_kwarg_metas["hidden_states"].shape),
         )
 
-    def get_fp8_context(self):
-        """Return a new fp8 context in cudagraph mode."""
-        from megatron.core.fp8_utils import get_fp8_context  # to avoid circular import
-
-        return get_fp8_context(self.base_module.config, self.base_module.layer_number - 1)
-
-    def get_fp4_context(self):
-        """Return a new fp4 context in cudagraph mode."""
-        from megatron.core.fp4_utils import get_fp4_context  # to avoid circular import
-
-        return get_fp4_context(self.base_module.config, self.base_module.layer_number - 1)
-
     def get_quantization_context(self):
         """Return appropriate quantization context (FP8 or FP4) in cudagraph mode."""
-        if self.fp8_enabled:
-            return self.get_fp8_context()
-        elif self.fp4_enabled:
-            return self.get_fp4_context()
+        if self.fp8_runtime_enabled:
+            from megatron.core.fp8_utils import get_fp8_context  # to avoid circular import
+
+            return get_fp8_context(self.base_module.config, self.base_module.layer_number - 1)
+        elif self.fp4_runtime_enabled:
+            from megatron.core.fp4_utils import get_fp4_context  # to avoid circular import
+
+            return get_fp4_context(self.base_module.config, self.base_module.layer_number - 1)
         else:
             return nullcontext()
 
-    def create_fwd_graph(self, args, kwargs, clone_inputs=True):
+    def get_connected_params(self, outputs):
+        """Iterate through the autograd graph of 'outputs' and returns all parameters connected.
+        In theory this should return all parameters that return a nonzero wgrad when computing
+        the backward pass of 'outputs'."""
+        # Flatten outputs and start traversal from roots that require gradients
+        args = (outputs,) if torch.is_tensor(outputs) else outputs
+        stack = [
+            t.grad_fn
+            for t in self.get_tensors(args, check_types=False)
+            if t.requires_grad and t.grad_fn
+        ]
+        visited, p_ids = set(), set()
+
+        while stack:
+            if (fn := stack.pop()) not in visited:
+                visited.add(fn)
+                # AccumulateGrad nodes (leafs) hold the 'variable' (Parameter) they accumulate into
+                if hasattr(fn, 'variable'):
+                    p_ids.add(id(fn.variable))
+                stack.extend(f for f, _ in fn.next_functions if f)
+
+        # Return module params that were found in the graph, preserving original order
+        return tuple(p for p in self.base_module.parameters() if id(p) in p_ids)
+
+    def create_fwd_graph(self, args, kwargs, outputs=None, clone_inputs=True):
         """Create a fwd cudagraph for this runner. Should be called inside
         'create_cudagraphs()'."""
 
-        # Freeze GC, to speed up capture time ~15-20x.
-        if FREEZE_GC:
-            gc.freeze()
+        global fwd_buffer_reuse_ref_count
+
+        self.args = args
+        self.kwargs = kwargs
+        self.outputs = outputs
 
         # save grads and other variables that may be affected by graph warmup
         if self.training and torch.is_grad_enabled():
-            save_main_grads = [
-                param.main_grad.clone()
-                for param in self.base_module.parameters()
-                if hasattr(param, 'main_grad')
-            ]
+            grad_backup = []
+            for param in self.base_module.parameters():
+                grad_backup.append(param.main_grad.clone() if hasattr(param, "main_grad") else None)
 
-        saved_fp8_tensors = None
+            saved_fp8_tensors = None
+            if self.fp8_enabled:
+                if is_te_min_version("1.13.0"):
+                    saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp8_recipe)
+                else:
+                    saved_fp8_tensors = save_fp8_tensors(
+                        [self.base_module], self.fp8_recipe.amax_history_len
+                    )
+            elif self.fp4_enabled:
+                if is_te_min_version("2.7.0.dev0"):
+                    saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp4_recipe)
+                else:
+                    raise ValueError("FP4 requires TE >= 2.7.0.dev0 for NVFP4BlockScaling support.")
 
-        if self.fp8_enabled:
-            if is_te_min_version("1.13.0"):
-                saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp8_recipe)
-            else:
-                saved_fp8_tensors = save_fp8_tensors(
-                    [self.base_module], self.fp8_recipe.amax_history_len
-                )
-        elif self.fp4_enabled:
-            if is_te_min_version("2.7.0.dev0"):
-                saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp4_recipe)
-            else:
-                raise ValueError("FP4 requires TE >= 2.7.0.dev0 for NVFP4BlockScaling support.")
+        # cache the moe aux loss if needed, which is accumulated inside the forward pass
+        from megatron.core.transformer.transformer_layer import MoETransformerLayer
 
-        if clone_inputs:
-            args, kwargs = self.zero_out_tensors(args, kwargs)
+        is_moe = isinstance(self.base_module, MoETransformerLayer)
+        if is_moe:
+            from megatron.core.transformer.moe.moe_utils import get_moe_layer_wise_logging_tracker
 
-        input_tensors = self.get_tensors(args, kwargs)
-        self.fwd_graph_input_surface = input_tensors + tuple(self.base_module.parameters())
+            tracker = get_moe_layer_wise_logging_tracker()
+            cached_aux_losses = {}
+            for name in tracker:
+                if "values" in tracker[name]:
+                    cached_aux_losses[name] = torch.clone(tracker[name]["values"])
 
         self.fwd_graph = torch.cuda.CUDAGraph()
 
@@ -732,183 +836,339 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True):
                     _ensure_generator_state_is_cudagraph_safe(gen)
                 )
 
-        # warmup again as case graph capture mode may execute a different codepath
-        for _ in range(self.num_warmup_steps):
-            with self.get_quantization_context():
-                outputs = self.base_module.forward(*args, **kwargs)
-            if self.training and torch.is_grad_enabled():
-                if isinstance(outputs, torch.Tensor):
-                    outputs = (outputs,)
-                outputs = self.get_tensors(outputs)
-                grad_inputs = torch.autograd.grad(
-                    outputs=tuple(o for o in outputs if o.requires_grad),
-                    inputs=tuple(i for i in self.fwd_graph_input_surface if i.requires_grad),
-                    grad_outputs=tuple(
-                        torch.zeros_like(o) if o.requires_grad else None for o in outputs
-                    ),
-                    only_inputs=True,
-                    allow_unused=True,
+        def _resolve_input_buffer(ten):
+            if not isinstance(ten, ArgMetadata):
+                return ten
+
+            # the input tensor is resued from another cudagraph's input or output
+            if (
+                hasattr(ten, "cg_buffer_metadata")
+                and ten.cg_buffer_metadata.fwd_cudagraph_buffer is not None
+            ):
+                global fwd_buffer_reuse_ref_count
+                buf = ten.cg_buffer_metadata.fwd_cudagraph_buffer
+
+                assert (
+                    ten.cg_buffer_metadata.is_cudagraph_input
+                    and buf.cg_buffer_metadata.capture_reuse_count > 0
                 )
 
-        with self.get_quantization_context():
-            torch.cuda.synchronize()
-            # Register default CUDA generators ourselves (fixed in-place to have normal tensors)
-            # before capture begins, to avoid inference-tensor state issues during capture.
-            with torch.inference_mode(mode=False):
-                for device_idx in range(torch.cuda.device_count()):
-                    default_gen = torch.cuda.default_generators[device_idx]
-                    self.fwd_graph.register_generator_state(
-                        _ensure_generator_state_is_cudagraph_safe(default_gen)
+                if (
+                    ten.cg_buffer_metadata.input_use_count > 1
+                    and ten.cg_buffer_metadata.input_use_count
+                    == buf.cg_buffer_metadata.capture_reuse_count
+                ):
+                    can_skip_replay_copy = False
+                else:
+                    can_skip_replay_copy = True
+
+                buf.cg_buffer_metadata.capture_reuse_count -= 1
+                if buf.cg_buffer_metadata.capture_reuse_count == 0:
+                    ten.cg_buffer_metadata.fwd_cudagraph_buffer = None
+                    fwd_buffer_reuse_ref_count -= 1
+            else:
+                # need to provide a fresh buffer from the reuse pool
+                buf = _CudagraphGlobalRecord.tensor_reuse_pool.get(ten)
+                can_skip_replay_copy = False
+
+            buf = buf.detach().requires_grad_(ten.requires_grad)
+            buf.can_skip_replay_copy = can_skip_replay_copy
+            return buf
+
+        if clone_inputs:
+            # if a buffer is used for multiple inputs, create it now
+            for ten in self.get_tensors(args, kwargs):
+                if (
+                    hasattr(ten, 'cg_buffer_metadata')
+                    and ten.cg_buffer_metadata.input_use_count > 1
+                    and ten.cg_buffer_metadata.fwd_cudagraph_buffer is None
+                ):
+                    buf = _CudagraphGlobalRecord.tensor_reuse_pool.get(ten)
+                    buf.cg_buffer_metadata = deepcopy(ten.cg_buffer_metadata)
+                    buf.cg_buffer_metadata.capture_reuse_count = (
+                        ten.cg_buffer_metadata.input_use_count
                     )
+                    ten.cg_buffer_metadata.fwd_cudagraph_buffer = buf
+                    fwd_buffer_reuse_ref_count += 1
 
-            with torch.cuda.graph(
-                self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local"
-            ):
-                outputs = self.base_module.forward(*args, **kwargs)
+            self.fwd_graph_input_args = tree_map(_resolve_input_buffer, args)
+            self.fwd_graph_input_kwargs = tree_map(_resolve_input_buffer, kwargs)
+        else:
+            self.fwd_graph_input_args, self.fwd_graph_input_kwargs = args, kwargs
+
+        self.fwd_graph_input_surface = self.get_tensors(
+            self.fwd_graph_input_args, self.fwd_graph_input_kwargs
+        )
+
+        ctx = torch.no_grad() if not self.grad_enabled else nullcontext()
+        with ctx:
+            # warmup again as case graph capture mode may execute a different codepath
+            _set_warmup_start()
+            for _ in range(self.num_warmup_steps):
+                with self.get_quantization_context():
+
+                    def clone_ten(ten):
+                        if not torch.is_tensor(ten):
+                            return ten
+                        return torch.zeros_like(ten).requires_grad_(ten.requires_grad)
+
+                    warmup_args = tree_map(clone_ten, self.fwd_graph_input_args)
+                    warmup_kwargs = tree_map(clone_ten, self.fwd_graph_input_kwargs)
+                    warmup_outputs = self.func(*warmup_args, **warmup_kwargs)
+
+                if self.grad_enabled:
+                    warmup_outputs = self.get_tensors(warmup_outputs)
+                    warmup_outputs = tuple(o for o in warmup_outputs if o.requires_grad)
+                    input_tensors = self.get_tensors(warmup_args, warmup_kwargs)
+                    torch.autograd.grad(
+                        outputs=warmup_outputs,
+                        inputs=tuple(i for i in input_tensors if i.requires_grad),
+                        grad_outputs=tuple(torch.zeros_like(o) for o in warmup_outputs),
+                        only_inputs=True,
+                        allow_unused=True,
+                    )
+            _set_warmup_end()
+
+            with self.get_quantization_context():
+                torch.cuda.synchronize()
+                # Register default CUDA generators ourselves (fixed in-place to have normal tensors)
+                # before capture begins, to avoid inference-tensor state issues during capture.
+                with torch.inference_mode(mode=False):
+                    for device_idx in range(torch.cuda.device_count()):
+                        default_gen = torch.cuda.default_generators[device_idx]
+                        self.fwd_graph.register_generator_state(
+                            _ensure_generator_state_is_cudagraph_safe(default_gen)
+                        )
+
+                # Freeze GC, to speed up capture time ~15-20x.
+                if FREEZE_GC:
+                    gc.freeze()
+
+                with torch.cuda.graph(
+                    self.fwd_graph, pool=self.mempool, capture_error_mode="thread_local"
+                ):
+                    fwd_graph_outputs = self.func(
+                        *self.fwd_graph_input_args, **self.fwd_graph_input_kwargs
+                    )
+
+                # Unfreeze GC.
+                if FREEZE_GC:
+                    gc.unfreeze()
+
+                    # gc.collect() drops references to unreachable tensors created during capture,
+                    # returning their storage to the allocator to avoid a slowdown during replay.
+                    # However, it forces expensive global garbage collection, so must be done
+                    # only on the last layer per-device to avoid slowing down graph creation.
+                    if self.is_last_layer:
+                        gc.collect()
 
         # save cudagraph output buffer
-        if isinstance(outputs, torch.Tensor):
-            outputs = (outputs,)
-        self.fwd_graph_outputs = outputs
-        self.fwd_graph_output_surface = self.get_tensors(outputs)
+        self.fwd_graph_outputs = fwd_graph_outputs
+        self.fwd_graph_output_surface = self.get_tensors(fwd_graph_outputs)
+
+        for fwd_graph_out, o in zip(
+            self.fwd_graph_output_surface, self.get_arg_metas(self.outputs)
+        ):
+            assert hasattr(o, "cg_buffer_metadata") and o.cg_buffer_metadata.is_cudagraph_output
+
+            if (
+                o.cg_buffer_metadata.is_cudagraph_input
+                and o.cg_buffer_metadata.fwd_cudagraph_buffer is None
+            ):
+                fwd_graph_out.cg_buffer_metadata = deepcopy(o.cg_buffer_metadata)
+                fwd_graph_out.cg_buffer_metadata.capture_reuse_count = (
+                    o.cg_buffer_metadata.cudagraph_reuse_ref_count
+                )
+                o.cg_buffer_metadata.fwd_cudagraph_buffer = fwd_graph_out
+                fwd_buffer_reuse_ref_count += 1
+
+        # if an input buffer requires a copy, and does not have metadata attached to it at this
+        # point, it will not be reused after this forward pass, so return it to the pool
+        for buf in self.fwd_graph_input_surface:
+            if (
+                hasattr(buf, "can_skip_replay_copy")
+                and not buf.can_skip_replay_copy
+                and not hasattr(buf, "cg_buffer_metadata")
+            ):
+                assert _CudagraphGlobalRecord.tensor_reuse_pool.owns(buf)
+                _CudagraphGlobalRecord.tensor_reuse_pool.insert(buf)
 
         if self.training and torch.is_grad_enabled():
             assert (
                 len(self.fwd_graph_output_surface) > 0
-            ), """Tried graphing a moudule that returned no tensors in training mode,
-                however the graphed module must output at least one tensor,
+            ), """Tried graphing a module that returned no tensors in training mode, 
+                however the graphed module must output at least one tensor, 
                 so that a corresponding backward node may be registered in the autograd graph."""
 
-            # restore cached grads
-            for param in self.base_module.parameters():
-                if hasattr(param, 'main_grad'):
-                    saved_grad = save_main_grads.pop(0)
-                    assert (
-                        param.main_grad.shape == saved_grad.shape
-                    ), "Error restoring grads while cudagraphing!"
-                    param.main_grad.copy_(saved_grad)
+            self.params_to_backprop = self.get_connected_params(fwd_graph_outputs)
+            self.num_wgrads = len(self.params_to_backprop)
+            self.num_dgrads = len(self.fwd_graph_input_surface)
+            self.fwd_graph_input_surface = self.fwd_graph_input_surface + self.params_to_backprop
 
-        if self.fp8_enabled or self.fp4_enabled:
-            restore_fp8_tensors([self.base_module], saved_fp8_tensors)
-
-        # Unfreeze GC.
-        if FREEZE_GC:
-            gc.unfreeze()
+            if self.fp8_enabled:
+                restore_fp8_tensors([self.base_module], saved_fp8_tensors)
+            # restore cached grads
+            for main_grad_copy, param in zip(grad_backup, self.base_module.parameters()):
+                if main_grad_copy is not None:
+                    param.main_grad.copy_(main_grad_copy)
 
-            # gc.collect() drops references to unreachable tensors created during capture,
-            # returning their storage to the allocator to avoid a slowdown during replay. However,
-            # it forces expensive global garbage collection, so must be done only on the last layer
-            # per-device to avoid slowing down graph creation.
-            if self.is_last_layer:
-                gc.collect()
+        if is_moe:
+            for name in tracker:
+                tracker[name]["values"].copy_(cached_aux_losses[name])
 
-    def create_bwd_graph(self, static_grad_outputs=None):
+    def create_bwd_graph(self):
         """Create a bwd cudagraph for this runner. Should be called inside
         'create_cudagraphs()'."""
 
-        # Freeze GC, to speed up capture time ~15-20x.
-        if FREEZE_GC:
-            gc.freeze()
+        # unlike 'fwd_buffer_reuse_ref_count', 'bwd_buffer_reuse_ref_count' may not decrement
+        # to 0 when activation checkpointing is used. See [interaction with recompute].
+        global bwd_buffer_reuse_ref_count
 
+        assert self.grad_enabled
         self.bwd_graph = torch.cuda.CUDAGraph()
 
         # For cases with multiple active RNG states, e.g. TP.
         for _, state in get_all_rng_states().items():
             self.bwd_graph.register_generator_state(state)
 
-        if static_grad_outputs is None:
-            static_grad_outputs = tuple(
-                torch.zeros_like(o) if o.requires_grad else None
-                for o in self.fwd_graph_output_surface
-            )
-        else:
-            # canoncalize as tuple
-            if torch.is_tensor(static_grad_outputs):
-                static_grad_outputs = (static_grad_outputs,)
+        self.static_grad_outputs = []
+        for o in self.get_arg_metas(self.outputs):
+            out_grad = None
+            if o.requires_grad:
+                # TODO: (jiemingz) [interaction with recompute]
+                # for activation recompute, the fwd pass is rerun in the backward pass and
+                # the metadata we attach in record_graph_capture is lost. As a result the next
+                # cudagraph expects the buffer to be provided 'fwd_cudagraph_buffer' but is missing.
+                # So, we cannot always assume this metadata exists. Consequently, there are extra
+                # copies between the outputs of the fwd-bwd pass and the bwd pass.
+                if (
+                    o.cg_buffer_metadata.is_cudagraph_input
+                    and o.cg_buffer_metadata.bwd_cudagraph_buffer is not None
+                ):
+                    o.cg_buffer_metadata.bwd_cudagraph_buffer.shape == o.shape
 
-        torch.cuda.synchronize()
-        with torch.cuda.graph(
-            self.bwd_graph, pool=self.bwd_mempool, capture_error_mode="thread_local"
-        ):
+                    out_grad = o.cg_buffer_metadata.bwd_cudagraph_buffer
+                    o.cg_buffer_metadata.bwd_cudagraph_buffer = None
+                    out_grad.cg_buffer_metadata.capture_reuse_count -= 1
+                    bwd_buffer_reuse_ref_count -= 1
+                else:
+                    out_grad = _CudagraphGlobalRecord.tensor_reuse_pool.get(o)
+                out_grad.requires_grad = True
+            self.static_grad_outputs.append(out_grad)
+
+        # Freeze GC, to speed up capture time ~15-20x.
+        if FREEZE_GC:
+            gc.freeze()
+
+        with torch.cuda.graph(self.bwd_graph, pool=self.mempool):
             grad_inputs = torch.autograd.grad(
                 outputs=tuple(o for o in self.fwd_graph_output_surface if o.requires_grad),
                 inputs=tuple(i for i in self.fwd_graph_input_surface if i.requires_grad),
-                grad_outputs=tuple(o for o in static_grad_outputs if o is not None),
+                grad_outputs=tuple(o for o in self.static_grad_outputs if o is not None),
                 retain_graph=self.backward_retain_grad,
                 only_inputs=True,
                 allow_unused=True,
             )
 
+        # Unfreeze GC.
+        if FREEZE_GC:
+            gc.unfreeze()
+
         # Constructs a tuple suitable for returning from Graphed.backward:
         # Pads out the actually-needed grads with Nones in gradient slots for inputs
-        # that don't require grad. I couldn't think of a one-liner for this pattern.
-        static_grad_inputs = []
-        grad_idx = 0
-        for arg in self.fwd_graph_input_surface:
-            has_wgrad_fusion = self.fuse_wgrad_accumulation and getattr(
-                arg, "grad_added_to_main_grad", False
-            )
-            if arg.requires_grad:
-                if has_wgrad_fusion:
-                    static_grad_inputs.append(None)
-                else:
-                    static_grad_inputs.append(grad_inputs[grad_idx])
-                grad_idx += 1
+        # that don't require grad
+        grad_inputs = list(grad_inputs)
+        self.static_grad_inputs = []
+        for input_tensor in self.get_arg_metas(self.args, self.kwargs):
+            if input_tensor.requires_grad:
+                input_grad = grad_inputs.pop(0)
+                input_grad.cg_buffer_metadata = deepcopy(input_tensor.cg_buffer_metadata)
+                if input_tensor.cg_buffer_metadata.is_cudagraph_output:
+                    if input_tensor.cg_buffer_metadata.bwd_cudagraph_buffer is None:
+                        input_tensor.cg_buffer_metadata.bwd_cudagraph_buffer = input_grad
+                        input_grad.cg_buffer_metadata.capture_reuse_count += 1
+                        bwd_buffer_reuse_ref_count += 1
+                self.static_grad_inputs.append(input_grad)
             else:
-                static_grad_inputs.append(None)
+                self.static_grad_inputs.append(None)
+
+        # at this point static_grad_inputs hold the input dgrads, add the wgrads next
+        assert self.num_wgrads == len(grad_inputs)
+        self.static_grad_inputs.extend(grad_inputs)
+        self.static_grad_inputs = tuple(self.static_grad_inputs)
+        self.static_grad_outputs = tuple(self.static_grad_outputs)
 
         self.groundtruth_grad_added_to_main_grad = {}
         if self.fuse_wgrad_accumulation:
-            for param in self.base_module.parameters():
+            for param in self.params_to_backprop:
                 if hasattr(param, "grad_added_to_main_grad"):
                     self.groundtruth_grad_added_to_main_grad[param] = param.grad_added_to_main_grad
 
-        self.static_grad_outputs = static_grad_outputs
-        self.static_grad_inputs = static_grad_inputs
+        # After backward pass grad_output buffers are no longer used and returned to the pool
+        for ten in self.static_grad_outputs:
+            if torch.is_tensor(ten):
+                # Check that the tensor is not in use. This scenario may occur when a cudagraph
+                # passes its input directly as an output, and places this output as the
+                # input of a subsequent cudgraph, leading to a grad output buffer to be still in use
+                # even after the backward pass.
+                reuse_count = (
+                    ten.cg_buffer_metadata.capture_reuse_count
+                    if hasattr(ten, "cg_buffer_metadata")
+                    else 0
+                )
 
-        # Unfreeze GC.
-        if FREEZE_GC:
-            gc.unfreeze()
+                if _CudagraphGlobalRecord.tensor_reuse_pool.owns(ten) and reuse_count == 0:
+                    _CudagraphGlobalRecord.tensor_reuse_pool.insert(ten)
 
-            if self.is_first_layer:
-                gc.collect()
+        # now weakref everything
+        if HAVE_TE_GRAPHS:
 
-    def get_input_grads_with_dummy_flags(self):
-        """Get the inputs grads that are returned by the bwd cudagraph call. If using grad accum
-        fusion, wgrads have already been accumulated, so return dummy wgrads."""
+            def replace_with_weak_ref(arg):
+                if not torch.is_tensor(arg):
+                    return arg
 
-        is_dummy_grad = [False] * len(self.static_grad_inputs)
-        if not self.fuse_wgrad_accumulation:
-            return self.static_grad_inputs, is_dummy_grad
-        else:
-            num_dgrads = len(self.static_grad_inputs) - len(list(self.base_module.parameters()))
-            dgrads = self.static_grad_inputs[:num_dgrads]
-            wgrads = self.static_grad_inputs[num_dgrads:]
-
-            wgrads_with_placeholders = []
-            is_dummy_grad = [False] * len(dgrads)
-            for idx, param in enumerate(self.base_module.parameters()):
-                wgrad_is_dummy = getattr(param, "grad_added_to_main_grad", False)
-                if wgrad_is_dummy:
-                    if getattr(param, "zero_out_wgrad", False):
-                        wgrad = torch.zeros(
-                            param.main_grad.shape,
-                            dtype=param.dtype,
-                            device=torch.cuda.current_device(),
-                            requires_grad=False,
-                        )
-                    else:
-                        wgrad = torch.empty(
-                            param.main_grad.shape,
-                            dtype=param.dtype,
-                            device=torch.cuda.current_device(),
-                            requires_grad=False,
-                        )
-                else:
-                    wgrad = wgrads[idx]
-                wgrads_with_placeholders.append(wgrad)
-                is_dummy_grad.append(wgrad_is_dummy)
-            return tuple(dgrads + wgrads_with_placeholders), is_dummy_grad
+                ref = make_weak_ref(arg)
+                ref.requires_grad = arg.requires_grad
+                if hasattr(arg, "can_skip_replay_copy"):
+                    ref.can_skip_replay_copy = arg.can_skip_replay_copy
+                return ref
+
+            self.fwd_graph_input_surface = tree_map(
+                replace_with_weak_ref, self.fwd_graph_input_surface
+            )
+            self.fwd_graph_input_args = tree_map(replace_with_weak_ref, self.fwd_graph_input_args)
+            self.fwd_graph_input_kwargs = tree_map(
+                replace_with_weak_ref, self.fwd_graph_input_kwargs
+            )
+            self.fwd_graph_output_surface = tree_map(
+                replace_with_weak_ref, self.fwd_graph_output_surface
+            )
+            # It is safe to weakref static_grad_inputs as any inuse input grads have a strong ref
+            # stored in 'bwd_cudagraph_buffer'
+            self.static_grad_inputs = tree_map(replace_with_weak_ref, self.static_grad_inputs)
+            self.static_grad_outputs = tree_map(replace_with_weak_ref, self.static_grad_outputs)
+
+        delattr(self, "args")
+        delattr(self, "kwargs")
+        delattr(self, "outputs")
+
+    def apply_cudagraph_record_metadata(self, args, kwargs, outputs):
+        """Attaches graph capture metadata to all passed in tensors."""
+
+        for t in self.get_tensors(args, kwargs):
+            if not hasattr(t, "cg_buffer_metadata"):
+                t.cg_buffer_metadata = CudagraphBufferMetadata()
+
+            t.cg_buffer_metadata.is_cudagraph_input = True
+            t.cg_buffer_metadata.input_use_count += 1
+
+            if t.cg_buffer_metadata.is_cudagraph_output:
+                t.cg_buffer_metadata.cudagraph_reuse_ref_count += 1
+
+        # mark all outputs, so that the fwd graph we may reuse cudagraph output buffers as inputs
+        for o in self.get_tensors(outputs):
+            o.cg_buffer_metadata = CudagraphBufferMetadata()
+            o.cg_buffer_metadata.is_cudagraph_output = True
 
     def record_graph_capture(self, args, kwargs):
         """Records the data needed to create this runner's forward cudagraph.
@@ -916,21 +1176,8 @@ def record_graph_capture(self, args, kwargs):
         The actual cudagraph will be created when 'create_cudagraphs()` is called. Subsequent
         passes should replay the graph."""
 
-        if not self.fwd_graph_recorded:
-            logger.debug(f"Recording forward graph creation...")
-            if self.is_transformer_decoder_layer and not self.is_first_layer:
-                # transformer layers hidden_states are already saved as the output of the previous
-                # layer's cudagraph so avoid saving again
-                kwargs_copy = dict(kwargs)
-                kwargs_copy['hidden_states'] = None
-                _CudagraphGlobalRecord.record_fwd_graph(self, args, kwargs_copy)
-            else:
-                _CudagraphGlobalRecord.record_fwd_graph(self, args, kwargs)
-
-            self.fwd_graph_recorded = True
-
         # Run the forward pass as normal in eager mode.
-        out = super(MegatronModule, self.base_module).__call__(*args, **kwargs)
+        out = self.func(*args, **kwargs)
 
         if type(out) != tuple:
             out = (out,)
@@ -947,9 +1194,38 @@ def record_graph_capture(self, args, kwargs):
             ]
         )
 
-        # autograd nodes return inputs as views, so clone the tensor as returning views may cause
-        # issues, for instance with pipeline parallelism
-        return tuple(o.clone() if torch.is_tensor(o) else o for o in out)
+        if not self.fwd_graph_recorded:
+            logger.debug(f"Recording forward graph creation...")
+
+            self.apply_cudagraph_record_metadata(args, kwargs, out)
+
+            def _replace_with_meta(arg):
+                return ArgMetadata(arg) if torch.is_tensor(arg) else arg
+
+            m_args = tree_map(_replace_with_meta, args)
+            m_kwargs = tree_map(_replace_with_meta, kwargs)
+            m_out = tree_map(_replace_with_meta, out)
+            _CudagraphGlobalRecord.record_fwd_graph(self, m_args, m_kwargs, m_out)
+
+            if HAVE_TE_GRAPHS:
+                if FP8GlobalStateManager.is_fp8_enabled():
+                    # check if the low precision recipe is either fp4 or fp8
+                    if is_te_min_version("2.7.0.dev0"):
+                        from transformer_engine.common.recipe import NVFP4BlockScaling
+
+                        recipe = FP8GlobalStateManager.get_fp8_recipe()
+                        if isinstance(recipe, NVFP4BlockScaling):
+                            self.fp4_runtime_enabled = True
+                        else:
+                            self.fp8_runtime_enabled = True
+                    else:
+                        self.fp8_runtime_enabled = True
+
+            self.fwd_graph_recorded = True
+
+        if len(out) == 1:
+            return out[0]
+        return tuple(out)
 
     def replay_graph_capture(self, is_first_microbatch, args, kwargs):
         """Replay the fwd cuda graph with autograd."""
@@ -962,15 +1238,17 @@ def replay_graph_capture(self, is_first_microbatch, args, kwargs):
             error_msg = "CUDA graph argument mismatch:\n" + "\n".join(mismatch_errors)
             raise AssertionError(error_msg)
 
-        inp_tensors = self.get_tensors(args, kwargs)
-        func_args = inp_tensors + tuple(self.parameters())
-        out = _CudagraphReplayNode.apply(self, is_first_microbatch, *func_args)
-        out = list(out)
+        inp_tensors = self.get_tensors(args, kwargs, check_types=False)
+        if self.grad_enabled:
+            func_args = inp_tensors + self.params_to_backprop
+        else:
+            func_args = inp_tensors
 
-        if torch.is_tensor(self.fwd_graph_outputs):
-            self.fwd_graph_outputs = [self.fwd_graph_outputs]
+        out = _CudagraphReplayNode.apply(self, is_first_microbatch, *func_args)
 
-        return tuple(out.pop(0) if torch.is_tensor(o) else o for o in self.fwd_graph_outputs)
+        out_iter = iter(self.to_list(out))
+        fwd_outputs = self.to_list(self.fwd_graph_outputs)
+        return tuple(next(out_iter) if torch.is_tensor(o) else o for o in fwd_outputs)
 
     def get_mismatch_errors(self, args, kwargs):
         """Return list of detailed errors for mismatched cudagraph args."""
@@ -1003,7 +1281,7 @@ def check(val, ref, context):
                     add_error(f"Tensor mismatch at {context}: {', '.join(mismatches)}")
 
             elif is_dataclass(ref.value):
-                for field in fields(ref.value):
+                for field in dataclasses.fields(ref.value):
                     check(
                         ArgMetadata(getattr(val.value, field.name)),
                         ArgMetadata(getattr(ref.value, field.name)),
@@ -1035,66 +1313,55 @@ def check(val, ref, context):
 
         return errors
 
-    def zero_out_tensors(self, args, kwargs=None):
-        """Replace all tensors inside arg, kwargs with zeroed copies."""
+    def get_arg_metas(self, args, kwargs=None):
+        """Replaces all passed in tensors with 'ArgMetadata' and returns them as a list."""
+        arg_metas = []
 
-        def clone_tensor(ten):
-            cloned = torch.zeros_like(ten)
-            cloned.requires_grad = ten.requires_grad
-            return cloned
+        def collect(item):
+            if isinstance(item, ArgMetadata):
+                arg_metas.append(item)
+            return item  # tree_map expects a return value to rebuild the tree
 
-        def process_arg(arg):
-            _check_supported_type(ArgMetadata(arg))
-            if torch.is_tensor(arg):
-                return clone_tensor(arg)
-            elif is_dataclass(arg):
-                for field in fields(arg):
-                    attr = getattr(arg, field.name)
-                    if torch.is_tensor(attr):
-                        setattr(arg, field.name, clone_tensor(attr))
-            return arg
-
-        args_replaced = []
-        for arg in args:
-            args_replaced.append(process_arg(arg))
-        if kwargs is None:
-            return args_replaced
-
-        kwargs_replaced = {}
-        for k, v in kwargs.items():
-            kwargs_replaced[k] = process_arg(v)
-
-        return args_replaced, kwargs_replaced
+        tree_map(collect, args)
+        if kwargs is not None:
+            tree_map(collect, kwargs)
 
-    @classmethod
-    def get_tensors(cls, args, kwargs=None):
-        """Filter and flatten all tensors from args and kwargs."""
+        return arg_metas
+
+    def get_tensors(self, args, kwargs=None, check_types=True):
+        """
+        Filter and flatten all tensors from args and kwargs using list comprehensions
+        and itertools.chain for faster flattening.
+        """
 
         def extract_tensors(arg):
-            _check_supported_type(ArgMetadata(arg))
+            if check_types:
+                _check_supported_type(ArgMetadata(arg))
             if torch.is_tensor(arg):
                 return [arg]
-            elif is_dataclass(arg):
-                tens = []
-                for field in fields(arg):
-                    attr = getattr(arg, field.name)
-                    if torch.is_tensor(attr):
-                        tens.append(attr)
-                return tens
-            else:
-                return []
 
-        tens = []
-        args, _ = tree_flatten(args)
-        for a in args:
-            tens.extend(extract_tensors(a))
+            if is_dataclass(arg):
+                return [
+                    attr
+                    for field in dataclasses.fields(arg)
+                    if torch.is_tensor(attr := getattr(arg, field.name))
+                ]
 
-        if kwargs is not None:
-            kwargs, _ = tree_flatten(kwargs)
-            for k in kwargs:
-                tens.extend(extract_tensors(k))
+            return []
 
-        return tuple(tens)
+        if torch.is_tensor(args):
+            return (args,)
+
+        args_tens = [tensor for arg in args for tensor in extract_tensors(arg)] if args else []
+        kwargs_tens = (
+            [tensor for val in kwargs.values() for tensor in extract_tensors(val)] if kwargs else []
+        )
+
+        return tuple(chain(args_tens, kwargs_tens))
+
+    def to_list(self, x):
+        """Helper function to wrap an input into a list"""
+        return [x] if torch.is_tensor(x) else list(x)
 
 
 class CudaGraphManager(torch.nn.Module):
@@ -1103,17 +1370,8 @@ class CudaGraphManager(torch.nn.Module):
     """A global mempool for when 'cuda_graph_use_single_mempool' is used."""
     global_mempool = None
 
-    """Forward pass mempools, used with cudagraph reuse mode."""
-    fwd_mempools = None
-
-    """Backward pass mempool, used with cudagraph reuse mode."""
-    bwd_mempool = None
-
     def __init__(
-        self,
-        config: TransformerConfig,
-        share_cudagraph_io_buffers: bool = True,
-        vp_stage: Optional[int] = None,
+        self, config: TransformerConfig, base_module=None, function_name=None, need_backward=True
     ):
         super().__init__()
         """Creates a CudaGraphManager to manage CUDA graphs for a Megatron module.
@@ -1121,14 +1379,21 @@ def __init__(
         Args:
             config: TransformerConfig object containing CUDA graph settings for memory
                 pooling, graph retention, gradient accumulation, FP8/FP4, and warmup steps.
-            share_cudagraph_io_buffers (bool, optional): (DEPRECATED, will be replaced by
-                config.cuda_graph_share_io_buffers) If None (default) or True, enables
-                buffer reuse optimizations for transformer and mamba layers. If False,
-                disables buffer reuse.
         """
         rng_tracker = get_cuda_rng_tracker()
-        self.share_cudagraph_io_buffers = share_cudagraph_io_buffers
-        self.vp_stage = vp_stage
+        self.need_backward = need_backward
+
+        if function_name is not None:
+            func = getattr(base_module, function_name)
+
+            def wrapped_func(*args, **kwargs):
+                out = self(base_module, args, kwargs)
+                return out
+
+            setattr(base_module, function_name, wrapped_func)
+        else:
+            func = None
+        self.func = func
 
         # need to delay the import here to avoid a circular import
         global HAVE_TE_GRAPHS
@@ -1144,57 +1409,28 @@ def __init__(
         ), "RNG tracker does not support cudagraphs!"
 
         assert config.cuda_graph_impl == "local", "Option cuda_graph_impl=local not enabled."
-        assert (
-            "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "")
-            or os.getenv("NCCL_GRAPH_REGISTER", "") == "0"
-        ), (
-            "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using "
-            "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True."
-        )
+        if torch.cuda.get_device_capability()[0] < 10:
+            assert (
+                "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "")
+                or os.getenv("NCCL_GRAPH_REGISTER", "") == "0"
+            ), (
+                "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using "
+                "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True."
+            )
 
-        self.cudagraph_runners = []
-        self.inference_cudagraphs_lookup_table = defaultdict(lambda: None)
+        self.cudagraph_runners: list[_CudaGraphRunner] = []
+        self.inference_cudagraphs_lookup_table: dict = defaultdict(lambda: None)
         self.is_first_microbatch = False
 
         # Without pipeline parallelism, microbatches execute one at a time.
         # Therefore modules will always execute in the same order, so cudagraphs
         # can both be reused and share a single mempool.
-        if parallel_state.get_pipeline_model_parallel_world_size() == 1:
-            self.reuse_cudagraphs = True
-            self.use_single_mempool = True
-        else:
-            if config.cuda_graph_use_single_mempool:
-                self.reuse_cudagraphs = False
-                self.use_single_mempool = True
-            else:
-                self.reuse_cudagraphs = True
-                self.use_single_mempool = False
-
-        # Mempools are static so that multiple cudagraph managers may share the same mempool
-        if self.use_single_mempool:
-            if CudaGraphManager.global_mempool is None:
-                CudaGraphManager.global_mempool = torch.cuda.graph_pool_handle()
-        else:
-            # All cudagraphs in the same microbatch use the same mempool. For pipeline parallelism,
-            # additonally all bwd passes share the same mempool
-            if CudaGraphManager.fwd_mempools is None:
-                CudaGraphManager.fwd_mempools = defaultdict(
-                    lambda: defaultdict(torch.cuda.graph_pool_handle)
-                )
-                CudaGraphManager.bwd_mempool = torch.cuda.graph_pool_handle()
-
-        # Cudagraph stream capture requires no operations on the default stream prior to the
-        # capture, so change to a side stream.
-        self.stream = torch.cuda.current_stream()
-        torch.cuda.set_stream(torch.cuda.Stream())
-
-    def set_is_first_microbatch(self, is_first_microbatch: bool):
-        """Update the is_first_microbatch flag for weight caching.
-
-        Args:
-            is_first_microbatch (bool): Whether this is the first microbatch in the step.
-        """
-        self.is_first_microbatch = is_first_microbatch
+        self.reuse_cudagraphs = parallel_state.get_pipeline_model_parallel_world_size() == 1
+        if CudaGraphManager.global_mempool is None:
+            CudaGraphManager.global_mempool = torch.cuda.graph_pool_handle()
+            # Cudagraph stream capture requires no operations on the default stream prior to the
+            # capture, so change to a side stream.
+            torch.cuda.set_stream(torch.cuda.Stream())
 
     def call_ddp_preforward_hook(self, module):
         """Call any DDP pre-forward hooks which are used to launch async data parallel
@@ -1211,31 +1447,14 @@ def call_ddp_preforward_hook(self, module):
                 # Only hooks from Mcore DDP, which take no args, should be called at this point.
                 hook(module)
 
-    def get_cudagraph_runner(self, megatron_module, args, kwargs):
+    def get_cudagraph_runner(self, megatron_module, args, kwargs, reuse_cudagraphs):
         '''Returns a valid cudagraph runner for the current forward call.
-        For single mempool mode, we create a cudagraph for each call, if the module is called
-        multiple times per step, for instance in the case of pipeline parallelism.
         The cudagraph corresponding to this call is the first element of 'self.cudagraph_runners'.
         We iterate through the list by 1 for each call, and the number of calls is equal to the
         length of 'self.cudagraph_runners'.
         Otherwise, we assign a mempool per microbatch, which allows cudagraphs to be reused
         over different microbatches by tracking their respective fwd and bwd passes.'''
-
-        if self.use_single_mempool:
-            fwd_mempool = CudaGraphManager.global_mempool
-            bwd_mempool = CudaGraphManager.global_mempool
-        else:
-            if megatron_module.config.virtual_pipeline_model_parallel_size is not None:
-                assert (
-                    self.vp_stage is not None
-                ), "vp_stage must be passed if virtual pipeline is enabled"
-                vpp_rank = self.vp_stage
-            else:
-                vpp_rank = 0
-            fwd_mempool = CudaGraphManager.fwd_mempools[vpp_rank][len(self.cudagraph_runners)]
-            bwd_mempool = CudaGraphManager.bwd_mempool
-
-        if self.reuse_cudagraphs:
+        if reuse_cudagraphs:
             is_inference_mode = 'inference_context' in kwargs.keys() and kwargs['inference_context']
             if is_inference_mode:
                 is_static_batching = kwargs['inference_context'].is_static_batching()
@@ -1248,15 +1467,20 @@ def get_cudagraph_runner(self, megatron_module, args, kwargs):
                     runner = self.inference_cudagraphs_lookup_table[padded_batch_dimensions]
             else:
                 # Todo: For training, we could also cache runners based on input shape.
-                runner = next(
-                    (
-                        r
-                        for r in self.cudagraph_runners
-                        if r.status == _GraphStatus.FWD_READY
+                # If autograd is currently disabled, it doesnt matter if a runner was created
+                # with or without autograd, so just get the first fwd ready runner.
+                require_grad = self.need_backward and torch.is_grad_enabled()
+
+                def is_valid(r):
+                    return (
+                        r.status == _GraphStatus.FWD_READY
                         and not r.get_mismatch_errors(args, kwargs)
-                    ),
-                    None,
-                )
+                        and (not require_grad or r.grad_enabled)
+                    )
+
+                # We must choose the first available runner, as the order of
+                # self.cudagraph_runners corresponds to the capture order.
+                runner = next((r for r in self.cudagraph_runners if is_valid(r)), None)
 
             if runner is None:
                 if _CudagraphGlobalRecord.cudagraph_created:
@@ -1268,11 +1492,11 @@ def get_cudagraph_runner(self, megatron_module, args, kwargs):
                 else:
                     runner = _CudaGraphRunner(
                         megatron_module,
-                        fwd_mempool,
-                        bwd_mempool,
+                        CudaGraphManager.global_mempool,
                         args,
                         kwargs,
-                        self.share_cudagraph_io_buffers,
+                        self.func,
+                        self.need_backward,
                     )
                     self.cudagraph_runners.append(runner)
                     if is_inference_mode:
@@ -1292,11 +1516,11 @@ def get_cudagraph_runner(self, megatron_module, args, kwargs):
             else:
                 runner = _CudaGraphRunner(
                     megatron_module,
-                    fwd_mempool,
-                    bwd_mempool,
+                    CudaGraphManager.global_mempool,
                     args,
                     kwargs,
-                    self.share_cudagraph_io_buffers,
+                    self.func,
+                    self.need_backward,
                 )
                 self.cudagraph_runners.append(runner)
 
@@ -1312,38 +1536,30 @@ def __call__(self, megatron_module, args, kwargs):
 
             kwargs (dict):  The keyword args to be passed to the module.
         """
-        # Set the is_first_microbatch flag on the megatron module if it's the first microbatch
-        if self.is_first_microbatch and hasattr(megatron_module, 'set_is_first_microbatch'):
-            megatron_module.set_is_first_microbatch()
+        is_inference_mode = 'inference_context' in kwargs.keys() and kwargs['inference_context']
+        is_in_checkpoint_fwd = is_checkpointing()
+        if HAVE_TE_GRAPHS:
+            is_in_checkpoint_fwd = is_in_checkpoint_fwd or is_fp8_activation_recompute_enabled()
 
         if _CudagraphGlobalRecord.cudagraph_created:
             if self.training and torch.is_grad_enabled():
-                # param.data_ptr() below is used to trigger any hooks that have attached to the
-                # parameter. Specifically, this is trying to trigger the param sync hook for the
-                # APEX optimizer, which triggers param syncs by hooking into any param references.
-                # However cudagraphs disables this, so we workaround by manually referencing
-                # params here. For more information see:
-                # https://github.com/NVIDIA/apex/blob/7001836/apex/contrib/optimizers/distributed_fused_adam.py#L885C9
-                for param in megatron_module.parameters():
-                    param.data_ptr()
-
                 # Trigger Mcore DDP pre-forward hooks
                 self.call_ddp_preforward_hook(megatron_module)
                 for module in megatron_module.modules():
                     self.call_ddp_preforward_hook(module)
 
-            runner = self.get_cudagraph_runner(megatron_module, args, kwargs)
+            runner = self.get_cudagraph_runner(megatron_module, args, kwargs, self.reuse_cudagraphs)
             out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs)
         else:
-            if 'inference_context' in kwargs.keys() and kwargs['inference_context']:
+            if is_inference_mode:
                 # Inference generation mode creates graphs immediately
-                runner = self.get_cudagraph_runner(megatron_module, args, kwargs)
+                runner = self.get_cudagraph_runner(megatron_module, args, kwargs, True)
                 runner.eval()
 
                 if not runner.fwd_graph_recorded:
                     # Reuse graph input-output buffers for inference
                     local_args, local_kwargs = args, kwargs
-                    if runner.reuse_input_output_buffer and not runner.is_first_layer:
+                    if not runner.is_first_layer:
                         # Find previous layer's runner in the global record
                         try:
                             previous_runner = next(
@@ -1364,10 +1580,9 @@ def __call__(self, megatron_module, args, kwargs):
                             # No match found for previous layer, continue with no buffer reuse
                             pass
 
-                    clone_inputs = not (
-                        runner.reuse_input_output_buffer and not runner.is_first_layer
+                    runner.create_fwd_graph(
+                        local_args, local_kwargs, outputs=None, clone_inputs=runner.is_first_layer
                     )
-                    runner.create_fwd_graph(local_args, local_kwargs, clone_inputs=clone_inputs)
                     runner.fwd_graph_recorded = True
                     runner.cudagraph_created = True
 
@@ -1378,10 +1593,10 @@ def __call__(self, megatron_module, args, kwargs):
 
                 # Now replay the graph
                 out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs)
-
-            elif self.training:
-                # Training mode
-                runner = self.get_cudagraph_runner(megatron_module, args, kwargs)
+            elif self.training or is_in_checkpoint_fwd:
+                runner = self.get_cudagraph_runner(
+                    megatron_module, args, kwargs, self.reuse_cudagraphs
+                )
                 # check if a layer is frozen during training.
                 if not torch.is_grad_enabled():
                     # If the layer is frozen, we need to set the runner to eval mode.
@@ -1390,13 +1605,17 @@ def __call__(self, megatron_module, args, kwargs):
             else:
                 # No cudagraphs were found in training mode with grad disabled, so fallback to
                 # eager since autograd is needed to correctly trace the backward graph.
-                return super(MegatronModule, megatron_module).__call__(*args, **kwargs)
+                if self.func is not None:
+                    return self.func(*args, **kwargs)
+                else:
+                    return super(MegatronModule, megatron_module).__call__(*args, **kwargs)
 
+        self.is_first_microbatch = False
         # If forward only, next replay should be a forward pass as well
-        if self.training and torch.is_grad_enabled():
-            runner.status = _GraphStatus.BWD_READY
-        else:
+        if is_inference_mode or not torch.is_grad_enabled():
             runner.status = _GraphStatus.FWD_READY
+        else:
+            runner.status = _GraphStatus.BWD_READY
 
         return out
 
diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py
index fc849da85c8..c30c107e791 100644
--- a/megatron/core/transformer/module.py
+++ b/megatron/core/transformer/module.py
@@ -169,9 +169,12 @@ def __init__(self, config: TransformerConfig, vp_stage: Optional[int] = None):
 
         # Enable cuda graphs.
         if config.cuda_graph_impl == "local":
-            from megatron.core.transformer.cuda_graphs import CudaGraphManager
+            if hasattr(self, "create_mcore_cudagraph_manager"):
+                self.create_mcore_cudagraph_manager(config)
+            else:
+                from megatron.core.transformer.cuda_graphs import CudaGraphManager
 
-            self.cudagraph_manager = CudaGraphManager(config, vp_stage=vp_stage)
+                self.cudagraph_manager = CudaGraphManager(config)
         elif config.cuda_graph_impl == "transformer_engine":
             # List to store CUDA graphs. A list of `N` CUDA graphs for this layer where N is
             # the number of microbatches. Multiple CUDA graphs per layer is required to support
@@ -336,11 +339,7 @@ def _should_call_te_cudagraph(self, *args, **kwargs):
         )
 
     def __call__(self, *args, **kwargs):
-
         if self._should_call_local_cudagraph(*args, **kwargs):
-            # Set the is_first_microbatch flag for weight caching
-            current_microbatch = getattr(self, 'current_microbatch', 0)
-            self.cudagraph_manager.set_is_first_microbatch(current_microbatch == 0)
             return self.cudagraph_manager(self, args, kwargs)
         elif self._should_call_te_cudagraph(*args, **kwargs):
             if not self.cuda_graphs:
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 2fad0f8e5b7..ef868ebbdb8 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -149,8 +149,11 @@ def __init__(
         super(MoELayer, self).__init__(
             config=config, layer_number=layer_number, pg_collection=pg_collection
         )
+        # If using mcore cudagraphs, recompute is handled by transformer_layer.MoETransformerLayer
         self.moe_layer_recompute = (
-            config.recompute_granularity == 'selective' and "moe" in config.recompute_modules
+            config.recompute_granularity == 'selective'
+            and "moe" in config.recompute_modules
+            and config.cuda_graph_impl != 'local'
         )
         self.shared_experts_recompute = (
             config.recompute_granularity == 'selective'
@@ -237,6 +240,7 @@ def __init__(
 
         # Cudagraph tensor store for resuming the forward pass from the end of the cudagraph.
         self.cudagraph_tensor_store = MoECudaGraphTensorStore()
+        self.fwd_execution_map = ["route", "expert_compute", "postprocess"]
 
     @maybe_skip_or_early_return_by_cudagraph("route")
     def route(self, hidden_states: torch.Tensor):
@@ -322,19 +326,23 @@ def routed_experts_compute(self, hidden_states: torch.Tensor, probs: torch.Tenso
 
         return output, mlp_bias
 
-    def combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Tensor]):
+    def combine(self, output: torch.Tensor):
         """Combines expert outputs via communication and adds shared expert output.
 
         This method uses the token dispatcher to combine the outputs from different
-        experts (e.g., via an All-to-All communication). It then adds the output
-        from the shared expert if it exists.
+        experts (e.g., via an All-to-All communication).
         """
         output = self.token_dispatcher.token_combine(output)
+        return output
+
+    def postprocess(self, output: torch.Tensor, shared_expert_output: Optional[torch.Tensor]):
+        """Project the output back from latent dimension to hidden dimension after combine
+        in latent dimension if needed. Combine expert output with shared_experts if needed."""
+
         output = self.token_dispatcher.combine_postprocess(output)
-        # Project the output back from latent dimension to hidden dimension after combine
-        # in latent dimension.
         if self.config.moe_latent_size:
             output, _ = self.fc2_latent_proj(output)
+
         if shared_expert_output is not None:
             output = output + shared_expert_output
         return output
@@ -346,7 +354,7 @@ def router_and_preprocess(self, hidden_states: torch.Tensor):
         hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map)
         return hidden_states, probs, residual
 
-    def forward(self, hidden_states: torch.Tensor):
+    def forward(self, hidden_states: torch.Tensor, intermediate_tensors=None):
         """Forward pass for the MoE layer.
 
         The forward pass comprises four main steps:
@@ -368,11 +376,16 @@ def forward(self, hidden_states: torch.Tensor):
             )
 
         # MoE forward: route -> dispatch -> compute -> combine
-        def custom_forward(hidden_states):
+        def custom_forward(hidden_states, intermediate_tensors):
             try:
-                shared_expert_output = self.shared_experts_compute(hidden_states)
-                probs, routing_map = self.route(hidden_states)
-                hidden_states, probs = self.preprocess(hidden_states, probs, routing_map)
+                if "route" in self.fwd_execution_map:
+                    shared_expert_output = self.shared_experts_compute(hidden_states)
+                    probs, routing_map = self.route(hidden_states)
+                    hidden_states, probs = self.preprocess(hidden_states, probs, routing_map)
+
+                    if intermediate_tensors is not None:
+                        return hidden_states, probs, shared_expert_output
+
             except MoECudaGraphPartialCaptureSignal as e:
                 # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator.
                 # It means we should early-return from the MoE layer forward pass.
@@ -381,10 +394,28 @@ def custom_forward(hidden_states):
                 # We need to return the intermediate tensors as CUDA graph outputs.
                 return e.get_early_return_outputs(hidden_states, shared_expert_output)
 
-            dispatched_input, probs = self.dispatch(hidden_states, probs)
-            output, mlp_bias = self.routed_experts_compute(dispatched_input, probs)
-            assert mlp_bias is None, f"mlp_bias is not supported for {type(self.token_dispatcher)}"
-            output = self.combine(output, shared_expert_output)
+            if "expert_compute" in self.fwd_execution_map:
+                if intermediate_tensors is not None:
+                    hidden_states, probs = intermediate_tensors
+
+                dispatched_input, probs = self.dispatch(hidden_states, probs)
+                output, mlp_bias = self.routed_experts_compute(dispatched_input, probs)
+                assert (
+                    mlp_bias is None
+                ), f"mlp_bias is not supported for {type(self.token_dispatcher)}"
+                output = self.combine(output)
+
+                if intermediate_tensors is not None:
+                    return output, mlp_bias
+
+            if "postprocess" in self.fwd_execution_map:
+                if intermediate_tensors is not None:
+                    output, shared_expert_output = intermediate_tensors
+
+                output = self.postprocess(output, shared_expert_output)
+
+                if intermediate_tensors is not None:
+                    return output
 
             return output, mlp_bias
 
@@ -400,7 +431,7 @@ def custom_forward(hidden_states):
             else:
                 outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
         else:
-            outputs = custom_forward(hidden_states)
+            outputs = custom_forward(hidden_states, intermediate_tensors)
 
         return outputs
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 5fdeda23dea..bd7c29551a8 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -1387,6 +1387,10 @@ def wrapped_func(moe_layer, *args, **kwargs):
             Otherwise, we execute the original function and check if we should raise a signal to
             early return in CUDA graph capture.
             """
+
+            if moe_layer.config.cuda_graph_impl != "transformer_engine":
+                return func(moe_layer, *args, **kwargs)
+
             # The non-cudagraph codepath just calls the original function.
             if not is_graph_capturing() and moe_layer.cudagraph_tensor_store.is_empty():
                 return func(moe_layer, *args, **kwargs)
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index 8322d44d3bb..f2e26c63cf5 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -436,9 +436,9 @@ def __init__(
             "no_sync": 4,
         }
         self.cuda_dtoh_point = "before_permutation_1"
-        if (
-            config.cuda_graph_impl == "transformer_engine"
-            and CudaGraphScope.moe_preprocess in config.cuda_graph_scope
+        if config.cuda_graph_impl != "none" and (
+            CudaGraphScope.moe_preprocess in config.cuda_graph_scope
+            or not self.config.cuda_graph_scope
         ):
             self.cuda_dtoh_point = "before_ep_alltoall"
         if MoEAlltoAllTokenDispatcher.cuda_dtoh_stream is None:
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index cabad4e15d7..77dc81cfd92 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -1676,55 +1676,59 @@ def __post_init__(self):
                 raise ValueError("CUDA graphs not supported with CPU offloading.")
 
             if self.cuda_graph_impl == "local":
-                assert not self.cuda_graph_scope or self.cuda_graph_scope == [
-                    CudaGraphScope.full_iteration
-                ], (
-                    "For local cuda graph implementation, the only valid value for "
-                    "cuda_graph_scope is full_iteration, or an empty list to denote layerwise "
-                    "graphs. To use other scopes, use cuda_graph_impl=transformer_engine."
-                )
+                # local impl doesn't currently distinguish between moe_preproocess or moe_router
+                # so just set both if either is specified.
+                if (
+                    CudaGraphScope.moe_router in self.cuda_graph_scope
+                    or CudaGraphScope.moe_preprocess in self.cuda_graph_scope
+                ):
+                    if CudaGraphScope.moe_router not in self.cuda_graph_scope:
+                        self.cuda_graph_scope.append(CudaGraphScope.moe_router)
+                    if CudaGraphScope.moe_preprocess not in self.cuda_graph_scope:
+                        self.cuda_graph_scope.append(CudaGraphScope.moe_preprocess)
 
+            # Check cuda graph scopes
             if self.cuda_graph_impl == "transformer_engine":
                 assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, (
                     "To use full iteration cuda graph, please use "
                     "cuda_graph_impl=local instead of cuda_graph_impl=transformer_engine."
                 )
+            assert (
+                CudaGraphScope.moe not in self.cuda_graph_scope
+                or CudaGraphScope.moe_router not in self.cuda_graph_scope
+            ), 'cuda_graph_scope must not contain both moe and moe_router.'
+            if CudaGraphScope.moe_preprocess in self.cuda_graph_scope:
+                assert (
+                    CudaGraphScope.moe_router in self.cuda_graph_scope
+                ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.'
+            if self.num_moe_experts is None or self.num_moe_experts <= 1:
                 assert (
                     CudaGraphScope.moe not in self.cuda_graph_scope
-                    or CudaGraphScope.moe_router not in self.cuda_graph_scope
-                ), 'cuda_graph_scope must not contain both moe and moe_router.'
-                if CudaGraphScope.moe_preprocess in self.cuda_graph_scope:
-                    assert (
-                        CudaGraphScope.moe_router in self.cuda_graph_scope
-                    ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.'
-                if self.num_moe_experts is None or self.num_moe_experts <= 1:
+                    and CudaGraphScope.moe_router not in self.cuda_graph_scope
+                ), 'moe cuda graph is only supported for MoE.'
+            else:
+                if self.moe_layer_freq == 1 or (
+                    isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq
+                ):
+                    assert CudaGraphScope.mlp not in self.cuda_graph_scope, (
+                        'mlp cuda graph is only supported for dense layers, '
+                        'but not found in the model.'
+                    )
+                if (
+                    self.moe_expert_capacity_factor is None
+                    or not self.moe_pad_expert_input_to_capacity
+                ):
                     assert (
                         CudaGraphScope.moe not in self.cuda_graph_scope
-                        and CudaGraphScope.moe_router not in self.cuda_graph_scope
-                    ), 'moe cuda graph is only supported for MoE.'
-                else:
-                    if self.moe_layer_freq == 1 or (
-                        isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq
+                    ), 'moe cuda graph is only supported with drop-padding MoE.'
+                    if self.moe_token_dispatcher_type == 'alltoall' and (
+                        self.moe_expert_capacity_factor is not None
+                        or self.moe_router_padding_for_fp8
                     ):
-                        assert CudaGraphScope.mlp not in self.cuda_graph_scope, (
-                            'mlp cuda graph is only supported for dense layers, '
-                            'but not found in the model.'
+                        assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, (
+                            'moe_preprocess cuda graph is not supported when there are '
+                            'DtoH copies and synchronizations in the preprocess step.'
                         )
-                    if (
-                        self.moe_expert_capacity_factor is None
-                        or not self.moe_pad_expert_input_to_capacity
-                    ):
-                        assert (
-                            CudaGraphScope.moe not in self.cuda_graph_scope
-                        ), 'moe cuda graph is only supported with drop-padding MoE.'
-                        if self.moe_token_dispatcher_type == 'alltoall' and (
-                            self.moe_expert_capacity_factor is not None
-                            or self.moe_router_padding_for_fp8
-                        ):
-                            assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, (
-                                'moe_preprocess cuda graph is not supported when there are '
-                                'DtoH copies and synchronizations in the preprocess step.'
-                            )
 
             if self.recompute_granularity:
                 if self.recompute_granularity != "selective":
@@ -1734,10 +1738,15 @@ def __post_init__(self):
                 else:
                     # The recompute module should be inside or outside of the graph scope.
                     # Recompute module coverring graph scope is not allowed.
-                    if "moe" in self.recompute_modules:
+                    if (
+                        self.cuda_graph_impl == "transformer_engine"
+                        and "moe" in self.recompute_modules
+                    ):
                         assert (
                             CudaGraphScope.moe_router not in self.cuda_graph_scope
-                        ), "moe recompute is not supported with moe_router CUDA graph."
+                        ), "moe recompute is not supported with moe_router CUDA graph with: "
+                        "--cuda-graph-impl transformer_engine."
+
                     # Graphed recompute module doesn't accept random number.
                     if (
                         not self.cuda_graph_scope
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 920c3b8fcba..f575794a819 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -268,6 +268,7 @@ def __init__(
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
     ):
+        self.submodules_config = submodules
         super().__init__(config=config, vp_stage=vp_stage)
 
         if pg_collection is None:
@@ -275,7 +276,6 @@ def __init__(
         self.pg_collection = pg_collection
         self.tp_group = pg_collection.tp
 
-        self.submodules_config = submodules
         self.layer_number = layer_number + get_transformer_layer_offset(
             self.config, vp_stage, get_pg_rank(pg_collection.pp)
         )
@@ -389,6 +389,7 @@ def can_recompute_pre_mlp_layernorm_for_cudagraph():
                     if (
                         not self.is_moe_layer
                         or CudaGraphScope.moe_router not in self.config.cuda_graph_scope
+                        or self.config.cuda_graph_impl == "local"
                     ):
                         # Not a MoE layer, or not capturing the router part.
                         return True
@@ -462,6 +463,27 @@ def can_recompute_pre_mlp_layernorm_for_cudagraph():
         # self.bias_dropout_add_exec_handler = nullcontext if use_nvfuser else torch.enable_grad
         self.bias_dropout_add_exec_handler = torch.enable_grad
 
+    def create_mcore_cudagraph_manager(self, config):
+        """Register the transformer layer for cudagraphs."""
+
+        from megatron.core.transformer.cuda_graphs import CudaGraphManager
+
+        # If full scope, just cudagraph the entire layer
+        if not self.config.cuda_graph_scope:
+            self.cudagraph_manager = CudaGraphManager(config)
+        elif (
+            CudaGraphScope.attn in self.config.cuda_graph_scope
+            and self.submodules_config.self_attention != IdentityOp
+        ):
+            self.cudagraph_manager = CudaGraphManager(config)
+        elif (
+            CudaGraphScope.mlp in self.config.cuda_graph_scope
+            and self.submodules_config.mlp != IdentityOp
+        ):
+            # Cudagraphing MoE layers are supposed handled by MoeTransforerLayer
+            assert not self.is_moe_layer
+            self.cudagraph_manager = CudaGraphManager(config)
+
     @staticmethod
     def _get_layer_offset(config: TransformerConfig):
         """
@@ -635,6 +657,23 @@ def _forward_attention(
 
         return hidden_states, context
 
+    def _forward_pre_mlp_layernorm(self, hidden_states):
+        from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
+            FineGrainedActivationOffloadingInterface as off_interface,
+        )
+
+        if self.recompute_pre_mlp_layernorm:
+            self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states:
+                pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint(
+                    self.pre_mlp_layernorm, hidden_states
+                )
+        else:
+            with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states:
+                pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
+
+        return pre_mlp_layernorm_output
+
     def _forward_mlp(self, hidden_states, inference_context=None):
         """
         Perform a forward pass through the feed-forward layer.
@@ -646,23 +685,11 @@ def _forward_mlp(self, hidden_states, inference_context=None):
             output (Tensor): Transformed hidden states of shape [s, b, h].
         """
 
-        from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
-            FineGrainedActivationOffloadingInterface as off_interface,
-        )
-
         # Residual connection.
         residual = hidden_states
 
         # Optional Layer norm post the cross-attention.
-        if self.recompute_pre_mlp_layernorm:
-            self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput()
-            with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states:
-                pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint(
-                    self.pre_mlp_layernorm, hidden_states
-                )
-        else:
-            with off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") as hidden_states:
-                pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
+        pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)
 
         nvtx_range_push(suffix="mlp")
         # Potentially chunk the MLP computation during prefill to minimize the peak activation size
@@ -714,12 +741,6 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 self._set_fc2_residual(residual)
             mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
 
-        if self.recompute_pre_mlp_layernorm:
-            # discard the output of the pre-mlp layernorm and register the recompute
-            # as a gradient hook of mlp_output_with_bias[0]
-            self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(
-                mlp_output_with_bias[0]
-            )
         nvtx_range_pop(suffix="mlp")
 
         if (
@@ -734,7 +755,7 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be
                 # recomputed in backward pass. For example, the router path and the shared expert
                 # path. So only register in one path is risky.
-                for tensor in mlp_output_with_bias[1:]:
+                for tensor in mlp_output_with_bias:
                     self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor)
             return list(mlp_output_with_bias) + [residual]
         else:
@@ -759,6 +780,13 @@ def _forward_post_mlp(self, mlp_output_with_bias, residual):
             self.config.inference_fuse_tp_communication
         )
 
+        if self.recompute_pre_mlp_layernorm:
+            # discard the output of the pre-mlp layernorm and register the recompute
+            # as a gradient hook of mlp_output_with_bias[0]
+            self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(
+                mlp_output_with_bias[0]
+            )
+
         # TODO: could we move `bias_dropout_add_exec_handler` itself
         # inside the module provided in the `bias_dropout_add_spec` module?
         nvtx_range_push(suffix="mlp_bda")
@@ -1055,7 +1083,12 @@ def _te_cuda_graph_replay(self, *args, **kwargs):
             self.mlp.cudagraph_tensor_store.clear()
             nvtx_range_pop(suffix="mlp")
 
+            # If we early returned, layernorm recompute hooks were attached to the output buffer
+            # of the cudagraph, so disable the recompute hooks inside _forward_post_mlp
+            recompute_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm
+            self.recompute_pre_mlp_layernorm = False
             output = self._forward_post_mlp(mlp_output_with_bias, residual)
+            self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm
         else:
             # If EP overlap is enabled, needs to return same outputs as submodule.attn
             if self.config.overlap_moe_expert_parallel_comm:
@@ -1166,6 +1199,7 @@ def __call__(self, *args, **kwargs):
                 kwargs["dynamic_inference_decode_only"] = kwargs[
                     'inference_context'
                 ].is_decode_only()
+
         return super().__call__(*args, **kwargs)
 
     def get_layer_norm_weights(self):
@@ -1175,3 +1209,153 @@ def get_layer_norm_weights(self):
             List[Tensor]: A list of layernorm weight tensors.
         """
         return
+
+
+class MoETransformerLayer(TransformerLayer):
+    """
+    A Transformer layer specialized for Mixture-of-Experts (MoE) architectures.
+
+    Implements specific functionality to support CUDA graph capture for MoE layers.
+    Due to the dynamic nature of MoE, capturing the entire layer in a single CUDA graph
+    can be challenging. This class supports "partial" CUDA graphs by decomposing the
+    MLP forward pass into router, expert-compute, and post-process stages.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.is_moe_layer = True
+        self.use_partial_cudagraphs = False
+        self.moe_layer_recompute = False
+        self.token_dispatcher_attrs = {}
+
+        super().__init__(*args, **kwargs)
+
+    def create_mcore_cudagraph_manager(self, config):
+        """
+        Initializes the CUDA graph manager(s) for the MoE layer.
+
+        Unlike the standard layer which typically uses a single manager, this method
+        can configure multiple graph managers if partial CUDA graphs are enabled via
+        `cuda_graph_scope`. This allows capturing the static parts of the MoE pass
+        while leaving the expert computation to execute eagerly.
+        """
+
+        from megatron.core.transformer.cuda_graphs import CudaGraphManager
+
+        if not self.config.cuda_graph_scope or CudaGraphScope.moe in self.config.cuda_graph_scope:
+            self.cudagraph_manager = CudaGraphManager(config)
+        elif (
+            CudaGraphScope.moe_router in self.config.cuda_graph_scope
+            or CudaGraphScope.moe_preprocess in self.config.cuda_graph_scope
+        ):
+            # full MoE layer recompute with partial_cudagraphs. If not partial cudagraphs, MoE
+            # layer recompute is handled by the moe_layer.MoELayer class
+            self.moe_layer_recompute = (
+                self.config.recompute_granularity == 'selective'
+                and "moe" in self.config.recompute_modules
+                and self.config.cuda_graph_impl == "local"
+            )
+
+            self.use_partial_cudagraphs = True
+            self.cudagraph_manager_router = CudaGraphManager(
+                self.config, self, function_name="_forward_mlp_router"
+            )
+            self.cudagraph_manager_postprocess = CudaGraphManager(
+                self.config, self, function_name="_forward_mlp_postprocess"
+            )
+
+    def _forward_mlp_router(self, hidden_states):
+        """
+        Executes the router phase of the MoE block.
+
+        This includes the pre-MLP layernorm and the routing logic.
+        This method is isolated so it can be captured by `cudagraph_manager_router`.
+        """
+
+        residual = hidden_states
+        self.mlp.fwd_execution_map = "route"
+        pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)
+        router_outputs = self.mlp(pre_mlp_layernorm_output, intermediate_tensors=())
+
+        for attr_name in self.mlp.token_dispatcher.cudagraph_attrs:
+            attr = getattr(self.mlp.token_dispatcher, attr_name)
+            if torch.is_tensor(attr):
+                if attr_name in self.token_dispatcher_attrs:
+                    self.token_dispatcher_attrs[attr_name].copy_(attr)
+                else:
+                    self.token_dispatcher_attrs[attr_name] = attr.detach()
+
+        return residual, *router_outputs
+
+    def _forward_mlp_expert_compute(self, hidden_states, probs):
+        """
+        Executes the actual computation of the experts.
+
+        This phase takes the routing information and inputs, dispatches them to the
+        appropriate experts, and computes the results. In partial graph modes, this
+        step runs eagerly between the router and postprocess graph replays.
+        """
+
+        for name, attr in self.token_dispatcher_attrs.items():
+            setattr(self.mlp.token_dispatcher, name, attr)
+
+        self.mlp.fwd_execution_map = "expert_compute"
+        return self.mlp(None, intermediate_tensors=(hidden_states, probs))
+
+    def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_bias):
+        """
+        Executes the post-processing phase of the MoE block.
+
+        Handles combining the expert outputs, applying biases, re-registering
+        activation recomputation hooks if necessary, and performing the final
+        Bias-Dropout-Add. This method is isolated so it can be captured by cudagraphs.
+
+        """
+
+        self.mlp.fwd_execution_map = "postprocess"
+        output = self.mlp(None, intermediate_tensors=(output, shared_expert_output))
+        return self._forward_post_mlp((output, mlp_bias), residual)
+
+    def _forward_mlp(self, hidden_states, inference_context=None):
+        """
+        Orchestrates the MLP forward pass, handling partial CUDA graph execution logic.
+
+        If `use_partial_cudagraphs` is True, this method stitches together the
+        router, expert_compute, and postprocess calls.
+        """
+
+        if inference_context is not None:
+            assert not self.use_partial_cudagraphs, (
+                "Partial cudagraphs for MoEs were detected during inference!"
+                "Please do not use --cuda-graph-scope moe_router moe_preprocess "
+                "alongside inference."
+            )
+
+        def _forward_mlp_partial_cudagraphs(hidden_states, inference_context=None):
+            residual, hidden_states, probs, shared_expert_output = self._forward_mlp_router(
+                hidden_states
+            )
+            expert_output, mlp_bias = self._forward_mlp_expert_compute(hidden_states, probs)
+            return self._forward_mlp_postprocess(
+                residual, expert_output, shared_expert_output, mlp_bias
+            )
+
+        if self.use_partial_cudagraphs:
+            if self.moe_layer_recompute:
+                if self.config.fp8 or self.config.fp4:
+                    from megatron.core.extensions.transformer_engine import te_checkpoint
+
+                    return te_checkpoint(
+                        _forward_mlp_partial_cudagraphs,
+                        False,
+                        tensor_parallel.random.get_cuda_rng_tracker,
+                        parallel_state.get_tensor_model_parallel_group(),
+                        hidden_states,
+                    )
+                else:
+                    return tensor_parallel.checkpoint(
+                        _forward_mlp_partial_cudagraphs, False, hidden_states
+                    )
+            else:
+                return _forward_mlp_partial_cudagraphs(hidden_states)
+        else:
+            return super()._forward_mlp(hidden_states)
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 47d0001d6e5..cbd1d31d867 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1352,13 +1352,14 @@ def validate_args(args, defaults={}):
         ):
             args.te_rng_tracker = True
             warn_rank_0("te_rng_tracker is not enabled, enabling it for CUDA graphs.", args.rank)
-        assert (
-            "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "")
-            or os.getenv("NCCL_GRAPH_REGISTER", "") == "0"
-        ), (
-            "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using "
-            "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True."
-        )
+        if args.cuda_graph_impl == "transformer_engine":
+            assert (
+                "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "")
+                or os.getenv("NCCL_GRAPH_REGISTER", "") == "0"
+            ), (
+                "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using "
+                "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True."
+            )
     if args.cuda_graph_scope == "full" or (
         isinstance(args.cuda_graph_scope, list) and "full" in args.cuda_graph_scope
     ):
diff --git a/megatron/training/training.py b/megatron/training/training.py
index b7040e7bbb9..be4f29a3476 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -140,7 +140,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.transformer.moe import upcycling_utils
-from megatron.core.transformer.moe.moe_utils import track_moe_metrics
+from megatron.core.transformer.moe.moe_utils import track_moe_metrics, clear_aux_losses_tracker
 from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper
 from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
 from megatron.core.parallel_state import (
@@ -2953,6 +2953,8 @@ def get_e2e_base_metrics():
             timers('interval-time', log_level=0).start(barrier=True)
             if args.log_energy:
                 energy_monitor.resume()
+            if args.num_experts is not None:
+                clear_aux_losses_tracker()
 
         # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC).
         # Some of these only happen at specific iterations. Capture updated FLOPs accumulator
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/golden_values_dev_dgx_h100.json
index ab11d31f2ca..c7ce9851234 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/golden_values_dev_dgx_h100.json
@@ -283,14 +283,14 @@
   ]
  },
  "throughput": [
-  76.30580996730768,
-  88.09632062440096,
-  88.06043831072262,
-  88.2961798635866,
-  88.30652818803674,
-  88.44774285517468,
-  88.336161355204,
-  88.45930829300391
+ 112.34141028159206,
+ 143.47299899774578,
+ 136.9190123220356,
+ 133.0550750138523,
+ 140.54753942350868,
+ 142.31278267940777,
+ 142.60535677655014,
+ 142.2477862300286
  ],
  "mem-max-allocated-bytes": 23014038016
 }
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/golden_values_dev_dgx_h100.json
index 8e7d12105ac..51437664cf7 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/golden_values_dev_dgx_h100.json
@@ -283,13 +283,13 @@
   ]
  },
  "throughput": [
-  3.446833367136259,
-  69.64151223259532,
-  69.9765204692347,
-  70.25474012041042,
-  69.64760269536946,
-  69.98609501222526,
-  70.21408666363853,
-  70.1614678530764
+ 3.8530017644378955,
+ 98.13326345491087,
+ 100.89859151541394,
+ 100.80208030416277,
+ 100.18034658518215,
+ 100.88831730291241,
+ 100.4922180479951,
+ 101.13060027776349
  ]
 }

From bb42a0081e85b5115f7c264bea78e82c0bc26a93 Mon Sep 17 00:00:00 2001
From: liuyun7345 <51505092+liuyun7345@users.noreply.github.com>
Date: Tue, 27 Jan 2026 00:43:31 +0800
Subject: [PATCH 22/79] fix(fsdp): add CLI argument for
 outer_dp_sharding_strategy (#3053)

Co-authored-by: Jianbin Chang <shjwudp@gmail.com>
---
 megatron/training/arguments.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index cbd1d31d867..2108a1230b9 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2821,6 +2821,13 @@ def _add_distributed_args(parser):
     group.add_argument('--data-parallel-sharding-strategy', type=str, default='no_shard',
                        choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'],
                        help='Sharding strategy of data parallelism.')
+    group.add_argument('--outer-dp-sharding-strategy', type=str, default='no_shard',
+                       choices=['no_shard', 'optim'],
+                       help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. '
+                            'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). '
+                            'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". '
+                            'This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). '
+                            'Default: "no_shard".')
     group.add_argument('--no-gradient-reduce-div-fusion', action='store_false', dest='gradient_reduce_div_fusion',
                        help='If not set, fuse the division in gradient reduce.')
     group.add_argument('--fsdp-double-buffer', action='store_true',

From 94d81860f1037308e6c2a165519db7be28e6f3b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Mon, 26 Jan 2026 20:19:37 +0100
Subject: [PATCH 23/79] ci: Log node name (#3081)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/actions/action.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/actions/action.yml b/.github/actions/action.yml
index a1cb5ddfefc..dfc6d79688e 100644
--- a/.github/actions/action.yml
+++ b/.github/actions/action.yml
@@ -52,6 +52,10 @@ inputs:
 runs:
   using: "composite"
   steps:
+    - name: Print node name
+      shell: bash -x -e -u -o pipefail {0}
+      run: echo "node_name=$NODE_NAME" | tee -a "$GITHUB_OUTPUT"
+
     - name: Checkout repository
       uses: actions/checkout@v2
 

From 23a76d1b35b20afe9b3029a2f3166e63564e9669 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Mon, 26 Jan 2026 22:44:41 +0100
Subject: [PATCH 24/79] docs: Release docs (#3055)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/workflows/release-docs.yml | 74 ++++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)
 create mode 100644 .github/workflows/release-docs.yml

diff --git a/.github/workflows/release-docs.yml b/.github/workflows/release-docs.yml
new file mode 100644
index 00000000000..d15ea74f052
--- /dev/null
+++ b/.github/workflows/release-docs.yml
@@ -0,0 +1,74 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+name: Release docs
+on:
+  workflow_dispatch:
+    inputs:
+      dry-run:
+        description: Whether to run the workflow in dry-run mode
+        required: true
+        type: boolean
+        default: true
+      version-number:
+        description: Version number to release this as (use `latest` for main branch)
+        required: true
+        type: string
+      notify-emails:
+        description: Email addresses to send the notification to. Format as "me@me.com,you@you.com".
+        required: true
+        type: string
+      aws-region:
+        description: AWS region
+        required: false
+        type: string
+        default: us-east-1
+
+jobs:
+  build-docs:
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_build_docs.yml@v0.67.0
+
+  publish-docs:
+    runs-on: ubuntu-latest
+    needs: [build-docs]
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          repository: NVIDIA-NeMo/FW-CI-templates
+          ref: v0.67.2
+          path: FW-CI-templates
+
+      - uses: ./FW-CI-templates/.github/actions/publish-docs
+        # This workflow runs either on main, or on a version tag. Any other git ref will lead
+        # to an error.
+        # If its on main, it will publish to "latest" directory in Akamai.
+        # If its on a versioned tag, it will extract the version number from the tag (strip `v` prefix)
+        # and publish to the versioned directory in Akamai.
+        with:
+          dry-run: ${{ inputs.dry-run }}
+          artifacts-name: docs-html
+          artifacts-path: _build/html
+          emails-csv: ${{ inputs.notify-emails && format('{0},{1}', vars.docs_release_emails, inputs.notify-emails) || vars.docs_release_emails }}
+          overwrite-latest-on-tag: false
+          run-on-version-tag-only: ${{ github.ref_name != 'main' }}
+          request-name: megatron-core-publish-docs-${{ github.run_id }}
+          aws-region: ${{ inputs.aws-region }}
+          aws-role-to-assume: ${{ secrets.AWS_ASSUME_ROLE_ARN }}
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          akamai-host: ${{ secrets.AKAMAI_HOST }}
+          akamai-client-token: ${{ secrets.AKAMAI_CLIENT_TOKEN }}
+          akamai-client-secret: ${{ secrets.AKAMAI_CLIENT_SECRET }}
+          akamai-access-token: ${{ secrets.AKAMAI_ACCESS_TOKEN }}
+          s3-target-root: ${{ secrets.S3_BUCKET_NAME }}
+          s3-target-path: megatron-core/developer-guide

From 35e85a6710d4a9ec6184b4ab6ae13bea8e833933 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Tue, 27 Jan 2026 02:12:26 +0000
Subject: [PATCH 25/79] Reapply "Support multimodule communication (#2031)"
 (#3068)

This reverts commit 0972f020ddf36d0e7583343e664536cb95248f5f.
---
 .../pipeline_parallel/bridge_communicator.py  |   3 -
 .../multimodule_communicator.py               | 531 ++++++++++++
 .../test_multimodule_communicator.py          | 782 ++++++++++++++++++
 3 files changed, 1313 insertions(+), 3 deletions(-)
 create mode 100644 megatron/core/pipeline_parallel/multimodule_communicator.py
 create mode 100644 tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py

diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py
index a67ded6bf08..f1e74a2f16d 100644
--- a/megatron/core/pipeline_parallel/bridge_communicator.py
+++ b/megatron/core/pipeline_parallel/bridge_communicator.py
@@ -628,9 +628,6 @@ def send_forward_recv_backward(
                 dist.broadcast(
                     shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
                 )
-                dist.broadcast(
-                    shape_tensor, src=self.current_rank, group=self.src_grid_broadcast_pg
-                )
 
                 # Broadcast the tensors to all ranks in the group
                 dist.broadcast(
diff --git a/megatron/core/pipeline_parallel/multimodule_communicator.py b/megatron/core/pipeline_parallel/multimodule_communicator.py
new file mode 100644
index 00000000000..1e8da3468e2
--- /dev/null
+++ b/megatron/core/pipeline_parallel/multimodule_communicator.py
@@ -0,0 +1,531 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.distributed as dist
+
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.pipeline_parallel.bridge_communicator import BridgeCommunicator
+from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator
+
+# Types
+Shape = Union[List[int], torch.Size]
+
+
+@dataclass
+class RankModuleInfo:
+    """Information about a rank in a module.
+
+    Attributes:
+        pp_rank: The stage index of the current rank within the module's pipeline.
+        pp_size: The total number of pipeline stages (ranks) in the module.
+        p2p_communicator: Intra-module point-to-point communicator.
+        bridge_comms_as_src_module: Bridge communicators for outgoing connections
+            from this module to downstream modules. One module may have multiple
+            bridge communicators if it has multiple outgoing connections.
+        bridge_comms_as_dest_module: Bridge communicators for incoming connections
+            to this module from upstream modules. One module may have multiple
+            bridge communicators if it has multiple incoming connections.
+        is_source_stage: True if this rank is at the absolute first stage in the
+            overall model (no incoming connections).
+        is_terminal_stage: True if this rank is at the absolute last stage in the
+            overall model (no outgoing connections).
+    """
+
+    pp_rank: int
+    pp_size: int
+    p2p_communicator: Optional[P2PCommunicator]
+    bridge_comms_as_src_module: Optional[List[BridgeCommunicator]]
+    bridge_comms_as_dest_module: Optional[List[BridgeCommunicator]]
+    is_source_stage: Optional[bool] = True
+    is_terminal_stage: Optional[bool] = True
+
+
+class MultiModulePipelineCommunicator:
+    """Communicator for a multi-module pipeline."""
+
+    def __init__(
+        self,
+        module_to_grid_map: Dict[str, HyperCommGrid],
+        topology: Dict[str, List[str]],
+        config: ModelParallelConfig,
+        dim_mapping: Dict[str, List[int]] = None,
+    ):
+        """
+        Initialize the MultiModulePipelineCommunicator.
+
+        Args:
+            module_to_grid_map (dict): A dictionary mapping module names to HyperCommGrids.
+                Example:
+                    module_to_grid_map = {
+                        'image_encoder': image_encoder_grid,
+                        'audio_encoder': audio_encoder_grid,
+                        'llm': llm_grid,
+                        'generator': generator_grid
+                    }
+            topology (dict): A dictionary mapping module names to lists of outgoing modules.
+                Example:
+                    topology = {
+                        'image_encoder': ['llm'],
+                        'audio_encoder': ['llm'],
+                        'llm': ['generator'],
+                        'generator': []
+                    }
+            config (ModelParallelConfig): A ModelParallelConfig object.
+            dim_mapping (Dict[str, List[int]]): Dimension mapping for sequence, batch, hidden.
+                Example:
+                    dim_mapping = {'s': 0, 'h': 2, 'b': 1}
+                Default: None
+        """
+        self.module_to_grid_map = module_to_grid_map
+        self.topology = topology
+        self.config = config
+        self.dim_mapping = dim_mapping
+        self.current_rank = dist.get_rank()
+
+        # Build bridge communicators for all modules
+        self.bridge_comms = []
+        self._build_bridge_comms()
+
+        self.rank_module_map = {}
+        self._build_rank_module_info_map()
+
+    def _build_bridge_comms(self):
+        """Construct and store BridgeCommunicator objects that describe the outgoing
+        communication relationships for all of the modules.
+        """
+        for src_module_name, src_grid in self.module_to_grid_map.items():
+            for dest_module_name in self.topology[src_module_name]:
+                dest_grid = self.module_to_grid_map[dest_module_name]
+                bridge_comm = BridgeCommunicator(
+                    src_grid=src_grid,
+                    dest_grid=dest_grid,
+                    dim_mapping=self.dim_mapping,
+                    comm_dtype=self.config.pipeline_dtype,
+                    src_module_name=src_module_name,
+                    dest_module_name=dest_module_name,
+                )
+                self.bridge_comms.append(bridge_comm)
+
+    @property
+    def is_pp_first_stage(self):
+        """Return True if the current rank has the absolute first stage in the overall model.
+
+        The absolute first stage is defined as:
+        1. The current rank must be in the first PP stage (pp_rank == 0) of some module
+        2. That module must be a source module (no incoming connections in topology)
+        """
+        for module_name, rank_module_info in self.rank_module_map.items():
+            # Check if this rank is at the first PP stage of this module
+            if rank_module_info.pp_rank == 0:
+                # Check if this module is a source module (no incoming connections)
+                if self._is_source_module(module_name):
+                    return True
+        return False
+
+    @property
+    def is_pp_last_stage(self):
+        """Return True if the current rank has the absolute last stage in the overall model.
+
+        The absolute last stage is defined as:
+        1. The current rank must be in the last PP stage of some module
+        2. That module must be a sink module (no outgoing connections in topology)
+        """
+        for module_name, rank_module_info in self.rank_module_map.items():
+            # Check if this rank is at the last PP stage of this module
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # Check if this module is a sink module (no outgoing connections)
+                if self._is_sink_module(module_name):
+                    return True
+        return False
+
+    def _is_source_module(self, module_name: str) -> bool:
+        """Check if a module is a source module (has no incoming connections)."""
+        # A module is a source if no other module lists it as a destination
+        for src_module, dest_modules in self.topology.items():
+            if module_name in dest_modules:
+                return False
+        return True
+
+    def _is_sink_module(self, module_name: str) -> bool:
+        """Check if a module is a sink module (has no outgoing connections)."""
+        return len(self.topology.get(module_name, [])) == 0
+
+    def is_current_rank_in_grid(self, grid: HyperCommGrid) -> bool:
+        """Check if the current rank is in the grid."""
+        return grid.rank_offset <= self.current_rank < grid.rank_offset + grid.size
+
+    @property
+    def num_warmup_microbatches(self):
+        """Calculate the number of warmup microbatches for the current rank.
+
+        Uses the same simple logic as P2PCommunicator:
+        total_pipeline_stages - current_rank_stage - 1
+
+        Returns:
+            int: Number of warmup microbatches for this rank
+        """
+        # Get total pipeline depth across all modules
+        total_stages = self.compute_total_pipeline_stages(self.topology, self.module_to_grid_map)
+
+        # Get current rank's position in the overall pipeline (0-indexed)
+        # Use compute_total_pipeline_stages with current rank to get cumulative position
+        if self.rank_module_map:
+            # Take the first module this rank belongs to
+            # TODO: ykarnati - improve this logic.
+            module_name = next(iter(self.rank_module_map.keys()))
+            current_stage = (
+                self.compute_total_pipeline_stages(
+                    self.topology,
+                    self.module_to_grid_map,
+                    rank=self.current_rank,
+                    module_name=module_name,
+                )
+                - 1
+            )  # Convert from 1-indexed to 0-indexed
+        else:
+            current_stage = 0
+
+        assert (
+            current_stage <= total_stages
+        ), f"current_stage: {current_stage} is greater than total_stages: {total_stages}"
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"current_stage: {current_stage} total_stages: {total_stages} "
+            f"num_warmup_microbatches: {total_stages - current_stage - 1}"
+        )
+        return total_stages - current_stage - 1
+
+    def _build_rank_module_info_map(self):
+        """For each module in the current rank, initialize the P2P communicator
+        and build the bridge communicator info for the module.
+        Each rank may hold multiple modules when colocated.
+        """
+        for module_name, module_grid in self.module_to_grid_map.items():
+            if self.is_current_rank_in_grid(module_grid):
+                # Initialize P2P communicator
+                pp_group = module_grid.get_pg('pp')
+                p2p_comm = P2PCommunicator(pp_group, self.config)
+                pp_size = dist.get_world_size(pp_group)
+                rank_in_pp_group = dist.get_group_rank(pp_group, self.current_rank)
+                pp_rank = rank_in_pp_group % pp_size
+
+                bridge_comms_as_dest_module = []
+                bridge_comms_as_src_module = []
+                # If first stage, check if the module has any incoming modules
+                # If so, initialize bridge communicator
+                if pp_rank == 0:
+                    for bridge_comm in self.bridge_comms:
+                        if (
+                            bridge_comm.is_current_rank_in_grid(bridge_comm.dest_grid)
+                            and bridge_comm.dest_module_name == module_name
+                        ):
+                            bridge_comms_as_dest_module.append(bridge_comm)
+                # If last stage, check if the module has any outgoing modules
+                # If so, initialize bridge communicator
+                if pp_rank == pp_size - 1:
+                    for bridge_comm in self.bridge_comms:
+                        if (
+                            bridge_comm.is_current_rank_in_grid(bridge_comm.src_grid)
+                            and bridge_comm.src_module_name == module_name
+                        ):
+                            bridge_comms_as_src_module.append(bridge_comm)
+                # Build RankModuleInfo for the module
+                rank_module_info = RankModuleInfo(
+                    pp_rank=pp_rank,
+                    pp_size=pp_size,
+                    p2p_communicator=p2p_comm,
+                    bridge_comms_as_dest_module=bridge_comms_as_dest_module,
+                    bridge_comms_as_src_module=bridge_comms_as_src_module,
+                )
+                self.rank_module_map[module_name] = rank_module_info
+
+    def recv_forward(
+        self, tensor_shape: Optional[Shape] = None, is_first_stage: bool = False
+    ) -> Dict[str, torch.Tensor]:
+        """Receive forward activation tensor.
+
+        Args:
+            tensor_shape: Expected activation tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[receive_forward] tensors_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
+        )
+        input_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+
+            if rank_module_info.pp_rank == 0:
+                # If first stage, and has incoming modules, receive forward activation
+                # from incoming modules.
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    input_dict[bridge_comm.src_module_name] = bridge_comm.recv_forward()
+            else:
+                # If not first stage, receive forward activation tensor from P2P communicator.
+                input_dict[module_name] = rank_module_info.p2p_communicator.recv_forward(
+                    tensor_shapes=tensor_shape, is_first_stage=False
+                )
+        return input_dict
+
+    def send_forward(self, output_dict: Dict[str, torch.Tensor], is_last_stage: bool = False):
+        """Send forward activation tensor.
+
+        Args:
+            output_dict: A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_forward] output_dict keys: {output_dict.keys()}, is_last_stage: {is_last_stage}"
+        )
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has outgoing modules, send forward activation
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    bridge_comm.send_forward(output_dict[module_name])
+            else:
+                # If not last stage, send forward activation by using P2P communicator.
+                rank_module_info.p2p_communicator.send_forward(
+                    output_dict[module_name], is_last_stage=False
+                )
+
+    def send_forward_recv_backward(
+        self,
+        output_dict: Dict[str, torch.Tensor],
+        tensor_shape: Optional[Shape] = None,
+        is_last_stage: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Send forward activation tensor and receive backward activation tensor.
+
+        Args:
+            output_dict: A dictionary mapping module names to tensors.
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_forward_recv_backward] output_dict keys: {output_dict.keys()}, "
+            f"tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
+        )
+        grad_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has outgoing modules, send forward activation and
+                # receive backward gradient by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    grad_dict[bridge_comm.src_module_name] = bridge_comm.send_forward_recv_backward(
+                        output_dict[module_name]
+                    )
+            else:
+                # If not last stage, send forward activation and receive backward gradient
+                # by using P2P communicator.
+                grad_dict[module_name] = (
+                    rank_module_info.p2p_communicator.send_forward_recv_backward(
+                        output_dict[module_name], tensor_shapes=tensor_shape, is_last_stage=False
+                    )
+                )
+        return grad_dict
+
+    def send_backward_recv_forward(
+        self,
+        grad_dict: Dict[str, torch.Tensor],
+        tensor_shape: Optional[Shape] = None,
+        is_first_stage: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Send backward activation tensor and receive forward activation tensor.
+
+        Args:
+            grad_dict: A dictionary mapping module names to tensors.
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_backward_recv_forward] grad_dict keys: {grad_dict.keys()}, "
+            f"tensor_shape: {tensor_shape}, is_first_stage: {is_first_stage}"
+        )
+        input_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == 0:
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    # If first stage, and has incoming modules, send backward gradient and
+                    # receive forward activation by using bridge communicator.
+                    input_dict[bridge_comm.src_module_name] = (
+                        bridge_comm.send_backward_recv_forward(
+                            grad_dict[bridge_comm.src_module_name]
+                        )
+                    )
+            else:
+                # If not first stage, send backward gradient and receive forward activation
+                # by using P2P communicator.
+                input_dict[module_name] = (
+                    rank_module_info.p2p_communicator.send_backward_recv_forward(
+                        grad_dict[module_name], tensor_shapes=tensor_shape, is_first_stage=False
+                    )
+                )
+        return input_dict
+
+    def recv_backward(
+        self, tensor_shape: Optional[Shape] = None, is_last_stage: bool = False
+    ) -> Dict[str, torch.Tensor]:
+        """Receive backward activation tensor.
+
+        Args:
+            tensor_shape: Expected gradient tensor shape
+
+        Returns:
+            A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[recv_backward] tensor_shape: {tensor_shape}, is_last_stage: {is_last_stage}"
+        )
+        grad_dict = {}
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                # If last stage, and has incoming modules, receive backward gradient
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_src_module:
+                    grad_dict[bridge_comm.src_module_name] = bridge_comm.recv_backward()
+            else:
+                # If not last stage, receive backward gradient by using P2P communicator.
+                grad_dict[module_name] = rank_module_info.p2p_communicator.recv_backward(
+                    tensor_shapes=tensor_shape, is_last_stage=False
+                )
+        return grad_dict
+
+    def send_backward(self, grad_dict: Dict[str, torch.Tensor], is_first_stage: bool = False):
+        """Send backward activation tensor.
+
+        Args:
+            grad_dict: A dictionary mapping module names to tensors.
+        """
+        logging.debug(
+            f"[Rank {dist.get_rank()} ][MultiModulePipelineCommunicator] "
+            f"[send_backward] grad_dict keys: {grad_dict.keys()}, is_first_stage: {is_first_stage}"
+        )
+        for module_name, rank_module_info in self.rank_module_map.items():
+            if rank_module_info.pp_rank == 0:
+                # If first stage, and has incoming modules, send backward activation
+                # by using bridge communicator.
+                for bridge_comm in rank_module_info.bridge_comms_as_dest_module:
+                    bridge_comm.send_backward(grad_dict[bridge_comm.src_module_name])
+            else:
+                # If not first stage, send backward activation by using P2P communicator.
+                rank_module_info.p2p_communicator.send_backward(
+                    grad_dict[module_name], is_first_stage=False
+                )
+
+    @staticmethod
+    def compute_total_pipeline_stages(
+        topology: Dict[str, List[str]],
+        module_to_grid_map: Dict[str, HyperCommGrid],
+        rank: Optional[int] = None,
+        module_name: Optional[str] = None,
+    ) -> int:
+        """Compute the total number of pipeline stages across a multi-module chain.
+
+        Interprets ``topology`` as a directed acyclic graph (DAG) where nodes are modules
+        and edges indicate forward data flow from source to destination modules. Each node
+        is assigned a weight equal to its pipeline parallel size (number of PP stages).
+
+        The total number of stages is defined as the length of the longest path in this DAG
+        under node weights.
+
+        If ``rank`` is None (default), returns the maximum over all terminal (sink) modules of
+        the sum of PP sizes along a path ending at that terminal. For example, given:
+
+            image_encoder ->\
+                              -> llm -> generator
+            audio_encoder  ->/
+
+        the total is: max(pp(image_encoder), pp(audio_encoder)) + pp(llm) + pp(generator).
+
+        If ``rank`` is provided, the result is the total number of pipeline stages up to (and
+        including) the PP stage that ``rank`` occupies inside its module. In this case, the
+        weight of the target module equals (pp_rank_index(rank) + 1) instead of the module's
+        full PP size; other modules still contribute their full PP sizes. If the rank belongs to
+        multiple modules (colocation), pass ``module_name`` to disambiguate; otherwise the
+        maximum across all candidate modules containing the rank is returned.
+
+        Args:
+            topology: Mapping from a module to its list of outgoing modules.
+            module_to_grid_map: Mapping from module name to its ``HyperCommGrid``.
+
+        Returns:
+            The total number of pipeline stages along the longest path given the constraints.
+
+        Raises:
+            ValueError: If the topology contains cycles; or has no terminal nodes when
+                ``rank`` is None
+        """
+        nodes = set(module_to_grid_map.keys())
+        # Build adjacency and reverse-adjacency (predecessors).
+        adj: Dict[str, List[str]] = {node: list(topology.get(node, [])) for node in nodes}
+        preds: Dict[str, List[str]] = {node: [] for node in nodes}
+        for src, outs in adj.items():
+            for dst in outs:
+                preds[dst].append(src)
+
+        # Identify terminal nodes (no outgoing edges) for the rank=None case.
+        sinks = [node for node, outs in adj.items() if not outs]
+        if rank is None and not sinks:
+            raise ValueError(
+                "Topology must be a DAG with at least one terminal (no outgoing) module."
+            )
+
+        def pp_size(name: str) -> int:
+            grid = module_to_grid_map[name]
+            pp_dim_index = grid.dim_names.index('pp')
+            return grid.shape[pp_dim_index]
+
+        def partial_weight_for_target(target: str) -> Optional[int]:
+            if rank is None:
+                return None
+            grid = module_to_grid_map.get(target)
+            rank_groups = grid._gen_rank_enum(['pp'])
+            stage_index: Optional[int] = None
+            for group in rank_groups:
+                if rank in group:
+                    stage_index = group.index(rank)
+                    break
+            return stage_index + 1
+
+        def longest_path_to(target: str) -> int:
+            visiting = set()
+            partial = partial_weight_for_target(target)
+
+            def weight(name: str) -> int:
+                if partial is not None and name == target:
+                    return partial
+                return pp_size(name)
+
+            def dfs(node: str) -> int:
+                if node in visiting:
+                    raise ValueError("Topology contains cycles; expected a DAG.")
+                visiting.add(node)
+                best = 0
+                for p in preds.get(node, []):
+                    val = dfs(p)
+                    if val > best:
+                        best = val
+                visiting.remove(node)
+                return weight(node) + best
+
+            return dfs(target)
+
+        if rank is None:
+            return max(longest_path_to(sink) for sink in sinks)
+
+        return longest_path_to(module_name)
diff --git a/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
new file mode 100644
index 00000000000..22f790cc0a9
--- /dev/null
+++ b/tests/unit_tests/pipeline_parallel/test_multimodule_communicator.py
@@ -0,0 +1,782 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import os
+import sys
+
+import pytest
+import torch
+import torch.distributed as dist
+from packaging import version
+
+from megatron.core import parallel_state
+from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.pipeline_parallel.multimodule_communicator import MultiModulePipelineCommunicator
+from tests.unit_tests.pipeline_parallel.test_bridge_communicator import (
+    _avg_params,
+    _create_transformer_block,
+    _get_pg_collection_from_grid,
+    create_hypercomm_grid,
+    get_transformer_block_and_grid,
+)
+from tests.unit_tests.test_utilities import Utils
+
+
+class TestMultiModulePipelineCommunicator:
+
+    @classmethod
+    def setup_class(cls):
+        """Set up distributed environment for the entire test class."""
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        if torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+        world_size = dist.get_world_size()
+        if world_size != 8:
+            pytest.skip(
+                f"These tests require 8 GPUs, but only {world_size} are available.",
+                allow_module_level=True,
+            )
+
+    def teardown_class(cls):
+        Utils.destroy_model_parallel()
+
+    def test_multimodule_communicator_init(self):
+        """Test MultiModulePipelineCommunicator initialization."""
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=2, cp=1, pp=1, dp=1)
+
+        # Define module-grid mapping
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        # Define module computation topology
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(bf16=True)
+        # Initialize communicator
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+        # Test attributes match expectations
+        assert mllm_comm.module_to_grid_map == module_to_grid_map
+        assert mllm_comm.topology == topology
+        assert mllm_comm.config == config
+        assert mllm_comm.current_rank == dist.get_rank()
+
+    def test_compute_total_pipeline_stages(self):
+        """Test compute_total_pipeline_stages for overall chain and until specific ranks."""
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Define module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+
+        # Overall total pipeline stages: max(1,1) + 2 + 1 = 4
+        total = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
+            topology, module_to_grid_map
+        )
+        assert total == 4
+
+        llm_pp_rank = MultiModulePipelineCommunicator.compute_total_pipeline_stages(
+            topology, module_to_grid_map, rank=2, module_name='llm'
+        )
+        assert llm_pp_rank == 2
+
+    def test_send_forward_recv_forward(self):
+        """Test send_forward and recv_forward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate forward communication for each module
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder sends output forward
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder sends output forward
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                # LLM stage receives both image and audio outputs
+                input_dict = mllm_comm.recv_forward()
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+                mllm_comm.send_forward(output_dict)
+            else:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+                mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator module receives final LLM output
+            input_dict = mllm_comm.recv_forward()
+            assert input_dict['llm'].shape == (1, 32, 128)
+
+    def test_send_forward_recv_forward_with_different_pp_size(self):
+        """Test for the case when pp(image_encoder) != pp(audio_encoder)."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=2, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=4, tp=1, cp=1, pp=4, dp=1)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+        }
+        topology = {'image_encoder': ['llm'], 'audio_encoder': ['llm'], 'llm': []}
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate forward communication for each module
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            if dist.get_rank() == 0:
+                # Image encoder sends output forward
+                mllm_comm.send_forward(output_dict)
+            else:
+                # Image stage receives image outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 8, 128))
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder sends output forward
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+            if dist.get_rank() == 4:
+                # LLM stage receives both image and audio outputs
+                input_dict = mllm_comm.recv_forward()
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+                mllm_comm.send_forward(output_dict)
+            elif dist.get_rank() == 5 or dist.get_rank() == 6:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+                mllm_comm.send_forward(output_dict)
+            elif dist.get_rank() == 7:
+                # LLM stage receives concatenated LLM outputs
+                input_dict = mllm_comm.recv_forward(tensor_shape=(2, 32, 128))
+                assert input_dict['llm'].shape == (2, 32, 128)
+
+    def test_send_backward_recv_backward(self):
+        """Test send_backward and recv_backward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate backward communication for each module
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator sends gradient backward
+            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
+            mllm_comm.send_backward(grad_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                # LLM receives expanded gradient and sends backward
+                received_grad = mllm_comm.recv_backward()
+                assert received_grad['llm'].shape == (2, 32, 128)
+                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                mllm_comm.send_backward(grad_dict)
+            else:
+                # LLM receives gradient and sends backward to both image/audio encoders
+                received_grad = mllm_comm.recv_backward(tensor_shape=(2, 32, 128))
+                assert received_grad['llm'].shape == (2, 32, 128)
+                grad_dict = {
+                    'image_encoder': torch.randn(2, 8, 128).cuda(),
+                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
+                }
+                mllm_comm.send_backward(grad_dict)
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder receives its gradient
+            received_grad = mllm_comm.recv_backward()
+            assert received_grad['image_encoder'].shape == (2, 8, 128)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder receives its gradient
+            received_grad = mllm_comm.recv_backward()
+            assert received_grad['audio_encoder'].shape == (2, 16, 128)
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    def test_send_forward_recv_backward_send_backward_recv_forward(self):
+        """Test send_forward_recv_backward and send_backward_recv_forward operations."""
+        if not dist.is_initialized():
+            pytest.skip("Distributed not initialized")
+
+        # Create process group grids for each module
+        image_encoder_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1)
+        audio_encoder_grid = create_hypercomm_grid(offset=1, tp=1, cp=1, pp=1, dp=1)
+        llm_grid = create_hypercomm_grid(offset=2, tp=2, cp=1, pp=2, dp=1)
+        generator_grid = create_hypercomm_grid(offset=6, tp=1, cp=1, pp=1, dp=2)
+
+        # Set up module-grid mapping and topology
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],
+            'audio_encoder': ['llm'],
+            'llm': ['generator'],
+            'generator': [],
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(module_to_grid_map, topology, config)
+
+        # Simulate bidirectional send/recv for forward and backward in pipeline
+
+        # Encoder stages send forward to the first stage of LLM, and receive backward from the first stage of LLM
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            output_dict = {'image_encoder': torch.randn(2, 8, 128).cuda()}
+            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+            assert received_grad['image_encoder'].shape == (2, 8, 128)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            output_dict = {'audio_encoder': torch.randn(2, 16, 128).cuda()}
+            received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+            assert received_grad['audio_encoder'].shape == (2, 16, 128)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                grad_dict = {
+                    'image_encoder': torch.randn(2, 8, 128).cuda(),
+                    'audio_encoder': torch.randn(2, 16, 128).cuda(),
+                }
+                input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
+                assert input_dict['image_encoder'].shape == (2, 8, 128)
+                assert input_dict['audio_encoder'].shape == (2, 16, 128)
+
+        # First stage of LLM sends forward to the second stage of LLM, and receive backward from the second stage of LLM
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                received_grad = mllm_comm.send_forward_recv_backward(
+                    output_dict, tensor_shape=(2, 32, 128)
+                )
+                assert received_grad['llm'].shape == (2, 32, 128)
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                grad_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                input_dict = mllm_comm.send_backward_recv_forward(
+                    grad_dict, tensor_shape=(2, 32, 128)
+                )
+                assert input_dict['llm'].shape == (2, 32, 128)
+
+        # Second stage of LLM sends forward to generator, and receive backward from generator
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 4 or dist.get_rank() == 5:
+                output_dict = {'llm': torch.randn(2, 32, 128).cuda()}
+                received_grad = mllm_comm.send_forward_recv_backward(output_dict)
+                assert received_grad['llm'].shape == (2, 32, 128)
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            grad_dict = {'llm': torch.randn(1, 32, 128).cuda()}
+            input_dict = mllm_comm.send_backward_recv_forward(grad_dict)
+            assert input_dict['llm'].shape == (1, 32, 128)
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    def test_send_forward_recv_forward_with_transformer_blocks(self):
+        """Test send_forward and recv_forward operations."""
+
+        # Set model/test dimensions for easier debugging and output comparison
+        hidden_size = 16
+        sequence_length = 2
+        micro_batch_size = 2
+
+        # For reproducibility, set a fixed seed
+        torch.manual_seed(12345)
+        dtype = torch.float32
+
+        # Create random input hidden states tensor
+        hidden_states = torch.randn(
+            (sequence_length, micro_batch_size, hidden_size), device="cuda"
+        ).to(dtype)
+        current_rank = dist.get_rank()
+
+        # ========== Initialize tensor model-parallel environment ==========
+        parallel_state_tp = 2
+        Utils.initialize_model_parallel(tensor_model_parallel_size=2)
+
+        # ========== Build reference 1D grid and transformer block for weight sharing ==========
+        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
+        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
+        ref_block = _create_transformer_block(
+            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
+        )
+        _avg_params(
+            ref_block, ref_grid.get_pg("dp")
+        )  # Ensure parameters are averaged across data parallel (DP)
+
+        # ========== Create different transformer blocks for each model stage ==========
+        # Image encoder
+        image_encoder_block, image_encoder_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=1,
+            grid_offset=0,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # Audio encoder
+        audio_encoder_block, audio_encoder_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=1,
+            grid_offset=1,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # LLM (Large Language Model) block with tensor & pipeline parallelism
+        llm_block, llm_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=2,
+            cp_size=1,
+            pp_size=2,
+            dp_size=1,
+            grid_offset=2,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        # Generator block (final stage) with DP=2
+        generator_block, generator_grid = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=1,
+            cp_size=1,
+            pp_size=1,
+            dp_size=2,
+            grid_offset=6,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        # ========== Define module-to-grid correspondence and pipeline topology ==========
+        module_to_grid_map = {
+            'image_encoder': image_encoder_grid,
+            'audio_encoder': audio_encoder_grid,
+            'llm': llm_grid,
+            'generator': generator_grid,
+        }
+        topology = {
+            'image_encoder': ['llm'],  # image_encoder sends output to llm
+            'audio_encoder': ['llm'],  # audio_encoder sends output to llm
+            'llm': ['generator'],  # llm sends output to generator
+            'generator': [],  # generator is the final module
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        # Define dimension mapping for sequence, batch, hidden
+        dim_mapping = {'s': 0, 'h': 2, 'b': 1}
+        seq_dim = dim_mapping['s']
+
+        # Communication handler for multi-module pipeline (send/recv abstraction)
+        mllm_comm = MultiModulePipelineCommunicator(
+            module_to_grid_map, topology, config, dim_mapping=dim_mapping
+        )
+
+        # ========== Run actual distributed pipeline blocks (per process, depending on role) ==========
+        if mllm_comm.is_current_rank_in_grid(image_encoder_grid):
+            # Image encoder rank: run forward and send output
+            image_encoder_output = image_encoder_block(
+                hidden_states=hidden_states, attention_mask=None
+            )
+            output_dict = {'image_encoder': image_encoder_output}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(audio_encoder_grid):
+            # Audio encoder rank: run forward and send output
+            audio_encoder_output = audio_encoder_block(
+                hidden_states=hidden_states, attention_mask=None
+            )
+            output_dict = {'audio_encoder': audio_encoder_output}
+            mllm_comm.send_forward(output_dict)
+        if mllm_comm.is_current_rank_in_grid(llm_grid):
+            if dist.get_rank() == 2 or dist.get_rank() == 3:
+                # LLM stage 0 (receives both image and audio, concatenates along seq_dim)
+                input_dict = mllm_comm.recv_forward()
+                llm_output = llm_block(
+                    hidden_states=torch.cat(
+                        [input_dict['image_encoder'], input_dict['audio_encoder']], dim=seq_dim
+                    ),
+                    attention_mask=None,
+                )
+                output_dict = {'llm': llm_output}
+                mllm_comm.send_forward(output_dict)
+            else:
+                # LLM stage 1 (receives output of previous LLM stage)
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(sequence_length * 2, micro_batch_size, hidden_size)
+                )
+                llm_output = llm_block(hidden_states=input_dict['llm'], attention_mask=None)
+                output_dict = {'llm': llm_output}
+                mllm_comm.send_forward(output_dict)
+
+        if mllm_comm.is_current_rank_in_grid(generator_grid):
+            # Generator block: only receives from llm and runs forward
+            input_dict = mllm_comm.recv_forward()
+            generator_output = generator_block(hidden_states=input_dict['llm'], attention_mask=None)
+
+        # ========== Build a reference (serial/global) pipeline for correctness checking ==========
+        global_image_encoder_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_audio_encoder_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_llm_block_pp_rank_0, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_llm_block_pp_rank_1, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_generator_block, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        # Run each stage sequentially as a global pipeline (for truth)
+        global_image_encoder_output = global_image_encoder_block(
+            hidden_states=hidden_states, attention_mask=None
+        )
+        global_audio_encoder_output = global_audio_encoder_block(
+            hidden_states=hidden_states, attention_mask=None
+        )
+        # Compare output between global and distributed blocks for image/audio stage
+        if current_rank == 0:
+            torch.testing.assert_close(
+                global_image_encoder_output, image_encoder_output, rtol=1e-3, atol=1e-3
+            )
+        if current_rank == 1:
+            torch.testing.assert_close(
+                global_audio_encoder_output, audio_encoder_output, rtol=1e-3, atol=1e-3
+            )
+
+        # Feed outputs to LLM stages (emulate pipeline cut with concatenation)
+        global_llm_input = torch.cat(
+            [global_image_encoder_output, global_audio_encoder_output], dim=seq_dim
+        )
+        global_llm_pp_rank_0_output = global_llm_block_pp_rank_0(
+            hidden_states=global_llm_input, attention_mask=None
+        )
+        if current_rank == 2 or current_rank == 3:
+            torch.testing.assert_close(
+                global_llm_pp_rank_0_output, llm_output, rtol=1e-3, atol=1e-3
+            )
+        global_llm_pp_rank_1_output = global_llm_block_pp_rank_1(
+            hidden_states=global_llm_pp_rank_0_output, attention_mask=None
+        )
+        if current_rank == 4 or current_rank == 5:
+            torch.testing.assert_close(
+                global_llm_pp_rank_1_output, llm_output, rtol=1e-3, atol=1e-3
+            )
+
+        # Generator output and comparison to distributed output (for each DP chunk)
+        global_generator_block_output = global_generator_block(
+            hidden_states=global_llm_pp_rank_1_output, attention_mask=None
+        )
+        global_generator_block_chunks = torch.split(
+            global_generator_block_output, global_generator_block_output.shape[1] // 2, dim=1
+        )
+        if current_rank == 6:
+            torch.testing.assert_close(
+                global_generator_block_chunks[0], generator_output, rtol=1e-3, atol=1e-3
+            )
+        if current_rank == 7:
+            torch.testing.assert_close(
+                global_generator_block_chunks[1], generator_output, rtol=1e-3, atol=1e-3
+            )
+
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse('2.3.0'),
+        reason="Feature requires PyTorch 2.3 or later",
+    )
+    @pytest.mark.parametrize(
+        "grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp",
+        [
+            (2, 1, 1, 2, 1, 1, 2),  # TP2PP1DP1 to TP2PP1DP1
+            (2, 1, 1, 2, 2, 1, 2),  # TP2PP1DP1 to TP2PP2DP1
+            (2, 2, 1, 2, 2, 1, 2),  # TP2PP2DP1 to TP2PP2DP1
+            (4, 1, 1, 4, 1, 1, 4),  # TP4DP1 to TP4DP1
+            (2, 1, 2, 4, 1, 1, 2),  # TP2DP2 to TP4DP1
+            (4, 1, 1, 2, 1, 2, 2),  # TP4DP1 to TP2DP2
+            (2, 1, 2, 1, 1, 4, 2),  # TP2DP2 to TP1DP4
+        ],
+    )
+    def test_send_forward_recv_forward_with_transformer_blocks_and_different_parallelisms(
+        self, grid1_tp, grid1_pp, grid1_dp, grid2_tp, grid2_pp, grid2_dp, parallel_state_tp
+    ):
+        """Test bridge communicator with two transformer blocks having different process group configurations."""
+        # Model and input configuration
+        hidden_size = 16
+        sequence_length = 2
+        micro_batch_size = 8
+        torch.manual_seed(12345)
+        dtype = torch.float32
+
+        # Create random input tensor on CUDA
+        hidden_states = torch.randn(
+            (sequence_length, micro_batch_size, hidden_size), device="cuda"
+        ).to(dtype)
+        hidden_states_ref = hidden_states.clone()
+        current_rank = dist.get_rank()
+
+        # Initialize model parallel with desired TP
+        Utils.initialize_model_parallel(tensor_model_parallel_size=parallel_state_tp)
+
+        # Build a reference grid and block for parameter sharing & DP averaging
+        ref_grid = create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=8)
+        ref_pg_collection = _get_pg_collection_from_grid(ref_grid)
+        ref_block = _create_transformer_block(
+            dtype=dtype, hidden_size=hidden_size, pg_collection=ref_pg_collection
+        )
+        _avg_params(
+            ref_block, ref_grid.get_pg("dp")
+        )  # Synchronize parameters across DP for reproducibility
+
+        # ====== Create two transformer block+grid pairs with different TP/DP settings ======
+        block_grid_1, grid_1 = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=grid1_tp,
+            pp_size=grid1_pp,
+            dp_size=grid1_dp,
+            grid_offset=0,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        block_grid_2, grid_2 = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=grid2_tp,
+            pp_size=grid2_pp,
+            dp_size=grid2_dp,
+            grid_offset=grid_1.size,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        dist.barrier()  # Synchronize ranks before communication
+
+        # Module-grid map and pipeline communication topology
+        module_to_grid_map = {'image_encoder': grid_1, 'llm': grid_2}
+        topology = {
+            'image_encoder': ['llm'],  # image_encoder sends forward results to llm
+            'llm': [],  # llm is the last stage here
+        }
+        config = ModelParallelConfig(pipeline_dtype=torch.float)
+        mllm_comm = MultiModulePipelineCommunicator(
+            module_to_grid_map, topology, config, dim_mapping={'s': 0, 'h': 2, 'b': 1}
+        )
+
+        output_grid_2 = None
+        # If current rank is in the first grid, run first block and send output
+        if grid_1 is not None and mllm_comm.is_current_rank_in_grid(grid_1):
+            rank_module_info = mllm_comm.rank_module_map['image_encoder']
+            if rank_module_info.pp_rank == 0:
+                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'image_encoder': hidden_states})
+            else:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(sequence_length, micro_batch_size, hidden_size)
+                )
+                hidden_states = input_dict['image_encoder']
+                hidden_states = block_grid_1(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'image_encoder': hidden_states})
+
+        # If current rank is in second grid, receive and run the second block
+        if grid_2 is not None and mllm_comm.is_current_rank_in_grid(grid_2):
+            rank_module_info = mllm_comm.rank_module_map['llm']
+            if rank_module_info.pp_rank == 0:
+                input_dict = mllm_comm.recv_forward()
+                hidden_states = input_dict['image_encoder']
+                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+                if rank_module_info.pp_rank == rank_module_info.pp_size - 1:
+                    output_grid_2 = hidden_states
+                else:
+                    mllm_comm.send_forward({'llm': hidden_states})
+            elif rank_module_info.pp_rank < rank_module_info.pp_size - 1:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(
+                        sequence_length,
+                        (grid1_dp * micro_batch_size) // grid2_dp,
+                        hidden_size,
+                    )
+                )
+                hidden_states = input_dict['llm']
+                hidden_states = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+                mllm_comm.send_forward({'llm': hidden_states})
+            else:
+                input_dict = mllm_comm.recv_forward(
+                    tensor_shape=(
+                        sequence_length,
+                        (grid1_dp * micro_batch_size) // grid2_dp,
+                        hidden_size,
+                    )
+                )
+                hidden_states = input_dict['llm']
+                output_grid_2 = block_grid_2(hidden_states=hidden_states, attention_mask=None)
+
+                # Compute expected output shape based on change in DP size (chunk/expand batch dimension appropriately)
+                factor = max(grid1_dp, grid2_dp) // min(grid1_dp, grid2_dp)
+                expected_output_shape = (
+                    sequence_length,
+                    (
+                        micro_batch_size * factor
+                        if grid1_dp > grid2_dp
+                        else micro_batch_size // factor
+                    ),
+                    hidden_size,
+                )
+                assert (
+                    output_grid_2.shape == expected_output_shape
+                ), f"Output2 shape mismatch: {output_grid_2.shape}"
+
+        # ====== Reference: global (replicated) pipeline forward for correctness checking ======
+        global_block_1, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+        global_block_2, _ = get_transformer_block_and_grid(
+            ref_block,
+            tp_size=parallel_state_tp,
+            use_global_parallel_state=True,
+            hidden_size=hidden_size,
+            dtype=dtype,
+        )
+
+        for i in range(grid1_pp):
+            hidden_states_ref = global_block_1(hidden_states=hidden_states_ref, attention_mask=None)
+
+        for i in range(grid2_pp):
+            hidden_states_ref = global_block_2(hidden_states=hidden_states_ref, attention_mask=None)
+
+        # Output comparison under different DP compositions between grids
+        if (
+            grid_2 is not None
+            and mllm_comm.is_current_rank_in_grid(grid_2)
+            and rank_module_info.pp_rank == rank_module_info.pp_size - 1
+        ):
+            if grid1_dp == grid2_dp:
+                # DP size matches: all outputs directly compared
+                torch.testing.assert_close(hidden_states_ref, output_grid_2, rtol=1e-3, atol=1e-3)
+            elif grid1_dp < grid2_dp:
+                # If grid2 expands DP: each output_grid_2 chunk corresponds to a split of the reference output
+                grid2_dp_ranks = grid_2._gen_rank_enum([x for x in grid_2.dim_names if x != "dp"])
+                global_block_2_chunks = torch.split(
+                    hidden_states_ref, hidden_states_ref.shape[1] // (grid2_dp // grid1_dp), dim=1
+                )
+                relevant_chunk = None
+                for i, dp_ranks in enumerate(grid2_dp_ranks):
+                    if current_rank in dp_ranks:
+                        relevant_chunk = global_block_2_chunks[i % len(global_block_2_chunks)]
+                torch.testing.assert_close(relevant_chunk, output_grid_2, rtol=1e-3, atol=1e-3)
+            else:
+                # If DP shrinks (grid1_dp > grid2_dp): just compare the relevant first chunk
+                output_grid_2_first_chunk = torch.chunk(output_grid_2, grid1_dp // grid2_dp, dim=1)[
+                    0
+                ]
+                torch.testing.assert_close(
+                    hidden_states_ref, output_grid_2_first_chunk, rtol=1e-3, atol=1e-3
+                )

From db6b895c69b71d2db7f5fd58d3f6de76ed40a619 Mon Sep 17 00:00:00 2001
From: litianjian <45817262+litianjian@users.noreply.github.com>
Date: Tue, 27 Jan 2026 04:16:54 +0800
Subject: [PATCH 26/79] Add router replay for MoE models (#2101)

Co-authored-by: litianjian <litianjian@bytedance.com>
Co-authored-by: Yan Bai <baiyan1996@icloud.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: Siddharth Singh <136645615+sidsingh-nvidia@users.noreply.github.com>
---
 docs/api-guide/router_replay.md               | 176 ++++++++++++++++++
 megatron/core/transformer/moe/moe_utils.py    |  19 +-
 megatron/core/transformer/moe/router.py       |   6 +
 .../core/transformer/moe/router_replay.py     | 161 ++++++++++++++++
 .../core/transformer/transformer_config.py    |   3 +
 megatron/training/arguments.py                |   3 +
 .../unit_tests/models/test_mamba_moe_model.py |   1 +
 .../transformer/moe/test_router_replay.py     |  95 ++++++++++
 8 files changed, 463 insertions(+), 1 deletion(-)
 create mode 100644 docs/api-guide/router_replay.md
 create mode 100644 megatron/core/transformer/moe/router_replay.py
 create mode 100644 tests/unit_tests/transformer/moe/test_router_replay.py

diff --git a/docs/api-guide/router_replay.md b/docs/api-guide/router_replay.md
new file mode 100644
index 00000000000..300a50db127
--- /dev/null
+++ b/docs/api-guide/router_replay.md
@@ -0,0 +1,176 @@
+# Design Document: MoE Router Replay Feature
+
+### 1. Overview
+
+This document provides a detailed description of the "Router Replay" feature implemented within the Megatron-LM Core for Mixture-of-Experts (MoE) models.
+
+This feature is designed to enhance determinism and analyzability in MoE model training and inference. It enables the model to load routing decisions from a predefined file and enforce their use during the forward pass, thereby bypassing the real-time routing computation.
+
+### 2. Motivation
+
+*   **Determinism & Reproducibility**: In distributed training, MoE routing decisions can exhibit minor variations due to factors like floating-point precision. By replaying a fixed routing table, the MoE computation path is guaranteed to be identical across runs, which facilitates debugging and reproducing experimental results.
+*   **Performance Profiling**: The router's own computation (e.g., logits calculation, top-k selection) incurs overhead. In replay mode, this part of the computation can be completely skipped, allowing for more precise isolation and profiling of performance bottlenecks within the Expert Layers themselves.
+*   **Debugging Aid**: When issues arise in the model, fixing the routing decisions helps to isolate variables, making it easier to determine whether the problem lies with the routing mechanism or the expert computations.
+
+### 3. Design and Architecture
+
+The design follows the principles of being non-intrusive and on-demand, with the core idea of activating the replay logic only when explicitly requested by the user.
+
+*   **Core Components**:
+    *   `RouterReplay` (located in `megatron/core/transformer/moe/router_replay.py`): A utility class for replaying MoE routing decisions. When enabled via the `moe_enable_routing_replay` flag, a separate instance of `RouterReplay` is created for each MoE layer's router. Each instance is responsible for loading routing data and providing the deterministic routing decisions for its corresponding layer during the forward pass.
+    *   `moe_enable_routing_replay` (located in `megatron/core/transformer/transformer_config.py`): A boolean global configuration flag that serves as the sole entry point for enabling this feature.
+
+*   **Workflow**:
+    The feature supports different modes, such as recording and replaying, controlled by a `RouterReplayAction`.
+
+    1.  **Enabling the Feature**: The user sets `moe_enable_routing_replay` to `True` in the model configuration.
+    2.  **Initialization**: When `moe_enable_routing_replay` is true, each `TopKRouter` creates its own `RouterReplay` instance.
+    3.  **Mode Configuration**: The user must programmatically set the desired router replay action (e.g., `record`, `forward_replay`, `backward_replay`) on the `RouterReplay` instances.
+    4.  **Execution Flow (within a mini-batch)**:
+        *   **Forward Pass**:
+            *   For each micro-batch, the `topk_routing_with_score_function` checks the `router_replay_action`.
+            *   **In `record` mode**: The dynamically computed `top-k` expert indices are captured and stored.
+            *   **In `forward_replay` mode**: The function retrieves pre-loaded expert indices from `target_topk_idx`. These indices are used for the forward computation and are also appended to the `replay_backward_list` to prepare for the backward pass.
+        *   **Backward Pass**:
+            *   For each micro-batch (processed in reverse order in pipeline parallelism), the `router_replay_action` is checked again.
+            *   **In `backward_replay` mode**: The function retrieves the expert indices for the corresponding micro-batch by popping them from the `replay_backward_list`. This mode is intended for training recomputation (e.g., activation checkpointing and pipeline recompute) so the same routing decisions are used during recompute/backward as in forward, ensuring determinism and correctness.
+
+### 4. Implementation Details
+
+The implementation cleanly separates the replay logic from the router's core computation.
+
+*   **`megatron/core/transformer/transformer_config.py`**:
+    *   Adds the configuration option `moe_enable_routing_replay: bool = False`.
+
+*   **`megatron/core/transformer/moe/moe_utils.py`**:
+    *   Introduces the `RouterReplay` class to manage the state for recording and replaying routing decisions for a single MoE layer.
+        *   `target_topk_idx`: An attribute holding the expert indices for the current micro-batch during forward replay mode.
+        *   `recorded_topk_idx`: An attribute for storing the computed expert indices when in record mode.
+        *   `replay_backward_list`: A list that accumulates the `top-k` indices used during the forward passes of a mini-batch. This list is consumed in FIFO order during the backward pass to ensure correctness under pipeline parallelism.
+        *   `set_target_indices()`: A method to load the replay indices into `target_topk_idx` for the forward pass.
+        *   `record_indices()`: A method to save the computed indices.
+    *   The `topk_routing_with_score_function` is modified to contain the core logic. It checks the `router_replay_action` on the `router_replay` instance and accordingly performs one of the following actions: computes and records indices, replays indices from `target_topk_idx` (for forward), replays indices from `replay_backward_list` (for backward), or falls through to the default dynamic routing.
+
+#### Training recompute usage
+- During forward replay, `set_target_indices()` prepares `replay_backward_list` so each micro-batch’s indices are available for recomputation.
+- During recompute/backward, set action to `REPLAY_BACKWARD` so indices are consumed in FIFO order to mirror the forward sequence.
+
+### 5. Usage Guide
+
+1.  **Enable & Instantiate**
+    - Create one `RouterReplay` instance per MoE router layer when building the model.
+    - Optionally use the global helpers to set/clear actions across all layers.
+2.  **Record Routing Decisions**
+    - Set action: `RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)`.
+    - Run the model; retrieve per-layer indices via `RouterReplay.get_recorded_data()` and persist.
+3.  **Forward Replay**
+    - Load indices and distribute: `RouterReplay.set_replay_data(list_of_tensors)`.
+    - Set action: `RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)`.
+    - Run the model; dynamic top‑k is bypassed and target indices are used.
+4.  **Backward Replay**
+    - For training recomputation (activation checkpointing or pipeline recompute), set action: `REPLAY_BACKWARD` during recomputation.
+    - Per micro‑batch indices are consumed from `replay_backward_list` in FIFO order.
+5.  **Cleanup**
+    - Use `RouterReplay.clear_global_indices()`, `RouterReplay.clear_global_router_replay_action()`, and `RouterReplay.clear_global_router_replay_instances()` to restore default behavior and prevent memory leaks.
+
+#### Quick usage with `topk_routing_with_score_function`
+
+```python
+import torch
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+from megatron.core.transformer.moe.moe_utils import topk_routing_with_score_function
+
+rr = RouterReplay()
+
+# Record
+RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+logits = torch.randn(8, 16)
+probs_rec, routing_map_rec = topk_routing_with_score_function(
+    logits=logits, topk=2, use_pre_softmax=False, score_function="softmax", router_replay=rr,
+)
+recorded = rr.get_recorded_indices()
+torch.save(recorded, "/tmp/replay.pt")
+
+# Forward replay
+rr.clear_router_replay_action()
+rr.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+target = torch.load("/tmp/replay.pt")
+rr.set_target_indices(target)
+probs_rep, routing_map_rep = topk_routing_with_score_function(
+    logits=logits, topk=2, use_pre_softmax=False, score_function="softmax", router_replay=rr,
+)
+
+RouterReplay.clear_global_router_replay_action()
+RouterReplay.clear_global_indices()
+RouterReplay.clear_global_router_replay_instances()
+```
+
+### 6. Minimal Demo
+
+Here is a minimal code example showing how to use RouterReplay for recording and replaying:
+
+```python
+import torch
+import torch.distributed as dist
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+
+
+# Initialize distributed training
+if not dist.is_initialized():
+    dist.init_process_group(backend="nccl")
+
+# Create a transformer config with RouterReplay enabled
+config = TransformerConfig(
+    num_experts=8,
+    expert_model_parallel_size=1,
+    num_top_k=2,
+    moe_enable_routing_replay=True
+)
+
+# Create a TopKRouter instance
+router = TopKRouter(config)
+
+# Generate sample input (batch_size, sequence_length, hidden_size)
+logits = torch.randn(16, 32, 8).to(torch.cuda.current_device())
+
+# -----------------
+# 1. Recording Mode
+# -----------------
+print("=== Recording Mode ===")
+# Set global router replay action to RECORD
+RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+
+# Perform routing
+routing_output = router.forward(logits)
+print(f"Recorded top-k indices shape: {routing_output.top_k_idx.shape}")
+
+# -----------------
+# 2. Forward Replay Mode
+# -----------------
+print("\n=== Forward Replay Mode ===")
+# Save recorded indices to a file
+torch.save(routing_output.top_k_idx, "/tmp/replay.pt")
+
+# Load indices from file and set as target for replay
+replay_indices = torch.load("/tmp/replay.pt")
+for router_instance in RouterReplay.global_router_replay_instances:
+    router_instance.target_topk_idx = replay_indices
+
+# Set global router replay action to REPLAY_FORWARD
+RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+
+# Perform routing again - this will use the replayed indices
+replay_routing_output = router.forward(logits)
+print(f"Replayed top-k indices shape: {replay_routing_output.top_k_idx.shape}")
+print(f"Are indices the same? {torch.equal(routing_output.top_k_idx, replay_routing_output.top_k_idx)}")
+
+
+# Clean up
+RouterReplay.clear_global_router_replay_action()
+RouterReplay.clear_global_indices()
+RouterReplay.clear_global_router_replay_instances()
+if dist.is_initialized():
+    dist.destroy_process_group()
+```
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index bd7c29551a8..dc7450d93d0 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -14,6 +14,7 @@
 from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphScope
+from megatron.core.transformer.moe.router_replay import RouterReplay
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import internal_api
 
@@ -567,6 +568,7 @@ def topk_routing_with_score_function(
     score_function: str = "softmax",
     expert_bias: Optional[torch.Tensor] = None,
     fused: bool = False,
+    router_replay: Optional['RouterReplay'] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute the routing probabilities and map for top-k selection with score function.
 
@@ -584,6 +586,11 @@ def topk_routing_with_score_function(
         expert_bias (torch.Tensor, optional): The bias added to logits for expert routing.
                                               Defaults to None.
         fused (bool, optional): Whether to use the fused version. Defaults to False.
+        router_replay (Optional['RouterReplay']): For debugging and development, allows for
+                                             deterministic routing by replaying a previously
+                                             recorded routing sequence.
+
+                                              Defaults to None.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]:
@@ -611,7 +618,7 @@ def topk_routing_with_score_function(
             expert_bias=expert_bias,
         )
 
-    def compute_topk(
+    def _compute_topk(
         scores: torch.Tensor,
         topk: int,
         num_groups: Optional[int] = None,
@@ -642,6 +649,16 @@ def compute_topk(
         else:
             return torch.topk(scores, k=topk, dim=1)
 
+    def compute_topk(scores, topk, num_groups=None, group_topk=None):
+        # Default behavior if no replay is active
+
+        if router_replay is None:
+            return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
+        else:
+            return router_replay.get_replay_topk(
+                scores, topk, num_groups, group_topk, _compute_topk
+            )
+
     if score_function == "softmax":
         if use_pre_softmax:
             scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index c22ca4e8446..4e3a08d66d8 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -21,6 +21,7 @@
     topk_routing_with_score_function,
     z_loss_func,
 )
+from megatron.core.transformer.moe.router_replay import RouterReplay
 from megatron.core.transformer.transformer_config import TransformerConfig
 
 
@@ -201,6 +202,10 @@ def __init__(
             self.global_tokens_per_expert = None
             self.ga_steps = None
 
+        self.router_replay = None
+        if self.config.moe_enable_routing_replay:
+            self.router_replay = RouterReplay()
+
     def _maintain_float32_expert_bias(self):
         """
         Maintain the expert bias in float32.
@@ -523,6 +528,7 @@ def routing(self, logits: torch.Tensor):
                 score_function=self.score_function,
                 expert_bias=self.expert_bias,
                 fused=self.config.moe_router_fusion,
+                router_replay=self.router_replay,
             )
 
         # Apply token dropping to probs and routing_map.
diff --git a/megatron/core/transformer/moe/router_replay.py b/megatron/core/transformer/moe/router_replay.py
new file mode 100644
index 00000000000..b6b8e26a0a6
--- /dev/null
+++ b/megatron/core/transformer/moe/router_replay.py
@@ -0,0 +1,161 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+from enum import Enum
+from typing import Callable, List, Optional, Tuple
+
+import torch
+
+
+class RouterReplayAction(Enum):
+    """
+    A Enum to define the actions for router replay.
+    """
+
+    RECORD = "record"  # Record the topk indices for replay
+    REPLAY_FORWARD = "replay_forward"  # Replay the recorded topk indices for forward pass
+    REPLAY_BACKWARD = "replay_backward"  # Replay topk indices for re-compute during backward pass
+
+
+class RouterReplay:
+    """
+    A class to manage the recording and replaying of MoE routing decisions.
+    It holds all router instances and provides static methods to globally
+    control recording and replaying.
+    """
+
+    # Static variable to hold all router instances, one per MoE layer.
+    global_router_replay_instances: List['RouterReplay'] = []
+
+    @staticmethod
+    def set_replay_data(all_layers_topk_indices: List[torch.Tensor]):
+        """
+        Distributes the topk indices for all layers to their respective RouterReplay instances.
+        :param all_layers_topk_indices: A list of tensors, where each tensor contains the
+                                        topk indices for a specific layer. The order
+                                        must match the instantiation order of the routers.
+        """
+        if len(all_layers_topk_indices) != len(RouterReplay.global_router_replay_instances):
+            raise ValueError(
+                f"The number of replay tensors ({len(all_layers_topk_indices)}) "
+                f"does not match instances ({len(RouterReplay.global_router_replay_instances)})."
+            )
+        for i, router_instance in enumerate(RouterReplay.global_router_replay_instances):
+            router_instance.set_target_indices(all_layers_topk_indices[i])
+
+    @staticmethod
+    def get_recorded_data() -> List[torch.Tensor]:
+        """
+        Collects the recorded topk indices from all RouterReplay instances.
+        :return: A list of tensors, each containing the recorded topk indices for a layer.
+        """
+        return [
+            router.get_recorded_indices() for router in RouterReplay.global_router_replay_instances
+        ]
+
+    @staticmethod
+    def clear_global_indices():
+        """Clears the recorded and target topk indices in all instances."""
+        for router in RouterReplay.global_router_replay_instances:
+            router.clear_indices()
+
+    @staticmethod
+    def set_global_router_replay_action(router_replay_action: RouterReplayAction):
+        """Sets the router replay action for all router instances."""
+        for router in RouterReplay.global_router_replay_instances:
+            router.set_router_replay_action(router_replay_action)
+
+    @staticmethod
+    def clear_global_router_replay_action():
+        """Clears the router replay action for all router instances."""
+        for router in RouterReplay.global_router_replay_instances:
+            router.clear_router_replay_action()
+
+    @staticmethod
+    def clear_global_router_replay_instances():
+        """Clear the global list of router replay instances to prevent memory leaks."""
+        RouterReplay.global_router_replay_instances.clear()
+
+    def __init__(self):
+        """Initializes a RouterReplay instance for a specific layer."""
+        self.target_topk_idx: Optional[torch.Tensor] = None  # Target topk indices for replay
+        self.recorded_topk_idx: Optional[torch.Tensor] = None  # Recorded topk indices for replay
+        self.router_replay_action: Optional[RouterReplayAction] = (
+            None  # Router replay action for this layer
+        )
+        self.replay_backward_list: List[torch.Tensor] = (
+            []
+        )  # List of tensors for backward pass replay
+        RouterReplay.global_router_replay_instances.append(self)
+
+    def set_target_indices(self, topk_indices: torch.Tensor):
+        """Sets the target topk indices for replay."""
+        self.target_topk_idx = topk_indices
+        self.replay_backward_list.append(topk_indices)
+
+    def get_recorded_indices(self) -> Optional[torch.Tensor]:
+        """Returns the recorded topk indices."""
+        return self.recorded_topk_idx
+
+    def record_indices(self, topk_indices: torch.Tensor):
+        """Records the topk indices."""
+        self.recorded_topk_idx = topk_indices
+
+    def clear_indices(self):
+        """Clears the recorded and target topk indices."""
+        self.recorded_topk_idx = None
+        self.target_topk_idx = None
+        self.replay_backward_list = []
+
+    def set_router_replay_action(self, router_replay_action: RouterReplayAction):
+        """Sets the router replay action for this layer."""
+        self.router_replay_action = router_replay_action
+
+    def clear_router_replay_action(self):
+        """Clears the router replay action for this layer."""
+        self.router_replay_action = None
+
+    def get_replay_topk(
+        self,
+        scores: torch.Tensor,
+        topk: int,
+        num_groups: Optional[int] = None,
+        group_topk: Optional[int] = None,
+        default_compute_topk: Callable[
+            [torch.Tensor, int, Optional[int], Optional[int]], Tuple[torch.Tensor, torch.Tensor]
+        ] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        A wrapper for top-k computation that handles different replay actions.
+
+        Args:
+            scores (torch.Tensor): The scores to compute top-k on.
+            topk (int): The number of top elements to select.
+            num_groups (Optional[int]): Number of expert groups for group-limited routing.
+            group_topk (Optional[int]): Number of groups to select for each token.
+            default_compute_topk (Callable): The default top-k computation function, which
+                                             should return a tuple of (values, indices).
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: A tuple containing the top-k values and indices.
+        """
+        if self.router_replay_action == RouterReplayAction.RECORD:
+            probs, top_indices = default_compute_topk(
+                scores, topk, num_groups=num_groups, group_topk=group_topk
+            )
+            self.record_indices(top_indices)
+            return probs, top_indices
+        elif self.router_replay_action == RouterReplayAction.REPLAY_FORWARD:
+            top_indices = self.target_topk_idx
+            # Ensure indices are on the correct device
+            top_indices = top_indices.to(scores.device)
+            # Gather the scores for the replayed indices to get the probabilities
+            probs = scores.gather(1, top_indices)
+            return probs, top_indices
+        elif self.router_replay_action == RouterReplayAction.REPLAY_BACKWARD:
+            top_indices = self.replay_backward_list.pop(0)
+            # Ensure indices are on the correct device
+            top_indices = top_indices.to(scores.device)
+            # Gather the scores for the replayed indices to get the probabilities
+            probs = scores.gather(1, top_indices)
+            return probs, top_indices
+        else:
+            return default_compute_topk(scores, topk, num_groups, group_topk)
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 77dc81cfd92..633c51e789e 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -549,6 +549,9 @@ class TransformerConfig(ModelParallelConfig):
     moe_router_topk: int = 2
     """Number of experts to route to for each token."""
 
+    moe_enable_routing_replay: bool = False
+    """If True, enable the routing replay feature for MoE layers."""
+
     moe_router_topk_limited_devices: Optional[int] = None
     """Number of EP ranks to consider for each token in group-limited routing,
     DEPRECATED and replaced by moe_router_num_groups and moe_router_group_topk.
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 2108a1230b9..6da61d9ba51 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -3301,6 +3301,9 @@ def _add_moe_args(parser):
                        help='Score function for MoE TopK routing. Can be "softmax" or "sigmoid".')
     group.add_argument('--moe-router-topk', type=int, default=2,
                        help='Number of experts to route to for each token. The default is 2.')
+    group.add_argument('--moe-enable-routing-replay', action='store_true',
+                       help='Enable routing replay for MoE routers. When enabled, the router will '
+                            'use a pre-defined routing table instead of computing it on the fly.')
     group.add_argument('--moe-router-pre-softmax', action='store_true',
                        help='Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k.')
     group.add_argument('--moe-router-num-groups', type=int, default=None,
diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py
index 5680751f63f..3c7ae93a17c 100644
--- a/tests/unit_tests/models/test_mamba_moe_model.py
+++ b/tests/unit_tests/models/test_mamba_moe_model.py
@@ -191,6 +191,7 @@
     "moe_token_dropping": False,
     "moe_use_legacy_grouped_gemm": False,
     "moe_z_loss_coeff": None,
+    "moe_enable_routing_replay": False,
     "mrope_section": None,
     "mtp_loss_scaling_factor": 0.1,
     "mtp_num_layers": None,
diff --git a/tests/unit_tests/transformer/moe/test_router_replay.py b/tests/unit_tests/transformer/moe/test_router_replay.py
new file mode 100644
index 00000000000..840fc0fd269
--- /dev/null
+++ b/tests/unit_tests/transformer/moe/test_router_replay.py
@@ -0,0 +1,95 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+import pytest
+import torch
+
+from megatron.core.transformer.moe.moe_utils import topk_routing_with_score_function
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+
+
+def setup_function():
+    RouterReplay.global_router_replay_instances.clear()
+
+
+def teardown_function():
+    RouterReplay.global_router_replay_instances.clear()
+
+
+def test_record_mode_with_topk_routing_softmax_post():
+    rr = RouterReplay()
+    rr.set_router_replay_action(RouterReplayAction.RECORD)
+    logits = torch.randn(4, 6)
+    probs, routing_map = topk_routing_with_score_function(
+        logits=logits, topk=2, use_pre_softmax=False, router_replay=rr, score_function="softmax"
+    )
+    recorded = rr.get_recorded_indices()
+    expected_idx = torch.topk(logits, k=2, dim=1).indices
+    assert recorded is not None
+    assert torch.equal(recorded, expected_idx)
+    assert probs.shape == (4, 6)
+    assert routing_map.shape == (4, 6)
+    assert routing_map.sum(dim=1).eq(2).all()
+
+
+def test_replay_forward_with_topk_routing_softmax_pre():
+    rr = RouterReplay()
+    rr.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+    logits = torch.randn(3, 5)
+    target = torch.tensor([[1, 2], [0, 3], [2, 4]], dtype=torch.long)
+    rr.set_target_indices(target)
+    probs, routing_map = topk_routing_with_score_function(
+        logits=logits, topk=2, use_pre_softmax=True, router_replay=rr, score_function="softmax"
+    )
+    assert routing_map.sum(dim=1).eq(2).all()
+    scores = torch.softmax(logits, dim=-1)
+    assert torch.equal(probs.gather(1, target), scores.gather(1, target))
+
+
+def test_replay_forward_with_topk_routing_softmax_post():
+    rr = RouterReplay()
+    rr.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+    logits = torch.randn(3, 6)
+    target = torch.tensor([[1, 2], [0, 5], [3, 4]], dtype=torch.long)
+    rr.set_target_indices(target)
+    probs, routing_map = topk_routing_with_score_function(
+        logits=logits, topk=2, use_pre_softmax=False, router_replay=rr, score_function="softmax"
+    )
+    selected = torch.softmax(logits.gather(1, target), dim=-1)
+    assert torch.equal(probs.gather(1, target), selected)
+    assert routing_map.sum(dim=1).eq(2).all()
+
+
+def test_global_set_get_clear_indices():
+    r1 = RouterReplay()
+    r2 = RouterReplay()
+    t1 = torch.tensor([[0, 1]], dtype=torch.long)
+    t2 = torch.tensor([[1, 0]], dtype=torch.long)
+    RouterReplay.set_replay_data([t1, t2])
+    assert torch.equal(r1.target_topk_idx, t1)
+    assert torch.equal(r2.target_topk_idx, t2)
+    r1.record_indices(t1)
+    r2.record_indices(t2)
+    rec = RouterReplay.get_recorded_data()
+    assert len(rec) == 2
+    assert torch.equal(rec[0], t1)
+    assert torch.equal(rec[1], t2)
+    RouterReplay.clear_global_indices()
+    assert r1.target_topk_idx is None and r2.target_topk_idx is None
+    assert r1.get_recorded_indices() is None and r2.get_recorded_indices() is None
+
+
+def test_global_action_set_and_clear():
+    r1 = RouterReplay()
+    r2 = RouterReplay()
+    RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+    assert r1.router_replay_action == RouterReplayAction.REPLAY_FORWARD
+    assert r2.router_replay_action == RouterReplayAction.REPLAY_FORWARD
+    RouterReplay.clear_global_router_replay_action()
+    assert r1.router_replay_action is None and r2.router_replay_action is None
+
+
+def test_set_replay_data_length_mismatch():
+    _ = RouterReplay()
+    with pytest.raises(ValueError):
+        RouterReplay.set_replay_data(
+            [torch.tensor([[0, 1]], dtype=torch.long), torch.tensor([[1, 0]], dtype=torch.long)]
+        )

From 528cb2e520bbad8c921e68787e6b1885b67c05e5 Mon Sep 17 00:00:00 2001
From: jeffnvidia <152798700+jeffnvidia@users.noreply.github.com>
Date: Mon, 26 Jan 2026 23:35:38 +0200
Subject: [PATCH 27/79] add all_gather process-group for overlapping in fsdp
 disributed training (#2663)

---
 .../megatron_fsdp/param_and_grad_buffer.py    | 28 ++++++++++++-
 .../fsdp/src/megatron_fsdp/utils.py           | 20 ++++++++-
 megatron/core/parallel_state.py               | 41 ++++++++++++++++++-
 megatron/training/arguments.py                |  3 ++
 megatron/training/initialize.py               |  1 +
 tests/unit_tests/test_parallel_state.py       | 28 +++++++++++++
 6 files changed, 117 insertions(+), 4 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
index 546bc0721e0..9a0ef354c26 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
@@ -1602,6 +1602,18 @@ def __init__(
             if self.dist_index.get_outer_fsdp_group() is not None:
                 # Outer/Inter-FSDP group when using hybrid FSDP
                 self.ubr_groups.append(self.dist_index.get_outer_fsdp_group())
+            if (
+                self.dist_index.get_fsdp_group(
+                    is_expert_parallel=False, independent_all_gather=True
+                )
+                is not None
+            ):
+                # All-gather group used when overlapping all-gather and gradient reduction.
+                self.ubr_groups.append(
+                    self.dist_index.get_fsdp_group(
+                        is_expert_parallel=False, independent_all_gather=True
+                    )
+                )
 
             if torch.distributed.get_rank() == 0:
                 logging.info(
@@ -1888,6 +1900,18 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
                     is_expert_parallel=group.is_expert_param
                 )
 
+            # When --create-all-gather-group is enabled, use a separate process group for
+            # all-gather operations (model_weight_buffer) to enable overlap with gradient reduction
+            # operations (main_grad_buffer). This avoids head-of-line blocking between forward
+            # all-gather and backward reduce-scatter on the same communicator.
+            model_wbuf_dp_group = main_buf_dp_group
+            if not group.is_expert_param and not should_create_hfsdp_wbuf_and_gbuf:
+                ag_group = self.dist_index.get_fsdp_group(
+                    is_expert_parallel=False, independent_all_gather=True
+                )
+                if ag_group is not None:
+                    model_wbuf_dp_group = ag_group
+
             gradient_scaling_factor = (
                 self.gradient_scaling_factor
                 if not group.is_expert_param
@@ -1928,10 +1952,10 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
                     self.ddp_config,
                     group.params,
                     is_data_distributed=is_model_weight_buffer_distributed
-                    and main_buf_dp_group.size() > 1,
+                    and model_wbuf_dp_group.size() > 1,
                     dtype=param_dtype,
                     device=self.device,
-                    data_parallel_group=main_buf_dp_group,
+                    data_parallel_group=model_wbuf_dp_group,
                     is_transpose_buffer=False,
                     temporary_bucket_allocator=self.weight_alloc,
                     bucket_id=group_id,
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py
index 01523929ae1..d5fbc91fcf8 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py
@@ -21,6 +21,13 @@
 from importlib.metadata import version
 from typing import Callable, Optional, Sequence, Union
 
+try:
+    import megatron.core.parallel_state as parallel_state
+
+    HAVE_MEGATRON_CORE = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_MEGATRON_CORE = False
+
 try:
     import einops
 
@@ -486,6 +493,13 @@ def __init__(
             if contains_submesh(self.device_mesh, self.dp_shard_dim)
             else None
         )
+        # AG group comes from parallel_state, not the mesh
+        # the purpose of this independent group is to overlap all-gather and gradient reduction.
+        self.fsdp_group_ag = None
+        if HAVE_MEGATRON_CORE and parallel_state.has_separate_all_gather_group():
+            self.fsdp_group_ag = parallel_state.get_data_parallel_group(
+                with_context_parallel=True, independent_all_gather=True
+            )
         # Retrieve the outer-FSDP process group from the DeviceMesh.
         self.outer_fsdp_group = (
             self.device_mesh[self.dp_outer_dim].get_group()
@@ -620,10 +634,14 @@ def get_dp_group(self, is_expert_parallel: bool = False) -> ProcessGroup:
             return self.hybrid_fsdp_group
         return self.fsdp_group
 
-    def get_fsdp_group(self, is_expert_parallel: bool = False) -> ProcessGroup:
+    def get_fsdp_group(
+        self, is_expert_parallel: bool = False, independent_all_gather: bool = False
+    ) -> ProcessGroup:
         """Get the FSDP process group."""
         if is_expert_parallel:
             return self.expt_fsdp_group
+        if independent_all_gather:
+            return self.fsdp_group_ag
         return self.fsdp_group
 
     def get_outer_fsdp_group(self) -> ProcessGroup:
diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py
index c5a73600ee1..7bb96407838 100644
--- a/megatron/core/parallel_state.py
+++ b/megatron/core/parallel_state.py
@@ -120,6 +120,7 @@
 
 # Data parallel group information with context parallel combined.
 _DATA_PARALLEL_GROUP_WITH_CP = None
+_DATA_PARALLEL_GROUP_WITH_CP_AG = None
 _DATA_PARALLEL_GROUP_WITH_CP_GLOO = None
 _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP = None
 
@@ -566,6 +567,7 @@ def initialize_model_parallel(
     create_gloo_process_groups: bool = True,
     high_priority_stream_groups: Optional[List[str]] = None,
     sharp_enabled_group: Optional[str] = None,
+    create_all_gather_group: Optional[bool] = False,
 ) -> None:
     """Initialize model data parallel groups.
 
@@ -680,6 +682,13 @@ def initialize_model_parallel(
             By default (None), it is enabled from dp group.
             Available options (choose one): [dp, dp_replica]
 
+        create_all_gather_group (bool, default = False):
+            Create a separate process group for all-gather operations to avoid
+            head-of-line blocking with reduce-scatter operations. When enabled,
+            creates an additional NCCL communicator with identical ranks as the
+            dp-cp group but with independent progress engines for better communication
+            overlap.
+
     Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
     use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
     the model pipeline. The present function will
@@ -816,6 +825,7 @@ def initialize_model_parallel(
     global _DATA_PARALLEL_GROUP_GLOO
     global _DATA_PARALLEL_GLOBAL_RANKS
     global _DATA_PARALLEL_GROUP_WITH_CP
+    global _DATA_PARALLEL_GROUP_WITH_CP_AG
     global _DATA_PARALLEL_GROUP_WITH_CP_GLOO
     global _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP
     global _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP
@@ -847,6 +857,15 @@ def initialize_model_parallel(
             pg_options=get_nccl_options("dp_cp", nccl_comm_cfgs),
             group_desc="DATA_PARALLEL_GROUP_WITH_CP",
         )
+        if create_all_gather_group:
+            group_with_cp_ag = create_group(
+                ranks_with_cp,
+                timeout=timeout,
+                pg_options=get_nccl_options("dp_cp", nccl_comm_cfgs),
+                group_desc="DATA_PARALLEL_GROUP_WITH_CP_AG",
+            )
+        else:
+            group_with_cp_ag = None
         if create_gloo_process_groups:
             group_with_cp_gloo = create_group(
                 ranks_with_cp,
@@ -858,6 +877,7 @@ def initialize_model_parallel(
             group_with_cp_gloo = None
         if rank in ranks_with_cp:
             _DATA_PARALLEL_GROUP_WITH_CP = group_with_cp
+            _DATA_PARALLEL_GROUP_WITH_CP_AG = group_with_cp_ag
             _DATA_PARALLEL_GROUP_WITH_CP_GLOO = group_with_cp_gloo
             _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP = ranks_with_cp
 
@@ -1387,7 +1407,9 @@ def get_pipeline_model_parallel_group(check_initialized=True):
     return _PIPELINE_MODEL_PARALLEL_GROUP
 
 
-def get_data_parallel_group(with_context_parallel=False, partial_data_parallel=False):
+def get_data_parallel_group(
+    with_context_parallel=False, partial_data_parallel=False, independent_all_gather=False
+):
     """Get the data-parallel group the caller rank belongs to."""
     if with_context_parallel:
         if partial_data_parallel:
@@ -1395,6 +1417,11 @@ def get_data_parallel_group(with_context_parallel=False, partial_data_parallel=F
                 _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP is not None
             ), "Intra partial data parallel group is not initialized"
             return _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP
+        if independent_all_gather:
+            assert (
+                _DATA_PARALLEL_GROUP_WITH_CP_AG is not None
+            ), "data parallel group with context parallel AG is not initialized"
+            return _DATA_PARALLEL_GROUP_WITH_CP_AG
         assert (
             _DATA_PARALLEL_GROUP_WITH_CP is not None
         ), "data parallel group with context parallel combined is not initialized"
@@ -1405,6 +1432,15 @@ def get_data_parallel_group(with_context_parallel=False, partial_data_parallel=F
         return _DATA_PARALLEL_GROUP
 
 
+def has_separate_all_gather_group() -> bool:
+    """Check if a separate all-gather process group has been created.
+
+    Returns True if a dedicated all-gather process group exists for improved
+    communication overlap, False otherwise.
+    """
+    return _DATA_PARALLEL_GROUP_WITH_CP_AG is not None
+
+
 def get_data_parallel_group_gloo(with_context_parallel=False, partial_data_parallel=False):
     """Get the Gloo data-parallel group the caller rank belongs to."""
     if with_context_parallel:
@@ -2065,6 +2101,9 @@ def destroy_model_parallel():
     global _DATA_PARALLEL_GROUP_WITH_CP
     _DATA_PARALLEL_GROUP_WITH_CP = None
 
+    global _DATA_PARALLEL_GROUP_WITH_CP_AG
+    _DATA_PARALLEL_GROUP_WITH_CP_AG = None
+
     global _CONTEXT_PARALLEL_GROUP
     _CONTEXT_PARALLEL_GROUP = None
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 6da61d9ba51..e27321f0096 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2815,6 +2815,9 @@ def _add_distributed_args(parser):
                        help='IB SHARP can be enabled from only one communication group. '
                        'By default, it is enabled from dp group. '
                        'Available options: [dp, dp_replica]')
+    group.add_argument('--create-all-gather-group', action='store_true',
+                   help='Create a separate process group for all-gather operations '
+                   'to overlap reduce-scatter and all-gather operations.')
     group.add_argument('--use-megatron-fsdp', action='store_true',
                        help='Use the Megatron FSDP code path in DDP.')
     group.add_argument('--init-model-with-meta-device', action='store_true')
diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py
index 00fa9ad5088..e300c03218b 100644
--- a/megatron/training/initialize.py
+++ b/megatron/training/initialize.py
@@ -389,6 +389,7 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s
                 create_gloo_process_groups=args.enable_gloo_process_groups,
                 high_priority_stream_groups=args.high_priority_stream_groups,
                 sharp_enabled_group=args.sharp_enabled_group,
+                create_all_gather_group=args.create_all_gather_group,
             )
             if args.rank == 0:
                 print(
diff --git a/tests/unit_tests/test_parallel_state.py b/tests/unit_tests/test_parallel_state.py
index 0c722ee0257..21dc740cdf4 100644
--- a/tests/unit_tests/test_parallel_state.py
+++ b/tests/unit_tests/test_parallel_state.py
@@ -530,3 +530,31 @@ def test_hybrid_dp_cp_groups(world_size, tp_size, cp_size, dp_size):
         assert group.size() == group_size
 
     Utils.destroy_model_parallel()
+
+
+def test_separate_all_gather_group():
+    """Test separate all-gather group for improved communication overlap."""
+    # Test without creating AG group (default)
+    Utils.initialize_model_parallel(context_parallel_size=world_size, create_all_gather_group=False)
+    assert not ps.has_separate_all_gather_group()
+    assert ps._DATA_PARALLEL_GROUP_WITH_CP_AG is None
+    Utils.destroy_model_parallel()
+
+    # Test with creating AG group
+    Utils.initialize_model_parallel(context_parallel_size=world_size, create_all_gather_group=True)
+    assert ps.has_separate_all_gather_group()
+    assert ps._DATA_PARALLEL_GROUP_WITH_CP_AG is not None
+
+    # Verify it returns the correct group
+    ag_group = ps.get_data_parallel_group(with_context_parallel=True, independent_all_gather=True)
+    regular_group = ps.get_data_parallel_group(
+        with_context_parallel=True, independent_all_gather=False
+    )
+    assert ag_group is not None
+    assert regular_group is not None
+    # They should have the same ranks but different communicators
+    ag_ranks = torch.distributed.get_process_group_ranks(ag_group)
+    regular_ranks = torch.distributed.get_process_group_ranks(regular_group)
+    assert ag_ranks == regular_ranks
+
+    Utils.destroy_model_parallel()

From b47c376f5ceb8fa9f2cb1ba8510112b668e41366 Mon Sep 17 00:00:00 2001
From: "Chenhan D. Yu" <5185878+ChenhanYu@users.noreply.github.com>
Date: Mon, 26 Jan 2026 16:32:33 -0800
Subject: [PATCH 28/79] Support NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 FP8/NVFP4
 PTQ in example (#3079)

Signed-off-by: jenchen13 <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Co-authored-by: jenchen13 <jennifchen@nvidia.com>
Co-authored-by: Jenny Chen <jc4686@columbia.edu>
Co-authored-by: Asha Anoosheh <ashaanoosheh1@gmail.com>
---
 examples/post_training/modelopt/README.md     | 14 +++--
 .../NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh    | 57 +++++++++++++++++++
 .../post_training/modelopt/convert_model.py   |  5 +-
 examples/post_training/modelopt/export.py     |  7 ++-
 examples/post_training/modelopt/quantize.py   | 12 +---
 .../modelopt/mamba/model_specs.py             | 17 +++++-
 6 files changed, 95 insertions(+), 17 deletions(-)
 create mode 100644 examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh
 mode change 100644 => 100755 examples/post_training/modelopt/export.py

diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md
index 48e679e4e31..6ebdd8ac5d6 100644
--- a/examples/post_training/modelopt/README.md
+++ b/examples/post_training/modelopt/README.md
@@ -32,6 +32,7 @@ knowledge distillation, pruning, speculative decoding, and more.
 | `meta-llama/Llama-4-{Scout,Maverick}-17B-{16,128}E-Instruct` | ✅ | ✅ | - | - |
 | `moonshotai/Kimi-K2-Instruct` | ✅ | ✅ | - | - |
 | `nvidia/NVIDIA-Nemotron-Nano-9B-v2` | ✅ | - | ✅ | ✅ |
+| `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` | ✅ | - | ✅ | ✅ |
 | `openai/gpt-oss-{20b, 120b}` | ✅ | **Online** | ✅ | ✅ |
 | `Qwen/Qwen3-{0.6B, 8B}` | ✅ | ✅ | ✅ | ✅ |
 | `Qwen/Qwen3-{30B-A3B, 235B-A22B}` | **WAR** | ✅ | ✅ | ✅ |
@@ -57,18 +58,19 @@ Provide the pretrained checkpoint path through variable `${HF_MODEL_CKPT}` and p
 Hugging Face-Like quantized checkpoint for TensorRT-LLM, vLLM, or SGLang deployement,
 provide `${EXPORT_DIR}` to `export.sh`.
 
-> **📙 NOTE:** ModelOpt supports different quantization formats. By default, we simulate the
-> low-precision numerical behavior (fake-quant) which can be run on GPUs with compute > 80.
+> **📙 NOTE:** ModelOpt supports different quantization formats which are listed in the [ModelOpt quant configs](https://github.com/NVIDIA/Model-Optimizer/blob/7971fff05882da7eae16eae6bc927d1481dcd63f/modelopt/torch/quantization/config.py#L626).
+> The quant config is specified by the full config name in all-caps, e.g. NVFP4_DEFAULT_CFG.
+> By default, we simulate the low-precision numerical behavior (fake-quant) which can be run on GPUs with compute > 80.
 > Real low-precision paramters (e.g. `E4M3` or `E2M1`)
 > and low-precision compute (e.g. `FP8Linear`) are also supported depending on GPU compute capability.
-> **See [Adanvanced Topics](./ADVANCED.md) for details**.
+> **See [Advanced Topics](./ADVANCED.md) for details**.
 
 ```sh
 \
     TP=1 \
     HF_MODEL_CKPT=<pretrained_model_name_or_path> \
     MLM_MODEL_SAVE=/tmp/Llama-3.2-1B-Instruct_quant \
-    ./quantize.sh meta-llama/Llama-3.2-1B-Instruct nvfp4
+    ./quantize.sh meta-llama/Llama-3.2-1B-Instruct NVFP4_DEFAULT_CFG 
 
 \
     PP=1 \
@@ -78,6 +80,8 @@ provide `${EXPORT_DIR}` to `export.sh`.
     ./export.sh meta-llama/Llama-3.2-1B-Instruct
 ```
 
+For KV cache quantization, add a flag like `MLM_EXTRA_ARGS="--export-kv-cache-quant fp8"` while specifying your desired KV cache precision (see `KV_QUANT_CFG_CHOICES` in `quantize.py`).
+
 ### ⭐ Online BF16 EAGLE3 Training
 
 Online EAGLE3 training has both the target (frozen) and draft models in the memory where the `hidden_states`
@@ -100,7 +104,7 @@ deployment.
     ./export.sh meta-llama/Llama-3.2-1B-Instruct
 ```
 
-See [Adanvanced Topics](./ADVANCED.md) for a `moonshotai/Kimi-K2-Instruct` EAGLE3 training example using `slurm`.
+See [Advanced Topics](./ADVANCED.md) for a `moonshotai/Kimi-K2-Instruct` EAGLE3 training example using `slurm`.
 
 ### ⭐ Offline BF16 EAGLE3 Training
 Unlike online EAGLE3 training, offline workflow precomputes target model `hidden_states` and dumps to disk.
diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh
new file mode 100644
index 00000000000..c294e03235c
--- /dev/null
+++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh
@@ -0,0 +1,57 @@
+#!/bin/bash
+
+if [ -z ${HF_MODEL_CKPT} ]; then
+    HF_MODEL_CKPT=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+    TOKENIZER_MODEL=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+else
+    TOKENIZER_MODEL=${HF_MODEL_CKPT}
+fi
+
+MODEL_ARGS=" \
+    --trust-remote-code \
+    --save-interval 100000 \
+    --micro-batch-size 1 \
+    --moe-token-dispatcher-type allgather \
+    --enable-experimental \
+    --moe-permute-fusion \
+    --use-fused-weighted-squared-relu \
+    --cross-entropy-loss-fusion \
+    --cross-entropy-fusion-impl native \
+    --moe-router-score-function sigmoid \
+    --moe-grouped-gemm \
+    --num-experts 128 \
+    --moe-router-topk 6 \
+    --moe-aux-loss-coeff 1e-4 \
+    --moe-router-topk-scaling-factor 2.5 \
+    --moe-router-enable-expert-bias \
+    --moe-router-dtype fp32 \
+    --moe-router-load-balancing-type seq_aux_loss \
+    --moe-shared-expert-intermediate-size 3712 \
+    \
+    --attention-backend flash \
+    --disable-gloo-process-groups \
+    --is-hybrid-model \
+    --mamba-num-heads 64 \
+    --mamba-head-dim 64 \
+    --hybrid-override-pattern MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME \
+    --use-mcore-models \
+    --untie-embeddings-and-output-weights \
+    --disable-bias-linear \
+    --init-method-std 0.0173 \
+    --position-embedding-type none \
+    --squared-relu \
+    --num-layers 52 \
+    --hidden-size 2688 \
+    --num-attention-heads 32 \
+    --group-query-attention \
+    --num-query-groups 2 \
+    --ffn-hidden-size 1856 \
+    --kv-channels 128 \
+    --normalization RMSNorm \
+    \
+    --tokenizer-type HuggingFaceTokenizer \
+    --bf16 \
+    --seq-length 8192 \
+    --max-position-embeddings 8192 \
+    --export-model-type MambaModel \
+    "
diff --git a/examples/post_training/modelopt/convert_model.py b/examples/post_training/modelopt/convert_model.py
index 53ae25f8d92..cf5f6e5bbbb 100644
--- a/examples/post_training/modelopt/convert_model.py
+++ b/examples/post_training/modelopt/convert_model.py
@@ -136,7 +136,10 @@ def check_arguments():
         print_rank_0(
             "Import model from Hugging Face checkpoint in dtype {}.".format(str(import_dtype))
         )
-        import_kwargs = {"dtype": import_dtype}
+        import_kwargs = {
+            "dtype": import_dtype,
+            "moe_router_dtype": args.moe_router_dtype,
+        }
         if modelopt_version_at_least("0.41.0"):
             import_kwargs.update({"trust_remote_code": args.trust_remote_code})
         import_mcore_gpt_from_hf(
diff --git a/examples/post_training/modelopt/export.py b/examples/post_training/modelopt/export.py
old mode 100644
new mode 100755
index 0aa625b875d..9dc66eecb6d
--- a/examples/post_training/modelopt/export.py
+++ b/examples/post_training/modelopt/export.py
@@ -5,6 +5,7 @@
 import os
 import sys
 import warnings
+from pathlib import Path
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))
 
@@ -74,8 +75,11 @@ def add_modelopt_export_args(parser):
     unwrapped_model = unwrap_model(model)[0]
     unwrapped_model.to_empty(device="cpu")
 
-    if args.load is not None:
+    if args.load is not None and Path(args.load).is_dir():
         _ = load_modelopt_checkpoint(model)
+    else:
+        raise ValueError(f"Invalid load checkpoint directory: {args.load}")
+
 
     # Decide whether we are exporting only the extra_modules (e.g. EAGLE3).
     # Only the last pp stage may have extra_modules, hence broadcast from the last rank.
@@ -90,6 +94,7 @@ def add_modelopt_export_args(parser):
         "export_extra_modules": export_extra_modules,
         "dtype": torch.bfloat16,
         "export_dir": args.export_dir,
+        "moe_router_dtype": unwrapped_model.config.moe_router_dtype,
     }
     if modelopt_version_at_least("0.41.0"):
         export_kwargs.update({"trust_remote_code": args.trust_remote_code})
diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py
index 635c18ee545..ceedce606a5 100644
--- a/examples/post_training/modelopt/quantize.py
+++ b/examples/post_training/modelopt/quantize.py
@@ -52,15 +52,9 @@
 
 warnings.filterwarnings("ignore")
 
-# TODO deprecate these aliases in the next release
-QUANT_CFG_CHOICES = {
-    "int8_sq": mtq.INT8_SMOOTHQUANT_CFG,
-    "fp8": mtq.FP8_DEFAULT_CFG,
-    "fp8_blockwise": mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
-    "int4_awq": mtq.INT4_AWQ_CFG,
-    "w4a8_awq": mtq.W4A8_AWQ_BETA_CFG,
-    "nvfp4": mtq.NVFP4_DEFAULT_CFG,
-}
+QUANT_CFG_CHOICES = {}
+
+# Auto-load all quant configs by full name
 for k in mtq.config.choices:
     QUANT_CFG_CHOICES[k] = getattr(mtq, k)
 
diff --git a/megatron/core/post_training/modelopt/mamba/model_specs.py b/megatron/core/post_training/modelopt/mamba/model_specs.py
index e8a14212bc3..0a38d05b980 100755
--- a/megatron/core/post_training/modelopt/mamba/model_specs.py
+++ b/megatron/core/post_training/modelopt/mamba/model_specs.py
@@ -2,6 +2,7 @@
 
 from megatron.core.extensions.transformer_engine import TEDotProductAttention
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.post_training.modelopt.layers import Norm
 from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
 from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
@@ -83,9 +84,23 @@ def get_mamba_stack_modelopt_spec(
         ),
     )
 
+    moe_layer = ModuleSpec(
+        module=TransformerLayer,
+        submodules=TransformerLayerSubmodules(
+            pre_mlp_layernorm=Norm,
+            mlp=get_moe_module_spec(
+                use_te=False, num_experts=8, moe_grouped_gemm=False  # Can be anything non None
+            ),
+            mlp_bda=get_bias_dropout_add,
+        ),
+    )
+
     return ModuleSpec(
         module=MambaStack,
         submodules=MambaStackSubmodules(
-            mamba_layer=mamba_layer, attention_layer=attention_layer, mlp_layer=mlp_layer
+            mamba_layer=mamba_layer,
+            attention_layer=attention_layer,
+            mlp_layer=mlp_layer,
+            moe_layer=moe_layer,
         ),
     )

From 703195367fe022e9c6efc8abf635179ca2de6c06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Tue, 27 Jan 2026 13:47:51 +0100
Subject: [PATCH 29/79] ci: Disable
 gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq (#3099)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .../gpt-dynamic-inference-with-coordinator.yaml      | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml
index e882d721860..9c5beb7e43a 100644
--- a/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml
+++ b/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml
@@ -3,7 +3,7 @@ format_version: 1
 maintainers: [mcore]
 loggers: [stdout]
 spec:
-  name: '{test_case}_{environment}_{platforms}'
+  name: "{test_case}_{environment}_{platforms}"
   model: gpt
   build: mcore-pyt-{environment}
   nodes: 1
@@ -69,11 +69,11 @@ products:
       - environment: [dev]
         scope: [mr]
         platforms: [dgx_h100]
-  - test_case: [gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq]
-    products:
-      - environment: [dev]
-        scope: [mr]
-        platforms: [dgx_h100]
+  # - test_case: [gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq]
+  #   products:
+  #     - environment: [dev]
+  #       scope: [mr]
+  #       platforms: [dgx_h100]
   - test_case: [gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq]
     products:
       - environment: [dev]

From dea21a0af1ce235fa1b2ddb8051232de1d0d643c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Tue, 27 Jan 2026 14:38:08 +0100
Subject: [PATCH 30/79] ci: Repeat func tests, save logs of unit tests and
 lessen debug output (#3089)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/actions/action.yml                                | 5 ++++-
 .../test_utils/python_scripts/launch_nemo_run_workload.py | 8 +++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/.github/actions/action.yml b/.github/actions/action.yml
index dfc6d79688e..6e9c72016f6 100644
--- a/.github/actions/action.yml
+++ b/.github/actions/action.yml
@@ -149,14 +149,17 @@ runs:
           ARGS=(
             --scope mr-github
             --enable-lightweight-mode
+            --n-repeat 1
           )
         elif [ "${{ steps.has-run-functional-tests-label.outputs.main }}" == "true" ]; then
           ARGS=(
             --scope mr-github
+            --n-repeat 5
           )
         else
           ARGS=(
             --scope mr-github-slim
+            --n-repeat 5
           )
         fi
 
@@ -258,5 +261,5 @@ runs:
       if: always()
       with:
         name: ${{ steps.check.outputs.logs_report }}
-        path: ${{ inputs.is_unit_test == 'true' && 'logs' || 'assets_dir' }}
+        path: ${{ inputs.is_unit_test == 'true' && 'assets_dir/logs' || 'assets_dir' }}
         include-hidden-files: true
diff --git a/tests/test_utils/python_scripts/launch_nemo_run_workload.py b/tests/test_utils/python_scripts/launch_nemo_run_workload.py
index 26a7dbd79f5..8d006f70d19 100644
--- a/tests/test_utils/python_scripts/launch_nemo_run_workload.py
+++ b/tests/test_utils/python_scripts/launch_nemo_run_workload.py
@@ -50,6 +50,9 @@ def is_flaky_failure(concat_allranks_logs: str) -> bool:
 @click.option("--environment", required=True, type=str, help="Environment of the workload")
 @click.option("--platform", required=True, type=str, help="Platform of the workload")
 @click.option("--container-image", required=True, type=str, help="Container image of the workload")
+@click.option(
+    "--n-repeat", required=False, type=int, help="Number of times to repeat the workload", default=1
+)
 @click.option("--data-dir", required=False, type=str, help="Data directory of the workload")
 @click.option("--tag", required=False, type=str, help="Tag of the workload")
 @click.option(
@@ -68,6 +71,7 @@ def main(
     environment,
     platform,
     container_image,
+    n_repeat: int = 1,
     data_dir: Optional[str] = None,
     tag: Optional[str] = None,
     enable_lightweight_mode: Optional[bool] = False,
@@ -92,6 +96,7 @@ def main(
     magic_values["assets_dir"] = "/opt/megatron-lm/assets_dir"
     magic_values["artifacts_dir"] = "/opt/megatron-lm/artifacts_dir"
     magic_values["environment"] = environment
+    magic_values["n_repeat"] = n_repeat
     magic_values["test_case"] = workload.spec["test_case"]
     magic_values["name"] = workload.spec["name"].format(**magic_values)
     workload.spec["script"] = workload.spec["script"].format(**magic_values)
@@ -113,9 +118,10 @@ def main(
             "PYTHONUNBUFFERED": "1",
             "OUTPUT_PATH": os.getcwd(),
             "ENABLE_LIGHTWEIGHT_MODE": str(enable_lightweight_mode).lower(),
-            "N_REPEAT": "1",
+            "N_REPEAT": str(n_repeat),
             "CLUSTER": "dgxh100_dgxc",
             "NCCL_DEBUG": "INFO",
+            "NCCL_DEBUG_FILE": "/opt/megatron-lm/assets_dir/logs/nccl_debug.log",
         },
         packager=run.Packager(),
         volumes=artifacts,

From dd83fc6822064f6ec238be9fe16dd6096d544189 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Tue, 27 Jan 2026 15:59:36 +0100
Subject: [PATCH 31/79] ci: Update improvement of step-time (#3104)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .../golden_values_dev_dgx_h100.json           | 28 +++++++++----------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
index f2b6084c49b..c5fdf3beffa 100644
--- a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
@@ -28,11 +28,11 @@
         "end_step": 5,
         "step_interval": 1,
         "values": {
-            "1": 76714901504.0,
-            "2": 76724633600.0,
-            "3": 76724633600.0,
-            "4": 76724633600.0,
-            "5": 76724633600.0
+            "1": 76691947520.0,
+            "2": 76708724736.0,
+            "3": 76708724736.0,
+            "4": 76708724736.0,
+            "5": 76708724736.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -40,11 +40,11 @@
         "end_step": 5,
         "step_interval": 1,
         "values": {
-            "1": 76714909696.0,
-            "2": 77061054464.0,
-            "3": 77061103616.0,
-            "4": 77061226496.0,
-            "5": 77061226496.0
+            "1": 76691955712.0,
+            "2": 77045972992.0,
+            "3": 77046243328.0,
+            "4": 77047095296.0,
+            "5": 77047095296.0
         }
     },
     "iteration-time": {
@@ -53,10 +53,10 @@
         "step_interval": 1,
         "values": {
             "1": "nan",
-            "2": 121.41938,
-            "3": 88.73186,
-            "4": 93.15825,
-            "5": 91.09737
+            "2": 135.42645,
+            "3": 78.78998,
+            "4": 79.18825,
+            "5": 80.10109
         }
     }
 }
\ No newline at end of file

From 2bdf7e14b76a4cf6cd30a9d0fa939608d66652c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Tue, 27 Jan 2026 16:01:14 +0100
Subject: [PATCH 32/79] ci: Add GPU health checks (#3100)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/actions/action.yml | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/.github/actions/action.yml b/.github/actions/action.yml
index 6e9c72016f6..f3e42e5843d 100644
--- a/.github/actions/action.yml
+++ b/.github/actions/action.yml
@@ -56,6 +56,38 @@ runs:
       shell: bash -x -e -u -o pipefail {0}
       run: echo "node_name=$NODE_NAME" | tee -a "$GITHUB_OUTPUT"
 
+    - name: GPU Sanity Check
+      shell: bash -x -e -u -o pipefail {0}
+      run: |
+        echo "Starting GPU Sanity Check..."
+
+        # 1. Check for active Compute Processes
+        # query-compute-apps returns a list of PIDs using the GPU. If empty, we are good.
+        OPEN_PROCESSES=$(docker run --rm --gpus all ubuntu nvidia-smi --query-compute-apps=pid,process_name --format=csv,noheader)
+
+        if [ -n "$OPEN_PROCESSES" ]; then
+          echo "::error::❌ GPU is not clean! Found active processes:"
+          echo "$OPEN_PROCESSES"
+        else
+          echo "✅ No active compute processes found."
+        fi
+
+        # 2. Check VRAM Usage (Optional but recommended)
+        # We allow a small buffer (e.g., < 300MiB) for driver overhead/Xorg, 
+        # though on headless K8s nodes this should be very close to 0.
+
+        MEMORY_USAGES=$(docker run --rm --gpus all ubuntu nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)
+
+        # Check each GPU visible to the container
+        for MEMORY in $MEMORY_USAGES; do
+          if [ "$MEMORY" -gt 300 ]; then
+            echo "::error::❌ GPU VRAM usage is suspiciously high: ${MEMORY} MiB"
+          fi
+        done
+
+        echo "✅ GPU Memory is clean (all < 300 MiB)."
+        echo "Ready to start workflow."
+
     - name: Checkout repository
       uses: actions/checkout@v2
 

From d68721bf3ce4690e34c3212b2698dcfa5db1d8a7 Mon Sep 17 00:00:00 2001
From: Jon Barker <jbarker@nvidia.com>
Date: Tue, 27 Jan 2026 09:08:48 -0700
Subject: [PATCH 33/79] Harden GRPO functional tests (#3065)

Co-authored-by: Jon Barker <jbarker@oci-hsg-cs-001-vscode-01.cm.cluster>
---
 .../compute_golden_statistics.py              | 836 ++++++++++++++++++
 .../test_grpo_training_loop.py                | 180 +++-
 .../golden_values_dev_dgx_h100.json           | 173 ----
 .../golden_values_dev_dgx_h100.json           | 287 ------
 .../model_config.yaml                         |  78 --
 .../env_config.yaml                           |   0
 .../golden_values_dev_dgx_h100.json           |  83 ++
 .../model_config.yaml                         |  92 +-
 .../env_config.yaml                           |   0
 .../golden_values_dev_dgx_h100.json           |  83 ++
 .../model_config.yaml                         | 103 +++
 tests/test_utils/recipes/gpt-grpo.yaml        |  16 +-
 12 files changed, 1322 insertions(+), 609 deletions(-)
 create mode 100644 tests/functional_tests/python_test_utils/compute_golden_statistics.py
 delete mode 100644 tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json
 delete mode 100644 tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json
 delete mode 100644 tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/model_config.yaml
 rename tests/functional_tests/test_cases/gpt/{gpt_grpo_tp1_pp1_dp8_583m_throughputtest => gpt_grpo_tp4_pp1_dp2_8b_throughput}/env_config.yaml (100%)
 create mode 100644 tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/golden_values_dev_dgx_h100.json
 rename tests/functional_tests/test_cases/gpt/{gpt_grpo_tp1_pp1_dp8_583m_throughputtest => gpt_grpo_tp4_pp1_dp2_8b_throughput}/model_config.yaml (50%)
 rename tests/functional_tests/test_cases/gpt/{gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github => gpt_grpo_tp4_pp1_dp2_8b_throughput_github}/env_config.yaml (100%)
 create mode 100644 tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/golden_values_dev_dgx_h100.json
 create mode 100644 tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml

diff --git a/tests/functional_tests/python_test_utils/compute_golden_statistics.py b/tests/functional_tests/python_test_utils/compute_golden_statistics.py
new file mode 100644
index 00000000000..d4863fa9476
--- /dev/null
+++ b/tests/functional_tests/python_test_utils/compute_golden_statistics.py
@@ -0,0 +1,836 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""
+Compute statistical bounds for golden values from multiple test runs.
+
+This script aggregates results from multiple parallel runs of a functional test
+and computes statistics (min, max, mean, std) for each metric at each step.
+The output can be used to determine appropriate tolerances for test validation.
+
+Usage:
+    # Step 1: Run batch tests (from megatron-rl directory):
+    ./tests/functional_tests/shell_test_utils/run_batch_ci_tests.sh \\
+        test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_correctness_and_throughput.sh 10
+
+    # Step 2: Wait for jobs to complete, then compute statistics:
+    python tests/functional_tests/python_test_utils/compute_golden_statistics.py \\
+        --results-dir batch_test_logs_gpt_grpo_*/ \\
+        --output golden_values_stats.json \\
+        --recommend-tolerances
+
+    # The script parses .out log files to find where each run wrote its results.
+    # Each .out file should contain: "This test wrote results into /opt/megatron-lm/runs/<uuid>"
+    # The container path /opt/megatron-lm maps to the workspace root on the host.
+
+    # Or specify individual JSON files directly:
+    python compute_golden_statistics.py \\
+        --result-files runs/abc123/golden_values.json runs/def456/golden_values.json \\
+        --output golden_values_stats.json
+"""
+
+import argparse
+import glob
+import json
+import logging
+import math
+import os
+import sys
+from pathlib import Path
+from statistics import mean, median, stdev
+from typing import Any, Dict, List, Optional, Tuple
+
+logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
+logger = logging.getLogger(__name__)
+
+
+def find_result_json_files(results_dir: str, workspace_root: Optional[str] = None) -> List[str]:
+    """
+    Find all result JSON files from a batch test run.
+
+    The batch test infrastructure (run_batch_ci_tests.sh) writes .out log files
+    to the results directory. Each .out file contains a line like:
+        "This test wrote results into /opt/megatron-lm/runs/<uuid>"
+
+    The container path /opt/megatron-lm maps to the workspace root on the host.
+    This function parses the .out files to find where the JSON results are.
+
+    Args:
+        results_dir: Path to batch_test_logs_* directory containing .out files
+        workspace_root: Root of the megatron workspace (defaults to cwd)
+    """
+    result_files = []
+    results_path = Path(results_dir)
+
+    if not results_path.exists():
+        logger.error(f"Results directory not found: {results_dir}")
+        return []
+
+    if workspace_root is None:
+        # Try to find workspace root by looking for common markers
+        workspace_root = os.getcwd()
+
+    # Find all .out files from batch test runs
+    out_files = list(results_path.glob("*.out"))
+
+    if not out_files:
+        logger.warning(f"No .out files found in {results_dir}")
+        # Fall back to searching for JSON files directly
+        return _find_json_files_directly(results_dir)
+
+    logger.info(f"Found {len(out_files)} .out files to parse")
+
+    for out_file in out_files:
+        json_path = _extract_result_path_from_log(out_file, workspace_root)
+        if json_path and os.path.exists(json_path):
+            result_files.append(json_path)
+        elif json_path:
+            logger.warning(f"Result file not found: {json_path} (from {out_file.name})")
+
+    return result_files
+
+
+def _extract_result_path_from_log(out_file: Path, workspace_root: str) -> Optional[str]:
+    """
+    Parse a .out log file to find the result JSON path.
+
+    Looks for the line: "This test wrote results into /opt/megatron-lm/runs/<uuid>"
+    and converts the container path to the host path.
+    """
+    try:
+        with open(out_file, 'r', errors='ignore') as f:
+            content = f.read()
+    except IOError as e:
+        logger.warning(f"Failed to read {out_file}: {e}")
+        return None
+
+    # Look for the output path marker
+    marker = "This test wrote results into "
+    for line in content.split('\n'):
+        if marker in line:
+            # Extract the path after the marker
+            idx = line.find(marker)
+            output_path = line[idx + len(marker) :].strip()
+
+            # Convert container path to host path
+            # /opt/megatron-lm/... -> <workspace_root>/...
+            if output_path.startswith("/opt/megatron-lm/"):
+                host_path = output_path.replace("/opt/megatron-lm/", "")
+                output_path = os.path.join(workspace_root, host_path)
+
+            # Find JSON result files in this directory (search recursively)
+            output_dir = Path(output_path)
+            if output_dir.exists() and output_dir.is_dir():
+                # Look for result JSON files with various naming patterns
+                # Search recursively since files may be in subdirectories (e.g., 1/, 2/)
+                patterns = [
+                    "**/golden_values*.json",
+                    "**/generations*.json",
+                    "**/test_results*.json",
+                ]
+
+                for pattern in patterns:
+                    json_files = list(output_dir.glob(pattern))
+                    if json_files:
+                        # Return the first match
+                        logger.debug(f"Found result file: {json_files[0]}")
+                        return str(json_files[0])
+
+                # Fallback: any JSON file in subdirectories
+                json_files = list(output_dir.glob("**/*.json"))
+                if json_files:
+                    logger.debug(f"Found result file (fallback): {json_files[0]}")
+                    return str(json_files[0])
+
+            logger.debug(f"Output directory not found or empty: {output_path}")
+            return None
+
+    logger.debug(f"No output path marker found in {out_file.name}")
+    return None
+
+
+def _find_json_files_directly(results_dir: str) -> List[str]:
+    """
+    Fallback: search for JSON files directly in the results directory.
+
+    This is used when .out files don't contain the expected markers.
+    """
+    result_files = []
+    results_path = Path(results_dir)
+
+    # Look for golden_values*.json files in subdirectories
+    patterns = ["**/golden_values*.json", "**/test_results*.json", "**/*_output.json"]
+
+    for pattern in patterns:
+        matches = list(results_path.glob(pattern))
+        result_files.extend([str(p) for p in matches])
+
+    # Remove duplicates while preserving order
+    seen = set()
+    unique_files = []
+    for f in result_files:
+        if f not in seen:
+            seen.add(f)
+            unique_files.append(f)
+
+    return unique_files
+
+
+def load_result_file(filepath: str) -> Optional[Dict[str, Any]]:
+    """Load a single result JSON file."""
+    try:
+        with open(filepath, 'r') as f:
+            content = f.read()
+
+        data = json.loads(content)
+
+        # Handle JSONL format (single line)
+        if isinstance(data, str):
+            data = json.loads(data)
+
+        return data
+    except (json.JSONDecodeError, IOError) as e:
+        logger.warning(f"Failed to load {filepath}: {e}")
+        return None
+
+
+def _detect_result_format(data: Dict[str, Any]) -> str:
+    """
+    Detect whether the result file is from a training test or inference test.
+
+    Returns:
+        "training" - TensorBoard metrics format: {"metric_name": {"values": {...}}}
+        "inference" - Generation output format: {"request_id": {"latency": ..., ...}}
+        "unknown" - Unrecognized format
+    """
+    if not data:
+        return "unknown"
+
+    # Check first key's value structure
+    first_key = next(iter(data.keys()))
+    first_value = data[first_key]
+
+    if isinstance(first_value, dict):
+        if 'values' in first_value:
+            return "training"
+        if 'latency' in first_value or 'generated_text' in first_value:
+            return "inference"
+
+    return "unknown"
+
+
+def _is_valid_numeric(value) -> bool:
+    """Check if a value is a valid (non-NaN) numeric value."""
+    if isinstance(value, str):
+        try:
+            value = float(value)
+        except ValueError:
+            return False
+
+    if isinstance(value, (int, float)):
+        return not math.isnan(value)
+
+    return False
+
+
+def _to_float(value) -> Optional[float]:
+    """Convert value to float, returning None for invalid/NaN values."""
+    if isinstance(value, str):
+        try:
+            value = float(value)
+        except ValueError:
+            return None
+
+    if isinstance(value, (int, float)):
+        if math.isnan(value):
+            return None
+        return float(value)
+
+    return None
+
+
+def _aggregate_training_results(
+    data: Dict[str, Any], aggregated: Dict[str, Dict[str, List[float]]], run_index: int
+) -> None:
+    """Aggregate results from training test format."""
+    for metric_name, metric_data in data.items():
+        if not isinstance(metric_data, dict) or 'values' not in metric_data:
+            continue
+
+        if metric_name not in aggregated:
+            aggregated[metric_name] = {}
+
+        values = metric_data['values']
+        for step, value in values.items():
+            # Skip non-numeric or NaN values
+            float_val = _to_float(value)
+            if float_val is None:
+                continue
+
+            if step not in aggregated[metric_name]:
+                aggregated[metric_name][step] = []
+
+            aggregated[metric_name][step].append(float_val)
+
+        # For metrics that use median-based comparison in the test (iteration-time,
+        # mem-allocated-bytes, mem-max-allocated-bytes), also store all values from
+        # this run so we can compute per-run medians later.
+        # IMPORTANT: Store values in step order to match the test's index-based slicing.
+        if metric_name in ['iteration-time', 'mem-allocated-bytes', 'mem-max-allocated-bytes']:
+            all_values_key = f"_all_values_run_{run_index}"
+            if all_values_key not in aggregated[metric_name]:
+                aggregated[metric_name][all_values_key] = []
+
+            # Sort by step number to ensure consistent ordering for index-based slicing
+            sorted_steps = sorted(
+                values.keys(), key=lambda x: int(x) if x.isdigit() else float('inf')
+            )
+            for step in sorted_steps:
+                float_val = _to_float(values[step])
+                if float_val is None:
+                    continue
+                aggregated[metric_name][all_values_key].append(
+                    float_val
+                )  # Just the value, not tuple
+
+
+def _aggregate_inference_results(
+    data: Dict[str, Any], aggregated: Dict[str, Dict[str, List[float]]], run_index: int
+) -> None:
+    """
+    Aggregate results from inference test format.
+
+    Extracts metrics like latency, step_count, and logprob statistics
+    from generation outputs.
+    """
+    # Metrics to extract per request
+    latencies = []
+    step_counts = []
+    prompt_logprob_means = []
+    generated_logprob_means = []
+
+    for request_id, request_data in data.items():
+        if not isinstance(request_data, dict):
+            continue
+
+        # Extract latency
+        if 'latency' in request_data:
+            latencies.append(float(request_data['latency']))
+
+        # Extract step count
+        if 'step_count' in request_data:
+            step_counts.append(float(request_data['step_count']))
+
+        # Extract mean of prompt logprobs (as a consistency metric)
+        if 'prompt_logprobs' in request_data and request_data['prompt_logprobs']:
+            logprobs = request_data['prompt_logprobs']
+            if isinstance(logprobs, list) and len(logprobs) > 0:
+                prompt_logprob_means.append(sum(logprobs) / len(logprobs))
+
+        # Extract mean of generated logprobs
+        if 'generated_log_probs' in request_data and request_data['generated_log_probs']:
+            logprobs = request_data['generated_log_probs']
+            if isinstance(logprobs, list) and len(logprobs) > 0:
+                generated_logprob_means.append(sum(logprobs) / len(logprobs))
+
+    # Store aggregated metrics using run_index as the "step"
+    run_key = str(run_index)
+
+    if latencies:
+        if 'latency' not in aggregated:
+            aggregated['latency'] = {}
+        if 'mean' not in aggregated['latency']:
+            aggregated['latency']['mean'] = []
+        aggregated['latency']['mean'].append(sum(latencies) / len(latencies))
+
+        if 'total' not in aggregated['latency']:
+            aggregated['latency']['total'] = []
+        aggregated['latency']['total'].append(sum(latencies))
+
+    if step_counts:
+        if 'step_count' not in aggregated:
+            aggregated['step_count'] = {}
+        if 'mean' not in aggregated['step_count']:
+            aggregated['step_count']['mean'] = []
+        aggregated['step_count']['mean'].append(sum(step_counts) / len(step_counts))
+
+    if prompt_logprob_means:
+        if 'prompt_logprob_mean' not in aggregated:
+            aggregated['prompt_logprob_mean'] = {}
+        if 'mean' not in aggregated['prompt_logprob_mean']:
+            aggregated['prompt_logprob_mean']['mean'] = []
+        aggregated['prompt_logprob_mean']['mean'].append(
+            sum(prompt_logprob_means) / len(prompt_logprob_means)
+        )
+
+    if generated_logprob_means:
+        if 'generated_logprob_mean' not in aggregated:
+            aggregated['generated_logprob_mean'] = {}
+        if 'mean' not in aggregated['generated_logprob_mean']:
+            aggregated['generated_logprob_mean']['mean'] = []
+        aggregated['generated_logprob_mean']['mean'].append(
+            sum(generated_logprob_means) / len(generated_logprob_means)
+        )
+
+
+def aggregate_results(result_files: List[str]) -> Dict[str, Dict[str, List[float]]]:
+    """
+    Aggregate results from multiple JSON files.
+
+    Supports both training test format (TensorBoard metrics) and
+    inference test format (generation outputs).
+
+    Returns:
+        Dict mapping metric_name -> step/key -> list of values across all runs
+    """
+    aggregated: Dict[str, Dict[str, List[float]]] = {}
+    loaded_count = 0
+    detected_format = None
+
+    for idx, filepath in enumerate(result_files):
+        data = load_result_file(filepath)
+        if data is None:
+            continue
+
+        loaded_count += 1
+
+        # Detect format from first file
+        file_format = _detect_result_format(data)
+        if detected_format is None:
+            detected_format = file_format
+            logger.info(f"Detected result format: {file_format}")
+
+        if file_format == "training":
+            _aggregate_training_results(data, aggregated, idx)
+        elif file_format == "inference":
+            _aggregate_inference_results(data, aggregated, idx)
+        else:
+            logger.warning(f"Unknown format in {filepath}, skipping")
+
+    logger.info(f"Successfully loaded {loaded_count} of {len(result_files)} result files")
+    return aggregated
+
+
+def compute_statistics(aggregated: Dict[str, Dict[str, List[float]]]) -> Dict[str, Any]:
+    """
+    Compute statistics for each metric at each step.
+
+    Returns:
+        Dict with structure:
+        {
+            "metric_name": {
+                "num_samples": N,
+                "values": {
+                    "step": {
+                        "min": ...,
+                        "max": ...,
+                        "mean": ...,
+                        "std": ...,
+                        "samples": [...]  # original values
+                    }
+                }
+            }
+        }
+    """
+    stats: Dict[str, Any] = {}
+
+    for metric_name, step_values in aggregated.items():
+        # Determine number of samples (should be consistent across steps)
+        # Skip internal keys used for median calculations
+        regular_steps = {k: v for k, v in step_values.items() if not k.startswith("_")}
+        sample_counts = [len(vals) for vals in regular_steps.values()]
+        num_samples = max(sample_counts) if sample_counts else 0
+
+        metric_stats = {"num_samples": num_samples, "values": {}}
+
+        for step, values in regular_steps.items():
+            if len(values) == 0:
+                continue
+
+            step_stats = {
+                "min": min(values),
+                "max": max(values),
+                "mean": mean(values),
+                "std": stdev(values) if len(values) > 1 else 0.0,
+                "count": len(values),
+            }
+
+            # Include original samples for debugging
+            step_stats["samples"] = values
+
+            metric_stats["values"][step] = step_stats
+
+        stats[metric_name] = metric_stats
+
+    return stats
+
+
+def compute_recommended_tolerances(
+    stats: Dict[str, Any],
+    aggregated: Dict[str, Dict[str, List[float]]],
+    confidence_multiplier: float = 3.0,
+    start_step: int = 1,
+) -> Dict[str, Dict[str, float]]:
+    """
+    Compute recommended tolerances for each metric based on observed variance.
+
+    For metrics that use median-based comparison in the test (iteration-time,
+    mem-allocated-bytes, mem-max-allocated-bytes), computes variance of per-run
+    medians rather than per-step variance.
+
+    Args:
+        stats: Output from compute_statistics()
+        aggregated: Raw aggregated data (needed for median calculations)
+        confidence_multiplier: Number of standard deviations for bounds (default 3.0 for ~99.7% coverage)
+        start_step: First step to include in tolerance calculation (skips warmup steps)
+
+    Returns:
+        Dict mapping metric_name -> {
+            "relative_tolerance": recommended relative tolerance,
+            "absolute_tolerance": recommended absolute tolerance (for near-zero values),
+            "max_observed_relative_variance": max(|value - mean| / |mean|) across all samples
+        }
+    """
+    tolerances = {}
+
+    # Metrics that use median-based comparison in the test (iteration-time)
+    median_based_metrics = ['iteration-time']
+    # Metrics that use max-based comparison in the test (memory)
+    max_based_metrics = ['mem-allocated-bytes', 'mem-max-allocated-bytes']
+
+    for metric_name, metric_data in stats.items():
+        max_relative_variance = 0.0
+        max_absolute_variance = 0.0
+        steps_included = 0
+
+        # For median-based metrics, compute variance of per-run medians
+        if metric_name in median_based_metrics and metric_name in aggregated:
+            run_medians = []
+
+            # Find all run data keys
+            for key in aggregated[metric_name].keys():
+                if key.startswith("_all_values_run_"):
+                    run_data = aggregated[metric_name][key]
+                    # Use index-based slicing to match test behavior:
+                    # [start_step:] skips the first `start_step` items
+                    filtered_values = run_data[start_step:]
+
+                    if filtered_values:
+                        run_median = median(filtered_values)
+                        run_medians.append(run_median)
+
+            if run_medians:
+                median_mean = mean(run_medians)
+
+                # Compute relative variance of medians
+                if abs(median_mean) > 1e-9:
+                    for m in run_medians:
+                        rel_var = abs(m - median_mean) / abs(median_mean)
+                        max_relative_variance = max(max_relative_variance, rel_var)
+                else:
+                    for m in run_medians:
+                        max_absolute_variance = max(max_absolute_variance, abs(m))
+
+                steps_included = len(run_medians)
+
+                logger.debug(
+                    f"{metric_name}: computed variance from {len(run_medians)} run medians, "
+                    f"mean={median_mean:.4f}, max_rel_var={max_relative_variance:.4%}"
+                )
+
+        # For max-based metrics (memory), compute variance of per-run max values
+        elif metric_name in max_based_metrics and metric_name in aggregated:
+            run_maxes = []
+
+            # Find all run data keys
+            for key in aggregated[metric_name].keys():
+                if key.startswith("_all_values_run_"):
+                    run_data = aggregated[metric_name][key]
+                    # Skip first value (warmup), take max of rest
+                    filtered_values = run_data[1:] if len(run_data) > 1 else run_data
+
+                    if filtered_values:
+                        run_max = max(filtered_values)
+                        run_maxes.append(run_max)
+
+            if run_maxes:
+                max_mean = mean(run_maxes)
+
+                # Compute relative variance of max values
+                if abs(max_mean) > 1e-9:
+                    for m in run_maxes:
+                        rel_var = abs(m - max_mean) / abs(max_mean)
+                        max_relative_variance = max(max_relative_variance, rel_var)
+                else:
+                    for m in run_maxes:
+                        max_absolute_variance = max(max_absolute_variance, abs(m))
+
+                steps_included = len(run_maxes)
+
+                logger.debug(
+                    f"{metric_name}: computed variance from {len(run_maxes)} run maxes, "
+                    f"mean={max_mean:.4f}, max_rel_var={max_relative_variance:.4%}"
+                )
+        else:
+            # Standard per-step variance calculation for other metrics
+            for step, step_stats in metric_data["values"].items():
+                # Skip warmup steps - try to parse step as int, skip if < start_step
+                try:
+                    step_num = int(step)
+                    if step_num < start_step:
+                        continue
+                except (ValueError, TypeError):
+                    # Non-numeric step key (e.g., "mean" for inference metrics) - include it
+                    pass
+
+                steps_included += 1
+                mean_val = step_stats["mean"]
+
+                # Compute observed relative variance
+                if abs(mean_val) > 1e-9:
+                    # For non-zero means, compute relative variance
+                    for sample in step_stats["samples"]:
+                        rel_var = abs(sample - mean_val) / abs(mean_val)
+                        max_relative_variance = max(max_relative_variance, rel_var)
+                else:
+                    # For near-zero means, track absolute variance
+                    for sample in step_stats["samples"]:
+                        max_absolute_variance = max(max_absolute_variance, abs(sample))
+
+        # Recommend tolerance with safety margin
+        # Use observed variance * confidence_multiplier, with a minimum of 0.1%
+        recommended_relative = max(max_relative_variance * confidence_multiplier, 0.001)
+
+        # Round to reasonable precision
+        recommended_relative = round(recommended_relative, 4)
+
+        tolerances[metric_name] = {
+            "relative_tolerance": recommended_relative,
+            "absolute_tolerance": max(max_absolute_variance * confidence_multiplier, 1e-6),
+            "max_observed_relative_variance": round(max_relative_variance, 6),
+            "max_observed_absolute_variance": round(max_absolute_variance, 6),
+            "steps_included": steps_included,
+        }
+
+    return tolerances
+
+
+def format_summary(stats: Dict[str, Any], tolerances: Dict[str, Dict[str, float]]) -> str:
+    """Format a human-readable summary of the statistics."""
+    lines = []
+    lines.append("=" * 70)
+    lines.append("Golden Values Statistics Summary")
+    lines.append("=" * 70)
+
+    for metric_name in sorted(stats.keys()):
+        metric_data = stats[metric_name]
+        tol = tolerances.get(metric_name, {})
+
+        lines.append(f"\n{metric_name}:")
+        lines.append(f"  Samples: {metric_data['num_samples']}")
+        lines.append(f"  Steps: {len(metric_data['values'])}")
+
+        if tol:
+            lines.append(
+                f"  Max observed relative variance: {tol.get('max_observed_relative_variance', 'N/A'):.4%}"
+            )
+            lines.append(
+                f"  Recommended relative tolerance: {tol.get('relative_tolerance', 'N/A'):.2%}"
+            )
+            lines.append(
+                f"  Recommended absolute tolerance: {tol.get('absolute_tolerance', 'N/A'):.2e}"
+            )
+
+        # Show a few example steps
+        values = metric_data["values"]
+        example_steps = list(values.keys())[:3]
+        if example_steps:
+            lines.append("  Example steps:")
+            for step in example_steps:
+                s = values[step]
+                lines.append(
+                    f"    Step {step}: mean={s['mean']:.6g}, std={s['std']:.6g}, "
+                    f"range=[{s['min']:.6g}, {s['max']:.6g}]"
+                )
+
+    lines.append("\n" + "=" * 70)
+    return "\n".join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute statistical bounds for golden values from multiple test runs.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+
+    input_group = parser.add_mutually_exclusive_group(required=True)
+    input_group.add_argument(
+        "--results-dir",
+        type=str,
+        help="Directory containing batch test results (searches for JSON files)",
+    )
+    input_group.add_argument(
+        "--result-files",
+        type=str,
+        nargs="+",
+        help="Explicit list of result JSON files to aggregate",
+    )
+
+    parser.add_argument(
+        "--output", "-o", type=str, required=True, help="Output path for statistics JSON file"
+    )
+
+    parser.add_argument(
+        "--recommend-tolerances",
+        action="store_true",
+        help="Compute and display recommended tolerances based on observed variance",
+    )
+
+    parser.add_argument(
+        "--confidence-multiplier",
+        type=float,
+        default=1.5,
+        help="Multiplier for observed max variance when computing recommended tolerance. "
+        "Example: if max observed variance is 5%% and multiplier is 1.5, recommended tolerance is 7.5%%. "
+        "Use higher values (2-3) for more safety margin. Default: 1.5",
+    )
+
+    parser.add_argument(
+        "--min-samples",
+        type=int,
+        default=2,
+        help="Minimum number of samples required to compute statistics (default: 2)",
+    )
+
+    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose output")
+
+    parser.add_argument(
+        "--workspace-root",
+        type=str,
+        default=None,
+        help="Root of the megatron workspace (where runs/ directory is located). "
+        "Defaults to current working directory.",
+    )
+
+    parser.add_argument(
+        "--start-step",
+        type=int,
+        default=0,
+        help="Number of initial steps to skip (index-based, matching test behavior). "
+        "Uses Python slicing [start_step:] so --start-step 10 skips first 10 items. "
+        "Default: 0 (include all). Set to match THROUGHPUT_TEST_PARAMS.--start_step from model_config.yaml.",
+    )
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+
+    # Find or use result files
+    if args.results_dir:
+        result_files = find_result_json_files(args.results_dir, args.workspace_root)
+        if not result_files:
+            logger.error(f"No result JSON files found in {args.results_dir}")
+            logger.info("Make sure the batch tests have completed and results are available.")
+            logger.info(
+                "The script looks for .out files and parses them to find the result JSON paths."
+            )
+            logger.info(
+                "Each .out file should contain: 'This test wrote results into /opt/megatron-lm/runs/<uuid>'"
+            )
+            sys.exit(1)
+        logger.info(f"Found {len(result_files)} result files from {args.results_dir}")
+    else:
+        result_files = args.result_files
+        # Verify files exist
+        for f in result_files:
+            if not os.path.exists(f):
+                logger.error(f"Result file not found: {f}")
+                sys.exit(1)
+
+    if args.verbose:
+        for f in result_files:
+            logger.debug(f"  - {f}")
+
+    # Aggregate results
+    aggregated = aggregate_results(result_files)
+
+    if not aggregated:
+        logger.error("No valid results found to aggregate")
+        sys.exit(1)
+
+    # Check minimum samples
+    for metric_name, step_values in aggregated.items():
+        for step, values in step_values.items():
+            if len(values) < args.min_samples:
+                logger.warning(
+                    f"{metric_name} step {step}: only {len(values)} samples "
+                    f"(minimum {args.min_samples} recommended)"
+                )
+
+    # Compute statistics
+    stats = compute_statistics(aggregated)
+
+    # Compute recommended tolerances (excluding warmup steps)
+    if args.start_step > 1:
+        logger.info(f"Excluding steps < {args.start_step} from tolerance calculation (warmup)")
+    tolerances = compute_recommended_tolerances(
+        stats, aggregated, args.confidence_multiplier, start_step=args.start_step
+    )
+
+    # Build output
+    output = {
+        "metadata": {
+            "num_runs": len(result_files),
+            "result_files": result_files,
+            "confidence_multiplier": args.confidence_multiplier,
+            "start_step": args.start_step,
+        },
+        "statistics": stats,
+        "recommended_tolerances": tolerances,
+    }
+
+    # Write output
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(output_path, 'w') as f:
+        json.dump(output, f, indent=2)
+
+    logger.info(f"Statistics written to {args.output}")
+
+    # Print summary
+    if args.recommend_tolerances or args.verbose:
+        print(format_summary(stats, tolerances))
+
+        print("\nRecommended tolerance settings:")
+        print("-" * 50)
+        # Training test metrics
+        training_metrics = [
+            "lm-loss",
+            "lm loss",
+            "iteration-time",
+            "mem-allocated-bytes",
+            "mem-max-allocated-bytes",
+        ]
+        # Inference test metrics
+        inference_metrics = [
+            "latency",
+            "step_count",
+            "prompt_logprob_mean",
+            "generated_logprob_mean",
+        ]
+
+        for metric_name in training_metrics + inference_metrics:
+            if metric_name in tolerances:
+                tol = tolerances[metric_name]
+                var_name = metric_name.upper().replace('-', '_').replace(' ', '_')
+                print(
+                    f"{var_name}_RELATIVE_TOLERANCE = "
+                    f"{tol['relative_tolerance']}  # {tol['relative_tolerance']:.2%}"
+                )
+                print(f"{var_name}_ABSOLUTE_TOLERANCE = " f"{tol['absolute_tolerance']:.2e}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/functional_tests/python_test_utils/test_grpo_training_loop.py b/tests/functional_tests/python_test_utils/test_grpo_training_loop.py
index 1b6eedd4fdb..6faca9b11b3 100644
--- a/tests/functional_tests/python_test_utils/test_grpo_training_loop.py
+++ b/tests/functional_tests/python_test_utils/test_grpo_training_loop.py
@@ -2,14 +2,93 @@
 
 import json
 import logging
-import math
 from statistics import median
+from typing import Any, Dict, List, Tuple
+
+import yaml
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-
-def test_grpo_training_loop(golden_values_path: str, test_values_path: str) -> None:
+# Tolerance settings for all metrics.
+# These tolerances account for hardware variance (different GPU silicon,
+# driver versions, CUDA/cuDNN differences) while still catching real regressions.
+# Tolerances can be tuned using compute_golden_statistics.py to analyze variance
+# across multiple runs on different hardware.
+
+# LM Loss tolerances
+LM_LOSS_RELATIVE_TOLERANCE = 0.01  # 1% relative tolerance
+LM_LOSS_ABSOLUTE_TOLERANCE = 1e-6  # For values near zero
+
+# Iteration time tolerances (performance metric, higher variance expected)
+ITERATION_TIME_RELATIVE_TOLERANCE = 0.15  # 15% relative tolerance
+
+# Memory allocation tolerances
+MEM_ALLOCATED_BYTES_RELATIVE_TOLERANCE = 0.10  # 10% relative tolerance
+MEM_MAX_ALLOCATED_BYTES_RELATIVE_TOLERANCE = 0.10  # 10% relative tolerance
+
+
+def validate_with_tolerance(
+    golden_values: Dict[str, Any],
+    current_values: Dict[str, Any],
+    relative_tolerance: float,
+    absolute_tolerance: float = 1e-9,
+    metric_name: str = "metric",
+) -> Tuple[bool, List[str]]:
+    """
+    Validate that current values are within tolerance of golden values.
+
+    Args:
+        golden_values: Dict mapping step -> expected value
+        current_values: Dict mapping step -> actual value
+        relative_tolerance: Maximum allowed relative difference (e.g., 0.01 for 1%)
+        absolute_tolerance: Tolerance for values near zero
+        metric_name: Name of metric for error messages
+
+    Returns:
+        Tuple of (passed: bool, mismatches: List[str])
+    """
+    mismatches = []
+
+    for step, golden_val in golden_values.items():
+        if step not in current_values:
+            mismatches.append(f"Step {step}: missing in current run")
+            continue
+
+        current_val = current_values[step]
+
+        # Handle the case where golden value is zero or near-zero
+        if golden_val == 0 or abs(golden_val) < absolute_tolerance:
+            if abs(current_val) > absolute_tolerance:
+                mismatches.append(f"Step {step}: expected ~0, got {current_val}")
+        else:
+            # Calculate relative difference
+            rel_diff = abs(current_val - golden_val) / abs(golden_val)
+            if rel_diff > relative_tolerance:
+                mismatches.append(
+                    f"Step {step}: {current_val} differs from golden {golden_val} "
+                    f"by {rel_diff:.4%} (tolerance: {relative_tolerance:.2%})"
+                )
+
+    # Check for extra steps in current that aren't in golden
+    extra_steps = set(current_values.keys()) - set(golden_values.keys())
+    if extra_steps:
+        logger.info(f"{metric_name}: Ignoring extra steps in current run: {extra_steps}")
+
+    return len(mismatches) == 0, mismatches
+
+
+def test_grpo_training_loop(
+    golden_values_path: str, test_values_path: str, model_config_path: str
+) -> None:
+    with open(model_config_path, 'r') as f:
+        model_config = yaml.safe_load(f)
+        metrics = model_config["METRICS"]
+        if "THROUGHPUT_TEST_PARAMS" in model_config:
+            throughput_test_params = model_config["THROUGHPUT_TEST_PARAMS"]
+            start_step = throughput_test_params["--start_step"]
+        else:
+            start_step = 1
 
     with open(golden_values_path, 'r') as f1, open(test_values_path, 'r') as f2:
         golden_values_content = f1.read()
@@ -41,54 +120,99 @@ def test_grpo_training_loop(golden_values_path: str, test_values_path: str) -> N
         )
     assert len(output_groundtruth) > 0, "No test performed for output"
 
-    if "iteration-time" in output_groundtruth.keys():
+    if "iteration-time" in metrics and "iteration-time" in output_current:
 
         # First warmup iteration is excluded from iteration-time statistics.
         iteration_time_sampled = median(
-            [l for l in output_current["iteration-time"]['values'].values()][1:]
+            [l for l in output_current["iteration-time"]['values'].values()][start_step:]
         )
         iteration_time_golden = median(
-            [l for l in output_groundtruth["iteration-time"]['values'].values()][1:]
+            [l for l in output_groundtruth["iteration-time"]['values'].values()][start_step:]
         )
 
-        # 10% is empirically observed to be within hardware variance.
-        assert (
-            0.9 * iteration_time_golden <= iteration_time_sampled <= 1.2 * iteration_time_golden
-        ), (
-            f"Iteration time {iteration_time_sampled} ms not within 10% below or 20% above "
-            f"golden value ~{iteration_time_golden} ms. "
+        lower_bound = (1 - ITERATION_TIME_RELATIVE_TOLERANCE) * iteration_time_golden
+        upper_bound = (1 + ITERATION_TIME_RELATIVE_TOLERANCE) * iteration_time_golden
+        assert lower_bound <= iteration_time_sampled <= upper_bound, (
+            f"Iteration time {iteration_time_sampled} ms not within "
+            f"{ITERATION_TIME_RELATIVE_TOLERANCE:.0%} of golden value ~{iteration_time_golden} ms. "
             f"Sampled: {output_current['iteration-time']} ms. "
             f"Please update golden values in the functional tests if this is expected."
         )
 
         output_groundtruth.pop('iteration-time')
 
-    if "lm-loss" in output_groundtruth.keys():
+    if "lm-loss" in metrics and "lm-loss" in output_current:
 
-        # Require exact matching of all lm-loss values.
+        # Validate lm-loss values with tolerance to account for hardware variance.
+        # Previously required exact matching, but this caused flaky failures due to
+        # floating-point differences across different GPU hardware.
         golden_lm_loss_values = output_groundtruth["lm-loss"]['values']
         current_lm_loss_values = output_current["lm-loss"]['values']
 
-        assert golden_lm_loss_values == current_lm_loss_values, (
-            f"LM loss values do not exactly match.\n"
-            f"Golden: {golden_lm_loss_values}\n"
-            f"Current: {current_lm_loss_values}\n"
-            f"Please update golden values in the functional tests if this is expected."
+        passed, mismatches = validate_with_tolerance(
+            golden_lm_loss_values,
+            current_lm_loss_values,
+            relative_tolerance=LM_LOSS_RELATIVE_TOLERANCE,
+            absolute_tolerance=LM_LOSS_ABSOLUTE_TOLERANCE,
+            metric_name="lm-loss",
         )
 
+        if not passed:
+            error_msg = (
+                f"LM loss values outside tolerance ({LM_LOSS_RELATIVE_TOLERANCE:.1%}):\n"
+                + "\n".join(f"  - {m}" for m in mismatches)
+                + f"\n\nGolden: {golden_lm_loss_values}\n"
+                + f"Current: {current_lm_loss_values}\n"
+                + "Please update golden values in the functional tests if this is expected."
+            )
+            assert False, error_msg
+
         output_groundtruth.pop('lm-loss')
 
-    if "num-zeros" in output_groundtruth.keys():
+    if "mem-allocated-bytes" in metrics and "mem-allocated-bytes" in output_current:
 
-        # Require exact matching of all lm-loss values.
-        golden_num_zeros_values = output_groundtruth["num-zeros"]['values']
-        current_num_zeros_values = output_current["num-zeros"]['values']
+        # Use max instead of median - we care about worst-case memory usage
+        # Skip first step (warmup) which may have different memory characteristics
+        current_values = [l for l in output_current["mem-allocated-bytes"]['values'].values()][1:]
+        golden_values = [l for l in output_groundtruth["mem-allocated-bytes"]['values'].values()][
+            1:
+        ]
+
+        mem_allocated_bytes_sampled = max(current_values)
+        mem_allocated_bytes_golden = max(golden_values)
+
+        upper_bound = (1 + MEM_ALLOCATED_BYTES_RELATIVE_TOLERANCE) * mem_allocated_bytes_golden
+        assert mem_allocated_bytes_sampled <= upper_bound, (
+            f"Max mem allocated bytes {mem_allocated_bytes_sampled} bytes exceeds "
+            f"{MEM_ALLOCATED_BYTES_RELATIVE_TOLERANCE:.0%} above golden max {mem_allocated_bytes_golden} bytes. "
+            f"Upper bound: {upper_bound} bytes. "
+            f"Please update golden values in the functional tests if this is expected."
+        )
 
-        assert golden_num_zeros_values == current_num_zeros_values, (
-            f"LM loss values do not exactly match.\n"
-            f"Golden: {golden_num_zeros_values}\n"
-            f"Current: {current_num_zeros_values}\n"
+        output_groundtruth.pop('mem-allocated-bytes')
+
+    if "mem-max-allocated-bytes" in metrics and "mem-max-allocated-bytes" in output_current:
+
+        # Use max - we care that peak memory doesn't exceed the golden peak
+        # Skip first step (warmup) which may have different memory characteristics
+        current_values = [l for l in output_current["mem-max-allocated-bytes"]['values'].values()][
+            1:
+        ]
+        golden_values = [
+            l for l in output_groundtruth["mem-max-allocated-bytes"]['values'].values()
+        ][1:]
+
+        mem_max_allocated_bytes_sampled = max(current_values)
+        mem_max_allocated_bytes_golden = max(golden_values)
+
+        upper_bound = (
+            1 + MEM_MAX_ALLOCATED_BYTES_RELATIVE_TOLERANCE
+        ) * mem_max_allocated_bytes_golden
+        assert mem_max_allocated_bytes_sampled <= upper_bound, (
+            f"Max mem-max-allocated bytes {mem_max_allocated_bytes_sampled} bytes exceeds "
+            f"{MEM_MAX_ALLOCATED_BYTES_RELATIVE_TOLERANCE:.0%} above golden max {mem_max_allocated_bytes_golden} bytes. "
+            f"Upper bound: {upper_bound} bytes. "
             f"Please update golden values in the functional tests if this is expected."
         )
 
-        output_groundtruth.pop('num-zeros')
+        output_groundtruth.pop('mem-max-allocated-bytes')
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json
deleted file mode 100644
index a19d42718aa..00000000000
--- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json
+++ /dev/null
@@ -1,173 +0,0 @@
-{
-    "mem-allocated-bytes": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 55289954304.0,
-            "2": 55292747776.0,
-            "3": 55292731392.0,
-            "4": 55292891136.0,
-            "5": 55292878848.0,
-            "6": 55292878848.0,
-            "7": 55292878848.0,
-            "8": 55292788736.0,
-            "9": 55292788736.0,
-            "10": 55292788736.0,
-            "11": 55292792832.0,
-            "12": 55292792832.0,
-            "13": 55292792832.0,
-            "14": 55292792832.0,
-            "15": 55292792832.0,
-            "16": 55292796928.0,
-            "17": 55292796928.0,
-            "18": 55292801024.0,
-            "19": 55292805120.0,
-            "20": 55292801024.0,
-            "21": 55292801024.0,
-            "22": 55292796928.0,
-            "23": 55292801024.0,
-            "24": 55292796928.0,
-            "25": 55292801024.0,
-            "26": 55292796928.0,
-            "27": 55292796928.0,
-            "28": 55292801024.0,
-            "29": 55292801024.0,
-            "30": 55292805120.0,
-            "31": 55292805120.0,
-            "32": 55292805120.0,
-            "33": 55292805120.0,
-            "34": 55292805120.0,
-            "35": 55292805120.0,
-            "36": 55292805120.0,
-            "37": 55292801024.0,
-            "38": 55292801024.0,
-            "39": 55292801024.0,
-            "40": 55292805120.0,
-            "41": 55292805120.0,
-            "42": 55292805120.0,
-            "43": 55292801024.0,
-            "44": 55292796928.0,
-            "45": 55292801024.0,
-            "46": 55292801024.0,
-            "47": 55292801024.0,
-            "48": 55292801024.0,
-            "49": 55292805120.0,
-            "50": 55292805120.0
-        }
-    },
-    "mem-max-allocated-bytes": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 55289958400.0,
-            "2": 57103880192.0,
-            "3": 57104392192.0,
-            "4": 57104416768.0,
-            "5": 57104416768.0,
-            "6": 57104416768.0,
-            "7": 57104416768.0,
-            "8": 57104416768.0,
-            "9": 57104416768.0,
-            "10": 57104416768.0,
-            "11": 57104416768.0,
-            "12": 57104416768.0,
-            "13": 57104416768.0,
-            "14": 57104416768.0,
-            "15": 57104416768.0,
-            "16": 57104416768.0,
-            "17": 57104416768.0,
-            "18": 57104416768.0,
-            "19": 57104416768.0,
-            "20": 57104416768.0,
-            "21": 57104416768.0,
-            "22": 57104416768.0,
-            "23": 57104416768.0,
-            "24": 57104416768.0,
-            "25": 57104416768.0,
-            "26": 57104416768.0,
-            "27": 57104416768.0,
-            "28": 57104416768.0,
-            "29": 57104416768.0,
-            "30": 57104416768.0,
-            "31": 57104416768.0,
-            "32": 57104416768.0,
-            "33": 57104416768.0,
-            "34": 57104416768.0,
-            "35": 57104416768.0,
-            "36": 57104416768.0,
-            "37": 57104416768.0,
-            "38": 57104416768.0,
-            "39": 57104416768.0,
-            "40": 57104416768.0,
-            "41": 57104416768.0,
-            "42": 57104416768.0,
-            "43": 57104416768.0,
-            "44": 57104416768.0,
-            "45": 57104416768.0,
-            "46": 57104416768.0,
-            "47": 57104416768.0,
-            "48": 57104416768.0,
-            "49": 57104416768.0,
-            "50": 57104416768.0
-        }
-    },
-    "iteration-time": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 38.24908,
-            "2": 4.52458,
-            "3": 3.69393,
-            "4": 3.38577,
-            "5": 3.41862,
-            "6": 3.27421,
-            "7": 3.32023,
-            "8": 3.83723,
-            "9": 4.07373,
-            "10": 3.47799,
-            "11": 3.27499,
-            "12": 3.37017,
-            "13": 3.3918,
-            "14": 3.25114,
-            "15": 3.29905,
-            "16": 3.29943,
-            "17": 3.50383,
-            "18": 3.56844,
-            "19": 3.30276,
-            "20": 3.34553,
-            "21": 3.29165,
-            "22": 3.30348,
-            "23": 3.33814,
-            "24": 3.31525,
-            "25": 3.29337,
-            "26": 3.26119,
-            "27": 3.5167,
-            "28": 3.2312,
-            "29": 3.45063,
-            "30": 3.3088,
-            "31": 3.32522,
-            "32": 3.28154,
-            "33": 3.23551,
-            "34": 3.20003,
-            "35": 3.25844,
-            "36": 3.67071,
-            "37": 3.1881,
-            "38": 3.30757,
-            "39": 3.32895,
-            "40": 3.29602,
-            "41": 3.25522,
-            "42": 3.28932,
-            "43": 3.32204,
-            "44": 3.26419,
-            "45": 3.75371,
-            "46": 3.23126,
-            "47": 3.25929,
-            "48": 3.19512,
-            "49": 3.32815,
-            "50": 3.25617
-        }
-    }
-}
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json
deleted file mode 100644
index 4db934b1330..00000000000
--- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json
+++ /dev/null
@@ -1,287 +0,0 @@
-{
-    "lm loss": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 0.04567,
-            "2": 0.0,
-            "3": 0.0,
-            "4": 0.0,
-            "5": 0.0,
-            "6": 0.0,
-            "7": 0.0,
-            "8": 0.04622,
-            "9": 0.0,
-            "10": 0.0,
-            "11": 0.0,
-            "12": 0.0,
-            "13": 0.0,
-            "14": 0.0,
-            "15": 0.0,
-            "16": 0.0,
-            "17": 0.0,
-            "18": 0.0,
-            "19": 0.0,
-            "20": 0.0,
-            "21": 0.0,
-            "22": 0.0,
-            "23": 0.0,
-            "24": 0.0,
-            "25": 0.03308,
-            "26": 0.0,
-            "27": 0.0,
-            "28": 0.09392,
-            "29": 0.0,
-            "30": 0.0,
-            "31": 0.0,
-            "32": 0.0,
-            "33": 0.0,
-            "34": 0.03909,
-            "35": 0.0,
-            "36": 0.0,
-            "37": 0.0,
-            "38": 0.0,
-            "39": 0.04574,
-            "40": 0.0,
-            "41": 0.0,
-            "42": 0.0,
-            "43": 0.0,
-            "44": 0.0,
-            "45": 0.0,
-            "46": 0.0,
-            "47": 0.0,
-            "48": 0.0,
-            "49": 0.0,
-            "50": 0.0
-        }
-    },
-    "num-zeros": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 43.0,
-            "2": 583687296.0,
-            "3": 583687296.0,
-            "4": 583687296.0,
-            "5": 583687296.0,
-            "6": 583687296.0,
-            "7": 583687296.0,
-            "8": 42.0,
-            "9": 583687296.0,
-            "10": 583687296.0,
-            "11": 583687296.0,
-            "12": 583687296.0,
-            "13": 583687296.0,
-            "14": 583687296.0,
-            "15": 583687296.0,
-            "16": 583687296.0,
-            "17": 583687296.0,
-            "18": 583687296.0,
-            "19": 583687296.0,
-            "20": 583687296.0,
-            "21": 583687296.0,
-            "22": 583687296.0,
-            "23": 583687296.0,
-            "24": 583687296.0,
-            "25": 56.0,
-            "26": 583687296.0,
-            "27": 583687296.0,
-            "28": 18.0,
-            "29": 583687296.0,
-            "30": 583687296.0,
-            "31": 583687296.0,
-            "32": 583687296.0,
-            "33": 583687296.0,
-            "34": 32.0,
-            "35": 583687296.0,
-            "36": 583687296.0,
-            "37": 583687296.0,
-            "38": 583687296.0,
-            "39": 27.0,
-            "40": 583687296.0,
-            "41": 583687296.0,
-            "42": 583687296.0,
-            "43": 583687296.0,
-            "44": 583687296.0,
-            "45": 583687296.0,
-            "46": 583687296.0,
-            "47": 583687296.0,
-            "48": 583687296.0,
-            "49": 583687296.0,
-            "50": 583687296.0
-        }
-    },
-    "mem-allocated-bytes": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 56705486848.0,
-            "2": 56707366912.0,
-            "3": 56707289088.0,
-            "4": 56707284992.0,
-            "5": 56707284992.0,
-            "6": 56707293184.0,
-            "7": 56707297280.0,
-            "8": 56707293184.0,
-            "9": 56707293184.0,
-            "10": 56707297280.0,
-            "11": 56707289088.0,
-            "12": 56707293184.0,
-            "13": 56707301376.0,
-            "14": 56707305472.0,
-            "15": 56707313664.0,
-            "16": 56707317760.0,
-            "17": 56707325952.0,
-            "18": 56707330048.0,
-            "19": 56707338240.0,
-            "20": 56707342336.0,
-            "21": 56707350528.0,
-            "22": 56707354624.0,
-            "23": 56707358720.0,
-            "24": 56707317760.0,
-            "25": 56707317760.0,
-            "26": 56707309568.0,
-            "27": 56707309568.0,
-            "28": 56707305472.0,
-            "29": 56707309568.0,
-            "30": 56707309568.0,
-            "31": 56707309568.0,
-            "32": 56707305472.0,
-            "33": 56707305472.0,
-            "34": 56707276800.0,
-            "35": 56707284992.0,
-            "36": 56707293184.0,
-            "37": 56707293184.0,
-            "38": 56707276800.0,
-            "39": 56707284992.0,
-            "40": 56707284992.0,
-            "41": 56707252224.0,
-            "42": 56707256320.0,
-            "43": 56707260416.0,
-            "44": 56707252224.0,
-            "45": 56707235840.0,
-            "46": 56707244032.0,
-            "47": 56707244032.0,
-            "48": 56707239936.0,
-            "49": 56707235840.0,
-            "50": 56707227648.0
-        }
-    },
-    "mem-max-allocated-bytes": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": 56705486848.0,
-            "2": 58520117248.0,
-            "3": 58520694784.0,
-            "4": 58520694784.0,
-            "5": 58520694784.0,
-            "6": 58520698880.0,
-            "7": 58520707072.0,
-            "8": 58520707072.0,
-            "9": 58520707072.0,
-            "10": 58520707072.0,
-            "11": 58520707072.0,
-            "12": 58520707072.0,
-            "13": 58520707072.0,
-            "14": 58520711168.0,
-            "15": 58520719360.0,
-            "16": 58520723456.0,
-            "17": 58520731648.0,
-            "18": 58520735744.0,
-            "19": 58520743936.0,
-            "20": 58520748032.0,
-            "21": 58520756224.0,
-            "22": 58520764416.0,
-            "23": 58520764416.0,
-            "24": 58520764416.0,
-            "25": 58520764416.0,
-            "26": 58520764416.0,
-            "27": 58520764416.0,
-            "28": 58520764416.0,
-            "29": 58520764416.0,
-            "30": 58520764416.0,
-            "31": 58520764416.0,
-            "32": 58520764416.0,
-            "33": 58520764416.0,
-            "34": 58520764416.0,
-            "35": 58520764416.0,
-            "36": 58520764416.0,
-            "37": 58520764416.0,
-            "38": 58520764416.0,
-            "39": 58520764416.0,
-            "40": 58520764416.0,
-            "41": 58520764416.0,
-            "42": 58520764416.0,
-            "43": 58520764416.0,
-            "44": 58520764416.0,
-            "45": 58520764416.0,
-            "46": 58520764416.0,
-            "47": 58520764416.0,
-            "48": 58520764416.0,
-            "49": 58520764416.0,
-            "50": 58520764416.0
-        }
-    },
-    "iteration-time": {
-        "start_step": 1,
-        "end_step": 50,
-        "step_interval": 1,
-        "values": {
-            "1": "nan",
-            "2": 64.88323,
-            "3": 9.98948,
-            "4": 10.5653,
-            "5": 9.49213,
-            "6": 9.7058,
-            "7": 10.3713,
-            "8": 9.69584,
-            "9": 10.08558,
-            "10": 9.64307,
-            "11": 9.39285,
-            "12": 9.22534,
-            "13": 9.45398,
-            "14": 9.3236,
-            "15": 9.30815,
-            "16": 9.42684,
-            "17": 9.27604,
-            "18": 9.46377,
-            "19": 9.24656,
-            "20": 9.22709,
-            "21": 9.15955,
-            "22": 9.39831,
-            "23": 9.1461,
-            "24": 9.14062,
-            "25": 9.43925,
-            "26": 9.27344,
-            "27": 9.13835,
-            "28": 9.11182,
-            "29": 9.28006,
-            "30": 9.29592,
-            "31": 9.99338,
-            "32": 10.28927,
-            "33": 9.71657,
-            "34": 10.01927,
-            "35": 9.49163,
-            "36": 9.72794,
-            "37": 9.31159,
-            "38": 9.29786,
-            "39": 9.318,
-            "40": 9.48741,
-            "41": 9.59212,
-            "42": 9.29507,
-            "43": 9.30203,
-            "44": 9.37176,
-            "45": 9.23509,
-            "46": 9.32089,
-            "47": 9.36602,
-            "48": 9.43024,
-            "49": 9.19031,
-            "50": 9.19624
-        }
-    }
-}
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/model_config.yaml
deleted file mode 100644
index 48732597428..00000000000
--- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/model_config.yaml
+++ /dev/null
@@ -1,78 +0,0 @@
-ENV_VARS:
-  CUDA_DEVICE_MAX_CONNECTIONS: 1
-  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
-  NCCL_ALGO: Ring
-  CUBLAS_WORKSPACE_CONFIG: :4096:8
-TEST_TYPE: frozen-start
-MODE: rl
-MODEL_ARGS:
-  --tiktoken-pattern: v2
-  --use-mcore-models: true
-  --tokenizer-type: TikTokenizer
-  --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/mcore_mistral/nemo_minitron-0.5b/v1/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json
-  --load: ${CHECKPOINT_LOAD_PATH}/model/mcore_mistral/nemo_minitron-0.5b/v1/
-  --auto-detect-ckpt-format: true
-  --max-tokens-to-oom: 3600000
-  --inference-max-seq-length: 1024
-  --attention-backend: flash
-  --mock-data: true
-  --micro-batch-size: 1
-  --no-load-optim: true
-  --no-use-tokenizer-model-from-checkpoint-args: true
-  --timing-log-level: 0
-  --distributed-backend: nccl
-  --log-interval: 1
-  --log-progress: true
-  --transformer-impl: transformer_engine
-  --tensor-model-parallel-size: 1
-  --pipeline-model-parallel-size: 1
-  --ckpt-format: torch_dist
-  --bf16: true
-  --log-memory-to-tensorboard: true
-  --log-num-zeros-in-grad: true
-  --log-validation-ppl-to-tensorboard: true
-  --log-timers-to-tensorboard: true
-  --num-layers: 24
-  --hidden-size: 1152
-  --num-attention-heads: 16
-  --max-position-embeddings: 1024
-  --seq-length: 1024
-  --timing-log-option: minmax
-  --log-throughput: true
-  --no-create-attention-mask-in-dataloader: true
-  --straggler-minmax-count: 16
-  --tensorboard-log-interval: 1
-  --empty-unused-memory-level: 2
-  --langrl-inference-server-type: inplace_megatron
-  --seed: 42
-  --calculate-per-token-loss: true
-  --rl-use-sequence-packing: true
-  --rl-sequence-packing-algo: fifo
-  --rl-offload-optimizer-during-inference: true
-  --timing-log-level: 1
-  --log-timers-to-tensorboard: true
-  --cuda-graph-impl: local
-  --micro-batch-size: 1
-  --global-batch-size: 16
-  --grpo-group-size: 2
-  --grpo-prompts-per-step: 8
-  --grpo-iterations: 1
-  --grpo-clamp-eps-lower: 0.2
-  --grpo-clamp-eps-upper: 0.2
-  --grpo-kl-beta: 0.0
-  --grpo-entropy-term-weight: 0.0
-  --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/env_config.yaml
-  --rl-partial-rollouts: true
-  --lr: 0.000001
-  --lr-warmup-samples: 0
-  --clip-grad: 1.0
-  --use-checkpoint-args: true
-  --dist-ckpt-strictness: log_unexpected
-  --perform-rl-step: true
-  --train-samples: 48828125
-  --exit-interval: 50
-  --tensorboard-dir: ${TENSORBOARD_PATH}
-  --save-interval: 1000000
-  --eval-interval: 1000000
-  --finetune: true
-  --inference-logging-step-interval: 1
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/env_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/env_config.yaml
similarity index 100%
rename from tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/env_config.yaml
rename to tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/env_config.yaml
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/golden_values_dev_dgx_h100.json
new file mode 100644
index 00000000000..4a8586c8e8a
--- /dev/null
+++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/golden_values_dev_dgx_h100.json
@@ -0,0 +1,83 @@
+{
+    "mem-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 20,
+        "step_interval": 1,
+        "values": {
+            "1": 60978671616.0,
+            "2": 60979740672.0,
+            "3": 60979740672.0,
+            "4": 60979740672.0,
+            "5": 60979998720.0,
+            "6": 60979732480.0,
+            "7": 60979994624.0,
+            "8": 60979728384.0,
+            "9": 60979986432.0,
+            "10": 60980248576.0,
+            "11": 60979982336.0,
+            "12": 60979720192.0,
+            "13": 60979716096.0,
+            "14": 60979716096.0,
+            "15": 60979716096.0,
+            "16": 60979716096.0,
+            "17": 60979712000.0,
+            "18": 60979970048.0,
+            "19": 60979716096.0,
+            "20": 60979716096.0
+        }
+    },
+    "mem-max-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 20,
+        "step_interval": 1,
+        "values": {
+            "1": 60978675712.0,
+            "2": 64214241280.0,
+            "3": 64214241280.0,
+            "4": 64214241280.0,
+            "5": 64214241280.0,
+            "6": 64214241280.0,
+            "7": 64214241280.0,
+            "8": 64214241280.0,
+            "9": 64214241280.0,
+            "10": 64214241280.0,
+            "11": 64214241280.0,
+            "12": 64214241280.0,
+            "13": 64214241280.0,
+            "14": 64214241280.0,
+            "15": 64214241280.0,
+            "16": 64214241280.0,
+            "17": 64214241280.0,
+            "18": 64214241280.0,
+            "19": 64214241280.0,
+            "20": 64214241280.0
+        }
+    },
+    "iteration-time": {
+        "start_step": 1,
+        "end_step": 20,
+        "step_interval": 1,
+        "values": {
+            "1": "nan",
+            "2": 37.77975,
+            "3": 15.85042,
+            "4": 14.84801,
+            "5": 14.16031,
+            "6": 14.7285,
+            "7": 14.32408,
+            "8": 14.76569,
+            "9": 13.73696,
+            "10": 14.6546,
+            "11": 14.12618,
+            "12": 14.29456,
+            "13": 14.27773,
+            "14": 14.10944,
+            "15": 13.7968,
+            "16": 13.90572,
+            "17": 13.58351,
+            "18": 14.3947,
+            "19": 13.78201,
+            "20": 13.44734
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml
similarity index 50%
rename from tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/model_config.yaml
rename to tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml
index 8c78989cef7..7f6fe4756e3 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml
@@ -3,74 +3,93 @@ ENV_VARS:
   NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
   NCCL_ALGO: Ring
   CUBLAS_WORKSPACE_CONFIG: :4096:8
+  N_REPEAT: 1
 TEST_TYPE: frozen-start
 MODE: rl
 MODEL_ARGS:
-  --tiktoken-pattern: v2
-  --use-mcore-models: true
-  --tokenizer-type: TikTokenizer
-  --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/mcore_mistral/nemo_minitron-0.5b/v1/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json
-  --load: ${CHECKPOINT_LOAD_PATH}/model/mcore_mistral/nemo_minitron-0.5b/v1/
-  --auto-detect-ckpt-format: true
-  --max-tokens-to-oom: 3600000
-  --inference-max-seq-length: 1024
-  --attention-backend: flash
-  --mock-data: true
-  --micro-batch-size: 1
-  --no-load-optim: true
-  --no-use-tokenizer-model-from-checkpoint-args: true
-  --timing-log-level: 0
-  --distributed-backend: nccl
-  --log-interval: 1
-  --log-progress: true
-  --transformer-impl: transformer_engine
-  --tensor-model-parallel-size: 1
-  --pipeline-model-parallel-size: 1
+  --tensor-model-parallel-size: 4
+  --inference-dynamic-batching-num-cuda-graphs: 1
+  --inference-dynamic-batching-unified-memory-level: 1
+  --inference-dynamic-batching-buffer-size-gb: 20
   --ckpt-format: torch_dist
+  --seq-length: 1024
+  --inference-max-seq-length: 1024
+  --load: ${CHECKPOINT_LOAD_PATH}/model/qwen3-8b-dist
+  --untie-embeddings-and-output-weights: true
+  --num-layers: 36
+  --hidden-size: 4096
+  --ffn-hidden-size: 12288
+  --num-attention-heads: 32
+  --kv-channels: 128
+  --max-position-embeddings: 1024
+  --group-query-attention: true
+  --num-query-groups: 8
+  --normalization: RMSNorm
+  --norm-epsilon: 0.000001
+  --qk-layernorm: true
+  --position-embedding-type: rope
+  --rotary-percent: 1.0
+  --rotary-base: 1000000
+  --use-rotary-position-embeddings: true
+  --swiglu: true
+  --disable-bias-linear: true
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --no-masked-softmax-fusion: true
+  --attention-softmax-in-fp32: true
+  --tokenizer-type: HuggingFaceTokenizer
+  --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/qwen3-8b-dist/tokenizer
+  --langrl-inference-server-type: inplace_megatron_chat
+  --langrl-inference-server-conversation-template: ${CHECKPOINT_LOAD_PATH}/model/qwen3-8b-dist/tokenizer
+  --vocab-size: 151936
+  --make-vocab-size-divisible-by: 128
+  --optimizer: adam
+  --adam-beta1: 0.9
+  --adam-beta2: 0.999
+  --adam-eps: 0.00000001
+  --lr: 0.000001  
+  --min-lr: 0.0000001
+  --lr-warmup-samples: 0
+  --clip-grad: 1.0
+  --weight-decay: 0.01
+  --deterministic-mode: true
+  --use-mcore-models: true
   --bf16: true
   --log-memory-to-tensorboard: true
   --log-num-zeros-in-grad: true
   --log-validation-ppl-to-tensorboard: true
   --log-timers-to-tensorboard: true
-  --num-layers: 24
-  --hidden-size: 1152
-  --num-attention-heads: 16
-  --max-position-embeddings: 1024
-  --seq-length: 1024
   --timing-log-option: minmax
   --log-throughput: true
   --no-create-attention-mask-in-dataloader: true
   --straggler-minmax-count: 16
   --tensorboard-log-interval: 1
+  --log-interval: 1
+  --log-progress: true
   --empty-unused-memory-level: 2
-  --langrl-inference-server-type: inplace_megatron
   --seed: 42
   --calculate-per-token-loss: true
   --rl-use-sequence-packing: true
   --rl-sequence-packing-algo: fifo
   --rl-offload-optimizer-during-inference: true
   --timing-log-level: 1
-  --log-timers-to-tensorboard: true
   --cuda-graph-impl: local
   --micro-batch-size: 1
-  --global-batch-size: 16
+  --global-batch-size: 2
   --grpo-group-size: 2
-  --grpo-prompts-per-step: 8
+  --grpo-prompts-per-step: 2
   --grpo-iterations: 1
   --grpo-clamp-eps-lower: 0.2
   --grpo-clamp-eps-upper: 0.2
   --grpo-kl-beta: 0.0
   --grpo-entropy-term-weight: 0.0
-  --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/env_config.yaml
+  --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/env_config.yaml
   --rl-partial-rollouts: true
-  --lr: 0.000001
-  --lr-warmup-samples: 0
-  --clip-grad: 1.0
   --use-checkpoint-args: true
   --dist-ckpt-strictness: log_unexpected
   --perform-rl-step: true
   --train-samples: 48828125
-  --exit-interval: 50
+  --exit-interval: 20
   --tensorboard-dir: ${TENSORBOARD_PATH}
   --save-interval: 1000000
   --eval-interval: 1000000
@@ -78,4 +97,7 @@ MODEL_ARGS:
   --inference-logging-step-interval: 1
 METRICS:
   - "mem-allocated-bytes"
-  - "mem-max-allocated-bytes"
\ No newline at end of file
+  - "mem-max-allocated-bytes"
+  - "iteration-time"
+THROUGHPUT_TEST_PARAMS:
+  --start_step: 10
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/env_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/env_config.yaml
similarity index 100%
rename from tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/env_config.yaml
rename to tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/env_config.yaml
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/golden_values_dev_dgx_h100.json
new file mode 100644
index 00000000000..e0bcb14f29b
--- /dev/null
+++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/golden_values_dev_dgx_h100.json
@@ -0,0 +1,83 @@
+{
+    "mem-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 20,
+        "step_interval": 1,
+        "values": {
+            "1": 60922077184.0,
+            "2": 60922073088.0,
+            "3": 60922068992.0,
+            "4": 60922073088.0,
+            "5": 60922331136.0,
+            "6": 60922073088.0,
+            "7": 60922335232.0,
+            "8": 60922068992.0,
+            "9": 60922073088.0,
+            "10": 60922073088.0,
+            "11": 60922077184.0,
+            "12": 60922093568.0,
+            "13": 60922351616.0,
+            "14": 60922085376.0,
+            "15": 60922085376.0,
+            "16": 60922085376.0,
+            "17": 60922089472.0,
+            "18": 60922085376.0,
+            "19": 60922085376.0,
+            "20": 60922089472.0
+        }
+    },
+    "mem-max-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 20,
+        "step_interval": 1,
+        "values": {
+            "1": 60922081280.0,
+            "2": 64156041216.0,
+            "3": 64156041216.0,
+            "4": 64156041216.0,
+            "5": 64156041216.0,
+            "6": 64156041216.0,
+            "7": 64156041216.0,
+            "8": 64156041216.0,
+            "9": 64156041216.0,
+            "10": 64156041216.0,
+            "11": 64156045312.0,
+            "12": 64156061696.0,
+            "13": 64156061696.0,
+            "14": 64156061696.0,
+            "15": 64156061696.0,
+            "16": 64156061696.0,
+            "17": 64156061696.0,
+            "18": 64156061696.0,
+            "19": 64156061696.0,
+            "20": 64156061696.0
+        }
+    },
+    "iteration-time": {
+        "start_step": 1,
+        "end_step": 20,
+        "step_interval": 1,
+        "values": {
+            "1": "nan",
+            "2": 110.96005,
+            "3": 49.76537,
+            "4": 46.36583,
+            "5": 46.63055,
+            "6": 50.62966,
+            "7": 46.52987,
+            "8": 44.32473,
+            "9": 46.39563,
+            "10": 44.4327,
+            "11": 43.93223,
+            "12": 46.84642,
+            "13": 43.45953,
+            "14": 42.21466,
+            "15": 42.70466,
+            "16": 42.45673,
+            "17": 43.68298,
+            "18": 41.36069,
+            "19": 42.64788,
+            "20": 45.08387
+        }
+    }
+}
diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml
new file mode 100644
index 00000000000..456280fdb04
--- /dev/null
+++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml
@@ -0,0 +1,103 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+  N_REPEAT: 1
+TEST_TYPE: frozen-start
+MODE: rl
+MODEL_ARGS:
+  --tensor-model-parallel-size: 4
+  --inference-dynamic-batching-num-cuda-graphs: 1
+  --inference-dynamic-batching-unified-memory-level: 1
+  --inference-dynamic-batching-buffer-size-gb: 20
+  --ckpt-format: torch_dist
+  --seq-length: 1024
+  --inference-max-seq-length: 1024
+  --load: ${CHECKPOINT_LOAD_PATH}/model/qwen3-8b-dist
+  --untie-embeddings-and-output-weights: true
+  --num-layers: 36
+  --hidden-size: 4096
+  --ffn-hidden-size: 12288
+  --num-attention-heads: 32
+  --kv-channels: 128
+  --max-position-embeddings: 1024
+  --group-query-attention: true
+  --num-query-groups: 8
+  --normalization: RMSNorm
+  --norm-epsilon: 0.000001
+  --qk-layernorm: true
+  --position-embedding-type: rope
+  --rotary-percent: 1.0
+  --rotary-base: 1000000
+  --use-rotary-position-embeddings: true
+  --swiglu: true
+  --disable-bias-linear: true
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --no-masked-softmax-fusion: true
+  --attention-softmax-in-fp32: true
+  --tokenizer-type: HuggingFaceTokenizer
+  --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/qwen3-8b-dist/tokenizer
+  --langrl-inference-server-type: inplace_megatron_chat
+  --langrl-inference-server-conversation-template: ${CHECKPOINT_LOAD_PATH}/model/qwen3-8b-dist/tokenizer
+  --vocab-size: 151936
+  --make-vocab-size-divisible-by: 128
+  --optimizer: adam
+  --adam-beta1: 0.9
+  --adam-beta2: 0.999
+  --adam-eps: 0.00000001
+  --lr: 0.000001  
+  --min-lr: 0.0000001
+  --lr-warmup-samples: 0
+  --clip-grad: 1.0
+  --weight-decay: 0.01
+  --deterministic-mode: true
+  --use-mcore-models: true
+  --bf16: true
+  --log-memory-to-tensorboard: true
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --timing-log-option: minmax
+  --log-throughput: true
+  --no-create-attention-mask-in-dataloader: true
+  --straggler-minmax-count: 16
+  --tensorboard-log-interval: 1
+  --empty-unused-memory-level: 2
+  --seed: 42
+  --calculate-per-token-loss: true
+  --rl-use-sequence-packing: true
+  --rl-sequence-packing-algo: fifo
+  --rl-offload-optimizer-during-inference: true
+  --timing-log-level: 1
+  --log-interval: 1
+  --log-progress: true
+  --cuda-graph-impl: local
+  --micro-batch-size: 1
+  --global-batch-size: 4
+  --grpo-group-size: 2
+  --grpo-prompts-per-step: 2
+  --grpo-iterations: 1
+  --grpo-clamp-eps-lower: 0.2
+  --grpo-clamp-eps-upper: 0.2
+  --grpo-kl-beta: 0.0
+  --grpo-entropy-term-weight: 0.0
+  --langrl-env-config: tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/env_config.yaml
+  --rl-partial-rollouts: true
+  --use-checkpoint-args: true
+  --dist-ckpt-strictness: log_unexpected
+  --perform-rl-step: true
+  --train-samples: 48828125
+  --exit-interval: 20
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --save-interval: 1000000
+  --eval-interval: 1000000
+  --finetune: true
+  --inference-logging-step-interval: 1
+METRICS:
+  - "mem-allocated-bytes"
+  - "mem-max-allocated-bytes"
+  - "iteration-time"
+THROUGHPUT_TEST_PARAMS:
+  --start_step: 10
diff --git a/tests/test_utils/recipes/gpt-grpo.yaml b/tests/test_utils/recipes/gpt-grpo.yaml
index 11e8eadea9b..e707c1c2431 100644
--- a/tests/test_utils/recipes/gpt-grpo.yaml
+++ b/tests/test_utils/recipes/gpt-grpo.yaml
@@ -45,7 +45,7 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
@@ -54,15 +54,15 @@ spec:
     bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh ${{ARGUMENTS[@]}}
 
 products:
-  # - test_case: [gpt_grpo_tp1_pp1_dp8_583m_throughputtest] # Offline until golden values are properly written to disk
-  #   products:
-  #     - environment: [dev]
-  #       scope: [mr]
-  #       platforms: [dgx_h100]
-  - test_case: [gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github]
+  - test_case: [gpt_grpo_tp4_pp1_dp2_8b_throughput]
     products:
       - environment: [dev]
-        scope: [mr-github-broken]
+        scope: [mr]
+        platforms: [dgx_h100]
+  - test_case: [gpt_grpo_tp4_pp1_dp2_8b_throughput_github]
+    products:
+      - environment: [dev]
+        scope: [mr-github]
         platforms: [dgx_h100]
   - test_case: [gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest]
     products:

From 4015ff17658cc7ac622cc0f5a3be254dc384f26f Mon Sep 17 00:00:00 2001
From: helen ngo <helenn@nvidia.com>
Date: Tue, 27 Jan 2026 11:17:57 -0500
Subject: [PATCH 34/79] Inference functional tests: Write outputs to
 INFERENCE_OUTPUT_PATH instead of TENSORBOARD_PATH (#3061)

---
 tests/functional_tests/shell_test_utils/run_ci_test.sh         | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../model_config.yaml                                          | 2 +-
 .../recipes/gpt-dynamic-inference-with-coordinator.yaml        | 3 ++-
 tests/test_utils/recipes/gpt-dynamic-inference.yaml            | 3 ++-
 tests/test_utils/recipes/gpt-static-inference.yaml             | 3 ++-
 tests/test_utils/recipes/mamba-dynamic-inference.yaml          | 3 ++-
 tests/test_utils/recipes/mamba-static-inference.yaml           | 3 ++-
 .../recipes/moe-dynamic-inference-with-coordinator.yaml        | 3 ++-
 tests/test_utils/recipes/moe-dynamic-inference.yaml            | 3 ++-
 tests/test_utils/recipes/moe-grpo.yaml                         | 2 +-
 tests/test_utils/recipes/moe-static-inference.yaml             | 3 ++-
 34 files changed, 42 insertions(+), 34 deletions(-)

diff --git a/tests/functional_tests/shell_test_utils/run_ci_test.sh b/tests/functional_tests/shell_test_utils/run_ci_test.sh
index 7b5d58bd1a6..3f1500b502c 100644
--- a/tests/functional_tests/shell_test_utils/run_ci_test.sh
+++ b/tests/functional_tests/shell_test_utils/run_ci_test.sh
@@ -327,7 +327,7 @@ for i in $(seq 1 $N_REPEAT); do
             if [[ "$TEST_TYPE" == "frozen-start" ]]; then
                 uv run --no-sync pytest -s -o log_cli=true --log-cli-level=info $ROOT_DIR/tests/functional_tests/python_test_utils/test_inference_regular_pipeline.py \
                     --golden-values-path $GOLDEN_VALUES_PATH \
-                    --test-values-path $TENSORBOARD_PATH \
+                    --test-values-path $INFERENCE_OUTPUT_PATH \
                     --model-config-path ${TRAINING_PARAMS_PATH} \
                     $ALLOW_NONDETERMINISTIC_ALGO_ARG
             fi
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/model_config.yaml
index abe9cac678f..743c4f50da3 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_fp8_logitsmatch/model_config.yaml
@@ -49,7 +49,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/model_config.yaml
index 9b3f434b26b..b5dc7cd5bd2 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_cuda_graphs_logitsmatch_decode_graphs_only/model_config.yaml
@@ -50,7 +50,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
index 34636bf08e5..aae99fd1c4c 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
@@ -46,7 +46,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --output-every-n-results: 32
   --prompt-file: ${DATA_PATH}/text/sharegpt-vicuna/filtered/processed.jsonl
   --prompt-file-num-truncate: 128 # originally 1024
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_logitsmatch_zmq/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_logitsmatch_zmq/model_config.yaml
index fc3dc87240c..d84dd24487f 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_logitsmatch_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_logitsmatch_zmq/model_config.yaml
@@ -45,7 +45,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-step: 32
   --use-flashinfer-fused-rope: true
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq/model_config.yaml
index 0ef4f95aaaf..aa4fde5e512 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_dp8_583m_throughputtest_zmq/model_config.yaml
@@ -48,7 +48,7 @@ MODEL_ARGS:
   --disable-chunked-prefill: true
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --output-every-n-results: 32
   --prompt-file: ${DATA_PATH}/text/sharegpt-vicuna/filtered/processed.jsonl
   --prompt-file-num-truncate: 1024
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp8_dp1_583m_logitsmatch_zmq/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp8_dp1_583m_logitsmatch_zmq/model_config.yaml
index a0e83307504..bd34c11fc24 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp8_dp1_583m_logitsmatch_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp8_dp1_583m_logitsmatch_zmq/model_config.yaml
@@ -48,7 +48,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-step: 32
   --use-flashinfer-fused-rope: true
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq/model_config.yaml
index 17cef6d8e16..13d409c5968 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq/model_config.yaml
@@ -48,7 +48,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-step: 32
   --use-flashinfer-fused-rope: true
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_583m_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_583m_logitsmatch/model_config.yaml
index c1cef970264..4458edf5772 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_583m_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_583m_logitsmatch/model_config.yaml
@@ -44,7 +44,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 10
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-step: 32
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_dp1_583m_logitsmatch_zmq/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_dp1_583m_logitsmatch_zmq/model_config.yaml
index 920da1d1682..8d5779a5099 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_dp1_583m_logitsmatch_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp8_pp1_dp1_583m_logitsmatch_zmq/model_config.yaml
@@ -46,7 +46,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --dist-ckpt-strictness: log_unexpected
   --inference-ckpt-non-strict: true # To handle the extra_state errors
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-step: 32
   --use-flashinfer-fused-rope: true
diff --git a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch/model_config.yaml
index efe4f7424f9..6d63b0e4228 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch/model_config.yaml
@@ -75,7 +75,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 80
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompt-file: ./tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch/test_prompts.jsonl
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-logging-step-interval: 1
diff --git a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_cudagraphs/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_cudagraphs/model_config.yaml
index 352b1426554..8f54fff0a2f 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_cudagraphs/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_cudagraphs/model_config.yaml
@@ -47,7 +47,7 @@ MODEL_ARGS:
   --inference-rng-tracker: true
   --inference-max-requests: 1
   --dist-ckpt-strictness: log_unexpected
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-logging-step-interval: 1
diff --git a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_fp8_cudagraphs/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_fp8_cudagraphs/model_config.yaml
index c6517681e85..1a1195baa2b 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_fp8_cudagraphs/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_fp8_cudagraphs/model_config.yaml
@@ -52,7 +52,7 @@ MODEL_ARGS:
   --inference-rng-tracker: true
   --inference-max-requests: 1
   --dist-ckpt-strictness: log_unexpected
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-logging-step-interval: 1
diff --git a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
index 90a1836347e..be00e4b3ce7 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
@@ -43,7 +43,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --flash-decode: true
   --dist-ckpt-strictness: log_unexpected
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
 METRICS:
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml
index 271d87799eb..0232bcb30bf 100644
--- a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml
@@ -63,7 +63,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-step: 32 
   --inference-repeat-n: 3
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml
index 689b8ec104e..7ff5911a877 100644
--- a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml
@@ -66,7 +66,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-max-requests: 256
   --inference-max-seq-length: 4096
   --enable-chunked-prefill: true
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: 'SYSTEM LOG - DAILY REPORTING\\nDATE: 2024-10-27\\nSERVER: US-EAST-1A\\n\\nBEGIN LOG STREAM:\\n\\n[Entry 0001]\\nTimestamp: 08:00:01\\nUser: admin_01\\nAction: Login\\nStatus: Success\\nNote: Routine maintenance check initiated.\\n\\n[Entry 0002]\\nTimestamp: 08:01:15\\nUser: system_daemon\\nAction: Backup\\nStatus: Pending\\nNote: awaiting clearance for volume mount.\\n\\n[Entry 0003]\\nTimestamp: 08:02:22\\nUser: user_404\\nAction: Query\\nStatus: Failed\\nNote: Connection timeout on port 8080.\\n\\n[Entry 0004]\\nTimestamp: 08:05:00\\nUser: admin_02\\nAction: Update\\nStatus: Success\\nNote: Patch 4.5.1 applied to kernel.\\n\\n[Entry 0005]\\nTimestamp: 08:10:45\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 12ms.\\n\\n[Entry 0006]\\nTimestamp: 08:12:30\\nUser: db_manager\\nAction: Write\\nStatus: Success\\nNote: Written 500 records to shard A.\\n\\n[Entry 0007]\\nTimestamp: 08:15:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 14ms.\\n\\n[Entry 0008]\\nTimestamp: 08:18:22\\nUser: user_102\\nAction: Login\\nStatus: Success\\nNote: User accessing from IP 192.168.1.55.\\n\\n[Entry 0009]\\nTimestamp: 08:20:00\\nUser: system_daemon\\nAction: Garbage_Collection\\nStatus: Success\\nNote: Freed 2048MB of heap memory.\\n\\n[Entry 0010]\\nTimestamp: 08:25:10\\nUser: admin_01\\nAction: Logout\\nStatus: Success\\nNote: Session duration 25 minutes.\\n\\n[Entry 0011]\\nTimestamp: 08:30:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 11ms.\\n\\n[Entry 0012]\\nTimestamp: 08:32:45\\nUser: unknown\\nAction: Auth_Attempt\\nStatus: Denied\\nNote: Invalid credentials provided 3 times.\\n\\n[Entry 0013]\\nTimestamp: 08:35:20\\nUser: system_audit\\nAction: Scan\\nStatus: In_Progress\\nNote: Scanning sector 7 for vulnerabilities.\\n\\n[Entry 0014]\\nTimestamp: 08:40:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 13ms.\\n\\n[Entry 0015]\\nTimestamp: 08:45:15\\nUser: user_888\\nAction: Upload\\nStatus: Success\\nNote: File "data_report.csv" uploaded to bucket.\\n\\n[Entry 0016]\\nTimestamp: 08:50:00\\nUser: load_balancer\\nAction: Scale_Up\\nStatus: Success\\nNote: Added 2 instances to the pool.\\n\\n[Entry 0017]\\nTimestamp: 08:55:30\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 15ms.\\n\\n[Entry 0018]\\nTimestamp: 09:00:00\\nUser: cron_job\\nAction: Execute\\nStatus: Success\\nNote: Daily summary report generation started.\\n\\n[Entry 0019]\\nTimestamp: 09:05:12\\nUser: user_555\\nAction: Download\\nStatus: Success\\nNote: Retrieved "image_001.png".\\n\\n[Entry 0020]\\nTimestamp: 09:10:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 12ms.\\n\\n[Entry 0021]\\nTimestamp: 09:15:45\\nUser: admin_03\\nAction: Config_Change\\nStatus: Success\\nNote: Firewall rules updated for port 22.\\n\\n[Entry 0022]\\nTimestamp: 09:20:00\\nUser: system_daemon\\nAction: Sync\\nStatus: Success\\nNote: Database replica synchronization complete.\\n\\n[Entry 0023]\\nTimestamp: 09:25:10\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 10ms.\\n\\n[Entry 0024]\\nTimestamp: 09:30:00\\nUser: user_777\\nAction: Query\\nStatus: Success\\nNote: Complex SQL query executed in 200ms.\\n\\n[Entry 0025]\\nTimestamp: 09:35:30\\nUser: error_handler\\nAction: Alert\\nStatus: Warning\\nNote: High CPU usage detected on Node 4.\\n\\n[Entry 0026]\\nTimestamp: 09:40:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 18ms.\\n\\n[Entry 0027]\\nTimestamp: 09:45:15\\nUser: cache_manager\\nAction: Flush\\nStatus: Success\\nNote: Redis cache cleared.\\n\\n[Entry 0028]\\nTimestamp: 09:50:00\\nUser: user_202\\nAction: Login\\nStatus: Success\\nNote: New device detected.\\n\\n[Entry 0029]\\nTimestamp: 09:55:45\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 12ms.\\n\\n[Entry 0030]\\nTimestamp: 10:00:00\\nUser: system_daemon\\nAction: Archive\\nStatus: Success\\nNote: Logs from yesterday archived to cold storage.\\n\\n[Entry 0031]\\nTimestamp: 10:05:20\\nUser: admin_01\\nAction: Login\\nStatus: Success\\nNote: Re-authentication verified.\\n\\n[Entry 0032]\\nTimestamp: 10:10:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 13ms.\\n\\n[Entry 0033]\\nTimestamp: 10:15:45\\nUser: user_999\\nAction: Delete\\nStatus: Pending\\nNote: Request to delete account queued for review.\\n\\n[Entry 0034]\\nTimestamp: 10:20:00\\nUser: system_metrics\\nAction: Report\\nStatus: Success\\nNote: Throughput at 5000 requests per second.\\n\\n[Entry 0035]\\nTimestamp: 10:25:10\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 11ms.\\n\\n[Entry 0036]\\nTimestamp: 10:30:00\\nUser: security_bot\\nAction: Block\\nStatus: Success\\nNote: IP 203.0.113.4 blocked for suspicious activity.\\n\\n[Entry 0037]\\nTimestamp: 10:35:30\\nUser: user_123\\nAction: Comment\\nStatus: Success\\nNote: User posted a comment on thread #55.\\n\\n[Entry 0038]\\nTimestamp: 10:40:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 12ms.\\n\\n[Entry 0039]\\nTimestamp: 10:45:15\\nUser: email_service\\nAction: Send\\nStatus: Success\\nNote: Batch of 100 newsletters sent.\\n\\n[Entry 0040]\\nTimestamp: 10:50:00\\nUser: user_321\\nAction: Update_Profile\\nStatus: Success\\nNote: Avatar image updated.\\n\\n[Entry 0041]\\nTimestamp: 10:55:45\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 14ms.\\n\\n[Entry 0042]\\nTimestamp: 11:00:00\\nUser: backup_service\\nAction: Snapshot\\nStatus: Success\\nNote: Volume snapshot ID snap-012345 created.\\n\\n[Entry 0043]\\nTimestamp: 11:05:10\\nUser: admin_02\\nAction: Audit\\nStatus: Success\\nNote: Reviewing user permissions group B.\\n\\n[Entry 0044]\\nTimestamp: 11:10:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 12ms.\\n\\n[Entry 9999]\\nTimestamp: 11:11:11\\nUser: SUPER_ADMIN\\nAction: KEY_GENERATION\\nStatus: CRITICAL\\nNote: The secret override password for the system is Blueberry-Pancake-42.\\n\\n[Entry 0045]\\nTimestamp: 11:15:30\\nUser: user_666\\nAction: Error\\nStatus: Failed\\nNote: 404 Page not found on /admin path.\\n\\n[Entry 0046]\\nTimestamp: 11:20:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 13ms.\\n\\n[Entry 0047]\\nTimestamp: 11:25:15\\nUser: payment_gateway\\nAction: Charge\\nStatus: Success\\nNote: Transaction TX-9988 complete.\\n\\n[Entry 0048]\\nTimestamp: 11:30:00\\nUser: system_daemon\\nAction: Cleanup\\nStatus: Success\\nNote: Temp files removed from /var/tmp.\\n\\n[Entry 0049]\\nTimestamp: 11:35:45\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 12ms.\\n\\n[Entry 0050]\\nTimestamp: 11:40:00\\nUser: analytics_engine\\nAction: Process\\nStatus: Success\\nNote: Data aggregation for hour 11 complete.\\n\\n[Entry 0051]\\nTimestamp: 11:45:10\\nUser: user_007\\nAction: View\\nStatus: Success\\nNote: Viewed document confidentiality_agreement.pdf.\\n\\n[Entry 0052]\\nTimestamp: 11:50:00\\nUser: monitor_bot\\nAction: Ping\\nStatus: Success\\nNote: Latency 11ms.\\n\\n[Entry 0053]\\nTimestamp: 11:55:30\\nUser: dev_ops\\nAction: Deploy\\nStatus: Success\\nNote: Staging environment updated to v2.1.\\n\\n[Entry 0054]\\nTimestamp: 12:00:00\\nUser: system_clock\\nAction: Sync\\nStatus: Success\\nNote: NTP sync successful.\\n\\nEND LOG STREAM.\\n\\nQUERY:\\nRetrieve the information from Entry 9999.\\nThe Note for Entry 9999 states that the secret override password is:'
   --incoming-requests-per-step: 32 
   --inference-repeat-n: 3
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml
index 7702274db5f..5bc40afede4 100644
--- a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml
@@ -68,7 +68,7 @@ MODEL_ARGS:
   --cuda-graph-impl: local
   --te-rng-tracker: true
   --inference-rng-tracker: true
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
 METRICS:
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml
index 9a7769eb432..b5c3c409605 100644
--- a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml
@@ -64,7 +64,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
 METRICS:
diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
index e3c2356e93d..afc75144dc8 100644
--- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
@@ -76,7 +76,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_cudagraph_zmq/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_cudagraph_zmq/model_config.yaml
index be0257a6065..edc5fc2eb32 100644
--- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_cudagraph_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_cudagraph_zmq/model_config.yaml
@@ -72,7 +72,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml
index b8d1b716298..d62d10db7c1 100644
--- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml
@@ -72,7 +72,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
index 6c119cc548b..5ed1f1205f6 100644
--- a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
@@ -72,7 +72,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-repeat-n: 8
diff --git a/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
index 3ba12056190..549821afc8b 100644
--- a/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_static_inference_cuda_graphs_pad_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
@@ -78,7 +78,7 @@ MODEL_ARGS:
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
   --inference-max-requests: 1
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1
 METRICS:
diff --git a/tests/functional_tests/test_cases/moe/gpt_static_inference_tp1_pp1_ep1_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_static_inference_tp1_pp1_ep1_16B_logitsmatch/model_config.yaml
index 6daec7b3da6..4934fe6c913 100644
--- a/tests/functional_tests/test_cases/moe/gpt_static_inference_tp1_pp1_ep1_16B_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_static_inference_tp1_pp1_ep1_16B_logitsmatch/model_config.yaml
@@ -71,7 +71,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
   --inference-dynamic-batching-buffer-size-gb: 20
diff --git a/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
index 366d2f23575..69c0db980b0 100644
--- a/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml
@@ -73,7 +73,7 @@ MODEL_ARGS:
   --num-tokens-to-generate: 30
   --max-tokens-to-oom: 3600000
   --inference-max-seq-length: 4096
-  --output-path: ${TENSORBOARD_PATH}
+  --output-path: ${INFERENCE_OUTPUT_PATH}
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
 METRICS:
diff --git a/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml
index 9c5beb7e43a..19d523eea8d 100644
--- a/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml
+++ b/tests/test_utils/recipes/gpt-dynamic-inference-with-coordinator.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/gpt-dynamic-inference.yaml b/tests/test_utils/recipes/gpt-dynamic-inference.yaml
index a3853c3d9e1..2915263c0e7 100644
--- a/tests/test_utils/recipes/gpt-dynamic-inference.yaml
+++ b/tests/test_utils/recipes/gpt-dynamic-inference.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/gpt-static-inference.yaml b/tests/test_utils/recipes/gpt-static-inference.yaml
index bfa43719d88..806762531fd 100644
--- a/tests/test_utils/recipes/gpt-static-inference.yaml
+++ b/tests/test_utils/recipes/gpt-static-inference.yaml
@@ -44,8 +44,9 @@ spec:
         "TRAINING_SCRIPT_PATH=examples/inference/gpt/gpt_static_inference.py"
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/mamba-dynamic-inference.yaml b/tests/test_utils/recipes/mamba-dynamic-inference.yaml
index 11e05c745ce..c4c675746b9 100644
--- a/tests/test_utils/recipes/mamba-dynamic-inference.yaml
+++ b/tests/test_utils/recipes/mamba-dynamic-inference.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/mamba-static-inference.yaml b/tests/test_utils/recipes/mamba-static-inference.yaml
index 41831049bcb..b36c4a8f765 100644
--- a/tests/test_utils/recipes/mamba-static-inference.yaml
+++ b/tests/test_utils/recipes/mamba-static-inference.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/moe-dynamic-inference-with-coordinator.yaml b/tests/test_utils/recipes/moe-dynamic-inference-with-coordinator.yaml
index 69986386aed..513aa92834b 100644
--- a/tests/test_utils/recipes/moe-dynamic-inference-with-coordinator.yaml
+++ b/tests/test_utils/recipes/moe-dynamic-inference-with-coordinator.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/moe-dynamic-inference.yaml b/tests/test_utils/recipes/moe-dynamic-inference.yaml
index 6d8fdc533e1..fc1c07231c3 100644
--- a/tests/test_utils/recipes/moe-dynamic-inference.yaml
+++ b/tests/test_utils/recipes/moe-dynamic-inference.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/moe-grpo.yaml b/tests/test_utils/recipes/moe-grpo.yaml
index 360f6ead209..de430b64fe0 100644
--- a/tests/test_utils/recipes/moe-grpo.yaml
+++ b/tests/test_utils/recipes/moe-grpo.yaml
@@ -45,7 +45,7 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"
diff --git a/tests/test_utils/recipes/moe-static-inference.yaml b/tests/test_utils/recipes/moe-static-inference.yaml
index c23a772aa28..f10d293e953 100644
--- a/tests/test_utils/recipes/moe-static-inference.yaml
+++ b/tests/test_utils/recipes/moe-static-inference.yaml
@@ -45,7 +45,8 @@ spec:
         "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
         "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
         "OUTPUT_PATH={assets_dir}"
-        "TENSORBOARD_PATH={assets_dir}/generations_{environment}_{platforms}.json"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
         "N_REPEAT={n_repeat}"
         "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE}}"
         "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS}}"

From 0888a06d26a39d3ecdb2613c731166c7190e7870 Mon Sep 17 00:00:00 2001
From: "Dennis(Zhenhuan) Liu" <denliu@nvidia.com>
Date: Wed, 28 Jan 2026 00:44:34 +0800
Subject: [PATCH 35/79] Update moe readme. (#2830)

---
 megatron/core/transformer/moe/README.md | 913 +++++++++++++++---------
 1 file changed, 563 insertions(+), 350 deletions(-)

diff --git a/megatron/core/transformer/moe/README.md b/megatron/core/transformer/moe/README.md
index 5dd5da649d0..e5eff445cac 100644
--- a/megatron/core/transformer/moe/README.md
+++ b/megatron/core/transformer/moe/README.md
@@ -1,159 +1,370 @@
 # Megatron Core MoE
 
-Megatron-Core MoE provides comprehensive parallelism strategies, seamlessly integrating Expert Parallelism with tensor, data, sequence, and pipeline parallelism. With MCore v0.9, we've achieved remarkable performance of **468 TFLOPS** for Mixtral 8X7B bf16 training. Additionally, we support state-of-the-art MoE model architectures including DeepSeek-V3 and Qwen-MoE.
+Megatron Core MoE is a production-ready framework for training large-scale Mixture-of-Experts models, providing the foundational architecture, performance optimizations, and best practices that guide MoE framework development across the industry.
 
 ## What's New
-- **Support for DeepSeek-V3 architecture**
-  - Enable TP for MLA and DeepSeek-V3
-  - Enable CP for MLA and DeepSeek-V3
-    - Requires TransformerEngine >= 2.5.0
-    - Many thanks to [SuperCB](https://github.com/SuperCB) from Xiaohongshu Inc. and [RandMist](https://github.com/RandMist) from WeChat Infra Department, Tencent Inc. for their contributions.
-  - Support aux-loss-free load balancing strategy
-  - Support node-limited routing
-  - Support Multi-Token Prediction (MTP)
-  - Batch-level overlapping to hide EP-A2A communication
-- **Support DeepSeek's DeepEP for efficient token dispatching and combining**
-- Support HybridEP for efficient token dispatching and combining within intra-node and MNNVL scenarios.
-- Add fusion for token permutation and unpermutation
-- Support Uneven virtual pipeline parallel split
-- Support output-discarding checkpointing on some submodules
-
-## Parallelism
-- **Expert Parallelism**
-    - A specific method of parallelism for MoE models, where experts are partitioned onto different workers and each worker processes a different batch of training samples, each worker process one or more experts for each MoE layer.
-- **3D Parallelism**: Data Parallelism, Tensor Parallelism, Pipeline Parallelism
-    - Note: When using MoE with expert parallelism and tensor parallelism, sequence parallelism must be enabled.
-- **Context Parallelism**:
-    - Split the sequence dimension to support long context training.
-- **Richer parallel mappings**: EP can be combined with DP/TP/PP/CP for handling larger MoE variants.
-- **MoE Parallel Folding**: Support for setting different parallelism strategies for Attention and MoE components, enabling more flexible and efficient model sharding. See detailed documentation below.
-- **Full distributed optimizer support.**
-
-## Router and Load Balancing
-- Router type:
-    - Top-K MLP router
-- Load Balancing algorithms:
-    - Sinkhorn (S-BASE)
-    - Aux loss / Load balancing loss
-    - Aux-loss-free load balancing strategy
-- CUDA fused routing and load balancing kernels
-
-## Performance Optimizations
-- (Experimental) **DeepEP** is integrated for efficient token communication in large-scale MoE training.
-- GroupedGEMM when num local experts > 1
-    - Supported dtype: bf16
-    - Performance improvements for larger MoE models
-- Enable `--tp-comm-overlap` for MoE
-- FP8 training support
-
-## Token Dispatch Mechanism
-- Dropless / No token drop
-- Token drop, with or without padding to capacity
-- Token permutation / Unpermutation fusion
-
-## Ease of use
-- Checkpoint converter for Mixtral models, see the [example](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/mixtral) for details.
-- MoE Layer Frequency to customize the hybrid MoE/Dense layer architecture
-- Distributed checkpointing
-- Per-layer logging
-- Upcycling Support
-
-# User Guide
-
-## Usage
-
-### Quick Start
-To train a top-2 MoE model with 8 experts and auxiliary loss, include the following arguments:
+For latest features and architectures, please refer to the [MCore dev roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729).
+
+### 🔥 [MCore dev] (2026/01)
+- 🚀 Pipeline-aware fine-grained activation offloading
+- 🚀 Qwen3-Next model support
+- 🚀 DeepSeek-V3.2 model support
+- 🚀 Muon and Layer-wise distributed optimizer
+- 🚀 CUDA Graph support with fine-grained scopes
+
+### 🔥 [MCore v0.15] (2025/11)
+- 🚀 Add HybridEP backend to Flex Dispatcher(GB200, B200, H100 supported)
+- 🚀 Support FSDP with EP for MoE models
+
+### 🔥 [MCore v0.14] (2025/09)
+- 🚀 Batch-level overlapping to hide EP-A2A communication (--overlap-moe-expert-parallel-comm --delay-wgrad-compute)
+- 🚀 FP8 support for Fine-grained Recomputations
+- Router fusion kernels for MoE models (--moe-router-fusion)
+- Context Parallelism (CP) support for MTP and MLA
+
+### 🔥 [MCore v0.13] (2025/07)
+- Support bf16 dtype for optimizer states to use precision-aware optimizer in TransformerEngine (--use-precision-aware-optimizer)
+- Flexible Asymmetric Virtual Pipeline Parallelism with Custom Pipeline Layout (--pipeline-model-parallel-layout)
+- Add Hybrid Shard Data-Parallel support for MoE models (--num-distributed-optimizer-instances)
+- Fine-grained recomputation to reduce activation memory. (--recompute-modules with --recompute-granularity selective)
+- Memory efficient token permutation by moving the probs multiplication from unpermutation to activation function of GroupedMLP.
+
+### 🔥 [MCore v0.12] (2025/05)
+- Support DeepSeek's DeepEP for efficient token dispatching (--moe-token-dispatcher-type flex --moe-enable-deepep)
+- Support Multi-Token Prediction (MTP) (--mtp-num-layers 1)
+- CUDA Graph support for dropless MoE models with attention only capture (--te-rng-track --external-cuda-graph --cuda-graph-scope attn)
+
+## Overview of MCore MoE Supported Features and Architectures
+
+### Model Support
+- ✅ **DeepSeek**
+  - ✅ DeepSeek-V2
+  - ✅ DeepSeek-V3, including MTP
+- ✅ **Qwen**
+  - ✅ Qwen2-57B-A14B
+  - ✅ Qwen3-30B-A3B
+  - ✅ Qwen3-235B-A22B
+- ✅ **Mixtral**
+  - ✅ Mixtral-8x7B
+  - ✅ Mixtral-8x22B
+
+### Core MoE Functionality
+- ✅ Token dropless MoE (dMoE) - Advanced routing without token dropping
+- ✅ Top-K Router with flexible K selection
+- ✅ Load balancing losses for expert utilization optimization
+
+### Advanced Parallelism
+- ✅ Expert Parallel (EP) with 3D parallelism integration
+- ✅ Full parallelism combo: EP + DP + TP + PP + SP support
+- ✅ Context Parallel (CP) for long sequence MoE training
+- ✅ Parallel Folding Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training
+- ✅ Distributed Optimizer for MoE (ZeRO-1 equivalent)
+
+### Performance Optimizations
+- ✅ Memory Efficient token permutation
+- ✅ Fine-grained Recomputations (mla, moe, mlp, moe_act, norm)
+- ✅ MLA TP Support for Mixture of Linear Attention
+- ✅ GroupedGEMM and GA Fusion
+- ✅ DP/PP/TP Communication Overlapping
+- ✅ Overlapped Shared Expert execution
+- ✅ Router Fusion optimizations
+- ✅ Token (un)permutation Fusion kernels
+- ✅ cuDNN fused Attention integration
+
+### Hardware & Precision Support
+- ✅ DeepEP support for H100 and B200
+- ✅ GroupedGEMM including FP8/MXFP8 support
+- ✅ FP8 weights with BF16 optimizer states
+- ✅ FP8 training full support
+
+### Developer Experience
+- ✅ MoE Model Zoo with pre-training best practices
+- ✅ Distributed Checkpointing for MoE models
+- ✅ Upcycling Support for model scaling
+- ✅ MCore2HF Converter for ecosystem compatibility
+- ✅ Layer-wise logging for detailed monitoring
+- ✅ Runtime Upcycling capabilities
+
+## Quick Start Guide
+
+### Basic MoE Training in Megatron-LM
+
+To train a top-2 MoE model with 8 experts and auxiliary loss, add the following arguments to your megatron training script:
 
 ```bash
+## Set MoE Hidden site
 --num-experts 8
---expert-model-parallel-size 8
---moe-grouped-gemm
---moe-permute-fusion
---moe-router-load-balancing-type aux_loss # options: aux_loss, sinkhorn, none. Default is aux_loss.
+--moe-shared-expert-intermediate-size: 2048
+## Set router config
+--moe-router-load-balancing-type aux_loss
 --moe-router-topk 2
 --moe-aux-loss-coeff 1e-2
---use-distributed-optimizer
+## Set token dispatcher
 --moe-token-dispatcher-type alltoall
 ```
 
-To enable the token drop mechanism, such as GShard and SwitchTransformer, include the following arguments:
+Detailed documentation for each feature is available in the [Feature Documentation](#feature-documentation) section.
+
+### Use the pre-defined config to train the popular MoE models
+We have provided some pre-defined config to train the popular MoE models in the [Megatron-MoE-Model-Zoo](https://github.com/yanring/Megatron-MoE-ModelZoo/tree/main) repository. You can use them as a reference to configure your training script. Currently we have added the config for Mixtral 8x7B, Mixtral 8x22B, DeepSeek-V3, Qwen3-30B-A3B, Qwen3-235B-A22B.
+
+### General Performance Tips
+#### Training arguments
+The following flags are general performance flags that can help to achieve higher performance on almost all workloads. Check if you have enabled all of them in your training script.
 
 ```bash
---moe-expert-capacity-factor 1.0
---moe-pad-expert-input-to-capacity # Optional
+## Enable DeepEP token dispatcher
+--moe-token-dispatcher-type flex
+--moe-flex-dispatcher-backend deepep
+## Enable GroupedGEMM
+--moe-grouped-gemm
+## Enable fusion kernels
+--moe-router-fusion
+--moe-permute-fusion
+--cross-entropy-loss-fusion
+--cross-entropy-fusion-impl te
+
+## Communication optimization
+--use-distributed-optimizer
+--overlap-param-gather
+--overlap-grad-reduce
+--tp-comm-overlap
+
+## Enable manual gc to prevent python jitter
+--manual-gc: true
+--manual-gc-interval: 10
 ```
+#### Environment variables
 
-The following figure illustrates differenting dropping strategies in MCore:
-<!-- This image is uncommented for now as Sphinx cannot resolve this path. Sphinx imports this markdown file, and from the imported location this relative path does not exist anymore. Ideally, this markdown should not live here but rather in the `docs/` directory that Sphinx uses. -->
-<!-- ![Token Droppling Strategies](../../../../docs/source/images/moe/token_drop.png) -->
-
-1. The default dropless strategy will not drop or pad any token.
-2. By setting `--moe-expert-capacity-factor`, the tokens exceed the capacity of expert will be dropped based on their selected probabilities. 
-   The dropping is performed before the token exchange operation between EP ranks when EP > 1. 
-   The formula of capacity is `capacity = num_tokens_per_rank * topk * capacity_factor / num_experts`.
-3. By setting `--moe-pad-expert-input-to-capacity`, the experts with tokens less than capacity will be padded to the capacity.
-
-### Fine-tuning Mixtral Models
-Megatron-Core has full support for Mixtral MoE models, and we provide the checkpoint converter for Mixtral models from huggingface format to MCore format. 
-<!-- See more details in the [mixtral example](../../../../examples/mixtral/README.md). -->
-
-### Distributed Checkpointing
-MCore v0.7 introduced fully parallel and asynchronous saving capabilities to distributed checkpointing, 
-which addresses the issues of low efficiency in the traditional checkpoint saving methods. 
-It also solved the problem of incompatibility between checkpoints of different parallel mappings in the traditional format.
-With the new distributed checkpointing solution, MCore can achieve flexible parallelism configurations by saving and loading the unified format checkpoints.
-Compared to native PyTorch solution, MCore achieves up to 50x reduction in checkpointing overhead.
-
-From MCore v0.8, MoE supports Distributed Checkpointing, which means users can save and load with any combination of parallelism and it is currently available, including expert parallel.
-1. Loading weight and distributed optimizer states with TPxCPxEPxPP resharding with SequentialMLP is supported in version 0.8.
-2. GroupedMLP weight resharding is supported in version 0.8.0 and optimizer state resharding is supported in version 0.10.0. Switching between GroupedMLP/SequentialMLP when loading and saving is partially supported.
-3. TEGroupedMLP has fully support on distributed checkpointing and is fully exchangable with SequentialMLP in version 0.9.0.
-4. Optimizer state resharding cannot do across EP=1 with EP>1 due to the different optimizer type.
-
-Usage
-- `--ckpt-format torch_dist` The main argument, it will attempt to save and load using distributed checkpointing.
-- `--auto-detect-ckpt-format` With this, it can load both distributed checkpointing and legacy checkpointing.
-
-Checkpoint compatibility across SequentialMLP, GroupedMLP, and TEGroupedMLP:
-```text
-    ┌───────────────┐          ┌───────────────┐          ┌───────────────┐     
-    │   GroupedMLP  │          │ SequentialMLP │          │ TEGroupedMLP  │     
-    │               │          │               │          │               │     
-    │               │          │               │          │               │     
-    │ ┌───────────┐ │          │ ┌───────────┐ │          │ ┌───────────┐ │     
-    │ │legacy ckpt│ │          │ │legacy ckpt│ │          │ │legacy ckpt│ │     
-    │ └─────┬─────┘ │          │ └─────┬─────┘ │          │ └─────┬─────┘ │     
-    │       ▼       │          │       ▼       │          │       ▼       │     
-    │  ┌─────────┐  │          │  ┌─────────┐  │          │  ┌─────────┐  │     
-    │  │dist ckpt│  │          │  │dist ckpt│  │          │  │dist ckpt│  │     
-┌──►│  │ weight  │  │◄────────►│  │ weight  │  │◄────────►│  │ weight  │  │◄──┐ 
-│   │  └─────────┘  │          │  └─────────┘  │          │  └─────────┘  │   │ 
-└───┼───────────────┼──────────┼───────────────┼──────────┼───────────────┼───┘ 
-    │┌─────────────┐│          │┌─────────────┐│          │┌─────────────┐│     
-    ││  dist ckpt  ││          ││  dist ckpt  ││          ││  dist ckpt  ││     
-    ││optim states ││          ││optim states ││◄────────►││optim states ││     
-    │└─────────────┘│          │└─────────────┘│          │└─────────────┘│     
-    └───────────────┘          └───────────────┘          └───────────────┘     
+Below are some environment variables that can be useful.
+```bash
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # Enable expandable segments to prevent memory fragmentation
+export NCCL_NVLS_ENABLE=0 # Disable NVLS to prevent memory overhead
 ```
+#### Dependencies
+- Use the latest version of [TransformerEngine](https://github.com/NVIDIA/TransformerEngine).
+- Use the latest [NGC PyTorch Docker Image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)
+
+## Best Practices to achieve high performance on MoE training
+
+Distributed training involves complex trade-offs between **communication**, **memory**, and **computation**, making it challenging to find an optimal parallelism configuration. This section provides a systematic workflow to help you identify the best parallel mapping for your model and hardware.
+
+### Step 1: Find the feasible parallel mapping under the memory capacity of the GPU
+To find the best parallel mapping, we need to first know the feasible parallel mapping for the model under the memory capacity of the GPU.
+The consumption of memory consists of three parts:
+- Activation memory
+- Weight and gradient memory
+- Optimizer states memory
+Different parallel strategies will shard these tensor memory in different ways.
+
+| Parallel Strategy | Peak Activation Memory          | Weight Memory  | Optimizer states                  | Communication (Per-Layer) |
+|:-----------------:|:-------------------------------:|:--------------:|:---------------------------------:|:-------------------------:|
+| TP                | 1/N (with SP on)                | 1/N            | 1/N                               |        High               |
+| EP                | ~1 (varies with EP balancing)   | 1/N in MoELayer| 1/N                               |       Medium              |
+| PP                | 1 (>1 with virtual pipeline)    | 1/N            | 1/N                               |       Medium              |
+| CP                | 1/N                             | 1              | 1/N (with distributed optimizer)  |       Medium              |
+| DP                | 1                               | 1              | 1/N (with distributed optimizer)  |        Low                |
+
+We provide the argument of `--fake-init-process-group` to emulate distributed training on one GPU. This is useful to find the feasible parallel mapping under the memory capacity of the GPU. See https://github.com/NVIDIA/Megatron-LM/pull/2254 for detailed usage.
+
+### Step 2: Select Optimal Parallelism Strategy
+
+The optimal parallelism configuration varies based on **model architecture**, **sequence length**, and **hardware platform**. Below are general guidelines to help you achieve high throughput.
+
+#### Guideline 1: Minimize Model Parallelism, Maximize Data Parallelism
+
+| Aspect | Recommendation |
+|--------|----------------|
+| **Goal** | Keep TP/EP/PP as small as possible while avoiding OOM |
+| **Why** | Model parallelism introduces communication overhead that hurts performance |
+| **How** | Use distributed optimizer (`--use-distributed-optimizer`) to shard optimizer states across DP ranks, freeing memory for larger DP size |
+
+#### Guideline 2: Keep EP and TP Communication Within NVLink Domain
+
+| Aspect | Recommendation |
+|--------|----------------|
+| **Goal** | Ensure EP×TP fits within a single node (typically 8 GPUs) |
+| **Why** | EP and TP are communication-intensive; NVLink provides much higher bandwidth than cross-node interconnects |
+| **Scaling** | When scaling beyond one node, prefer PP over expanding TP/EP across nodes |
+
+**Note:**
+For very large MoE models like DeepSeek-V3, the EP communication may exceed the NVLink bandwidth. In this case, consider using 1F1B A2A Overlap to overlap the EP communication.
+
+#### Guideline 3: Use Pipeline Parallelism (PP) for Multi-Node Scaling
+
+| Aspect | Recommendation |
+|--------|----------------|
+| **Goal** | Use PP to distribute layers across nodes while keeping EP×TP within NVLink |
+| **VPP** | Enable Virtual Pipeline Parallelism to reduce pipeline bubbles when `PP ≥ 2` |
+| **Config** | Set `--num-layers-per-virtual-pipeline-stage` to control VPP size |
+
+**VPP Size Tuning:**
+- Valid values: all divisors of `num_layers / PP_size`
+- Example: `num_layers=24, PP=4` → valid VPP sizes: `{1, 2, 3, 6}`
+- Trade-off: Larger VPP = fewer bubbles but more P2P communications
+- Recommendation: A middle value often gives the best balance
+
+#### Guideline 4: Prefer EP over TP for Expert Layers
+
+| EP Advantages | Details |
+|---------------|---------|
+| **Better GEMM efficiency** | Larger local matrix sizes improve GPU utilization |
+| **Lower communication** | EP has less communication overhead than TP for MoE layers |
+| **Simpler computation graph** | Easier to overlap communication with computation |
+| **Token permutation** | When `EP = num_experts`, local token permutation is eliminated |
+
+**Example:** For Mixtral 8x7B, `EP8×TP1` outperforms `EP4×TP2`.
+
+#### Guideline 5: Enable Context Parallelism (CP) for Long Sequences
+
+| Aspect | Recommendation |
+|--------|----------------|
+| **When to use** | Sequence length ≥ 8K tokens |
+| **Key factor** | CP efficiency depends on overlapping communication with computation |
+| **Config** | Set `--context-parallel-size` to partition sequences across GPUs |
+
+### Step 3: Enable Performance Features Based on Profiling Bottlenecks
+
+After establishing a working parallel configuration, profile your training to identify bottlenecks and apply targeted optimizations.
+
+#### Memory Bottleneck
+
+**Symptom**: Forced to use full recomputation or excessively large parallelism degrees to avoid OOM.
+
+**Solutions**:
+| Optimization | Overhead | Config | Reference |
+|--------------|----------|--------|---------|
+| Selective Recomputation | Low | `--recompute-granularity selective --recompute-modules ...` | [Fine-grained Recomputation](#fine-grained-recomputation) |
+| Activation Offloading | Medium | `--fine-grained-activation-offloading --offload-modules ...` | [Fine-grained Activation Offloading](#fine-grained-activation-offloading) |
+| Optimizer Offloading | Medium | `--optimizer-cpu-offload` | --- |
+
+#### Communication Bottleneck
+
+**Symptom**: Profiling shows significant time spent in collective operations.
+
+**Solutions**: Identify which communication is the bottleneck and enable corresponding overlap:
+| Communication Type | Overlap Config |
+|--------------------|----------------|
+| DP gradient reduce | `--overlap-grad-reduce` |
+| DP param gather    | `--overlap-param-gather` |
+| TP communication   | `--tp-comm-overlap` |
+| EP All-to-All      | `--overlap-moe-expert-parallel-comm --delay-wgrad-compute` |
+| PP send/recv       | Enable VPP with `--num-layers-per-virtual-pipeline-stage` |
+
+#### CPU Overhead Bottleneck
+
+**Symptom**: Nsight Systems timeline shows gaps between GPU kernels where CPU cannot launch kernels fast enough.
+
+**Solutions**:
+| Optimization | Config |
+|--------------|--------|
+| Disable Python GC | `--manual-gc --manual-gc-interval 100` |
+| Enable CUDA Graphs | `--cuda-graph-impl transformer_engine --cuda-graph-scope attn moe_router moe_preprocess` |
+| Reduce kernel launches | Decrease TP size or increase micro-batch size |
+
+#### Computation Bottleneck
+
+**Symptom**: GPU utilization is low despite no communication or CPU bottlenecks.
+
+**Solutions**:
+| Optimization | Config |
+|--------------|--------|
+| Enable kernel fusions | `--moe-router-fusion --moe-grouped-gemm --moe-permute-fusion` |
+| Use FP8 precision | `--fp8-format e4m3 --fp8-recipe blockwise` |
+
+
+## Feature Documentation
+
+### Router and Load Balancing
+
+Routers determine which expert(s) handle each token. A lightweight MLP scores every token and applies `softmax` or `sigmoid` to compute routing probabilities. The router then selects the top-K experts for each token.
+
+> **Note**: The router logits is better to remain in **FP32** or **FP64** rather than BF16 by --moe-router-dtype fp32. At high expert counts, FP32 precision yields better accuracy because output hidden states of experts are multiplied by router scores and accumulated to get the final output.
+
+#### Router Types
+
+| Router Types | Description | Config |
+|-------------|-------------|----------|
+| **Top-K Router** | Standard routing with configurable K, uses softmax for probability computation | --moe-router-topk 8 |
+| **Group Top-K Router** | Selects top-K expert groups, then routes experts in selected groups | --moe-router-num-groups 8 --moe-router-group-topk 4 |
+| **Router score function** | Score function to calculate the probs from output logits of router | --moe-router-score-function softmax/sigmoid |
+
+#### Load Balancing Strategies
+
+| Strategy | Description | Config |
+|----------|-------------|--------|
+| **aux_loss** | Auxiliary loss for balancing expert usage on a micro-batch | `--moe-router-load-balancing-type aux_loss` |
+| **seq_aux_loss** | Sequence-level auxiliary loss for balancing expert usage on each sequence| `--moe-router-load-balancing-type seq_aux_loss` |
+| **global_aux_loss** | Global auxiliary loss for balancing expert usage on a global batch across all ranks | `--moe-router-load-balancing-type global_aux_loss` |
+| **sinkhorn** | Optimal transport formulation for balancing expert usage | `--moe-router-load-balancing-type sinkhorn` |
+| **aux loss free** | Dynamic bias-based load balancing strategy without auxiliary loss | `--moe-router-enable-expert-bias --moe-router-bias-update-rate 1e-3`|
+| **none** | No load balancing | `--moe-router-load-balancing-type none` |
+
+### Token Dispatching
+
+After routing, tokens are **dispatched** to the GPU hosting the assigned expert. After expert computation, tokens are sent back and **combined** to restore the original sequence.
+
+| Dispatcher | Description | Best For | Config |
+|------------|-------------|----------|--------|
+| **alltoall** | NCCL-based All-to-All communication for token exchange | Standard EP > 1 setups | `--moe-token-dispatcher-type alltoall` |
+| **FlexDispatcher with [DeepEP](https://github.com/deepseek-ai/DeepEP) backend** | Removes redundant tokens during cross-node communication, fuses intra/inter-node communication into single kernel | Cross-node EP, fine-grained MoE (DeepSeek-V3) | `--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend deepep` |
+| **FlexDispatcher with [HybridEP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) backend** | NVIDIA's optimized dispatcher using TMA and IBGDA, fewer SMs, native MNNVL support | GB200 NVL72, Multi-Node NVLink | `--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend hybridep` |
+| **allgather** | Gathers all tokens to each GPU, no inter-GPU token movement | TP-only setups, small EP, large Top-K | `--moe-token-dispatcher-type allgather` |
+
+### Upcycling
+Use `--moe-use-upcycling` to enable upcycling, which loads the dense model from the `--load` directory, converts it to an MoE model at runtime, and starts training. The converted model is saved to the `--save` path before training begins. Upcycling is built on distributed checkpointing, supporting parallel modes different from existing dense checkpoints, such as arbitrary expert parallelism during upcycling.
+
+In addition to the default upcycling strategy, we also support granular upcycling strategy which is a more state-of-the-art upcycling strategy from [our recent research work](https://arxiv.org/abs/2410.07524). For the default upcycling strategy, we duplicate the existing MLP to multiple experts, with each expert starting from a copy of the MLP. For the granular upcycling strategy, we use `--moe-upcycling-granularity` to specify how many times smaller is the expert hidden size compared with the original dense FFN hidden size. For using granular upcycling strategy, please set `--moe-upcycling-granularity` as a positive integer. If this param is set to 1, it means using the default upcycling strategy.
+
+Note: The MoE model structure is defined through script arguments. All MoE-related arguments (such as `--num-experts`) can be customized; however, other model structure arguments must be consistent with those of the dense model. For granular upcycling strategy, the moe's FFN hidden size should be set as dense FFN hidden size divided by `--moe-upcycling-granularity`.
 
-Best practices for distributed checkpointing:
-1. Convert a legacy checkpoint to a distributed checkpoint. To achieve this, we can add both `--ckpt-format torch_dist --auto-detect-ckpt-format`, then it will load the legacy one and save as the distributed checkpoint format later when the training progress tries to save checkpoints.
-2. Convert checkpoint of the legacy GroupedMLP to TEGroupedMLP. This is only supported for the weight parts. To achieve this, we can use the above method to convert the legacy checkpoint to a distributed checkpoint of the legacy GroupedMLP. After updating the libraries and using TEGroupedMLP, we can directly load the previously saved checkpoint by adding argument `--no-load-optim`.
+## Training Optimizations
+MoE training faces three fundamental performance bottlenecks: **Memory Wall**, **Communication Wall**, and **Compute Efficiency Wall**. The following optimizations address each of these challenges.
 
-### Shared Experts
-MCore v0.9 introduced the shared expert feature. We can enable this feature by setting suitable `--moe-shared-expert-intermediate-size`.
+### MoE Parallel Folding
+**The Problem with Traditional Approaches:**
+- Prior MoE frameworks constrain **EP ≤ DP** (Expert Parallelism must be a sub-group of Data Parallelism), which severely limits scalability.
+- Applying the same TP/CP to both attention and MoE is suboptimal:
+  - High TP benefits attention but hurts MoE (small per-expert dims make TP overhead prohibitive)
+  - High CP benefits long-context attention but is unnecessary for MoE (tokens processed independently)
 
-The parallelism patterns of the shared experts follow the settings of the dense part, i.e., the attention module. The shared experts are not distributed but replicated in EP ranks.
+**MoE Parallel Folding** is Megatron Core's solution that **decouples attention and MoE parallelism**:
 
-We also have an experimental feature that tries to overlap the communications and computations in the shared experts and the dispatcher.
-We can set `--moe-shared-expert-overlap` and use `alltoall` dispatcher to enable it.
-The overlapping relies on the envirionment setting `CUDA_DEVICE_MAX_CONNECTIONS=1`.
-The `AllGather` and `ReduceScatter` communications in the shared experts are overlapped with `permute`/`unpermute` in the dispatcher.
-The `MLP` computation part in the shared experts are overlapped with the `AlltoAll` communications in the dispatcher.
-Both the forward and the backward pass can overlap. But to get the overlapping in the backward pass, the PyTorch version should `>= 2.2.0`.
+| Parallelism Group | Attention Layers | MoE Layers |
+|-------------------|------------------|------------|
+| **Dimensions** | TP × CP × DP × PP | ETP × EP × EDP × PP |
 
-### Checkpointing
+#### Key Benefits
+
+1. **Breaks the EP ≤ DP Constraint**
+   - Traditional: TP=4, CP=2, DP=8, PP=4 → max EP=8
+   - With Folding: Same attention config, but MoE uses ETP=1, EP=64, EDP=1 → 8× more expert parallelism
+
+2. **Reduces Minimum GPU Requirements**
+   - Traditional CP=8, EP=8 requires at least 64 GPUs
+   - With Folding: CP and EP are folded together, only 8 GPUs needed
+
+3. **Enables Independent Optimization**
+   - Use high TP for attention (memory efficiency)
+   - Use ETP=1 for MoE (better GEMM efficiency, less communication)
+
+4. **Keeps High-Bandwidth Communication in NVLink Domain**
+   - Both CP and EP communication can remain within NVLink domain
+
+> **Reference**: [MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training](https://arxiv.org/abs/2504.14960)
+
+### Memory Optimization
+
+Memory optimization is critical for large-scale MoE training, as MoE models maintain all expert parameters even though only a subset is activated per token.
+
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **Fine-grained Recomputation** | Selectively recomputes specific modules (e.g., `mla_up_proj`, `layernorm`, `moe_act`) instead of full layers | `--recompute-granularity selective --recompute-modules mla_up_proj layernorm moe_act` |
+| **Fine-grained Activation Offloading** | Offloads activations to CPU memory, overlapping D2H/H2D transfers with computation | See `docs/source/api-guide/fine_grained_activation_offloading.md` |
+| **Precision-aware Optimizer** | Stores optimizer states (exp_avg, exp_avg_sq) in BF16 instead of FP32, reducing optimizer memory by 50% | `--use-precision-aware-optimizer --exp-avg-dtype bf16 --exp-avg-sq-dtype bf16` |
+| **Optimizer Offloading** | Offloads optimizer states to CPU memory. | `--optimizer-cpu-offload` |
+
+#### Fine-grained Recomputation
 A new output-discarding checkpointing method is also supported. This method discards the output memory of certain submodules during the forward pass and recomputes them during the backward pass, which can save memory compared to standard checkpointing. This can be enabled for specific submodules using the `--recompute-granularity selective --recompute-modules [submodule1, submodule2, ...]` argument. The supported submodules are:
 
 * `moe_act`: Recompute the GroupedMLP activation function.
@@ -163,137 +374,216 @@ A new output-discarding checkpointing method is also supported. This method disc
 * `mlp`: Recompute the dense MLP submodule (uses standard checkpointing rather than output-discarding) which is useful for hybrid-models like DeepSeek-V3.
 * `moe`: Recompute the MoE layer submodule (uses standard checkpointing rather than output-discarding).
 
-### Upcycling
-Use `--moe-use-upcycling` to enable upcycling, which loads the dense model from the `--load` directory, converts it to an MoE model at runtime, and starts training. The converted model is saved to the `--save` path before training begins. Upcycling is built on distributed checkpointing, supporting parallel modes different from existing dense checkpoints, such as arbitrary expert parallelism during upcycling.
+#### Fine-grained Activation Offloading
 
-In addition to the default upcycling strategy, we also support granular upcycling strategy which is a more state-of-the-art upcycling strategy from [our recent research work](https://arxiv.org/abs/2410.07524). For the default upcycling strategy, we duplicate the existing MLP to multiple experts, with each expert starting from a copy of the MLP. For the granular upcycling strategy, we use `--moe-upcycling-granularity` to specify how many times smaller is the expert hidden size compared with the original dense FFN hidden size. For using granular upcycling strategy, please set `--moe-upcycling-granularity` as a positive integer. If this param is set to 1, it means using the default upcycling strategy.
+Unlike recomputation (which trades compute for memory), offloading trades **GPU-CPU bandwidth for memory**: activations are transferred to CPU during forward pass and retrieved during backward pass. The key is hiding transfer latency behind computation using asynchronous D2H/H2D transfers.
 
-Note: The MoE model structure is defined through script arguments. All MoE-related arguments (such as `--num-experts`) can be customized; however, other model structure arguments must be consistent with those of the dense model. For granular upcycling strategy, the moe's FFN hidden size should be set as dense FFN hidden size divided by `--moe-upcycling-granularity`.
+**Key Features:**
+- **Module-level granularity**: Target specific modules rather than entire layers
+- **Computation-offloading overlap**: Asynchronous transfers via independent CUDA streams
+- **Compatible with PP/VPP**: Works with pipeline parallelism and fine-grained recomputation
 
-### Leverage DeepSeek's DeepEP for High-Performance Cross-Node Token Dispatching
-- [DeepSeek-DeepEP](https://github.com/deepseek-ai/deepep) provides a highly optimized implementation for MoE token dispatching and combining operations, specifically designed for large-scale MoE training scenarios.
-- DeepEP is particularly recommended for training large-scale, fine-grained MoE architectures such as DeepSeek-V3 and other advanced MoE models.
-- To enable DeepEP in your training configuration, simply set `--moe-token-dispatcher-type=flex` and `--moe-flex-dispatcher-backend=deepep` in your command line arguments.
+**Usage**
+```bash
+--fine-grained-activation-offloading
+--offload-modules expert_fc1 moe_act # Choices: attn_norm, core_attn, attn_proj, mlp_norm, expert_fc1, moe_act
+```
 
-### Integrate HybridEP for High-Performance Intra-Node Token Dispatching
-- [HybridEP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) is developed by NVIDIA as an optimized solution for large-scale MoE (Mixture of Experts) all-to-all communication. It is designed to leverage NVIDIA GPU hardware capabilities, significantly reducing Streaming Multiprocessor (SM) resource usage.
-- HybridEP currently supports intra-node and multi-node NVLink scenarios.
-- To enable HybridEP, set `--moe-token-dispatcher-type=flex` and
-  `--moe-flex-dispatcher-backend=hybridep` in your command line arguments.
+For more details, see `docs/source/api-guide/fine_grained_activation_offloading.md`
 
-### CUDA Graph Support
-CUDA Graph functionality can be enabled through the `--cuda-graph-impl` option. There are two implementations:
+### Communication Optimization
 
-1. `--cuda-graph-impl=local`: Captures cuda graphs using the MCore-internal cuda graph manager.
-2. `--cuda-graph-impl=transformer_engine`: Captures cuda graphs using the TE `make_graphed_callables()` interface.
+Distributed training introduces communication overhead from various parallelism strategies. Megatron Core supports overlapping communication with computation to hide latency and improve throughput.
 
-To use `--cuda-graph-impl=transformer_engine`, the user should call related methods `TECudaGraphHelper.create_cudagraphs()` and `TECudaGraphHelper.cuda_graph_set_manual_hooks()` in the training script. Please refer to the usage in `megatron/training/training.py`.
+#### Data Parallel (DP) Communication Overlap
 
-For MoE models, certain configurations may prevent CUDA Graph capture of MoE layers. Specifically, when `--moe-expert-capacity-factor` and `--moe-pad-expert-input-to-capacity` are not set, the resulting dynamic shapes make MoE layers uncapturable. In such cases, you can still leverage CUDA Graphs for the attention layers (operations in `TransformerLayer._forward_attention()`) by setting `--cuda-graph-scope=attn`, while leaving the MoE layers (operations in `TransformerLayer._forward_mlp()`) unmodified. See the argument description for more usage of `--cuda-graph-scope`.
+With distributed optimizer, DP introduces **reduce-scatter** (gradients) and **all-gather** (parameters) communications, chunked by Transformer layer granularity.
+
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **Gradient Reduce Overlap** | Overlaps gradient reduce-scatter with backward computation | `--overlap-grad-reduce` |
+| **Param Gather Overlap** | Overlaps parameter all-gather with forward computation | `--overlap-param-gather` |
+| **BF16 Gradient Reduce** | Reduces gradients in BF16 instead of FP32 for better performance | `--grad-reduce-in-fp32 false` (via mixed precision config) |
+| **FP8 Param Gather** | Conducts parameter all-gather in FP8, reducing overhead by 50% | `--fp8-param-gather` |
+
+#### Tensor Parallel (TP) Communication Overlap
+
+TP with sequence parallelism introduces activation all-gather and reduce-scatter operations. Communications are overlapped in **bulk** (no dependency) or **pipelined** (with dependency) fashion.
+
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **TP Comm Overlap** | Enables bulk and pipelined TP communication overlap | `--tp-comm-overlap` |
+
+> **Requirements**: `tensor_model_parallel_size >= 2` and `--sequence-parallel`
+
+#### Pipeline Parallel (PP) Communication Overlap
+
+PP introduces P2P activation sends/receives between pipeline stages. Overlap is automatic in the 1F1B pipelining phase when VPP is enabled.
+
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **P2P Comm Overlap** | Overlaps PP P2P communications with non-dependent computations | `--overlap-p2p-comm` (auto-enabled with VPP) |
+| **VPP for Better Overlap** | Increases overlap opportunities by reducing layers per virtual stage | `--num-layers-per-virtual-pipeline-stage` |
+
+#### Expert Parallel (EP) Communication Overlap
+
+EP All-to-All can consume 30-40% of training time without optimization. These features hide or reduce EP communication overhead.
+
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **EP A2A Overlap** | Overlaps All-to-All with computation by merging FWD-BWD passes of adjacent microbatches | `--overlap-moe-expert-parallel-comm --delay-wgrad-compute` |
+| **Shared Expert Overlap** | Runs shared expert computation concurrently with EP token transfer | `--moe-shared-expert-overlap` |
+
+> **Requirements for EP A2A Overlap**: `expert_model_parallel_size > 1`, CUDA_DEVICE_MAX_CONNECTIONS > 1.
+
+### Compute Optimization
 
+Fine-grained MoE produces many small operations that can underutilize GPU resources. These optimizations reduce kernel launch overhead and improve GPU utilization.
 
-### Batch-Level EP-A2A hidding
-Enable A2A overlap across different batches inspired by the DSv3 DualPipe implmentation. \
-**Features** 
-- Hide ep a2a communication by batch-level overlapping
-- Split weight gradient and activation gradient computations for better overlap with communications
-- Support interleaved pipelined parallelism
-- Support FP8 training
-- Support MTP (`-mtp-num-layers 1` only, multiple MTP layers are not supported yet.)
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **Grouped GEMM** | Batches multiple expert GEMM operations into a single kernel call, improving GPU utilization | `--moe-grouped-gemm` |
+| **Router Fusion** | Fuses router projection, top-k selection, softmax, and auxiliary loss into fewer kernels | `--moe-router-fusion` |
+| **Permute Fusion** | Fuses token permutation/unpermutation operations into optimized single kernels | `--moe-permute-fusion` |
+| **FP8 Training** | Uses FP8 Tensor Core operations for faster GEMMs on Hopper/Blackwell GPUs | `--fp8 --fp8-recipe blockwise` |
 
 
-**Usage** 
+### FP8 Training
+
+FP8 training provides benefits across all three performance walls:
+
+| Wall | FP8 Benefit | Impact |
+|------|-------------|--------|
+| **Memory** | 50% activation reduction | Stores linear layer inputs in FP8 instead of BF16 |
+| **Memory** | Eliminate BF16 weight copies | Native FP8 casts directly from FP32 to FP8 |
+| **Communication** | 50% EP dispatch volume | Dispatches tokens in FP8 instead of BF16 |
+| **Communication** | 50% parameter all-gather | With FP8 primary weights (except MXFP8) |
+| **Compute** | Faster Tensor Core GEMMs | FP8 ops on Hopper/Blackwell are faster than BF16 |
+
+#### FP8 Recipes
+
+| Recipe | Scaling Granularity | Format | Platform | Use Case |
+|--------|---------------------|--------|----------|----------|
+| **Per-tensor** | Whole tensor | E4M3/E5M2 hybrid | Hopper, Blackwell | Conservative, initial experimentation |
+| **Blockwise** | 1×128 (activations), 128×128 (weights) | E4M3 | Hopper | **Production-proven** (DeepSeek-V3, Minimax-M2) |
+| **MXFP8** | 1×32 | E4M3 + E8M0 scaling | Blackwell | Native hardware support on GB200 |
+
+> **Recommendation**: Use **blockwise FP8** on Hopper for production training. It has been validated at scale on DeepSeek-V3 class models.
+
+#### MoE-Specific FP8 Optimizations
+
+| Optimization | Description | Config |
+|--------------|-------------|--------|
+| **Routing Map Padding** | Pads routing map (not tokens) to align M dimension to 16/32, avoiding per-tensor padding overhead | `--moe-router-padding-for-fp8` |
+| **FP8 Primary Weights** | Casts FP32 master weights directly to FP8, eliminating BF16 intermediate copy | `--fp8-param-gather` (Need additional `--reuse-grad-buf-for-mxfp8-param-ag` for MXFP8) |
+
+
+#### Example Configuration
+
 ```bash
-# Add the following flags to your training scripts
---overlap-moe-expert-parallel-comm
-# [optional] only works with specific TE version
---delay-wgrad-compute
+# Blockwise FP8 on Hopper (recommended for production)
+--fp8-format e4m3
+--fp8-recipe blockwise
+--fp8-param-gather
+--moe-router-padding-for-fp8
+
+# MXFP8 on Blackwell
+--fp8-format e4m3
+--fp8-recipe mxfp8
+--moe-router-padding-for-fp8
+--fp8-param-gather 
+--reuse-grad-buf-for-mxfp8-param-ag
 ```
 
-### Fine-grained Activation Offloading (collaborated with rednote)
-Offload the input activation at the granularity of modules
+> **Note**: For blockwise and MXFP8 recipes with current scaling, training loss curves show negligible difference compared to BF16 baselines.
 
-**Usage**
-```bash
-# Enable fine-grained activation offloading
---fine-grained-activation-offloading
 
-# Specify which modules are going to offload its input
-# Choices: "attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act".
---offload-modules expert_fc1
-```
-For more details, please refer to the ```docs/user-guide/features/fine_grained_activation_offloading.md```
-
-### MoE Related Arguments
-| Item | Description |
-| --- | --- |
-| --num-experts | Number of Experts in MoE (None means no MoE) |
-| --expert-model-parallel-size | Degree of expert model parallelism. Default is 1. |
-| --moe-ffn-hidden-size | MoE Feed-Forward Network hidden size. Default is None. |
-
-<details>
-<summary> View all MoE related arguments. </summary>
-
-| Item | Description |
-| --- | --- |
-| --num-experts | Number of Experts in MoE (None means no MoE) |
-| --expert-model-parallel-size | Degree of expert model parallelism. Default is 1. |
-| --moe-ffn-hidden-size | MoE Feed-Forward Network hidden size. Default is None. |
-| --expert-tensor-parallel-size | Degree of tensor model parallelism of expert layer. Default is same to --tensor-model-parallel-size. |
-| --moe-layer-freq | Frequency between MoE layers and Dense layers. Accepts either: 1) An integer N for 1:N ratio (one expert layer for every N-1 dense layers), 2) A string "N" for the same ratio, or 3) A string with Python list expression for custom patterns like `([1]*3+[0]*1)*3` which gives [1,1,1,0,1,1,1,0,1,1,1,0] where 1=expert layer and 0=dense layer. Examples: `([0]+[1]*23)` for 1 dense layer followed by 23 experts layers, `([1]*3+[0]*2)*2` for three expert layers followed by two dense layers, repeated twice. Default is 1. |
-| --moe-grouped-gemm | When there are multiple experts per rank, launch multiple local GEMM kernels in multiple streams to improve the utilization and performance with GroupedLinear in TransformerEngine. |
-| --moe-router-load-balancing-type | Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2 and DeepSeekV3, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss". |
-| --moe-router-dtype | Data type for routing computation and expert output weighted averaging. Options are 'fp32' and 'fp64'. This can improve numerical stability, particularly when using a large number of experts. The throughput/memory impact should be negligible when used with --moe-permute-fusion. Default is None (no dtype promotion). |
-| --moe-router-topk | Number of experts to route to for each token. The default is 2. |  
-| --moe-router-score-function | Score function for MoE routing. Can be "softmax" or "sigmoid". Default is "softmax". |
-| --moe-router-pre-softmax | Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k. |
-| --moe-router-num-groups | Number of groups to divide experts into for group-limited routing. When using group-limited routing: 1) Experts are divided into equal-sized groups, 2) For each token, a subset of groups are selected based on routing scores (sum of top-2 expert scores within each group), 3) From these selected groups, moe_router_topk experts are chosen.  Two common use cases: 1) Device-limited routing: Set equal to expert parallel size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) 2) Node-limited routing: Set equal to number of nodes in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437)) |
-| --moe-router-group-topk | Number of selected groups for group-limited routing. |
-| --moe-router-topk-scaling-factor | Scaling factor for routing score in top-k selection, only works when --moe-router-pre-softmax enabled. Defaults to None, which means no scaling. |
-| --moe-router-enable-expert-bias | TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. The routing decision is based on the sum of the routing scores and the expert bias. See https://arxiv.org/abs/2408.15664 for details. |
-| --moe-router-fusion | Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in TransformerEngine 2.7.0 and above. |
-| --moe-router-bias-update-rate | The expert bias is updated based on the number of assigned tokens to each expert in a global batch, where the bias is increased for experts with less assigned tokens and decreased for experts with more assigned tokens. Default is 1e-3 same as that used in DeepSeekV3. |
-| --moe-router-force-load-balancing | (Experimental) Force override routing to balance token distribution using random logits for MoE routers, supporting naive top-k and group-limited top-k. This experimental feature is for benchmarking purposes only! |
-| --moe-router-padding-for-quantization | Pad the routing_map to make sure the number of tokens each expert received is a multiple of 16/32 for FP8/FP4 precision. It is suggested to enable this for dropless training with FP8 precision when num_local_experts > 1. This is a more efficient way to pad for FP8 which eliminates the explicit padding in the GroupedMLP layer. |
-| --moe-aux-loss-coeff | Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended. Default is 0.0. |
-| --moe-z-loss-coeff | Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended. Default is None. |
-| --moe-input-jitter-eps | Add noise to the input tensor by applying jitter with a specified epsilon value. Default is None. |
-| --moe-token-dispatcher-type | Determines the token dispatcher type. Choices are "allgather", "alltoall". Default is "allgather". We recommend using 'alltoall' if expert parallelism is applied. We have upgraded the "alltoall" dispatcher in place during MCore v0.9, while the original implementation renamed as "alltoall_seq" is retained until MCore v0.13.|
-| --moe-flex-dispatcher-backend | (Experimental) Select the backend for the flex token dispatcher. Supported options: "deepep", "hybridep". Enables efficient token dispatching and combining for MoE models. |
-| --moe-per-layer-logging | Enable per-layer logging for MoE, currently supports auxiliary loss and z loss. |
-| --moe-expert-capacity-factor | The capacity factor for each expert, None means no token will be dropped. Default is None. |
-| --moe-pad-expert-input-to-capacity | Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set. |
-| --moe-token-drop-policy | The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with the lowest probabilities will be dropped. If "position", tokens at the end of each batch will be dropped. |
-| --moe-layer-recompute | Enable activation checkpointing for moe_layer, should be used when memory is not sufficient. |
-| --moe-permute-fusion | Fuse token rearrangement ops during token dispatching. |
-| --moe-shared-expert-intermediate-size | Set shared expert total ffn hidden size. It should be equal to `num_shared_experts * ffn_size_of_each_shared_expert` if there are multiple shared experts. None means no shared expert. |
-| --moe-shared-expert-overlap | (Experimental, may change) If this is set, the communications/computations in the shared experts and the dispatcher will overlap (The `alltoall` dispatcher is needed.) Otherwise, the shared expert runs after the routed experts. |
-| --moe-use-upcycling | Load the dense model checkpoint, convert it into an MoE model at runtime and start training. The converted model will be saved to the path specified by `--save` before training begins. Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.|
-| --overlap-moe-expert-parallel-comm | Enable batch-level overlapping in 1f1b stage. |
-| --delay-wgrad-compute | Enable split dgrad and wgrad for `overlap-moe-expert-parallel-comm` execution. Increasing room to hide communication latency by more finegrained control. |
-| --pipeline-model-parallel-layout | (Experimental, may change) A string containing a Python list expression that defines a custom pipeline model parallel layout. |
-| --moe-upcycling-granularity | This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.|
+### CUDA Graph
+CUDA Graph functionality can be enabled through the `--cuda-graph-impl` option. There are two implementations:
 
-</details>
+1. `--cuda-graph-impl=local`: Captures cuda graphs using the MCore-internal cuda graph manager.
+2. `--cuda-graph-impl=transformer_engine`: Captures cuda graphs using the TE `make_graphed_callables()` interface.
 
-## MoE training example:
-<details>
-<summary>Click here. </summary>
+To use `--cuda-graph-impl=transformer_engine`, the user should call related methods `TECudaGraphHelper.create_cudagraphs()` and `TECudaGraphHelper.cuda_graph_set_manual_hooks()` in the training script. Please refer to the usage in `megatron/training/training.py`.
+
+For MoE models, certain configurations may prevent CUDA Graph capture of MoE layers. Specifically, when `--moe-expert-capacity-factor` and `--moe-pad-expert-input-to-capacity` are not set, the resulting dynamic shapes make MoE layers uncapturable. In such cases, you can still leverage CUDA Graphs for the attention layers (operations in `TransformerLayer._forward_attention()`) by setting `--cuda-graph-scope=attn`, while leaving the MoE layers (operations in `TransformerLayer._forward_mlp()`) unmodified. See the argument description for more usage of `--cuda-graph-scope`.
 
+## MoE Arguments Reference
+### Core Arguments
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --num-experts | Number of Experts in MoE | None |
+| --expert-model-parallel-size | Degree of expert model parallelism | 1 |
+| --moe-ffn-hidden-size | MoE FFN hidden size | FFN hidden size of the dense model |
+| --expert-tensor-parallel-size | Expert layer tensor parallelism | Same as TP(Recommeded to set to 1 for fine-grained MoE models) |
+| --moe-layer-freq | MoE layer frequency pattern | 1 |
+
+### Router Arguments
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --moe-router-load-balancing-type | Load balancing: aux_loss, sinkhorn, seq_aux_loss, none | aux_loss |
+| --moe-router-topk | Number of experts per token | 2 |
+| --moe-router-score-function | Score function: softmax, sigmoid | softmax |
+| --moe-router-pre-softmax | Softmax before top-k | False |
+| --moe-router-num-groups | Groups for group-limited routing | None |
+| --moe-router-group-topk | Selected groups in group-limited routing | None |
+| --moe-router-enable-expert-bias | Dynamic per-expert bias | False |
+| --moe-router-bias-update-rate | Bias update rate | 1e-3 |
+| --moe-router-fusion | Enable router fusion | False |
+| --moe-router-dtype | Router precision: fp32, fp64 | None |
+| --moe-router-padding-for-fp8 | Pad for FP8 alignment | False |
+
+### Loss and Regularization
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --moe-aux-loss-coeff | Auxiliary loss coefficient | 0.0 |
+| --moe-z-loss-coeff | Z-loss coefficient | None |
+| --moe-input-jitter-eps | Input jitter epsilon | None |
+
+### Token Dispatching
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --moe-token-dispatcher-type | Dispatcher: allgather, alltoall, flex | allgather |
+| --moe-enable-deepep | Enable DeepEP (with flex) | False |
+| --moe-expert-capacity-factor | Capacity factor | None |
+| --moe-pad-expert-input-to-capacity | Pad to capacity | False |
+| --moe-token-drop-policy | Drop policy: probs, position | probs |
+| --moe-permute-fusion | Fuse permutation ops | False |
+
+### Performance Optimization
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --moe-grouped-gemm | Use GroupedGEMM | False |
+| --overlap-moe-expert-parallel-comm | Batch-level EP overlap | False |
+| --delay-wgrad-compute | Split dgrad/wgrad compute | False |
+| --moe-shared-expert-intermediate-size | Shared expert FFN size | None |
+| --moe-shared-expert-overlap | Overlap shared expert | False |
+
+### Memory and Checkpointing
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --moe-layer-recompute | Recompute MoE layer | False |
+| --moe-use-upcycling | Enable upcycling | False |
+| --moe-upcycling-granularity | Upcycling granularity | 1 |
+
+### Miscellaneous
+| Argument | Description | Default |
+|----------|-------------|---------|
+| --moe-per-layer-logging | Per-layer logging | False |
+| --moe-router-force-load-balancing | Force load balancing (experimental) | False |
+
+## Examples
 ```bash
 #!/bin/bash
 
 # Runs Mixtral 8x7B model on 32 H100/A100 GPUs
-# The Dropless MoE suffers from an imbalanced token distribution at the early stage of training (the first few hundred iterations), which may lead to poor performance and out-of-memory (OOM) issues.
-# To check the performance of a Dropless MoE model, we should run the model for at least 500 iterations or resume from trained checkpoints.
 
 export CUDA_DEVICE_MAX_CONNECTIONS=1
 
 GPUS_PER_NODE=8
-# Change for multinode config
 MASTER_ADDR=${MASTER_ADDR:-"localhost"}
 MASTER_PORT=${MASTER_PORT:-"6000"}
-NNODES=${NNODES:-"1"}
+NNODES=${NNODES:-"4"}
 NODE_RANK=${RANK:-"0"}
 WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
 
@@ -333,11 +623,12 @@ MODEL_ARGS=(
 MOE_ARGS=(
     --num-experts 8
     --expert-model-parallel-size 8
-    --moe-router-load-balancing-type aux_loss # options: aux_loss, sinkhorn, None. Default is aux_loss.
+    --moe-router-load-balancing-type aux_loss
     --moe-router-topk 2
     --moe-aux-loss-coeff 1e-2
     --moe-grouped-gemm
     --moe-permute-fusion
+    --moe-token-dispatcher-type alltoall
 )
 
 DATA_ARGS=(
@@ -372,24 +663,17 @@ MODEL_PARALLEL_ARGS=(
 )
 
 LOGGING_ARGS=(
-    --log-interval 1 \
-    --save-interval 10000 \
-    --eval-interval 1000 \
-    --eval-iters 10 \
-    --save $CHECKPOINT_PATH \
-    --load $CHECKPOINT_PATH \
-    --tensorboard-dir "${CHECKPOINT_PATH}/tensorboard" \
-    --no-load-optim \
-    --no-load-rng
+    --log-interval 1
+    --save-interval 10000
+    --eval-interval 1000
+    --eval-iters 10
+    --save $CHECKPOINT_PATH
+    --load $CHECKPOINT_PATH
+    --tensorboard-dir "${CHECKPOINT_PATH}/tensorboard"
+    --ckpt-format torch_dist
+    --auto-detect-ckpt-format
 )
 
-if [ -n "${WANDB_API_KEY}" ]; then
-    LOGGING_ARGS+=(
-        --wandb-project ${WANDB_PROJECT:-"Mixtral-Finetuning"}
-        --wandb-exp-name ${WANDB_NAME:-"Mixtral_8x7B"} 
-    )
-fi
-
 torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
     ${MODEL_ARGS[@]} \
     ${MOE_ARGS[@]} \
@@ -398,107 +682,36 @@ torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
     ${MODEL_PARALLEL_ARGS[@]} \
     ${LOGGING_ARGS[@]}
 ```
+
 </details>
 
-# Performance Best Practice
+## Contributing
 
-## Tuning Guide of Parallel Mappings
+We welcome contributions! Please see [CONTRIBUTING.md](../../../../CONTRIBUTING.md) for guidelines.
 
-To find a good parallel mapping that help you achieve a high throughput of a new model, there are some general rule that could help. Here is an overview of properties in different aspects for each parallel strategy.
+## Support
 
-| Parallel Strategy | Peak Activation Memory          | Weight Memory  | Optimizer states                  | Communication (Per-Layer) |
-|:-----------------:|:-------------------------------:|:--------------:|:---------------------------------:|:-------------------------:|
-| TP                | 1/N (with SP on)                | 1/N            | 1/N                               |        High               |
-| EP                | 1                               | 1/N in MoELayer| 1/N                               |       Medium              |
-| PP                | 1 (>1 with virtual pipeline)    | 1/N            | 1/N                               |       Medium              |
-| CP                | 1/N                             | 1              | 1/N (with distributed optimizer)  |       Medium              |
-| DP                | 1                               | 1              | 1/N (with distributed optimizer)  |        Low                |
+- GitHub Issues: [Report bugs or request features](https://github.com/NVIDIA/Megatron-LM/issues)
+- Documentation: [Full documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)
+
+
+## Citation
 
-For a specific model, the best parallel mapping varies based on the model architecture, trained sequence length and the hardware platform.
-Here we provide some general rules to get better performance:
-1. Keep the model parallelism size as small as possible. 
-    - For the large language models, model parallelism is often required to prevent OOM, but it will bring communication overhead and hurt performance. 
-    - With distributed optimizer, master weights and optimizer states will be sharded across all DP ranks with slight communication overhead.
-    So try to reduce the model parallelism size and increase data parallelism size when there are lots of free GPU memory during training.
-2. Ensure the EPxTP communication within the NVLink domain.
-    - Communications of EP and TP should remain within the NVLink domain as much as possible, as both are communication-intensive.
-    - If the model is too large and requires scaling across multiple nodes, consider PP before TP and EP. See item 3 for details.
-3. Use Pipeline Parallelism to scale the model further.
-    - Enable Virtual Pipeline Parallelism(VPP) to reduce pp bubbles when PP_size >= 2 by setting `num_layers_per_virtual_pipeline_stage`.
-    - VPP_size tuning: the legal values of vpp_size are all common divisors of num_layers/pp_size, E.g., num_layers=24, pp_size=4, then we can pick vpp_size from {1, 2, 3, 6}. The larger the vpp_size, the lower the pipeline bubbles, while the larger number of P2P communications between each PP stages. Empirically a value in the middle often gives the best trade-off. `VPP_size=num_layers / PP_size / num_layers_per_virtual_pipeline_stage`
-4. Prefer EP over TP for the expert layer when possible:
-    - TP saves more memory than EP, but EP can achieve better GEMM efficiency and less communication overhead than TP.
-    - If EP size increased to the number of expert, the local token permutation/un-permutation for experts computation are omitted.
-    - Simplify the computation graph of MoE layers, more convenient for performing potential comm-computation overlapping.
-    - In practice, EP8TP1 is better than EP4TP2 for 8x7B.
-5. Enable Context Parallelism for long context training.
-    - The efficiency of CP largely depends on whether its communication can be overlapped with computation. 
-    - Empirically, use CP when sequence length >= 8K.
-
-## MoE Parallel Folding
-
-MoE Parallel Folding separates the MoE related parallel groups from Dense groups.
-1. Traditional MoE parallel groups are entangled with dense by using a 5-dimension parallel group generator with default order `tp-cp-ep-dp-pp`. The EP group in MoE is a sub-group of DP in Attention.
-2. With MoE Parallel Folding, we use a parallel group generator with `tp-cp-dp-pp` for Attention, and another with `tp-ep-dp-pp` for MoE. The EPxTP group in MoE is a sub-group of DPxCPxTP in Attention.
-
-By setting `--expert-tensor-parallel-size`, we can set MoE-specific TP size.
-
-### Advantages of MoE Parallel Folding
-1. The CP and EP group are folded together by default, such that:
-    1. It reduces the minimal required GPUs to turn on both CP and EP. For example, the traditional way with (CP=8, EP=8) needs at least 64 GPUs, for now it only requires 8 GPUs.
-    2. The CP and EP communication can be both put in the NVLink domain.
-2. We can set different TP sizes for Attention and MoE part.
-    1. For MoE, EP is often more efficient than TP. But in the traditional way, only using EP can get OOM for most models.
-    2. With MoE parallel folding, we can turn on TP for Attention part and setting TP=1 for MoE models, which often gets better MFU.
-
-## End-to-End Training Practice
-**Use the latest NVIDIA PyTorch or NeMo Docker Image**
-- [NGC PyTorch Image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)
-- [NGC NeMo Image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)
-
-**Token Dispatcher Choices**
-- Token Dispatcher sends tokens to the designated expert, involves tensor rearangement and communications.
-- Dispatcher `allgather` is the default option. It achieves better performance and efficiency when only tensor parallelism is used or when the Top-k value is very large.
-- Dispatcher `alltoall` is recommended if expert parallelism is applied.
-- Dispatcher `flex` is a new dispatcher decouples communication group from model parallelism. It supports two backends(DeepEP and HybridEP) selectable via `--moe-flex-dispatcher-backend`.
-
-**Enable Communication Overlap**
-- Enable `--overlap-param-gather` and `--overlap-grad-reduce` with distributed optimizer.
-- Enable `--tp-comm-overlap` when TP>1.
-- Enable p2p comm overlap when PP > 1 by setting `num_layers_per_virtual_pipeline_stage`.
-
-**Enable GroupedGEMM when num_local_experts>1 with `--moe-grouped-gemm`**
-- GroupedGEMM has higher efficiency than vanilla sequential GEMMs for each expert.
-- Recommend to use the TE version of Grouped GEMM (by upgrading to MCore v0.8 and TE v1.9), which support Gradient Accumulation Fusion and FP8 Training.
-
-**OOM Caused by Token Distribution Imbalance when Training From Scratch**  
-MoE suffers from a severe load imbalance issue when the router is under-trained, leading to the model easily running out of memory (OOM), which typically occurs in the first 100~300 steps when training from scratch. 
-Therefore, there are two recommended ways during the first 200 steps to avoid the OOM problem, which can be removed after the token distribution is more stable:
-1. Increase the `expert-tensor-parallel-size` and decrease `expert-model-parallel-size` to replace EP with TP in MoELayer, this can prevent the load imbalancing between EP ranks. Since current ETP implementation has some memory overhead, you can further enable activation recomputation only for MoE Layer by adding `--moe-layer-recompute`.
-2. Setting capacity factor to a relatively small number like 1.0 by adding `--moe-token-capacity-factor 1.0`.
-
-**Leverage DeepSeek's DeepEP for High-Performance Cross-Node Token Dispatching**
-- The primary advantage of DeepEP is its cross-node token communication efficiency, which delivers substantial performance improvements when deploying expert parallelism across multiple nodes with large TopK values.
-- To enable DeepEP in your training configuration, simply set `--moe-token-dispatcher-type=flex` and `--moe-enable-deepep` in your command line arguments.
-
-**FP8 Training Best Practice**
-- Using latest version of [TransformerEngine](https://github.com/NVIDIA/TransformerEngine).
-- Enable router padding with `--moe-router-padding-for-quantization` to reduce padding overhead.
-- Enable native FP8 weights with `--fp8-param-gather` to reduce weights memory cost.
-
-## Reference Best Parallel Mapping
-
-Here are the reference parallel mappings of MCore v0.8 for Mixtral 8x7B and 8x22B models:
-|        Model            | Vocab Size| Dispatcher | Precision | #GPUs | SEQ LEN | TP | EP | PP | VP | MBS | GBS |
-|:-----------------------:|:---------:|:----------:|:---------:|:-----:|:-------:|:--:|:--:|:--:|:--:|:---:|:---:|
-| Mixtral 8x7B(Dropless)  |   32K     | All-to-All | BF16      | 64    | 4096    | 1  | 8  | 4  | 8  | 1   | 256 |
-| Mixtral 8x22B(Dropless) |   32K     | All-to-All | BF16      | 128   | 4096    | 4  | 2  | 8  | 7  | 1   | 256 |
-
-Detailed Benchmark Information:  
-Server:
-- 8xH100 80GB HBM3 
-- NVLink 4th Generation
-- InfiniBand 8x400 Gbit/s
-
-Docker Image:
-- PyTorch 24.09 with TransformerEngine v1.11
+If you use Megatron-Core MoE in your research, please cite:
+
+```bibtex
+
+@article{megatron-lm,
+  title={Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},
+  author={Shoeybi, Mohammad and Patwary, Mostofa and Puri, Raul and LeGresley, Patrick and Casper, Jared and Catanzaro, Bryan},
+  journal={arXiv preprint arXiv:1909.08053},
+  year={2019}
+}
+
+@article{moe-parallel-folding,
+    title={MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core}, 
+    author={Liu, Dennis and Yan, Zijie and Yao, Xin and Liu, Tong and Korthikanti, Vijay and Wu, Evan and Fan, Shiqing and Deng, Gao and Bai, Hongxiao and Chang, Jianbin and Aithal, Ashwath and Andersch, Michael and Shoeybi, Mohammad and Yao, Jiajie and Zhou, Chandler and Wu, David and Li, Xipeng and Yang, June},
+    year={2025},
+    journal={arXiv preprint arXiv:2504.14960},
+}
+```

From 2b02a28f542f671e737a17549e2efe9387be2697 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Tue, 27 Jan 2026 17:46:09 +0100
Subject: [PATCH 36/79] build: Bump to TE2.12 (#3086)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 pyproject.toml                                |  2 +-
 .../golden_values_dev_dgx_h100.json           |  8 +++---
 .../moe/test_a2a_token_dispatcher.py          |  1 +
 .../transformer/moe/test_token_dispatcher.py  |  2 ++
 uv.lock                                       | 25 ++++++++++++++-----
 5 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f5b563bb51d..7f5ceac6203 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -179,7 +179,7 @@ override-dependencies = [
 flash_mla = [
     { git = "https://github.com/deepseek-ai/FlashMLA", rev = "9edee0c022cd0938148a18e334203b0aab43aa19" },
 ]
-transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "release_v2.11" }
+transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "d9b7fc5770a88af06e2e9c2bd97b550614c3a69f" }
 nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "01a9a8ba360f7b2908728ad0516e0ad9d936966d" }
 emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.1.0" }
 
diff --git a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
index c5fdf3beffa..339197f8f0f 100644
--- a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
@@ -53,10 +53,10 @@
         "step_interval": 1,
         "values": {
             "1": "nan",
-            "2": 135.42645,
-            "3": 78.78998,
-            "4": 79.18825,
-            "5": 80.10109
+            "2": 135.80412,
+            "3": 81.98981,
+            "4": 82.00576,
+            "5": 82.33207
         }
     }
 }
\ No newline at end of file
diff --git a/tests/unit_tests/transformer/moe/test_a2a_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_a2a_token_dispatcher.py
index e20e35b84e1..aad1fcaca5f 100644
--- a/tests/unit_tests/transformer/moe/test_a2a_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_a2a_token_dispatcher.py
@@ -95,6 +95,7 @@ def test_capacity_padding_forward_backward(self, tp_size, ep_size, permute_fusio
         )
         container.dispatcher_drop_and_pad_test()
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(
         not is_te_min_version("1.7.0"), reason="TE 1.7.0 is required for MoE with FP8."
     )
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index fd6fb7f6d09..24617952b94 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -364,6 +364,7 @@ def setup_method(self, method):
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @pytest.mark.internal
     @pytest.mark.parametrize("tp_size,ep_size", [(8, 1), (1, 8), (2, 4), (1, 1)])
@@ -382,6 +383,7 @@ def test_forward_backward(self, tp_size, ep_size, permute_fusion):
 
         container.dispatcher_dropless_test()
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @pytest.mark.internal
     @pytest.mark.parametrize("permute_fusion", permute_fusion_params)
diff --git a/uv.lock b/uv.lock
index b95e1cef2cf..705340af107 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2376,7 +2376,7 @@ requires-dist = [
     { name = "torch" },
     { name = "tqdm", marker = "extra == 'dev'" },
     { name = "tqdm", marker = "extra == 'lts'" },
-    { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=release_v2.11" },
+    { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=d9b7fc5770a88af06e2e9c2bd97b550614c3a69f" },
     { name = "transformers", marker = "extra == 'mlm'" },
     { name = "wandb", marker = "extra == 'mlm'" },
     { name = "wget", marker = "extra == 'dev'" },
@@ -3096,6 +3096,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e5/15/97e6e4ddfe5fc35bcee74a45b7c33fb73abb83713c7dfa26420b971a86c3/nv_one_logger_training_telemetry-2.3.1-py3-none-any.whl", hash = "sha256:5319443829b59378a498c3c62ac98973e14f31be675c229ff2b14e2fe109aa0b", size = 44140, upload-time = "2025-10-29T21:21:40.72Z" },
 ]
 
+[[package]]
+name = "nvdlfw-inspect"
+version = "0.2.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pyyaml" },
+    { name = "torch", marker = "sys_platform == 'never'" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8a/86/94188e03e5d4dd7b73c390b0cddcde5618b3799c18e327b2bf15763f6137/nvdlfw_inspect-0.2.2-py3-none-any.whl", hash = "sha256:8a4dc2814c5a4cd19ae304170b9bfa514538ef3c3eb243a45a82404ec3cb279d", size = 30964, upload-time = "2025-12-03T10:52:01.933Z" },
+]
+
 [[package]]
 name = "nvidia-cudnn-frontend"
 version = "1.17.0"
@@ -4137,12 +4149,12 @@ name = "pytest"
 version = "8.3.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" },
-    { name = "exceptiongroup", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
     { name = "iniconfig" },
     { name = "packaging" },
     { name = "pluggy" },
-    { name = "tomli", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" },
+    { name = "tomli", marker = "python_full_version < '3.11'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/ae/3c/c9d525a414d506893f0cd8a8d0de7706446213181570cdbd766691164e40/pytest-8.3.5.tar.gz", hash = "sha256:f4efe70cc14e511565ac476b57c279e12a855b11f48f212af1080ef2263d3845", size = 1450891, upload-time = "2025-03-02T12:54:54.503Z" }
 wheels = [
@@ -5728,11 +5740,12 @@ wheels = [
 
 [[package]]
 name = "transformer-engine"
-version = "2.11.0+c188b533"
-source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=release_v2.11#c188b533cc3721ca9c6bbfd26148f5cf60108c25" }
+version = "2.12.0+d9b7fc57"
+source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=d9b7fc5770a88af06e2e9c2bd97b550614c3a69f#d9b7fc5770a88af06e2e9c2bd97b550614c3a69f" }
 dependencies = [
     { name = "einops" },
     { name = "importlib-metadata" },
+    { name = "nvdlfw-inspect" },
     { name = "onnx" },
     { name = "onnxscript" },
     { name = "packaging" },

From 6cf285b23889ab4b6b225d2896b3ed1744685119 Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Tue, 27 Jan 2026 08:53:38 -0800
Subject: [PATCH 37/79] Logging cleanup (only log on rank 0 if possible)
 (#3036)

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 megatron/core/_rank_utils.py                  |  43 +++++
 megatron/core/datasets/megatron_tokenizer.py  |  12 +-
 .../core/dist_checkpointing/exchange_utils.py |  22 +--
 megatron/core/dist_checkpointing/optimizer.py |   8 +-
 .../core/dist_checkpointing/serialization.py  |  10 +-
 .../strategies/async_utils.py                 |  26 ++-
 .../megatron_fsdp/param_and_grad_buffer.py    | 115 ++++++++-----
 .../core/distributed/param_and_grad_buffer.py |   6 +
 megatron/core/hyper_comm_grid.py              |   5 +-
 megatron/core/nccl_allocator.py               |  55 +++---
 megatron/core/optimizer/distrib_optimizer.py  |  33 +++-
 .../modelopt/gpt/state_dict_hooks.py          |   6 +-
 megatron/core/rerun_state_machine.py          | 157 +++++++++---------
 .../copy_services/gloo_copy_service.py        |  17 +-
 .../copy_services/nccl_copy_service.py        |  19 ++-
 .../text/libraries/huggingface_tokenizer.py   |  14 +-
 .../text/libraries/tiktoken_tokenizer.py      |  19 ++-
 megatron/core/transformer/cuda_graphs.py      |  23 +--
 megatron/core/utils.py                        |  28 +---
 megatron/training/arguments.py                | 102 +++++-------
 megatron/training/checkpointing.py            |  31 ++--
 megatron/training/global_vars.py              |   6 +-
 megatron/training/initialize.py               |  48 +++---
 megatron/training/tokenizer/tokenizer.py      |   4 +-
 megatron/training/training.py                 |   4 +-
 megatron/training/utils.py                    |  21 ++-
 pretrain_gpt.py                               |  11 +-
 pretrain_mamba.py                             |  11 +-
 28 files changed, 498 insertions(+), 358 deletions(-)
 create mode 100644 megatron/core/_rank_utils.py

diff --git a/megatron/core/_rank_utils.py b/megatron/core/_rank_utils.py
new file mode 100644
index 00000000000..6b1a35ca798
--- /dev/null
+++ b/megatron/core/_rank_utils.py
@@ -0,0 +1,43 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Low-level rank utilities with minimal dependencies to avoid circular imports."""
+
+import logging
+import os
+from typing import Any
+
+import torch
+
+
+def safe_get_rank() -> int:
+    """Safely get the rank of the current process.
+
+    Returns the rank from torch.distributed if initialized, otherwise falls back
+    to the RANK environment variable, defaulting to 0.
+
+    Returns:
+        int: The rank of the current process.
+    """
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+
+    # If torch.distributed is not initialized, try to read environment variables.
+    try:
+        return int(os.environ.get("RANK", 0))
+    except (ValueError, TypeError):
+        return 0
+
+
+def log_single_rank(logger: logging.Logger, *args: Any, rank: int = 0, **kwargs: Any) -> None:
+    """Log a message only on a single rank.
+
+    If torch distributed is initialized, write log on only one rank.
+
+    Args:
+        logger: The logger to write the logs.
+        *args: All logging.Logger.log positional arguments.
+        rank: The rank to write on. Defaults to 0.
+        **kwargs: All logging.Logger.log keyword arguments.
+    """
+    if safe_get_rank() == rank:
+        logger.log(*args, **kwargs)
diff --git a/megatron/core/datasets/megatron_tokenizer.py b/megatron/core/datasets/megatron_tokenizer.py
index 08b602c4766..a8c2a431e59 100644
--- a/megatron/core/datasets/megatron_tokenizer.py
+++ b/megatron/core/datasets/megatron_tokenizer.py
@@ -23,11 +23,15 @@ class MegatronLegacyTokenizer(ABC):
     """
 
     def __init__(self, *tokenizer_paths: str, **tokenizer_options: Any):
+        from megatron.core.utils import log_single_rank
+
         # Deprecation warning
-        logger.warning(
-            "You’re using the legacy tokenizer system, which is deprecated "
-            "and will be removed in a future release. Please migrate to the new tokenizer system "
-            "(`megatron.core.tokenizers.MegatronTokenizer`)."
+        log_single_rank(
+            logger,
+            logging.WARNING,
+            "You're using the legacy tokenizer system, which is deprecated "
+            "and will be removed in a future release. Please migrate to the new "
+            "tokenizer system (`megatron.core.tokenizers.MegatronTokenizer`).",
         )
         self.unique_identifiers = OrderedDict()
         self.unique_identifiers["class"] = type(self).__name__
diff --git a/megatron/core/dist_checkpointing/exchange_utils.py b/megatron/core/dist_checkpointing/exchange_utils.py
index c91fc091f6f..79f906b237a 100644
--- a/megatron/core/dist_checkpointing/exchange_utils.py
+++ b/megatron/core/dist_checkpointing/exchange_utils.py
@@ -11,7 +11,7 @@
 import numpy as np
 import torch
 
-from ..utils import get_pg_rank, get_pg_size
+from ..utils import get_pg_rank, get_pg_size, log_single_rank
 from .core import CheckpointingException
 from .dict_utils import nested_values
 from .mapping import ShardedStateDict, ShardedTensor, is_main_replica
@@ -408,11 +408,11 @@ def exchange_loaded_tensors_gather_object(
     # Error checks
     if len(all_loaded_tensors) != sum(map(len, all_loaded_tensors_list)):
         err_msg = "Duplicate shard ids loaded by different ranks"
-        if torch.distributed.get_rank() == 0:
-            logger.error(
-                f"{err_msg}. Shards ids by rank:"
-                f" {[lt.keys() for lt in all_loaded_tensors_list]}"
-            )
+        log_single_rank(
+            logger,
+            logging.ERROR,
+            f"{err_msg}. Shards ids by rank:" f" {[lt.keys() for lt in all_loaded_tensors_list]}",
+        )
         raise CheckpointingException(err_msg)
 
     return all_loaded_tensors
@@ -439,11 +439,11 @@ def exchange_loaded_objects_gather_object(
     # Error checks
     if len(all_loaded_objects) != sum(map(len, all_loaded_objects_list)):
         err_msg = "Duplicate shard ids loaded by different ranks"
-        if torch.distributed.get_rank() == 0:
-            logger.error(
-                f"{err_msg}. Shards ids by rank:"
-                f" {[lt.keys() for lt in all_loaded_objects_list]}"
-            )
+        log_single_rank(
+            logger,
+            logging.ERROR,
+            f"{err_msg}. Shards ids by rank:" f" {[lt.keys() for lt in all_loaded_objects_list]}",
+        )
         raise CheckpointingException(err_msg)
 
     return all_loaded_objects
diff --git a/megatron/core/dist_checkpointing/optimizer.py b/megatron/core/dist_checkpointing/optimizer.py
index fb08ea52b1b..69227f1ab66 100644
--- a/megatron/core/dist_checkpointing/optimizer.py
+++ b/megatron/core/dist_checkpointing/optimizer.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from megatron.core.utils import to_local_if_dtensor
+from megatron.core.utils import log_single_rank, to_local_if_dtensor
 
 from .dict_utils import nested_values
 from .mapping import (
@@ -70,10 +70,12 @@ def get_param_id_to_sharded_param_map(
             logger.debug(f'{ten} is not tracked by the optimizer')
 
     if not id_to_sharded_param_map:
-        logger.warning(
+        log_single_rank(
+            logger,
+            logging.WARNING,
             "Sharded parameters mapping is empty. It means tensors in model state dict"
             " do not correspond to tensors in optimizer parameters map."
-            " Make sure to call state_dict with `keep_vars=True`."
+            " Make sure to call state_dict with `keep_vars=True`.",
         )
     return id_to_sharded_param_map
 
diff --git a/megatron/core/dist_checkpointing/serialization.py b/megatron/core/dist_checkpointing/serialization.py
index 0469949c67d..94c7a6cf663 100644
--- a/megatron/core/dist_checkpointing/serialization.py
+++ b/megatron/core/dist_checkpointing/serialization.py
@@ -15,6 +15,7 @@
 import torch
 
 from megatron.core.msc_utils import MultiStorageClientFeature
+from megatron.core.utils import log_single_rank
 
 from . import ShardedTensor
 from .core import CheckpointingConfig, save_config
@@ -181,9 +182,12 @@ def load_common_state_dict(checkpoint_dir: Union[str, Path]) -> StateDict:
     """
     if isinstance(checkpoint_dir, Path):
         checkpoint_dir = str(checkpoint_dir)
-        logger.warning(
-            "DEPRECATED: Passing 'checkpoint_dir' as a Path object in load_common_state_dict will "
-            "no longer be supported in a future release. Please pass it as a string instead."
+        log_single_rank(
+            logger,
+            logging.WARNING,
+            "DEPRECATED: Passing 'checkpoint_dir' as a Path object in "
+            "load_common_state_dict will no longer be supported in a future release. "
+            "Please pass it as a string instead.",
         )
     sharded_strategy, common_strategy = verify_checkpoint_and_load_strategy(checkpoint_dir)
     return common_strategy.load_common(checkpoint_dir)
diff --git a/megatron/core/dist_checkpointing/strategies/async_utils.py b/megatron/core/dist_checkpointing/strategies/async_utils.py
index 94af4beef54..85941fd1ed6 100644
--- a/megatron/core/dist_checkpointing/strategies/async_utils.py
+++ b/megatron/core/dist_checkpointing/strategies/async_utils.py
@@ -16,6 +16,8 @@
 import torch
 from torch import multiprocessing as mp
 
+from megatron.core.utils import log_single_rank
+
 from ..utils import debug_time
 
 logger = logging.getLogger(__name__)
@@ -167,7 +169,7 @@ def sync_all_async_calls(self, is_alive: int) -> bool:
     @abstractmethod
     def close(self, abort=False):
         """Terminate the async caller at exit of an application or some termination conditions"""
-        logger.info(f"AsyncCaller: {torch.distributed.get_rank()}, Destroying Async Caller")
+        logger.debug(f"AsyncCaller: {torch.distributed.get_rank()}, Destroying Async Caller")
 
     def __del__(self):
         raise NotImplementedError("This should be implemented")
@@ -265,7 +267,11 @@ def close(self, abort=False):
         if self.process:
             logger.debug(f"rank: {torch.distributed.get_rank()}, joining self.process")
             if abort:
-                logger.warning(f"Temporal worker aborted in rank {torch.distributed.get_rank()}")
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    f"Temporal worker aborted in rank {torch.distributed.get_rank()}",
+                )
                 self.process.kill()
             else:
                 self.process.join()
@@ -319,7 +325,7 @@ def schedule_async_call(self, async_req: AsyncRequest) -> None:
         self.start_time = time()
         if self.process is None:
             ctx = mp.get_context('spawn')
-            logger.info(
+            logger.debug(
                 f"PersistentAsyncCaller: {torch.distributed.get_rank()}, Starting Async Caller"
             )
             self.process: mp.Process = ctx.Process(
@@ -333,7 +339,7 @@ def schedule_async_call(self, async_req: AsyncRequest) -> None:
                 ),
             )
             self.process.start()
-            logger.info(
+            logger.debug(
                 f"PersistentAsyncCaller: {torch.distributed.get_rank()}, Started Async Caller"
             )
 
@@ -419,12 +425,16 @@ def close(self, abort=False):
             abort (bool, optional): Default to False. Needs to be manually set to true when
                 the checkpoint async process needs to be aborted.
         """
-        logger.info(
+        logger.debug(
             f"PersistentAsyncCaller: {torch.distributed.get_rank()}, Destroying Async Caller"
         )
         if self.process:
             if abort:
-                logger.warning(f"Persistent worker aborted in rank {torch.distributed.get_rank()}")
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    f"Persistent worker aborted in rank {torch.distributed.get_rank()}",
+                )
                 self.process.kill()
             else:
                 self.queue.put('DONE')
@@ -469,7 +479,7 @@ def async_loop(
         # Set logger.
         logger = logging.getLogger(__name__)
         logger.setLevel(log_level)
-        logger.info(f"PersistentAsyncCaller: persistent ckpt worker for {rank} has started")
+        logger.debug(f"PersistentAsyncCaller: persistent ckpt worker for {rank} has started")
 
         # Set CUDA device to appropriate local_rank to ensure allocations / CUDA contexts
         # in this new process are on the right device, and device 0 on the node does not
@@ -496,7 +506,7 @@ def async_loop(
                 comp_q.put(item.call_idx)
                 queue.task_done()
 
-        logger.info(f"PersistentAsyncCaller: persistent ckpt worker for {rank}  has terminated")
+        logger.debug(f"PersistentAsyncCaller: persistent ckpt worker for {rank}  has terminated")
 
 
 class _ActiveAsyncRequest(NamedTuple):
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
index 9a0ef354c26..66f7f1aec3b 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
@@ -63,8 +63,9 @@
         DistributedDataParallelConfig,
     )
     from megatron.core.tensor_parallel import get_cuda_rng_tracker
-    from megatron.core.utils import is_submodule
+    from megatron.core.utils import is_submodule, log_single_rank
 
+    HAVE_MCORE = True
     logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.")
 
 except ImportError:
@@ -72,6 +73,18 @@
     from .distributed_data_parallel_config import DistributedDataParallelConfig
     from .utils import get_cuda_rng_tracker, is_submodule
 
+    HAVE_MCORE = False
+
+    def log_single_rank(
+        logger_: logging.Logger, level: int, msg: str, *args, rank: int = 0, **kwargs
+    ):
+        """Fallback log_single_rank when Megatron Core is not available."""
+        if torch.distributed.is_initialized():
+            if torch.distributed.get_rank() == rank:
+                logger_.log(level, msg, *args, **kwargs)
+        else:
+            logger_.log(level, msg, *args, **kwargs)
+
     logger.info("Megatron Core is not installed, Megatron-FSDP will run without Megatron Core.")
 
 try:
@@ -215,11 +228,12 @@ def __exit__(self, *args):
         self.mem_allocator.__exit__(*args)
         for group in self.groups[1:]:
             backend = group._get_backend(torch.device("cuda", torch.cuda.current_device()))
-            if torch.distributed.get_rank() == 0:
-                logger.info(
-                    f"[MultiGroupUBRAllocator] Registering mem pool to group {group}, "
-                    f"group.group_desc:{group.group_desc}"
-                )
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f"[MultiGroupUBRAllocator] Registering mem pool to group {group}, "
+                f"group.group_desc:{group.group_desc}",
+            )
             backend.register_mem_pool(self.pool)
 
 
@@ -1586,14 +1600,17 @@ def __init__(
             NCCL_MEMORY_POOL = nccl_allocator.create_nccl_mem_pool(
                 symmetric=not self.ddp_config.disable_symmetric_registration
             )
-            if torch.distributed.get_rank() == 0:
-                logging.info(
-                    f"[Rank {torch.distributed.get_rank()}] Created NCCL memory pool for \
-                        UserBuffer Registration"
-                )
-                logging.info(
-                    f"[Rank {torch.distributed.get_rank()}] FSDP double buffer is enabled."
-                )
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f"[Rank {torch.distributed.get_rank()}] Created NCCL memory pool for "
+                "UserBuffer Registration",
+            )
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f"[Rank {torch.distributed.get_rank()}] FSDP double buffer is enabled.",
+            )
             # Select the communicator groups to register FSDP buffers.
             self.ubr_groups = [self.dist_index.get_fsdp_group(is_expert_parallel=False)]
             if self.dist_index.get_fsdp_group(is_expert_parallel=True) is not None:
@@ -1615,27 +1632,29 @@ def __init__(
                     )
                 )
 
-            if torch.distributed.get_rank() == 0:
-                logging.info(
-                    f"[ParamAndGradBuffer] FSDP UBRegistration Groups ({len(self.ubr_groups)}):"
-                )
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f"[ParamAndGradBuffer] FSDP UBRegistration Groups ({len(self.ubr_groups)}):",
+            )
             # All ranks in each group must participate in the collective to avoid deadlock.
             for i, group in enumerate(self.ubr_groups):
-                if torch.distributed.get_rank() == 0:
-                    logging.info(
-                        f"Group [{i+1}/{len(self.ubr_groups)}] \
-                            group.group_desc: {group.group_desc}, group.size(): {group.size()}"
-                    )
+                log_single_rank(
+                    logger,
+                    logging.INFO,
+                    f"Group [{i+1}/{len(self.ubr_groups)}] "
+                    f"group.group_desc: {group.group_desc}, group.size(): {group.size()}",
+                )
                 torch.distributed.barrier(group=group, async_op=False)
-                if torch.distributed.get_rank() == 0:
-                    logging.info(
-                        f"Call Success with the group [{i+1}/{len(self.ubr_groups)}] \
-                            group.group_desc: {group.group_desc}"
-                    )
+                log_single_rank(
+                    logger,
+                    logging.INFO,
+                    f"Call Success with the group [{i+1}/{len(self.ubr_groups)}] "
+                    f"group.group_desc: {group.group_desc}",
+                )
             # Call barrier from the global communitcator group
             torch.distributed.barrier(async_op=False)
-            if torch.distributed.get_rank() == 0:
-                logging.info(f"Call Success with the global communicator group")
+            log_single_rank(logger, logging.INFO, "Call Success with the global communicator group")
 
         # If using nccl_ub, it returns a function that registers buffers to the NCCL memory pool
         # Buffer is registered to data_parallel_group and expert_data_parallel_group if it exists
@@ -1754,21 +1773,23 @@ def manual_buffer_registration(self):
         torch.cuda.synchronize()
 
         for group in self.ubr_groups:
-            if torch.distributed.get_rank() == 0:
-                logging.info(
-                    f"[MCORE][FSDP][Manual REG] Registering mem pool to group {group},"
-                    f"group.group_desc:{group.group_desc}, group.size(): {group.size()}"
-                )
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f"[MCORE][FSDP][Manual REG] Registering mem pool to group {group},"
+                f"group.group_desc:{group.group_desc}, group.size(): {group.size()}",
+            )
             nccl_allocator.register_mem_pool(
                 NCCL_MEMORY_POOL,
                 group,
                 symmetric=not self.ddp_config.disable_symmetric_registration,
             )
-            if torch.distributed.get_rank() == 0:
-                logging.info(
-                    f"[MCORE][FSDP][Manual REG] Registered mem pool to group {group},"
-                    f"group.group_desc:{group.group_desc}, group.size(): {group.size()}"
-                )
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f"[MCORE][FSDP][Manual REG] Registered mem pool to group {group},"
+                f"group.group_desc:{group.group_desc}, group.size(): {group.size()}",
+            )
 
     def _log_parameter_groups(self):
         """Compact log of FSDP parameter groups and their parameters."""
@@ -1817,8 +1838,7 @@ def _bytes_to_mb(bytes_val: int) -> str:
             f"Total pad: {_bytes_to_mb(total_padded_bytes)}"
         )
 
-        if torch.distributed.get_rank() == 0:
-            logger.info("\n".join(log_lines))
+        log_single_rank(logger, logging.INFO, "\n".join(log_lines))
 
     def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
         """
@@ -2108,8 +2128,7 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
                 f"CUDA params numel: {cuda_params_numel / 1_000_000:.2f} M, "
                 f"CPU params numel: {cpu_params_numel / 1_000_000:.2f} M"
             )
-            if torch.distributed.get_rank() == 0:
-                logger.info(log_str)
+            log_single_rank(logger, logging.INFO, log_str)
 
         # Initialize the model weight buffer data of each parameter group.
         # Specifically, replace the Torch module's parameter data with tensors
@@ -3751,8 +3770,12 @@ def check_gpu_memory(threshold=0.9):
 
     near_full = allocated_ratio >= threshold or reserved_ratio >= threshold
 
-    if near_full and torch.distributed.get_rank() == 0:
-        logger.info(f"GPU Memory: Allocated: {allocated_ratio:.2%}, Reserved: {reserved_ratio:.2%}")
+    if near_full:
+        log_single_rank(
+            logger,
+            logging.INFO,
+            f"GPU Memory: Allocated: {allocated_ratio:.2%}, Reserved: {reserved_ratio:.2%}",
+        )
     return near_full
 
 
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index e7e32ddc081..7abdaab103b 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -16,6 +16,7 @@
 from megatron.core import parallel_state
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.rerun_state_machine import get_rerun_state_machine
+from megatron.core.utils import log_single_rank
 
 from ..fp8_utils import (
     is_float8tensor,
@@ -162,6 +163,11 @@ def __init__(
         global dist_reduce_scatter_func
         if self.ddp_config.reduce_scatter_with_fp32_accumulation:
             dist_reduce_scatter_func = reduce_scatter_with_fp32_accumulation
+            log_single_rank(
+                logger,
+                logging.INFO,
+                "Using reduce_scatter_with_fp32_accumulation as reduce-scatter implementation",
+            )
 
         # per_param_grad_ready_counts is a dict mapping parameters to number of times
         # `register_grad_ready` is called for that parameter *when
diff --git a/megatron/core/hyper_comm_grid.py b/megatron/core/hyper_comm_grid.py
index 401d4a1c927..f624ba7bdb3 100644
--- a/megatron/core/hyper_comm_grid.py
+++ b/megatron/core/hyper_comm_grid.py
@@ -158,7 +158,10 @@ def create_pg(self, dims: Union[str, list[str]], **kwargs: Any) -> dist.ProcessG
         rank_enum = self._gen_rank_enum(ordered_dims)
         pg, _ = dist.new_subgroups_by_enumeration(rank_enum, backend=self.backend, **kwargs)
 
-        logging.info(f"Generated process group for {unique_group_key} with enumeration {rank_enum}")
+        if dist.get_rank() == 0:
+            logging.info(
+                f"Generated process group for {unique_group_key} with enumeration {rank_enum}"
+            )
         self._pgs[unique_group_key] = pg
         return pg
 
diff --git a/megatron/core/nccl_allocator.py b/megatron/core/nccl_allocator.py
index 8eb4047634c..d475684f0d3 100644
--- a/megatron/core/nccl_allocator.py
+++ b/megatron/core/nccl_allocator.py
@@ -10,7 +10,9 @@
 # pylint: disable=unused-import
 from torch.utils import cpp_extension
 
-from megatron.core.utils import is_torch_min_version
+from megatron.core.utils import is_torch_min_version, log_single_rank
+
+logger = logging.getLogger(__name__)
 
 # MCORE NCCL Allocator copies and modifies the APEX NCCL allocator.
 # The original APEX NCCL allocator is available at:
@@ -153,7 +155,7 @@ def init() -> None:
     # Disables the use of the tensor register allocator hook
     os.environ["TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK"] = "0"
     _build_nccl_allocator()
-    logging.info(f"[MCORE][NCCL_ALLOCATOR] Initialized NCCL Allocator")
+    log_single_rank(logger, logging.INFO, "[MCORE][NCCL_ALLOCATOR] Initialized NCCL Allocator")
 
 
 # register_mem_pool/deregister_mem_pool are used for manual (de)registration of the memory pool.
@@ -169,9 +171,11 @@ def register_mem_pool(pool, group, symmetric=True):
             backend.register_mem_pool(pool, symm=symmetric)
         except TypeError:
             # Older PyTorch/APIs without 'symm' keyword.
-            logging.warning(
-                f"[MCORE][NCCL_ALLOCATOR] Failed in symmetric registration."
-                f"Falling back to registration api without 'symm' keyword!!"
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "[MCORE][NCCL_ALLOCATOR] Failed in symmetric registration. "
+                "Falling back to registration api without 'symm' keyword!!",
             )
             backend.register_mem_pool(pool)
     else:
@@ -230,9 +234,11 @@ def __enter__(self):
                     backend.deregister_mem_pool(self.pool)
                 except RuntimeError:
                     desc = getattr(self.group, "group_desc", None)
-                    logging.warning(
+                    log_single_rank(
+                        logger,
+                        logging.WARNING,
                         f"[MCORE][NCCL_ALLOCATOR] Failed to deregister mem pool from"
-                        f"{repr(self.group)}({desc}) group!!"
+                        f"{repr(self.group)}({desc}) group!!",
                     )
 
     def __exit__(self, *args):
@@ -246,18 +252,22 @@ def __exit__(self, *args):
                         backend.register_mem_pool(self.pool, symm=self.symmetric)
                     except TypeError:
                         # Older PyTorch/APIs without 'symm' keyword.
-                        logging.warning(
-                            f"[MCORE][NCCL_ALLOCATOR] Failed in symmetric registration."
-                            f"Falling back to non-symmetric registration!!"
+                        log_single_rank(
+                            logger,
+                            logging.WARNING,
+                            "[MCORE][NCCL_ALLOCATOR] Failed in symmetric registration. "
+                            "Falling back to non-symmetric registration!!",
                         )
                         backend.register_mem_pool(self.pool)
                 else:
                     backend.register_mem_pool(self.pool)
             except RuntimeError:
                 desc = getattr(self.group, "group_desc", None)
-                logging.warning(
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
                     f"[MCORE][NCCL_ALLOCATOR] Failed to register mem pool to"
-                    f"{repr(self.group)}({desc}) group!!"
+                    f"{repr(self.group)}({desc}) group!!",
                 )
 
         self.mem_context.__exit__(*args)
@@ -315,9 +325,11 @@ def __enter__(self):
                     backend.deregister_mem_pool(self.pool)
                 except RuntimeError:
                     desc = getattr(group, "group_desc", None)
-                    logging.warning(
+                    log_single_rank(
+                        logger,
+                        logging.WARNING,
                         f"[MCORE][MultiGroupMemPoolAllocator] Failed to deregister mem pool from"
-                        f"{repr(group)}({desc}) group!!"
+                        f"{repr(group)}({desc}) group!!",
                     )
 
     def __exit__(self, *args):
@@ -331,18 +343,23 @@ def __exit__(self, *args):
                         backend.register_mem_pool(self.pool, symm=self.symmetric)
                     except TypeError:
                         # Older PyTorch/APIs without 'symm' keyword.
-                        logging.warning(
-                            f"[MCORE][MultiGroupMemPoolAllocator] Failed in symmetric registration."
-                            f"Falling back to non-symmetric registration!!"
+                        log_single_rank(
+                            logger,
+                            logging.WARNING,
+                            "[MCORE][MultiGroupMemPoolAllocator] "
+                            "Failed in symmetric registration. "
+                            "Falling back to non-symmetric registration!!",
                         )
                         backend.register_mem_pool(self.pool)
                 else:
                     backend.register_mem_pool(self.pool)
             except RuntimeError:
                 desc = getattr(group, "group_desc", None)
-                logging.warning(
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
                     f"[MCORE][MultiGroupMemPoolAllocator] Failed to register mem pool to"
-                    f"{repr(group)}({desc}) group!!"
+                    f"{repr(group)}({desc}) group!!",
                 )
         self.mem_context.__exit__(*args)
 
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index 6e093f96f7e..e2b1b0dbd73 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -5,6 +5,7 @@
 
 import gc
 import itertools
+import logging
 from collections import ChainMap
 from dataclasses import replace
 from logging import getLogger
@@ -13,6 +14,8 @@
 import torch
 import torch.nn.functional
 
+from megatron.core.utils import log_single_rank
+
 from ..dist_checkpointing.optimizer import KEEP_VARS_HINT
 
 HAVE_APEX_OR_TE = True
@@ -840,24 +843,32 @@ def make_needed_groups(param_group):
         # Grad scaler.
         if 'grad_scaler' not in state_dict:
             if self.config.fp16:
-                logger.info(
-                    '***WARNING*** found an old checkpoint, will not ' 'load grad scaler ...'
+                log_single_rank(
+                    logger,
+                    logging.INFO,
+                    '***WARNING*** found an old checkpoint, will not load grad scaler ...',
                 )
         else:
             if self.grad_scaler:
                 self.grad_scaler.load_state_dict(state_dict['grad_scaler'])
             else:
-                logger.info(
+                log_single_rank(
+                    logger,
+                    logging.INFO,
                     '***WARNING*** fould the grad scaler in the '
                     'checkpoint but it is None in the class. '
-                    'Skipping loading grad scaler ...'
+                    'Skipping loading grad scaler ...',
                 )
 
         if 'param_state' in state_dict:
             assert 'param_state_sharding_type' in state_dict, state_dict.keys()
             param_state = state_dict['param_state']
             sharding_type = state_dict['param_state_sharding_type']
-            logger.info(f'Loading distributed optimizer sharded state of type {sharding_type}')
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f'Loading distributed optimizer sharded state of type {sharding_type}',
+            )
             if sharding_type == 'dp_zero_gather_scatter':
                 self.load_parameter_state_from_dp_zero(param_state)
             elif sharding_type == 'fully_reshardable':
@@ -1202,10 +1213,12 @@ def sharded_state_dict(
         Regular state dict parameters are saved on DP rank 0 and loaded on all ranks.
         """
         if sharding_type is not None:
-            logger.warning(
+            log_single_rank(
+                logger,
+                logging.WARNING,
                 'DistributedOptimizer.sharded_state_dict parameter `sharding_type`'
                 ' is deprecated and will be removed.'
-                ' Use `metadata["distrib_optim_sharding_type"] instead`.'
+                ' Use `metadata["distrib_optim_sharding_type"] instead`.',
             )
         else:
             sharding_type = (metadata or {}).get(
@@ -1222,10 +1235,12 @@ def sharded_state_dict(
             return state_dict
 
         if not is_loading and sharding_type == 'fully_sharded_bucket_space':
-            logger.warning(
+            log_single_rank(
+                logger,
+                logging.WARNING,
                 '`fully_sharded_bucket_space` sharding for DistributedOptimizer'
                 ' checkpoint is deprecated and will be removed in the future.'
-                ' Please switch to `full_sharded_model_space`.'
+                ' Please switch to `full_sharded_model_space`.',
             )
 
         state_dict = self.state_dict()
diff --git a/megatron/core/post_training/modelopt/gpt/state_dict_hooks.py b/megatron/core/post_training/modelopt/gpt/state_dict_hooks.py
index 596d210f581..22624d9ab2a 100644
--- a/megatron/core/post_training/modelopt/gpt/state_dict_hooks.py
+++ b/megatron/core/post_training/modelopt/gpt/state_dict_hooks.py
@@ -1,8 +1,9 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
+import logging
 from logging import getLogger
 
-import torch
+from megatron.core.utils import log_single_rank
 
 logger = getLogger(__name__)
 
@@ -58,7 +59,6 @@ def mcore_gpt_load_te_state_dict_pre_hook(
                 key_rewrite_list += [(key, key.replace(old_name, new_name))]
 
     for old_key, new_key in key_rewrite_list:
-        if torch.distributed.get_rank() == 0:
-            logger.info("replace {} with {}".format(old_key, new_key))
+        log_single_rank(logger, logging.INFO, "replace {} with {}".format(old_key, new_key))
         state_dict[new_key] = state_dict[old_key]
         state_dict.pop(old_key)
diff --git a/megatron/core/rerun_state_machine.py b/megatron/core/rerun_state_machine.py
index 9ce7259d09c..8fce2beaa85 100644
--- a/megatron/core/rerun_state_machine.py
+++ b/megatron/core/rerun_state_machine.py
@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 
+from megatron.core._rank_utils import log_single_rank, safe_get_rank
 from megatron.core.dist_checkpointing.mapping import ShardedObject
 
 """DISCLAIMER: THIS IS AN EXPERIMENTAL FEATURE.
@@ -234,14 +235,12 @@ def __init__(
 
         self.saved_results: dict[Call, Any] = {}
         self.stats: dict[Caller, QuickStats] = defaultdict(lambda: QuickStats())
-        if _safe_get_rank() == 0:
-            logger.warning(f"RerunStateMachine initialized in mode {mode}")
+        log_single_rank(logger, logging.WARNING, f"RerunStateMachine initialized in mode {mode}")
 
     def set_mode(self, mode: RerunMode) -> None:
         """Method to set the operating mode"""
 
-        if _safe_get_rank() == 0:
-            logger.warning(f"Setting RerunStateMachine mode {mode}")
+        log_single_rank(logger, logging.WARNING, f"Setting RerunStateMachine mode {mode}")
         self.mode = mode
 
     def get_mode(self) -> RerunMode:
@@ -328,7 +327,7 @@ def train_step(data_iterator, ...):
             if not will_rerun:
                 self.state = RerunState.NOT_RUNNING_YET
                 return False
-            if self.mode == RerunMode.VALIDATE_RESULTS and _safe_get_rank() == 0:
+            if self.mode == RerunMode.VALIDATE_RESULTS and safe_get_rank() == 0:
                 logger.warning("Need to rerun step to check reproducibility of initial result")
             self.state = RerunState.RERUNNING_IN_PLACE
             self._restore_state()
@@ -353,10 +352,11 @@ def train_step(data_iterator, ...):
                 [self.continue_requested, self.checkpoint_requested]
             )
             if will_continue:
-                if _safe_get_rank() == 0:
-                    logger.warning(
-                        "Continuing normal execution because failed validation was not fatal"
-                    )
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    "Continuing normal execution because failed validation was not fatal",
+                )
                 self.state = RerunState.NOT_RUNNING_YET
                 return False
             if will_checkpoint:
@@ -376,18 +376,20 @@ def train_step(data_iterator, ...):
                 [self.restart_again_requested, self.continue_requested]
             )
             if will_restart_again:
-                if _safe_get_rank() == 0:
-                    logger.warning(
-                        "Need to restart job from the same checkpoint "
-                        "because it was scheduled on the same node/GPU"
-                    )
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    "Need to restart job from the same checkpoint "
+                    "because it was scheduled on the same node/GPU",
+                )
                 self.state = RerunState.RERUNNING_AGAIN_FROM_CHECKPOINT
             else:
                 if will_continue:
-                    if _safe_get_rank() == 0:
-                        logger.warning(
-                            "Continuing normal execution because failed validation was not fatal"
-                        )
+                    log_single_rank(
+                        logger,
+                        logging.WARNING,
+                        "Continuing normal execution because failed validation was not fatal",
+                    )
                     self.state = RerunState.NOT_RUNNING_YET
             return False
         raise RuntimeError("Should not be here")
@@ -423,33 +425,37 @@ def train_step(data_iterator, ...):
         if self.mode in [RerunMode.DISABLED, RerunMode.REPORT_DETERMINISM_STATS]:
             return False, False, 0
         if self.state == RerunState.RERUNNING_IN_PLACE:
-            if _safe_get_rank() == 0:
-                logger.warning(
-                    "Exiting now. A checkpoint at the last iteration is being saved "
-                    "if further examination is needed"
-                )
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "Exiting now. A checkpoint at the last iteration is being saved "
+                "if further examination is needed",
+            )
             return True, True, EXIT_CODE_FAILED_ON_RESULT_VALIDATION
         elif self.state == RerunState.WILL_RERUN_FROM_CHECKPOINT:
-            if _safe_get_rank() == 0:
-                logger.warning(
-                    "Saving a checkpoint and exiting now. Please resume the job "
-                    "from the checkpoint to rerun the last iteration "
-                    "and establish a diagnostic"
-                )
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "Saving a checkpoint and exiting now. Please resume the job "
+                "from the checkpoint to rerun the last iteration "
+                "and establish a diagnostic",
+            )
             return True, True, EXIT_CODE_RESUME_TO_DISAMBIGUATE
         elif self.state == RerunState.RERUNNING_FROM_CHECKPOINT:
-            if _safe_get_rank() == 0:
-                logger.warning(
-                    "Exiting now. A checkpoint at the last iteration already exists "
-                    "if further examination is needed"
-                )
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "Exiting now. A checkpoint at the last iteration already exists "
+                "if further examination is needed",
+            )
             return False, True, EXIT_CODE_FAILED_ON_RESULT_VALIDATION
         elif self.state == RerunState.RERUNNING_AGAIN_FROM_CHECKPOINT:
-            if _safe_get_rank() == 0:
-                logger.warning(
-                    "Exiting now. Please resume the job from the same checkpoint "
-                    "to rerun the last iteration and establish a diagnostic"
-                )
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "Exiting now. Please resume the job from the same checkpoint "
+                "to rerun the last iteration and establish a diagnostic",
+            )
             return False, True, EXIT_CODE_RESUME_TO_DISAMBIGUATE
         return False, False, 0
 
@@ -514,7 +520,7 @@ def train_step(data_iterator, ...):
                 self._log_validation_error_to_file(
                     status=RerunValidationStatus.RERUN_DISABLED, result=result, message=message
                 )
-                rank: int = _safe_get_rank()
+                rank: int = safe_get_rank()
                 node: str = os.uname()[1]
                 device: int = torch.cuda.current_device()
                 full_message: str = (
@@ -556,7 +562,7 @@ def train_step(data_iterator, ...):
             return
 
         def log_failure(message: str, fatal: bool = True) -> None:
-            rank: int = _safe_get_rank()
+            rank: int = safe_get_rank()
             node: str = os.uname()[1]
             device: int = torch.cuda.current_device()
             if fatal:
@@ -575,8 +581,7 @@ def log_failure(message: str, fatal: bool = True) -> None:
         # the check_for_nan_in_loss_and_grad option but never call validate_result.
         if not self.logged_sdc_enabled:
             self.logged_sdc_enabled = True
-            if _safe_get_rank() == 0:
-                logger.warning("Result validation enabled")
+            log_single_rank(logger, logging.WARNING, "Result validation enabled")
 
         # If this the initial run of the iteration, and no unexpected result has already been
         # identified?
@@ -597,7 +602,7 @@ def log_failure(message: str, fatal: bool = True) -> None:
                 )
                 logger.error(
                     f"Unexpected result {result} "
-                    f"on rank {_safe_get_rank()} "
+                    f"on rank {safe_get_rank()} "
                     f"at iteration #{self.current_iteration} "
                     f"invocation #{validation_call.sequence} "
                     f"(message='{message}')"
@@ -789,11 +794,12 @@ def save_my_model_checkpoint(data_iterator, ...):
                 return None
 
             if ckpt_format != "torch_dist":
-                if _safe_get_rank() == 0:
-                    logger.warning(
-                        "RerunStateMachine checkpoints ONLY SUPPORTED "
-                        "for checkpoint format torch_dist"
-                    )
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    "RerunStateMachine checkpoints ONLY SUPPORTED "
+                    "for checkpoint format torch_dist",
+                )
                 return None
 
         data_iterators: list[RerunDataIterator] = self._sanitize_data_iterators(data_iterator)
@@ -871,13 +877,17 @@ def load_checkpoint(checkpoint, ...)
         """
 
         if self.mode == RerunMode.DISABLED:
-            if _safe_get_rank() == 0:
-                logger.warning(
-                    "RerunStateMachine disabled via CLI, ignoring machine state saved in checkpoint"
-                )
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "RerunStateMachine disabled via CLI, ignoring machine state saved in checkpoint",
+            )
             return
-        if _safe_get_rank() == 0:
-            logger.warning("Getting RerunStateMachine state from checkpoint. Will rerun step.")
+        log_single_rank(
+            logger,
+            logging.WARNING,
+            "Getting RerunStateMachine state from checkpoint. Will rerun step.",
+        )
         self.mode = state_dict["mode"]
         self.current_iteration = state_dict["current_iteration"]
         self.state = state_dict["state"]
@@ -923,7 +933,7 @@ def _get_validation_call_info(self, message: str) -> Call:
         assert frame is not None
         filename: str = inspect.getframeinfo(frame).filename
         lineno: int = frame.f_lineno
-        rank: int = _safe_get_rank()
+        rank: int = safe_get_rank()
         caller = Caller(message=message, rank=rank)
         self.validation_counts[caller] += 1
         sequence: int = self.validation_counts[caller]
@@ -994,7 +1004,7 @@ def _log_validation_error_to_file(
         if self.result_rejected_tracker_filename is not None:
             # Append to log.
             try:
-                rank: int = _safe_get_rank()
+                rank: int = safe_get_rank()
                 node: str = os.uname()[1]
                 device: int = torch.cuda.current_device()
                 with open(self.result_rejected_tracker_filename, "a") as f:
@@ -1057,13 +1067,21 @@ def get_skipped_iterations_from_tracker_file(cls, tracker_file_name: str) -> lis
                             if len(iterations_seen_by_job[job][iteration]) > 1:
                                 iterations_to_ignore.add(iteration)
         except Exception as e:
-            logger.error(f"Could not parse iterations to skip in tracker file! ({e})")
+            log_single_rank(
+                logger, logging.ERROR, f"Could not parse iterations to skip in tracker file! ({e})"
+            )
         iterations_to_skip = sorted(iterations_to_potentially_skip - iterations_to_ignore)
-        logger.warning(f"Will skip these iterations from tracker file: {iterations_to_skip}")
+        log_single_rank(
+            logger,
+            logging.WARNING,
+            f"Will skip these iterations from tracker file: {iterations_to_skip}",
+        )
         if len(iterations_to_ignore) > 0:
-            logger.warning(
+            log_single_rank(
+                logger,
+                logging.WARNING,
                 "Will not skip these iterations due to multiple rank errors: "
-                f"{sorted(iterations_to_ignore)}"
+                f"{sorted(iterations_to_ignore)}",
             )
         return iterations_to_skip
 
@@ -1254,7 +1272,7 @@ def maybe_inject(self) -> bool:
         if not self.should_inject_errors or self.injected_error_type is not None:
             return False
         r: int = (
-            random.randint(0, self.error_injection_rate - 1) + _safe_get_rank()
+            random.randint(0, self.error_injection_rate - 1) + safe_get_rank()
         ) % self.error_injection_rate
         if r != 0:
             return False
@@ -1341,7 +1359,7 @@ def get_rerun_state_machine() -> RerunStateMachine:
     """Helper function to return the singleton instance of the rerun machine."""
 
     if _GLOBAL_RERUN_STATE_MACHINE is None:
-        logger.warning("Implicit initialization of Rerun State Machine!")
+        log_single_rank(logger, logging.WARNING, "Implicit initialization of Rerun State Machine!")
         initialize_rerun_state_machine()
         assert _GLOBAL_RERUN_STATE_MACHINE is not None
     return _GLOBAL_RERUN_STATE_MACHINE
@@ -1355,19 +1373,6 @@ def _set_rerun_state_machine(rerun_state_machine) -> None:
     _GLOBAL_RERUN_STATE_MACHINE = rerun_state_machine
 
 
-def _safe_get_rank() -> int:
-    """Internal function that safely checks and returns the rank of the caller."""
-
-    if torch.distributed.is_initialized():
-        return torch.distributed.get_rank()
-
-    # If torch.distributed is not initialized, try to read environment variables.
-    try:
-        return int(os.environ.get("RANK", 0))
-    except (ValueError, TypeError):
-        return 0
-
-
 def _compare_floats(a: torch.Tensor, b: torch.Tensor) -> float:
     """Internal function that implements the default compare_func.
 
diff --git a/megatron/core/resharding/copy_services/gloo_copy_service.py b/megatron/core/resharding/copy_services/gloo_copy_service.py
index 95f9d454682..c9c83ca74a5 100644
--- a/megatron/core/resharding/copy_services/gloo_copy_service.py
+++ b/megatron/core/resharding/copy_services/gloo_copy_service.py
@@ -44,7 +44,10 @@ def __init__(self):
         self.send_ops: List[SendOp] = []
         self.recv_ops: List[Tuple[RecvOp, torch.Tensor]] = []
         self._copy_stream = torch.cuda.Stream()
-        logger.info(f"GlooCopyService initialized on rank {self.rank} with {self.world_size} ranks")
+        if self.rank == 0:
+            logger.info(
+                f"GlooCopyService initialized on rank {self.rank} with {self.world_size} ranks"
+            )
 
     def submit_send(self, src_tensor: torch.Tensor, dest_rank: int):
         self.send_ops.append(SendOp(task_id=None, tensor=src_tensor, dest_rank=dest_rank))
@@ -71,10 +74,11 @@ def submit_recv_with_id(self, task_id: int, dest_tensor: torch.Tensor, src_rank:
 
     def run(self):
         total_ops = len(self.send_ops) + len(self.recv_ops)
-        logger.info(
-            f"GlooCopyService rank {self.rank}: executing batched communication: "
-            f"{len(self.send_ops)} sends + {len(self.recv_ops)} recvs = {total_ops} ops"
-        )
+        if self.rank == 0:
+            logger.info(
+                f"GlooCopyService rank {self.rank}: executing batched communication: "
+                f"{len(self.send_ops)} sends + {len(self.recv_ops)} recvs = {total_ops} ops"
+            )
 
         p2p_ops: List[dist.P2POp] = []
 
@@ -141,6 +145,7 @@ def run(self):
         if self._copy_stream is not None:
             torch.cuda.current_stream().wait_stream(self._copy_stream)
 
-        logger.info("GlooCopyService: batched communication completed")
+        if self.rank == 0:
+            logger.info("GlooCopyService: batched communication completed")
         self.send_ops.clear()
         self.recv_ops.clear()
diff --git a/megatron/core/resharding/copy_services/nccl_copy_service.py b/megatron/core/resharding/copy_services/nccl_copy_service.py
index 43556f02986..8724279b991 100644
--- a/megatron/core/resharding/copy_services/nccl_copy_service.py
+++ b/megatron/core/resharding/copy_services/nccl_copy_service.py
@@ -45,7 +45,8 @@ def __init__(self):
         # Dedicated stream for local (same-rank) copies to avoid unnecessary
         # serialization with work on the default stream.
         self._copy_stream = torch.cuda.Stream()
-        logger.info(f"NCCLCopyService initialized with {self.world_size} ranks")
+        if self.rank == 0:
+            logger.info(f"NCCLCopyService initialized with {self.world_size} ranks")
 
     def submit_send(self, src_tensor: torch.Tensor, dest_rank: int):
         self.send_ops.append(SendOp(task_id=None, tensor=src_tensor, dest_rank=dest_rank))
@@ -64,12 +65,13 @@ def submit_recv_with_id(self, task_id: int, dest_tensor: torch.Tensor, src_rank:
 
     def run(self):
         total_ops = len(self.send_ops) + len(self.recv_ops)
-        logger.info(
-            "Executing batched communication: %d sends + %d recvs = %d ops",
-            len(self.send_ops),
-            len(self.recv_ops),
-            total_ops,
-        )
+        if self.rank == 0:
+            logger.info(
+                "Executing batched communication: %d sends + %d recvs = %d ops",
+                len(self.send_ops),
+                len(self.recv_ops),
+                total_ops,
+            )
 
         local_sends = [op for op in self.send_ops if op.dest_rank == self.rank]
         remote_sends = [op for op in self.send_ops if op.dest_rank != self.rank]
@@ -121,6 +123,7 @@ def run(self):
         # Make sure the copy stream is finished
         torch.cuda.current_stream().wait_stream(self._copy_stream)
 
-        logger.info("Batched communication completed")
+        if self.rank == 0:
+            logger.info("Batched communication completed")
         self.send_ops.clear()
         self.recv_ops.clear()
diff --git a/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py b/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py
index 458689fa1f4..965f43733a6 100644
--- a/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py
+++ b/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py
@@ -10,6 +10,8 @@
 except ModuleNotFoundError:
     HAVE_TRANSFORMERS = False
 
+from megatron.core.utils import log_single_rank
+
 from .abstract_tokenizer import MegatronTokenizerTextAbstract
 
 logger = logging.getLogger(__name__)
@@ -166,9 +168,11 @@ def __init__(
             tokenizer.resize_token_embeddings(tokenizer_default.vocab_size)
             """
 
-            logger.warning(
+            log_single_rank(
+                logger,
+                logging.WARNING,
                 f'{new_tokens_in_vocab} \n will be added to the vocabulary.\n'
-                f'Please resize your model accordingly.'
+                f'Please resize your model accordingly.',
             )
         self.add_special_tokens(special_tokens_dict)
         self.space_sensitive = self.text_to_tokens('x y') != self.text_to_tokens(
@@ -196,7 +200,11 @@ def add_special_tokens(self, special_tokens_dict: dict) -> int:
         num_tokens_added = self.tokenizer.add_special_tokens(special_tokens_dict)
 
         if num_tokens_added > 0:
-            logger.info(f'{num_tokens_added} special tokens added, resize your model accordingly.')
+            log_single_rank(
+                logger,
+                logging.INFO,
+                f'{num_tokens_added} special tokens added, resize your model accordingly.',
+            )
         for k in self.tokenizer.SPECIAL_TOKENS_ATTRIBUTES:
             setattr(self, k, getattr(self.tokenizer, k, None))
         return num_tokens_added
diff --git a/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py b/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py
index e9d486d4e60..20e13206ceb 100644
--- a/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py
+++ b/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py
@@ -2,6 +2,7 @@
 
 import base64
 import json
+import logging
 import os
 from pathlib import Path
 from typing import Dict, List, Optional
@@ -23,6 +24,9 @@
 SPECIAL_TOKEN_TEMPLATE = "<SPECIAL_{id}>"
 
 
+logger = logging.getLogger(__name__)
+
+
 def reload_mergeable_ranks(
     path: str, max_vocab: Optional[int] = None, num_special_tokens: Optional[int] = None
 ) -> Dict[bytes, int]:
@@ -39,15 +43,18 @@ def reload_mergeable_ranks(
     """
 
     assert path.endswith(".json")
+    from megatron.core.utils import log_single_rank
 
     # reload vocab
     with open(path, "r") as f:
         vocab = json.load(f)
     assert isinstance(vocab, list)
-    print(f"Vocab size: {len(vocab)}")
+    log_single_rank(logger, logging.INFO, f"Vocab size: {len(vocab)}")
     if max_vocab is not None:
         vocab = vocab[:max_vocab]
-        print(f"Cutting vocab to first {len(vocab)} tokens.")
+        from megatron.core.utils import log_single_rank
+
+        log_single_rank(logger, logging.INFO, f"Cutting vocab to first {len(vocab)} tokens")
 
     # build ranks
     ranks: Dict[bytes, int] = {}
@@ -124,10 +131,14 @@ def __init__(
             for i in range(len(special_tokens), num_special_tokens)
         ]
         self.special_filler = special_filler
+        from megatron.core.utils import log_single_rank
+
         if special_filler:
-            print(
+            log_single_rank(
+                logger,
+                logging.INFO,
                 "Adding special tokens: "
-                f"{', '.join(special_tokens)}, {special_filler[0]}, ..., {special_filler[-1]}"
+                f"{', '.join(special_tokens)}, {special_filler[0]}, ..., {special_filler[-1]}",
             )
         self.special_tokens = special_tokens + special_filler
         assert (
diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py
index 3a0632d4ee7..3643c42c3ce 100644
--- a/megatron/core/transformer/cuda_graphs.py
+++ b/megatron/core/transformer/cuda_graphs.py
@@ -441,17 +441,18 @@ def format_mem_bytes(mem_bytes):
             ),
         }
 
-        if torch.distributed.get_rank() == 0:
-            logger.info(
-                "> built %d cuda graph(s) in %.2f sec, with total memory usage: "
-                "allocated %s, reserved %s."
-                % (
-                    len(cls.cudagraph_record),
-                    capture_stats["time"],
-                    format_mem_bytes(capture_stats["allocated_bytes"]),
-                    format_mem_bytes(capture_stats["reserved_bytes"]),
-                )
-            )
+        log_single_rank(
+            logger,
+            logging.INFO,
+            "> built %d cuda graph(s) in %.2f sec, with total memory usage: "
+            "allocated %s, reserved %s."
+            % (
+                len(cls.cudagraph_record),
+                capture_stats["time"],
+                format_mem_bytes(capture_stats["allocated_bytes"]),
+                format_mem_bytes(capture_stats["reserved_bytes"]),
+            ),
+        )
 
         # Mark cuda graphs as created.
         for g in cls.cudagraph_record:
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index ee5ae0ec92a..d7b702f25ec 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -44,6 +44,7 @@
     HAVE_TRITON = False
 
 from megatron.core import config
+from megatron.core._rank_utils import log_single_rank
 from megatron.core.package_info import __version__ as mcore_version
 
 try:
@@ -154,7 +155,9 @@ def validator(func: Callable, max_lifetime: int = 3) -> Callable:
             PkgVersion(introduced_with_version).minor + max_lifetime
             < PkgVersion(mcore_version).minor
         ):
-            logger.warning(
+            log_single_rank(
+                logger,
+                logging.WARNING,
                 "%s has reached end of life. Please migrate to a non-experimental function.",
                 func.__name__,
             )
@@ -219,7 +222,9 @@ def validator(cls: Callable, max_lifetime: int = 3) -> Callable:
             PkgVersion(introduced_with_version).minor + max_lifetime
             < PkgVersion(mcore_version).minor
         ):
-            logger.warning(
+            log_single_rank(
+                logger,
+                logging.WARNING,
                 "%s has reached end of life. Please migrate to a non-experimental function.",
                 cls.__name__,
             )
@@ -828,25 +833,6 @@ def scaled_init_method_normal(sigma, num_layers, multiplier=2.0):
     return functools.partial(torch.nn.init.normal_, mean=0.0, std=std)
 
 
-def log_single_rank(logger: logging.Logger, *args: Any, rank: int = 0, **kwargs: Any):
-    """If torch distributed is initialized, write log on only one rank
-
-    Args:
-        logger (logging.Logger): The logger to write the logs
-
-        args (Tuple[Any]): All logging.Logger.log positional arguments
-
-        rank (int, optional): The rank to write on. Defaults to 0.
-
-        kwargs (Dict[str, Any]): All logging.Logger.log keyword arguments
-    """
-    if torch.distributed.is_initialized():
-        if torch.distributed.get_rank() == rank:
-            logger.log(*args, **kwargs)
-    else:
-        logger.log(*args, **kwargs)
-
-
 def log_on_each_pipeline_stage(
     logger: logging.Logger,
     *args: Any,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index e27321f0096..9eac769567e 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -123,7 +123,7 @@ def parse_args(extra_args_provider=None, ignore_unknown_args=False):
     if not args.enable_msc:
         MultiStorageClientFeature.disable()
         assert MultiStorageClientFeature.is_enabled() is False
-        print('WARNING: The MSC feature is disabled.')
+        warn_rank_0('The MSC feature is disabled.')
 
     return args
 
@@ -463,17 +463,16 @@ def validate_args(args, defaults={}):
             assert args.micro_batch_size == 1, \
                 "micro_batch_size must be 1 when using sequence packing. To increase compute per micro batch increase the sequence length."
 
-    if args.rank == 0:
-        print('using world size: {}, data-parallel size: {}, '
-              'context-parallel size: {}, '
-              'hierarchical context-parallel sizes: {}, '
-              'tensor-model-parallel size: {}, '
-              'pipeline-model-parallel size: {}'.format(
-                  args.world_size, args.data_parallel_size,
-                  args.context_parallel_size,
-                  args.hierarchical_context_parallel_sizes,
-                  args.tensor_model_parallel_size,
-                  args.pipeline_model_parallel_size), flush=True)
+    print_rank_0('using world size: {}, data-parallel size: {}, '
+                 'context-parallel size: {}, '
+                 'hierarchical context-parallel sizes: {}, '
+                 'tensor-model-parallel size: {}, '
+                 'pipeline-model-parallel size: {}'.format(
+                     args.world_size, args.data_parallel_size,
+                     args.context_parallel_size,
+                     args.hierarchical_context_parallel_sizes,
+                     args.tensor_model_parallel_size,
+                     args.pipeline_model_parallel_size))
 
     # Checks.
 
@@ -499,9 +498,8 @@ def validate_args(args, defaults={}):
     del args.model_parallel_size
 
     if args.checkpoint_activations:
-        if args.rank == 0:
-            print('--checkpoint-activations is no longer valid, use --recompute-activations, '
-                  'or, for more control, --recompute-granularity and --recompute-method.')
+        print_rank_0('--checkpoint-activations is no longer valid, use --recompute-activations, '
+                     'or, for more control, --recompute-granularity and --recompute-method.')
         exit()
     del args.checkpoint_activations
 
@@ -537,19 +535,16 @@ def validate_args(args, defaults={}):
         # arguments that are passed to the program. We check this by
         # ensuring the arg is set to None.
         if getattr(args, key, None) is not None:
-            if args.rank == 0:
-                print('WARNING: overriding default arguments for {key}:{v} \
-                       with {key}:{v2}'.format(key=key, v=defaults[key],
-                                               v2=getattr(args, key)),
-                                               flush=True)
+            warn_rank_0('Overriding default arguments for {key}:{v} '
+                        'with {key}:{v2}'.format(key=key, v=defaults[key],
+                                                 v2=getattr(args, key)))
         else:
             setattr(args, key, defaults[key])
 
     if args.data_path is not None and args.split is None:
         legacy_default_split_value = '969, 30, 1'
-        if args.rank == 0:
-            print('WARNING: Please specify --split when using --data-path. Using legacy default value '
-                  f'of "{legacy_default_split_value}"')
+        warn_rank_0('Please specify --split when using --data-path. Using legacy default value '
+                    f'of "{legacy_default_split_value}"')
         args.split = legacy_default_split_value
 
     use_data_path = (args.data_path is not None) or (args.data_args_path is not None)
@@ -577,9 +572,7 @@ def validate_args(args, defaults={}):
     assert args.micro_batch_size > 0
     if args.global_batch_size is None:
         args.global_batch_size = args.micro_batch_size * args.data_parallel_size
-        if args.rank == 0:
-            print('setting global batch size to {}'.format(
-                args.global_batch_size), flush=True)
+        print_rank_0('setting global batch size to {}'.format(args.global_batch_size))
     assert args.global_batch_size > 0
 
     # Uneven virtual pipeline parallelism
@@ -673,8 +666,9 @@ def validate_args(args, defaults={}):
                 'since non-interleaved schedule does not support overlapping p2p communication '
                 'and aligned param AG')
 
-    if args.rank == 0:
-        print(f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}")
+    print_rank_0(
+        f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}"
+    )
 
     if args.overlap_param_gather:
         assert args.use_distributed_optimizer or args.use_megatron_fsdp, \
@@ -788,9 +782,8 @@ def validate_args(args, defaults={}):
         # where NaNs in grads / loss are signal to the loss scaler.
         if not args.loss_scale:
             args.check_for_nan_in_loss_and_grad = False
-            if args.rank == 0:
-                print('WARNING: Setting args.check_for_nan_in_loss_and_grad to False since '
-                      'dynamic loss scaling is being used')
+            warn_rank_0('Setting args.check_for_nan_in_loss_and_grad to False since '
+                        'dynamic loss scaling is being used')
     if args.bf16:
         assert not args.fp16
         args.params_dtype = torch.bfloat16
@@ -804,9 +797,7 @@ def validate_args(args, defaults={}):
             args.accumulate_allreduce_grads_in_fp32 = False
         elif not args.accumulate_allreduce_grads_in_fp32 and args.main_grads_dtype == torch.float32:
             args.accumulate_allreduce_grads_in_fp32 = True
-            if args.rank == 0:
-                print('accumulate and all-reduce gradients in fp32 for '
-                      'bfloat16 data type.', flush=True)
+            print_rank_0('accumulate and all-reduce gradients in fp32 for bfloat16 data type.')
     if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope:
         if not args.inference_dynamic_batching:
             assert not args.check_for_nan_in_loss_and_grad, \
@@ -815,9 +806,7 @@ def validate_args(args, defaults={}):
             assert args.fp8 is None, \
             "fp8 is not supported with inference dynamic batching and full_iteration CUDA graph"
 
-    if args.rank == 0:
-        print('using {} for parameters ...'.format(args.params_dtype),
-              flush=True)
+    print_rank_0('using {} for parameters ...'.format(args.params_dtype))
 
     if args.dataloader_type is None:
         args.dataloader_type = 'single'
@@ -963,10 +952,9 @@ def validate_args(args, defaults={}):
     # Persistent fused layer norm.
     if not is_torch_min_version("1.11.0a0"):
         args.no_persist_layer_norm = True
-        if args.rank == 0:
-            print('Persistent fused layer norm kernel is supported from '
-                  'pytorch v1.11 (nvidia pytorch container paired with v1.11). '
-                  'Defaulting to no_persist_layer_norm=True')
+        print_rank_0('Persistent fused layer norm kernel is supported from '
+                     'pytorch v1.11 (nvidia pytorch container paired with v1.11). '
+                     'Defaulting to no_persist_layer_norm=True')
 
     # Activation recomputing.
     if args.distribute_saved_activations:
@@ -1118,8 +1106,7 @@ def validate_args(args, defaults={}):
         args.num_experts = None
     if args.num_experts is not None and args.moe_ffn_hidden_size is None:
         args.moe_ffn_hidden_size = args.ffn_hidden_size
-        if args.rank == 0:
-            print("Warning: moe_ffn_hidden_size is not set, using ffn_hidden_size for MoE instead.")
+        warn_rank_0("moe_ffn_hidden_size is not set, using ffn_hidden_size for MoE instead.")
 
     # Context parallel
     if args.context_parallel_size > 1:
@@ -1241,11 +1228,10 @@ def validate_args(args, defaults={}):
 
     if args.use_dist_ckpt and args.async_save:
         if not args.use_persistent_ckpt_worker:
-            if args.rank == 0:
-                print(
-                    'Warning: --async-save is not supported without --use-persistent-ckpt-worker. '
-                    'Disabling --async-save.'
-                )
+            warn_rank_0(
+                '--async-save is not supported without --use-persistent-ckpt-worker. '
+                'Disabling --async-save.'
+            )
             args.async_save = False
 
     # Inference args
@@ -1269,18 +1255,15 @@ def validate_args(args, defaults={}):
         assert args.save is not None, "When using upcycling, the --save option must be specified."
         if not args.no_load_optim:
             args.no_load_optim = True
-            if args.rank == 0:
-                print('Warning: enabling --no-load-optim for upcycling.')
+            warn_rank_0('enabling --no-load-optim for upcycling.')
         if not args.no_load_rng:
             args.no_load_rng = True
-            if args.rank == 0:
-                print('Warning: enabling --no-load-rng for upcycling.')
+            warn_rank_0('enabling --no-load-rng for upcycling.')
 
     # --skip-train checks.
     if args.skip_train and not args.no_load_optim:
         args.no_load_optim = True
-        if args.rank == 0:
-            print('Warning: enabling --no-load-optim when skipping training.')
+        warn_rank_0('enabling --no-load-optim when skipping training.')
 
     # Muon optimizer check
     if 'muon' in args.optimizer:
@@ -1311,7 +1294,7 @@ def validate_args(args, defaults={}):
         assert args.replication_jump is not None, "--replication requires the value of --replication-jump!"
         assert args.non_persistent_ckpt_type == "local", f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}"
     elif args.replication_jump:
-        print("Warning: --replication-jump was specified despite not using replication. Ignoring.")
+        warn_rank_0("--replication-jump was specified despite not using replication. Ignoring.")
         args.replication_jump = None
 
     if args.delay_wgrad_compute:
@@ -1395,17 +1378,16 @@ def validate_args(args, defaults={}):
 
 def _print_args(title, args):
     """Print arguments."""
-    if args.rank == 0:
-        print(f'------------------------ {title} ------------------------',
-              flush=True)
+    from megatron.training.utils import is_rank0
+    if is_rank0():
+        print(f'------------------------ {title} ------------------------', flush=True)
         str_list = []
         for arg in vars(args):
             dots = '.' * (48 - len(arg))
             str_list.append('  {} {} {}'.format(arg, dots, getattr(args, arg)))
         for arg in sorted(str_list, key=lambda x: x.lower()):
             print(arg, flush=True)
-        print(f'-------------------- end of {title} ---------------------',
-              flush=True)
+        print(f'-------------------- end of {title} ---------------------', flush=True)
 
 
 def _check_arg_is_not_none(args, arg):
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index b73466ccfde..4a0218c7106 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -526,8 +526,8 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
             raise NotImplementedError(f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})")
 
     ckpt_format = args.ckpt_format if ckpt_type == CheckpointType.GLOBAL else 'torch'
-    print_rank_0('saving checkpoint at iteration {:7d} to {} in {} format'.format(
-        iteration, save_dir, ckpt_format))
+    print_rank_0(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving checkpoint "
+                 f"at iteration {iteration:7d} to {save_dir} in {ckpt_format} format")
 
     # Collect rng state across data parallel ranks.
     if tp_group is None and pp_group is None:
@@ -707,8 +707,8 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
 
         if ckpt_type == CheckpointType.LOCAL:
             def iter_finalize_fn():
-                print_rank_0('  successfully saved local checkpoint from iteration {:7d}'
-                             .format(iteration))
+                print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully "
+                             f"saved local checkpoint from iteration {iteration:7d}")
                 if args.log_progress and args.async_save:
                     append_to_progress_log(f'Saved async local checkpoint\tIteration: {iteration}',
                                            barrier=False)
@@ -724,9 +724,10 @@ def iter_finalize_fn():
                     f.write("release" if release else str(iteration))
                 tensor_rank_to_print = (tensor_rank if tensor_rank is not None else mpu.get_tensor_model_parallel_rank()) + 1
                 pipeline_rank_to_print = (pipeline_rank if pipeline_rank is not None else mpu.get_pipeline_model_parallel_rank()) + 1
-                print_rank_0(f'  successfully saved checkpoint from iteration {int(iteration):7d} to {args.save} '
-                             f'[ t {tensor_rank_to_print}/{mpu.get_tensor_model_parallel_world_size()}, '
-                             f'p {pipeline_rank_to_print}/{mpu.get_pipeline_model_parallel_world_size()} ]')
+                print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully saved "
+                             f"checkpoint from iteration {int(iteration):7d} to {args.save} "
+                             f"[ t {tensor_rank_to_print}/{mpu.get_tensor_model_parallel_world_size()}, "
+                             f"p {pipeline_rank_to_print}/{mpu.get_pipeline_model_parallel_world_size()} ]")
                 if args.log_progress and args.async_save:
                     append_to_progress_log(f'Saved async checkpoint\tIteration: {iteration}',
                                            barrier=False)
@@ -736,8 +737,9 @@ def delete_checkpoint(args, iteration_to_delete):
                                                           return_base_dir=True)
                     try:
                         shutil.rmtree(checkpoint_name)  # TODO: Make this work with MSC remote paths?
-                        print_rank_0(f'  successfully deleted checkpoint from iteration {iteration_to_delete:7d} '
-                                     f'at {args.save}')
+                        print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully "
+                                     f"deleted checkpoint from iteration {iteration_to_delete:7d} "
+                                     f"at {args.save}")
                         if args.log_progress:
                             append_to_progress_log(f'Deleted checkpoint\tIteration: {iteration_to_delete}', barrier=False)
                     except Exception as e:
@@ -788,8 +790,8 @@ def wandb_finalize_fn():
 
     if args.async_save:
         schedule_async_save(async_save_request)
-        print_rank_0('  scheduled an async checkpoint save at iteration {:7d} to {}' \
-                     .format(iteration, save_dir))
+        print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] scheduled "
+                     f"an async checkpoint save at iteration {iteration:7d} to {save_dir}")
 
     # Wait so everyone is done (not necessary)
     if torch.distributed.is_initialized():
@@ -852,7 +854,8 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path)
         return
 
     dp_rank = mpu.get_data_parallel_rank()
-    print(f"saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}")
+    if dp_rank == 0:
+        print(f"saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}")
     train_dataloader_state_dict = train_iterator.iterable.save_state()
     data_state_save_path = get_checkpoint_name(
         dataloader_save_path, iteration,
@@ -1833,7 +1836,7 @@ def load_model_state_dict(module, state_dict, strict: bool):
             if 'rerun_state_machine' in state_dict:
                 get_rerun_state_machine().load_state_dict(state_dict['rerun_state_machine'])
         except Exception as e:
-            print(f"Unable to restore RerunMachine from checkpoint: {e}. Skipping.")
+            print_rank_0(f"Unable to restore RerunMachine from checkpoint: {e}. Skipping.")
 
     # rng states.
     if not release and not args.finetune and not args.no_load_rng and not ignore_rng_state:
@@ -1848,7 +1851,7 @@ def load_model_state_dict(module, state_dict, strict: bool):
                     if f"({pp_rank}, {tp_rank})" in state_dict['rng_state']:
                         rng_state = state_dict['rng_state'][f"({pp_rank}, {tp_rank})"]
                     else:
-                        print("WARNING: RNG state not found for current TP/PP rank")
+                        print_rank_0("WARNING: RNG state not found for current TP/PP rank")
                         rng_state = next(iter(state_dict['rng_state'].values()))
                 else:
                     rng_state = state_dict['rng_state']
diff --git a/megatron/training/global_vars.py b/megatron/training/global_vars.py
index a718877b40c..76e8df7cee3 100644
--- a/megatron/training/global_vars.py
+++ b/megatron/training/global_vars.py
@@ -249,13 +249,13 @@ def _set_adlr_autoresume(args):
     _ensure_var_is_not_initialized(_GLOBAL_ADLR_AUTORESUME, 'adlr autoresume')
 
     if args.adlr_autoresume:
-        if args.rank == 0:
-            print('enabling autoresume ...', flush=True)
+        from megatron.training.utils import print_rank_0
+        print_rank_0('enabling autoresume ...')
         sys.path.append(os.environ.get('SUBMIT_SCRIPTS', '.'))
         try:
             from userlib.auto_resume import AutoResume
         except ImportError:
-            print('ADLR autoresume is not available, exiting ...')
+            print_rank_0('ADLR autoresume is not available, exiting ...')
             sys.exit()
 
         _GLOBAL_ADLR_AUTORESUME = AutoResume
diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py
index e300c03218b..c150ac3d5ca 100644
--- a/megatron/training/initialize.py
+++ b/megatron/training/initialize.py
@@ -26,11 +26,13 @@
 from megatron.core.utils import get_te_version, is_te_min_version, is_torch_min_version
 from megatron.legacy import fused_kernels
 from megatron.training import get_adlr_autoresume, get_args, get_tensorboard_writer
+from megatron.training.utils import print_rank_0, warn_rank_0
 from megatron.training import inprocess_restart
 from megatron.training.arguments import parse_args, validate_args
 from megatron.training.async_utils import init_persistent_async_worker
 from megatron.training.checkpointing import load_args_from_checkpoint
 from megatron.training.global_vars import set_global_variables
+from megatron.training.utils import is_rank0
 from megatron.training.yaml_arguments import validate_yaml
 
 logger = logging.getLogger(__name__)
@@ -117,8 +119,7 @@ def state_restore_func(state_dict):
     )
     
     if args.batch_invariant_mode:
-        if args.rank == 0:
-            print("Enabling batch invariant mode globally", flush=True)
+        print_rank_0("Enabling batch invariant mode globally")
         enable_batch_invariant_mode()
 
     # torch.distributed initialization
@@ -128,8 +129,7 @@ def finish_mpu_init():
         _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, store)
 
         # Random seeds for reproducibility.
-        if args.rank == 0:
-            print("> setting random seeds to {} ...".format(args.seed))
+        print_rank_0("> setting random seeds to {} ...".format(args.seed))
         _set_random_seed(
             args.seed,
             args.data_parallel_random_init,
@@ -212,13 +212,10 @@ def _compile_dependencies():
     )
     # Print a warning.
     if not ((args.fp16 or args.bf16) and custom_kernel_constraint and args.masked_softmax_fusion):
-        if args.rank == 0:
-            print(
-                "WARNING: constraints for invoking optimized"
-                " fused softmax kernel are not met. We default"
-                " back to unfused kernel invocations.",
-                flush=True,
-            )
+        warn_rank_0(
+            "Constraints for invoking optimized fused softmax kernel are not met. "
+            "We default back to unfused kernel invocations."
+        )
 
     # Always build on rank zero first.
     if torch.distributed.get_rank() == 0:
@@ -322,18 +319,13 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s
     device_count = torch.cuda.device_count()
     if torch.distributed.is_initialized():
 
-        if args.rank == 0:
-            print(
-                "torch distributed is already initialized, " "skipping initialization ...",
-                flush=True,
-            )
+        print_rank_0("torch distributed is already initialized, skipping initialization ...")
         args.rank = torch.distributed.get_rank()
         args.world_size = torch.distributed.get_world_size()
 
     else:
 
-        if args.rank == 0:
-            print("> initializing torch distributed ...", flush=True)
+        print_rank_0("> initializing torch distributed ...")
         # Manually set the device ids.
         if device_count > 0:
             torch.cuda.set_device(args.local_rank)
@@ -391,15 +383,14 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s
                 sharp_enabled_group=args.sharp_enabled_group,
                 create_all_gather_group=args.create_all_gather_group,
             )
-            if args.rank == 0:
-                print(
-                    f"> initialized tensor model parallel with size "
-                    f"{mpu.get_tensor_model_parallel_world_size()}"
-                )
-                print(
-                    f"> initialized pipeline model parallel with size "
-                    f"{mpu.get_pipeline_model_parallel_world_size()}"
-                )
+            print_rank_0(
+                f"> initialized tensor model parallel with size "
+                f"{mpu.get_tensor_model_parallel_world_size()}"
+            )
+            print_rank_0(
+                f"> initialized pipeline model parallel with size "
+                f"{mpu.get_pipeline_model_parallel_world_size()}"
+            )
 
 
 def _init_autoresume():
@@ -551,5 +542,6 @@ def setup_logging() -> None:
         logging_level = args.logging_level
 
     if logging_level is not None:
-        logger.info(f'Setting logging level to {logging_level}')
+        if is_rank0():
+            logger.info(f'Setting logging level to {logging_level}')
         logging.getLogger().setLevel(logging_level)
diff --git a/megatron/training/tokenizer/tokenizer.py b/megatron/training/tokenizer/tokenizer.py
index 08d8aacd9ea..33340a5e978 100644
--- a/megatron/training/tokenizer/tokenizer.py
+++ b/megatron/training/tokenizer/tokenizer.py
@@ -20,8 +20,8 @@
 
 def build_tokenizer(args, **kwargs):
     """Initialize tokenizer."""
-    if args.rank == 0:
-        print('> building {} tokenizer ...'.format(args.tokenizer_type), flush=True)
+    from megatron.training.utils import print_rank_0
+    print_rank_0('> building {} tokenizer ...'.format(args.tokenizer_type))
 
     # Select and instantiate the tokenizer.
     if args.tokenizer_type == 'BertWordPieceLowerCase':
diff --git a/megatron/training/training.py b/megatron/training/training.py
index be4f29a3476..84f6d6f771d 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -2026,7 +2026,7 @@ def training_log(
                 writer.add_scalar('iteration-time', elapsed_time_per_iteration, iteration)
             if wandb_writer:
                 wandb_writer.log({'iteration-time': elapsed_time_per_iteration}, iteration)
-        log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+        log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}]"
         log_string += ' iteration {:8d}/{:8d} |'.format(iteration, args.train_iters)
         log_string += ' consumed samples: {:12d} |'.format(args.consumed_train_samples)
         if has_rl_utils and args.rl_use_sequence_packing:
@@ -3006,7 +3006,7 @@ def get_e2e_base_metrics():
     if args.log_energy:
         energy_monitor.lap()
         total_energy = energy_monitor.get_total()
-        print_rank_0(f"Total training energy (GPU): {total_energy / 1e6} MJ")
+        print_rank_0(f"Total training energy (GPU): {total_energy / 1e6:.3f} MJ")
         energy_monitor.shutdown()
 
     # If any exit conditions (signal handler, duration, iterations) have been reached, exit.
diff --git a/megatron/training/utils.py b/megatron/training/utils.py
index cdf9631c3b5..d9681728467 100644
--- a/megatron/training/utils.py
+++ b/megatron/training/utils.py
@@ -12,6 +12,7 @@
 import torch
 
 from megatron.core.msc_utils import MultiStorageClientFeature, open_file
+from megatron.core._rank_utils import safe_get_rank as _safe_get_rank
 
 try:
     from transformer_engine.pytorch.optimizers import multi_tensor_applier, multi_tensor_l2norm
@@ -394,11 +395,9 @@ def print_rank_0(message, rank=None):
     if rank is not None:
         if rank == 0:
             print(message, flush=True)
-    elif torch.distributed.is_initialized():
-        if torch.distributed.get_rank() == 0:
-            print(message, flush=True)
     else:
-        print(message, flush=True)
+        if _safe_get_rank() == 0:
+            print(message, flush=True)
 
 
 def warn_rank_0(message, rank=None):
@@ -406,20 +405,20 @@ def warn_rank_0(message, rank=None):
     if rank is not None:
         if rank == 0:
             warnings.warn(message)
-    elif torch.distributed.is_initialized():
-        if torch.distributed.get_rank() == 0:
-            warnings.warn(message)
     else:
-        warnings.warn(message)
+        if _safe_get_rank() == 0:
+            warnings.warn(message)
 
 
 def is_rank0():
-    """Returns true if called in the rank0, false otherwise"""
-    return torch.distributed.is_initialized() and torch.distributed.get_rank() == 0
+    """Returns true if called in the rank0, false otherwise."""
+    return _safe_get_rank() == 0
 
 
 def is_last_rank():
-    return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
+    """Returns true if called on last rank, false otherwise."""
+    assert torch.distributed.is_initialized()
+    return _safe_get_rank() == (torch.distributed.get_world_size() - 1)
 
 
 def print_rank_last(message):
diff --git a/pretrain_gpt.py b/pretrain_gpt.py
index c2e79715e48..8eff08d24b2 100644
--- a/pretrain_gpt.py
+++ b/pretrain_gpt.py
@@ -2,11 +2,20 @@
 
 """Pretrain and SFT GPT."""
 
-# Capture the true program start time BEFORE any heavy imports
+# Capture the true program start time BEFORE any heavy imports.
 import time
 _PROGRAM_START_TIME = time.time()
 
 import json
+
+# Suppress warnings on all ranks but rank 0.
+import os
+import warnings
+rank = int(os.environ.get('RANK', 0))
+if rank != 0:
+    warnings.filterwarnings("ignore", category=UserWarning)
+    warnings.filterwarnings("ignore", category=FutureWarning)
+
 from functools import partial
 from typing import List, Optional, Tuple
 
diff --git a/pretrain_mamba.py b/pretrain_mamba.py
index 6fcc0d25c45..e1379be63e9 100644
--- a/pretrain_mamba.py
+++ b/pretrain_mamba.py
@@ -1,11 +1,20 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 """Pretrain and SFT Mamba."""
 
-# Capture the true program start time BEFORE any heavy imports
+# Capture the true program start time BEFORE any heavy imports.
 import time
 _PROGRAM_START_TIME = time.time()
 
 import json
+
+# Suppress warnings on all ranks but rank 0.
+import os
+import warnings
+rank = int(os.environ.get('RANK', 0))
+if rank != 0:
+    warnings.filterwarnings("ignore", category=UserWarning)
+    warnings.filterwarnings("ignore", category=FutureWarning)
+
 from functools import partial
 from typing import List, Optional, Tuple
 

From 65217aa2b62494d82b0e61378032f1f35020aa43 Mon Sep 17 00:00:00 2001
From: Philip Petrakian <ppetrakian@nvidia.com>
Date: Tue, 27 Jan 2026 09:58:56 -0800
Subject: [PATCH 38/79] Move all bert and t5 tests to nightly (#3106)

---
 tests/test_utils/recipes/bert.yaml |  8 ++++----
 tests/test_utils/recipes/t5.yaml   | 10 +++++-----
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tests/test_utils/recipes/bert.yaml b/tests/test_utils/recipes/bert.yaml
index 49fed7f5542..89499f93c5e 100644
--- a/tests/test_utils/recipes/bert.yaml
+++ b/tests/test_utils/recipes/bert.yaml
@@ -59,22 +59,22 @@ products:
   - test_case: [bert_mcore_tp2_pp2]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [bert_mcore_tp2_pp2_local_spec]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [bert_mcore_tp2_pp2_resume_torch_dist]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [bert_mcore_tp2_pp2_resume_torch_dist_local_spec]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [bert_mcore_tp1_pp2]
     products:
diff --git a/tests/test_utils/recipes/t5.yaml b/tests/test_utils/recipes/t5.yaml
index 96b560c6427..1761cd3f1e6 100644
--- a/tests/test_utils/recipes/t5.yaml
+++ b/tests/test_utils/recipes/t5.yaml
@@ -59,27 +59,27 @@ products:
   - test_case: [t5_11b_mcore_tp4_pp1]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [t5_mcore_te_tp4_pp1]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [t5_mcore_te_tp4_pp1_resume_torch_dist]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [t5_mcore_tp4_pp1]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [t5_mcore_tp4_pp1_resume_torch_dist]
     products:
       - environment: [dev]
-        scope: [mr]
+        scope: [nightly]
         platforms: [dgx_h100]
   - test_case: [t5_mcore_te_tp1_pp1_vp1_resume_torch]
     products:

From 4fb549f72373a3e31c66361778b0f8ca4d6f25ea Mon Sep 17 00:00:00 2001
From: Philip Petrakian <ppetrakian@nvidia.com>
Date: Tue, 27 Jan 2026 16:10:30 -0800
Subject: [PATCH 39/79] Create greptile.json (#3087)

---
 greptile.json | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)
 create mode 100644 greptile.json

diff --git a/greptile.json b/greptile.json
new file mode 100644
index 00000000000..e08c2def387
--- /dev/null
+++ b/greptile.json
@@ -0,0 +1,39 @@
+{
+    "labels": [],
+    "comment": "Disclaimer: This is AI-generated.",
+    "commentTypes": ["logic", "syntax", "style"],
+    "instructions": "Only comment if the PR description is unchanged from the default template, if a docstring is missing, or if there is a typo.",
+    "ignoreKeywords": "rename\nlinter\nprettier\ngreptile-ignor",
+    "ignorePatterns": "greptile.json\ntesting/**/*.py\n*.md\n*.txt\n*.json",
+    "patternRepositories": ["NVIDIA/Megatron-LM"],
+    "triggerOnUpdates": true,
+    "shouldUpdateDescription": false,
+    "disabledLabels": ["docs"],
+    "includeAuthors": [],
+    "excludeAuthors": ["github-actions"],
+    "strictness": 3,
+    "fixWithAI": false,
+    "includeBranches": ["main"],
+    "statusCheck": false,
+    "skipReview": "AUTOMATIC",
+    "summarySection": {
+      "included": false,
+      "collapsible": false,
+      "defaultOpen": false
+    },
+    "issuesTableSection": {
+      "included": false,
+      "collapsible": false,
+      "defaultOpen": false
+    },
+    "confidenceScoreSection": {
+      "included": false,
+      "collapsible": false,
+      "defaultOpen": false
+    },
+    "sequenceDiagramSection": {
+      "included": false,
+      "collapsible": false,
+      "defaultOpen": false
+    }
+  }
\ No newline at end of file

From 6273d74a5630c65a82897bb0638eb6f46875177c Mon Sep 17 00:00:00 2001
From: GitHub Actions <github-actions[bot]@users.noreply.github.com>
Date: Wed, 28 Jan 2026 00:12:30 +0000
Subject: [PATCH 40/79] Update copy-pr-bot.yaml [skip ci]

---
 .github/copy-pr-bot.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml
index 2ece09a999f..e8e5bcd69cb 100644
--- a/.github/copy-pr-bot.yaml
+++ b/.github/copy-pr-bot.yaml
@@ -1,4 +1,4 @@
 enabled: true
 auto_sync_draft: false
 auto_sync_ready: true
-trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "pablo-garay", "parthmannan", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]
+trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]

From 33224ccb75e117da95738658a59fdc55b79e6463 Mon Sep 17 00:00:00 2001
From: Kunlun Li <94586211+kunlunl@users.noreply.github.com>
Date: Wed, 28 Jan 2026 08:15:02 +0800
Subject: [PATCH 41/79] Fix bug of reuse_grad_buf_for_mxfp8_param_ag (#2802)

Signed-off-by: kunlunl <kunlunl@nvidia.com>
---
 megatron/training/training.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/megatron/training/training.py b/megatron/training/training.py
index 84f6d6f771d..fbc267fba82 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -1638,10 +1638,19 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
         # For the mxfp8_param with reuse_grad_buf_for_mxfp8_param_ag and dp_ag_overlap,
         # we need to call the _copy_main_params_to_param_buffer() after the grad buffer
         # is zeroed by zero_grad_buffer() because param and grad buffer are shared.
+        #
+        # However, we should skip this on the first iteration when forward_pre_hook is disabled,
+        # because:
+        # 1. The first iteration's params are already in param.data (from init or checkpoint).
+        # 2. Without forward_pre_hook, finish_param_sync() won't be called to zero the grad buffer,
+        #    so the main grads will be polluted by the main params.
         if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather:
-            for optim_instance in optimizer.chained_optimizers:
-                if isinstance(optim_instance, DistributedOptimizer):
-                    optim_instance._copy_main_params_to_param_buffer()
+            # Check if forward_pre_hook is enabled by checking if hooks are registered.
+            forward_pre_hook_enabled = len(model[0].remove_forward_pre_hook_handles) > 0
+            if forward_pre_hook_enabled:
+                for optim_instance in optimizer.chained_optimizers:
+                    if isinstance(optim_instance, DistributedOptimizer):
+                        optim_instance._copy_main_params_to_param_buffer()
 
         # Forward pass.
         if save_dgrads_in_this_iteration:

From fb6a592350c5ac4c15a10eab2eb4f0f05576a8e0 Mon Sep 17 00:00:00 2001
From: Parth Mannan <38387286+parthmannan@users.noreply.github.com>
Date: Tue, 27 Jan 2026 15:47:25 -0800
Subject: [PATCH 42/79] Fix for Hybrid CP (#3091)

---
 megatron/training/arguments.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 9eac769567e..9dcc8f65e81 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2851,13 +2851,13 @@ def _add_distributed_args(parser):
                        '--hierarchical-context-parallel-sizes 2 4 indicates every two adjacent gpus '
                        'forms the first level of cp groups and the cp ranks with the same odevity '
                        'forms the second level of cp groups.')
-    group.add_argument('--max-seqlen-per-cp-rank', type=int, default=None,
-                       help='Maximum sequence length per CP rank. This is used to calculate the '
-                       'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.')
+    group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None,
+                       help='Maximum sequence length per DPxCP rank. This is used to calculate the '
+                       'number of sub-samples assigned to each DPxCP rank when using Hybrid Context Parallel.')
     group.add_argument('--hybrid-context-parallel', action='store_true', default=False,
                        help='Enables hybrid context parallel. This is used to balance the workload '
                        'of each CP rank when we use packed samples with variable sequence lengths. '
-                       'Requires --max-seqlen-per-cp-rank to be set.')
+                       'Requires --max-seqlen-per-dp-cp-rank to be set.')
     group.add_argument('--nccl-communicator-config-path', type=str, default=None,
                        help='Path to the yaml file with NCCL communicator '
                        'configurations. The number of min/max thread groups and thread '

From f6c8a61be78ef71f4442be376b70825f92d34e95 Mon Sep 17 00:00:00 2001
From: Jon Barker <jbarker@nvidia.com>
Date: Tue, 27 Jan 2026 17:42:57 -0700
Subject: [PATCH 43/79] Fix GRPO re-fit functional test (#3113)

---
 .../model_config.yaml                                  | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml
index 47228df80b4..ed5d123892e 100644
--- a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml
@@ -129,3 +129,13 @@ MODEL_ARGS:
   --tensorboard-dir: ${TENSORBOARD_PATH}
   --straggler-minmax-count: 16
   --empty-unused-memory-level: 2
+
+METRICS:
+  - "iteration-time"
+  - "lm loss"
+  - "num-zeros"
+  - "mem-allocated-bytes"
+  - "mem-max-allocated-bytes"
+
+THROUGHPUT_TEST_PARAMS:
+  --start_step: 1

From 991138eedc580ba89ce8e44e771783f083cd9dee Mon Sep 17 00:00:00 2001
From: megnvidia <mmiranda@nvidia.com>
Date: Tue, 27 Jan 2026 16:39:18 -0800
Subject: [PATCH 44/79] Minimize README contents (#3020)

Signed-off-by: meg miranda <mmiranda@nvidia.com>
Co-authored-by: Philip Petrakian <pgpetrak@gmail.com>
---
 README.md                        | 408 +++----------------------------
 docs/get-started/install.md      |  87 +++++++
 docs/get-started/overview.md     |  84 +++++++
 docs/get-started/quickstart.md   |  22 +-
 docs/get-started/releasenotes.md |  10 +
 docs/index.md                    |  11 +
 6 files changed, 239 insertions(+), 383 deletions(-)
 create mode 100644 docs/get-started/install.md
 create mode 100644 docs/get-started/overview.md
 create mode 100644 docs/get-started/releasenotes.md

diff --git a/README.md b/README.md
index fb74c9420e8..6fa300a6d4d 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 <div align="center">
 
-Megatron-LM & Megatron Core
-===========================
+Megatron-LM and Megatron Core
+=============================
 
 <h4>GPU-optimized library for training transformer models at scale</h4>
 
@@ -11,19 +11,36 @@ Megatron-LM & Megatron Core
 
 <div align="left">
 
-## ⚡ Quick Start
+## About
 
-```bash
-# 1. Install Megatron Core with required dependencies
-pip install --no-build-isolation megatron-core[mlm,dev]
+**Megatron-Core (MCore)**: Composable library with GPU-optimized building blocks for custom training frameworks.
+You can install this library using pip or use it within the Megatron-LM GitHub repository.
 
-# 2. Clone repository for examples
-git clone https://github.com/NVIDIA/Megatron-LM.git
-cd Megatron-LM
-pip install --no-build-isolation .[mlm,dev]
-```
+**Megatron-LM**: Reference implementation that includes end-to-end examples utilizing Megatron Core.
+
+**Megatron-Bridge**: Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and example model training recipes.
+
+For more information, refer to [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
+
+## Quick Start
+
+Install Megatron Core with pip:
+
+1. Install Megatron Core with required dependencies:
+
+    ```bash
+    pip install --no-build-isolation megatron-core[mlm,dev]
+    ```
+
+2. Clone repository for examples:
+
+    ```bash
+    git clone https://github.com/NVIDIA/Megatron-LM.git
+    cd Megatron-LM
+    pip install --no-build-isolation .[mlm,dev]
+    ```
 
-**→ [Complete Installation Guide](#installation)** - Docker, pip variants (dev,lts,etc.), and system requirements
 
 # Latest News
 
@@ -44,58 +61,10 @@ pip install --no-build-isolation .[mlm,dev]
 
 </details>
 
-<details>
-<summary>Table of Contents</summary>
-
-**Getting Started**
-
-- [Quick Start](#-quick-start)
-- [Latest News](#latest-news)
-- [Megatron Overview](#megatron-overview)
-  - [Project Structure](#project-structure)
-  - [Megatron-LM: Reference Implementation](#megatron-lm-reference-implementation)
-  - [Megatron Core: Production Library](#megatron-core-production-library)
-- [Installation](#installation)
-  - [Docker (Recommended)](#-docker-recommended)
-  - [Pip Installation](#pip-installation)
-  - [System Requirements](#system-requirements)
-
-**Core Features**
-
-- [Performance Benchmarking](#performance-benchmarking)
-  - [Weak Scaling Results](#weak-scaling-results)
-  - [Strong Scaling Results](#strong-scaling-results)
-- [Ecosystem Libraries](#ecosystem-libraries)
-
-**Training**
-
-- [Training](#training)
-  - [Getting Started](#getting-started)
-  - [Data Preparation](#data-preparation)
-- [Parallelism Strategies](#parallelism-strategies)
-  - [Data Parallelism (DP)](#data-parallelism-dp)
-  - [Tensor Parallelism (TP)](#tensor-parallelism-tp)
-  - [Pipeline Parallelism (PP)](#pipeline-parallelism-pp)
-  - [Context Parallelism (CP)](#context-parallelism-cp)
-  - [Expert Parallelism (EP)](#expert-parallelism-ep)
-  - [Parallelism Selection Guide](#parallelism-selection-guide)
-- [Performance Optimizations](#performance-optimizations)
-
-**Resources**
-
-- [Examples](./examples/) - Training scripts and tutorials
-- [Documentation](https://docs.nvidia.com/Megatron-Core/) - Official docs
-- [Roadmaps](#roadmaps) - Development roadmaps and feature tracking
-- [Community & Support](#community--support) - Get help and contribute
-  - [Getting Help](#getting-help)
-  - [Contributing](#contributing)
-  - [Citation](#citation)
 
-</details>
 
-# Megatron Overview
 
-## Project Structure
+# Project Structure
 
 ```
 Megatron-LM/
@@ -120,135 +89,11 @@ Megatron-LM/
 └── docs/                        # Documentation
 ```
 
-### Megatron-LM: Reference Implementation
-
-**Reference implementation** that includes Megatron Core plus everything needed to train models.
-
-**Best for:**
-
-- **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware
-- **Research teams** exploring new architectures and training techniques
-- **Learning distributed training** concepts and best practices
-- **Quick experimentation** with proven model configurations
-
-**What you get:**
-
-- Pre-configured training scripts for GPT, LLaMA, DeepSeek, Qwen, and more.
-- End-to-end examples from data prep to evaluation
-- Research-focused tools and utilities
-
-### Megatron Core: Composable Library
-
-**Composable library** with GPU-optimized building blocks for custom training frameworks.
-
-**Best for:**
-
-- **Framework developers** building on top of modular and optimized components
-- **Research teams** needing custom training loops, optimizers, or data pipelines
-- **ML engineers** requiring fault-tolerant training pipelines
-
-**What you get:**
-
-- Composable transformer building blocks (attention, MLP, etc.)
-- Advanced parallelism strategies (TP, PP, DP, EP, CP)
-- Pipeline schedules and distributed optimizers
-- Mixed precision support (FP16, BF16, FP8)
-- GPU-optimized kernels and memory management
-- High-performance dataloaders and dataset utilities
-- Model architectures (LLaMA, Qwen, GPT, Mixtral, Mamba, etc.)
-
-## Ecosystem Libraries
-
-**Libraries used by Megatron Core:**
-
-- **[Megatron Energon](https://github.com/NVIDIA/Megatron-Energon)** 📣 **NEW!** - Multi-modal data loader (text, images, video, audio) with distributed loading and dataset blending
-- **[Transformer Engine](https://github.com/NVIDIA/TransformerEngine)** - Optimized kernels and FP8 mixed precision support
-- **[Resiliency Extension (NVRx)](https://github.com/NVIDIA/nvidia-resiliency-ext)** - Fault tolerant training with failure detection and recovery
-
-**Libraries using Megatron Core:**
-
-- **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes
-- **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
-- **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
-- **[Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/).
-
-**Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
-
-# Installation
-
-## 🐳 Docker (Recommended)
-
-We strongly recommend using the previous releases of [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) rather than the latest one for optimal compatibility with Megatron Core release and testing matrix. Our releases are always based on the previous month's NGC container, so this ensures compatibility and stability.
-
-**Note:** The NGC PyTorch container constraints the python environment globally via `PIP_CONSTRAINT`. In the following examples we will unset the variable.
-
-This container comes with all dependencies pre-installed with compatible versions and optimized configurations for NVIDIA GPUs:
-
-- PyTorch (latest stable version)
-- CUDA, cuDNN, NCCL (latest stable versions)
-- Support for FP8 on NVIDIA Hopper, Ada, and Blackwell GPUs
-- For best performance, use NVIDIA Turing GPU architecture generations and later
-
-```bash
-# Run container with mounted directories
-docker run --runtime --nvidia --gpus all -it --rm \
-  -v /path/to/megatron:/workspace/megatron \
-  -v /path/to/dataset:/workspace/dataset \
-  -v /path/to/checkpoints:/workspace/checkpoints \
-  -e PIP_CONSTRAINT= \
-  nvcr.io/nvidia/pytorch:25.04-py3
-```
-
-## Pip Installation
-
-Megatron Core offers support for two NGC PyTorch containers:
-
-- `dev`: Moving head that supports the most recent upstream dependencies
-- `lts`: Long-term support of NGC PyTorch 24.01
-
-Both containers can be combined with `mlm` which adds package dependencies for Megatron-LM on top of Megatron Core.
-
-```bash
-# Install the latest release dependencies
-pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
-pip install --no-build-isolation megatron-core[dev]
-# For running an M-LM application:
-pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
-pip install --no-build-isolation megatron-core[mlm,dev]
-```
-
-```bash
-# Install packages for LTS support NGC PyTorch 24.01
-pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
-pip install --no-build-isolation megatron-core[lts]
-# For running an M-LM application:
-pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
-pip install --no-build-isolation megatron-core[mlm,lts]
-```
-
-For a version of Megatron Core with only torch, run:
-
-```bash
-pip install megatron-core
-```
-
-## System Requirements
 
-### Hardware Requirements
-
-- **FP8 Support**: NVIDIA Hopper, Ada, Blackwell GPUs
-- **Recommended**: NVIDIA Turing architecture or later
-
-### Software Requirements
-
-- **CUDA/cuDNN/NCCL**: Latest stable versions
-- **PyTorch**: Latest stable version
-- **Transformer Engine**: Latest stable version
-- **Python**: 3.12 recommended
 
 # Performance Benchmarking
 
-For our latest performance benchmarking results, please refer to [NVIDIA NeMo Framework Performance Summary](https://docs.nvidia.com/nemo-framework/user-guide/latest/performance/performance_summary.html).
+For our latest performance benchmarking results, please refer to [NVIDIA NeMo Framework Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html).
 
 Our codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to **47% Model FLOP Utilization (MFU)** on H100 clusters.
 
@@ -281,199 +126,12 @@ We also strong scaled the standard GPT-3 model (our version has slightly more th
 
 ![Strong scaling](images/strong_scaling.png)
 
-# Training
-
-## Getting Started
-
-### Simple Training Example
-
-```bash
-# Distributed training example (2 GPUs, mock data)
-torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py
-```
-
-### LLaMA-3 Training Example
-
-```bash
-# 8 GPUs, FP8 precision, mock data
-./examples/llama/train_llama3_8b_fp8.sh
-```
-
-## Data Preparation
 
-### JSONL Data Format
 
-```json
-{"text": "Your training text here..."}
-{"text": "Another training sample..."}
-```
-
-### Basic Preprocessing
-
-```bash
-python tools/preprocess_data.py \
-    --input data.jsonl \
-    --output-prefix processed_data \
-    --tokenizer-type HuggingFaceTokenizer \
-    --tokenizer-model /path/to/tokenizer.model \
-    --workers 8 \
-    --append-eod
-```
-
-### Key Arguments
-
-- `--input`: Path to input JSON/JSONL file
-- `--output-prefix`: Prefix for output binary files (.bin and .idx)
-- `--tokenizer-type`: Tokenizer type (`HuggingFaceTokenizer`, `GPT2BPETokenizer`, etc.)
-- `--tokenizer-model`: Path to tokenizer model file
-- `--workers`: Number of parallel workers for processing
-- `--append-eod`: Add end-of-document token
-
-<!-- **→ [Complete Data Preparation Guide](./docs/data-preparation.md)** - Comprehensive guide covering advanced preprocessing, dataset collection, deduplication, and optimization strategies -->
-
-# Parallelism Strategies
-
-## Data Parallelism (DP)
-
-### Standard Data Parallel
-
-```bash
-# Standard DDP - replicate model on each GPU
-torchrun --nproc_per_node=8 pretrain_gpt.py \
-    --data-parallel-sharding-strategy no_shard
-```
-
-### Fully Sharded Data Parallel (FSDP)
-
-```bash
-# Megatron's optimized FSDP (~15% faster than PyTorch FSDP2)
---use-custom-fsdp
-
-# PyTorch FSDP2
---use-torch-fsdp2
-
-# Sharding strategies
---data-parallel-sharding-strategy optim              # Shard optimizer states (ZeRO-1)
---data-parallel-sharding-strategy optim_grads        # Shard gradients + optimizer (ZeRO-2)
---data-parallel-sharding-strategy optim_grads_params # Shard parameters + gradients + optimizer (ZeRO-3)
-```
-
-## Tensor Parallelism (TP)
-
-Split individual model layers across GPUs:
-
-```bash
---tensor-model-parallel-size 4  # 4-way tensor parallelism
---sequence-parallel             # Enable sequence parallelism (recommended with TP)
-```
-
-## Pipeline Parallelism (PP)
-
-Split model depth across GPUs:
-
-```bash
---pipeline-model-parallel-size 8     # 8 pipeline stages
---virtual-pipeline-model-parallel-size 4  # Virtual pipeline for better load balancing
-```
-
-## Context Parallelism (CP)
-
-Split long sequences across GPUs for handling long contexts:
-
-```bash
---context-parallel-size 2                    # 2-way context parallelism
---cp-comm-type p2p                          # Communication: p2p, a2a, allgather, a2a+p2p
---hierarchical-context-parallel-sizes 2 4   # Hierarchical context parallelism
-```
-
-## Expert Parallelism (EP)
 
-For Mixture of Experts (MoE) models:
 
-```bash
---expert-model-parallel-size 4  # 4-way expert parallelism
---num-experts 8                 # 8 experts per MoE layer
---moe-grouped-gemm              # Optimize expert computation
-```
-
-## Combining Parallelism Strategies
-
-### Parallelism Selection Guide
-
-Based on [NVIDIA NeMo production configurations](https://github.com/NVIDIA/NeMo/tree/main/scripts/performance/recommended_model_configs):
-
-| Model | Size | GPUs | TP | PP | CP | EP | Notes |
-|-------|------|------|----|----|----|----|-------|
-| **LLaMA-3** | 8B | 8 | 1 | 1 | 2 | 1 | CP for long seqlen (8K) |
-| **LLaMA-3** | 70B | 64 | 4 | 4 | 2 | 1 | TP+PP |
-| **LLaMA-3.1** | 405B | 1024 | 8 | 8 | 2 | 1 | 3D parallelism for scale |
-| **GPT-3** | 175B | 128-512 | 4 | 8 | 1 | 1 | Large model config |
-| **Mixtral** | 8x7B | 64 | 1 | 4 | 1 | 8 | EP for MoE |
-| **Mixtral** | 8x22B | 256 | 4 | 4 | 8 | 8 | Combined TP+EP for large MoE |
-| **DeepSeek-V3** | 671B | 1024 | 2 | 16 | 1 | 64 | Large MoE config |
-
-### MoE-Specific Requirements
-
-**Important**: When combining Expert Parallelism (EP) with Tensor Parallelism (TP), **Sequence Parallelism (SP) must be enabled**.
-
-## Performance Optimizations
-
-| Feature | Flag | Benefit |
-|---------|------|---------|
-| **FlashAttention** | `--attention-backend` | Faster attention and lower memory usage |
-| **FP8 Training** | `--fp8-hybrid` | Faster training |
-| **Activation Checkpointing** | `--recompute-activations` | Reduced memory usage |
-| **Data Parallelism Communication Overlap** | `--overlap-grad-reduce` | Faster distributed training |
-| **Distributed Optimizer** | `--use-distributed-optimizer` | Reduced checkpointing time |
-
-**→ [NVIDIA NeMo Framework Performance Tuning Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/performance/performance-guide.html#performance-tuning-guide)** - Comprehensive performance optimization guide covering advanced tuning techniques, communication overlaps, memory optimizations, and profiling options.
-
-### FlashAttention
-
-[FlashAttention](https://github.com/Dao-AILab/flash-attention) is a fast and memory-efficient attention algorithm. We recommend the default usage, which uses cuDNN for attention via Transformer Engine and provides up to 50% speedups on forward and 84% on backward propagation with FP8 kernels. The `flash-attn` package is also supported via `--use-flash-attn`.
-
-### Mixed Precision Training
-
-```bash
---fp16                    # Standard FP16
---bf16                    # BFloat16 (recommended for large models)
---fp8-hybrid              # FP8 training (Hopper, Ada, and Blackwell GPUs)
-```
 
-### Activation Checkpointing and Recomputation
-
-```bash
-# For limited memory
---recompute-activations
-
-# For extreme memory constraints
---recompute-granularity full \
---recompute-method uniform
-```
-
-### Data Parallelism Communication Overlap
-
-```bash
---overlap-grad-reduce
---overlap-param-gather
-```
-
-### Distributed Optimizer
-
-```bash
---use-distributed-optimizer
-```
-
-# Roadmaps
-
-Stay up-to-date with our development roadmaps and planned features:
-
-- **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive MoE feature development including DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements
-- **[GPT-OSS Implementation Tracker](https://github.com/NVIDIA/Megatron-LM/issues/1739)** - Advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions
-
-*More roadmap trackers will be added soon.*
-
-# Community & Support
+# Resources
 
 ## Getting Help
 
@@ -493,6 +151,8 @@ We ❤️ contributions! Ways to contribute:
 
 ## Citation
 
+If you use Megatron in your research or project, we appreciate that you use the following citations:
+
 ```bibtex
 @article{megatron-lm,
   title={Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},
diff --git a/docs/get-started/install.md b/docs/get-started/install.md
new file mode 100644
index 00000000000..dd000500f58
--- /dev/null
+++ b/docs/get-started/install.md
@@ -0,0 +1,87 @@
+# Megatron Core Installation
+
+Installation is supported using Docker and pip.
+
+## System Requirements
+
+### Hardware Requirements
+
+- **FP8 Support**: NVIDIA Hopper, Ada, Blackwell GPUs
+- **Recommended**: NVIDIA Turing architecture or later
+
+### Software Requirements
+
+- **CUDA/cuDNN/NCCL**: Latest stable versions
+- **PyTorch**: Latest stable version
+- **Transformer Engine**: Latest stable version
+- **Python**: 3.12 recommended
+
+
+## Docker Installation (Recommended)
+
+We strongly recommend using the previous releases of [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) rather than the latest one for optimal compatibility with Megatron Core release and testing matrix. Our releases are always based on the previous month's NGC container, so this ensures compatibility and stability.
+
+**Note:** The NGC PyTorch container constraints the python environment globally via `PIP_CONSTRAINT`. In the following examples we will unset the variable.
+
+This container comes with all dependencies pre-installed with compatible versions and optimized configurations for NVIDIA GPUs:
+
+- PyTorch (latest stable version)
+- CUDA, cuDNN, NCCL (latest stable versions)
+- Support for FP8 on NVIDIA Hopper, Ada, and Blackwell GPUs
+- For best performance, use NVIDIA Turing GPU architecture generations and later
+
+```bash
+# Run container with mounted directories
+docker run --runtime --nvidia --gpus all -it --rm \
+  -v /path/to/megatron:/workspace/megatron \
+  -v /path/to/dataset:/workspace/dataset \
+  -v /path/to/checkpoints:/workspace/checkpoints \
+  -e PIP_CONSTRAINT= \
+  nvcr.io/nvidia/pytorch:25.04-py3
+```
+
+## Pip Installation
+
+Megatron Core installation offers support for two NGC PyTorch containers:
+
+- `dev`: Moving head that supports the most recent upstream dependencies
+- `lts`: Long-term support of NGC PyTorch 24.01
+
+Both containers can be combined with `mlm`, which adds package dependencies for Megatron-LM on top of Megatron Core.
+
+
+1. Install the latest release dependencies
+
+    ```bash
+    pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
+    pip install --no-build-isolation megatron-core[dev]
+    ```
+
+2. Next choose one of the following options:
+
+* For running an Megatron LM application
+
+        ```bash
+        pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
+        pip install --no-build-isolation megatron-core[mlm,dev]
+        ```
+* Install packages for LTS support NGC PyTorch 24.01
+
+        ```bash
+        pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
+        pip install --no-build-isolation megatron-core[lts]
+        ```
+
+* For running an Megatron LM application
+
+        ```bash
+        pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2"
+        pip install --no-build-isolation megatron-core[mlm,lts]
+        ```
+
+* For a version of Megatron Core with only Torch, run
+
+        ```bash
+        pip install megatron-core
+        ```
+
diff --git a/docs/get-started/overview.md b/docs/get-started/overview.md
new file mode 100644
index 00000000000..883f40e0c61
--- /dev/null
+++ b/docs/get-started/overview.md
@@ -0,0 +1,84 @@
+# Overview
+
+Megatron-Core and Megatron-LM are open-source tools that are typically used together to train LLMs at scale across GPUs. Megatron-Core expands the capability of Megatron-LM. Megatron Bridge connects Megatron-Core and Megatron-LM to other popular training models, such as Hugging Face.
+
+## Megatron Core
+
+NVIDIA Megatron Core is a library of essential building blocks for highly efficient large-scale generative AI training. It can be used to train models with unparalleled speed at scale across thousands of GPUs. It provides an extensive set of tools for multimodal and speech AI. It expands Megatron LM capabilities.
+
+Megatron-Core contains GPU-optimized techniques featuring advanced parallelism strategies, optimizations like FP8 training, and support for the latest LLM, MoE, and multimodal architectures. It abstracts these techniques into composable and modular APIs.
+
+Megatron-Core is compatible with all NVIDIA Tensor Core GPUs and popular LLM architectures such as GPT, BERT, T5, and RETRO.
+
+
+**Composable library** with GPU-optimized building blocks for custom training frameworks.
+
+**Best for:**
+
+- **Framework developers** building on top of modular and optimized components
+- **Research teams** needing custom training loops, optimizers, or data pipelines
+- **ML engineers** requiring fault-tolerant training pipelines
+
+**What you get:**
+
+- Composable transformer building blocks (attention, MLP)
+- Advanced parallelism strategies (TP, PP, DP, EP, CP)
+- Pipeline schedules and distributed optimizers
+- Mixed precision support (FP16, BF16, FP8)
+- GPU-optimized kernels and memory management
+- High-performance dataloaders and dataset utilities
+- Model architectures (LLaMA, Qwen, GPT, Mixtral, Mamba)
+
+## Megatron-LM
+
+Megatron-LM is a reference implementation, with a lightweight large-scale LLM training framework. It offers a customizable native PyTorch training loop with fewer abstraction layers. It was designed for scaling transformer models to the multi-billion and trillion-parameter regimes under realistic memory and compute constraints. **It serves as a straightforward entry point for exploring Megatron-Core.**
+
+It uses advanced parallelization techniques including model parallelism (tensor and pipeline), to allow models with billions of parameters to fit and train across large GPU clusters. It enables breakthroughs in large-scale NLP tasks. It splits model computations across many GPUs, overcoming single-GPU memory limits for training huge models, like GPT-style transformers.  
+
+
+**Reference implementation** that includes Megatron Core plus everything needed to train models.
+
+**Best for:**
+
+- **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware
+- **Research teams** exploring new architectures and training techniques
+- **Learning distributed training** concepts and best practices
+- **Quick experimentation** with proven model configurations
+
+**What you get:**
+
+- Pre-configured training scripts for GPT, LLaMA, DeepSeek, Qwen, and more.
+- End-to-end examples from data prep to evaluation
+- Research-focused tools and utilities
+
+
+
+## Megatron Bridge
+
+Megatron Bridge provides out-of-the-box bridges and training recipes for models built on top of base model architectures from Megatron Core.  
+
+Megatron Bridge provides a robust, parallelism-aware pathway to convert models and checkpoints. This bidirectional converter performs on-the-fly, model-parallel-aware, per-parameter conversion, and full in-memory loading.
+
+After training or modifying a Megatron model, you can convert it again for deployment or sharing.  
+
+[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)
+
+
+
+## Ecosystem Libraries
+
+**Libraries used by Megatron Core:**
+
+- **[Megatron Energon](https://github.com/NVIDIA/Megatron-Energon)** - Multi-modal data loader (text, images, video, audio) with distributed loading and dataset blending
+- **[Transformer Engine](https://github.com/NVIDIA/TransformerEngine)** - Optimized kernels and FP8 mixed precision support
+- **[Resiliency Extension (NVRx)](https://github.com/NVIDIA/nvidia-resiliency-ext)** - Fault tolerant training with failure detection and recovery
+
+**Libraries using Megatron Core:**
+
+- **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes
+- **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
+- **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
+- **[Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/).
+
+**Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
+
diff --git a/docs/get-started/quickstart.md b/docs/get-started/quickstart.md
index 36a923e6ad2..61868e7877c 100644
--- a/docs/get-started/quickstart.md
+++ b/docs/get-started/quickstart.md
@@ -1,18 +1,22 @@
 # Quick Start
 
-## Installation
+## Quick Installation
 
 Install Megatron Core with pip:
 
-```bash
-# 1. Install Megatron Core with required dependencies
-pip install --no-build-isolation megatron-core[mlm,dev]
+1. Install Megatron Core with required dependencies:
 
-# 2. Clone repository for examples
-git clone https://github.com/NVIDIA/Megatron-LM.git
-cd Megatron-LM
-pip install --no-build-isolation .[mlm,dev]
-```
+    ```bash
+    pip install --no-build-isolation megatron-core[mlm,dev]
+    ```
+
+2. Clone repository for examples:
+
+    ```bash
+    git clone https://github.com/NVIDIA/Megatron-LM.git
+    cd Megatron-LM
+    pip install --no-build-isolation .[mlm,dev]
+    ```
 
 That's it! You're ready to start training.
 
diff --git a/docs/get-started/releasenotes.md b/docs/get-started/releasenotes.md
new file mode 100644
index 00000000000..e2d77cf0070
--- /dev/null
+++ b/docs/get-started/releasenotes.md
@@ -0,0 +1,10 @@
+# Release Notes
+
+
+## Roadmaps
+
+Stay up-to-date with our development roadmaps and planned features:
+
+- **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive MoE feature development including DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements
+- **[GPT-OSS Implementation Tracker](https://github.com/NVIDIA/Megatron-LM/issues/1739)** - Advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions
+
diff --git a/docs/index.md b/docs/index.md
index 88760513f23..448a75e4c93 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -14,12 +14,23 @@ Megatron Core offers a flexible, reusable foundation for building large-scale tr
 * High-performance dataloaders and dataset utilities
 * Model architectures (LLaMA, Qwen, DeepSeek, GPT, Mamba, etc.)
 
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: About Megatron Core
+
+get-started/overview
+get-started/releasenotes
+```
+
 ```{toctree}
 :maxdepth: 2
 :hidden:
 :caption: Get Started
 
 get-started/quickstart
+get-started/install
 ```
 
 ```{toctree}

From 964c902b48fb3b9d092a5b70719110f042eecc31 Mon Sep 17 00:00:00 2001
From: Jianbin Chang <shjwudp@gmail.com>
Date: Wed, 28 Jan 2026 10:58:38 +0800
Subject: [PATCH 45/79] Add end-to-end tests for M-FSDP and ND-Parallel (#3031)

---
 .../fsdp/src/megatron_fsdp/megatron_fsdp.py   | 154 ++++++++----
 .../megatron_fsdp/param_and_grad_buffer.py    |   2 +-
 tests/test_utils/recipes/unit-tests.yaml      |   2 +-
 .../test_mcore_fully_sharded_data_parallel.py | 233 +++++++++++++++++-
 .../test_mfsdp_fully_shard.py                 |   0
 .../distributed/megatron_fsdp/utils.py        | 196 +++++++++++++++
 6 files changed, 532 insertions(+), 55 deletions(-)
 rename tests/unit_tests/distributed/{ => megatron_fsdp}/test_mcore_fully_sharded_data_parallel.py (66%)
 rename tests/unit_tests/distributed/{fsdp => megatron_fsdp}/test_mfsdp_fully_shard.py (100%)
 create mode 100644 tests/unit_tests/distributed/megatron_fsdp/utils.py

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
index f3708a35dd8..dd14efaccb3 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
@@ -272,8 +272,10 @@ def __init__(
             "optim",
         ]
         if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params":
-            # Default to overlapped NCCL communication when fully-sharding.
+            # Default to overlapped parameter gather when fully-sharding.
             self.ddp_config.overlap_param_gather = True
+        if self.ddp_config.data_parallel_sharding_strategy in ["optim_grads_params", "optim_grads"]:
+            # Default to overlapped gradient reduce-scatter when sharding gradients.
             self.ddp_config.overlap_grad_reduce = True
         if not self.is_delay_grad_reduce:
             # Gradient reduce-scatter must be overlapped when using sharding optimizer
@@ -493,6 +495,7 @@ def _register_fsdp_hooks(self, root_module):
         self.forward_pre_hooks = {}
         self.forward_hooks = {}
         self.backward_pre_hooks = {}
+        self.grad_acc_hooks = {}
 
         """
         An FSDP unit is a module designed to manage the lifecycle of model parameters
@@ -567,46 +570,78 @@ def _grad_acc(param):
 
         self._params_require_handle_grad = set()
 
-        def _post_backward(module, *unused):
+        def _post_backward_release_module(module, *unused):
             """
-            Deallocate the module parameters after the backward pass,
-            and reduce-scatter the gradients before the optimizer step.
+            Post-backward hook for an FSDP unit to release parameters and process
+            its gradients after the backward pass.
+
+            This hook:
+            - Validates that the module is an FSDP unit and that the data-parallel
+            sharding strategy is ``"optim_grads_params"``.
+            - Releases the module's parameters for the backward phase to free memory.
+            - Marks the module as IDLE in the training state machine.
             """
-            if isinstance(module, tuple(fsdp_unit_modules)):
-                if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params":
-                    # Deallocate the module parameters after the backward pass,
-                    # because we have our data-parallel gradients computed.
-                    release_module_parameters(module, bwd=True)
-                    module._training_state = TrainingState.IDLE
-                param_list = list(module.parameters())
-            else:
-                param_list = list(module.parameters(recurse=False))
+            assert isinstance(module, tuple(fsdp_unit_modules))
+            assert self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params"
 
-            if self.enable_fine_grained_param_gather_hook:
-                param_list = list(module.parameters(recurse=False))
+            # Release parameters for this module after backward.
+            release_module_parameters(module, bwd=True)
 
-            # If the parameter is shared, we do not accumulate gradients
-            # here, as the gradients will be accumulated in the
-            # root post-backward hook.
-            param_list = [p for p in param_list if not getattr(p, "_is_shared", False)]
+            # Transition this module back to the IDLE training state.
+            module._training_state = TrainingState.IDLE
 
-            # Write computed gradients into the allocated main gradient bucket for reduce-scatter.
+        @torch.compiler.disable
+        def _process_post_backward_gradients(param_list):
+            """
+            Process gradients for a list of parameters after the backward pass.
+
+            This helper accumulates gradients into the main_grad buffer and, when
+            appropriate, launches asynchronous reduce-scatter operations according
+            to the data-parallel sharding strategy and training phase.
+
+            Args:
+                param_list (List[torch.nn.Parameter]): Parameters whose gradients
+                    should be processed.
+
+            Behavior:
+                - Skips processing for shared parameters (those with ``_is_shared=True``),
+                since their gradients are handled by the root post-backward hook.
+                - Determines whether to reduce gradients based on:
+                    * Data-parallel sharding strategy (``"optim_grads"`` or
+                        ``"optim_grads_params"``).
+                    * Whether this is the last microbatch of the iteration.
+                    * Whether ``model_auto_sync`` is enabled.
+                - When reduction conditions are met, performs an asynchronous
+                reduce-scatter of gradients prior to the optimizer step, which
+                requires a subsequent call to ``finish_grad_sync()`` to complete.
+                - Marks parameters as processed by adding them to
+                    ``_params_require_handle_grad``.
+
+            Notes:
+                - With gradient-sharding strategies, gradient reduction occurs on
+                every backward propagation.
+                - Without gradient sharding, gradient reduction is deferred until
+                the last microbatch or when auto-sync is enabled.
+                - In hybrid FSDP configurations, an outer FSDP group gradient reduction
+                may be triggered.
+            """
+            # Filter out shared parameters whose gradients are handled by the root hook.
+            param_list = [p for p in param_list if not getattr(p, "_is_shared", False)]
             for param in param_list:
                 _grad_acc(param)
-                self._params_require_handle_grad.discard(param)
 
+            # Only reduce if gradients are sharded, or on the final microbatch, or when
+            # model_auto_sync is enabled.
             grad_reduce_every_bprop = self.ddp_config.data_parallel_sharding_strategy in [
                 "optim_grads",
                 "optim_grads_params",
             ]
-            # Only reduce if we are sharding gradients, or are on the final microbatch.
-            # If is_last_microbatch is not specified, then we should reduce gradients
-            # if model_auto_sync is enabled, otherwise wait until is_last_microbatch
-            # is actually specified by the user, context manager, or FW before reduction.
             is_last_microbatch = getattr(self, "is_last_microbatch", False)
+
             if grad_reduce_every_bprop or is_last_microbatch or self.model_auto_sync:
-                # Reduce-scatter the gradients asynchronously before the optimizer step.
-                # Requires calling finish_grad_sync() to wait for the reduce-scatter to complete.
+                # Launch asynchronous reduce-scatter of gradients before the optimizer
+                # step. This requires a later call to finish_grad_sync() to wait for
+                # completion.
                 self.grad_reduce_pipeline.reduce_gradients(
                     param_list,
                     suggested_queue_capacity=self.suggested_RS_queue_capacity,
@@ -616,6 +651,10 @@ def _post_backward(module, *unused):
                     ),
                 )
 
+            # Mark parameters as processed.
+            for param in param_list:
+                self._params_require_handle_grad.discard(param)
+
         @torch.compiler.disable
         def _pre_forward_param_unshard(
             module: nn.Module, args: Tuple[Any, ...], kwargs: Dict[str, Any]
@@ -656,8 +695,11 @@ def _register_post_backward_hook(
             kwargs: Dict[str, Any],
         ):
             """
-            Pre-forward hook utilized to attach a gradient reduction post-backward
-            hook to the module.
+            Register a post-backward hook for the given module by inserting an autograd
+            Function in front of it. Note that a post-backward hook implemented in this
+            way is not compatible with in-place modifications of the module's inputs,
+            since such operations can trigger an autograd error that
+            "the output is a view and is being modified in-place".
             """
             if not torch.is_grad_enabled():
                 # No gradients / backward pass, don't attach the post-backward hook.
@@ -678,11 +720,10 @@ def _register_post_backward_hook(
                 return args, kwargs
 
             """
-            Bootstrapped identity autograd function that attaches a post-backward
-            "hook" to the module to trigger model compute parameter deallocation
-            and gradient reduce-scatter immediately after the module backward pass
-            has completed to shard this layer's model and gradient memory after
-            the current backward pass stage is complete.
+            Identity autograd Function that attaches a post-backward "hook" to the
+            module, triggering parameter deallocation immediately after the module's
+            backward pass has completed in order to shard this layer's model memory
+            once the current backward stage is done.
             """
             inp_tensors = RegisterFSDPBackwardFunction.apply(
                 functools.partial(post_backward_hook, module), *inp_tensors
@@ -701,7 +742,10 @@ def _register_post_backward_hook(
 
         def _root_post_backward(*unused):
             # Make sure all the gradients are handled.
-            for param in self._params_require_handle_grad:
+            ordered_params = sorted(
+                list(self._params_require_handle_grad), key=lambda p: self.param_to_name[p]
+            )
+            for param in ordered_params:
                 _grad_acc(param)
 
             # Reduce the remaining gradients.
@@ -716,7 +760,7 @@ def _root_post_backward(*unused):
             is_last_microbatch = getattr(self, "is_last_microbatch", False)
             if grad_reduce_every_bprop or is_last_microbatch or self.model_auto_sync:
                 self.grad_reduce_pipeline.reduce_gradients(
-                    list(self._params_require_handle_grad),
+                    ordered_params,
                     suggested_queue_capacity=self.suggested_RS_queue_capacity,
                     outer_fsdp_group_grad_reduce=(
                         self.dist_index.use_hybrid_fsdp
@@ -883,25 +927,11 @@ def _register_pre_backward_param_unshard_hook(module):
                 create_custom_backward_hook(module, _pre_backward_param_unshard)
             )
 
-        def _register_grad_acc_and_reduce_hook(module):
-            """
-            Register the post-backward hook to deallocate model parameters and
-            reduce-scatter gradients immediately after the module backward pass
-            has completed to conserve memory for the subsequent backward pass.
-            """
-            self.forward_pre_hooks[f"module {name} register post-backward hook"] = (
-                module.register_forward_pre_hook(
-                    functools.partial(_register_post_backward_hook, _post_backward),
-                    with_kwargs=True,
-                )
-            )
-
         fsdp_modules = []
         for name, module in root_module.named_modules():
             if self.enable_fine_grained_param_gather_hook:
                 _register_pre_forward_param_unshard_hook(module)
                 _register_pre_backward_param_unshard_hook(module)
-                _register_grad_acc_and_reduce_hook(module)
 
             # Skip if the module is already registered in fsdp_modules.
             if any(is_submodule(module, fsdp_module) for fsdp_module in fsdp_modules):
@@ -933,8 +963,28 @@ def _register_grad_acc_and_reduce_hook(module):
                     module.register_forward_hook(_release_module_fp8_transpose_cache, prepend=False)
                 )
 
-            if not self.enable_fine_grained_param_gather_hook:
-                _register_grad_acc_and_reduce_hook(module)
+            # Register the post-backward hook to deallocate model parameters
+            # and reduce-scatter gradients after the backward pass.
+            if isinstance(module, tuple(fsdp_unit_modules)):
+                if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params":
+                    self.forward_pre_hooks[f"module {name} register post-backward hook"] = (
+                        module.register_forward_pre_hook(
+                            functools.partial(
+                                _register_post_backward_hook, _post_backward_release_module
+                            ),
+                            with_kwargs=True,
+                        )
+                    )
+                grad_acc_param_list = list(module.parameters())
+            else:
+                grad_acc_param_list = list(module.parameters(recurse=False))
+
+            for param in grad_acc_param_list:
+                self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = (
+                    param.register_post_accumulate_grad_hook(
+                        lambda p: _process_post_backward_gradients([p])
+                    )
+                )
 
         # Register root module pre- and post-backward hooks in cases where the
         # forward function of root module is not called, but rather the forward
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
index 66f7f1aec3b..215e7ab8776 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
@@ -3222,7 +3222,7 @@ def _bucket_group_gradient_reduce(
                     reduce_op = gradient_reduce_preprocessing(
                         bucket.data, scaling_factor, gbuf.ddp_config
                     )
-                    if not gbuf.is_data_distributed:
+                    if ddp_config.data_parallel_sharding_strategy == "no_shard":
                         # All-reduce the gradients on every rank. No scattering
                         # or sharding necessary.
                         torch.distributed.all_reduce(
diff --git a/tests/test_utils/recipes/unit-tests.yaml b/tests/test_utils/recipes/unit-tests.yaml
index 8e9421ac02e..c3527272782 100644
--- a/tests/test_utils/recipes/unit-tests.yaml
+++ b/tests/test_utils/recipes/unit-tests.yaml
@@ -136,7 +136,7 @@ products:
         scope: [unit-tests]
         n_repeat: [1]
         time_limit: [1800]
-  - test_case: [tests/unit_tests/distributed/fsdp/**/*.py]
+  - test_case: [tests/unit_tests/distributed/megatron_fsdp/**/*.py]
     products:
       - environment: [lts, dev]
         tag: [latest]
diff --git a/tests/unit_tests/distributed/test_mcore_fully_sharded_data_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py
similarity index 66%
rename from tests/unit_tests/distributed/test_mcore_fully_sharded_data_parallel.py
rename to tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py
index 3f0cce4e40b..77274ec4d50 100644
--- a/tests/unit_tests/distributed/test_mcore_fully_sharded_data_parallel.py
+++ b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py
@@ -6,6 +6,7 @@
 import torch
 from packaging import version
 from torch import testing
+from torch.testing import assert_close
 
 import megatron.core.parallel_state as mpu
 from megatron.core.distributed import DistributedDataParallelConfig
@@ -16,6 +17,12 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer import TransformerConfig
 from megatron.core.utils import is_torch_min_version
+from tests.unit_tests.distributed.megatron_fsdp.utils import (
+    make_gpt_mock_data_iterator,
+    make_moe_args_model_and_optimizer,
+    pretrain_forward_backward,
+    set_manual_seed,
+)
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -363,7 +370,7 @@ def train_step(model, optimizer, inputs):
         # And proceed one more step to check the results
         if fsdp_manual_registration:
             out1, loss1 = train_step(baseline_fsdp_model, optimizer1, input_data)
-            target_fsdp_model.manual_buffer_registration()
+            target_fsdp_model.param_and_grad_buffer.manual_buffer_registration()
             out2, loss2 = train_step(target_fsdp_model, optimizer2, input_data)
 
         testing.assert_close(out1, out2, rtol=0, atol=0)
@@ -493,3 +500,227 @@ def test_fsdp_with_hybrid_sharding(self, num_fsdp_group):
                 atol=0,
                 msg=f"Parameter gradients for {name1} and {name2} don't match",
             )
+
+
+@pytest.fixture(scope="class")
+def ref_cache():
+    """
+    Shared read/write cache for an class.
+    Keys: arbitrary strings, values: anything (tensors, dicts, etc.).
+    """
+    return {}
+
+
+class TestMegatronFSDPE2E:
+
+    @staticmethod
+    def _training_loop(seed=42, **kwargs):
+        """
+        Run a small deterministic (optional) training loop using a mocked MoE/GPT model and optimizer.
+        This helper initializes model-parallel state, creates a model and optimizer via
+        make_moe_args_model_and_optimizer, constructs a mock GPT data iterator, and runs
+        NUM_TRAINING_STEPS iterations of forward/backward/optimization. Losses from each
+        training step are collected and returned.
+        Args:
+            seed (int, optional): RNG seed for reproducibility. Default: 42.
+            **kwargs: Configuration overrides (all optional). Recognized keys:
+                - vocab_size (int): Vocabulary size for the mock model. Default: 100.
+                - seq_length (int): Sequence length used for the mock data. Default: 128.
+                - micro_batch_size (int): Per-microbatch size. Default: 2.
+                - global_batch_size (int): Global batch size across data-parallel ranks. Default: 32.
+                - train_iters (int): Number of training iterations to run. Default: 20.
+                - tensor_model_parallel_size (int): Tensor model parallel world size. Default: 1.
+                - pipeline_model_parallel_size (int): Pipeline model parallel world size. Default: 1.
+                - num_layers_per_virtual_pipeline_stage (int or None): Virtual pipeline configuration.
+                - expert_model_parallel_size (int): Expert model parallel size for MoE. Default: 1.
+                - expert_tensor_parallel_size (int): Expert tensor parallel size for MoE. Default: 1.
+                - num_distributed_optimizer_instances (int): Number of distributed optimizer instances. Default: 1.
+        Returns:
+            list: A list of length train_iters containing the per-step language-model loss values
+            (the value appended from output[-1] each iteration). Loss objects are returned as produced
+            by the training utilities (typically tensors or scalars).
+        Side effects:
+            - Calls Utils.initialize_model_parallel(...) and Utils.destroy_model_parallel().
+            - Sets global RNG state via set_manual_seed(seed).
+            - Constructs models/optimizers via make_moe_args_model_and_optimizer and a data iterator
+              via make_gpt_mock_data_iterator.
+            - Runs optimizer.zero_grad(), pretrain_forward_backward(...), and optim.step() repeatedly.
+            - Calculates the number of micro-batches per step as:
+                global_batch_size // micro_batch_size // data_parallel_world_size.
+              This requires that global_batch_size be divisible by micro_batch_size * data_parallel_world_size.
+        Raises:
+            ValueError: If batch-size arithmetic or other setup assumptions (e.g., divisibility) are violated.
+        """
+        # Configuration parameters with defaults
+        VOCAB_SIZE = kwargs.get("vocab_size", 100)
+        MAX_SEQ_LEN = kwargs.get("seq_length", 128)
+        MICRO_BATCH_SIZE = kwargs.get("micro_batch_size", 2)
+        GLOBAL_BATCH_SIZE = kwargs.get("global_batch_size", 32)
+        NUM_TRAINING_STEPS = kwargs.get("train_iters", 20)
+        TP = kwargs.get("tensor_model_parallel_size", 1)
+        PP = kwargs.get("pipeline_model_parallel_size", 1)
+        VPP = kwargs.get("num_layers_per_virtual_pipeline_stage", None)
+        EP = kwargs.get("expert_model_parallel_size", 1)
+        ETP = kwargs.get("expert_tensor_parallel_size", 1)
+        OUTER_DP = kwargs.get("num_distributed_optimizer_instances", 1)
+
+        # Initialize model parallel groups
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=TP,
+            pipeline_model_parallel_size=PP,
+            expert_model_parallel_size=EP,
+            expert_tensor_parallel_size=ETP,
+            num_distributed_optimizer_instances=OUTER_DP,
+        )
+        DP_GROUP = mpu.get_data_parallel_group()
+
+        # Set manual seed for reproducibility
+        set_manual_seed(seed)
+
+        # Create model and optimizer
+        model_chunks, optim = make_moe_args_model_and_optimizer(
+            ut_filename="test_mcore_fully_sharded_data_parallel.py",
+            micro_batch_size=MICRO_BATCH_SIZE,
+            global_batch_size=GLOBAL_BATCH_SIZE,
+            vocab_size=VOCAB_SIZE,
+            padded_vocab_size=VOCAB_SIZE,
+            seq_length=MAX_SEQ_LEN,
+            sequence_parallel=TP > 1,
+            tensor_model_parallel_size=TP,
+            pipeline_model_parallel_size=PP,
+            num_layers_per_virtual_pipeline_stage=VPP,
+            train_iters=NUM_TRAINING_STEPS,
+            **kwargs,
+        )
+
+        # Prepare data iterator
+        data_iterator = make_gpt_mock_data_iterator(
+            dp_group=DP_GROUP,
+            vocab_size=VOCAB_SIZE,
+            sequence_length=MAX_SEQ_LEN,
+            batch_size=MICRO_BATCH_SIZE,
+            num_samples=GLOBAL_BATCH_SIZE * NUM_TRAINING_STEPS,
+        )
+
+        outputs = []
+
+        # Training loop
+        for _ in range(NUM_TRAINING_STEPS):
+            optim.zero_grad()
+            output = pretrain_forward_backward(
+                model=model_chunks,
+                data_iterator=data_iterator,
+                sequence_length=MAX_SEQ_LEN,
+                micro_batch_size=MICRO_BATCH_SIZE,
+                num_micro_batches=GLOBAL_BATCH_SIZE // MICRO_BATCH_SIZE // DP_GROUP.size(),
+            )
+            optim.step()
+
+            # Collect loss
+            outputs.append(output[-1])
+
+        Utils.destroy_model_parallel()
+
+        return outputs
+
+    @pytest.mark.skipif(
+        not is_torch_min_version("2.4.0"), reason="Test needs to be updated for torch >= 2.4.0"
+    )
+    @pytest.mark.parametrize(
+        "nd_topology",
+        [
+            pytest.param({"TP": 2}, id="TP2"),
+            pytest.param({"EP": 2, "ETP": 2}, id="EP2_ETP2"),
+            pytest.param({"OUTER_DP": 2, "EP": 2}, id="OUTER_DP2_EP2"),
+        ],
+    )
+    @pytest.mark.parametrize(
+        ("fsdp_sharding_strategy", "use_double_buffer"),
+        [
+            ("optim_grads_params", False),
+            ("optim_grads_params", True),
+            ("optim_grads", False),
+            ("optim", True),
+        ],
+    )
+    def test_compatible_with_nd_parallel(
+        self, ref_cache, nd_topology, fsdp_sharding_strategy, use_double_buffer
+    ):
+        nd_topology_str = "_".join([f"{k}{v}" for k, v in nd_topology.items()])
+        if nd_topology_str not in ref_cache:
+            ref_cache[nd_topology_str] = TestMegatronFSDPE2E._training_loop(
+                use_distributed_optimizer=True
+            )
+
+        outputs = TestMegatronFSDPE2E._training_loop(
+            use_megatron_fsdp=True,
+            data_parallel_sharding_strategy=fsdp_sharding_strategy,
+            init_model_with_meta_device=True,
+            ckpt_format="fsdp_dtensor",
+            gradient_accumulation_fusion=False,
+            fsdp_double_buffer=use_double_buffer,
+        )
+        reference_outputs = ref_cache[nd_topology_str]
+
+        if torch.distributed.get_rank() == 0:
+            for step, (output, ref_output) in enumerate(zip(outputs, reference_outputs)):
+                loss = output["lm loss"]
+                ref_loss = ref_output["lm loss"]
+                assert_close(
+                    loss,
+                    ref_loss,
+                    atol=0,
+                    rtol=0.05,
+                    msg=(
+                        f"Loss mismatch at step {step}, FSDP Loss = {loss.item()}, "
+                        f"Reference Loss = {ref_loss.item()}"
+                        f", Compare = {compare_losses(loss.item(), ref_loss.item())}"
+                    ),
+                )
+
+
+def compare_losses(loss_a: float, loss_b: float, reference: str = "b"):
+    """
+    Compare two loss values with absolute and relative differences.
+
+    Parameters
+    ----------
+    loss_a : float
+        First loss value (e.g., baseline model).
+    loss_b : float
+        Second loss value (e.g., new model).
+    reference : {"a", "b"}, default "b"
+        Which loss to treat as the reference when computing the
+        relative difference. If "b", relative diff is vs loss_b;
+        if "a", vs loss_a.
+
+    Returns
+    -------
+    dict with keys:
+        "abs_diff" : float
+            |loss_a - loss_b|
+        "rel_diff" : float
+            |loss_a - loss_b| / reference_loss
+        "better" : str
+            "a" if loss_a < loss_b, "b" if loss_b < loss_a, "equal" otherwise.
+    """
+    abs_diff = abs(loss_a - loss_b)
+
+    if reference == "a":
+        ref = loss_a
+    else:
+        ref = loss_b
+
+    if ref == 0:
+        rel_diff = float("inf")  # or None, depending on your preference
+    else:
+        rel_diff = abs_diff / ref
+
+    if loss_a < loss_b:
+        better = "a"
+    elif loss_b < loss_a:
+        better = "b"
+    else:
+        better = "equal"
+
+    return {"abs_diff": abs_diff, "rel_diff": rel_diff, "better": better}
diff --git a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_fully_shard.py
similarity index 100%
rename from tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py
rename to tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_fully_shard.py
diff --git a/tests/unit_tests/distributed/megatron_fsdp/utils.py b/tests/unit_tests/distributed/megatron_fsdp/utils.py
new file mode 100644
index 00000000000..18a2da63786
--- /dev/null
+++ b/tests/unit_tests/distributed/megatron_fsdp/utils.py
@@ -0,0 +1,196 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+import sys
+from functools import partial
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+
+from gpt_builders import gpt_builder
+from megatron.core.distributed import finalize_model_grads
+from megatron.core.enums import ModelType
+from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
+from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.utils import get_attr_wrapped_model
+from megatron.training.arguments import parse_args, validate_args
+from megatron.training.global_vars import destroy_global_vars, set_global_variables
+from megatron.training.training import setup_model_and_optimizer
+from megatron.training.utils import is_first_or_last_pipeline_stage
+from model_provider import model_provider
+
+
+def pretrain_forward_backward(
+    *, model, data_iterator, sequence_length=128, micro_batch_size=2, num_micro_batches=1
+):
+    forward_backward_func = get_forward_backward_func()
+    output = forward_backward_func(
+        forward_step_func=_forward_step_func,
+        data_iterator=data_iterator,
+        model=model,
+        num_microbatches=num_micro_batches,
+        seq_length=sequence_length,
+        micro_batch_size=micro_batch_size,
+        forward_only=False,
+    )
+    return output
+
+
+def make_gpt_mock_data_iterator(
+    dp_group, num_samples=1000, vocab_size=50257, sequence_length=128, batch_size=8, seed=42
+):
+    dataset = GPTMockDataset(
+        num_samples=num_samples, sequence_length=sequence_length, vocab_size=vocab_size, seed=seed
+    )
+    sampler = DistributedSampler(dataset, num_replicas=dp_group.size(), rank=dp_group.rank())
+    dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)
+    for batch in dataloader:
+        batch["position_ids"] = torch.arange(sequence_length, dtype=torch.int64)
+        yield batch
+
+
+def make_moe_args_model_and_optimizer(ut_filename, **overrides):
+    sys.argv = [ut_filename]
+    base_args = dict(
+        num_layers=4,
+        mtp_num_layers=1,
+        hidden_size=128,
+        num_attention_heads=2,
+        max_position_embeddings=128,
+        bf16=False,
+        add_bias_linear=False,
+        swiglu=True,
+        position_embedding_type="rope",
+        rotary_percent=1.0,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        num_experts=4,
+        moe_shared_expert_intermediate_size=256,
+        moe_layer_freq=[0, 0, 1, 1],
+        moe_permute_fusion=True,
+        moe_router_fusion=True,
+        moe_router_topk=2,
+        moe_router_dtype="fp32",
+        create_attention_mask_in_dataloader=True,
+        lr=3e-5,
+        min_lr=3e-5,
+        use_distributed_optimizer=True,
+        finalize_model_grads_func=finalize_model_grads,
+    )
+
+    base_args.update(overrides)
+    args = parse_args()
+    for key, value in base_args.items():
+        setattr(args, key, value)
+
+    validate_args(args)
+
+    destroy_global_vars()
+    destroy_num_microbatches_calculator()
+    set_global_variables(args, build_tokenizer=False)
+
+    model, optimizer, _ = setup_model_and_optimizer(
+        model_provider_func=partial(model_provider, gpt_builder),
+        model_type=ModelType.encoder_or_decoder,
+    )
+    return model, optimizer
+
+
+def set_manual_seed(seed=42):
+    torch.manual_seed(seed)
+    model_parallel_cuda_manual_seed(seed)
+
+
+class GPTMockDataset(Dataset):
+    """
+    Mock dataset for torchtitan GPT training tests
+    Generates synthetic tokenized sequences on-the-fly
+    """
+
+    def __init__(
+        self,
+        num_samples=10000,
+        micro_batch_size=1,
+        sequence_length=2048,
+        vocab_size=128256,
+        seed=42,
+    ):
+        """
+        Initialize mock dataset
+
+        Args:
+            num_samples: Total number of samples
+            sequence_length: Length of each sequence
+            vocab_size: Size of vocabulary
+            seed: Random seed for reproducibility
+        """
+        self.num_samples = num_samples
+        self.micro_batch_size = micro_batch_size
+        self.sequence_length = sequence_length
+        self.vocab_size = vocab_size
+        self.seed = seed
+
+        # Set numpy seed for deterministic generation
+        np.random.seed(seed)
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        """
+        Generate a single training sample
+
+        Returns:
+            dict with 'tokens' and 'labels'
+        """
+        # Use idx as seed for reproducible but varied samples
+        rng = np.random.RandomState(self.seed + idx)
+
+        # Generate random token sequence
+        tokens = rng.randint(0, self.vocab_size, size=self.sequence_length, dtype=np.int64)
+
+        # Labels are tokens shifted by 1 (next token prediction)
+        labels = 1 + tokens
+
+        return {
+            'tokens': torch.from_numpy(tokens.copy()),
+            'labels': torch.from_numpy(labels.copy()),
+            "attention_mask": torch.ones(
+                (1, self.sequence_length, self.sequence_length), dtype=bool
+            ),
+            "loss_mask": torch.ones(self.sequence_length),
+        }
+
+
+def _forward_step_func(data_iterator, model, device="cuda"):
+
+    def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
+
+        losses = output_tensor.float()
+        loss_mask = loss_mask.view(-1).float()
+        loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
+        # If you have data parallel reduce loss across data parallel groups.
+        # If pipeline parallel, loss computation is done only in last stage.
+
+        return loss, {'lm loss': loss}
+
+    vp_stage = get_attr_wrapped_model(model, "vp_stage")
+
+    if not is_first_or_last_pipeline_stage(vp_stage):
+        tokens, labels, loss_mask, attention_mask, position_ids = None, None, None, None, None
+    else:
+        data = next(data_iterator)
+        tokens = data["tokens"].to(device, non_blocking=True)
+        labels = data["labels"].to(device, non_blocking=True)
+        loss_mask = data["loss_mask"].to(device, non_blocking=True)
+        attention_mask = (
+            None
+            if "attention_mask" not in data
+            else data["attention_mask"].to(device, non_blocking=True)
+        )
+        position_ids = data["position_ids"].to(device, non_blocking=True)
+
+    output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+
+    return output_tensor, partial(loss_func, loss_mask)

From 38cd9fcfe2ea4c28ef9b0a082160dc2046c182e0 Mon Sep 17 00:00:00 2001
From: Faradawn Yang <73060648+faradawn@users.noreply.github.com>
Date: Tue, 27 Jan 2026 19:49:20 -0800
Subject: [PATCH 46/79] Fix Multimodal Dockerfile (#3006)

Co-authored-by: Maanu Grover <109391026+maanug-nv@users.noreply.github.com>
---
 examples/multimodal/Dockerfile | 81 ++++++++++++++++++++++++----------
 examples/multimodal/README.md  |  4 ++
 2 files changed, 61 insertions(+), 24 deletions(-)

diff --git a/examples/multimodal/Dockerfile b/examples/multimodal/Dockerfile
index 7b54091ae63..ccb5741da62 100644
--- a/examples/multimodal/Dockerfile
+++ b/examples/multimodal/Dockerfile
@@ -1,26 +1,59 @@
-FROM nvcr.io/nvidia/pytorch:24.02-py3
+# Base image: NVIDIA PyTorch container with CUDA, cuDNN, NCCL, Python, and uv pre-installed
+FROM nvcr.io/nvidia/pytorch:25.11-py3
 
-RUN apt update && \
-    apt -y upgrade && \
-    apt install -y --no-install-recommends \
-        software-properties-common \
-        build-essential \
-        python3-pip \
-        python3-dev \
-        bash \
-        git \
-        vim \
-        tmux \
-        python-is-python3 \
-        default-jre
+# Install JRE for pycocoevalcap
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        default-jre && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
-RUN pip install --upgrade pip
-RUN pip install einops einops-exts sentencepiece braceexpand webdataset packaging
-RUN pip install transformers datasets accelerate timm
-RUN pip install pytest-cov pytest_mock nltk wrapt
-RUN pip install zarr "tensorstore==0.1.45"
-RUN pip install black isort click==8.0.2
-RUN pip install pycocoevalcap megatron-energon mistral-common tiktoken
-RUN pip install git+https://github.com/openai/CLIP.git
-# Use --no-deps for the following to avoid outdated and unnecessary dependencies.
-RUN pip install open_clip_torch open-flamingo[eval] --no-deps
+# Using --break-system-packages to allow system-wide installation in managed environment
+RUN uv pip install --system --no-cache --break-system-packages \
+    einops \
+    einops-exts \
+    sentencepiece \
+    braceexpand \
+    webdataset \
+    packaging \
+    transformers \
+    datasets \
+    accelerate \
+    timm \
+    pytest-cov \
+    pytest_mock \
+    nltk \
+    wrapt \
+    zarr \
+    tensorstore \
+    black \
+    isort \
+    click \
+    pycocoevalcap \
+    megatron-energon \
+    mistral-common \
+    tiktoken \
+    # Additional dependencies for megatron-core[mlm]
+    flask-restful \
+    wandb
+
+# Install CLIP from GitHub
+RUN uv pip install --system --no-cache --break-system-packages \
+    git+https://github.com/openai/CLIP.git
+
+# Install packages with --no-deps to avoid outdated and unnecessary dependencies
+RUN uv pip install --system --no-cache --break-system-packages --no-deps \
+    open_clip_torch \
+    "open-flamingo[eval]"
+
+# Copy Megatron-LM source and install megatron-core
+# This assumes the build context is the Megatron-LM root directory
+# Build with: docker build -t megatron-multimodal -f examples/multimodal/Dockerfile .
+WORKDIR /workspace/megatron-lm
+COPY . .
+
+# Install megatron-core in editable mode for development
+RUN uv pip install --system --no-cache --break-system-packages --no-build-isolation -e ".[mlm]"
+
+# Set working directory to examples for convenience
+WORKDIR /workspace/megatron-lm
diff --git a/examples/multimodal/README.md b/examples/multimodal/README.md
index a65839f8f15..e7fe2e62b8e 100644
--- a/examples/multimodal/README.md
+++ b/examples/multimodal/README.md
@@ -13,6 +13,10 @@ Multimodal support in megatron is still under active development. This example i
 ### Docker container
 
 You can build a docker container using `examples/multimodal/Dockerfile` to run this example.
+```
+# At the Megatron-LM root directory, execute the following
+docker build -t megatron-multimodal -f examples/multimodal/Dockerfile .
+```
 
 ### Language model
 

From b6b49e7e60777e9bfc0947550c967fd858e33dde Mon Sep 17 00:00:00 2001
From: Jianbin Chang <shjwudp@gmail.com>
Date: Wed, 28 Jan 2026 13:21:55 +0800
Subject: [PATCH 47/79] [M-FSDP] Fix double buffering not working with
 activation recompute (#2689)

Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>
---
 .../fsdp/src/megatron_fsdp/megatron_fsdp.py   | 41 +++++++++++++++----
 .../megatron_fsdp/param_and_grad_buffer.py    | 36 ++++++++++++++--
 2 files changed, 65 insertions(+), 12 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
index dd14efaccb3..bd13e76379e 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
@@ -510,10 +510,31 @@ def _register_fsdp_hooks(self, root_module):
         """
         fsdp_unit_modules = self.fsdp_unit_modules
 
-        def release_module_parameters(module, bwd, *unused):
+        def release_module_parameters(module, bwd, lazy=False, *unused):
+            """
+            Release the parameters of a given module after completing the forward
+            and backward passes.
+
+            Args:
+                module: The module whose parameters should be released.
+                bwd (bool): Indicates if the release is triggered during the backward pass.
+                lazy (bool, optional): Determines when the parameter buffer (bucket) is released.
+                    - If False, the buffer is released immediately.
+                    - If True, the release is deferred until just before the all-gather pipeline
+                    requests a new buffer. The delayed release is performed by invoking
+                    `recycle_unused_buckets`.
+                *unused: Placeholder for any unused arguments.
+
+            Notes:
+                - The function maps each parameter to its corresponding buffer group,
+                then releases the associated bucket through the all-gather pipeline.
+                - If `ddp_config.keep_fp8_transpose_cache` is False, it also clears
+                the FP8 transpose cache associated with the module’s parameters.
+            """
             for param in module.parameters():
                 bucket_id = self.param_and_grad_buffer.param_to_param_group[param]
-                self.all_gather_pipeline.release_bucket(bucket_id, bwd)
+                self.all_gather_pipeline.release_bucket(bucket_id, bwd, lazy=lazy)
+
             if not self.ddp_config.keep_fp8_transpose_cache:
                 release_params_fp8_transpose_cache(module.parameters())
 
@@ -850,19 +871,23 @@ def _root_pre_backward(module: nn.Module, *unused):
         def _post_forward(module: nn.Module, input: Any, output: Any):
             # When composed with module-hook-based activation recomputation, the
             # post-backward hook is responsible for resharding the module parameters
-            # after the forward pass. Skip resharding the module parameters in this case.
+            # after the forward pass. In this case, the resharding is performed lazily.
             if module._training_state == TrainingState.PRE_BACKWARD:
-                # Skip weight deallocation until the backward pass is complete
-                # during activation recomputation / gradient checkpointing.
-                return output
+                # Delay parameter resharding because this is currently running inside
+                # the activation recomputation forward. The corresponding backward
+                # pass may still need these parameters, and delaying avoids an
+                # unnecessary all-gather.
+                lazy_release = True
+            else:
+                lazy_release = False
+                module._training_state = TrainingState.IDLE
 
             assert isinstance(
                 module, tuple(fsdp_unit_modules)
             ), "_post_forward hook should only be registered on FSDP unit modules."
 
             # Release the module parameters after the forward pass to save memory.
-            release_module_parameters(module, bwd=False)
-            module._training_state = TrainingState.IDLE
+            release_module_parameters(module, bwd=False, lazy=lazy_release)
 
             return output
 
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
index 215e7ab8776..3ec117ebd9e 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
@@ -3646,14 +3646,42 @@ def wait_bucket_ready(self, bucket_id, bwd, empty_ok=False):
         mark_bucket_ready_to_use()
 
     @torch.no_grad()
-    def release_bucket(self, bucket_id, bwd):
-        """Release the bucket."""
-        # TODO(mxfp8): In some cases, there won't be ag before bwd?
-        bucket_key = self.get_bucket_key(bucket_id, bwd)
+    def release_bucket(self, bucket_id, bwd, lazy: bool = False):
+        """
+        Release the specified parameter bucket, freeing its associated buffer storage.
 
+        This function marks or frees the memory of a parameter bucket depending on
+        whether lazy release is enabled. It ensures that buckets are not released
+        while still being communicated or in use by the pipeline.
+
+        Args:
+            bucket_id (int): Identifier of the bucket to be released.
+            bwd (bool): Indicates if the release is triggered during the backward pass.
+            lazy (bool, optional): Determines when the parameter buffer (bucket) is released.
+                - If False, the buffer is released immediately.
+                - If True, the release is deferred until just before the all-gather pipeline
+                requests a new buffer. The delayed release is performed by invoking
+                `recycle_unused_buckets`.
+
+        Raises:
+            ValueError: If the specified bucket is currently in communication and
+                cannot be safely released.
+
+        Notes:
+            - Buckets marked as lazy will be released later when the pipeline determines
+            they are no longer needed.
+            - If the bucket has a transpose weight buffer (used in FP8 backward passes),
+            this buffer is freed; otherwise, the model weight buffer is released.
+        """
+        bucket_key = self.get_bucket_key(bucket_id, bwd)
         if self.bucket_status[bucket_key] == BucketStatus.EMPTY:
             return
 
+        if lazy:
+            # Mark the bucket can be released later.
+            self.bucket_can_be_released[bucket_key] = True
+            return
+
         self.wait_bucket_ready(bucket_id, bwd, empty_ok=True)
         if self.bucket_status[bucket_key] == BucketStatus.COMMUNICATING:
             raise ValueError(f"Bucket {bucket_id} is communicating.")

From d4f9347335d4a98b88c00613686b8e64f81dfd2b Mon Sep 17 00:00:00 2001
From: Maanu Grover <109391026+maanug-nv@users.noreply.github.com>
Date: Tue, 27 Jan 2026 23:13:17 -0800
Subject: [PATCH 48/79] [training migration] Add CheckpointConfig dataclass
 (#2431)

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/arguments.py       | 121 +-----------------
 megatron/training/training_config.py | 176 +++++++++++++++++++++++++++
 2 files changed, 181 insertions(+), 116 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 9dcc8f65e81..f3293080e42 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2499,135 +2499,34 @@ def _add_learning_rate_args(parser):
 
 
 def _add_checkpointing_args(parser):
-    group = parser.add_argument_group(title='checkpointing')
-
-    group.add_argument('--save', type=str, default=None,
-                       help='Output directory to save checkpoints to.')
-    group.add_argument('--save-interval', '--persistent-save-interval', type=int, default=None,
-                       help='Number of iterations between persistent checkpoint saves.')
-    group.add_argument('--save-wgrads-interval', type=int, default=None,
-                       help='Number of iterations between wgrad (main_grad) saves.')
-    group.add_argument('--save-dgrads-interval', type=int, default=None,
-                       help='Number of iterations between dgrad saves.')
-    group.add_argument('--save-retain-interval', type=int, default=None,
-                       help='Number of iterations between retained checkpoints (other'
-                       'checkpoints _except the last checkpoint_ are automatically deleted).')
+    from megatron.training.training_config import CheckpointConfig
+
+    ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_tokenizer_assets", "save_optim", "save_rng", "load_optim", "load_rng"])
+    group = ckpt_factory.build_group(parser, "checkpointing")
+
     group.add_argument('--no-save-optim', action='store_true', default=None,
                        help='Do not save current optimizer.')
     group.add_argument('--no-save-rng', action='store_true', default=None,
                        help='Do not save current rng state.')
-    group.add_argument('--load', type=str, default=None,
-                       help='Directory containing a model checkpoint.')
     group.add_argument('--no-load-optim', action='store_true', default=None,
                        help='Do not load optimizer when loading checkpoint.')
-    group.add_argument('--load-main-params-from-ckpt', action='store_true', default=None,
-                       help='Load main parameters from checkpoint directly.')
     group.add_argument('--no-load-rng', action='store_true', default=None,
                        help='Do not load rng state when loading checkpoint.')
-    group.add_argument('--no-strict-fsdp-dtensor-load', action='store_false', dest='strict_fsdp_dtensor_load',
-                       help='Do not strict loading for fsdp_dtensor checkpoint format.')
-    group.add_argument('--non-persistent-save-interval', type=int, default=None,
-                       help='Number of iterations between non-persistent saves.')
-    group.add_argument('--non-persistent-ckpt-type', type=str, default=None,
-                       choices=['global', 'local', 'in_memory', None],
-                       help='Type of non-persistent model checkpoints. '
-                           '"global" - Saved as a standard checkpoint (e.g., on Lustre) with old checkpoints being removed. '
-                           '"local" - Each rank saves a portion of the checkpoint locally (e.g., on SSD/ramdisk). '
-                           'None - No non-persistent checkpointing (default option).')
-    group.add_argument('--non-persistent-global-ckpt-dir', type=str, default=None,
-                       help='Directory containing global non-persistent model checkpoints.')
-    group.add_argument('--non-persistent-local-ckpt-dir', type=str, default=None,
-                       help='Directory containing local non-persistent model checkpoints.')
-    group.add_argument('--non-persistent-local-ckpt-algo', type=str, default='fully_parallel',
-                       choices=['fully_parallel', 'atomic'],
-                       help='Algorithm for local non-persistent checkpointing.')
-    group.add_argument('--finetune', action='store_true',
-                       help='Load model for finetuning. Do not load optimizer '
-                       'or rng state from checkpoint and set iteration to 0. '
-                       'Assumed when loading a release checkpoint.')
-    group.add_argument('--pretrained-checkpoint', type=str, default=None,
-                       help='Directory containing a pretrained model checkpoint for finetuning.')
-    group.add_argument('--ckpt-step', type=int, default=None,
-                       help='Checkpoint step to load model from.')
     group.add_argument('--no-initialization', action='store_false',
                        help='Do not perform initialization when building model, '
                        'can reduce startup time when definitely loading from a '
                        'checkpoint',
                        dest='perform_initialization')
-    group.add_argument('--use-checkpoint-args', action='store_true',
-                       help='Override model-related command-line arguments with arguments from checkpoint')
-    group.add_argument('--use-mp-args-from-checkpoint-args', action='store_true',
-                       help='Copy model parallelism command-line arguments from checkpoint')
-    group.add_argument('--no-use-tokenizer-model-from-checkpoint-args', action='store_false',
-                       dest='use_tokenizer_model_from_checkpoint_args',
-                       help='If set, do not use tokenizer model path from checkpoint')
-    group.add_argument('--exit-on-missing-checkpoint', action='store_true',
-                       help="If '--load' is set, but checkpoint is not found "
-                       "(e.g., path typo), then exit instead of random "
-                       "initialization.")
     group.add_argument('--use-dist-ckpt', action='store_true',
                        dest='use_dist_ckpt_deprecated',
                        help='Deprecated: see --ckpt-format.')
-    group.add_argument('--use-persistent-ckpt-worker', action='store_true',
-                       help='Enables a persitent checkpoint worker for async save')
 
-    group.add_argument('--auto-detect-ckpt-format', action='store_true',
-                       help='Determine if the checkpoint format is in legacy or distributed format.'
-                            ' If False, expects distributed checkpoint iff args.ckpt_format != "torch".'
-                            ' Might slow down loading a bit (double rank0 ckpt load).')
     group.add_argument('--dist-ckpt-format',
                        dest='dist_ckpt_format_deprecated',
                        help='Deprecated: see --ckpt-format.')
-    group.add_argument('--ckpt-format', default='torch_dist',
-                       choices=['torch', 'torch_dist', 'torch_dcp', 'fsdp_dtensor'],
-                       help='Checkpoint format to use. torch is the format used by torch.save/load.'
-                       ' torch_dist is a megatron built-in distributed checkpointing format.'
-                       ' torch_dcp is the torch.distributed.checkpoint format.'
-                       ' fsdp_dtensor is a torch DCP native, Megatron FSDP training-specific checkpoint format.')
-    group.add_argument('--ckpt-convert-format', default=None,
-                       choices=['torch', 'torch_dist'],
-                       help='Checkpoint format for conversion.')
-    group.add_argument('--ckpt-convert-save', default=None,
-                       help='Save directory for converted checkpoint.')
-    group.add_argument('--ckpt-convert-update-legacy-dist-opt-format', action='store_true',
-                       help='When loading a checkpoint, update the legacy format '
-                       'for the distributed optimizer, which previously used a '
-                       'merged param/grad buffer and a different bucket mapping. '
-                       'The legacy format was deprecated on Feb 13, 2024.')
     group.add_argument('--ckpt-fully-parallel-save', action='store_true',
                        dest='ckpt_fully_parallel_save_deprecated',
                        help='Deprecated: see --no-ckpt-fully-parallel-save.')
-    group.add_argument('--no-ckpt-fully-parallel-save', action='store_false',
-                       dest='ckpt_fully_parallel_save',
-                       help='Disable applying full save parallelization across DP for'
-                            ' distributed checkpoints. Depending on ckpt format'
-                            ' might decrease the number of files in the checkpoint.'
-                            ' Makes DistributedOptimizer checkpoint non-reshardable.')
-    group.add_argument('--async-save', action='store_true', default=None,
-                       help='Apply async checkpointing save. Currently works only with'
-                            '`torch_dist` distributed checkpoint format.')
-    group.add_argument('--ckpt-fully-parallel-load', action='store_true',
-                       help='Apply full load parallelization across DP for'
-                            ' distributed checkpoints.')
-    group.add_argument('--ckpt-assume-constant-structure', action='store_true',
-                       help='If the model and optimizer state dict structure is'
-                            'constant throughout a *single training job*, it allows for'
-                            'different checkpointing performance optimizations.')
-    group.add_argument('--dist-ckpt-strictness', type=str, default='assume_ok_unexpected',
-                       choices=[e.value for e in StrictHandling],
-                       help='Determine handling of key mismatch during checkpoint load.'
-                            ' Check StrictHandling docs for flags meaning.'
-                            ' NOTE: This flag controls only distributed checkpoint'
-                            ' load from storage, not loading state dict into the model.')
-    group.add_argument('--dist-ckpt-optim-fully-reshardable', action='store_true',
-                       help='Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP)'
-                            ' as opposed to plain DP reshardability.')
-    group.add_argument('--distrib-optim-fully-reshardable-mem-efficient', action='store_true',
-                       help='During distributed optimizer checkpoint save and load tries to use as'
-                            ' little memory as possible by using Gloo (instead of NCCL) and only one'
-                            ' rank for saving. Turn on only if experiencing host or device memory'
-                            ' issues. Has affect only with `--dist-ckpt-optim-fully-reshardable`'
-                            ' flag.')
     return parser
 
 
@@ -2866,16 +2765,6 @@ def _add_distributed_args(parser):
     group.add_argument('--use-tp-pp-dp-mapping', action='store_true', default=False,
                         help='If set, distributed ranks initialize order is changed '
                         'from tp-cp-ep-dp-pp to tp-cp-ep-pp-dp.')
-    group.add_argument('--replication', action='store_true', default=False,
-                       help="If set, replication of local checkpoints is enabled. "
-                       "Needs to be enabled on all ranks.")
-    group.add_argument('--replication-jump', default=None, type=int,
-                       help="Specifies `J`, the spacing between ranks storing replicas of a given rank's data. "
-                       "Replicas for rank `n` may be on ranks `n+J`, `n+2J`, ..., or `n-J`, `n-2J`, etc. "
-                       "This flag has an effect only if --replication is used. "
-                       "and must be consistent across all ranks.")
-    group.add_argument('--replication-factor', default=2, type=int,
-                       help="Number of machines storing the replica of a given rank's data.")
     group.add_argument('--fake-process-group', action='store_true', default=False,
                        help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \
                        This is quite useful for profiling memory usage of distributed training with just one GPU. \
diff --git a/megatron/training/training_config.py b/megatron/training/training_config.py
index 617c5cf5dfa..27c3f384c2f 100644
--- a/megatron/training/training_config.py
+++ b/megatron/training/training_config.py
@@ -320,3 +320,179 @@ class LoggerConfig:
 
     save_config_filepath: str | None = None
     """If set, save the task configuration (ConfigContainer) to this file."""
+
+
+@dataclass(kw_only=True)
+class CheckpointConfig:
+    """Configuration settings for model checkpointing (saving and loading)."""
+
+    save: str | None = None
+    """Output directory to save checkpoints to."""
+
+    save_interval: int | None = field(default=None, metadata={"argparse_meta": {"arg_names": ["--save-interval", "--persistent-save-interval"]}})
+    """Number of iterations between persistent checkpoint saves."""
+
+    save_wgrads_interval: int | None = None
+    """Number of iterations between wgrad (main_grad) saves."""
+
+    save_dgrads_interval: int | None = None
+    """Number of iterations between dgrad saves."""
+
+    save_retain_interval: int | None = None
+    """Number of iterations between retained checkpoints
+    (other checkpoints except the last checkpoint are automatically deleted).
+    """
+
+    most_recent_k: int | None = -1
+    """Number of latest checkpoint to be saved."""
+
+    save_optim: bool = True
+    """Do not save current optimizer."""
+
+    save_rng: bool = True
+    """Do not save current rng state."""
+
+    load: str | None = None
+    """Directory containing a model checkpoint."""
+
+    load_optim: bool = True
+    """Do not load optimizer when loading checkpoint."""
+
+    load_main_params_from_ckpt: bool = False
+    """Load main parameters from checkpoint. When loading a model from a checkpoint without loading
+    the optimizer, the model parameters are updated but for fp16 optimizer with main parameters,
+    the main parameters need to also be updated.
+    """
+
+    load_rng: bool = True
+    """Do not load rng state when loading checkpoint."""
+
+    non_persistent_save_interval: int | None = None
+    """Number of iterations between non-persistent saves."""
+
+    non_persistent_ckpt_type: Literal["global", "local", "in_memory"] | None = None
+    """Type of non-persistent model checkpoints.
+    "global" - Saved as a standard checkpoint (e.g., on Lustre) with old checkpoints being removed.
+    "local" - [TBD] Each rank saves a portion of the checkpoint locally (e.g., on SSD/ramdisk).
+    "in_memory" - [TBD] A special kind of local checkpoint that avoids serialization.
+    None - No non-persistent checkpointing (default option)."""
+
+    non_persistent_global_ckpt_dir: str | None = None
+    """Directory containing global non-persistent model checkpoints."""
+
+    non_persistent_local_ckpt_dir: str | None = None
+    """Directory containing local non-persistent model checkpoints."""
+
+    non_persistent_local_ckpt_algo: Literal["fully_parallel", "atomic"] = "fully_parallel"
+    """Algorithm for local non-persistent checkpointing."""
+
+    finetune: bool = False
+    """Load model for finetuning. Do not load optimizer or rng state from checkpoint and set iteration to 0.
+    Assumed when loading a release checkpoint."""
+
+    pretrained_checkpoint: str | None = None
+    """Directory containing a pretrained model checkpoint for finetuning."""
+
+    ckpt_step: int | None = None
+    """Checkpoint step to load model from."""
+
+    use_checkpoint_args: bool = False
+    """Override model-related command-line arguments with arguments from checkpoint"""
+
+    use_mp_args_from_checkpoint_args: bool = False
+    """Copy model parallelism command-line arguments from checkpoint"""
+
+    use_tokenizer_model_from_checkpoint_args: bool = True
+    """If set, do not use tokenizer model path from checkpoint"""
+
+    exit_on_missing_checkpoint: bool = False
+    """If 'load' is set, but checkpoint is not found (e.g., path typo), then exit instead of random initialization."""
+
+    ckpt_format: Literal["torch", "torch_dist", "torch_dcp", "fsdp_dtensor"] = "torch_dist"
+    """ Checkpoint format to use. torch is the format used by torch.save/load.
+    torch_dist is a megatron built-in distributed checkpointing format.
+    torch_dcp is the torch.distributed.checkpoint format.
+    fsdp_dtensor is a torch DCP native, Megatron FSDP training-specific checkpoint format.
+    """
+
+    auto_detect_ckpt_format: bool = False
+    """Determine if the checkpoint format is in legacy or distributed format. If False,
+    expects distributed checkpoint iff args.ckpt_format != "torch". Might slow down 
+    loading a bit (double rank0 ckpt load).
+    """
+
+    ckpt_convert_format: Literal["torch", "torch_dist"] | None = None
+    """Checkpoint format for conversion."""
+
+    ckpt_convert_save: str | None = None
+    """Save directory for converted checkpoint."""
+
+    ckpt_convert_update_legacy_dist_opt_format: bool = False
+    """When loading a checkpoint, update the legacy format for the distributed optimizer,
+    which previously used a merged param/grad buffer and a different bucket mapping.
+    The legacy format was deprecated on Feb 13, 2024.
+    """
+
+    ckpt_fully_parallel_save: bool = True
+    """Disable applying full save parallelization across DP for distributed checkpoints.
+    Depending on ckpt format might decrease the number of files in the checkpoint.
+    Makes DistributedOptimizer checkpoint non-reshardable."""
+
+    async_save: bool = False
+    """Apply async checkpointing save. Currently works only with `torch_dist` distributed checkpoint format."""
+
+    use_persistent_ckpt_worker: bool = False
+    """Use a persistent background worker for async checkpoint saves. When enabled, creates a dedicated
+    worker thread/process for handling async saves. When disabled, uses temporal workers that are
+    created and destroyed for each save operation."""
+
+    ckpt_fully_parallel_load: bool = False
+    """Apply full load parallelization across DP for distributed checkpoints."""
+
+    ckpt_assume_constant_structure: bool = False
+    """Assume the checkpoint structure is constant across saves to enable optimizations."""
+
+    strict_fsdp_dtensor_load: bool = True
+    """Whether to enforce strict loading for FSDP DTensor checkpoints. When False, allows partial loading."""
+
+    dist_ckpt_strictness: Literal[
+        "assume_ok_unexpected",
+        "log_unexpected",
+        "log_all",
+        "raise_unexpected",
+        "raise_all",
+        "return_unexpected",
+        "return_all",
+        "ignore_all",
+    ] = "assume_ok_unexpected"
+    """Determine handling of key mismatch during checkpoint load. Check StrictHandling docs for flags meaning.
+    NOTE: This flag controls only distributed checkpoint load from storage, not loading state dict into the model."""
+
+    dist_ckpt_save_pre_mcore_014: bool = False
+    """Revert checkpointing simplifications introduced in Megatron-Core v0.14.
+    This option affects only checkpoint saving format and will be removed soon
+    (checkpoint load format is determined based on checkpoint metadata)."""
+
+    dist_ckpt_optim_fully_reshardable: bool = False
+    """Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP) as opposed to plain DP reshardability."""
+
+    distrib_optim_fully_reshardable_mem_efficient: bool = False
+    """During distributed optimizer checkpoint save and load tries to use as little memory as possible
+    by using Gloo (instead of NCCL) and only one rank for saving. Turn on only if experiencing host or device memory
+    issues. Has affect only with `dist_ckpt_optim_fully_reshardable` flag."""
+
+    save_tokenizer_assets: bool = True
+    """Save tokenizer files to checkpoint directory. When enabled, saves all tokenizer artifacts
+    (vocab files, special tokens, tokenizer config) to make checkpoints self-contained and portable.
+    Set to False for performance-sensitive scenarios where tokenizer files are not needed."""
+
+    replication: bool = False
+    """If set, replication of local checkpoints is enabled. Needs to be enabled on all ranks."""
+
+    replication_jump: int | None = None
+    """Specifies `J`, the spacing between ranks storing replicas of a given rank's data. Replicas
+    for rank `n` may be on ranks `n+J`, `n+2J`, ..., or `n-J`, `n-2J`, etc. This flag has an
+    effect only if --replication is used. and must be consistent across all ranks."""
+
+    replication_factor: int = 2
+    """Number of machines storing the replica of a given rank's data."""

From d5cac8065642a5ce71b6050aced757a8f27638ec Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Wed, 28 Jan 2026 09:11:16 +0000
Subject: [PATCH 49/79] chore: rotate oncall schedule

---
 .github/oncall_schedule.json | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/oncall_schedule.json b/.github/oncall_schedule.json
index 5a9f35f5b5a..5fa49e966bc 100644
--- a/.github/oncall_schedule.json
+++ b/.github/oncall_schedule.json
@@ -1,8 +1,4 @@
 [
-    {
-        "user": "maanug-nv",
-        "date": "2026-01-21"
-    },
     {
         "user": "dimapihtar",
         "date": "2026-01-28"
@@ -46,5 +42,9 @@
     {
         "user": "maanug-nv",
         "date": "2026-04-08"
+    },
+    {
+        "user": "BoxiangW",
+        "date": "2026-04-15"
     }
 ]

From 008926ac9b15de1946952743d4365ab9104babfb Mon Sep 17 00:00:00 2001
From: Maanu Grover <109391026+maanug-nv@users.noreply.github.com>
Date: Wed, 28 Jan 2026 00:46:52 -0800
Subject: [PATCH 50/79] [training migration] Add StragglerDetectionConfig
 dataclass (#2435)

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/arguments.py         | 14 +++++---------
 megatron/training/resilience_config.py | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index f3293080e42..8827a7bdf55 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1893,15 +1893,11 @@ def _add_network_size_args(parser):
 
 
 def _add_straggler_detector_args(parser):
-    group = parser.add_argument_group(title='straggler')
-    group.add_argument('--log-straggler', action='store_true',
-                       help='If set, tracks and logs straggler per GPU.')
-    group.add_argument('--disable-straggler-on-startup', action='store_true',
-                       help='If set, StragglerDetector is disabled on startup.')
-    group.add_argument('--straggler-ctrlr-port', type=int, default=65535,
-                       help='Port number to toggle StragglerDetector on/off at runtime')
-    group.add_argument('--straggler-minmax-count', type=int, default=1,
-                       help='Number of ranks to report with high/low estimated throughput')
+    from megatron.training.resilience_config import StragglerDetectionConfig
+
+    straggler_factory = ArgumentGroupFactory(StragglerDetectionConfig)
+    group = straggler_factory.build_group(parser, "straggler")
+
     return parser
 
 def _add_workload_inspector_server_args(parser):
diff --git a/megatron/training/resilience_config.py b/megatron/training/resilience_config.py
index 13929c25660..dd0bd716521 100644
--- a/megatron/training/resilience_config.py
+++ b/megatron/training/resilience_config.py
@@ -22,3 +22,21 @@ class RerunStateMachineConfig:
 
     check_for_spiky_loss: bool = False
     """Check for spiky loss."""
+
+
+@dataclass(kw_only=True)
+class StragglerDetectionConfig:
+    """Configuration settings for detecting and logging GPU stragglers."""
+
+    log_straggler: bool = False
+    """If set, tracks and logs straggler per GPU."""
+
+    straggler_ctrlr_port: int = 65535
+    """Port number to toggle StragglerDetector on/off at runtime"""
+
+    straggler_minmax_count: int = 1
+    """Number of ranks to report with high/low estimated throughput"""
+
+    disable_straggler_on_startup: bool = False
+    """If set, StragglerDetector is disabled on startup."""
+

From 1453f9442e5aad7ceb04af08039c7aefbee5fed3 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com>
Date: Wed, 28 Jan 2026 03:31:30 -0800
Subject: [PATCH 51/79] Standardize RL unit tests (#3088)

---
 megatron/rl/rl_utils.py                       |   7 +-
 tests/unit_tests/rl/test_rl_utils.py          | 418 ++++++++++++++++
 .../{ => rl}/test_sequence_packing_utils.py   |   0
 tests/unit_tests/test_rl_utils.py             | 460 ------------------
 4 files changed, 422 insertions(+), 463 deletions(-)
 create mode 100644 tests/unit_tests/rl/test_rl_utils.py
 rename tests/unit_tests/{ => rl}/test_sequence_packing_utils.py (100%)
 delete mode 100644 tests/unit_tests/test_rl_utils.py

diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index 25e63408f48..5d6b3b77653 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -916,9 +916,10 @@ def prepare_trajectories(
             env_id = rollout.env_id
             env_id_counts[env_id] += 1
 
-    logger.info(f"[{dist.get_rank()}] Rollout counts:")
-    for env_id, count in env_id_counts.items():
-        logger.info(f"[{dist.get_rank()}] \t{env_id}: {count}")
+    if torch.distributed.is_initialized():
+        logger.info(f"[{dist.get_rank()}] Rollout counts:")
+        for env_id, count in env_id_counts.items():
+            logger.info(f"[{dist.get_rank()}] \t{env_id}: {count}")
 
     generation_masks = torch.tensor(generation_masks, dtype=torch.bool, device='cpu')
     trajs = torch.tensor(trajs, device='cpu')
diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py
new file mode 100644
index 00000000000..d3570bee108
--- /dev/null
+++ b/tests/unit_tests/rl/test_rl_utils.py
@@ -0,0 +1,418 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import itertools
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from megatron.core.enums import ModelType
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
+from megatron.core.pipeline_parallel.utils import is_pp_last_stage
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer import TransformerConfig
+from megatron.rl import rl_utils
+from megatron.rl.agent.api import TokenRollout
+from megatron.training.arguments import parse_args, validate_args
+from megatron.training.global_vars import destroy_global_vars, set_global_variables
+from tests.unit_tests.test_utilities import Utils
+
+BATCH = 2
+SEQ = 4
+VOCAB = 754
+
+
+class MockModel(LanguageModule):
+    def __init__(self, batch=BATCH, seq=SEQ, vocab=VOCAB):
+        self.batch = batch
+        self.seq = seq
+        self.vocab = vocab
+        self.pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        self.config = TransformerConfig(
+            num_attention_heads=8, num_layers=8, pipeline_dtype=torch.bfloat16
+        )
+        self.model_type = ModelType.encoder_or_decoder
+
+    def __call__(self, x, position_ids, attention_mask, **kwargs):
+        del position_ids
+        del attention_mask
+        batch, seq = x.shape
+        mock_model_outputs = torch.ones((batch, seq, self.vocab), device=x.device)
+        return mock_model_outputs
+
+    def load_state_dict(self, params):
+        del params
+
+    def train(self, mode=True):
+        del mode
+
+    def state_dict(self):
+        return {}
+
+    def set_input_tensor(self, input_tensor):
+        pass
+
+
+class MockTokenizer:
+    def __init__(self):
+        self.pad = 42
+        self.eod = 43
+        self.vocab_size = VOCAB
+        self.bos = None
+
+    def detokenize(self, tokens):
+        return [str(tok) for tok in tokens]
+
+
+@pytest.fixture
+def initialize_model_parallel(request, monkeypatch):
+    """Fixture to initialize and destroy model parallel.
+
+    Parameters are passed via request.param as a tuple: (tp, pp)
+    Skips if world_size < tp * pp.
+    """
+    monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
+
+    tp, pp = request.param
+    world_size = Utils.world_size
+    Utils.initialize_model_parallel(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
+    dp = world_size // (tp * pp)
+    yield world_size, dp, tp, pp
+    Utils.destroy_model_parallel()
+    destroy_global_vars()
+    destroy_num_microbatches_calculator()
+
+
+@pytest.fixture(autouse=True)
+def cleanup_global_state():
+    """Ensure global state is correctly cleaned up after every test."""
+    yield
+    destroy_global_vars()
+    destroy_num_microbatches_calculator()
+
+
+class TestRLUtils:
+    """Test class for RL utilities."""
+
+    def create_test_args(self, **kwargs):
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+
+        args = parse_args(ignore_unknown_args=True)
+        args.num_layers = 8
+        args.num_attention_heads = 8
+        args.vocab_size = VOCAB
+        args.hidden_size = 128
+        args.max_position_embeddings = 256
+        args.seq_length = 256
+
+        args.micro_batch_size = 1
+
+        for key, value in kwargs.items():
+            setattr(args, key, value)
+
+        args = validate_args(args)
+        set_global_variables(args, False)
+        return args
+
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
+            for tp, pp in itertools.product([1, 2, 4, 8], [1, 2, 4, 8])
+            if tp * pp <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    @pytest.mark.parametrize("use_sequence_packing", [False])
+    def test_get_logprobs(self, initialize_model_parallel, use_sequence_packing):
+        """Test that getting logprobs at least does not crash."""
+        self.create_test_args(rl_use_sequence_packing=use_sequence_packing)
+
+        model = MockModel()
+        tokens = torch.ones((BATCH, SEQ), dtype=torch.long)
+        logprobs = rl_utils.get_logprobs(
+            model, tokens, position_ids=None, sequence_packing=use_sequence_packing
+        )
+        if is_pp_last_stage(model.pg_collection.pp):
+            # We chop off 1 element from the sequence dimension.
+            assert logprobs.shape == (BATCH, SEQ - 1)
+            # As we return ones as logits, all logprobs should be the same.
+            assert torch.all(logprobs == logprobs[0, 0]).item()
+        else:
+            assert logprobs.shape == (BATCH, SEQ, VOCAB)
+
+    def test_grpo_loss_calculation_all_pi_eq(self):
+        # All policies are equal: clamping is inactive, ratios are ones.
+        current_logprobs = torch.ones(BATCH, SEQ)
+        old_logprobs = torch.ones(BATCH, SEQ)
+        ref_logprobs = torch.ones(BATCH, SEQ)
+        advantages = torch.zeros(BATCH)
+        loss, kl_term, ratios, entropy_term, _, _ = rl_utils.calculate_grpo_loss(
+            current_logprobs=current_logprobs,
+            old_logprobs=old_logprobs,
+            ref_logprobs=ref_logprobs,
+            advantages=advantages,
+            clamp_eps_lower=0.1,
+            clamp_eps_upper=0.1,
+            kl_beta=0.1,
+            entropy_weight=0.0,
+        )
+        torch.testing.assert_close(loss, torch.zeros_like(loss))
+        torch.testing.assert_close(kl_term, torch.zeros_like(kl_term))
+        torch.testing.assert_close(ratios, torch.ones_like(ratios))
+        torch.testing.assert_close(entropy_term, -torch.ones_like(ratios) * torch.e)
+
+    def test_grpo_loss_calculation_2x_ratios(self):
+        # All policies are equal: clamping is inactive, ratios are ones.
+        current_logprobs = torch.ones(BATCH, SEQ)
+        old_logprobs = torch.ones(BATCH, SEQ) - torch.log(torch.tensor([2.0]))
+        ref_logprobs = torch.ones(BATCH, SEQ)
+        advantages = torch.ones(BATCH)
+        loss, kl_term, ratios, _, _, _ = rl_utils.calculate_grpo_loss(
+            current_logprobs=current_logprobs,
+            old_logprobs=old_logprobs,
+            ref_logprobs=ref_logprobs,
+            advantages=advantages,
+            clamp_eps_lower=2.1,
+            clamp_eps_upper=2.1,
+            kl_beta=0.0,
+            entropy_weight=0.0,
+        )
+        # Clamping does not affect us, as 2.1 [eps] > 2 [ratio].
+        # kl_beta = 0 -> we only have the non-kl term of the loss active.
+        torch.testing.assert_close(loss, -torch.ones_like(loss) * 2)
+        # pi and pi_{ref} are the same here.
+        torch.testing.assert_close(kl_term, torch.zeros_like(kl_term))
+        # Current probs are 2x more probable than old pi.
+        torch.testing.assert_close(ratios, torch.ones_like(ratios) * 2)
+
+    def test_entropy_calculation(self):
+        # All policies are equal: clamping is inactive, ratios are ones.
+        current_logprobs = torch.ones(BATCH, SEQ)
+        old_logprobs = torch.ones(BATCH, SEQ)
+        ref_logprobs = torch.ones(BATCH, SEQ)
+        advantages = torch.zeros(BATCH)
+        loss, _, ratios, entropy_term, _, _ = rl_utils.calculate_grpo_loss(
+            current_logprobs=current_logprobs,
+            old_logprobs=old_logprobs,
+            ref_logprobs=ref_logprobs,
+            advantages=advantages,
+            clamp_eps_lower=0.1,
+            clamp_eps_upper=0.1,
+            kl_beta=0.0,
+            entropy_weight=1.0,
+        )
+        torch.testing.assert_close(loss, torch.ones_like(ratios) * torch.e)
+        torch.testing.assert_close(entropy_term, -torch.ones_like(ratios) * torch.e)
+
+    def test_grpo_loss_truncation(self):
+        # All ratios are 2
+        _, _, _, _, truncated_from_above, truncated_from_below = rl_utils.calculate_grpo_loss(
+            current_logprobs=torch.ones(BATCH, SEQ),
+            old_logprobs=0.5 * torch.ones(BATCH, SEQ),
+            ref_logprobs=torch.ones(BATCH, SEQ),
+            advantages=torch.zeros(BATCH),
+            clamp_eps_lower=0.1,
+            clamp_eps_upper=0.1,
+            kl_beta=0.1,
+            entropy_weight=0.0,
+        )
+        assert truncated_from_above.float().mean() == 1
+        assert truncated_from_below.float().sum() == 0
+
+        # All ratios are 0.01
+        _, _, _, _, truncated_from_above, truncated_from_below = rl_utils.calculate_grpo_loss(
+            current_logprobs=0.01 * torch.ones(BATCH, SEQ),
+            old_logprobs=torch.ones(BATCH, SEQ),
+            ref_logprobs=torch.ones(BATCH, SEQ),
+            advantages=torch.zeros(BATCH),
+            clamp_eps_lower=0.1,
+            clamp_eps_upper=0.1,
+            kl_beta=0.1,
+            entropy_weight=0.0,
+        )
+        assert truncated_from_above.float().sum() == 0
+        assert truncated_from_below.float().mean() == 1
+
+        # Mixed ratios: [[2., 0.5], [20., 1.]]
+        current_logprobs = torch.tensor([[1.0, 1.0], [1.0, 1.0]])
+        old_logprobs = torch.tensor([[0.5, 2.0], [0.05, 1.0]])
+        _, _, _, _, truncated_from_above, truncated_from_below = rl_utils.calculate_grpo_loss(
+            current_logprobs=current_logprobs,
+            old_logprobs=old_logprobs,
+            ref_logprobs=old_logprobs,
+            advantages=torch.zeros(BATCH),
+            clamp_eps_lower=0.1,
+            clamp_eps_upper=0.1,
+            kl_beta=0.1,
+            entropy_weight=0.0,
+        )
+        torch.testing.assert_close(
+            truncated_from_above, torch.tensor([[True, False], [True, False]])
+        )
+        torch.testing.assert_close(
+            truncated_from_below, torch.tensor([[False, True], [False, False]])
+        )
+
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
+            for tp, pp in itertools.product([1, 2, 4, 8], [1, 2, 4, 8])
+            if tp * pp <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    def test_prepare_data_for_update(self, initialize_model_parallel):
+        """Test that getting logprobs at least does not crash."""
+        world_size, dp, tp, pp = initialize_model_parallel
+        self.create_test_args(
+            micro_batch_size=2,
+            seq_length=4,
+            curr_iteration=1,
+            tensor_model_parallel_size=tp,
+            pipeline_model_parallel_size=pp,
+        )
+
+        model = MockModel()
+        tokenizer = MockTokenizer()
+
+        r1 = TokenRollout(
+            trajectory=[1, 2, 3],
+            reward=3.14,
+            generation_mask=[False, True, True],
+            logprobs=[0.1, 0.2, 0.3],
+            env_id='MEGAENV',
+            problem_id="2",
+        )
+        r2 = TokenRollout(
+            trajectory=[1, 2, 3, 4],
+            reward=0.14,
+            generation_mask=[False, True, True, True],
+            logprobs=[0.1, 0.2, 0.3, -1.2],
+            env_id='MEGAENV',
+            problem_id="2",
+        )
+        rollouts = [[r1, r2] for _ in range(dp)]
+        try:
+            rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
+        except AssertionError as e:
+            # We expect trajectories to come padded there.
+            assert str(e).startswith('Rollout is not the correct length')
+
+        r1 = TokenRollout(
+            trajectory=torch.tensor([1, 2, 3, tokenizer.eod], dtype=torch.float).cuda(),
+            reward=3.14,
+            generation_mask=torch.tensor([False, True, True, True], dtype=torch.float).cuda(),
+            logprobs=torch.tensor([-0.2, -0.3, -3.2]).cuda(),
+            env_id='MEGAENV',
+            problem_id="2",
+        )
+        r2 = TokenRollout(
+            trajectory=torch.tensor([1, 2, 234, tokenizer.eod], dtype=torch.float).cuda(),
+            reward=0.14,
+            generation_mask=torch.tensor([False, True, True, True], dtype=torch.float).cuda(),
+            logprobs=torch.tensor([-0.2, -0.3, -1.2]),
+            env_id='MEGAENV',
+            problem_id="2",
+        )
+        rollouts = [[r1, r2] for _ in range(dp)]
+        data_iter = rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
+
+        _, _, old_logprobs, _, _, _, _ = next(data_iter)
+        # All logits are ones in the MockModel.
+        # All probabilities should be uniform.
+        torch.testing.assert_close(old_logprobs.exp(), torch.ones_like(old_logprobs) / VOCAB)
+
+    @pytest.mark.parametrize("use_sequence_packing", [True, False])
+    def test_prepare_trajectories(self, use_sequence_packing):
+        """Test that rollouts are properly prepared for training."""
+        seq_length = 8
+        self.create_test_args(
+            rl_use_sequence_packing=use_sequence_packing,
+            rl_sequence_packing_bin_size=20,
+            rl_skip_bos_token=False,
+            micro_batch_size=1,
+            seq_length=seq_length,
+        )
+        tokenizer = MockTokenizer()
+
+        # Create rollouts of varying lengths
+        r1 = TokenRollout(
+            trajectory=[1, 2, 3, tokenizer.eod],
+            reward=3.14,
+            generation_mask=[False, True, True, True],
+            logprobs=[0.1, 0.2, 0.3, 0.35],
+            env_id='MEGAENV',
+            problem_id="1",
+        )
+        r2 = TokenRollout(
+            trajectory=[4, 5, 6, 7, tokenizer.eod],
+            reward=0.14,
+            generation_mask=[False, True, True, True, True],
+            logprobs=[0.4, 0.5, 0.6, 0.7, 0.75],
+            env_id='MEGAENV',
+            problem_id="2",
+        )
+        r3 = TokenRollout(
+            trajectory=[8, 9, tokenizer.eod],
+            reward=2.71,
+            generation_mask=[False, True, True],
+            logprobs=[0.8, 0.9, 0.95],
+            env_id='MEGAENV',
+            problem_id="3",
+        )
+
+        rollouts = [[r1, r2, r3]]
+
+        trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories(
+            rollouts, tokenizer, seq_length
+        )
+
+        expected_trajs = torch.tensor(
+            [
+                [1, 2, 3, tokenizer.eod] + [tokenizer.pad] * 4,
+                [4, 5, 6, 7, tokenizer.eod] + [tokenizer.pad] * 3,
+                [8, 9, tokenizer.eod] + [tokenizer.pad] * 5,
+            ],
+            dtype=torch.long,
+            device=trajs.device,
+        )
+        assert torch.equal(trajs, expected_trajs)
+
+        expected_genmask = torch.tensor(
+            [
+                [False, True, True, True] + [False] * 4,
+                [False, True, True, True, True] + [False] * 3,
+                [False, True, True] + [False] * 5,
+            ],
+            dtype=torch.bool,
+            device=genmask.device,
+        )
+        assert torch.equal(genmask, expected_genmask)
+
+        if use_sequence_packing:
+            expected_logprobs = torch.tensor(
+                [
+                    [0.1, 0.2, 0.3, 0.35] + [0.0] * 4,
+                    [0.4, 0.5, 0.6, 0.7, 0.75] + [0.0] * 3,
+                    [0.8, 0.9, 0.95] + [0.0] * 5,
+                ],
+                dtype=torch.float32,
+                device=inference_logprobs.device,
+            )
+            torch.testing.assert_close(inference_logprobs, expected_logprobs, rtol=0, atol=0)
+        else:
+            expected_logprobs = [
+                [0.1, 0.2, 0.3, 0.35],
+                [0.4, 0.5, 0.6, 0.7, 0.75],
+                [0.8, 0.9, 0.95],
+            ]
+            assert len(inference_logprobs) == len(expected_logprobs)
+            for got, exp in zip(inference_logprobs, expected_logprobs):
+                got_t = got if torch.is_tensor(got) else torch.tensor(got, dtype=torch.float32)
+                exp_t = torch.tensor(exp, dtype=torch.float32, device=got_t.device)
+                torch.testing.assert_close(got_t, exp_t, rtol=0, atol=0)
diff --git a/tests/unit_tests/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py
similarity index 100%
rename from tests/unit_tests/test_sequence_packing_utils.py
rename to tests/unit_tests/rl/test_sequence_packing_utils.py
diff --git a/tests/unit_tests/test_rl_utils.py b/tests/unit_tests/test_rl_utils.py
deleted file mode 100644
index f28240591fe..00000000000
--- a/tests/unit_tests/test_rl_utils.py
+++ /dev/null
@@ -1,460 +0,0 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
-import os
-from types import SimpleNamespace
-from unittest.mock import patch
-
-import pytest
-import torch
-
-from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
-from megatron.core.enums import ModelType
-from megatron.core.models.common.language_module.language_module import LanguageModule
-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.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer import TransformerConfig
-from megatron.core.transformer.module import Float16Module
-from megatron.rl import rl_utils, sequence_packing_utils
-from megatron.rl.agent.api import TokenRollout
-from megatron.training import arguments, global_vars
-from tests.unit_tests.test_utilities import Utils
-
-BATCH = 2
-SEQ = 4
-VOCAB = 754
-
-
-class MockModel(LanguageModule):
-    def __init__(self, batch=BATCH, seq=SEQ, vocab=VOCAB):
-        self.batch = batch
-        self.seq = seq
-        self.vocab = vocab
-        self.pg_collection = SimpleNamespace(pp=None)
-        self.config = TransformerConfig(num_attention_heads=1, num_layers=1)
-        self.model_type = ModelType.encoder_or_decoder
-
-    def __call__(self, x, position_ids, attention_mask, **kwargs):
-        del position_ids
-        del attention_mask
-        batch, seq = x.shape
-        mock_model_outputs = torch.ones((batch, seq, self.vocab), device=x.device)
-        return mock_model_outputs
-
-    def load_state_dict(self, params):
-        del params
-
-    def train(self, mode=True):
-        del mode
-
-    def state_dict(self):
-        return {}
-
-    def set_input_tensor(self, input_tensor):
-        pass
-
-
-class MockTokenizer:
-    def __init__(self):
-        self.pad = 42
-        self.eod = 43
-        self.vocab_size = VOCAB
-        self.bos = None
-
-    def detokenize(self, tokens):
-        return [str(tok) for tok in tokens]
-
-
-@pytest.fixture(scope='module', autouse=True)
-def mock_pipeline_stuff():
-    with patch('megatron.rl.rl_utils.is_pp_last_stage', return_value=True):
-        yield
-
-
-def test_get_logprobs():
-    """Test that getting logprobs at least does not crash."""
-    # We use args inside of get_logprobs, we need to initialize them.
-    args = arguments.parse_args(ignore_unknown_args=True)
-    global_vars.set_args(args)
-
-    tokens = torch.ones((BATCH, SEQ), dtype=torch.long)
-    logprobs = rl_utils.get_logprobs(MockModel(), tokens, position_ids=None)
-    # We chop off 1 element from the sequence dimension.
-    assert logprobs.shape == (BATCH, SEQ - 1)
-    # As we return ones as logits, all logprobs should be the same.
-    assert torch.all(logprobs == logprobs[0, 0]).item()
-
-
-def test_get_logprobs_with_sequence_packing():
-    """Test that getting logprobs at least does not crash."""
-    # We use args inside of get_logprobs, we need to initialize them.
-    args = arguments.parse_args(ignore_unknown_args=True)
-    setattr(args, 'rl_use_sequence_packing', True)
-    global_vars.set_args(args)
-
-    tokens = torch.ones((BATCH, SEQ), dtype=torch.long)
-    logprobs = rl_utils.get_logprobs(MockModel(), tokens, position_ids=None)
-    # We chop off 1 element from the sequence dimension.
-    assert logprobs.shape == (BATCH, SEQ - 1)
-    # As we return ones as logits, all logprobs should be the same.
-    assert torch.all(logprobs == logprobs[0, 0]).item()
-
-
-@patch('torch.distributed.get_rank', return_value=0)
-def test_prepare_trajectories(mock_rank):
-    # Make sure sequence packing is disabled for this test
-    import megatron.training.global_vars as global_vars
-
-    old_args = global_vars.get_args() if global_vars.get_args() is not None else None
-
-    # Create minimal args without sequence packing
-    args = type('Args', (), {})()
-    args.rl_use_sequence_packing = False
-    args.rl_inference_logprobs_is_correction = True
-    args.rl_skip_bos_token = False
-    global_vars.set_args(args)
-
-    tokenizer = MockTokenizer()
-    r1 = TokenRollout(
-        trajectory=[1, 2, tokenizer.eod],
-        reward=3.14,
-        generation_mask=[False, True, True],
-        logprobs=[0.1, 0.2, 0.3],
-        env_id='MEGAENV',
-        problem_id="2",
-    )
-    r2 = TokenRollout(
-        trajectory=[1, 2, tokenizer.eod],
-        reward=0.14,
-        generation_mask=[False, True, True],
-        logprobs=[0.1, 0.2, 0.3],
-        env_id='MEGAENV',
-        problem_id="2",
-    )
-    rollouts = [[r1, r2]]
-    seq_len = 7
-
-    trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories(rollouts, tokenizer, seq_len)
-
-    # Check that inference logprobs are being returned.
-    torch.testing.assert_close(inference_logprobs[0], torch.tensor([0.1, 0.2, 0.3]))
-    torch.testing.assert_close(inference_logprobs[1], torch.tensor([0.1, 0.2, 0.3]))
-
-    expected_mask = torch.tensor(
-        [
-            [False, True, True, False, False, False, False],
-            [False, True, True, False, False, False, False],
-        ]
-    )
-    torch.testing.assert_close(genmask, expected_mask)
-
-    expected_trajs = torch.tensor([[1, 2, 43, 42, 42, 42, 42], [1, 2, 43, 42, 42, 42, 42]])
-    torch.testing.assert_close(trajs, expected_trajs)
-
-
-@patch('torch.distributed.get_rank', return_value=0)
-def test_prepare_trajectories_with_packing(mock_rank):
-    """Test that rollouts data is properly prepared with sequence packing enabled."""
-    # Initialize args for sequence packing
-    args = arguments.parse_args(ignore_unknown_args=True)
-    setattr(args, 'micro_batch_size', 1)
-    setattr(args, 'global_batch_size', 1)
-    setattr(args, 'rl_use_sequence_packing', True)
-    global_vars.set_args(args)
-
-    tokenizer = MockTokenizer()
-    r1 = TokenRollout(
-        trajectory=[1, 2, tokenizer.eod],
-        reward=3.14,
-        generation_mask=[False, True, True],
-        logprobs=[0.1, 0.2, 0.3],
-        env_id='MEGAENV',
-        problem_id="2",
-    )
-    r2 = TokenRollout(
-        trajectory=[1, 2, 3, tokenizer.eod],
-        reward=0.14,
-        generation_mask=[False, True, True, True],
-        logprobs=[0.1, 0.2, 0.3, -1.2],
-        env_id='MEGAENV',
-        problem_id="2",
-    )
-    rollouts = [[r1, r2]]
-    seq_len = 7
-
-    trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories(rollouts, tokenizer, seq_len)
-
-    # With sequence packing, inference logprobs should be padded to same length
-    assert isinstance(inference_logprobs, torch.Tensor)
-    assert inference_logprobs.shape == (2, 7)  # 2 sequences, each padded to seq_len
-
-    # Check values (padded with zeros)
-    torch.testing.assert_close(
-        inference_logprobs[0], torch.tensor([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 0.0])
-    )
-    torch.testing.assert_close(
-        inference_logprobs[1], torch.tensor([0.1, 0.2, 0.3, -1.2, 0.0, 0.0, 0.0])
-    )
-
-    expected_mask = torch.tensor(
-        [
-            [False, True, True, False, False, False, False],
-            [False, True, True, True, False, False, False],
-        ]
-    )
-    torch.testing.assert_close(genmask, expected_mask)
-
-    expected_trajs = torch.tensor([[1, 2, 43, 42, 42, 42, 42], [1, 2, 3, 43, 42, 42, 42]])
-    torch.testing.assert_close(trajs, expected_trajs)
-
-
-def test_grpo_loss_calculation_all_pi_eq():
-    # All policies are equal: clamping is inactive, ratios are ones.
-    current_logprobs = torch.ones(BATCH, SEQ)
-    old_logprobs = torch.ones(BATCH, SEQ)
-    ref_logprobs = torch.ones(BATCH, SEQ)
-    advantages = torch.zeros(BATCH)
-    loss, kl_term, ratios, entropy_term, _, _ = rl_utils.calculate_grpo_loss(
-        current_logprobs=current_logprobs,
-        old_logprobs=old_logprobs,
-        ref_logprobs=ref_logprobs,
-        advantages=advantages,
-        clamp_eps_lower=0.1,
-        clamp_eps_upper=0.1,
-        kl_beta=0.1,
-        entropy_weight=0.0,
-    )
-    torch.testing.assert_close(loss, torch.zeros_like(loss))
-    torch.testing.assert_close(kl_term, torch.zeros_like(kl_term))
-    torch.testing.assert_close(ratios, torch.ones_like(ratios))
-    torch.testing.assert_close(entropy_term, -torch.ones_like(ratios) * torch.e)
-
-
-def test_grpo_loss_calculation_2x_ratios():
-    # All policies are equal: clamping is inactive, ratios are ones.
-    current_logprobs = torch.ones(BATCH, SEQ)
-    old_logprobs = torch.ones(BATCH, SEQ) - torch.log(torch.Tensor([2]))
-    ref_logprobs = torch.ones(BATCH, SEQ)
-    advantages = torch.ones(BATCH)
-    loss, kl_term, ratios, _, _, _ = rl_utils.calculate_grpo_loss(
-        current_logprobs=current_logprobs,
-        old_logprobs=old_logprobs,
-        ref_logprobs=ref_logprobs,
-        advantages=advantages,
-        clamp_eps_lower=2.1,
-        clamp_eps_upper=2.1,
-        kl_beta=0.0,
-        entropy_weight=0.0,
-    )
-    # Clamping does not affect us, as 2.1 [eps] > 2 [ratio].
-    # kl_beta = 0 -> we only have the non-kl term of the loss active.
-    torch.testing.assert_close(loss, -torch.ones_like(loss) * 2)
-    # pi and pi_{ref} are the same here.
-    torch.testing.assert_close(kl_term, torch.zeros_like(kl_term))
-    # Current probs are 2x more probable than old pi.
-    torch.testing.assert_close(ratios, torch.ones_like(ratios) * 2)
-
-
-def test_entropy_calculation():
-    # All policies are equal: clamping is inactive, ratios are ones.
-    current_logprobs = torch.ones(BATCH, SEQ)
-    old_logprobs = torch.ones(BATCH, SEQ)
-    ref_logprobs = torch.ones(BATCH, SEQ)
-    advantages = torch.zeros(BATCH)
-    loss, _, ratios, entropy_term, _, _ = rl_utils.calculate_grpo_loss(
-        current_logprobs=current_logprobs,
-        old_logprobs=old_logprobs,
-        ref_logprobs=ref_logprobs,
-        advantages=advantages,
-        clamp_eps_lower=0.1,
-        clamp_eps_upper=0.1,
-        kl_beta=0.0,
-        entropy_weight=1.0,
-    )
-    torch.testing.assert_close(loss, torch.ones_like(ratios) * torch.e)
-    torch.testing.assert_close(entropy_term, -torch.ones_like(ratios) * torch.e)
-
-
-def test_grpo_loss_truncation():
-
-    # All ratios are 2
-    _, _, _, _, truncated_from_above, truncated_from_below = rl_utils.calculate_grpo_loss(
-        current_logprobs=torch.ones(BATCH, SEQ),
-        old_logprobs=0.5 * torch.ones(BATCH, SEQ),
-        ref_logprobs=torch.ones(BATCH, SEQ),
-        advantages=torch.zeros(BATCH),
-        clamp_eps_lower=0.1,
-        clamp_eps_upper=0.1,
-        kl_beta=0.1,
-        entropy_weight=0.0,
-    )
-    assert truncated_from_above.float().mean() == 1
-    assert truncated_from_below.float().sum() == 0
-
-    # All ratios are 0.01
-    _, _, _, _, truncated_from_above, truncated_from_below = rl_utils.calculate_grpo_loss(
-        current_logprobs=0.01 * torch.ones(BATCH, SEQ),
-        old_logprobs=torch.ones(BATCH, SEQ),
-        ref_logprobs=torch.ones(BATCH, SEQ),
-        advantages=torch.zeros(BATCH),
-        clamp_eps_lower=0.1,
-        clamp_eps_upper=0.1,
-        kl_beta=0.1,
-        entropy_weight=0.0,
-    )
-    assert truncated_from_above.float().sum() == 0
-    assert truncated_from_below.float().mean() == 1
-
-    current_logprobs = torch.tensor([[1.0, 1.0], [1.0, 1.0]])
-    old_logprobs = torch.tensor([[0.5, 2.0], [0.05, 1.0]])
-    _, _, _, _, truncated_from_above, truncated_from_below = rl_utils.calculate_grpo_loss(
-        current_logprobs=current_logprobs,
-        old_logprobs=old_logprobs,
-        ref_logprobs=old_logprobs,
-        advantages=torch.zeros(BATCH),
-        clamp_eps_lower=0.1,
-        clamp_eps_upper=0.1,
-        kl_beta=0.1,
-        entropy_weight=0.0,
-    )
-    # ratios: [[2., 0.5],[20., 1.]]
-    torch.testing.assert_close(truncated_from_above, torch.tensor([[True, False], [True, False]]))
-    torch.testing.assert_close(truncated_from_below, torch.tensor([[False, True], [False, False]]))
-
-
-@pytest.mark.skipif(True, reason="broken")
-def test_prepare_data_for_update():
-    """Test that getting logprobs at least does not crash."""
-    Utils.initialize_model_parallel()
-
-    args = arguments.parse_args(ignore_unknown_args=True)
-    setattr(args, 'data_parallel_size', 1)
-    setattr(args, 'micro_batch_size', 2)
-    setattr(args, 'global_batch_size', 2)
-    setattr(args, 'seq_length', 4)
-    setattr(args, 'curr_iteration', 1)
-    global_vars.unset_global_variables()
-    global_vars.set_global_variables(args, build_tokenizer=False)
-
-    model = MockModel()
-    tokenizer = MockTokenizer()
-
-    try:
-        r1 = TokenRollout(
-            trajectory=[1, 2, 3],
-            reward=3.14,
-            generation_mask=[False, True, True],
-            logprobs=[0.1, 0.2, 0.3],
-            env_id='MEGAENV',
-            problem_id="2",
-        )
-        r2 = TokenRollout(
-            trajectory=[1, 2, 3, 4],
-            reward=0.14,
-            generation_mask=[False, True, True, True],
-            logprobs=[0.1, 0.2, 0.3, -1.2],
-            env_id='MEGAENV',
-            problem_id="2",
-        )
-        rollouts = [[r1, r2]]
-        try:
-            data_iter = rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
-        except AssertionError as e:
-            # We expect trajectories to come padded there.
-            assert str(e).startswith('Rollout is not the correct length')
-
-        r1 = TokenRollout(
-            trajectory=torch.Tensor([1, 2, 3, tokenizer.eod]).cuda(),
-            reward=3.14,
-            generation_mask=torch.Tensor([False, True, True, True]).cuda(),
-            logprobs=torch.Tensor([-0.2, -0.3, -3.2]).cuda(),
-            env_id='MEGAENV',
-            problem_id="2",
-        )
-        r2 = TokenRollout(
-            trajectory=torch.Tensor([1, 2, 234, tokenizer.eod]).cuda(),
-            reward=0.14,
-            generation_mask=torch.Tensor([False, True, True, True]).cuda(),
-            logprobs=torch.Tensor([-0.2, -0.3, -1.2]),
-            env_id='MEGAENV',
-            problem_id="2",
-        )
-        rollouts = [[r1, r2]]
-        data_iter = rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
-
-        _, _, old_logprobs, _, _, _, _ = next(data_iter)
-        # All logits are ones in the MockModel.
-        # All probabilities should be uniform.
-        torch.testing.assert_close(old_logprobs.exp(), torch.ones_like(old_logprobs) / VOCAB)
-    finally:
-        Utils.destroy_model_parallel()
-
-
-@patch('torch.distributed.get_rank', return_value=0)
-def test_prepare_trajectories_with_sequence_packing(mock_rank):
-    """Test prepare_trajectories with sequence packing enabled."""
-    # Set up args with sequence packing
-    args = arguments.parse_args(ignore_unknown_args=True)
-    setattr(args, 'rl_use_sequence_packing', True)
-    setattr(args, 'rl_sequence_packing_bin_size', 16)
-    setattr(args, 'data_parallel_size', 1)
-    setattr(args, 'micro_batch_size', 2)
-    setattr(args, 'global_batch_size', 2)
-    setattr(args, 'seq_length', 16)
-    setattr(args, 'curr_iteration', 1)
-    global_vars.unset_global_variables()
-    global_vars.set_global_variables(args, build_tokenizer=False)
-
-    tokenizer = MockTokenizer()
-
-    # Create rollouts of varying lengths
-    r1 = TokenRollout(
-        trajectory=[1, 2, tokenizer.eod],
-        reward=3.14,
-        generation_mask=[False, True, True],
-        logprobs=[0.1, 0.2, 0.3],
-        env_id='MEGAENV',
-        problem_id="1",
-    )
-    r2 = TokenRollout(
-        trajectory=[4, 5, 6, 7, tokenizer.eod],
-        reward=0.14,
-        generation_mask=[False, True, True, True, True],
-        logprobs=[0.4, 0.5, 0.6, 0.7, 0.8],
-        env_id='MEGAENV',
-        problem_id="2",
-    )
-    r3 = TokenRollout(
-        trajectory=[8, 9, tokenizer.eod],
-        reward=2.71,
-        generation_mask=[False, True, True],
-        logprobs=[0.9, 1.0, 1.1],
-        env_id='MEGAENV',
-        problem_id="3",
-    )
-
-    rollouts = [[r1, r2, r3]]
-    seq_len = 16
-
-    # Call prepare_trajectories with sequence packing
-    trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories(rollouts, tokenizer, seq_len)
-
-    # With sequence packing enabled but called from prepare_trajectories,
-    # it might still return individual sequences (not packed into bins yet)
-    # because the actual packing happens later in prepare_data_for_update
-    assert trajs.shape[0] == 3  # Three sequences
-    assert trajs.shape[1] == seq_len
-
-    # Verify that each sequence is properly padded
-    # Sequence 1: [1, 2, eod, pad] + padding
-    assert trajs[0, 0] == 1
-    assert trajs[0, 1] == 2
-    assert trajs[0, 2] == tokenizer.eod
-    assert trajs[0, 3] == tokenizer.pad
-
-    # Sequence 2: [4, 5, 6, 7, eod, pad] + padding
-    assert trajs[1, 0] == 4
-    assert trajs[1, 1] == 5
-    assert trajs[1, 4] == tokenizer.eod
-    assert trajs[1, 5] == tokenizer.pad

From 71c49b56da0c9d2d4fa32a2ac11bde5bd01261c9 Mon Sep 17 00:00:00 2001
From: HaochenYuan <106647990+HaochenYuan@users.noreply.github.com>
Date: Wed, 28 Jan 2026 21:03:20 +0800
Subject: [PATCH 52/79] Fix for PR-2142 (#3096)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
---
 .../core/extensions/transformer_engine.py     |   2 +-
 .../common/model_chunk_schedule_plan.py       |   2 +
 .../core/models/gpt/fine_grained_callables.py |  21 +-
 megatron/core/models/gpt/gpt_model.py         |  37 +++-
 megatron/core/models/mamba/mamba_model.py     |   2 +
 megatron/core/ssm/mamba_block.py              |   2 +
 megatron/core/transformer/mlp.py              |   2 +-
 megatron/core/transformer/moe/moe_layer.py    |  31 ++-
 megatron/core/transformer/moe/moe_utils.py    |  91 +++++++--
 megatron/core/transformer/moe/router.py       | 160 ++++++++++-----
 .../core/transformer/transformer_block.py     |  14 +-
 .../core/transformer/transformer_layer.py     |  49 +++--
 .../a2a_overlap/test_schedule_chunk_1f1b.py   | 116 ++++++++++-
 .../a2a_overlap/test_schedule_layer_1f1b.py   |   4 +-
 .../transformer/moe/test_aux_loss.py          | 182 ++++++++++++++++++
 .../transformer/moe/test_routers.py           |  47 +++++
 16 files changed, 664 insertions(+), 98 deletions(-)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 63694fc1172..ef8527e9e5e 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -2161,7 +2161,7 @@ def forward_post_hook(module, *_) -> None:
                     "TEFusedMLP module does not support submodules with post-backward hooks"
                 )
 
-        def forward(self, hidden_states: torch.Tensor) -> Tuple[Tensor, Optional[Tensor]]:
+        def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Optional[Tensor]]:
             """Forward."""
 
             # Construct fused impl if needed
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 71aa1ab97f0..033e8e808f9 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -281,6 +281,7 @@ def __init__(
         extra_block_kwargs=None,
         runtime_gather_output: Optional[bool] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask=None,
     ):
         """Initialize the schedule plan of all Transformer layers' sub-modules.
 
@@ -323,6 +324,7 @@ def __init__(
         self._model_chunk_state.mtp_hidden_states = None
         self._model_chunk_state.loss_mask = loss_mask
         self._model_chunk_state.packed_seq_params = packed_seq_params
+        self._model_chunk_state.padding_mask = padding_mask
         self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
         self._model_chunk_state.runtime_gather_output = runtime_gather_output
         self._model_chunk_state.model = model
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index bbeee561110..7cee9d2973c 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -131,13 +131,19 @@ def forward_impl(self):
         if not self.gpt_model.pre_process:
             self.chunk_state.decoder_input = self.gpt_model.decoder.input_tensor
         # Run GPTModel._preprocess
-        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
-            self.gpt_model._preprocess(
-                input_ids=self.chunk_state.input_ids,
-                position_ids=self.chunk_state.position_ids,
-                decoder_input=self.chunk_state.decoder_input,
-                packed_seq_params=self.chunk_state.packed_seq_params,
-            )
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = self.gpt_model._preprocess(
+            input_ids=self.chunk_state.input_ids,
+            position_ids=self.chunk_state.position_ids,
+            decoder_input=self.chunk_state.decoder_input,
+            packed_seq_params=self.chunk_state.packed_seq_params,
+            padding_mask=self.chunk_state.padding_mask,
         )
 
         # Saved for later use
@@ -146,6 +152,7 @@ def forward_impl(self):
         self.chunk_state.rotary_pos_cos = rotary_pos_cos
         self.chunk_state.rotary_pos_sin = rotary_pos_sin
         self.chunk_state.sequence_len_offset = sequence_len_offset
+        self.chunk_state.padding_mask = padding_mask
         return decoder_input
 
 
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index e70221d2cfa..e287344c13d 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -288,6 +288,7 @@ def _preprocess(
         decoder_input: Tensor = None,
         inference_context: BaseInferenceContext = None,
         packed_seq_params: PackedSeqParams = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Preprocesses inputs for the transformer decoder.
 
@@ -304,7 +305,20 @@ def _preprocess(
         if decoder_input is not None:
             pass
         elif self.pre_process:
+            if padding_mask is not None:
+                assert padding_mask.shape == input_ids.shape, (
+                    f"padding_mask shape {padding_mask.shape} does not match "
+                    f"input_ids shape {input_ids.shape}"
+                )
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+            if padding_mask is not None and self.config.sequence_parallel:
+                padding_mask = (
+                    tensor_parallel.scatter_to_sequence_parallel_region(
+                        padding_mask.transpose(0, 1).contiguous()
+                    )
+                    .transpose(0, 1)
+                    .contiguous()
+                )
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
@@ -423,6 +437,7 @@ def _preprocess(
             rotary_pos_cos,
             rotary_pos_sin,
             sequence_len_offset,
+            padding_mask,
         )
         if rotary_pos_cos_sin is not None:
             # only in the case of flashinfer fused rope will we
@@ -466,6 +481,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -476,6 +492,9 @@ def forward(
         Args:
             runtime_gather_output (bool): Gather output at runtime. Default None means
                 `parallel_output` arg in the constructor will be used.
+            padding_mask (Tensor, optional): Padding mask for MoE routing.
+                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Only used for MoE layers to exclude padding tokens from routing computations.
         """
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()
@@ -488,13 +507,19 @@ def forward(
             decoder_input=decoder_input,
             inference_context=inference_context,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
-        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset) = (
-            preproc_output[:5]
-        )
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = preproc_output[:6]
 
-        rotary_pos_cos_sin = preproc_output[5] if len(preproc_output) == 6 else None
+        rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
 
         # Run decoder.
         hidden_states = self.decoder(
@@ -507,6 +532,7 @@ def forward(
             rotary_pos_cos_sin=rotary_pos_cos_sin,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
+            padding_mask=padding_mask,
             **(extra_block_kwargs or {}),
         )
 
@@ -723,6 +749,7 @@ def build_schedule_plan(
         runtime_gather_output: Optional[bool] = None,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Builds a computation schedule plan for the model.
 
@@ -748,6 +775,7 @@ def build_schedule_plan(
             inference_params (InferenceParams, optional):
                 Parameters for inference. Defaults to None.
             loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
+            padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
 
         Returns:
             TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -769,6 +797,7 @@ def build_schedule_plan(
             extra_block_kwargs,
             runtime_gather_output,
             loss_mask,
+            padding_mask,
         )
 
     def sharded_state_dict(
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 0d71ead4b0f..8d45e1d0147 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -185,6 +185,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
+        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the Mamba model. This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -254,6 +255,7 @@ def forward(
             inference_context=inference_context,
             rotary_pos_emb=rotary_pos_emb,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
         if not self.post_process:
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 9e41aca8253..ef41faae143 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -211,6 +211,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
+        padding_mask=None,
     ):
         """
         Forward function of the MambaStack class.
@@ -293,6 +294,7 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
                         )
                     else:  # MambaLayer
                         hidden_states = layer(
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index 2eae0178eea..2bc3949a421 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -148,7 +148,7 @@ def __init__(
             tp_group=tp_group,
         )
 
-    def forward(self, hidden_states, per_token_scale=None):
+    def forward(self, hidden_states, per_token_scale=None, **kwargs):
         """Perform the forward pass through the MLP block."""
         # [s, b, 4 * h/p]
         nvtx_range_push(suffix="linear_fc1")
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index ef868ebbdb8..98d4d5fa505 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -243,13 +243,13 @@ def __init__(
         self.fwd_execution_map = ["route", "expert_compute", "postprocess"]
 
     @maybe_skip_or_early_return_by_cudagraph("route")
-    def route(self, hidden_states: torch.Tensor):
+    def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Compute token routing for preprocessing.
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
         """
-        probs, routing_map = apply_module(self.router)(hidden_states)
+        probs, routing_map = apply_module(self.router)(hidden_states, padding_mask)
         return probs, routing_map
 
     @maybe_skip_or_early_return_by_cudagraph("preprocess")
@@ -354,7 +354,12 @@ def router_and_preprocess(self, hidden_states: torch.Tensor):
         hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map)
         return hidden_states, probs, residual
 
-    def forward(self, hidden_states: torch.Tensor, intermediate_tensors=None):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        intermediate_tensors=None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ):
         """Forward pass for the MoE layer.
 
         The forward pass comprises four main steps:
@@ -364,8 +369,10 @@ def forward(self, hidden_states: torch.Tensor, intermediate_tensors=None):
         4. Combine: The outputs from the experts are combined and returned.
 
         Args:
-            hidden_states (torch.Tensor): The input tensor to the MoE layer.
-
+            hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size].
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
         Returns:
             A tuple containing the output tensor and the MLP bias, if any.
         """
@@ -374,13 +381,16 @@ def forward(self, hidden_states: torch.Tensor, intermediate_tensors=None):
                 "During training, performance may degrade if MoE and tensor parallelism"
                 "are enabled without also enabling sequence parallelism."
             )
+        # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states
+        if padding_mask is not None:
+            padding_mask = padding_mask.transpose(0, 1).bool()
 
         # MoE forward: route -> dispatch -> compute -> combine
-        def custom_forward(hidden_states, intermediate_tensors):
+        def custom_forward(hidden_states, intermediate_tensors, padding_mask=None):
             try:
                 if "route" in self.fwd_execution_map:
                     shared_expert_output = self.shared_experts_compute(hidden_states)
-                    probs, routing_map = self.route(hidden_states)
+                    probs, routing_map = self.route(hidden_states, padding_mask)
                     hidden_states, probs = self.preprocess(hidden_states, probs, routing_map)
 
                     if intermediate_tensors is not None:
@@ -427,11 +437,14 @@ def custom_forward(hidden_states, intermediate_tensors):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     parallel_state.get_tensor_model_parallel_group(),
                     hidden_states,
+                    padding_mask,
                 )
             else:
-                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+                outputs = tensor_parallel.checkpoint(
+                    custom_forward, False, hidden_states, padding_mask
+                )
         else:
-            outputs = custom_forward(hidden_states, intermediate_tensors)
+            outputs = custom_forward(hidden_states, intermediate_tensors, padding_mask)
 
         return outputs
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index dc7450d93d0..65d8fed1015 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -12,6 +12,7 @@
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.router_replay import RouterReplay
@@ -50,6 +51,7 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
@@ -101,10 +103,19 @@ def switch_load_balancing_loss_func(
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
         fused (bool): Whether to use the fused version of the auxiliary loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
 
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
+    # Apply padding mask to probs if provided
+    if padding_mask is not None:
+        # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
+        mask_expanded = padding_mask.unsqueeze(-1)
+        probs = probs * mask_expanded
+
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -124,19 +135,35 @@ def switch_load_balancing_loss_func(
     return aux_loss
 
 
-def z_loss_func(logits: torch.Tensor, z_loss_coeff: float) -> torch.Tensor:
+def z_loss_func(
+    logits: torch.Tensor, z_loss_coeff: float, padding_mask: Optional[torch.Tensor] = None
+) -> torch.Tensor:
     """Encourages the router's logits to remain small to enhance stability.
     Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
     Args:
         logits (torch.Tensor): The logits of the router.
         z_loss_coeff (float): The coefficient for the z-loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                               Shape [num_tokens]. True = padding (exclude),
+                                               False = valid (include). Defaults to None.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
     """
-
-    z_loss = torch.mean(torch.square(torch.logsumexp(logits, dim=-1))) * z_loss_coeff
+    logsum = torch.logsumexp(logits, dim=-1)
+    z_loss_values = torch.square(logsum)
+
+    if padding_mask is not None:
+        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
+        valid_mask = ~padding_mask
+        # Only compute z_loss for valid (non-padding) tokens
+        z_loss_values = z_loss_values * valid_mask
+        # Compute mean over valid tokens only
+        num_valid_tokens = valid_mask.sum()
+        z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff
+    else:
+        z_loss = torch.mean(z_loss_values) * z_loss_coeff
     return z_loss
 
 
@@ -186,6 +213,28 @@ def get_capacity(
     return capacity
 
 
+def get_tokens_per_expert_and_token_count(
+    routing_map: torch.Tensor,
+    reduce_group: torch.distributed.ProcessGroup,
+    topk: int = None,
+    with_padding_mask: bool = False,
+) -> torch.Tensor:
+    """
+    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with padding mask.
+    """
+    local_tokens_per_expert = routing_map.sum(dim=0)
+    global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
+        local_tokens_per_expert, reduce_group
+    )
+    if with_padding_mask:
+        local_num_tokens = local_tokens_per_expert.sum() / topk
+        total_num_tokens = global_tokens_per_expert.sum() / topk
+    else:
+        local_num_tokens = routing_map.shape[0]
+        total_num_tokens = local_num_tokens * reduce_group.size()
+    return global_tokens_per_expert, local_num_tokens, total_num_tokens
+
+
 class MoEAuxLossAutoScaler(torch.autograd.Function):
     """An AutoScaler that triggers the backward pass and scales the grad for auxiliary loss."""
 
@@ -701,7 +750,11 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
 
 
 def compute_routing_scores_for_aux_loss(
-    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False
+    logits: torch.Tensor,
+    topk: int,
+    score_function: str,
+    fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute routing scores based on the score function.
 
@@ -710,6 +763,9 @@ def compute_routing_scores_for_aux_loss(
         topk (int): The number of top-k indices to compute.
         score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
         fused (bool, optional): Whether to use the fused version. Defaults to False.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: The routing map and the normalized routing scores.
@@ -719,20 +775,27 @@ def compute_routing_scores_for_aux_loss(
             raise ValueError(
                 "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
             )
-        return fused_compute_score_for_moe_aux_loss(
+        routing_map, scores = fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
-
-    if score_function == "softmax":
-        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
-    elif score_function == "sigmoid":
-        scores = torch.sigmoid(logits)
-        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
     else:
-        raise ValueError(f"Invalid score_function: {score_function}")
+        if score_function == "softmax":
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        elif score_function == "sigmoid":
+            scores = torch.sigmoid(logits)
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        else:
+            raise ValueError(f"Invalid score_function: {score_function}")
+
+        _, top_indices = torch.topk(scores, k=topk, dim=1)
+        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
 
-    _, top_indices = torch.topk(scores, k=topk, dim=1)
-    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    # Apply padding mask to scores if provided
+    if padding_mask is not None:
+        # Invert padding_mask and make True indicates valid tokens
+        valid_mask = (~padding_mask).unsqueeze(-1)
+        routing_map = routing_map * valid_mask
+        scores = scores * valid_mask
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 4e3a08d66d8..4be97401748 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -1,12 +1,11 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
 from megatron.core.jit import jit_fuser
-from megatron.core.tensor_parallel import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,
@@ -14,6 +13,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
+    get_tokens_per_expert_and_token_count,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -273,22 +273,29 @@ def is_aux_loss_enabled(self) -> bool:
         return False
 
     def _apply_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        with_padding_mask: bool = False,
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
-        tokens_per_expert = routing_map.sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                topk=self.topk,
+                with_padding_mask=with_padding_mask,
+            )
         )
-        num_tokens = routing_map.shape[0]
-        total_num_tokens = num_tokens * self.tp_cp_group.size()
 
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
-            tokens_per_expert=tokens_per_expert,
+            tokens_per_expert=global_tokens_per_expert,
             total_num_tokens=total_num_tokens,
             topk=self.topk,
             num_experts=self.config.num_moe_experts,
@@ -296,7 +303,12 @@ def _apply_aux_loss(
             fused=self.config.moe_router_fusion,
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group
+            probs,
+            aux_loss_coeff,
+            aux_loss,
+            "load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -307,6 +319,7 @@ def _apply_seq_aux_loss(
         routing_map: torch.Tensor,
         seq_length: int,
         bsz: int,
+        with_padding_mask: bool = False,
     ):
         """Apply the sequence-level auxiliary loss for the given scores and routing map.
 
@@ -320,17 +333,21 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
-        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+        routing_map = routing_map.reshape(seq_length, -1)
+
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk * bsz,
+            )
         )
 
-        total_num_tokens = seq_length * self.tp_cp_group.size()
-
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
-                tokens_per_expert=tokens_per_expert,
+                tokens_per_expert=global_tokens_per_expert,
                 total_num_tokens=total_num_tokens,
                 topk=self.topk,
                 num_experts=self.config.num_moe_experts,
@@ -339,31 +356,43 @@ def _apply_seq_aux_loss(
             )
             / bsz
         )
+
         probs = self.attach_and_log_load_balancing_loss(
-            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group
+            probs,
+            seq_aux_loss_coeff,
+            aux_loss,
+            "seq_load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
     def _apply_global_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        with_padding_mask: bool = False,
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
         if global_aux_loss_coeff == 0:
             return probs
 
-        tokens_per_expert = routing_map.sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_dp_cp_group
+        # Use unified function to compute tokens_per_expert and num_tokens
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_dp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk,
+            )
         )
 
-        self.global_tokens_per_expert += tokens_per_expert
+        self.global_tokens_per_expert += global_tokens_per_expert
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        num_tokens = scores_for_aux_loss.shape[0]
-        total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
-
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
             tokens_per_expert=averated_tokens_per_expert,
@@ -380,6 +409,7 @@ def _apply_global_aux_loss(
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
             reduce_group_has_dp=True,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -391,18 +421,22 @@ def attach_and_log_load_balancing_loss(
         aux_loss_name: str,
         reduce_group: torch.distributed.ProcessGroup,
         reduce_group_has_dp: bool = False,
+        valid_token_count: Optional[Union[int, torch.Tensor]] = None,
     ):
         """Attach aux loss function to activation and add to logging.
 
         Args:
-            activation (torch.Tensor): The activation tensor to attach the loss to.
-            aux_loss_coeff (float): The coefficient for the auxiliary loss.
-            aux_loss (torch.Tensor): The auxiliary loss tensor.
-            aux_loss_name (str): The name of the auxiliary loss for logging.
-            reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss.
+            activation (torch.Tensor): Activation tensor to attach the aux loss to.
+            aux_loss_coeff (float): Coefficient for the aux loss.
+            aux_loss (torch.Tensor): Computed aux loss.
+            aux_loss_name (str): Name of the aux loss for logging.
+            reduce_group (torch.distributed.ProcessGroup): Process group for reduction.
             reduce_group_has_dp (bool): Whether the reduce group has data parallel ranks.
                 Set this to True if the reduce group has data parallel ranks. This flag is used to
                 ensure the correct reduction in aux loss tracking.
+            valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding
+                padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor).
+                If None, uses activation.shape[0]. Defaults to None.
         """
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
         # add the aux loss logging value to other layer's since it is difficult to get the
@@ -427,17 +461,22 @@ def attach_and_log_load_balancing_loss(
             # which scales both the main_loss gradient and aux_loss gradient by
             # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads function.
             # To correct this scaling, we need to scale the aux_loss by num_local_tokens here.
-            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0])
+            # Use valid_token_count (excluding padding) if provided, otherwise use total tokens.
+            num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0]
+            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens)
         else:
             activation = MoEAuxLossAutoScaler.apply(activation, aux_loss)
         return activation
 
-    def apply_z_loss(self, logits):
+    def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
         """Encourages the router's logits to remain small to enhance stability.
         Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
         Args:
             logits (torch.Tensor): The logits of the router.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [num_tokens]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             torch.Tensor: The logits after applying the z-loss.
@@ -445,7 +484,7 @@ def apply_z_loss(self, logits):
         if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled():
             # Skip Z loss calculations when using torch.no_grad() or checkpointing.
             moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()
-            z_loss = z_loss_func(logits, moe_z_loss_coeff)
+            z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask)
             if self.calculate_per_token_loss:
                 # The expected final scaling for z_loss gradients is
                 # 1/(num_micro_batches * dp_size).
@@ -454,7 +493,9 @@ def apply_z_loss(self, logits):
                 # which scales both the main_loss gradient and z_loss gradient by
                 # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads().
                 # To correct this scaling, we need to scale the z_loss by num_local_tokens here.
-                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * logits.shape[0])
+                # Count valid tokens: sum of inverted mask (False -> True = valid)
+                num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0]
+                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens)
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
@@ -488,20 +529,27 @@ def apply_input_jitter(self, input: torch.Tensor):
             return input
 
     @jit_fuser
-    def _apply_expert_bias(self, routing_map: torch.Tensor):
+    def _apply_expert_bias(
+        self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
+    ):
         """
         Update expert bias and tokens_per_expert
         Prevent extra local tokens accumulation on evaluation or activation recomputation
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
+                if padding_mask is not None:
+                    routing_map = routing_map & (~padding_mask)
                 self.local_tokens_per_expert += routing_map.sum(dim=0)
 
-    def routing(self, logits: torch.Tensor):
+    def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Top-k routing function
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
@@ -511,8 +559,12 @@ def routing(self, logits: torch.Tensor):
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
 
+        # Flatten padding_mask to [num_tokens] if provided
+        if padding_mask is not None:
+            padding_mask = padding_mask.reshape(-1)
+
         # Apply Z-Loss
-        logits = self.apply_z_loss(logits)
+        logits = self.apply_z_loss(logits, padding_mask=padding_mask)
 
         # Calculate probs and routing_map for token dispatching
         if self.routing_type == "sinkhorn":
@@ -546,18 +598,35 @@ def routing(self, logits: torch.Tensor):
         if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
-                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion
+                logits,
+                self.topk,
+                self.score_function,
+                fused=self.config.moe_router_fusion,
+                padding_mask=padding_mask,
+            )
+            probs = self._apply_aux_loss(
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
-            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)
             probs = self._apply_seq_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                seq_length,
+                bsz,
+                with_padding_mask=padding_mask is not None,
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
 
         # Optionally apply expert bias
-        self._apply_expert_bias(routing_map)
+        self._apply_expert_bias(routing_map, padding_mask=padding_mask)
 
         return probs, routing_map
 
@@ -567,12 +636,15 @@ def reset_global_aux_loss_tracker(self):
             self.global_tokens_per_expert.zero_()
             self.ga_steps.zero_()
 
-    def forward(self, input: torch.Tensor):
+    def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """
         Forward pass of the router.
 
         Args:
             input (torch.Tensor): Input tensor.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
         """
         self._maintain_float32_expert_bias()
 
@@ -584,7 +656,7 @@ def forward(self, input: torch.Tensor):
             # Apply force load balancing with random logits for benchmark
             logits = apply_random_logits(logits)
 
-        probs, routing_map = self.routing(logits)
+        probs, routing_map = self.routing(logits, padding_mask=padding_mask)
 
         return probs, routing_map
 
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index ea4464b4784..831b5546d53 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -453,12 +453,18 @@ def _checkpointed_forward(
         attention_bias: Tensor,
         packed_seq_params: PackedSeqParams,
         use_inner_quantization_context: bool,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Forward method with activation checkpointing."""
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+                hidden_states,
+                attention_mask,
+                context,
+                context_mask,
+                rotary_pos_emb,
+                padding_mask=None,
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
@@ -489,6 +495,7 @@ def custom_forward(
                             attention_bias=attention_bias,
                             inference_context=None,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
                         )
                 return hidden_states, context
 
@@ -508,6 +515,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                 )
             else:
                 return tensor_parallel.checkpoint(
@@ -518,6 +526,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                 )
 
         if self.config.recompute_method == 'uniform':
@@ -623,6 +632,7 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         dynamic_inference_decode_only: Optional[bool] = None,
@@ -732,6 +742,7 @@ def forward(
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,
                     use_inner_quantization_context=use_inner_quantization_context,
+                    padding_mask=padding_mask,
                 )
             else:
                 for l_no, layer in enumerate(self.layers):
@@ -764,6 +775,7 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
+                            padding_mask=padding_mask,
                         )
 
                     if (
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index f575794a819..a5eaec92866 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import functools
 import logging
 import warnings
 from abc import ABC
@@ -510,7 +511,11 @@ def forward(self, *args, **kwargs):
         # runners in the cuda graph manager
         kwargs.pop("dynamic_inference_decode_only", None)
         hidden_states, context = self._forward_attention(*args, **kwargs)
-        output = self._forward_mlp(hidden_states, kwargs.get("inference_context", None))
+        output = self._forward_mlp(
+            hidden_states,
+            kwargs.get("inference_context", None),
+            padding_mask=kwargs.get("padding_mask", None),
+        )
         return output, context
 
     def _forward_attention(
@@ -527,6 +532,7 @@ def _forward_attention(
         inference_context: Optional[Any] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[Any] = None,
     ):
@@ -674,13 +680,18 @@ def _forward_pre_mlp_layernorm(self, hidden_states):
 
         return pre_mlp_layernorm_output
 
-    def _forward_mlp(self, hidden_states, inference_context=None):
+    def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None):
         """
         Perform a forward pass through the feed-forward layer.
 
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
-
+                Shape [seq_length, batch_size, hidden_size].
+            inference_context: Inference context for optimizations.
+            padding_mask (Tensor, optional): Padding mask for MoE routing.
+                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Only used for MoE layers to exclude padding tokens from aux loss computations.
+                The MoELayer will internally transform this to [seq_length, bsz] format.
         Returns:
             output (Tensor): Transformed hidden states of shape [s, b, h].
         """
@@ -716,10 +727,13 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     self.pg_collection.tp,
                     pre_mlp_layernorm_output,
+                    padding_mask=padding_mask,
                 )
             else:
                 mlp_output_with_bias = tensor_parallel.checkpoint(
-                    self.mlp, False, pre_mlp_layernorm_output
+                    functools.partial(self.mlp, padding_mask=padding_mask),
+                    False,
+                    pre_mlp_layernorm_output,
                 )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
@@ -739,7 +753,7 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 # Set the residual for fused reduce-scatter + add + layer-norm + all-gather
                 # operation in MLP's fc2.
                 self._set_fc2_residual(residual)
-            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
 
         nvtx_range_pop(suffix="mlp")
 
@@ -1263,7 +1277,7 @@ def create_mcore_cudagraph_manager(self, config):
                 self.config, self, function_name="_forward_mlp_postprocess"
             )
 
-    def _forward_mlp_router(self, hidden_states):
+    def _forward_mlp_router(self, hidden_states, padding_mask=None):
         """
         Executes the router phase of the MoE block.
 
@@ -1274,7 +1288,9 @@ def _forward_mlp_router(self, hidden_states):
         residual = hidden_states
         self.mlp.fwd_execution_map = "route"
         pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)
-        router_outputs = self.mlp(pre_mlp_layernorm_output, intermediate_tensors=())
+        router_outputs = self.mlp(
+            pre_mlp_layernorm_output, intermediate_tensors=(), padding_mask=padding_mask
+        )
 
         for attr_name in self.mlp.token_dispatcher.cudagraph_attrs:
             attr = getattr(self.mlp.token_dispatcher, attr_name)
@@ -1315,7 +1331,7 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b
         output = self.mlp(None, intermediate_tensors=(output, shared_expert_output))
         return self._forward_post_mlp((output, mlp_bias), residual)
 
-    def _forward_mlp(self, hidden_states, inference_context=None):
+    def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None):
         """
         Orchestrates the MLP forward pass, handling partial CUDA graph execution logic.
 
@@ -1330,9 +1346,11 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 "alongside inference."
             )
 
-        def _forward_mlp_partial_cudagraphs(hidden_states, inference_context=None):
+        def _forward_mlp_partial_cudagraphs(
+            hidden_states, inference_context=None, padding_mask=None
+        ):
             residual, hidden_states, probs, shared_expert_output = self._forward_mlp_router(
-                hidden_states
+                hidden_states, padding_mask=padding_mask
             )
             expert_output, mlp_bias = self._forward_mlp_expert_compute(hidden_states, probs)
             return self._forward_mlp_postprocess(
@@ -1350,12 +1368,17 @@ def _forward_mlp_partial_cudagraphs(hidden_states, inference_context=None):
                         tensor_parallel.random.get_cuda_rng_tracker,
                         parallel_state.get_tensor_model_parallel_group(),
                         hidden_states,
+                        padding_mask=padding_mask,
                     )
                 else:
                     return tensor_parallel.checkpoint(
-                        _forward_mlp_partial_cudagraphs, False, hidden_states
+                        functools.partial(
+                            _forward_mlp_partial_cudagraphs, padding_mask=padding_mask
+                        ),
+                        False,
+                        hidden_states,
                     )
             else:
-                return _forward_mlp_partial_cudagraphs(hidden_states)
+                return _forward_mlp_partial_cudagraphs(hidden_states, padding_mask=padding_mask)
         else:
-            return super()._forward_mlp(hidden_states)
+            return super()._forward_mlp(hidden_states, padding_mask=padding_mask)
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index 81e61a3404a..6c59dd3f9e3 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -23,7 +23,7 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-def build_model(config):
+def build_model(config, use_padding_mask=False):
     seq_len = 32
     max_seq_len = 300
     # ids = random.sample([i for i in range(max_seq_len)], seq_len)
@@ -39,6 +39,12 @@ def build_model(config):
         "attention_mask": torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda(),
     }
 
+    # Optionally add padding_mask with same shape as input_ids
+    if use_padding_mask:
+        padding_mask = torch.zeros((1, seq_len), dtype=torch.bool).cuda()
+        padding_mask[0, -8:] = True
+        data["padding_mask"] = padding_mask
+
     # build layer spec
     transformer_layer_spec = get_gpt_decoder_block_spec(config=config, use_transformer_engine=True)
     mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec.layer_specs[-1], True)
@@ -48,7 +54,7 @@ def build_model(config):
         config=config,
         transformer_layer_spec=transformer_layer_spec,
         mtp_block_spec=mtp_block_spec,
-        vocab_size=100,
+        vocab_size=128,
         pre_process=True,
         post_process=True,
         max_sequence_length=max_seq_len,
@@ -174,3 +180,109 @@ def test_1f1b_schedule_model_chunk(self, mtp_layers, dispatcher_type, fp8_flag,
                 gpt_models[i] = None
             gc.collect()
             torch.cuda.empty_cache()
+
+    @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0")
+    @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types())
+    @pytest.mark.parametrize("layers", [[2, 1], [1, 1]])
+    @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+    def test_1f1b_schedule_model_chunk_with_padding_mask(self, dispatcher_type, layers, tp_size):
+        """
+        Verifies all-to-all overlap optimization with padding_mask produces
+        the same results as the reference implementation with various TP/EP/CP combinations.
+        """
+        # Re-initialize model parallel with the specified configuration
+        Utils.destroy_model_parallel()
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=1,
+            expert_model_parallel_size=4,
+            expert_tensor_parallel_size=1,
+        )
+        set_streams()
+
+        microbatches = 1
+
+        gpt_models = []
+        schedule_plans = []
+        ref_captures = []
+        datas = []
+
+        # create TransformerConfig
+        extra_kwargs = {
+            "moe_token_dispatcher_type": dispatcher_type,
+            "tensor_model_parallel_size": tp_size,
+            "sequence_parallel": tp_size > 1,
+        }
+        if dispatcher_type == "flex":
+            extra_kwargs["moe_flex_dispatcher_backend"] = "deepep"
+            extra_kwargs["moe_router_dtype"] = "fp32"
+        with deterministic_mode():
+            for layer_num in layers:
+                output_tensors = []
+                # build config
+                config = get_test_config(num_layers=layer_num, extra_kwargs=extra_kwargs)
+                # build model with padding_mask
+                gpt_model, schedule_plan, data = build_model(config, use_padding_mask=True)
+                gpt_model.cuda()
+                gpt_models.append(gpt_model)
+                datas.append(data)
+                schedule_plans.append(schedule_plan)
+
+                # run reference
+                for _ in range(microbatches):
+                    loss = gpt_model.forward(**data)
+                    loss = float16_to_fp32(loss)
+                    loss.backward(torch.ones_like(loss))
+                    output_tensors.append(loss)
+
+                capture = {"outputs": output_tensors}
+                for name, param in gpt_model.named_parameters():
+                    capture[name] = param.grad
+                ref_captures.append(capture)
+                gpt_model.zero_grad()
+            assert gpt_models[0].embedding is not None
+            assert gpt_models[1].embedding is not None
+            # run a2a overlap
+            capture_0 = {"outputs": []}
+            capture_1 = {"outputs": []}
+            a2a_captures = [capture_0, capture_1]
+            for i in range(microbatches):
+                # 1st forward
+                if i > 0:
+                    assert (
+                        schedule_plans[0].pre_process is None
+                    ), "pre_process should be released after backward"
+                    schedule_plans[0] = gpt_models[0].build_schedule_plan(**datas[0])
+                    schedule_plans[1] = gpt_models[1].build_schedule_plan(**datas[1])
+                f_input_0 = TransformerModelChunkSchedulePlan.run(schedule_plans[0], None)
+                capture_0["outputs"].append(f_input_0)
+                # overlap
+                f_input_1 = TransformerModelChunkSchedulePlan.run(
+                    schedule_plans[1], schedule_plans[0], b_grad=torch.ones_like(f_input_0)
+                )
+                capture_1["outputs"].append(f_input_1)
+                # last backward
+                TransformerModelChunkSchedulePlan.run(
+                    None, schedule_plans[1], b_grad=torch.ones_like(f_input_1)
+                )
+            for i in range(len(gpt_models)):
+                for name, param in gpt_models[i].named_parameters():
+                    a2a_captures[i][name] = param.grad
+
+            # compare results
+            for i in range(len(ref_captures)):
+                comp_res = compare_captures(ref_captures[i], a2a_captures[i], True, True)
+                assert comp_res[0], f"[rank {torch.distributed.get_rank()}] {comp_res[1]}"
+
+            # release resources is necessary, otherwise later testcases will oom
+            for i in range(len(schedule_plans)):
+                schedule_plans[i] = None
+                ref_captures[i] = None
+                a2a_captures[i] = None
+                for k in datas[i]:
+                    datas[i][k] = None
+                datas[i] = None
+                gpt_models[i].zero_grad()
+                gpt_models[i] = None
+            gc.collect()
+            torch.cuda.empty_cache()
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
index 0fd2c445c9f..c6c4a75af99 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
@@ -502,8 +502,8 @@ def test_mtp_layer_overlap(self, dispatcher_type, fp8_flag):
             position_ids = torch.tensor(data, dtype=torch.int64).repeat((1, 1)).cuda()
             attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda()
             # get rotary pos emb
-            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _ = gpt_model._preprocess(
-                input_ids, position_ids
+            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _, _padding_mask = (
+                gpt_model._preprocess(input_ids, position_ids)
             )
             # reset model
             params = reset_model(gpt_model)
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 621e200c2cb..ccd11bf29af 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -577,3 +577,185 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
                 reduce_from_tensor_model_parallel_region(aux_loss, router.tp_cp_group)
             assert aux_loss.item() == 1, f"{aux_loss_type}: {aux_loss.item()}"
             clear_aux_losses_tracker()
+
+
+class TestPaddingMaskAuxLoss:
+    """Test padding mask support in various aux loss types."""
+
+    def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1, sequence_parallel=False):
+        """Initialize model parallel with given configuration.
+
+        Args:
+            tp_size: Tensor parallel size.
+            ep_size: Expert parallel size.
+            cp_size: Context parallel size.
+        """
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=1,
+            context_parallel_size=cp_size,
+            expert_model_parallel_size=ep_size,
+        )
+        _set_random_seed(seed_=123, data_parallel_random_init=False)
+
+        # Store parallel configuration
+        self.tp_size = tp_size
+        self.ep_size = ep_size
+        self.cp_size = cp_size
+
+        # Default configuration
+        self.default_transformer_config = TransformerConfig(
+            num_layers=1,
+            hidden_size=12,
+            num_attention_heads=8,
+            num_moe_experts=32,
+            use_cpu_initialization=True,
+            moe_router_load_balancing_type="aux_loss",
+            moe_router_topk=8,
+            moe_aux_loss_coeff=1.0,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            add_bias_linear=False,
+            tensor_model_parallel_size=tp_size,
+            expert_model_parallel_size=ep_size,
+            context_parallel_size=cp_size,
+            sequence_parallel=sequence_parallel and tp_size > 1,
+        )
+
+    def new_router(self, **kwargs):
+        """Create a new router with updated configuration."""
+        pg_collection = get_default_pg_collection()
+        new_transformer_config = dataclasses.replace(self.default_transformer_config, **kwargs)
+        router = TopKRouter(config=new_transformer_config, pg_collection=pg_collection)
+        router.set_layer_number(0)
+        return router
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size):
+        """Test that padding tokens are correctly excluded from aux loss calculation."""
+        # Initialize model parallel with given configuration
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+
+        try:
+            clear_aux_losses_tracker()
+
+            router = self.new_router(
+                moe_router_load_balancing_type=aux_loss_type,
+                moe_aux_loss_coeff=1.0,
+                moe_router_dtype="fp64",
+            ).cuda()
+
+            seq_len = 32
+            batch_size = 2
+            hidden_size = router.config.hidden_size
+
+            # Create input with padding
+            hidden_states_full = torch.randn(
+                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+            )
+
+            # Create padding mask: first half valid, second half padding
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = True
+
+            # Test with padding mask
+            router.weight.grad = None
+            scores_with_mask, routing_map_with_mask = router(
+                hidden_states_full, padding_mask=padding_mask
+            )
+            scores_with_mask.backward(torch.zeros_like(scores_with_mask))
+
+            loss_name = {
+                "aux_loss": "load_balancing_loss",
+                "seq_aux_loss": "seq_load_balancing_loss",
+                "global_aux_loss": "global_load_balancing_loss",
+            }[aux_loss_type]
+
+            tracker = get_moe_layer_wise_logging_tracker()
+            aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
+            grad_with_mask = router.weight.grad.clone()
+
+            # Test without padding (with only half of the tokens)
+            clear_aux_losses_tracker()
+            router.weight.grad = None
+            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
+            scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
+            scores_without_mask.backward(torch.zeros_like(scores_without_mask))
+
+            aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
+            grad_without_mask = router.weight.grad.clone()
+
+            # The aux loss with mask should be equal to the aux loss without mask
+            assert torch.equal(aux_loss_with_mask, aux_loss_without_mask)
+            assert torch.equal(grad_with_mask, grad_without_mask)
+
+            clear_aux_losses_tracker()
+        finally:
+            # Always cleanup model parallel
+            Utils.destroy_model_parallel()
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
+        """Test that padding mask works correctly with z_loss."""
+        # Initialize model parallel with given configuration
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+
+        try:
+            clear_aux_losses_tracker()
+
+            router = self.new_router(
+                moe_router_load_balancing_type="aux_loss",
+                moe_aux_loss_coeff=0.0,
+                moe_z_loss_coeff=1.0,
+                moe_router_dtype="fp32",
+            ).cuda()
+
+            seq_len = 32
+            batch_size = 2
+            hidden_size = router.config.hidden_size
+
+            # Create input
+            hidden_states_full = torch.randn(
+                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+            )
+
+            # Create padding mask: first half valid, second half padding
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = True
+
+            # Test with padding mask
+            router.weight.grad = None
+            scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
+            scores_with_mask.sum().backward()
+
+            tracker = get_moe_layer_wise_logging_tracker()
+            z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
+            grad_with_mask = router.weight.grad.clone()
+
+            # Test without padding (with only half of the tokens)
+            clear_aux_losses_tracker()
+            router.weight.grad = None
+            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
+            scores_without_mask, _ = router(hidden_states_valid)
+            scores_without_mask.sum().backward()
+
+            z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
+            grad_without_mask = router.weight.grad.clone()
+
+            # The z_loss with mask should be close to the z_loss without mask
+            assert torch.equal(z_loss_with_mask, z_loss_without_mask)
+            assert torch.equal(grad_with_mask, grad_without_mask)
+
+            clear_aux_losses_tracker()
+        finally:
+            # Always cleanup model parallel
+            Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index 904595928de..4d6b5ee2c3e 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -127,6 +127,53 @@ def test_aux_loss(self):
         out.sum().mul_(0).backward()
         assert self.sequential_mlp.router.weight.grad.abs().sum() > 0
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_router_with_padding_mask(self):
+        """Test that padding mask correctly excludes padding tokens from routing."""
+        self.router = self.router.cuda()
+        seq_len = 32
+        batch_size = 2
+        hidden_size = self.router.config.hidden_size
+
+        # Create input with shape [seq_len, batch_size, hidden_size]
+        hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16()
+
+        # Create padding mask: first half valid, second half padding
+        # padding_mask shape: [seq_len, batch_size]
+        # Convention: True = padding (exclude), False = valid (include)
+        padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+        padding_mask[seq_len // 2 :, :] = True  # Second half is padding
+
+        # Test forward pass with padding mask
+        with torch.no_grad():
+            probs_with_mask, routing_map_with_mask = self.router(
+                hidden_states, padding_mask=padding_mask
+            )
+
+            # Test forward pass without padding mask (only valid tokens)
+            hidden_states_valid = hidden_states[: seq_len // 2, :, :]
+            probs_without_mask, routing_map_without_mask = self.router(hidden_states_valid)
+
+            # The valid part of routing with mask should match routing without mask
+            probs_valid_part = probs_with_mask.reshape(seq_len, batch_size, -1)[
+                : seq_len // 2, :, :
+            ]
+            probs_valid_part = probs_valid_part.reshape(-1, probs_valid_part.shape[-1])
+
+            # Check that shapes are as expected
+            assert probs_with_mask.shape == (
+                seq_len * batch_size,
+                self.router.config.num_moe_experts,
+            )
+            assert routing_map_with_mask.shape == (
+                seq_len * batch_size,
+                self.router.config.num_moe_experts,
+            )
+
+            # Verify that probs for valid tokens are similar
+            assert torch.equal(probs_valid_part, probs_without_mask)
+
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_router_dtype(self):

From 93ddc24720bb8d62ae68ced044ad59e1a95ead0f Mon Sep 17 00:00:00 2001
From: Tong Liu <tongliu@nvidia.com>
Date: Wed, 28 Jan 2026 20:50:31 +0800
Subject: [PATCH 53/79] Use the latest hybrid-ep (#3093)

---
 docker/Dockerfile.ci.dev | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev
index b43b7286506..bb9ca5fbe9a 100644
--- a/docker/Dockerfile.ci.dev
+++ b/docker/Dockerfile.ci.dev
@@ -70,7 +70,7 @@ RUN bash -ex <<"EOF"
 
     git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git
     pushd DeepEP
-        git checkout 83e0d156807f31abed4ea55c2fa6eb4b62a11b82
+        git checkout eb9cee7de5a24193bf09500668d3a619d3d3f3fb
         patch -p1 < /workspace/deepep.patch
     popd
     TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" uv pip install --no-build-isolation -v DeepEP/.

From fc6969fbb34164b911e66b39db3d1ed65a5b6837 Mon Sep 17 00:00:00 2001
From: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
Date: Wed, 28 Jan 2026 19:05:47 +0200
Subject: [PATCH 54/79] remove retro (#3001)

Signed-off-by: dimapihtar <dpihtar@gmail.com>
---
 .gitlab/labeler-config.yml                    |   3 -
 docs/api-guide/models/models.md               |   2 +-
 docs/models/llms.md                           |   7 -
 examples/gpt3/gpt_config.yaml                 |   1 -
 examples/retro/README.md                      |  74 --
 examples/retro/preprocess_data.sh             | 144 ----
 examples/retro/train_retro_2b_distributed.sh  |  98 ---
 .../blended_megatron_dataset_config.py        |   3 -
 megatron/core/datasets/retro/__init__.py      |   5 -
 .../core/datasets/retro/config/__init__.py    |  16 -
 .../datasets/retro/config/bert_embedders.py   |  49 --
 megatron/core/datasets/retro/config/config.py | 135 ----
 .../retro/config/gpt_chunk_datasets.py        |  15 -
 .../core/datasets/retro/config/tokenizers.py  |  15 -
 megatron/core/datasets/retro/db/__init__.py   |   9 -
 megatron/core/datasets/retro/db/build.py      | 649 ------------------
 megatron/core/datasets/retro/db/dataset.py    | 114 ---
 megatron/core/datasets/retro/db/utils.py      | 398 -----------
 megatron/core/datasets/retro/external_libs.py |  13 -
 .../core/datasets/retro/index/__init__.py     |  11 -
 megatron/core/datasets/retro/index/build.py   | 339 ---------
 megatron/core/datasets/retro/index/factory.py |  40 --
 megatron/core/datasets/retro/index/index.py   | 150 ----
 .../datasets/retro/index/indexes/__init__.py  |  10 -
 .../retro/index/indexes/faiss_base.py         | 179 -----
 .../retro/index/indexes/faiss_par_add.py      | 253 -------
 megatron/core/datasets/retro/index/utils.py   | 126 ----
 .../core/datasets/retro/index/validate.py     | 194 ------
 .../core/datasets/retro/query/__init__.py     |   1 -
 .../datasets/retro/query/gpt_chunk_dataset.py | 109 ---
 .../retro/query/multi_split_gpt_dataset.py    | 115 ----
 megatron/core/datasets/retro/query/query.py   | 449 ------------
 .../datasets/retro/query/retro_dataset.py     | 251 -------
 megatron/core/datasets/retro/query/utils.py   |  35 -
 megatron/core/datasets/retro/utils.py         | 386 -----------
 megatron/core/enums.py                        |   2 -
 megatron/core/models/retro/__init__.py        |  13 -
 megatron/core/models/retro/base_attention.py  |  47 --
 megatron/core/models/retro/config.py          |  88 ---
 .../core/models/retro/decoder_attention.py    | 319 ---------
 megatron/core/models/retro/decoder_spec.py    | 202 ------
 .../core/models/retro/encoder_attention.py    | 231 -------
 megatron/core/models/retro/encoder_spec.py    | 178 -----
 megatron/core/models/retro/model.py           | 107 ---
 megatron/core/models/retro/utils.py           |  24 -
 .../core/tokenizers/megatron_tokenizer.py     |   3 +-
 .../core/tokenizers/text/models/__init__.py   |   1 -
 .../tokenizers/text/models/retro_tokenizer.py |  12 -
 megatron/legacy/data/orqa_wiki_dataset.py     |   3 +-
 megatron/legacy/model/enums.py                |   3 -
 megatron/legacy/model/gpt_model.py            |   6 -
 megatron/legacy/model/language_model.py       |  18 +-
 megatron/legacy/model/transformer.py          | 305 +-------
 megatron/training/arguments.py                | 150 ----
 megatron/training/checkpointing.py            |   1 -
 megatron/training/one_logger_utils.py         |   9 +-
 megatron/training/training.py                 |   7 -
 megatron/training/yaml_arguments.py           |  17 -
 pretrain_retro.py                             | 258 -------
 .../models/test_t5_model.py                   |  28 +-
 .../test_pipeline_parallel_layout.py          |   1 -
 tests/unit_tests/dist_checkpointing/utils.py  |   1 -
 .../pipeline_parallel/test_pipeline_layout.py |   1 -
 tests/unit_tests/test_checkpointing.py        |   1 -
 .../transformer/test_retro_attention.py       | 204 ------
 tools/bert_embedding/embed.py                 | 225 +++++-
 tools/retro/README.md                         | 256 -------
 tools/retro/build_db.md                       | 421 ------------
 tools/retro/cli/__init__.py                   |   3 -
 tools/retro/cli/__main__.py                   |   9 -
 tools/retro/cli/cli.py                        | 301 --------
 tools/retro/config_utils.py                   | 632 -----------------
 tools/retro/docker/Dockerfile                 |  19 -
 tools/retro/preprocess_data.py                | 296 --------
 tools/retro/sft/README.md                     |   3 -
 tools/retro/sft/dataset_conv.py               | 446 ------------
 tools/retro/sft/open_inst.sh                  |   1 -
 tools/retro/sft/sft_retro.py                  | 278 --------
 tools/retro/sft/sft_retro_lm.sh               | 150 ----
 tools/retro/text_generation/evaluate.py       | 200 ------
 tools/retro/text_generation/metrics.py        |  80 ---
 tools/retro/text_generation/retro_api.py      | 221 ------
 tools/retro/text_generation/retro_generate.sh | 125 ----
 .../retro/text_generation/retro_generation.py | 250 -------
 .../text_generation/retro_text_generation.py  | 263 -------
 85 files changed, 235 insertions(+), 10582 deletions(-)
 delete mode 100644 examples/retro/README.md
 delete mode 100644 examples/retro/preprocess_data.sh
 delete mode 100644 examples/retro/train_retro_2b_distributed.sh
 delete mode 100644 megatron/core/datasets/retro/__init__.py
 delete mode 100644 megatron/core/datasets/retro/config/__init__.py
 delete mode 100644 megatron/core/datasets/retro/config/bert_embedders.py
 delete mode 100644 megatron/core/datasets/retro/config/config.py
 delete mode 100644 megatron/core/datasets/retro/config/gpt_chunk_datasets.py
 delete mode 100644 megatron/core/datasets/retro/config/tokenizers.py
 delete mode 100644 megatron/core/datasets/retro/db/__init__.py
 delete mode 100644 megatron/core/datasets/retro/db/build.py
 delete mode 100644 megatron/core/datasets/retro/db/dataset.py
 delete mode 100644 megatron/core/datasets/retro/db/utils.py
 delete mode 100644 megatron/core/datasets/retro/external_libs.py
 delete mode 100644 megatron/core/datasets/retro/index/__init__.py
 delete mode 100644 megatron/core/datasets/retro/index/build.py
 delete mode 100644 megatron/core/datasets/retro/index/factory.py
 delete mode 100644 megatron/core/datasets/retro/index/index.py
 delete mode 100644 megatron/core/datasets/retro/index/indexes/__init__.py
 delete mode 100644 megatron/core/datasets/retro/index/indexes/faiss_base.py
 delete mode 100644 megatron/core/datasets/retro/index/indexes/faiss_par_add.py
 delete mode 100644 megatron/core/datasets/retro/index/utils.py
 delete mode 100644 megatron/core/datasets/retro/index/validate.py
 delete mode 100644 megatron/core/datasets/retro/query/__init__.py
 delete mode 100644 megatron/core/datasets/retro/query/gpt_chunk_dataset.py
 delete mode 100644 megatron/core/datasets/retro/query/multi_split_gpt_dataset.py
 delete mode 100644 megatron/core/datasets/retro/query/query.py
 delete mode 100644 megatron/core/datasets/retro/query/retro_dataset.py
 delete mode 100644 megatron/core/datasets/retro/query/utils.py
 delete mode 100644 megatron/core/datasets/retro/utils.py
 delete mode 100644 megatron/core/models/retro/__init__.py
 delete mode 100644 megatron/core/models/retro/base_attention.py
 delete mode 100644 megatron/core/models/retro/config.py
 delete mode 100644 megatron/core/models/retro/decoder_attention.py
 delete mode 100644 megatron/core/models/retro/decoder_spec.py
 delete mode 100644 megatron/core/models/retro/encoder_attention.py
 delete mode 100644 megatron/core/models/retro/encoder_spec.py
 delete mode 100644 megatron/core/models/retro/model.py
 delete mode 100644 megatron/core/models/retro/utils.py
 delete mode 100644 megatron/core/tokenizers/text/models/retro_tokenizer.py
 delete mode 100644 pretrain_retro.py
 delete mode 100644 tests/unit_tests/transformer/test_retro_attention.py
 delete mode 100644 tools/retro/README.md
 delete mode 100644 tools/retro/build_db.md
 delete mode 100644 tools/retro/cli/__init__.py
 delete mode 100644 tools/retro/cli/__main__.py
 delete mode 100644 tools/retro/cli/cli.py
 delete mode 100644 tools/retro/config_utils.py
 delete mode 100644 tools/retro/docker/Dockerfile
 delete mode 100644 tools/retro/preprocess_data.py
 delete mode 100644 tools/retro/sft/README.md
 delete mode 100644 tools/retro/sft/dataset_conv.py
 delete mode 100644 tools/retro/sft/open_inst.sh
 delete mode 100644 tools/retro/sft/sft_retro.py
 delete mode 100644 tools/retro/sft/sft_retro_lm.sh
 delete mode 100755 tools/retro/text_generation/evaluate.py
 delete mode 100755 tools/retro/text_generation/metrics.py
 delete mode 100644 tools/retro/text_generation/retro_api.py
 delete mode 100755 tools/retro/text_generation/retro_generate.sh
 delete mode 100644 tools/retro/text_generation/retro_generation.py
 delete mode 100755 tools/retro/text_generation/retro_text_generation.py

diff --git a/.gitlab/labeler-config.yml b/.gitlab/labeler-config.yml
index 0e218e4bae7..2c37345c0e6 100644
--- a/.gitlab/labeler-config.yml
+++ b/.gitlab/labeler-config.yml
@@ -14,9 +14,6 @@ BERT:
 GPT:
   - megatron/core/models/gpt/**
 
-RETRO:
-  - megatron/core/models/retro/**
-
 Dist-Ckpt:
   - megatron/core/dist_checkpointing
 
diff --git a/docs/api-guide/models/models.md b/docs/api-guide/models/models.md
index d3cb3e4cb4a..69dfc80211d 100644
--- a/docs/api-guide/models/models.md
+++ b/docs/api-guide/models/models.md
@@ -1,6 +1,6 @@
 # models package
 
-This package contains most of the popular LLMs . Currently we have support for GPT, Bert, T5 and Retro . This is an ever growing list so keep an eye out.
+This package contains most of the popular LLMs . Currently we have support for GPT, Bert, and T5 . This is an ever growing list so keep an eye out.
 
 ## Subpackages
 
diff --git a/docs/models/llms.md b/docs/models/llms.md
index 1464b934f9d..6789a4c551c 100644
--- a/docs/models/llms.md
+++ b/docs/models/llms.md
@@ -31,12 +31,6 @@ See the [Megatron Bridge supported models list](https://github.com/NVIDIA-NeMo/M
 |-------|-------------|--------------|
 | **T5** | Text-to-Text Transfer Transformer | Unified text-to-text framework, sequence-to-sequence |
 
-## Retrieval-Augmented Models
-
-| Model | Description | Key Features |
-|-------|-------------|--------------|
-| **RETRO** | Retrieval-Enhanced Transformer | Retrieval-augmented generation, knowledge grounding |
-
 ## Example Scripts
 
 Training examples for these models can be found in the `examples/` directory:
@@ -46,7 +40,6 @@ Training examples for these models can be found in the `examples/` directory:
 - `examples/mamba/` - Mamba training scripts
 - `examples/bert/` - BERT training scripts
 - `examples/t5/` - T5 training scripts
-- `examples/retro/` - RETRO training scripts
 
 ## Model Implementation
 
diff --git a/examples/gpt3/gpt_config.yaml b/examples/gpt3/gpt_config.yaml
index 2fd40e62143..18d305d9cb1 100644
--- a/examples/gpt3/gpt_config.yaml
+++ b/examples/gpt3/gpt_config.yaml
@@ -257,7 +257,6 @@ vocab_extra_ids: 0
 seq_length: 4096
 encoder_seq_length: null
 decoder_seq_length: null
-retriever_seq_length: 256
 sample_rate: 1.0
 mask_prob: 0.15
 short_seq_prob: 0.1
diff --git a/examples/retro/README.md b/examples/retro/README.md
deleted file mode 100644
index f78bcdeb56b..00000000000
--- a/examples/retro/README.md
+++ /dev/null
@@ -1,74 +0,0 @@
-# RETRO MODEL
-
-## Table of contents
-- [1. Training Setup](#1-training-setup)
-- [2. Data Preprocessing](#2-data-preprocessing)
-- [3. Configurations](#3-configurations)
-
-## 1. Training setup
-<a id="markdown-training-setup" name="training-setup"></a>
-
-To run the model using a docker container run it as follows
-```
-PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:23.09-py3
-CHECKPOINT_PATH="" #<Specify path>
-TENSORBOARD_LOGS_PATH=""#<Specify path>
-
-docker run \
-  --gpus=all \
-  --ipc=host \
-  --workdir /workspace/megatron-lm \
-  -v /path/to/data:/path/to/data \
-  -v /path/to/megatron-lm:/workspace/megatron-lm \
-  megatron-lm nvcr.io/nvidia/pytorch:23.09-py3 \
-  bash examples/retro/train_retro_2b_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH"
-
-```
-NOTE: Depending on the environment you are running it the above command might look slightly different.
-
-NOTE: Due to how Retro preprocess and caches elements of the pretraining dataset before training begins, some arguments are auto-loaded from the Retro preprocessing configuration. These loaded arguments include:
-
-- `--data-path`
-- `--data-cache-path`
-- `--eval-interval`
-- `--eval-iters`
-- `--global-batch-size`
-- `--tokenizer-type`
-- `--tokenizer-model`
-- `--vocab-file`
-- `--merge-file`
-- `--seed`
-- `--seq-length`
-- `--train-samples`
-
-
-## 2. Data Preprocessing
-<a id="markdown-data-preprocessing" name="data-preprocessing"></a>
-
-Retro preprocesses and caches data prior to pretraining, to greatly speed up pretraining. During data preprocessing, the retrieval database is built, and neighbor IDs are queried for each sample within the pretraining dataset. Please see `preprocess_data.sh` for an example script to preprocess data for Retro. The reference documentation for data preprocessing can be found [here](tools/retro/README.md).
-
-
-## 3. Configurations
-<a id="markdown-configurations" name="configurations"></a>
-The example in this folder shows you how to run a 2B model. Below are a few other example configurations.
-
-### 857M
-```
-       --num-layers 24 \
-       --hidden-size 1024 \
-       --num-attention-heads 16 \
-       --seq-length 2048 \
-       --tensor-model-parallel-size 1 \
-       --pipeline-model-parallel-size 1 \
-
-```
-
-### 4B
-```
-       --num-layers 48 \
-       --hidden-size 2560 \
-       --num-attention-heads 32 \
-       --tensor-model-parallel-size 1 \
-       --pipeline-model-parallel-size 1 \
-
-```
diff --git a/examples/retro/preprocess_data.sh b/examples/retro/preprocess_data.sh
deleted file mode 100644
index 5d2e66ba0e7..00000000000
--- a/examples/retro/preprocess_data.sh
+++ /dev/null
@@ -1,144 +0,0 @@
-#!/bin/bash
-
-set -u
-
-unset NCCL_DEBUG
-
-######## Megatron, Retro dirs. ########
-
-REPO_DIR="<path/to/megatron/repo>"
-RETRO_PROJECT_DIR="<path/to/retro/project/directory>"
-
-######## Task (e.g., db, index, query). ########
-
-# This script takes a single argument, which specifies the retro task to be
-# performed. The available tasks are: db-build, index-train, index-add, and
-# query-neighbors.
-
-# ~~ Examples ~~
-# RETRO_TASKS="db-build"          # Build the retrieval database
-# RETRO_TASKS="index-train"       # Train the index
-# RETRO_TASKS="index-add"         # Add data to the index
-# RETRO_TASKS="query-neighbors"   # Perform query pretraining for neighbors
-
-# You can also provide the task as a command-line argument when executing the
-# script. Example: ./preprocess_data.sh index-add
-RETRO_TASKS=$1
-
-######## Data. ########
-DATA_BLEND="<see --data-path in arguments.py>"
-
-######## Index. ########
-
-RETRO_INDEX_STR="OPQ32_64,IVF65536_HNSW8,PQ32"
-RETRO_INDEX_NTRAIN=66625331
-RETRO_INDEX_TRAIN_LOAD_FRACTION=0.97
-RETRO_INDEX_ADD_LOAD_FRACTION=0.95
-
-######## GPT. ########
-
-RETRO_GPT_SEED=1234
-RETRO_GPT_SPLIT="98,2,0"
-RETRO_GPT_DATA_PATH=${DATA_BLEND}
-RETRO_GPT_TRAIN_SAMPLES=200000
-RETRO_GPT_EVAL_INTERVAL=2000
-RETRO_GPT_EVAL_ITERS=50
-RETRO_GPT_LR_DECAY_SAMPLES=175000
-RETRO_GPT_LR_WARMUP_SAMPLES=10000
-RETRO_GPT_SEQ_LENGTH=2048
-RETRO_GPT_GLOBAL_BATCH_SIZE=256
-RETRO_GPT_CHUNK_LENGTH=64
-
-######## Query. ########
-
-RETRO_QUERY_NUM_NEIGHBORS_QUERY=200
-RETRO_QUERY_NUM_NEIGHBORS_SAVE=20
-RETRO_QUERY_EF_SEARCH=32
-RETRO_QUERY_NPROBE=4096
-
-######## Args. ########
-
-ARGS=" \
-    --distributed-timeout-minutes 600 \
-    --tensor-model-parallel-size 1 \
-    --pipeline-model-parallel-size 1 \
-    --num-layers 24 \
-    --hidden-size 1024 \
-    --num-attention-heads 16 \
-    --micro-batch-size 1 \
-    --global-batch-size ${RETRO_GPT_GLOBAL_BATCH_SIZE} \
-    --seq-length 512 \
-    --max-position-embeddings 512 \
-    --load ${RETRO_PROJECT_DIR}/checkpoints/bert \
-    --exit-on-missing-checkpoint \
-    --no-load-optim \
-    --data-path [null] \
-    --tokenizer-type BertWordPieceLowerCase \
-    --vocab-file ${RETRO_PROJECT_DIR}/tokenizer/bert-large-uncased-vocab.txt \
-    --split ${RETRO_GPT_SPLIT} \
-    --distributed-backend nccl \
-    --lr 0.0001 \
-    --lr-decay-style linear \
-    --min-lr 1.0e-5 \
-    --train-samples ${RETRO_GPT_TRAIN_SAMPLES} \
-    --lr-decay-samples ${RETRO_GPT_LR_DECAY_SAMPLES} \
-    --lr-warmup-samples ${RETRO_GPT_LR_WARMUP_SAMPLES} \
-    --weight-decay 1e-2 \
-    --clip-grad 1.0 \
-    --eval-interval ${RETRO_GPT_EVAL_INTERVAL} \
-    --eval-iters ${RETRO_GPT_EVAL_ITERS} \
-    --bf16 \
-    --no-data-sharding \
-    --no-gradient-accumulation-fusion \
-    --no-async-tensor-model-parallel-allreduce \
-    --bert-embedder-type megatron \
-    --output-bert-embeddings \
-    \
-    --retro-project-dir ${RETRO_PROJECT_DIR} \
-    --retro-tasks ${RETRO_TASKS} \
-    --retro-bert-vocab-file tokenizer/bert-large-uncased-vocab.txt \
-    --retro-bert-tokenizer-type BertWordPieceLowerCase \
-    \
-    --retro-gpt-seed ${RETRO_GPT_SEED} \
-    --retro-gpt-tokenizer-type GPTSentencePieceTokenizer \
-    --retro-gpt-tokenizer-model /path/to/tokenizer/model \
-    --retro-gpt-seq-length ${RETRO_GPT_SEQ_LENGTH} \
-    --retro-gpt-chunk-length ${RETRO_GPT_CHUNK_LENGTH} \
-    --retro-gpt-global-batch-size ${RETRO_GPT_GLOBAL_BATCH_SIZE} \
-    --retro-gpt-eval-interval ${RETRO_GPT_EVAL_INTERVAL} \
-    --retro-gpt-eval-iters ${RETRO_GPT_EVAL_ITERS} \
-    --retro-gpt-split ${RETRO_GPT_SPLIT} \
-    --retro-gpt-data-path ${RETRO_GPT_DATA_PATH} \
-    --retro-gpt-train-samples ${RETRO_GPT_TRAIN_SAMPLES} \
-    \
-    --retro-index-str ${RETRO_INDEX_STR} \
-    --retro-index-ntrain ${RETRO_INDEX_NTRAIN} \
-    --retro-index-train-load-fraction ${RETRO_INDEX_TRAIN_LOAD_FRACTION} \
-    --retro-index-add-load-fraction ${RETRO_INDEX_ADD_LOAD_FRACTION} \
-    --no-retro-index-delete-training-embeddings \
-    --no-retro-index-delete-added-codes \
-    \
-    --retro-query-num-neighbors-query ${RETRO_QUERY_NUM_NEIGHBORS_QUERY} \
-    --retro-query-num-neighbors-save ${RETRO_QUERY_NUM_NEIGHBORS_SAVE} \
-    --retro-query-ef-search ${RETRO_QUERY_EF_SEARCH} \
-    --retro-query-nprobe ${RETRO_QUERY_NPROBE} \
-"
-
-######## Command. ########
-
-NPROCS=8 # Number of GPUs.
-CMD="\
-    cd ${REPO_DIR} && pwd && \
-    export PYTHONPATH=$PYTHONPATH:${REPO_DIR} && \
-    python -m torch.distributed.run \
-    --nproc_per_node ${NPROCS} \
-    --nnodes 1 \
-    --node_rank ${NODE_RANK} \
-    --master_addr ${MASTER_ADDR} \
-    --master_port 6000 \
-    tools/retro/preprocess_data.py ${ARGS} \
-"
-echo "~~~~~~~~~~~~~~~~~~~~~~~~~~"
-echo "CMD = '$CMD'."
-echo "~~~~~~~~~~~~~~~~~~~~~~~~~~"
-eval $CMD
diff --git a/examples/retro/train_retro_2b_distributed.sh b/examples/retro/train_retro_2b_distributed.sh
deleted file mode 100644
index c8276b56f43..00000000000
--- a/examples/retro/train_retro_2b_distributed.sh
+++ /dev/null
@@ -1,98 +0,0 @@
-#!/bin/bash
-
-# Runs the "307M" parameter Retro model.
-
-export CUDA_DEVICE_MAX_CONNECTIONS=1
-
-GPUS_PER_NODE=8
-# Change for multinode config
-MASTER_ADDR=localhost
-MASTER_PORT=6000
-NUM_NODES=1
-NODE_RANK=0
-WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
-
-CHECKPOINT_PATH=$1 #<Specify path>
-TENSORBOARD_LOGS_PATH=$2 #<Specify path>
-
-DISTRIBUTED_ARGS=(
-    --nproc_per_node $GPUS_PER_NODE 
-    --nnodes $NUM_NODES 
-    --master_addr $MASTER_ADDR 
-    --master_port $MASTER_PORT
-)
-
-######## GPT or Retro? ########
-
-# 0 : GPT.
-# 1 : Retro
-
-ADD_RETRIEVER=1
-
-######## Megatron, Retro dirs. ########
-
-RETRO_PROJECT_DIR="<path/to/retro/project/directory>"
-
-######## Model, training args. ########
-
-# ** Note: --seq-length auto loaded from Retro project dir.
-RETRO_MODEL_ARGS=(
-    --num-layers 32
-    --hidden-size 2048
-    --num-attention-heads 32
-)
-
-# ** Note: --data-path, --tokenizer-type, and --tokenizer-model auto loaded from Retro project dir.
-DATA_ARGS=(
-    --split 98,2,0
-)
-
-MODEL_PARALLEL_ARGS=(
-    --tensor-model-parallel-size 8
-    --pipeline-model-parallel-size 1 
-)
-
-# ** Note: --eval-interval, --eval-iters auto loaded from Retro project dir.
-EVAL_AND_LOGGING_ARGS=(
-    --log-interval 100
-    --save-interval 10000 
-    --eval-interval 1000 
-    --save $CHECKPOINT_PATH 
-    --load $CHECKPOINT_PATH 
-    --eval-iters 10
-    --tensorboard-dir $TENSORBOARD_LOGS_PATH 
-)
-
-TRAINING_ARGS=" \
-    --retro-project-dir ${RETRO_PROJECT_DIR} \
-    --transformer-impl transformer_engine \
-    --num-workers 8 \
-    --micro-batch-size 4 \
-    --lr-decay-samples 166400000 \
-    --lr-warmup-samples 162761 \
-    --lr 6.0e-4 \
-    --min-lr 6.0e-5 \
-    --lr-decay-style cosine \
-    --clip-grad 1.0 \
-    --weight-decay 0.1 \
-    --adam-beta1 0.9 \
-    --adam-beta2 0.95 \
-    --init-method-std 0.023 \
-    --log-params-norm \
-    --log-num-zeros-in-grad \
-    --bf16 \
-    --no-data-sharding \
-"
-
-if [ "$ADD_RETRIEVER" = "1" ]; then
-    TRAINING_ARGS+=" --retro-add-retriever"
-fi
-
-######## Command. ########
-
-torchrun ${DISTRIBUTED_ARGS[@]} pretrain_retro.py \
-    ${RETRO_MODEL_ARGS[@]} \
-    ${TRAINING_ARGS} \
-    ${MODEL_PARALLEL_ARGS[@]} \
-    ${DATA_ARGS[@]} \
-    ${EVAL_AND_LOGGING_ARGS[@]}
diff --git a/megatron/core/datasets/blended_megatron_dataset_config.py b/megatron/core/datasets/blended_megatron_dataset_config.py
index cee7f333bb8..a86efbe4963 100644
--- a/megatron/core/datasets/blended_megatron_dataset_config.py
+++ b/megatron/core/datasets/blended_megatron_dataset_config.py
@@ -181,9 +181,6 @@ def convert_split_vector_to_split_matrix(
 
     [0.99, 0.01, 0.0] -> [(0, 0.99), (0.99, 1.0), None]
 
-    Ex. a conversion for Retro when Retro pretraining uses a [0.99, 0.01, 0.0] split and Retro
-    preprocessing used a [0.98, 0.02, 0.0] split:
-
     [0.99, 0.01, 0.0], [0.98, 0.02, 0.0] -> [(0, 0.98), (0.99, 1.0), None]
 
     Args:
diff --git a/megatron/core/datasets/retro/__init__.py b/megatron/core/datasets/retro/__init__.py
deleted file mode 100644
index 7ce970c6e9f..00000000000
--- a/megatron/core/datasets/retro/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-from .config import RetroGPTChunkDatasets
-from .query.multi_split_gpt_dataset import MultiSplitGPTDataset, MultiSplitGPTDatasetConfig
-from .query.retro_dataset import get_retro_datasets
diff --git a/megatron/core/datasets/retro/config/__init__.py b/megatron/core/datasets/retro/config/__init__.py
deleted file mode 100644
index 3635bedb3f4..00000000000
--- a/megatron/core/datasets/retro/config/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-Exports:
-
-  - Embedder: Base class for all Bert embedders.
-  - RetroBertEmbedders: Container class for in-memory and on-disk embedders.
-  - RetroPreprocessingConfig: Configuration class for all of Retro preprocessing.
-  - RetroGPTChunkDatasets: Container class for train, valid, and test datasets.
-  - RetroTokenizers: Container class for GPT and Bert tokenizers.
-"""
-
-from .bert_embedders import Embedder, RetroBertEmbedders
-from .config import RetroPreprocessingConfig
-from .gpt_chunk_datasets import RetroGPTChunkDatasets
-from .tokenizers import RetroTokenizers
diff --git a/megatron/core/datasets/retro/config/bert_embedders.py b/megatron/core/datasets/retro/config/bert_embedders.py
deleted file mode 100644
index c34cd3d79dd..00000000000
--- a/megatron/core/datasets/retro/config/bert_embedders.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Container dataclass for holding both in-memory and on-disk Bert embedders."""
-
-import abc
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-
-
-class Embedder(abc.ABC):
-    """Base class for all Bert embedders.
-
-    All embedders should be able to embed either an entire text dataset (to a 2D
-    numpy array), or a single text string (to a 1D numpy array).
-    """
-
-    @abc.abstractmethod
-    def embed_text_dataset(self, text_dataset: torch.utils.data.Dataset) -> np.ndarray:
-        """Embed a text dataset.
-
-        Args:
-            text_dataset (torch.utils.data.Dataset): Text dataset to embed.
-                Each sample of the text dataset should output a dict with a key 'text'
-                and a string value.
-
-        Returns:
-            A 2D ndarray with shape (len(text_dataset), dimension(embedder)).
-        """
-
-    @abc.abstractmethod
-    def embed_text(self, text: str) -> np.ndarray:
-        """Embed a simple string of text.
-
-        Args:
-            text (str): A single text sample.
-
-        Returns:
-            A 1D ndarray with shape (dimensions(embedder),).
-        """
-
-
-@dataclass
-class RetroBertEmbedders:
-    """Container dataclass for in-memory and on-disk Bert embedders."""
-
-    disk: Embedder
-    mem: Embedder
diff --git a/megatron/core/datasets/retro/config/config.py b/megatron/core/datasets/retro/config/config.py
deleted file mode 100644
index ac9ca841242..00000000000
--- a/megatron/core/datasets/retro/config/config.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Retro preprocessing config."""
-
-from dataclasses import dataclass
-
-from megatron.core.transformer import TransformerConfig
-
-from .bert_embedders import RetroBertEmbedders
-from .gpt_chunk_datasets import RetroGPTChunkDatasets
-from .tokenizers import RetroTokenizers
-
-
-@dataclass
-class RetroPreprocessingConfig(TransformerConfig):
-    """Configuration object for Retro preprocessing.
-
-    *Note* : Arguments prefixed with '--retro-gpt-*' or '--retro-bert-*' are
-    included and named as such to more easily handle managing both models
-    running at the same time. Megatron is not optimized to run two models at
-    once, so this naming convention makes it clearer.
-
-    Args:
-
-        retro_project_dir (str): Retro project directory, which contains the preprocessed data for for pretraining. This directory is built during preprocessing (see tools/retro/README.md), and contains subdirectories for the chunk database and pretraining neighbors.
-        retro_tasks (str): Comma-separated list of tasks to run. Run entire preprocesing pipeline by using '--retro-tasks build'. Alternatively, run individual stages with tasks (in this order) 'db-build', 'index-build', or 'query-pretraining-neighbors'. For example, '--retro-tasks db-build,index-build,query-pretraining-neighbors' is equivalent to '--retro-tasks build'; or the argument can contain a subset of these tasks. Stages must always be run in the correct order (listed above).
-        retro_task_validate (float): If defined, validate a randomly sampled subset of the existing results of the given task. Each task implements a 'validate' method that is responsible for sampling a `retro_task_validate` fraction of the existing results, and then checking for bitwise equality with the current code base. (E.g., `--retro-task-validate 0.01`.)
-        retro_block_size (int): Number of chunks to process at a time when generating Bert embeddings and querying the search index. Partial results for each block are generally saved to disk in separate files.
-        retro_doc_block_size (int): Number of documents to processe at time when processing token datasets into chunk databases. The partial chunk database for each block is saved into a separate file.
-        retro_gpt_seed (int): Random seed used for python, numpy, pytorch, and cuda.
-        retro_gpt_data_path (str): Path to the training dataset. Accepted format: 1) a single data path, 2) multiple datasets in the form: dataset1-weight dataset1-path dataset2-weight dataset2-path ... It is used with --split when a single dataset used for all three: train, valid and test. It is exclusive to the other --*-data-path args.
-        retro_gpt_data_cache_path (str): Path to a directory to hold cached index files.
-        retro_gpt_split (str): Comma-separated list of proportions for training, validation, and test split. For example the split `90,5,5` will use 90%% of data for training, 5%% for validation and 5%% for test.
-        retro_gpt_train_samples (int): Total number of samples to train over all training runs.
-        retro_gpt_eval_interval (int): GPT evaluation interval.
-        retro_gpt_eval_iters (int): GPT evaluation iterations.
-        retro_gpt_tokenizer_type (str): GPT tokenizer type.
-        retro_gpt_tokenizer_model (str): GPT tokenizer model file.
-        retro_gpt_vocab_file (str): GPT vocab file.
-        retro_gpt_merge_file (str): GPT merge file.
-        retro_gpt_seq_length (int): GPT sequence length.
-        retro_gpt_global_batch_size (int): GPT global batch size.
-        retro_gpt_chunk_length (int): GPT chunk length.
-        retro_bert_tokenizer_type (str): Bert tokenizer type (for when using '--bert-embedder-type megatron').
-        retro_bert_vocab_file (str): Bert vocab file.
-        retro_bert_batch_size (int): Micro-batch size for processing Bert embeddings.
-        retro_bert_max_chunk_length (int): Maximum sequence length for Bert embeddings. (Named 'chunk' here in reference to these Bert sequences being converted from GPT chunks.)
-        retro_index_type (str): A 'faiss-base' index is a simple, un-optimized wrapper around a Faiss index. A 'faiss-par-add' index optimizes the 'add()' method by making it multi-node and multi-process, but with bit-wise equivalent results.
-        retro_index_str (str): Index string used for calling faiss.index_factory(). For example, 'IVF262144_HNSW32,Flat' or 'OPQ32_256,IVF4194304_HNSW32,PQ32'.
-        retro_index_ntrain (int): Number of database chunks to use for training the index. This value must be less or equal to the total number of chunks in the database.
-        retro_index_train_load_fraction (float): Fraction of sampled chunks to use for training the index. Useful when our total sampled embeddings use too much memory; lowering the load fraction is less costly than re-embedding a new sampled dataset from scratch.
-        retro_index_add_load_fraction (float): Fraction of database chunks to use for adding to the index. Useful when our total index size would use too much memory; lowering the load fraction is less costly than re-designing our token datasets.
-        retro_index_delete_training_embeddings (bool): Delete training embeddings for the search index. Useful for debugging.
-        retro_index_delete_added_codes (bool): Delete added codes for the search index. Useful for debugging.
-        retro_query_ef_search (int): Index ef-search parameter for Hierarchical Navigable Small Worlds (HNSW) during querying.
-        retro_query_nprobe (int): Index nprobe parameter for Inverted File (IVF) during querying.
-        retro_query_num_neighbors_query (int): Number of neighbors to retrieve when calling index.search().
-        retro_query_num_neighbors_save (int): Number of neighbors to save to disk after the index's returned neighbors. If longer than target value, neighbors truncated; and if shorter than target value, neighbors are padded with -1's.
-        retro_bert_embedders (RetroBertEmbedders): Set of Bert embedders used for embedding chunks. Contains entries: 1) 'mem' for an in-memory embedder, and 2) 'disk' for an embedder that saves results in blocks to disk.
-        retro_gpt_chunk_datasets (RetroGPTChunkDatasets): GPT datasets for 'train', 'valid', and 'test'.
-        retro_tokenizers (RetroTokenizers): GPT ('gpt') and Bert ('bert') tokenizers.
-    """
-
-    # Basic.
-    retro_project_dir: str = None
-    retro_tasks: str = 'build'
-    retro_task_validate: float = None
-    retro_block_size: int = 100000
-    retro_doc_block_size: int = 100000
-
-    # GPT.
-    retro_gpt_seed: int = 1234
-    retro_gpt_data_path: list = None  # basic list here, for parsing purposes
-    retro_gpt_data_cache_path: str = None
-    retro_gpt_split: str = '969,30,1'
-    retro_gpt_train_samples: int = None
-    retro_gpt_eval_interval: int = None
-    retro_gpt_eval_iters: int = None
-    retro_gpt_tokenizer_type: str = None
-    retro_gpt_tokenizer_model: str = None
-    retro_gpt_vocab_file: str = None
-    retro_gpt_merge_file: str = None
-    retro_gpt_seq_length: int = None
-    retro_gpt_global_batch_size: int = None
-    retro_gpt_chunk_length: int = 64
-
-    # Bert.
-    retro_bert_tokenizer_type: str = None
-    retro_bert_vocab_file: str = None
-    retro_bert_batch_size: int = 128
-    retro_bert_max_chunk_length: int = 256
-
-    # Index.
-    retro_index_type: str = 'faiss-par-add'
-    retro_index_str: str = None
-    retro_index_ntrain: int = None
-    retro_index_train_load_fraction: float = 1.0
-    retro_index_add_load_fraction: float = 1.0
-    retro_index_delete_training_embeddings: bool = True
-    retro_index_delete_added_codes: bool = True
-
-    # Query.
-    retro_query_ef_search: int = 256
-    retro_query_nprobe: int = 65536
-    retro_query_num_neighbors_query: int = 200
-    retro_query_num_neighbors_save: int = 20
-
-    # Tools.
-    retro_bert_embedders: RetroBertEmbedders = None
-    retro_gpt_chunk_datasets: RetroGPTChunkDatasets = None
-    retro_tokenizers: RetroTokenizers = None
-
-    def __post_init__(self) -> None:
-        """Validate Retro config."""
-
-        # Validate required attributes.
-        assert self.retro_project_dir is not None
-        assert self.retro_tasks is not None
-        assert self.retro_gpt_data_path is not None or self.retro_gpt_data_cache_path is not None
-        assert self.retro_gpt_train_samples is not None
-        assert self.retro_gpt_eval_interval is not None
-        assert self.retro_gpt_eval_iters is not None
-        assert self.retro_gpt_tokenizer_type is not None
-        assert self.retro_gpt_tokenizer_model is not None or (
-            self.retro_gpt_vocab_file is not None and self.retro_gpt_merge_file is not None
-        )
-        assert self.retro_gpt_seq_length is not None
-        assert self.retro_gpt_global_batch_size is not None
-        assert self.retro_bert_tokenizer_type is not None
-        assert self.retro_bert_vocab_file is not None
-        assert self.retro_index_str is not None
-        assert self.retro_index_ntrain is not None
-
-        # Split retro tasks.
-        self.retro_tasks = self.retro_tasks.split(",")
diff --git a/megatron/core/datasets/retro/config/gpt_chunk_datasets.py b/megatron/core/datasets/retro/config/gpt_chunk_datasets.py
deleted file mode 100644
index 831b1d812bf..00000000000
--- a/megatron/core/datasets/retro/config/gpt_chunk_datasets.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Container dataclass for GPT chunk datasets (train, valid, and test)."""
-
-from dataclasses import dataclass
-
-
-@dataclass
-class RetroGPTChunkDatasets:
-    """Container dataclass for GPT chunk datasets."""
-
-    # Each dict contains 'dataset', 'neighbor_dir', and 'num_active_chunks'.
-    train: dict = None
-    valid: dict = None
-    test: dict = None
diff --git a/megatron/core/datasets/retro/config/tokenizers.py b/megatron/core/datasets/retro/config/tokenizers.py
deleted file mode 100644
index 69ca94b3cfd..00000000000
--- a/megatron/core/datasets/retro/config/tokenizers.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Container class for GPT and Bert tokenizers."""
-
-from dataclasses import dataclass
-
-from megatron.core.tokenizers import MegatronTokenizerBase
-
-
-@dataclass
-class RetroTokenizers:
-    """Container class for GPT and Bert tokenizers."""
-
-    gpt: MegatronTokenizerBase = None
-    bert: MegatronTokenizerBase = None
diff --git a/megatron/core/datasets/retro/db/__init__.py b/megatron/core/datasets/retro/db/__init__.py
deleted file mode 100644
index f1f460b3b02..00000000000
--- a/megatron/core/datasets/retro/db/__init__.py
+++ /dev/null
@@ -1,9 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-Exports:
-
-  - build_db: Build a chunk database from a list of indexed datasets.
-"""
-
-from .build import build_db
diff --git a/megatron/core/datasets/retro/db/build.py b/megatron/core/datasets/retro/db/build.py
deleted file mode 100644
index 0cd94729385..00000000000
--- a/megatron/core/datasets/retro/db/build.py
+++ /dev/null
@@ -1,649 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Build a chunk database from a list of indexed datasets.
-
-Building a chunk database consists of.
-
-  - Breaking each document of each indexed dataset into consecutive
-      retro_gpt_chunk_length chunks.
-  - Re-tokenize each chunk into Bert, and discard any chunks with empty Bert
-      tokens.
-  - Save chunk offsets to disk for each indexed dataset.
-"""
-
-import os
-import types
-from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Tuple
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.indexed_dataset import IndexedDataset
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.utils import (
-    extract_data_config,
-    get_blocks_by_rank,
-    log_retro_rank_0,
-    retro_makedir,
-)
-
-from .utils import (
-    get_indexed_dataset_infos,
-    get_indexed_dataset_infos_path,
-    get_individual_chunk_db,
-    get_individual_db_dir,
-    get_individual_db_paths,
-    get_individual_doc_offsets,
-    get_merged_db_path_map,
-    init_indexed_dataset_infos,
-    save_indexed_dataset_infos,
-)
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-
-def build_partial_db(
-    config: types.SimpleNamespace,
-    dataset_idx: int,
-    n_datasets: int,
-    indexed_dataset: IndexedDataset,
-    block_id: int,
-    n_blocks: int,
-    block: dict,
-    proc_id: int,
-    n_procs: int,
-) -> Tuple[int, list, list, dict]:
-    """Process a document index range of the indexed dataset.
-
-    The chunk database is built in parallel blocks, since de-tokenizing &
-    re-tokenizing for Bert-length computation is expensive. This method
-    iterates each document and extracts sequential 'chunk-length' sequences
-    from each document.
-
-    Args:
-        config (types.SimpleNamespace): Subset of Retro config, containing
-            'chunk_length', 'gpt_eod', 'gpt_detokenize', 'bert_tokenize', and 'task_validate'.
-        dataset_idx (int): Index of this dataset out of all blended datasets.
-        n_datasets (int): Total number of blended datasets.
-        indexed_dataset (IndexedDataset): Indexed dataset to be chunked.
-        block_id (int): Block index out of all blocks to be processed.
-        n_blocks (int):  Total number of blocks to be processed.
-        block (dict): Range information such as start/end points for chunking idnexed dataset.
-        proc_id (int): Process ID for tracking parallel process order.
-        n_procs (int): Total number of parallel processes.
-
-    Returns:
-        A tuple containing:
-
-        - Process ID.
-        - List of valid chunks.
-        - List of invalid chunks (i.e., chunks that converted to empty Bert embeddings.).
-        - Dict mapping document ID to number of valid chunks.
-    """
-
-    if not HAVE_TQDM:
-        raise ImportError("tqdm is required to use the RetroDataset. Please install tqdm.")
-
-    # Document start/end indexes.
-    doc_range = block["range"]
-    n_docs = doc_range[1] - doc_range[0]
-    n_docs_per_proc = int(np.ceil(n_docs / n_procs))
-    doc_start_id = doc_range[0] + proc_id * n_docs_per_proc
-    doc_end_id = min(doc_range[1], doc_start_id + n_docs_per_proc)
-
-    # Print progress.
-    progress_proc_ids = set(range(n_procs)) if torch.distributed.get_rank() == 0 else set()
-    if proc_id in progress_proc_ids:
-        log_retro_rank_0(
-            " > building partial chunk db, proc %d / %d, docs %d:%d / %d."
-            % (proc_id, n_procs, doc_start_id, doc_end_id, n_docs)
-        )
-
-    # Progress bars (snapshot of overall progress).
-    doc_id_iter = range(doc_start_id, doc_end_id)
-    pbar = (
-        tqdm(doc_id_iter, "parse doc chunks", miniters=len(doc_id_iter) // 20)
-        if proc_id in progress_proc_ids
-        else doc_id_iter
-    )
-
-    # Iterate documents & parse chunks.
-    chunk_db_valid: List[Tuple] = []
-    chunk_db_invalid: List[Tuple] = []
-    doc_size_map = {}
-    for doc_id in pbar:
-        # Progress description.
-        try:
-            pbar.set_description(
-                "%sds %d / %d, block %d / %d, proc %d / %d."
-                % (
-                    "" if config.task_validate is None else "[validate] ",
-                    dataset_idx,
-                    n_datasets,
-                    block_id,
-                    n_blocks,
-                    proc_id,
-                    n_procs,
-                )
-            )
-        except Exception:
-            pass
-
-        # Remove EOD token.
-        doc = indexed_dataset.get(doc_id)
-        if doc[-1].item() == config.gpt_eod:
-            doc = doc[:-1]
-        doc_len = len(doc)
-
-        # Chunk start/end indexes.
-        chunk_start_idxs = list(range(0, doc_len, config.chunk_length))
-        chunk_end_idxs = [min(doc_len, s + config.chunk_length) for s in chunk_start_idxs]
-
-        # Re-tokenize each chunk to Bert/Wordpiece (empty bert -> 'invalid').
-        doc_size_map[doc_id] = 0
-        for i, chunk_start_idx in enumerate(chunk_start_idxs):
-            # Re-tokenize.
-            chunk_end_idx = chunk_end_idxs[i]
-            gpt_token_ids = indexed_dataset.get(
-                idx=doc_id, offset=chunk_start_idx, length=chunk_end_idx - chunk_start_idx
-            )
-            text = config.gpt_detokenize(gpt_token_ids.tolist())
-            bert_token_ids = config.bert_tokenize(text)
-
-            # 'Valid' for non-empty Bert chunks; 'invalid' otherwise.
-            if len(bert_token_ids) == 0:
-                _chunk_db = chunk_db_invalid
-            else:
-                _chunk_db = chunk_db_valid
-                doc_size_map[doc_id] += 1
-            _chunk_db.append((doc_id, chunk_start_idx, chunk_end_idx, len(bert_token_ids)))
-
-    return proc_id, chunk_db_valid, chunk_db_invalid, doc_size_map
-
-
-def build_block_db(
-    config: RetroPreprocessingConfig,
-    dataset_idx: int,
-    n_datasets: int,
-    indexed_dataset: IndexedDataset,
-    n_procs: int,
-    executor: ProcessPoolExecutor,
-    n_missing_blocks: int,
-    block_idx: int,
-    block: dict,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Split each document within block into consecutive retro_gpt_chunk_length size chunks.
-
-    Args:
-        config (RetroPreprocessingConfig): For DB building, we make use of attributes
-            'chunk_length', 'gpt_eod', 'gpt_detokenize', 'bert_tokenize', and 'task_validate'.
-        dataset_idx (int): Index of this dataset out of all blended datasets.
-        n_datasets (int): Total number of blended datasets.
-        indexed_dataset (IndexedDataset): Indexed dataset to be chunked.
-        n_procs (int): Total number of parallel processes.
-            executor (ProcessPoolExecutor): Executor for launching parallel processes.
-        n_missing_blocks (int):  Total number of blocks to be processed.
-        block_idx (int): Block index out of all blocks to be processed.
-        block (dict): Range information such as start/end points for chunking idnexed dataset.
-
-    Returns:
-        A tuple containing:
-
-        - List of valid chunks.
-        - List of invalid chunks (i.e., chunks that converted to empty Bert embeddings.).
-        - Dict mapping document ID to number of valid chunks.
-    """
-
-    # Build partial dbs.
-    log_retro_rank_0(" > build partial dbs.")
-    futures = []
-    for proc_id in range(n_procs):  # not true process id
-        futures.append(
-            executor.submit(
-                build_partial_db,
-                types.SimpleNamespace(
-                    chunk_length=config.retro_gpt_chunk_length,
-                    gpt_eod=config.retro_tokenizers.gpt.eod,
-                    gpt_detokenize=config.retro_tokenizers.gpt.detokenize,
-                    bert_tokenize=config.retro_tokenizers.bert.tokenize,
-                    task_validate=config.retro_task_validate,
-                ),
-                dataset_idx,
-                n_datasets,
-                indexed_dataset,
-                block_idx,
-                n_missing_blocks,
-                block,
-                proc_id,
-                n_procs,
-            )
-        )
-    partial_chunk_dbs = []
-    for future in as_completed(futures):
-        partial_chunk_dbs.append(future.result())
-
-    # Concatenate chunks.
-    partial_chunk_dbs.sort(key=lambda item: item[0])  # sort by proc_id
-    chunk_db_valid = [
-        item for partial_chunk_db in partial_chunk_dbs for item in partial_chunk_db[1]
-    ]
-    chunk_db_invalid = [
-        item for partial_chunk_db in partial_chunk_dbs for item in partial_chunk_db[2]
-    ]
-
-    # Convert to numpy.
-    log_retro_rank_0(" > converting chunk db to numpy.")
-    chunk_db_valid = np.array(chunk_db_valid, dtype="uint32")
-    chunk_db_invalid = np.array(chunk_db_invalid, dtype="uint32")
-
-    # Document offsets.
-    doc_sizes = [
-        (d, s) for partial_chunk_db in partial_chunk_dbs for d, s in partial_chunk_db[3].items()
-    ]
-    doc_sizes.sort(key=lambda item: item[0])
-    doc_offsets = np.cumsum([item[1] for item in doc_sizes]).astype("uint64")
-    doc_offsets = np.stack(
-        (np.array([item[0] for item in doc_sizes], dtype="uint64"), doc_offsets), axis=1
-    )
-
-    return chunk_db_valid, chunk_db_invalid, doc_offsets
-
-
-def save_block_db(
-    block: dict, chunk_db_valid: np.ndarray, chunk_db_invalid: np.ndarray, doc_offsets: np.ndarray
-) -> None:
-    """Save block of chunked tokens to disk. These blocks are later used for
-    training and adding to the vector index.
-
-    Args:
-        block (dict): Range information such as start/end points for chunking idnexed dataset.
-        chunk_db_valid (np.ndarray): Array of valid chunk indexes.
-        chunk_db_invalid (np.ndarray): Array of invalid chunk indexes.
-        doc_offsets (np.ndarray): Array of document offsets by chunks.
-    """
-    if not HAVE_H5PY:
-        raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-    log_retro_rank_0(" > saving individual db.")
-    with h5py.File(block["path"], "w") as f:
-        dset = f.create_dataset("chunks_valid", data=chunk_db_valid)
-        dset = f.create_dataset("chunks_invalid", data=chunk_db_invalid)
-        dset = f.create_dataset("doc_offsets", data=doc_offsets)
-
-
-def build_individual_db(
-    config: RetroPreprocessingConfig, dataset_idx: int, n_datasets: int, dataset_info: dict
-) -> None:
-    """Process a single indexed dataset & extract chunks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        dataset_idx (int): Dataset index within blended dataset.
-        n_datasets (int): Total number of datasets within blended dataset.
-        dataset_info (dict): Metadata for dataset
-            (see `save_indexed_dataset_infos()` in `utils.py` for more detail).
-    """
-
-    # Make directory.
-    db_dir = get_individual_db_dir(config.retro_project_dir, dataset_info["prefix"])
-    retro_makedir(config, db_dir)
-
-    # Indexed dataset.
-    indexed_dataset = dataset_info["dataset"]
-
-    # Missing DB blocks (split by documents).
-    blocks = get_blocks_by_rank(
-        db_dir,
-        len(indexed_dataset),
-        config.retro_doc_block_size,
-        validate=lambda f: f["chunks_valid"].shape == (0,) or f["chunks_valid"].shape[1] == 4,
-        sample=config.retro_task_validate,
-    )
-    if config.retro_task_validate is None:
-        active_blocks = blocks.missing
-    else:
-        assert blocks.n_missing_world == 0
-        active_blocks = blocks.existing
-
-    # Prevent missing-path-write race condition.
-    torch.distributed.barrier()
-
-    # Nothing to do?
-    if config.retro_task_validate is None and not active_blocks:
-        return
-
-    # Num processes.
-    if blocks.n_missing_world == 1:
-        n_procs = 128
-    elif blocks.n_missing_world <= 2:
-        n_procs = 64
-    elif blocks.n_missing_world <= 4:
-        n_procs = 32
-    elif blocks.n_missing_world <= 8:
-        n_procs = 16
-    else:
-        n_procs = 8
-
-    # Process documents in parallel.
-    with ProcessPoolExecutor(max_workers=n_procs) as executor:
-        for block_idx, block in enumerate(active_blocks):
-            if block is not None:
-                # Build block DB.
-                chunk_db_valid, chunk_db_invalid, doc_offsets = build_block_db(
-                    config=config,
-                    dataset_idx=dataset_idx,
-                    n_datasets=n_datasets,
-                    indexed_dataset=indexed_dataset,
-                    n_procs=n_procs,
-                    executor=executor,
-                    n_missing_blocks=len(active_blocks),
-                    block_idx=block_idx,
-                    block=block,
-                )
-
-                if config.retro_task_validate is None:
-                    # Save block DB.
-                    save_block_db(
-                        block=block,
-                        chunk_db_valid=chunk_db_valid,
-                        chunk_db_invalid=chunk_db_invalid,
-                        doc_offsets=doc_offsets,
-                    )
-
-                else:
-                    # Load existing block DB.
-                    with h5py.File(block["path"]) as f:
-                        existing_chunks_valid = np.copy(f["chunks_valid"])
-                        existing_chunks_invalid = np.copy(f["chunks_invalid"])
-                        existing_doc_offsets = np.copy(f["doc_offsets"])
-
-                    # Check equality.
-                    log_retro_rank_0(" > validate.")
-                    assert np.array_equal(existing_chunks_valid, chunk_db_valid)
-                    assert np.array_equal(existing_chunks_invalid, chunk_db_invalid)
-                    assert np.array_equal(existing_doc_offsets, doc_offsets)
-
-            # Wait for all ranks to finish block.
-            log_retro_rank_0(" > waiting for all ranks to finish block.")
-            torch.distributed.barrier()
-
-    log_retro_rank_0(" > finished saving individual db.")
-
-
-def build_individual_dbs(
-    config: RetroPreprocessingConfig, indexed_dataset_infos: List[Dict]
-) -> None:
-    """Iterate each indexed dataset & process its chunks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset.
-    """
-
-    # Build individual DBs.
-    log_retro_rank_0(" > build individual chunk dbs.")
-    for ds_idx, ds_info in enumerate(indexed_dataset_infos):
-        # Progress.
-        log_retro_rank_0(
-            " > building individual db, dataset %d / %d ... '%s'."
-            % (ds_idx, len(indexed_dataset_infos), ds_info["prefix"])
-        )
-
-        # Process single dataset.
-        build_individual_db(config, ds_idx, len(indexed_dataset_infos), ds_info)
-
-
-def update_chunk_counts(
-    config: RetroPreprocessingConfig, indexed_dataset_infos: List[Dict]
-) -> None:
-    """Set n_chunks_train & n_chunks sampled for each individual DB.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset
-            (i.e., 'prefix', 'ratio', 'n_chunks', etc.).
-    """
-
-    if torch.distributed.get_rank() != 0:
-        return
-
-    # Data ratio sum (for setting index training chunks).
-    data_ratio_sum = sum([d["ratio"] for d in indexed_dataset_infos])
-
-    # Training split size (split at document level).
-    train_fraction = float(extract_data_config(config).split.split(",")[0]) / 100
-    assert train_fraction > 0 and train_fraction <= 1
-
-    # Set n_chunks (including n_chunks_sampled for unambiguity).
-    log_retro_rank_0(" > compute n_chunks.")
-    for ds_index, ds_info in enumerate(indexed_dataset_infos):
-        db_paths = get_individual_db_paths(config.retro_project_dir, ds_info["prefix"])
-
-        # Update counts.
-        ds_info["n_docs"] = len(ds_info["dataset"].document_indices) - 1
-        ds_info["n_docs_train"] = int(train_fraction * ds_info["n_docs"])
-        ds_info["n_chunks"] = 0  # previously, 'n_chunks_valid'
-        ds_info["n_chunks_train"] = 0
-        ds_info["n_chunks_invalid"] = 0
-        for db_path in tqdm(
-            db_paths, "%d/%d, %s" % (ds_index, len(indexed_dataset_infos), ds_info["prefix"])
-        ):
-            with h5py.File(db_path, "r") as f:
-                ds_info["n_chunks"] += len(f["chunks_valid"])
-                ds_info["n_chunks_invalid"] += len(f["chunks_invalid"])
-                ds_info["n_chunks_train"] += (
-                    (np.copy(f["chunks_valid"][:, 0]) < ds_info["n_docs_train"]).sum().item()
-                )
-
-        ds_info["n_chunks_sampled"] = int(
-            config.retro_index_ntrain * ds_info["ratio"] / data_ratio_sum
-        )
-
-        # Verify counts.
-        assert ds_info["n_chunks_train"] <= ds_info["n_chunks"], "n_train (%d) > n_total (%d)." % (
-            ds_info["n_chunks_train"],
-            ds_info["n_chunks"],
-        )
-        assert (
-            ds_info["n_chunks_sampled"] <= ds_info["n_chunks_train"]
-        ), "n_sampled (%d) > n_train (%d)." % (
-            ds_info["n_chunks_sampled"],
-            ds_info["n_chunks_train"],
-        )
-
-
-def merge_dbs(project_dir: str, indexed_dataset_infos: List[Dict], db_type: str) -> None:
-    """Merge individual DBs into single DB.
-
-    Args:
-        project_dir (str): Retro project dir.
-        indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset
-            (i.e., 'prefix', 'ratio', 'n_chunks', etc.).
-        db_type (str): DB type (e.g., 'sampled', 'train', or 'valid').
-    """
-
-    if not HAVE_H5PY:
-        raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-    if torch.distributed.get_rank() != 0:
-        return
-
-    log_retro_rank_0(" > build %s chunk db." % db_type)
-
-    # Count chunks.
-    if db_type == "sampled":
-        n_chunks_key = "n_chunks_sampled"
-        n_docs_key = None
-    elif db_type == "train":
-        n_chunks_key = "n_chunks_train"
-        n_docs_key = "n_docs_train"
-    elif db_type == "valid":
-        n_docs_key = None
-    else:
-        raise Exception("handle db_type '%s'." % db_type)
-
-    if db_type == "valid":
-        n_chunks = sum(m["n_chunks"] - m["n_chunks_train"] for m in indexed_dataset_infos)
-    else:
-        n_chunks = sum(m[n_chunks_key] for m in indexed_dataset_infos)
-        n_docs = None if n_docs_key is None else sum(m[n_docs_key] for m in indexed_dataset_infos)
-
-    # DB path.
-    db_path = get_merged_db_path_map(project_dir)[db_type]
-
-    # Delete existing chunk db if incorrect size.
-    if os.path.exists(db_path):
-        try:
-            f = h5py.File(db_path)
-            n_alloc = len(f["chunks"])  # total allocated
-            n_written = f["n_written"][0].item()  # total written
-            f.close()
-
-            if n_chunks != n_alloc or n_chunks != n_written:
-                os.remove(db_path)
-
-        except Exception as e:
-            if isinstance(e, OSError):
-                os.remove(db_path)
-            elif isinstance(e, KeyError):
-                f.close()
-                os.remove(db_path)
-            else:
-                raise e
-
-    # Build merged chunk db.
-    if not os.path.exists(db_path):
-        os.makedirs(os.path.dirname(db_path), exist_ok=True)
-        f = h5py.File(db_path, "w")
-
-        # Initialize output arrays.
-        merged_chunk_db: np.ndarray = f.create_dataset("chunks", (n_chunks, 5), dtype="uint32")
-        merged_doc_offsets: np.ndarray = (
-            None
-            if n_docs_key is None
-            else f.create_dataset("doc_offsets", (n_docs, 3), dtype="uint64")
-        )
-        n_written = f.create_dataset("n_written", (1,), dtype="uint64")
-        n_written[0] = 0
-
-        # Iterate indexed datasets & collect chunks.
-        chunk_start_index = 0
-        doc_start_index = 0
-        doc_start_offset = 0
-        for ds_idx, ds_info in enumerate(indexed_dataset_infos):
-            log_retro_rank_0(
-                " > merging dbs; '%s', dataset %d / %d ... '%s'."
-                % (db_type, ds_idx, len(indexed_dataset_infos), ds_info["prefix"])
-            )
-            individual_chunk_db: np.ndarray = get_individual_chunk_db(project_dir, ds_idx, ds_info)
-            individual_doc_offsets: np.ndarray = (
-                None
-                if n_docs_key is None
-                else get_individual_doc_offsets(project_dir, ds_idx, ds_info)
-            )
-
-            if db_type == "valid":
-                individual_chunk_db = individual_chunk_db[ds_info["n_chunks_train"] :]
-                if n_docs_key is None:
-                    individual_doc_offsets = None
-                else:
-                    train_doc_offset = individual_doc_offsets[ds_info["n_docs_train"] - 1, 2]
-                    individual_doc_offsets = np.copy(
-                        individual_doc_offsets[ds_info["n_docs_train"] :]
-                    )
-                    individual_doc_offsets[:, 2] -= train_doc_offset
-
-                    log_retro_rank_0("~~~")
-                    log_retro_rank_0(individual_doc_offsets)
-                    log_retro_rank_0(train_doc_offset)
-                    raise Exception("test me.")
-            else:
-                individual_chunk_db = individual_chunk_db[: ds_info[n_chunks_key]]
-                individual_doc_offsets = (
-                    None
-                    if n_docs_key is None
-                    else np.copy(individual_doc_offsets[: ds_info[n_docs_key]])
-                )
-
-            merged_chunk_db[chunk_start_index : chunk_start_index + len(individual_chunk_db)] = (
-                individual_chunk_db
-            )
-            chunk_start_index += len(individual_chunk_db)
-            n_written[0] = chunk_start_index
-            if n_docs_key is not None:
-                individual_doc_offsets[:, 2] += doc_start_offset
-                doc_end_index = doc_start_index + individual_doc_offsets.shape[0]
-                merged_doc_offsets[doc_start_index:doc_end_index] = individual_doc_offsets
-                doc_start_index = doc_end_index
-                doc_start_offset = individual_doc_offsets[-1, 2].item()
-
-        f.close()
-
-
-def build_merged_dbs(project_dir: str, indexed_dataset_infos: List[Dict]) -> None:
-    """Merge individual dataset components into single database.
-
-    This method merges databases for DB types:
-    - 'sampled': used for training the vector index.
-    - 'train': used for adding to the trained vector index.
-    - 'valid': can be used for validating/testing the vector index.
-
-    Args:
-        project_dir (str): Retro project dir.
-        indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset
-            (i.e., 'prefix', 'ratio', 'n_chunks', etc.).
-    """
-    merge_dbs(project_dir, indexed_dataset_infos, "sampled")
-    merge_dbs(project_dir, indexed_dataset_infos, "train")
-    merge_dbs(project_dir, indexed_dataset_infos, "valid")
-
-
-def build_db(config: RetroPreprocessingConfig) -> None:
-    """Extract token chunks from each indexed dataset.
-
-    Iterate each document of each indexed dataset, extract that document's chunks,
-        and save to a 'DB' (hdf5 file).
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    project_dir = config.retro_project_dir
-
-    # Indexed dataset info.
-    if config.retro_task_validate is None:
-        indexed_dataset_infos = init_indexed_dataset_infos(config)
-    else:
-        indexed_dataset_infos = get_indexed_dataset_infos(config.retro_project_dir)
-    # Build individual dbs.
-    build_individual_dbs(config, indexed_dataset_infos)
-
-    # If validating, return here.
-    if config.retro_task_validate is not None:
-        return
-
-    # Single-process going forward.
-    if torch.distributed.get_rank() != 0:
-        return
-
-    # Update n_chunks & save indexed dataset infos.
-    if not os.path.exists(get_indexed_dataset_infos_path(project_dir)):
-        update_chunk_counts(config, indexed_dataset_infos)
-        save_indexed_dataset_infos(project_dir, indexed_dataset_infos)
-    indexed_dataset_infos = get_indexed_dataset_infos(project_dir)
-
-    # Builded merged dbs.
-    build_merged_dbs(project_dir, indexed_dataset_infos)
diff --git a/megatron/core/datasets/retro/db/dataset.py b/megatron/core/datasets/retro/db/dataset.py
deleted file mode 100644
index 61b62601d8c..00000000000
--- a/megatron/core/datasets/retro/db/dataset.py
+++ /dev/null
@@ -1,114 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""A DBDataset is for iterating the chunks of the chunk database.
-
-This dataset is used for both training a vector index, and adding vectors to a
-trained index.
-"""
-
-from typing import List
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.indexed_dataset import IndexedDataset
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-
-class DBDataset(torch.utils.data.Dataset):
-    """Dataset for iterating chunks.
-
-    Args:
-        db_path (str): Path of HDF5-format chunk database.
-        indexed_datasets (List[IndexedDataset]): Indexed datasets used to build database.
-        chunks (np.ndarray): Array of chunk indexes, for indexing into indexed datasets.
-            Format [dataset_idx, doc_id, start_idx, end_idx, bert_length].
-        chunk_length (int): Max GPT chunk length (e.g., 64).
-        eod_token_id (int): EOD token ID.
-    """
-
-    def __init__(
-        self,
-        db_path: str,
-        indexed_datasets: List[IndexedDataset],
-        chunks: np.ndarray,
-        chunk_length: int,
-        eod_token_id: int,
-    ):
-        assert chunks.shape[1] == 5, (
-            "expected 5 columns (dataset_idx, "
-            "doc_idx, token_start_idx, token_end_idx, bert_chunk_length); "
-            "found %d columns." % chunks.shape[1]
-        )
-
-        self.db_path = db_path
-        self.indexed_datasets = indexed_datasets
-        self.chunks = chunks
-        self.doc_chunk_map = None
-
-        self.max_chunk_length = chunk_length
-        self.eod_token_id = eod_token_id
-
-    def __len__(self) -> int:
-        """Length of DB dataset.
-
-        Returns:
-            Number of chunks contained in the dataset.
-        """
-        return self.chunks.shape[0]
-
-    def __getitem__(self, chunk_id: int) -> dict:
-        """DB dataset sample.
-
-        Args:
-            chunk_id (int): Index of chunk within dataset.
-
-        Returns:
-            A dict containing:
-            - 'doc_id': Document index within indexed dataset.
-            - 'text': GPT token IDs.
-        """
-
-        # Chunk start/end indexes.
-        indexed_dataset_id, doc_id, token_start_idx, token_end_idx, _ = [
-            value.item() for value in self.chunks[chunk_id]
-        ]
-        chunk_length = token_end_idx - token_start_idx
-        indexed_dataset = self.indexed_datasets[indexed_dataset_id]
-
-        # Chunk token ids.
-        token_ids = indexed_dataset.get(doc_id, offset=token_start_idx, length=chunk_length)
-
-        # Extend chunks to max_chunk_length by padding with EOD tokens.
-        if chunk_length != self.max_chunk_length:
-            assert chunk_length < self.max_chunk_length, "invalid chunk len."
-            token_ids = token_ids.tolist()
-            token_ids += [self.eod_token_id] * (self.max_chunk_length - chunk_length)
-
-        return {"doc_id": doc_id, "text": np.array(token_ids, dtype=np.int64)}
-
-    def load_doc_tuples(self) -> None:
-        """Load the dataset & document ids.
-
-        Load the dataset id & document id of each chunk in the database, to
-        be used for causality filtering during querying.
-        """
-        if not HAVE_TQDM:
-            raise ImportError("tqdm is required to use the DBDataset. Please install tqdm.")
-
-        self.doc_tuples = np.zeros(shape=(len(self), 2), dtype="uint32")
-        block_size = int(1e6)
-        for start_idx in tqdm(
-            range(0, len(self), block_size),
-            "load doc tuples",
-            miniters=(len(self) // block_size) // 10,
-            disable=torch.distributed.get_rank() != 0,
-        ):
-            end_idx = min(len(self), start_idx + block_size)
-            self.doc_tuples[start_idx:end_idx] = self.chunks[start_idx:end_idx, :2]
diff --git a/megatron/core/datasets/retro/db/utils.py b/megatron/core/datasets/retro/db/utils.py
deleted file mode 100644
index 7906f4bf9ec..00000000000
--- a/megatron/core/datasets/retro/db/utils.py
+++ /dev/null
@@ -1,398 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Utilities for building a chunk database."""
-
-import glob
-import json
-import os
-from typing import Dict, List, Optional
-
-import numpy as np
-
-from megatron.core.datasets.indexed_dataset import IndexedDataset
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.models.retro.utils import get_gpt_data_dir
-
-from .dataset import DBDataset
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-
-def get_db_dir(project_dir: str) -> str:
-    """Sub-directory for DB data.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-
-    Returns:
-        Path of the DB sub-directory within the project.
-    """
-    return os.path.join(project_dir, "db")
-
-
-def init_indexed_dataset_infos(config: RetroPreprocessingConfig) -> List[Dict]:
-    """Gather meta-info about each indexed dataset.
-
-    The returned info array allows for easy access to the configuration, and
-    helps remove ambiguity.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        List of processing metadata for each dataset, including:
-        - ratio: Data split weight.
-        - prefix: Relative path to dataset under DB sub-directory.
-    """
-
-    data_dir = get_gpt_data_dir(config.retro_project_dir)
-    data_blend: List[str] = config.retro_gpt_data_path
-    assert len(data_blend) % 2 == 0, "currently, only blended dataset is supported."
-
-    # Dataset infos.
-    infos = []
-    for i in range(0, len(data_blend), 2):
-        ratio = float(data_blend[i])
-        prefix = data_blend[i + 1]
-        path = os.path.join(data_dir, prefix + ".bin")
-        assert os.path.exists(path), "couldn't find '%s'." % path
-        infos.append({"ratio": ratio, "prefix": prefix})
-
-    # Load indexed datasets.
-    load_indexed_datasets(config.retro_project_dir, infos)
-
-    return infos
-
-
-def get_indexed_dataset_infos_path(project_dir: str) -> str:
-    """Path to indexed dataset meta-infos.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-
-    Returns:
-        Path to the `indexed_dataset_infos.json` file.
-    """
-    return os.path.join(get_db_dir(project_dir), "indexed_dataset_infos.json")
-
-
-def save_indexed_dataset_infos(project_dir: str, indexed_dataset_infos: List[Dict]) -> None:
-    """Save dataset order & meta-info.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        indexed_dataset_infos (List[Dict]): List of metadata for each dataset,
-            with each entry containing:
-
-        - ratio: Data split weight.
-        - prefix: Relative path to dataset under DB sub-directory.
-        - n_docs: Number of documents.
-        - n_docs_train: Number of documents used for pretraining.
-        - n_chunks: Number of valid chunks.
-        - n_chunks_train: Number of valid chunks used for pretraining.
-        - n_chunks_invalid: Number of invalid chunks.
-        - n_chunks_sampled: Number of valid chunks used for vector index training.
-    """
-
-    # Remove 'dataset' field.
-    clean_infos = []
-    for info in indexed_dataset_infos:
-        info = dict(info)
-        del info["dataset"]
-        clean_infos.append(info)
-
-    # Save.
-    with open(get_indexed_dataset_infos_path(project_dir), "w") as f:
-        json.dump(clean_infos, f, indent=4)
-
-
-def load_indexed_datasets(project_dir: str, indexed_dataset_infos: List[Dict]) -> None:
-    """Loaded indexed datasets into memory-mapped datasets.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        indexed_dataset_infos (List[Dict]): List of metadata for each dataset
-            (see `save_indexed_dataset_infos()` for more details.
-    """
-    data_dir = get_gpt_data_dir(project_dir)
-    for info in indexed_dataset_infos:
-        info["dataset"] = IndexedDataset(os.path.join(data_dir, info["prefix"]), mmap=True)
-
-
-def get_indexed_dataset_infos(project_dir: str) -> List[Dict]:
-    """Load indexed dataset meta-infos.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-
-    Returns:
-        List of metadata for each dataset (see `save_indexed_dataset_infos()` for more details.
-    """
-
-    # Load json.
-    path = get_indexed_dataset_infos_path(project_dir)
-    with open(path) as f:
-        infos = json.load(f)
-
-    # Load indexed datasets.
-    load_indexed_datasets(project_dir, infos)
-
-    return infos
-
-
-def get_individual_db_dir(project_dir: str, prefix: str) -> str:
-    """Individual DB's directory.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        prefix (str): Unique relative path to dataset within project dir.
-
-    Returns:
-        Path to the given datasets's chunk database.
-    """
-    return os.path.join(get_db_dir(project_dir), "individual", prefix)
-
-
-def get_individual_db_paths(project_dir: str, prefix: str) -> List[str]:
-    """Get paths of all database blocks of an individual dataset.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        prefix (str): Unique relative path to dataset within project dir.
-
-    Returns:
-        Paths to each HDF5 chunk database files that comprises this datasets full chunk database.
-    """
-    return sorted(glob.glob(get_individual_db_dir(project_dir, prefix) + "/*hdf5"))
-
-
-def get_individual_chunk_db(project_dir: str, ds_id: int, ds_info: dict) -> np.ndarray:
-    """Load individual dataset's chunk DB.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        ds_id (int): Index of dataset within blended dataset.
-        ds_info (dict): Preprocessing metadata for dataset
-            (see `save_indexed_dataset_infos()` for more detail).
-
-    Returns:
-        Array of chunk start/end indexes for this dataset,
-            where the chunk indexes can be used for indexing into
-            the corresponding indexed dataset.
-    """
-
-    if not HAVE_H5PY:
-        raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-    paths = get_individual_db_paths(project_dir, ds_info["prefix"])
-    # *Note*: convert to dataset, rather than copying to memory.
-    db = np.zeros((ds_info["n_chunks"], 5), dtype="uint32")
-    db[:, 0] = ds_id
-    start_idx = 0
-    for path in paths:
-        f = h5py.File(path, "r")
-        n_chunks_current = f["chunks_valid"].shape[0]
-        db[start_idx : (start_idx + n_chunks_current), 1:] = f["chunks_valid"]
-        start_idx += n_chunks_current
-        f.close()
-
-    assert start_idx == ds_info["n_chunks"]
-
-    return db
-
-
-def get_individual_doc_offsets(project_dir: str, ds_id: int, ds_info: dict) -> np.ndarray:
-    """Load individual dataset's document offsets.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        ds_id (int): Index of dataset within blended dataset.
-        ds_info (dict): Preprocessing metadata for dataset
-            (see `save_indexed_dataset_infos()` for more detail).
-
-    Returns:
-        Array of document offsets by chunk index for this dataset.
-    """
-
-    if not HAVE_H5PY:
-        raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-    paths = get_individual_db_paths(project_dir, ds_info["prefix"])
-    # *Note*: convert to dataset, rather than copying to memory.
-    doc_offsets = np.zeros((ds_info["n_docs"], 3), dtype="uint64")
-    doc_offsets[:, 0] = ds_id
-    start_idx = 0
-    start_offset = 0
-    for path in paths:
-        with h5py.File(path) as f:
-            current_doc_offsets = np.copy(f["doc_offsets"])
-            current_doc_offsets[:, 1] += start_offset
-            current_ndocs = current_doc_offsets.shape[0]
-            doc_offsets[start_idx : (start_idx + current_ndocs), 1:] = current_doc_offsets
-            start_idx += current_ndocs
-            start_offset = current_doc_offsets[-1, 1].item()
-
-    return doc_offsets
-
-
-def get_merged_db_path_map(project_dir: str) -> dict:
-    """Paths to merged datasets.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-
-    Returns:
-        A dict of chunk databases, one for each of:
-        - sampled: Chunks used for training the vector index.
-        - train: Chunks used for pretraining 'train' dataset.
-        - valid: Chunks used for pretraining 'valid' dataset.
-    """
-    base_dir = get_db_dir(project_dir)
-    return {
-        "sampled": os.path.join(base_dir, "merged", "sampled.hdf5"),
-        "train": os.path.join(base_dir, "merged", "train.hdf5"),
-        "valid": os.path.join(base_dir, "merged", "valid.hdf5"),
-    }
-
-
-def get_merged_dataset(
-    project_dir: str,
-    chunk_length: int,
-    eod_token_id: int,
-    db_type: str,
-    indexed_dataset_infos: Optional[List[Dict]] = None,
-) -> DBDataset:
-    """Get merged dataset.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        chunk_length (int): GPT chunk length (e.g., 64).
-        eod_token_id (int): EOD token ID.
-        db_type (str): DB type (e.g., 'sampled', 'train', or 'valid').
-        indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list
-            of dataset metadata (see `save_indexed_dataset_infos()` for more detail).
-            If not provided, the indexed dataset infos will be loaded from disk.
-
-    Returns:
-        A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
-    """
-    if not HAVE_H5PY:
-        raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-    if not indexed_dataset_infos:
-        indexed_dataset_infos = get_indexed_dataset_infos(project_dir)
-
-    # Load chunks.
-    db_path = get_merged_db_path_map(project_dir)[db_type]
-    f = h5py.File(db_path, "r")
-    chunks = f["chunks"]
-
-    # DB dataset.
-    indexed_datasets = [info["dataset"] for info in indexed_dataset_infos]
-    dataset = DBDataset(
-        db_path=db_path,
-        indexed_datasets=indexed_datasets,
-        chunks=chunks,
-        chunk_length=chunk_length,
-        eod_token_id=eod_token_id,
-    )
-
-    return dataset
-
-
-def get_merged_sampled_dataset(
-    project_dir: str,
-    chunk_length: int,
-    eod_token_id: int,
-    indexed_dataset_infos: Optional[List[Dict]] = None,
-) -> DBDataset:
-    """Get sampled dataset (for training the vector index).
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        chunk_length (int): GPT chunk length (e.g., 64).
-        eod_token_id (int): EOD token ID.
-        indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list
-            of dataset metadata (see `save_indexed_dataset_infos()` for more detail).
-            If not provided, the indexed dataset infos will be loaded from disk.
-
-    Returns:
-        A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
-    """
-    return get_merged_dataset(
-        project_dir, chunk_length, eod_token_id, "sampled", indexed_dataset_infos
-    )
-
-
-def get_merged_train_dataset(
-    project_dir: str,
-    chunk_length: int,
-    eod_token_id: int,
-    indexed_dataset_infos: Optional[List[Dict]] = None,
-) -> DBDataset:
-    """Get training dataset (for adding to the vector index).
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        chunk_length (int): GPT chunk length (e.g., 64).
-        eod_token_id (int): EOD token ID.
-        indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list of
-            dataset metadata (see `save_indexed_dataset_infos()` for more detail).
-            If not provided, the indexed dataset infos will be loaded from disk.
-
-    Returns:
-        A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
-    """
-    return get_merged_dataset(
-        project_dir, chunk_length, eod_token_id, "train", indexed_dataset_infos
-    )
-
-
-def get_merged_valid_dataset(
-    project_dir: str,
-    chunk_length: int,
-    eod_token_id: int,
-    indexed_dataset_infos: Optional[List[Dict]] = None,
-) -> DBDataset:
-    """Get validation dataset (for testing the vector index).
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        chunk_length (int): GPT chunk length (e.g., 64).
-        eod_token_id (int): EOD token ID.
-        indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list
-            of dataset metadata (see `save_indexed_dataset_infos()` for more detail).
-            If not provided, the indexed dataset infos will be loaded from disk.
-
-    Returns:
-        A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
-    """
-    return get_merged_dataset(
-        project_dir, chunk_length, eod_token_id, "valid", indexed_dataset_infos
-    )
-
-
-def get_merged_datasets(project_dir: str, chunk_length: int, eod_token_id: int) -> dict:
-    """Get all merged datasets.
-
-    Args:
-        project_dir (str): Path to Retro project dir.
-        chunk_length (int): GPT chunk length (e.g., 64).
-        eod_token_id (int): EOD token ID.
-
-    Returns:
-        A dict mapping DB type ('sampled', 'train', or 'valid') to the corresponding DBDataset,
-            which is a dataset that wraps the HDF5 chunk index array.
-    """
-    fns = {
-        "sampled": get_merged_sampled_dataset,
-        "train": get_merged_train_dataset,
-        "valid": get_merged_valid_dataset,
-    }
-    datasets = {key: fn(project_dir, chunk_length, eod_token_id) for key, fn in fns.items()}
-    return datasets
diff --git a/megatron/core/datasets/retro/external_libs.py b/megatron/core/datasets/retro/external_libs.py
deleted file mode 100644
index 3ac29bda2eb..00000000000
--- a/megatron/core/datasets/retro/external_libs.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Required external libraries for Retro preprocessing."""
-
-import importlib
-
-required_libs = ["faiss", "h5py", "transformers"]  # for huggingface bert
-
-for lib in required_libs:
-    try:
-        globals()[lib] = importlib.import_module(lib)
-    except ImportError as e:
-        pass
diff --git a/megatron/core/datasets/retro/index/__init__.py b/megatron/core/datasets/retro/index/__init__.py
deleted file mode 100644
index d069f55f228..00000000000
--- a/megatron/core/datasets/retro/index/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-Exports:
-
-  - train_index: Train an index on representative vectors.
-  - add_to_index: Add vectors to a trained index.
-  - build_index: Wrapper function that calls above two functions.
-"""
-
-from .build import add_to_index, build_index, train_index
diff --git a/megatron/core/datasets/retro/index/build.py b/megatron/core/datasets/retro/index/build.py
deleted file mode 100644
index f02b4288f9e..00000000000
--- a/megatron/core/datasets/retro/index/build.py
+++ /dev/null
@@ -1,339 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Construct an index.
-
-Constructing an index generally happens in two phases:
-
-  - index.train(): Train an index on a representative set of vectors.
-  - index.add(): Add vectors to an index, to be available for retrieval.
-"""
-
-import os
-import shutil
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.db.utils import (
-    get_merged_sampled_dataset,
-    get_merged_train_dataset,
-)
-from megatron.core.datasets.retro.utils import GPTToTextDataset
-
-from .factory import IndexFactory
-from .utils import (
-    get_training_data_block_dir,
-    get_training_data_block_paths,
-    get_training_data_merged_path,
-    get_training_data_root_dir,
-)
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-##################################################
-# Train index.
-##################################################
-
-
-def get_empty_index_path(config: RetroPreprocessingConfig) -> str:
-    """Path of empty index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to the empty (trained, but without added samples) vector index.
-    """
-    index = IndexFactory.get_index(config.retro_index_type)
-    empty_index_path = index.get_empty_index_path(config)
-    return empty_index_path
-
-
-def get_block_nload(block_path: str, load_fraction: float) -> int:
-    """Compute number of blocks to load.
-
-    This is computed by multiplying the total number of available blocks with the
-    fraction of blocks to load.
-
-    Args:
-        block_path (str): Path to HDF5 file containing block of data. File must contain key 'data'.
-        load_fraction (float): Fraction (0 < load_fraction <= 1) of block samples to load.
-
-    Returns:
-        Number of block samples to load.
-    """
-    if not HAVE_H5PY:
-        raise ImportError(
-            "h5py is required to use the merge_embedding_blocks function. Please install h5py."
-        )
-
-    with h5py.File(block_path) as fi:
-        return int(load_fraction * fi["data"].shape[0])
-
-
-def merge_embedding_blocks(config: RetroPreprocessingConfig) -> None:
-    """Merge individual embedding blocks into a single binary mmap file.
-
-    The embeddings are initially stored in block-sized (e.g., ~100k embeddings per
-    block) HDF5 files. These individual block files must be merged into a single
-    file before training, to be based as a numpy mmap array to the index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    if not HAVE_TQDM:
-        raise ImportError(
-            "tqdm is required to use the merge_embedding_blocks function. Please install tqdm."
-        )
-
-    if not HAVE_H5PY:
-        raise ImportError(
-            "h5py is required to use the merge_embedding_blocks function. Please install h5py."
-        )
-
-    if torch.distributed.get_rank() != 0:
-        return
-
-    # Get block, merged paths.
-    load_fraction = config.retro_index_train_load_fraction
-    block_paths = get_training_data_block_paths(config)
-    bin_path = get_training_data_merged_path(config)
-
-    # Skip, if already built.
-    if os.path.exists(bin_path):
-        return
-
-    # Merge blocks.
-    with open(bin_path, "wb") as fo:
-        byte_offset = 0
-        for block_idx, block_path in enumerate(
-            tqdm(
-                block_paths,
-                "merge train embeddings",
-                miniters=len(block_paths) // 10,
-                disable=torch.distributed.get_rank() != 0,
-            )
-        ):
-            with h5py.File(block_path) as fi:
-                nload = get_block_nload(block_path, load_fraction)
-                block = np.array(fi["data"][:nload], copy=False)
-
-                fo.write(block.tobytes())
-
-                byte_offset += block.size * block.itemsize
-                fo.seek(byte_offset)
-
-
-def get_text_dataset_for_training(config: RetroPreprocessingConfig) -> GPTToTextDataset:
-    """Convert GPT token chunk dataset to a text dataset for passing to the
-    embedder.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        The text dataset consisting of tokens converted from sampled chunk database.
-    """
-    gpt_dataset = get_merged_sampled_dataset(
-        project_dir=config.retro_project_dir,
-        chunk_length=config.retro_gpt_chunk_length,
-        eod_token_id=config.retro_tokenizers.gpt.eod,
-    )
-    text_dataset = GPTToTextDataset(gpt_dataset, config.retro_tokenizers.gpt)
-    return text_dataset
-
-
-def embed_training_chunks(config: RetroPreprocessingConfig) -> None:
-    """Embed DB chunks.
-
-    Store chunks in blocks on disk. These blocks will later be merged into
-    a single dataset for training the index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    merged_train_data_path = get_training_data_merged_path(config)
-    if os.path.exists(merged_train_data_path):
-        return
-
-    # Get training text dataset.
-    text_dataset = get_text_dataset_for_training(config)
-
-    # Embed dataset.
-    embedder = config.retro_bert_embedders.disk
-    embedder.embed_text_dataset("index", get_training_data_block_dir(config), text_dataset)
-
-    # Merge embeddings.
-    merge_embedding_blocks(config)
-
-
-def train_on_embeddings(config: RetroPreprocessingConfig) -> None:
-    """Train index on embedded DB chunks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-    index = IndexFactory.get_index(config.retro_index_type)
-    index.train(config)
-
-
-def remove_embeddings(config: RetroPreprocessingConfig) -> None:
-    """Remove embeddings after training.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-    torch.distributed.barrier()
-    if torch.distributed.get_rank() != 0:
-        return
-    empty_index_path = get_empty_index_path(config)
-    assert os.path.isfile(empty_index_path)
-    shutil.rmtree(get_training_data_root_dir(config), ignore_errors=True)
-
-
-def _train_index(config: RetroPreprocessingConfig) -> None:
-    """Train index on DB chunks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    # Check if trained index already exists.
-    if not os.path.isfile(get_empty_index_path(config)):
-        # Embed training chunks.
-        embed_training_chunks(config)
-
-        # Train index on embeddings.
-        train_on_embeddings(config)
-
-    # Wait for (single-process) training to complete.
-    torch.distributed.barrier()
-
-    # Remove embeddings.
-    if config.retro_index_delete_training_embeddings:
-        remove_embeddings(config)
-
-
-def train_index(config: RetroPreprocessingConfig) -> None:
-    """Entry point for training the index.
-
-    We select whether to train a new index, or validate an existing index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    # Train new index.
-    if config.retro_task_validate is None:
-        _train_index(config)
-
-    # Validate existing trained index.
-    else:
-        from .validate import validate_training_embeddings
-
-        validate_training_embeddings(config)
-
-
-##################################################
-# Add to index.
-##################################################
-
-
-def get_text_dataset_for_adding(config: RetroPreprocessingConfig) -> GPTToTextDataset:
-    """Convert GPT token chunk dataset to a text dataset for passing to the
-    embedder.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        The text dataset that consists of tokens converted from the 'train' chunk database.
-        These are the chunks used for retrieval by the pretraining 'train' dataset.
-    """
-    gpt_dataset = get_merged_train_dataset(
-        project_dir=config.retro_project_dir,
-        chunk_length=config.retro_gpt_chunk_length,
-        eod_token_id=config.retro_tokenizers.gpt.eod,
-    )
-    text_dataset = GPTToTextDataset(gpt_dataset, config.retro_tokenizers.gpt)
-    return text_dataset
-
-
-def _add_to_index(config: RetroPreprocessingConfig) -> str:
-    """Add DB chunks to index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to the populated index.
-    """
-
-    # Get index.
-    index = IndexFactory.get_index(config.retro_index_type)
-
-    # Get text dataset.
-    text_dataset = get_text_dataset_for_adding(config)
-
-    # Add to index.
-    output_index_path = index.add(config, text_dataset)
-
-    return output_index_path
-
-
-def add_to_index(config: RetroPreprocessingConfig) -> None:
-    """Entry point for adding to the index.
-
-    We select whether to add to a new index, or validate an existing index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    # Add to new index.
-    if config.retro_task_validate is None:
-        _add_to_index(config)
-
-    # Validate existing encodings.
-    else:
-        from .validate import validate_added_encodings
-
-        validate_added_encodings(config)
-
-
-##################################################
-# Build index (train + add).
-##################################################
-
-
-def build_index(config: RetroPreprocessingConfig) -> None:
-    """Build index.
-
-    Building index involves sequentially running stages above:
-    - Train index (on sampled training chunks).
-    - Add to index (on all training chunks).
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    # Train index.
-    train_index(config)
-
-    # Add to index.
-    add_to_index(config)
diff --git a/megatron/core/datasets/retro/index/factory.py b/megatron/core/datasets/retro/index/factory.py
deleted file mode 100644
index f88084ddb13..00000000000
--- a/megatron/core/datasets/retro/index/factory.py
+++ /dev/null
@@ -1,40 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""The IndexFactory constructs an index from an index type string."""
-
-from megatron.core.datasets.retro.index.index import Index
-
-from .indexes import FaissBaseIndex, FaissParallelAddIndex
-
-
-class IndexFactory:
-    """Get index.
-
-    Index type generally read from argument '--retro-index-ty'.
-    """
-
-    @classmethod
-    def get_index_class(cls, index_type: str) -> type:
-        """Get an index class, given a type string.
-
-        Args:
-            index_type (str): One of 'faiss-base' (naive Faiss index wrapper) or 'faiss-par-add' (Faiss index wrapper with near embarrassingly parallel index.add().
-
-        Returns:
-            An `Index` sub-type corresponding to the `index_type`.
-        """
-        return {"faiss-base": FaissBaseIndex, "faiss-par-add": FaissParallelAddIndex}[index_type]
-
-    @classmethod
-    def get_index(cls, index_type: str) -> Index:
-        """Construct an index from an index type string.
-
-        Args:
-            index_type (str): One of 'faiss-base' (naive Faiss index wrapper) or 'faiss-par-add' (Faiss index wrapper with near embarrassingly parallel index.add().
-
-        Returns:
-            An `Index` instance corresponding to the `index_type`.
-        """
-        index_class = cls.get_index_class(index_type)
-        index = index_class()
-        return index
diff --git a/megatron/core/datasets/retro/index/index.py b/megatron/core/datasets/retro/index/index.py
deleted file mode 100644
index 129c239de34..00000000000
--- a/megatron/core/datasets/retro/index/index.py
+++ /dev/null
@@ -1,150 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Base class for all vector indexes.
-
-A vector index is a type of retrieval database that is queried using vectors,
-and returns vectors that are 'similar' (e.g., by cosine distance) to the query
-vector. The construction and usage of an index generally has the following
-pattern:
-
-  - Train the index on representative vectors.
-  - Add vectors to the index (i.e., vectors available for retrieval)
-  - Query index with new vector, to retrieve similar vector indexes.
-"""
-
-import abc
-import os
-from typing import Tuple
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.retro.config import Embedder, RetroPreprocessingConfig
-from megatron.core.datasets.retro.utils import GPTToTextDataset
-
-from .utils import get_index_dir
-
-try:
-    import faiss
-
-    HAVE_FAISS = True
-except ImportError:
-    HAVE_FAISS = False
-
-
-class Index(abc.ABC):
-    """Abstract base class for indexes.
-
-    *Note* : While currently only Faiss-based classes are implemented, in the
-    future, this class will be extended with other types of indexes that have
-    different performance-accuracy trade-offs.
-
-    The primary methods to override are:
-    - train() : Train index on the sampled training chunks.
-    - add() : Add all training chunks to index.
-    """
-
-    @classmethod
-    def make_object_verbose(cls, index: "faiss.Index", verbose: bool) -> None:
-        """Make index object verbose.
-
-        Args:
-            index (faiss.Index): Faiss object to set verbose.
-            verbose (bool): Sets whether index should log status updates during training and adding.
-        """
-        if not HAVE_FAISS:
-            raise ImportError("faiss is required to use the Index class. Please install faiss.")
-
-        assert isinstance(verbose, bool)
-        faiss.ParameterSpace().set_index_parameter(index, "verbose", verbose)
-
-    def get_empty_index_path(self, config: RetroPreprocessingConfig) -> str:
-        """Get file path to empty index (i.e., trained, but unpopulated).
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-
-        Returns:
-            File path to empty index
-                (i.e., this index has had index.train() called, but not yet index.add()).
-        """
-        return os.path.join(
-            get_index_dir(config), "empty_%.3f.faissindex" % config.retro_index_train_load_fraction
-        )
-
-    def get_empty_index(self, config: RetroPreprocessingConfig) -> "faiss.Index":
-        """Get empty index (i.e., trained, but unpopulated).
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-
-        Returns:
-            Empty Faiss index, loaded from storage.
-        """
-        if not HAVE_FAISS:
-            raise ImportError("faiss is required to use the Index class. Please install faiss.")
-        return faiss.read_index(self.get_empty_index_path(config))
-
-    def get_added_index_path(self, config: RetroPreprocessingConfig) -> str:
-        """Get file path to index that has been populated with vectors.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-
-        Returns:
-            File path to added index
-                (i.e., this index has had both index.train() and index.add() called).
-        """
-        return os.path.join(
-            get_index_dir(config),
-            "added_%.3f_%.3f.faissindex"
-            % (config.retro_index_train_load_fraction, config.retro_index_add_load_fraction),
-        )
-
-    def get_added_index(self, config: RetroPreprocessingConfig) -> "faiss.Index":
-        """Get index that has been populated with vectors.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-
-        Returns:
-            'Added' (i.e., populated) Faiss index, loaded from storage.
-        """
-        if not HAVE_FAISS:
-            raise ImportError("faiss is required to use the Index class. Please install faiss.")
-        return faiss.read_index(self.get_added_index_path(config))
-
-    @abc.abstractmethod
-    def train(self, config: RetroPreprocessingConfig) -> None:
-        """Train index on a representative set of vectors.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-        """
-
-    @abc.abstractmethod
-    def add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
-        """Add vectors to index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-            text_dataset (GPTToTextDataset): Text dataset that will be embedded
-                and added to the index.
-        """
-
-    def embed_text_dataset_block(
-        self, embedder: Embedder, text_dataset: GPTToTextDataset, _range: Tuple[int, int]
-    ) -> np.ndarray:
-        """Embed a range of a text dataset.
-
-        Args:
-            embedder (Embedder): Embedder used for embedding a text dataset.
-            text_dataset (GPTToTextDataset): Text dataset that will be embedded.
-            _range (Tuple[int, int]): Start/end sample indices within
-                text dataset used for embedding.
-
-        Returns:
-            An array of embeddings, with shape (len(text_dataset), dimension(embedder)).
-        """
-        sub_dataset = torch.utils.data.Subset(text_dataset, range(*_range))
-        return embedder.embed_text_dataset(sub_dataset)
diff --git a/megatron/core/datasets/retro/index/indexes/__init__.py b/megatron/core/datasets/retro/index/indexes/__init__.py
deleted file mode 100644
index c445909fea5..00000000000
--- a/megatron/core/datasets/retro/index/indexes/__init__.py
+++ /dev/null
@@ -1,10 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-Exports:
-- FaissBaseIndex: Unoptimized Faiss index wrapper
-- FaissParallelAddIndex: Optimized index.add() for Faiss index.
-"""
-
-from .faiss_base import FaissBaseIndex
-from .faiss_par_add import FaissParallelAddIndex
diff --git a/megatron/core/datasets/retro/index/indexes/faiss_base.py b/megatron/core/datasets/retro/index/indexes/faiss_base.py
deleted file mode 100644
index 6db0a420dff..00000000000
--- a/megatron/core/datasets/retro/index/indexes/faiss_base.py
+++ /dev/null
@@ -1,179 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-This class implements a simple, un-optimized wrapper around a Faiss index, that
-implements the Index interface (see ..index.py). While this class is
-instantiable, it is meant to be extended with optimizations in classes that
-inherit from this class (see FaissParAddIndex, for an example).
-"""
-
-import os
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.index.index import Index
-from megatron.core.datasets.retro.index.utils import (
-    get_training_data_merged_path,
-    num_samples_to_block_ranges,
-)
-from megatron.core.datasets.retro.utils import GPTToTextDataset, log_retro_rank_0
-
-try:
-    import faiss
-
-    HAVE_FAISS = True
-except ImportError:
-    HAVE_FAISS = False
-
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-
-class FaissBaseIndex(Index):
-    """Base class for Faiss-base indexes.
-
-    This class wraps a Faiss index, and adds additional functionality for training
-    and adding codes. This base class performs a naive sequential code adding,
-    while the optimized FaissParallelAddIndex class performs a parallel
-    index.add().
-    """
-
-    def _train(self, config: RetroPreprocessingConfig) -> None:
-        """Train index (rank 0's method).
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-        """
-
-        if not HAVE_FAISS:
-            raise ImportError(
-                "faiss is required to use the FaissBaseIndex class. Please install faiss."
-            )
-
-        assert torch.distributed.get_rank() == 0
-
-        # Set num threads (torch.distributed reset it to 1).
-        faiss.omp_set_num_threads(64)
-
-        empty_index_path = self.get_empty_index_path(config)
-
-        # Index already exists? -> return.
-        if os.path.isfile(empty_index_path):
-            return
-
-        # Load data.
-        merged_path = get_training_data_merged_path(config)
-        inp = np.memmap(merged_path, dtype="f4", mode="r").reshape((-1, config.hidden_size))
-
-        # Init index.
-        index = faiss.index_factory(config.hidden_size, config.retro_index_str)
-
-        # Move to GPU.
-        log_retro_rank_0("> move faiss index to gpu.")
-        index_ivf = faiss.extract_index_ivf(index)
-        clustering_index = faiss.index_cpu_to_all_gpus(faiss.IndexFlatL2(index_ivf.d))
-        index_ivf.clustering_index = clustering_index
-        log_retro_rank_0("> finished moving to gpu.")
-        self.make_object_verbose(index, True)
-        self.make_object_verbose(index_ivf, True)
-        self.make_object_verbose(index_ivf.quantizer, True)
-        self.make_object_verbose(index_ivf.clustering_index, True)
-
-        # Train index.
-        index.train(inp)
-
-        # Save index.
-        faiss.write_index(index, empty_index_path)
-
-    def train(self, config: RetroPreprocessingConfig) -> None:
-        """Train index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-        """
-
-        # Single process only.
-        if torch.distributed.get_rank() == 0:
-            self._train(config)
-
-        torch.distributed.barrier()
-
-    def _add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
-        """Add to index (rank 0's method).
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-            text_dataset (GPTToTextDataset): Text dataset that will be embedded
-                and added to the index.
-        """
-
-        if not HAVE_FAISS:
-            raise ImportError(
-                "faiss is required to use the FaissBaseIndex class. Please install faiss."
-            )
-
-        if not HAVE_TQDM:
-            raise ImportError(
-                "tqdm is required to use the FaissBaseIndex class. Please install tqdm."
-            )
-
-        assert torch.distributed.get_rank() == 0
-
-        dataset_sample_ranges = num_samples_to_block_ranges(len(text_dataset))
-
-        # Set num threads (torch.distributed reset it to 1).
-        faiss.omp_set_num_threads(64)
-
-        # Bert embedder.
-        embedder = config.bert_embedders.mem
-
-        # Empty/added index paths.
-        empty_index_path = self.get_empty_index_path()
-        added_index_path = self.get_added_index_path()
-
-        # Skip adding, if index exists.
-        if os.path.isfile(added_index_path):
-            return
-
-        # Read trained index.
-        index = faiss.read_index(empty_index_path)
-
-        # Iterate data blocks & add.
-        for sample_range in tqdm(dataset_sample_ranges, "faiss_base.add"):
-            # Embed text.
-            embeds = self.embed_text_dataset_block(embedder, text_dataset, sample_range)
-
-            # Add to index.
-            index.add(embeds)
-
-        # Write index.
-        faiss.write_index(index, added_index_path)
-
-    def add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> str:
-        """Add to index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-            text_dataset (GPTToTextDataset): Text dataset that will be embedded
-                and added to the index.
-
-        Returns:
-            File path to the populated index.
-        """
-
-        # Single process only.
-        if torch.distributed.get_rank() == 0:
-            self._add(config, text_dataset)
-
-        # Wait for rank 0.
-        torch.distributed.barrier()
-
-        # Get output index path, for return.
-        return self.get_added_index_path(config)
diff --git a/megatron/core/datasets/retro/index/indexes/faiss_par_add.py b/megatron/core/datasets/retro/index/indexes/faiss_par_add.py
deleted file mode 100644
index ccd79f31d4b..00000000000
--- a/megatron/core/datasets/retro/index/indexes/faiss_par_add.py
+++ /dev/null
@@ -1,253 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Multi-process & multi-node version of Faiss's index.add().
-
-This class inherits from FaissBaseIndex, and optimizes the 'add()' method by
-making it multi-node and multi-process, with bit-wise equivalence to
-FaissBaseIndex. This allows 'add()' to scale out to very large datasets, since
-the vast majority of the computational effort is embarrassingly parallel.
-"""
-
-import os
-import shutil
-from typing import Tuple
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.retro.config import Embedder, RetroPreprocessingConfig
-from megatron.core.datasets.retro.index.utils import get_added_code_paths, get_added_codes_dir
-from megatron.core.datasets.retro.utils import (
-    GPTToTextDataset,
-    get_blocks_by_rank,
-    log_retro_rank_0,
-    retro_makedir,
-)
-
-from .faiss_base import FaissBaseIndex
-
-try:
-    import psutil
-
-    HAVE_PSUTIL = True
-except ImportError:
-    HAVE_PSUTIL = False
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-try:
-    import faiss
-
-    HAVE_FAISS = True
-except ImportError:
-    HAVE_FAISS = False
-
-
-class FaissParallelAddIndex(FaissBaseIndex):
-    """
-    This class parallelizes both 1) encoding vectors, and 2) adding codes to the
-    index. This class is more performant than naive use of Faiss, because most
-    of the computational work is in encoding the vectors, which is an
-    embarassingly parallel operation.
-    """
-
-    def encode_block(
-        self, index: "faiss.Index", embedder: Embedder, text_dataset: GPTToTextDataset, block: dict
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        """Encode sub-dataset block, to be later added to index.
-
-        Encode the data subset, generally in blocks of 1M vectors each. For
-        each block, the empty/trained index is loaded, codes are computed
-        via index.sa_encode(), and the resulting codes are saved to disk.
-
-        Args:
-            index (faiss.Index): Faiss index object.
-            embedder (Embedder): Embedder used to embed text dataset.
-            text_dataset (GPTToTextDataset): Text dataset to be embedded and encoded.
-            block (dict): Range information specifying start/end indices within text dataset.
-
-        Returns:
-            A tuple of (embeddings, encodings) for the given block subset of the text dataset.
-        """
-
-        # Embed block.
-        embeddings = self.embed_text_dataset_block(embedder, text_dataset, block["range"])
-
-        # Encode block.
-        log_retro_rank_0("encode.")
-        codes = index.sa_encode(embeddings)
-
-        # Return embeddings for validation purposes.
-        return embeddings, codes
-
-    def save_block(self, config: RetroPreprocessingConfig, block: dict, codes: np.ndarray) -> None:
-        """Save block of codes to disk.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-            block (dict): Range information specifying the start/end indices within
-                the encoded text dataset. Here, the 'path' item is used for writing
-                the encodings to storage.
-            codes (np.ndarray): Block of encodings to be saved to storage.
-        """
-        # Save neighbors.
-        log_retro_rank_0("save codes.")
-        retro_makedir(config, os.path.dirname(block["path"]))
-        with h5py.File(block["path"], "w") as f:
-            f.create_dataset("data", data=codes)
-
-    def encode(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
-        """Encode text dataset, to be later added to index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-            text_dataset (GPTToTextDataset): Text dataset to be encoded by the index.
-        """
-
-        codes_dir = get_added_codes_dir(config)
-        retro_makedir(config, codes_dir)
-
-        # Index.
-        index = self.get_empty_index(config)
-
-        # Bert embedder.
-        embedder = config.retro_bert_embedders.mem
-
-        # Missing code blocks.
-        def validate(f: h5py.File) -> None:
-            """Validation method for validating loaded encodings.
-
-            Args:
-                f (h5py.File): File that contains encodings.
-            """
-            assert len(f["data"].shape) == 2
-
-        blocks = get_blocks_by_rank(
-            codes_dir, len(text_dataset), config.retro_block_size, validate=validate
-        )
-
-        # Encode each block.
-        for block_index, block in enumerate(blocks.missing):
-            if block is not None:
-                # Progress.
-                log_retro_rank_0(
-                    "encode block %d / %d ... %s."
-                    % (block_index, len(blocks.missing), block["path"])
-                )
-
-                # Encode and save.
-                _, codes = self.encode_block(index, embedder, text_dataset, block)
-                self.save_block(config, block, codes)
-
-            # Synchronize progress across all ranks. (for easier observation)
-            log_retro_rank_0(" > waiting for other ranks to finish block.")
-            torch.distributed.barrier()
-
-    def add_codes(self, config: RetroPreprocessingConfig) -> None:
-        """Read codes from disk, and add them to the index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-        """
-
-        if not HAVE_PSUTIL:
-            raise ImportError(
-                "psutil is required to use the FaissParallelAddIndex class. Please install psutil."
-            )
-
-        if not HAVE_TQDM:
-            raise ImportError(
-                "tqdm is required to use the FaissParallelAddIndex class. Please install tqdm."
-            )
-
-        if not HAVE_FAISS:
-            raise ImportError(
-                "faiss is required to use the FaissParallelAddIndex class. Please install faiss."
-            )
-
-        if not HAVE_H5PY:
-            raise ImportError(
-                "h5py is required to use the FaissParallelAddIndex class. Please install h5py."
-            )
-
-        if torch.distributed.get_rank() != 0:
-            return
-
-        added_index_path = self.get_added_index_path(config)
-        if os.path.exists(added_index_path):
-            return
-
-        # Index.
-        log_retro_rank_0("read empty index.")
-        index = self.get_empty_index(config)
-        index_ivf = faiss.extract_index_ivf(index)
-
-        # Add codes.
-        log_retro_rank_0("add codes.")
-        code_paths = get_added_code_paths(config)
-        pbar = tqdm(code_paths)
-        for code_path in pbar:
-            pbar.set_description(
-                "add codes, mem %.3f gb, %.1f%%"
-                % (psutil.virtual_memory()[3] / 1024**3, psutil.virtual_memory()[2])
-            )
-            with h5py.File(code_path) as f:
-                nload = int(config.retro_index_add_load_fraction * f["data"].shape[0])
-                offset = int(os.path.basename(code_path).split("-")[0])
-                xids = np.arange(offset, offset + nload)
-                codes = np.copy(f["data"][:nload])
-                index_ivf.add_sa_codes(codes, xids)
-
-        # Update index's ntotal.
-        index.ntotal = index_ivf.ntotal
-
-        # Write index.
-        log_retro_rank_0("write added index.")
-        faiss.write_index(index, added_index_path)
-
-    def remove_codes(self, config: RetroPreprocessingConfig) -> None:
-        """Remove added codes after adding to index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-        """
-        if torch.distributed.get_rank() != 0:
-            return
-        assert os.path.isfile(self.get_added_index_path(config))
-
-        if config.retro_index_delete_added_codes:
-            raise Exception("remove?")
-            shutil.rmtree(get_added_codes_dir(config), ignore_errors=True)
-
-    def add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
-        """Add vectors to index.
-
-        Args:
-            config (RetroPreprocessingConfig): Retro preprocessing config.
-            text_dataset (GPTToTextDataset): Text dataset that will be embedded
-                and added to the index.
-        """
-
-        # Encode chunks.
-        self.encode(config, text_dataset)
-
-        # Add codes to index.
-        self.add_codes(config)
-
-        # Wait for (single-process) adding to complete.
-        torch.distributed.barrier()
-
-        # Remove codes.
-        self.remove_codes(config)
diff --git a/megatron/core/datasets/retro/index/utils.py b/megatron/core/datasets/retro/index/utils.py
deleted file mode 100644
index 58229439ae6..00000000000
--- a/megatron/core/datasets/retro/index/utils.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Utilities for building an index."""
-
-import glob
-import os
-from typing import List, Tuple
-
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.utils import retro_makedir
-
-
-def get_index_dir(config: RetroPreprocessingConfig) -> str:
-    """Create sub-directory for this index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to index sub-directory within Retro project.
-    """
-
-    # Directory path.
-    index_dir_path = os.path.join(
-        config.retro_project_dir, "index", config.retro_index_type, config.retro_index_str
-    )
-
-    # Make directory.
-    retro_makedir(config, index_dir_path)
-
-    return index_dir_path
-
-
-def num_samples_to_block_ranges(
-    config: RetroPreprocessingConfig, num_samples: int
-) -> List[Tuple[int, int]]:
-    """Split a range (length num_samples) into sequence of block ranges
-    of size block_size.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        num_samples (int): Split `num_samples` into consecutive block ranges, where each block is size `config.retro_block_size`.
-
-    Returns:
-        A list of tuples where each item is the (start, end) index for a given block.
-    """
-    block_size = config.retro_block_size
-    start_idxs = list(range(0, num_samples, block_size))
-    end_idxs = [min(num_samples, s + block_size) for s in start_idxs]
-    ranges = list(zip(start_idxs, end_idxs))
-    return ranges
-
-
-def get_training_data_root_dir(config: RetroPreprocessingConfig) -> str:
-    """Get root directory for embeddings (blocks and merged data).
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to the training data directory, which contains both training embedding blocks and the final merged training embeddings.
-    """
-    return os.path.join(config.retro_project_dir, "index", "train_emb")
-
-
-def get_training_data_block_dir(config: RetroPreprocessingConfig) -> str:
-    """Get directory for of saved embedding blocks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to the directory containing the training embedding blocks, which will be later merged into a single embedding array.
-    """
-    return os.path.join(get_training_data_root_dir(config), "blocks")
-
-
-def get_training_data_block_paths(config: RetroPreprocessingConfig) -> List[str]:
-    """Get paths to saved embedding blocks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Paths of all training embedding blocks.
-    """
-    return sorted(glob.glob(get_training_data_block_dir(config) + "/*.hdf5"))
-
-
-def get_training_data_merged_path(config: RetroPreprocessingConfig) -> str:
-    """Get path to merged training embeddings.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to the merged training embedding binary file.
-    """
-    return os.path.join(
-        get_training_data_root_dir(config),
-        "train_%.3f.bin" % config.retro_index_train_load_fraction,
-    )
-
-
-def get_added_codes_dir(config: RetroPreprocessingConfig) -> str:
-    """Get directory of saved encodings.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Path to the directory containing the vector encodings for adding to the index.
-    """
-    return os.path.join(get_index_dir(config), "add_codes")
-
-
-def get_added_code_paths(config: RetroPreprocessingConfig) -> List[str]:
-    """Get paths to all saved encodings.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        Paths of all vector encoding blocks, for adding to the index.
-    """
-    return sorted(glob.glob(get_added_codes_dir(config) + "/*.hdf5"))
diff --git a/megatron/core/datasets/retro/index/validate.py b/megatron/core/datasets/retro/index/validate.py
deleted file mode 100644
index 5f75147a8af..00000000000
--- a/megatron/core/datasets/retro/index/validate.py
+++ /dev/null
@@ -1,194 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Validate an index's data.
-
-This module contains functionality for checking for bitwise equality across code
-changes. The training and adding steps of index construction can be validated
-separately. The following high-level checks are supported:
-
-  - Training: Validate that saved training embeddings are bitwise equal with a
-      sample set of freshly computed embeddings. (*Note*:
-      `--no-retro-index-delete-training-embeddings` must be used.)
-  - Adding: Validate that the saved encodings are bitwise equal with a sample of
-      sample set of freshly computed encodings. (*Note*:
-      `--no-retro-index-delete-added-codes` must be used.)
-"""
-
-import numpy as np
-import torch
-from torch.utils.data import Subset
-
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.utils import get_blocks_by_rank, log_retro_rank_0
-
-from .build import get_text_dataset_for_adding, get_text_dataset_for_training
-from .factory import IndexFactory
-from .utils import get_added_codes_dir, get_training_data_block_dir
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-
-##################################################
-# Validate trained index.
-##################################################
-
-
-def validate_training_embeddings(config: RetroPreprocessingConfig) -> None:
-    """Validate training embeddings.
-
-    Steps:
-    - Randomly sample subset of text dataset blocks.
-    - Embed each block.
-    - Compare against saved embeddings.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    if not HAVE_H5PY:
-        raise ImportError(
-            "h5py is required to use the validate_training_embeddings function. "
-            "Please install h5py."
-        )
-
-    # Training text dataset.
-    text_dataset = get_text_dataset_for_training(config)
-
-    # Sample existing blocks.
-    blocks = get_blocks_by_rank(
-        dirname=get_training_data_block_dir(config),
-        n_samples=len(text_dataset),
-        block_size=config.retro_block_size,
-        validate=None,
-        sample=config.retro_task_validate,
-    )
-
-    assert blocks.n_missing_world == 0
-
-    # Embed & validate blocks.
-    embedder = config.retro_bert_embedders.mem
-    for block_idx, block in enumerate(blocks.existing):
-        # Missing block lists are extended with None to have equal-length
-        # lists. Skip the Nones.
-        if block is not None:
-            # Progress. (*note*: move world progress to here.)
-            log_retro_rank_0(
-                "embed training block %d / %d ... %s."
-                % (block_idx, len(blocks.existing), block["path"])
-            )
-
-            # Load existing block embeddings.
-            with h5py.File(block["path"]) as f:
-                existing_embeddings = np.copy(f["data"])
-
-            # Embed block.
-            sub_dataset = Subset(text_dataset, range(*block["range"]))
-            embeddings = embedder.embed_text_dataset(sub_dataset, "train")
-
-            # Check equality.
-            log_retro_rank_0(" > validate.")
-            assert np.array_equal(existing_embeddings, embeddings)
-
-        # Synchronize progress across all ranks. (for easier observation)
-        log_retro_rank_0(" > waiting for other ranks to finish block.")
-        torch.distributed.barrier()
-
-    log_retro_rank_0(" > finished validating training embeddings.")
-
-
-##################################################
-# Validate filled index.
-##################################################
-
-
-def validate_added_encodings(config: RetroPreprocessingConfig) -> None:
-    """Validate added encodings.
-
-    Steps:
-    - Randomly sample subset of text dataset blocks.
-    - Encode each block.
-    - Compare against saved encodings.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    # Index.
-    index = IndexFactory.get_index(config.retro_index_type)
-    inner_index = index.get_empty_index(config)
-
-    # Text dataset.
-    text_dataset = get_text_dataset_for_adding(config)
-
-    # Sample existing blocks.
-    def validate(f: h5py.File) -> None:
-        """Validation method for validating encoding blocks.
-
-        Args:
-            f (h5py.File): File with block of encodings.
-        """
-        assert len(f["data"].shape) == 2
-
-    blocks = get_blocks_by_rank(
-        dirname=get_added_codes_dir(config),
-        n_samples=len(text_dataset),
-        block_size=config.retro_block_size,
-        validate=validate,
-        sample=config.retro_task_validate,
-    )
-
-    assert blocks.n_missing_world == 0
-
-    # Encode and validate blocks.
-    embedder = config.retro_bert_embedders.mem
-    for block_idx, block in enumerate(blocks.existing):
-        if block is not None:
-            # Progress.
-            log_retro_rank_0(
-                "encode block %d / %d ... %s." % (block_idx, len(blocks.existing), block["path"])
-            )
-
-            # Load existing codes.
-            with h5py.File(block["path"]) as f:
-                existing_codes = np.copy(f["data"])
-
-            # Encode block.
-            embeddings, codes = index.encode_block(inner_index, embedder, text_dataset, block)
-
-            # Check equality.
-            log_retro_rank_0(" > validate.")
-            assert np.array_equal(existing_codes, codes)
-
-        # Synchronize progress across all ranks. (for easier observation)
-        log_retro_rank_0(" > waiting for other ranks to finish block.")
-        torch.distributed.barrier()
-
-    log_retro_rank_0(" > finished validating added encodings.")
-
-
-##################################################
-# Validate index (trained + filled).
-##################################################
-
-
-def validate_index(config: RetroPreprocessingConfig) -> None:
-    """Validate index.
-
-    Validating index involves sequentially running stages above:
-    - Validate trained index.
-    - Validate filled index.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    # Validate training embeddings.
-    validate_training_embeddings(config)
-
-    # Validate added codes.
-    validate_added_encodings(config)
diff --git a/megatron/core/datasets/retro/query/__init__.py b/megatron/core/datasets/retro/query/__init__.py
deleted file mode 100644
index ac9483373c9..00000000000
--- a/megatron/core/datasets/retro/query/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
diff --git a/megatron/core/datasets/retro/query/gpt_chunk_dataset.py b/megatron/core/datasets/retro/query/gpt_chunk_dataset.py
deleted file mode 100644
index 6191a30a31f..00000000000
--- a/megatron/core/datasets/retro/query/gpt_chunk_dataset.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-A GPTChunkDataset is a wrapper around a regular GPTDataset, that sequentially
-chunks the sample tokens into `retro_chunk_length` sized smaller samples.
-
-For example, if the GPTDataset has 100 samples and a sequence length of 2048, and
-retro_chunk_length is 64, then the GPTChunkDataset will contain 100*(2048/64) =
-3200 samples, each with length 64.
-"""
-
-import torch
-
-from megatron.core.datasets.gpt_dataset import GPTDataset
-from megatron.core.datasets.retro.utils import get_num_chunks_per_sample
-
-from .utils import get_neighbor_dir
-
-
-class GPTChunkDataset(torch.utils.data.Dataset):
-    """Pretraining chunk dataset wraps a standard GPT dataset.
-
-    This dataset conceptually divides each sample (e.g., length 2048)
-    into chunks (e.g., length 64) and restructures them into a list of
-    chunks (e.g., length num_samples * num_chunks_per_sample).
-
-    Args:
-        sample_dataset (GPTDataset): Original GPT dataset, with `sequence_length` size samples.
-        sample_length (int): Alias for `sequence_length`.
-        chunk_length (int): Retro chunk length (e.g., 64).
-    """
-
-    def __init__(self, sample_dataset: GPTDataset, sample_length: int, chunk_length: int):
-
-        super().__init__()
-
-        self.sample_dataset = sample_dataset
-        self.chunk_length = chunk_length
-        self.n_chunks_per_sample = get_num_chunks_per_sample(sample_length, chunk_length)
-        self.n_samples = len(sample_dataset)
-        self.n_chunks = self.n_samples * self.n_chunks_per_sample
-
-    def __len__(self) -> int:
-        """Get dataset length.
-
-        Returns:
-            Dataset length.
-        """
-        return self.n_chunks
-
-    def __getitem__(self, idx: int) -> dict:
-        """Get sample, including represented document IDs.
-
-        Args:
-            idx (int): Sample index.
-
-        Returns:
-            A sample, which contains both the chunk-length token sample ('text') along with all document_ids ('doc_ids') contained withing the full `sequence_length` sample.
-        """
-
-        # Convert global chunk index to global sample index & local chunk index.
-        sample_idx = idx // self.n_chunks_per_sample
-        chunk_idx = idx % self.n_chunks_per_sample
-
-        # Extract sample data.
-        sample = self.sample_dataset[sample_idx]
-        sample_token_ids = sample["text"]
-        sample_doc_ids = sample["document_ids"]
-
-        # Chunk start/end token idxs.
-        token_start_idx = chunk_idx * self.chunk_length
-        token_end_idx = token_start_idx + self.chunk_length
-        chunk_token_ids = sample_token_ids[token_start_idx:token_end_idx]
-
-        # Sample.
-        return {"doc_ids": sample_doc_ids, "text": chunk_token_ids}
-
-
-def build_gpt_chunk_datasets_from_gpt_datasets(
-    project_dir: str, gpt_datasets: dict, sample_length: int, chunk_length: int
-) -> dict:
-    """Get train, valid, test GPT chunk datasets.
-
-    Args:
-        project_dir (str): Retro project dir.
-        gpt_datasets (dict): Mapping of 'train', 'valid', and 'test' GPT datasets (original, unchunked datasets).
-        sample_length (int): Alias of `sequence_length`.
-        chunk_length (int): Retro chunk length (e.g., 64).
-
-    Returns:
-        A <dict> ?
-    """
-
-    # GPT chunk datasets.
-    chunk_datasets = {
-        key: (
-            {
-                "dataset": GPTChunkDataset(sample_ds, sample_length, chunk_length),
-                "neighbor_dir": get_neighbor_dir(project_dir, key, sample_ds),
-                "num_active_chunks": num_active_samples
-                * get_num_chunks_per_sample(sample_length, chunk_length),
-            }
-            if sample_ds
-            else None
-        )
-        for key, (sample_ds, num_active_samples) in gpt_datasets.items()
-    }
-
-    return chunk_datasets
diff --git a/megatron/core/datasets/retro/query/multi_split_gpt_dataset.py b/megatron/core/datasets/retro/query/multi_split_gpt_dataset.py
deleted file mode 100644
index 52b0b6bac4a..00000000000
--- a/megatron/core/datasets/retro/query/multi_split_gpt_dataset.py
+++ /dev/null
@@ -1,115 +0,0 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-
-"""A MultiSplitGPTDataset can handle multiple intersecting split strings, as well
-as returning all of the document IDs of a sample."""
-
-import logging
-from dataclasses import dataclass
-from typing import Dict, List
-
-import numpy
-
-from megatron.core.datasets.blended_megatron_dataset_config import (
-    convert_split_vector_to_split_matrix,
-    parse_and_normalize_split,
-)
-from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig
-from megatron.core.datasets.indexed_dataset import IndexedDataset
-from megatron.core.datasets.utils import Split
-from megatron.core.utils import log_single_rank
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class MultiSplitGPTDatasetConfig(GPTDatasetConfig):
-    """Configuration object for Megatron Core blended and Retro datasets.
-
-    Args:
-        return_document_ids (bool): Whether to return the document ids when querying the dataset.
-            Turn this option on during preprocessing.
-        split_preprocessing (str): The Retro preprocessing split string.
-            It follows the same pattern convention as 'split'.
-            Not to be used with 'blend_per_split'.
-    """
-
-    return_document_ids: bool = None
-
-    split_preprocessing: str = None
-
-    def __post_init__(self) -> None:
-        """Validate config attributes."""
-
-        super().__post_init__()
-        assert self.split is not None, "the Retro data pipeline does not support 'blend_per_split'"
-        assert self.return_document_ids is not None, "this attribute must be user defined"
-        assert self.split_preprocessing is not None, "this attribute must be user defined"
-        split_vector = parse_and_normalize_split(self.split)
-        split_preprocessing_vector = parse_and_normalize_split(self.split_preprocessing)
-        if not numpy.allclose(split_vector, split_preprocessing_vector):
-            self.split_matrix = convert_split_vector_to_split_matrix(
-                split_vector, split_preprocessing_vector
-            )
-            log_single_rank(
-                logger,
-                logging.WARNING,
-                f"split =/= split_preprocessing. Let split_matrix = {self.split_matrix}",
-            )
-
-
-class MultiSplitGPTDataset(GPTDataset):
-    """Retro's customized GPT dataset.
-
-    Args:
-        indexed_dataset (IndexedDataset): The IndexedDataset around which
-            to build the MegatronDataset.
-        dataset_path (str): The real path on disk to the dataset, for bookkeeping.
-        indexed_indices (numpy.ndarray): The set of the documents indices to expose.
-        num_samples (int): The number of samples to draw from the indexed dataset.
-        index_split (Split): The indexed_indices Split.
-        config (MultiSplitGPTDatasetConfig): The Retro-specific container for all
-            config sourced parameters.
-    """
-
-    def __init__(
-        self,
-        indexed_dataset: IndexedDataset,
-        dataset_path: str,
-        indexed_indices: numpy.ndarray,
-        num_samples: int,
-        index_split: Split,
-        config: MultiSplitGPTDatasetConfig,
-    ) -> None:
-        super().__init__(
-            indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config
-        )
-
-    def __getitem__(self, idx: int) -> Dict[str, numpy.ndarray]:
-        """Get dataset sample.
-
-        Args:
-            idx (int): The index into the dataset.
-
-        Returns:
-            Dict[str, numpy.ndarray]: The text ids and (optionally)
-                the document ids wrapped in a dictionary.
-        """
-        text, document_ids = self._query_document_sample_shuffle_indices(idx)
-        if self.config.return_document_ids:
-            return {"text": text, "document_ids": document_ids}
-        else:
-            return {"text": text}
-
-    @staticmethod
-    def _key_config_attributes() -> List[str]:
-        """Add custom attributes for building unique dataset hash.
-
-        The preprocessing split used for preprocessing will constrain
-            the samples available for pretraining.
-
-        Returns:
-            List[str]: The key config attributes.
-        """
-        return super(MultiSplitGPTDataset, MultiSplitGPTDataset)._key_config_attributes() + [
-            "split_preprocessing"
-        ]
diff --git a/megatron/core/datasets/retro/query/query.py b/megatron/core/datasets/retro/query/query.py
deleted file mode 100644
index 42d93d5aafb..00000000000
--- a/megatron/core/datasets/retro/query/query.py
+++ /dev/null
@@ -1,449 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Entry point for querying an index using a GPTChunkDataset.
-
-Querying involves:
-
-  - Iterate all chunks in the GPTChunkDataset.
-  - Query index for neighbor chunk IDs (i.e., chunks from the chunk database).
-  - Save neighbor chunk IDs to disk, for use in building a RetroDataset sample
-      during pretraining.
-"""
-
-import os
-import time
-import typing
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.db.dataset import DBDataset
-from megatron.core.datasets.retro.db.utils import (
-    get_merged_train_dataset as get_db_merged_train_dataset,
-)
-from megatron.core.datasets.retro.index.factory import IndexFactory
-from megatron.core.datasets.retro.index.index import Index
-from megatron.core.datasets.retro.index.utils import get_index_dir
-from megatron.core.datasets.retro.query.gpt_chunk_dataset import GPTChunkDataset
-from megatron.core.datasets.retro.utils import (
-    GPTToTextDataset,
-    get_blocks_by_rank,
-    log_retro_rank_0,
-    retro_makedir,
-)
-
-try:
-    import psutil
-
-    HAVE_PSUTIL = True
-except ImportError:
-    HAVE_PSUTIL = False
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-try:
-    import faiss
-
-    HAVE_FAISS = True
-except ImportError:
-    HAVE_FAISS = False
-
-
-def get_index(config: RetroPreprocessingConfig, ondisk: bool = False) -> "faiss.Index":
-    """Read index from disk.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        ondisk (bool): If `ondisk = True`, memory map the index.
-            (For debugging purposes only; very non-performant.)
-
-    Returns:
-        A Faiss index, loaded from storage.
-    """
-    if not HAVE_FAISS:
-        raise ImportError(
-            "faiss is required to use the query_neighbors function. " "Please install faiss."
-        )
-
-    # Load index.
-    index_wrapper = IndexFactory.get_index(config.retro_index_type)
-    index_dir = get_index_dir(config)
-    added_index_path = index_wrapper.get_added_index_path(config)
-    if ondisk:
-        index = faiss.read_index(added_index_path, faiss.IO_FLAG_MMAP)
-    else:
-        index = faiss.read_index(added_index_path)
-
-    # Search parameters.
-    faiss.ParameterSpace().set_index_parameter(index, "efSearch", config.retro_query_ef_search)
-    faiss.ParameterSpace().set_index_parameter(index, "nprobe", config.retro_query_nprobe)
-
-    return index
-
-
-def embed_block(
-    config: RetroPreprocessingConfig, gpt_dataset: GPTChunkDataset, block: dict
-) -> np.ndarray:
-    """Embed block of chunks.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        gpt_dataset (GPTChunkDataset): Chunk dataset to be embedded.
-        block (dict): Range information containing start/end indices of subset of chunk dataset.
-
-    Returns:
-        Embeddings array, with shape (len(block["range"]), dimension(embedder)).
-    """
-    text_block_dataset = torch.utils.data.Subset(
-        GPTToTextDataset(gpt_dataset, config.retro_tokenizers.gpt), range(*block["range"])
-    )
-    return config.retro_bert_embedders.mem.embed_text_dataset(text_block_dataset)
-
-
-def query_embeddings(
-    config: RetroPreprocessingConfig,
-    db_dataset: DBDataset,
-    index: Index,
-    embeddings: np.ndarray,
-    chunk_id_range: range,
-    sample_map: dict,
-    n_chunks_per_sample: int,
-    verbose: bool = True,
-) -> typing.Tuple[np.ndarray, np.ndarray]:
-    """Query neighbors of a block of embeddings.
-
-    Querying includes:
-      - Query index for neighbor chunk IDs.
-      - Filter chunk IDs that have the same document ID as the queried embedding.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        db_dataset (DBDataset): Dataset containing chunk database entries.
-        index (Index): Vector index populated with chunk database indices.
-        embeddings (np.ndarray): Embeddings from GPT chunk dataset.
-        chunk_id_range (range): Chunk ID range from GPT chunk dataset.
-        sample_map (dict): Mapping of sample_idx to dataset_idx and document_ids.
-            Used for document filtering.
-        n_chunks_per_sample (int): Number of chunks per sample
-            (e.g., sequence_length / chunk_length).
-        verbose (bool): Log querying progress.
-
-    Returns:
-        A tuple of original (unfiltered) neighbor IDs, and filtered (by document ID) neighbor IDs.
-    """
-
-    # Query neighbor ids.
-    if verbose:
-        log_retro_rank_0("search.")
-    t = time.time()
-    assert index.ntotal > 0, "check we don't accidentally have an empty index."
-    _, query_neighbor_ids = index.search(embeddings, config.retro_query_num_neighbors_query)
-    if verbose:
-        log_retro_rank_0("  time : %.3f sec." % (time.time() - t))
-
-    # Filter banned neighbor ids.
-    if verbose:
-        log_retro_rank_0("filter banned neighbor ids.")
-    filtered_neighbor_ids = np.full(
-        shape=(len(query_neighbor_ids), config.retro_query_num_neighbors_save),
-        fill_value=-1,
-        dtype="int64",
-    )
-    min_chunk_id, max_chunk_id = chunk_id_range
-    for chunk_id in range(min_chunk_id, max_chunk_id):
-        sample_id = chunk_id // n_chunks_per_sample
-        sample = sample_map[sample_id]
-        sample_dataset_idx = sample["dataset_idx"].item()
-        sample_doc_ids = sample["doc_ids"].tolist()
-        sample_doc_tuples = [(sample_dataset_idx, d) for d in sample_doc_ids]
-
-        # Get valid neighbors (!= -1).
-        query_row = [i for i in query_neighbor_ids[chunk_id - min_chunk_id] if i >= 0]
-
-        # Filter row.
-        filtered_row = [
-            i
-            for i in query_row
-            if tuple(db_dataset.doc_tuples[i].tolist()) not in sample_doc_tuples
-        ]
-        filtered_row = filtered_row[: config.retro_query_num_neighbors_save]
-        filtered_row += [-1] * (config.retro_query_num_neighbors_save - len(filtered_row))
-        filtered_neighbor_ids[chunk_id - min_chunk_id] = filtered_row
-
-    return query_neighbor_ids, filtered_neighbor_ids
-
-
-def query_embedding_block(
-    config: RetroPreprocessingConfig,
-    db_dataset: DBDataset,
-    index: Index,
-    embeddings: np.ndarray,
-    chunk_id_range: range,
-    sample_map: dict,
-    n_chunks_per_sample: int,
-) -> typing.Tuple[np.ndarray, np.ndarray]:
-    """Query a block of embeddings.
-
-    The block is broken into smaller sub-blocks, for easier tracking of progress.
-    Both the raw neighbor IDs and the filtered neighbor IDs (i.e., chunks with the
-    same document ID are removed) are collected.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        db_dataset (DBDataset): Dataset containing chunk database entries.
-        index (Index): Vector index populated with chunk database indices.
-        embeddings (np.ndarray): Embeddings from GPT chunk dataset.
-        chunk_id_range (range): Chunk ID range from GPT chunk dataset.
-        sample_map (dict): Mapping of sample_idx to dataset_idx and document_ids.
-            Used for document filtering.
-        n_chunks_per_sample (int): Number of chunks per sample
-            (e.g., sequence_length / chunk_length).
-
-    Returns:
-        A tuple of original (unfiltered) neighbor IDs, and filtered (by document ID) neighbor IDs.
-    """
-
-    if not HAVE_TQDM:
-        raise ImportError(
-            "tqdm is required to use the query_embeddings function. Please install tqdm."
-        )
-
-    query_neighbor_ids = []
-    filtered_neighbor_ids = []
-
-    # Query in sub-blocks.
-    partial_block_size = 1000
-    for partial_start_idx in tqdm(
-        range(0, len(embeddings), partial_block_size),
-        "  search",
-        miniters=(len(embeddings) // partial_block_size) // 10,
-        disable=torch.distributed.get_rank() != 0,
-    ):
-        partial_end_idx = min(len(embeddings), partial_start_idx + partial_block_size)
-        partial_embeddings = embeddings[partial_start_idx:partial_end_idx]
-        partial_chunk_id_range = (
-            chunk_id_range[0] + partial_start_idx,
-            chunk_id_range[0] + partial_end_idx,
-        )
-        partial_query_neighbor_ids, partial_filtered_neighbor_ids = query_embeddings(
-            config,
-            db_dataset,
-            index,
-            partial_embeddings,
-            partial_chunk_id_range,
-            sample_map,
-            n_chunks_per_sample,
-            verbose=False,
-        )
-        query_neighbor_ids.append(partial_query_neighbor_ids)
-        filtered_neighbor_ids.append(partial_filtered_neighbor_ids)
-
-    # Concatenate.
-    query_neighbor_ids = np.concatenate(query_neighbor_ids, axis=0)
-    filtered_neighbor_ids = np.concatenate(filtered_neighbor_ids, axis=0)
-
-    return query_neighbor_ids, filtered_neighbor_ids
-
-
-def query_block_neighbors(
-    config: RetroPreprocessingConfig,
-    db_dataset: DBDataset,
-    query_dataset: GPTChunkDataset,
-    index: Index,
-    block: dict,
-) -> None:
-    """Query neighbors of a dataset block (i.e., range).
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        db_dataset (DBDataset): Dataset containing chunk database entries.
-        query_dataset (GPTChunkDataset): GPT chunk dataset to be queried.
-        index (Index): Vector index populated with chunk database indices.
-        block (dict): Range information containing start/end indices
-            for querying GPT chunk dataset.
-    """
-
-    if not HAVE_H5PY:
-        raise ImportError(
-            "h5py is required to use the query_block_neighbors function. Please install h5py."
-        )
-
-    n_chunks_per_sample = query_dataset.n_chunks_per_sample
-
-    # Sample map.
-    sample_ids = sorted(
-        list(set(chunk_id // n_chunks_per_sample for chunk_id in range(*block["range"])))
-    )
-    sample_map = {}
-    for i in sample_ids:
-        sample = query_dataset.sample_dataset[i]
-        sample_map[i] = {"dataset_idx": sample["dataset_id"], "doc_ids": sample["document_ids"]}
-
-    # Embed block.
-    embeddings = embed_block(config, query_dataset, block)
-
-    # Query embeddings.
-    _, filtered_neighbor_ids = query_embedding_block(
-        config, db_dataset, index, embeddings, block["range"], sample_map, n_chunks_per_sample
-    )
-
-    if config.retro_task_validate is None:
-        # Save neighbors.
-        log_retro_rank_0("save neighbors.")
-        retro_makedir(config, os.path.dirname(block["path"]))
-        f = h5py.File(block["path"], "w")
-        f.create_dataset("neighbors", data=filtered_neighbor_ids)
-        f.close()
-
-    else:
-        # Validate neighbors.
-        with h5py.File(block["path"]) as f:
-            existing_neighbor_ids = np.copy(f["neighbors"])
-            assert np.array_equal(existing_neighbor_ids, filtered_neighbor_ids)
-
-
-def query_dataset_neighbors(
-    config: RetroPreprocessingConfig,
-    db_dataset: DBDataset,
-    query_dataset: GPTChunkDataset,
-    num_active_chunks: int,
-    prefix: str,
-    neighbor_dir: str,
-    index: Index,
-) -> None:
-    """Query neighbors of each chunk within a dataset.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        db_dataset (DBDataset): Dataset containing chunk database entries.
-        query_dataset (GPTChunkDataset): GPT chunk dataset to be queried.
-        num_active_chunks (int): The 'active' chunks are the subset of the GPT chunk dataset
-            that aren't being queried. This argument is used when validating the correctness
-            of a subset of the GPT chunk dataset.
-        prefix (str): Extra string for logging progress.
-        neighbor_dir (str): File path to directory for saving neighbor IDs.
-        index (Index): Vector index populated with chunk database indices.
-    """
-    if not HAVE_H5PY:
-        raise ImportError(
-            "h5py is required to use the query_dataset_neighbors function. Please install h5py."
-        )
-
-    def validate(f: h5py.File) -> None:
-        """Validation method for validating saved neighbor IDs.
-
-        Args:
-            f (h5py.File): File containing save neighbor IDs.
-        """
-        assert (
-            f["neighbors"].shape[1] == config.retro_query_num_neighbors_save
-        ), "neighbors.shape == %s; num_neighbors_target == %d." % (
-            str(f["neighbors"].shape),
-            config.retro_num_neighbors_target,
-        )
-
-    if config.retro_task_validate is None:
-        retro_makedir(config, neighbor_dir)
-        blocks = get_blocks_by_rank(
-            neighbor_dir, num_active_chunks, config.retro_block_size, validate=validate
-        )
-        active_blocks = blocks.missing
-    else:
-        blocks = get_blocks_by_rank(
-            neighbor_dir,
-            num_active_chunks,
-            config.retro_block_size,
-            validate=validate,
-            sample=config.retro_task_validate,
-        )
-        assert blocks.n_missing_world == 0
-        active_blocks = blocks.existing
-
-    if not HAVE_PSUTIL:
-        raise ImportError(
-            "psutil is required to use the query_dataset_neighbors function. Please install psutil."
-        )
-
-    # Query each block.
-    for block_index, block in enumerate(active_blocks):
-        if block is not None:
-            # Progress.
-            log_retro_rank_0(
-                "%squery '%s' block %d / %d ... %s ... mem %.3f gb, %.1f%%."
-                % (
-                    "" if config.retro_task_validate is None else "[validate] ",
-                    prefix,
-                    block_index,
-                    len(active_blocks),
-                    os.path.basename(block["path"]),
-                    psutil.virtual_memory()[3] / 1024**3,
-                    psutil.virtual_memory()[2],
-                )
-            )
-
-            # Query block neighbors.
-            query_block_neighbors(config, db_dataset, query_dataset, index, block)
-
-        # Synchronize progress across all ranks. (for easier observation)
-        log_retro_rank_0(" > waiting for other ranks to finish block.")
-        torch.distributed.barrier()
-
-
-def query_neighbors(config: RetroPreprocessingConfig) -> None:
-    """Query pretraining datasets (train & valid).
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-    """
-
-    if not HAVE_FAISS:
-        raise ImportError(
-            "faiss is required to use the query_neighbors function. Please install faiss."
-        )
-
-    # Num threads.
-    faiss.omp_set_num_threads(64)
-
-    # Load chunk db dataset.
-    log_retro_rank_0("load chunk db dataset.")
-    db_dataset = get_db_merged_train_dataset(
-        project_dir=config.retro_project_dir,
-        chunk_length=config.retro_gpt_chunk_length,
-        eod_token_id=config.retro_tokenizers.gpt.eod,
-    )
-    db_dataset.load_doc_tuples()
-
-    # Load index.
-    log_retro_rank_0(" > get index.")
-    index = get_index(config)
-
-    # Query each (i.e., train, valid, test) dataset.
-    log_retro_rank_0(" > query.")
-    for prefix, info in vars(config.retro_gpt_chunk_datasets).items():
-        if info is None:
-            continue
-        log_retro_rank_0(
-            " > query '%s' dataset ... %d samples." % (prefix, info["num_active_chunks"])
-        )
-        query_dataset_neighbors(
-            config,
-            db_dataset,
-            info["dataset"],
-            info["num_active_chunks"],
-            prefix,
-            info["neighbor_dir"],
-            index,
-        )
diff --git a/megatron/core/datasets/retro/query/retro_dataset.py b/megatron/core/datasets/retro/query/retro_dataset.py
deleted file mode 100644
index 3316f8dbbc9..00000000000
--- a/megatron/core/datasets/retro/query/retro_dataset.py
+++ /dev/null
@@ -1,251 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-A RetroDataset wraps both:
-
-  - A GPTDataset (which is nested as GPTChunkDataset -> MultiSplitGPTDataset ->
-      GPTDataset).
-  - Neighbor IDs of chunks in the chunk database, that were saved during
-      preprocessing.
-
-Both the GPT sample data and the neighbor IDs are returned within a sample from
-this dataset.
-"""
-
-import os
-from typing import Dict, Optional, Tuple
-
-import numpy as np
-import torch
-
-from megatron.core.datasets.retro.db.dataset import DBDataset
-from megatron.core.datasets.retro.db.utils import get_merged_train_dataset as get_db_dataset
-from megatron.core.datasets.retro.utils import BlockPathMap, log_retro_rank_0
-from megatron.core.models.retro import RetroConfig
-
-from .gpt_chunk_dataset import GPTChunkDataset, build_gpt_chunk_datasets_from_gpt_datasets
-from .utils import get_query_dir
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-
-class RetroDataset(torch.utils.data.Dataset):
-    """Dataset of retro samples.
-
-    Each sample contains the original GPT sample, along with the token IDs
-    of each neighbor of each chunk within the sequence. Neighbor array has
-    shape (num_chunks_per_sample, num_neighbors, num_retrieved_tokens).
-
-    ** Note: chunk dataset wraps original GPT dataset (see gpt_chunk_dataset.py).
-
-    Args:
-        num_queried_samples (int): Total number of queried samples.
-        num_neighbors (int): Total number of saved neighbors.
-        num_retrieved_chunks (int): Number of retrieved chunks
-            (e.g., 2 for neighbor + continuation).
-        block_size (int): Number of neighbor entries per file.
-        db_dataset (DBDataset): Chunk database used for retrieval.
-        chunk_dataset (GPTChunkDataset): GPT chunk dataset, which is a wrapper
-            around a standard GPT dataset that breaks each sample into chunks.
-        neighbor_path_map (BlockPathMap): Mapping of neighbor ID to file path.
-    """
-
-    def __init__(
-        self,
-        num_queried_samples: int,
-        num_neighbors: int,
-        num_retrieved_chunks: int,
-        block_size: int,
-        db_dataset: DBDataset,
-        chunk_dataset: GPTChunkDataset,
-        neighbor_path_map: BlockPathMap,
-    ):
-        super().__init__()
-
-        self.num_queried_samples = num_queried_samples
-        self.num_neighbors = num_neighbors
-        self.num_retrieved_chunks = num_retrieved_chunks
-        self.block_size = block_size
-        self.db_dataset = db_dataset
-        self.chunk_dataset = chunk_dataset
-        self.neighbor_path_map = neighbor_path_map
-
-    def __len__(self) -> int:
-        """Dataset length.
-
-        Returns:
-            Number of samples in dataset.
-        """
-        return len(self.chunk_dataset.sample_dataset)
-
-    def __getitem__(self, sample_idx: int) -> dict:
-        """Get dataset sample.
-
-        Args:
-            sample_idx (int): Index of sample in dataset.
-
-        Returns:
-            A dict consisting of GPT sample (attribute 'text') and corresponding neighbor chunk IDs
-            ('neighbor_chunks', for indexing chunk database) and neighbor token IDs
-            (corresponding chunk database GPT tokens).
-        """
-        if not HAVE_H5PY:
-            raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-        n_chunks_per_sample = self.chunk_dataset.n_chunks_per_sample
-
-        # Wrap sample idx around number of queried samples.
-        sample_idx = sample_idx % self.num_queried_samples
-
-        # Get standard sample.
-        sample = self.chunk_dataset.sample_dataset[sample_idx]
-
-        # Sample idx to chunk idxs.
-        chunk_idxs = list(
-            range(sample_idx * n_chunks_per_sample, (sample_idx + 1) * n_chunks_per_sample)
-        )
-
-        # Collect retrieved tokens.
-        all_retrieved_chunk_ids = []
-        all_retrieved_token_ids = []
-        for chunk_idx in chunk_idxs:
-            # Neighbor chunk ids.
-            neighbor_path = self.neighbor_path_map[chunk_idx]
-            with h5py.File(neighbor_path, "r") as f:
-                neighbor_chunk_ids = f["neighbors"][
-                    chunk_idx % self.block_size, : self.num_neighbors
-                ].tolist()
-
-            # Retrieved (neighbor + continuation) token ids.
-            retrieved_chunk_ids = []
-            retrieved_token_ids = []
-            for neighbor_chunk_id in neighbor_chunk_ids:
-                current_chunk_ids = [
-                    i % len(self.db_dataset)
-                    for i in range(neighbor_chunk_id, neighbor_chunk_id + self.num_retrieved_chunks)
-                ]
-                current_token_ids = [self.db_dataset[ci]["text"] for ci in current_chunk_ids]
-                retrieved_chunk_ids.append(current_chunk_ids)
-                retrieved_token_ids.append(current_token_ids)
-
-            # Collect retrieved tokens.
-            all_retrieved_chunk_ids.append(retrieved_chunk_ids)
-            all_retrieved_token_ids.append(retrieved_token_ids)
-
-        # Reshape retrieved tokens.
-        all_retrieved_chunk_ids = np.array(all_retrieved_chunk_ids).reshape(
-            (n_chunks_per_sample, self.num_neighbors, -1)
-        )
-        all_retrieved_token_ids = np.array(all_retrieved_token_ids).reshape(
-            (n_chunks_per_sample, self.num_neighbors, -1)
-        )
-
-        # Sample.
-        sample: Dict[str, np.ndarray] = {
-            **sample,
-            "neighbor_chunks": all_retrieved_chunk_ids,
-            "neighbor_tokens": all_retrieved_token_ids,
-        }
-
-        return sample
-
-
-def get_retro_datasets(
-    config: RetroConfig, gpt_datasets: dict, sample_length: int, eod_token_id: int
-) -> Tuple[Optional[RetroDataset], Optional[RetroDataset], Optional[RetroDataset]]:
-    """Get train, valid, test retro datasets.
-
-    Args:
-        config (RetroConfig): Retro preprocessing config.
-        gpt_datasets (dict): Mapping of data split key
-            ('train', 'valid', or 'test') to the original sequence-length
-            GPT dataset (i.e., not the chunk dataset).
-        sample_length (int): Alias to `sequence_length`.
-        eod_token_id (int): GPT EOD token ID.
-
-    Returns:
-        A tuple of 'train', 'valid', and 'test' `RetroDataset`s.
-    """
-
-    # DB dataset.
-    db_dataset = get_db_dataset(
-        project_dir=config.retro_project_dir,
-        chunk_length=config.retro_chunk_length,
-        eod_token_id=eod_token_id,
-    )
-
-    # GPT chunk datasets.
-    chunk_ds_info_map = build_gpt_chunk_datasets_from_gpt_datasets(
-        project_dir=config.retro_project_dir,
-        gpt_datasets=gpt_datasets,
-        sample_length=sample_length,
-        chunk_length=config.retro_chunk_length,
-    )
-
-    # Retro datasets.
-    retro_dataset_map: Dict[str, Optional[RetroDataset]] = {}
-    query_dir = get_query_dir(config.retro_project_dir)
-    for data_key, chunk_ds_info in chunk_ds_info_map.items():
-        # Skip unused datasets.
-        if chunk_ds_info is None:
-            retro_dataset_map[data_key] = None
-            continue
-
-        # For consistency with preprocessing, the neighbor_dir is overwritten
-        # (from its setting in `build_gpt_chunk_datasets_from_gpt_datasets()`
-        # above). This is one piece -- along with setting data_path and
-        # train_samples from config.json -- of ensuring consistency between
-        # preprocessing and pretraining.
-        chunk_dataset = chunk_ds_info["dataset"]
-        chunk_ds_info["neighbor_dir"] = os.path.join(
-            query_dir, config.retro_neighbor_dirs[data_key]
-        )
-        neighbor_dir = chunk_ds_info["neighbor_dir"]
-        neighbor_path_map = BlockPathMap.from_dir(
-            dir=neighbor_dir, block_size=config.retro_block_size
-        )
-
-        # Verify num chunks.
-        n_active_chunks = chunk_ds_info["num_active_chunks"]
-        n_neighbor_chunks = neighbor_path_map.max_idx
-
-        if not os.path.isdir(neighbor_dir):
-            if torch.distributed.get_rank() == 0:
-                raise Exception(
-                    "neighbor directory '%s' not found; please "
-                    "compare --train-samples, --seq-length, --seed, "
-                    "--eval-iters, and --eval-interval, with "
-                    "retro preprocessing args." % neighbor_dir
-                )
-            torch.distributed.barrier()
-            exit()
-
-        if config.retro_verify_neighbor_count and n_active_chunks != n_neighbor_chunks:
-            if torch.distributed.get_rank() == 0:
-                log_retro_rank_0("neighbor_dir : %s" % neighbor_dir)
-                log_retro_rank_0("neighbor_path_map : %s" % neighbor_path_map)
-                raise Exception(
-                    "num sampled chunks (%d) != num neighbor chunks "
-                    "(%d); did you complete querying the entire "
-                    "pretraining dataset?" % (n_active_chunks, n_neighbor_chunks)
-                )
-            torch.distributed.barrier()
-            exit()
-
-        # Retro dataset.
-        retro_dataset_map[data_key] = RetroDataset(
-            num_queried_samples=gpt_datasets[data_key][1],
-            num_neighbors=config.retro_num_neighbors,
-            num_retrieved_chunks=config.retro_num_retrieved_chunks,
-            block_size=config.retro_block_size,
-            db_dataset=db_dataset,
-            chunk_dataset=chunk_dataset,
-            neighbor_path_map=neighbor_path_map,
-        )
-
-    return (retro_dataset_map["train"], retro_dataset_map["valid"], retro_dataset_map["test"])
diff --git a/megatron/core/datasets/retro/query/utils.py b/megatron/core/datasets/retro/query/utils.py
deleted file mode 100644
index b4e0c67009a..00000000000
--- a/megatron/core/datasets/retro/query/utils.py
+++ /dev/null
@@ -1,35 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Utilities for querying the pretraining dataset."""
-
-import os
-
-from megatron.core.datasets.megatron_dataset import MegatronDataset
-
-
-def get_query_dir(project_dir: str) -> str:
-    """Get root directory of all saved query data.
-
-    Args:
-        project_dir (str): Retro project dir.
-
-    Returns:
-        Path to query sub-directory in Retro project.
-    """
-    return os.path.join(project_dir, "query")
-
-
-def get_neighbor_dir(project_dir: str, key: str, dataset: MegatronDataset) -> str:
-    """Get directory containing neighbor IDs for a dataset (i.e., train, valid, or test).
-
-    Args:
-        project_dir (str): Retro project dir.
-        key (str): Dataset split key; 'train', 'valid', or 'test'.
-        dataset (MegatronDataset): Dataset containing unique hash for finding corresponding neighbors.
-
-    Returns:
-        Path to directory containing this dataset's neighbors within Retro project.
-    """
-    return os.path.join(
-        get_query_dir(project_dir), os.path.basename(f"{key}_{dataset.unique_description_hash}")
-    )
diff --git a/megatron/core/datasets/retro/utils.py b/megatron/core/datasets/retro/utils.py
deleted file mode 100644
index 5d9900697fc..00000000000
--- a/megatron/core/datasets/retro/utils.py
+++ /dev/null
@@ -1,386 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Utilities for Retro preprocessing."""
-
-import glob
-import logging
-import os
-from types import SimpleNamespace
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict
-
-import numpy as np
-import torch
-from torch.distributed import ProcessGroup
-
-from megatron.core import parallel_state
-from megatron.core.datasets.retro.config import RetroPreprocessingConfig
-from megatron.core.datasets.retro.query.multi_split_gpt_dataset import (
-    MultiSplitGPTDataset,
-    MultiSplitGPTDatasetConfig,
-)
-from megatron.core.utils import log_single_rank
-
-logger = logging.getLogger(__name__)
-
-try:
-    from tqdm import tqdm
-
-    HAVE_TQDM = True
-except ImportError:
-    HAVE_TQDM = False
-
-try:
-    import h5py
-
-    HAVE_H5PY = True
-except ImportError:
-    HAVE_H5PY = False
-
-
-class Block(TypedDict):
-    """Specific block arg type to mute mypy."""
-
-    range: Tuple[int, int]
-    path: str
-
-
-def log_retro_rank_0(message: str) -> None:
-    """Log on rank 0.
-
-    Args:
-        message (str): Message to log.
-    """
-    log_single_rank(logger, logging.INFO, "[RETRO] " + message)
-
-
-def retro_makedir(config: RetroPreprocessingConfig, path: str) -> None:
-    """Make a directory, conditional on not being in validation mode.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-        path (str): Path to directory.
-    """
-    if config.retro_task_validate is None:
-        os.makedirs(path, exist_ok=True)
-
-
-def extract_data_config(config: RetroPreprocessingConfig) -> MultiSplitGPTDatasetConfig:
-    """Extract data config from dataset.
-
-    Args:
-        config (RetroPreprocessingConfig): Retro preprocessing config.
-
-    Returns:
-        The config object used to build the dataset.
-    """
-    return config.retro_gpt_chunk_datasets.train["dataset"].sample_dataset.config
-
-
-def get_num_chunks_per_sample(sample_length: int, chunk_length: int) -> int:
-    """Compute seq_length // chunk_length.
-
-    Args:
-        sample_length (int): Alias of `sequence_length`.
-        chunk_length (int): Retro chunk length (e.g., 64).
-
-    Returns:
-        Number of chunks per sample (i.e., `sequence_length` / `chunk_length`).
-    """
-    assert sample_length % chunk_length == 0
-    return sample_length // chunk_length
-
-
-class GPTToTextDataset(torch.utils.data.Dataset):
-    """Dataset to convert GPT tokens to text.
-
-    Args:
-        gpt_dataset (MultiSplitGPTDataset): GPT dataset, which outputs GPT token samples.
-        gpt_tokenizer (Any): GPT tokenizer.
-    """
-
-    def __init__(self, gpt_dataset: MultiSplitGPTDataset, gpt_tokenizer: Any):
-        super().__init__()
-
-        self.gpt_dataset = gpt_dataset
-        self.gpt_tokenizer = gpt_tokenizer
-
-    def __len__(self) -> int:
-        """Dataset length.
-
-        Returns:
-            Number of samples in the dataset.
-        """
-        return len(self.gpt_dataset)
-
-    def __getitem__(self, idx: int) -> dict:
-        """Get dataset sample.
-
-        Args:
-            idx (int): Index of sample.
-
-        Returns:
-            A dict containing attribute 'text' of type string.
-        """
-        gpt_token_ids = self.gpt_dataset[idx]["text"].tolist()
-        text = self.gpt_tokenizer.detokenize(gpt_token_ids)
-        return {"text": text}
-
-
-def get_blocks(
-    dirname: str, n_samples: int, block_size: int, validate: Optional[Callable] = None
-) -> SimpleNamespace:
-    """Divide range [0, num_samples) to sequence of block ranges.
-
-    This is a core method within the concept of block processing. The idea
-    is to divide a range (size n_samples) into a sequence of blocks. Each
-    block corresponds to a file within 'dirname' with name
-    '{start_idx}-{end_idx}.hdf5'. This method checks for the existence of
-    these files, and returns two lists, one for existing blocks and one for
-    missing blocks.
-
-    Args:
-        dirname (str): Path to directory containing block files.
-        n_samples (int): Ideal number of samples.
-            The total number of saved block data is <=n_samples.
-        block_size (int): Max number of samples per block file (e.g., 100000).
-        validate (Callable): Method for validating each block file during load.
-
-    Returns:
-        A namespace consisting of 2 lists: existing blocks, and missing blocks.
-        The total number of samples between the existing and missing blocks should
-        equal n_samples above.
-    """
-
-    if not HAVE_TQDM:
-        raise ImportError("tqdm is required to use the RetroDataset. Please install tqdm.")
-
-    if not HAVE_H5PY:
-        raise ImportError("h5py is required to use the RetroDataset. Please install h5py.")
-
-    assert os.path.isdir(dirname), "missing directory '%s.'" % dirname
-
-    # Block ranges.
-    block_start_idxs = list(range(0, n_samples, block_size))
-    block_end_idxs = [min(n_samples, i + block_size) for i in block_start_idxs]
-    block_ranges = list(zip(block_start_idxs, block_end_idxs))
-
-    # All block files (existing + missing).
-    n_digits = int(np.ceil(np.log(n_samples) / np.log(10)) + 1)
-
-    all_blocks: List[Block] = [
-        {
-            "range": r,
-            "path": os.path.join(
-                dirname, "%s-%s.hdf5" % tuple([str(i).zfill(n_digits) for i in r])
-            ),
-        }
-        for r in block_ranges
-    ]
-    all_block_path_set = set(block["path"] for block in all_blocks)
-
-    # Validate function.
-    validate = (lambda f: None) if validate is None else validate
-
-    # Delete corrupt files.
-    if torch.distributed.get_rank() == 0:
-        existing_block_paths = [
-            block["path"] for block in all_blocks if os.path.exists(block["path"])
-        ]
-        for index, path in enumerate(tqdm(existing_block_paths, "validating block.")):
-            assert path in all_block_path_set, "unexpected filename, '%s'." % path
-
-            try:
-                f = h5py.File(path, "r")
-            except Exception:
-                os.remove(path)
-                continue
-
-            try:
-                validate(f)
-            except Exception:
-                os.remove(path)
-            finally:
-                f.close()
-
-    # Wait for files to be deleted.
-    torch.distributed.barrier()
-
-    # Collect blocks.
-    blocks = SimpleNamespace(
-        existing=[b for b in all_blocks if os.path.exists(b["path"])],
-        missing=[b for b in all_blocks if not os.path.exists(b["path"])],
-    )
-
-    return blocks
-
-
-def get_blocks_by_rank(
-    dirname: str,
-    n_samples: int,
-    block_size: int,
-    validate: Optional[Callable] = None,
-    sample: Optional[float] = None,
-    process_group: Optional[ProcessGroup] = None,
-) -> SimpleNamespace:
-    """Divide existing and missing blocks evenly across all ranks.
-
-    See 'get_blocks()' above for description. The returned lists of existing and
-    missing blocks are split evenly across ranks via interleaving. This way,
-    each rank has a roughly equal number of blocks to process for a
-    downstream operation.
-
-    Args:
-        dirname (str): Path to directory containing block files.
-        n_samples (int): Ideal number of samples. The total number of saved block data
-            is <=n_samples.
-        block_size (int): Max number of samples per block file (e.g., 100000).
-        validate (Callable): Method for validating each block file during load.
-        sample (Optional[float]): If provided, sample a random subset of the blocks.
-            Used for validating preprocessing correctness.
-        process_group (Optional[ProcessGroup]): Process group for distributed operations.
-            If None, uses data parallel group.
-
-    Returns:
-        A namespace consisting of 2 lists: existing blocks, and missing blocks.
-        Each of these two lists is potentially a sub-sample of the total set of
-        existing and missing blocks, depending on whether sampling is used.
-        Additionally, the attributes n_existing_world and n_missing_world are the
-        total number of existing and missing blocks, independent of samples.
-        Therefore, (n_existing_world + n_missing_world) * block_size == n_samples.
-    """
-
-    if process_group is None:
-        process_group = parallel_state.get_data_parallel_group()
-
-    # Get world blocks.
-    blocks = get_blocks(dirname, n_samples, block_size, validate)
-
-    # This rank's existing and missing files.
-    rank_existing_blocks = blocks.existing[
-        process_group.rank() : len(blocks.existing) : process_group.size()
-    ]
-    rank_missing_blocks = blocks.missing[
-        process_group.rank() : len(blocks.missing) : process_group.size()
-    ]
-
-    # Extend rank's existing and missing blocks (with None) such that all ranks
-    # have equal length lists. This allows for easier tracking of global progress.
-    def get_world_max(n: int) -> int:
-        """Get max value across ranks.
-
-        Args:
-            n (int): Value on this rank.
-
-        Returns:
-            Max value across all ranks.
-        """
-        n_tensor = torch.cuda.LongTensor([n])
-        torch.distributed.all_reduce(n_tensor, op=torch.distributed.ReduceOp.MAX)
-        return n_tensor.item()
-
-    max_n_existing = get_world_max(len(rank_existing_blocks))
-    max_n_missing = get_world_max(len(rank_missing_blocks))
-
-    rank_existing_blocks += [None] * (max_n_existing - len(rank_existing_blocks))
-    rank_missing_blocks += [None] * (max_n_missing - len(rank_missing_blocks))
-
-    # Collect blocks.
-    blocks = SimpleNamespace(
-        n_existing_world=len(blocks.existing),
-        n_missing_world=len(blocks.missing),
-        existing=rank_existing_blocks,
-        missing=rank_missing_blocks,
-    )
-
-    if sample is not None:
-        # Sample existing and missing blocks evenly across all ranks. The
-        # returned lists of blocks are randomly sampled (without replacement)
-        # to yield `sample * len(blocks)` number of blocks.
-
-        # Randomly sample blocks.
-        def sample_blocks(_blocks: List[Optional[Dict]]) -> List[Optional[Dict]]:
-            """Sample a random subset of all blocks.
-
-            Args:
-                _blocks (List[Optional[Dict]]): List of all blocks.
-
-            Returns:
-                A random subset of the blocks.
-            """
-            n_blocks_sample = int(np.ceil(sample * len(_blocks)))
-            sampled_blocks: List[Optional[Dict]] = [b for b in _blocks if b is not None]
-
-            np.random.seed(None)
-            np.random.shuffle(sampled_blocks)
-
-            sampled_blocks = sampled_blocks[:n_blocks_sample]
-            sampled_blocks += [None] * (n_blocks_sample - len(sampled_blocks))
-
-            return sampled_blocks
-
-        blocks.existing = sample_blocks(blocks.existing)
-        blocks.missing = sample_blocks(blocks.missing)
-
-    return blocks
-
-
-class BlockPathMap:
-    """Map an index to its containing block path.
-
-    The common use for this class is to have a directory of files containing
-    blocks of processed data, of uniform block size (e.g., 100k samples per
-    file). Each file must follow a naming convention of 'startIdx-endIdx.[ext]',
-    where 'endIdx' minus 'startIdx' must equal the block size, with the possible
-    exception of the final block. Given an input index, this class maps the
-    index to the containing block file.
-
-    Args:
-        block_paths (List[str]): List of paths to saved block files.
-        block_size (int): Max number of samples per block file (e.g., 100000).
-    """
-
-    @classmethod
-    def from_dir(cls, dir: str, block_size: int, ext: str = "hdf5") -> Any:
-        """Get list of block files, and create map.
-
-        Args:
-            dir (str): Path to directory containing saved block files.
-            block_size (int): Max number of samples per block file (e.g., 100000).
-            ext (str): Block file extension (e.g., 'hdf5').
-
-        Returns:
-            A mapping of sample index to block file path.
-        """
-        assert os.path.isdir(dir), f"directory not found, '{dir}'."
-        return cls(sorted(glob.glob(dir + f"/*.{ext}")), block_size)
-
-    def __init__(self, block_paths: List[str], block_size: int):
-        self.max_idx = 0
-        self.block_path_map = {}
-        for block_path in block_paths:
-            name = os.path.splitext(os.path.basename(block_path))[0]
-            start_idx, end_idx = [int(i) for i in name.split("-")]
-            self.block_path_map[start_idx] = block_path
-            self.max_idx = max(self.max_idx, end_idx)
-        self.block_size = block_size
-
-    def __str__(self) -> str:
-        """Stringify the mapping.
-
-        Returns:
-            A string representation of this block path map.
-        """
-        return "%d paths" % len(self.block_path_map)
-
-    def __getitem__(self, idx: int) -> str:
-        """Get block path from index.
-
-        Args:
-            idx (int): Index of sample.
-
-        Returns:
-            The path to the block file containing the sample index.
-        """
-        block_start_idx = self.block_size * (idx // self.block_size)
-        block_path = self.block_path_map[block_start_idx]
-        return block_path
diff --git a/megatron/core/enums.py b/megatron/core/enums.py
index fcca219badd..9b76bc52a87 100644
--- a/megatron/core/enums.py
+++ b/megatron/core/enums.py
@@ -7,8 +7,6 @@ class ModelType(enum.Enum):
     """Model type."""
 
     encoder_or_decoder = 1
-    retro_encoder = 2
-    retro_decoder = 3
 
     @property
     def encoder_and_decoder(self):
diff --git a/megatron/core/models/retro/__init__.py b/megatron/core/models/retro/__init__.py
deleted file mode 100644
index ea7cea6d8fb..00000000000
--- a/megatron/core/models/retro/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""
-Exports:
-
-  - RetroConfig: configuration dataclass for RetroModel.
-  - RetroModel: The Retro model.
-  - get_retro_decoder_block_spec: Get spec for Retro decoder transformer block.
-"""
-
-from .config import RetroConfig
-from .decoder_spec import get_retro_decoder_block_spec
-from .model import RetroModel
diff --git a/megatron/core/models/retro/base_attention.py b/megatron/core/models/retro/base_attention.py
deleted file mode 100644
index fa07a51ddcd..00000000000
--- a/megatron/core/models/retro/base_attention.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Base class for decoder and encoder attention modules."""
-
-from megatron.core.models.retro.config import RetroConfig
-from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer.attention import CrossAttention, CrossAttentionSubmodules
-from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.transformer.module import MegatronModule
-
-
-class BaseRetroCrossAttention(MegatronModule):
-    """Base class for Retro cross attention, for both encoder & decoder layers.
-
-    This class collects the retro arguments below (i.e., num neighbors, chunk
-    length, and retrieve length) for use in Retro's custom cross attention
-    operators.
-
-    Args:
-        config (RetroConfig): Retro config.
-        submodules (CrossAttentionSubmodules): Cross attention submodules.
-        layer_number (int): Layer number within transformer block.
-        attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
-        pg_collection (ProcessGroupCollection): Model communication process groups.
-    """
-
-    def __init__(
-        self,
-        config: RetroConfig,
-        submodules: CrossAttentionSubmodules,
-        layer_number: int = 1,
-        attn_mask_type: AttnMaskType = AttnMaskType.padding,
-        pg_collection: ProcessGroupCollection = None,
-    ):
-        super().__init__(config=config)
-
-        self.attn = CrossAttention(
-            config=config,
-            submodules=submodules,
-            layer_number=layer_number,
-            attn_mask_type=attn_mask_type,
-            pg_collection=pg_collection,
-        )
-
-        self.retro_num_neighbors = config.retro_num_neighbors
-        self.retro_chunk_length = config.retro_chunk_length
-        self.retro_retrieved_length = config.retro_retrieved_length
diff --git a/megatron/core/models/retro/config.py b/megatron/core/models/retro/config.py
deleted file mode 100644
index 1b486767264..00000000000
--- a/megatron/core/models/retro/config.py
+++ /dev/null
@@ -1,88 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-
-"""Configuration dataclass for a RetroModel."""
-
-import os
-from dataclasses import dataclass
-
-from megatron.core.transformer import TransformerConfig
-from megatron.core.transformer.enums import AttnBackend
-from megatron.core.utils import is_te_min_version
-
-
-@dataclass
-class RetroConfig(TransformerConfig):
-    """Configuration object for Retro models."""
-
-    # Retro.
-    retro_project_dir: str = None
-    """Retro project directory, which contains the preprocessed data for for pretraining. This
-       directory is built during preprocessing (see tools/retro/README.md), and contains
-       subdirectories for the chunk database and pretraining neighbors.
-    """
-
-    retro_block_size: int = None
-    """Number of records to load per data file, as saved during preprocessing. Block processing is
-       used for efficient data preprocessing.
-    """
-
-    retro_chunk_length: int = None
-    """Chunk length used for performing chunked- cross-attention (CCA)."""
-
-    retro_encoder_num_layers: int = 2
-    """Number of layers to use for the retrieval encoder."""
-
-    retro_encoder_hidden_dropout: float = 0.1
-    """Hidden dropout for retrieval encoder."""
-
-    retro_encoder_attention_dropout: float = 0.1
-    """Attention dropout for retrieval encoder."""
-
-    retro_neighbor_dirs: dict = None
-    """Directory names of saved neighbor id files for train, valid, and test datasets."""
-
-    retro_num_neighbors: int = 2
-    """Number of neighbors to retrieve during pretraining."""
-
-    retro_num_retrieved_chunks: int = 2
-    """Number of chunks to retrieve from the retrieval database."""
-
-    retro_retrieved_length: int = None
-    """Cached value of retro_num_retrieved_chunks * retro_chunk_length (i.e., the total number of
-       retrieved tokens; neighbor + continuation).
-    """
-
-    retro_split_preprocessing: str = None
-    """Data split used during data preprocessing."""
-
-    retro_verify_neighbor_count: bool = True
-    """Verify that len(GPT dataset) == len(saved neighbors)."""
-
-    def __post_init__(self) -> None:
-        """Validate Retro config."""
-
-        super().__post_init__()
-
-        self.attention_backend = AttnBackend.unfused
-
-        # Validate Transformer Engine version.
-        if is_te_min_version("1.3"):
-            try:
-                assert os.getenv("NVTE_FLASH_ATTN") == "0"
-                assert os.getenv("NVTE_FUSED_ATTN") == "0"
-            except Exception as e:
-                raise Exception(
-                    "When using Transformer Engine >= 1.3, environment vars NVTE_FLASH_ATTN "
-                    "and NVTE_FUSED_ATTN most both be defined and set to '0'. "
-                    "Currently, NVTE_FLASH_ATTN == %s, NVTE_FUSED_ATTN == %s."
-                    % (
-                        os.getenv("NVTE_FLASH_ATTN", "[unset]"),
-                        os.getenv("NVTE_FUSED_ATTN", "[unset]"),
-                    )
-                )
-
-        # Preprocessing split should be defined.
-        assert self.retro_split_preprocessing is not None
-
-        # Pre-compute retrieved length.
-        self.retro_retrieved_length = self.retro_num_retrieved_chunks * self.retro_chunk_length
diff --git a/megatron/core/models/retro/decoder_attention.py b/megatron/core/models/retro/decoder_attention.py
deleted file mode 100644
index 5aedb053112..00000000000
--- a/megatron/core/models/retro/decoder_attention.py
+++ /dev/null
@@ -1,319 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-
-"""Retro's cross attention modules for the decoder block."""
-
-from functools import partial
-from typing import Callable, Optional
-
-import numpy as np
-import torch
-from torch import Tensor
-
-from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
-from megatron.core.inference.contexts import BaseInferenceContext
-from megatron.core.models.retro.base_attention import BaseRetroCrossAttention
-from megatron.core.models.retro.config import RetroConfig
-from megatron.core.models.retro.utils import get_all_true_mask
-from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer import ModuleSpec
-from megatron.core.transformer.attention import CrossAttentionSubmodules
-from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.transformer.transformer_block import TransformerBlock
-from megatron.core.utils import deprecate_inference_params
-
-
-class RetroDecoderCrossAttention(BaseRetroCrossAttention):
-    """Retro decoder's chunked cross attention operator.
-
-    See this paper for more details: https://arxiv.org/abs/2112.04426.
-    Neighboring chunks retrieved from the chunk database are used here for
-    chunked-cross attention.
-
-    ** Note about 'encoder_block_spec' **
-
-    Retro is an encoder-decoder model that uses its encoder for encoding
-    neighboring chunks that are retrieved from a chunk database. These
-    encoded neighbors are then used in the decoder stack for performing
-    chunked-cross attention (see paper link above).
-
-    In contrast to the T5 model, the encoder and decoder are computationally
-    intertwined, since the input to the encoder is the output of the self-
-    attention of the first decoder layer. As such, the encoder block itself
-    is instantiated within the first Retro decoder layer, in order to receive
-    the self-attention's output. (Note, that only the first decoder layer
-    instantiates an encoder block, and the remaining decoder layers use the
-    encoder output from the first decoder layer.)
-
-    Args:
-        config (RetroConfig): Retro config.
-        submodules (CrossAttentionSubmodules): Cross attention submodules.
-        layer_number (int): Layer number within transformer block.
-        attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
-        encoder_block_spec (ModuleSpec): The first Retro decoder layer is
-            provided with a transformer block spec to construct the neighbor encoder.
-        pg_collection (ProcessGroupCollection): Model communication process groups.
-    """
-
-    def __init__(
-        self,
-        config: RetroConfig,
-        submodules: CrossAttentionSubmodules,
-        layer_number: int = 1,
-        attn_mask_type: AttnMaskType = AttnMaskType.padding,
-        encoder_block_spec: ModuleSpec = None,
-        pg_collection: ProcessGroupCollection = None,
-    ):
-        super().__init__(
-            config=config,
-            submodules=submodules,
-            layer_number=layer_number,
-            attn_mask_type=attn_mask_type,
-            pg_collection=pg_collection,
-        )
-
-        if encoder_block_spec:
-            self.encoder = TransformerBlock(
-                config=config,
-                spec=encoder_block_spec,
-                pre_process=True,
-                post_process=False,
-                pg_collection=pg_collection,
-            )
-            # self._encoder_key = 'encoder' # ... necessary?
-        else:
-            self.encoder = None
-
-    def forward(
-        self,
-        hidden_states: Tensor,
-        attention_mask: Tensor,
-        key_value_states: Tensor = None,
-        inference_context: BaseInferenceContext = None,
-        # rotary_pos_emb: Tensor = None, # ... unsupported for retro.
-        *,
-        inference_params: Optional[BaseInferenceContext] = None,
-    ) -> dict:
-        """Cross attention for Retro decoder.
-
-        Notation:
-            ns : Sequence length.
-            bs : Batch size.
-            d  : Hidden size.
-            l  : Number of chunks per sample (i.e., seq_length/chunk_length).
-            m  : Number of tokens per chunk.
-            k  : Number of neighbors.
-            r  : Number of retrieved tokens (neighbors + continuation).
-
-        Args:
-            hidden_states (Tensor): Transformer layer hidden states.
-            attention_mask (Tensor): Attention mask.
-            key_value_states (Tensor): Neighbor embeddings if first decoder layer,
-                else encoder output.
-            inference_context (BaseInferenceContext): Inference context.
-
-        Returns:
-            A dict consisting of the attention output and context, along with
-                other scalars necessary for performing the downstream bias-dropout-add.
-        """
-
-        # hidden_states: [ ns, bs, d ]
-        # key_value_states: [ r, k*bs*l, d ]
-
-        inference_context = deprecate_inference_params(inference_context, inference_params)
-
-        ns, bs, d = hidden_states.shape
-        l = int(np.ceil(ns / self.retro_chunk_length))
-
-        # Retrieve neighbors.
-        if self.encoder:
-            # Sequence length remainder.
-            first_ns = ns % self.retro_chunk_length
-
-            # Case 1: Sequence length not divisible by chunk length.
-            if first_ns > 0:
-                # Split sequence into first partial chunk & remaining chunks.
-                first_chunk, rest_chunk = (hidden_states[:first_ns], hidden_states[first_ns:])
-
-                # Pad partial chunk with zeros.
-                first_chunk = torch.nn.functional.pad(
-                    first_chunk, (0, 0, 0, 0, 0, self.retro_chunk_length - first_ns), "constant", 0
-                )
-
-                # Concatenate padded chunk with remaining chunks.
-                chunked_output = torch.cat((first_chunk, rest_chunk), dim=0)  # [ l*m, bs, d ]
-
-            # Case 2: Sequence length is divisible by chunk length.
-            else:
-                chunked_output = hidden_states  # [ l*m, bs, d ]
-
-            # Chunk & permute hidden states.
-            # - hidden_states:  [ l*m, bs, d ]
-            # - chunked_output: [ m, bs*l, d ]
-            chunked_output = (
-                chunked_output.reshape(l, self.retro_chunk_length, bs, d)
-                .permute(1, 2, 0, 3)
-                .reshape(self.retro_chunk_length, bs * l, d)
-                .contiguous()
-            )
-
-            # flash attn: [ b, h, sq, sk ]
-            # fused attn: [ b, 1, 1, sq ]
-            chunked_output_mask = get_all_true_mask(
-                size=(1, 1, chunked_output.shape[0], key_value_states.shape[0]),
-                device=chunked_output.device,
-            )
-
-            # Encode neighbors. (Note: 'key_value_states' re-assigned here.)
-            key_value_states = self.encoder(
-                hidden_states=key_value_states,
-                attention_mask=attention_mask,
-                context=chunked_output,
-                context_mask=chunked_output_mask,
-                inference_context=inference_context,
-            )  # [ r, k*bs*l, d ]
-            key_value_states = key_value_states.reshape(
-                self.retro_retrieved_length * self.retro_num_neighbors, bs * l, d
-            )  # [ r*k, bs*l, d ]
-
-        # Attend starting at last token of first chunk.
-        pad = (ns - 1) % self.retro_chunk_length
-        attending_chunks = hidden_states[pad:]
-
-        # Pad attending tokens to sequence length.
-        padded_chunks = torch.nn.functional.pad(
-            attending_chunks, (0, 0, 0, 0, 0, self.retro_chunk_length - 1), "constant", 0
-        )
-
-        # Permute attending chunks.
-        # - padded_chunks:         [ l*m, bs, d ]
-        # - padded_chunked_output: [ m, bs*l, d ] (matches 'chunked_output' above)
-        padded_chunked_output = padded_chunks.reshape(l, self.retro_chunk_length, bs, d).permute(
-            1, 2, 0, 3
-        )
-        padded_chunked_output = padded_chunked_output.reshape(
-            self.retro_chunk_length, bs * l, d
-        ).contiguous()
-
-        # flash attn: [ b, h, sq, sk ]
-        # fused attn: [ b, 1, 1, sq ]
-        padded_chunked_output_mask = get_all_true_mask(
-            size=(1, 1, padded_chunked_output.shape[0], key_value_states.shape[0]),
-            device=padded_chunked_output.device,
-        )
-
-        # Attend to encoded neighbors.
-        attention_output, attention_bias = self.attn(
-            hidden_states=padded_chunked_output,
-            attention_mask=padded_chunked_output_mask,
-            key_value_states=key_value_states,
-        )
-
-        # Return dimensions for bias-dropout step.
-        return {
-            "ns": ns,
-            "bs": bs,
-            "d": d,
-            "l": l,
-            "pad": pad,
-            "attention_output": attention_output,  # [ m, bs*l, d ]
-            "attention_bias": attention_bias,  # [ d ]
-            "context": key_value_states,  # [ r*k, bs*l, d ]
-        }
-
-
-class RetroDecoderBiasDropoutAdd(MegatronModule):
-    """Retro decoder's bias-dropout-add operator.
-
-    This operator takes care of reshaping and permuting the output from the
-    chunk dimension to the sequence dimension.
-
-    Args:
-        config (RetroConfig): Retro config.
-    """
-
-    def __init__(self, config: RetroConfig):
-        super().__init__(config=config)
-        self.retro_chunk_length = config.retro_chunk_length
-
-    @classmethod
-    def _forward(
-        cls,
-        x_with_bias: dict,
-        residual: Tensor,
-        prob: float,
-        retro_chunk_length: int,
-        bias_dropout_add: Callable,
-    ) -> Tensor:
-        """Per-chunk bias-dropout-add.
-
-        Args:
-            x_with_bias (dict): Attention output and bias, along with other Retro
-                relevant parameters.
-            residual (Tensor): Transformer layer residual.
-            prob (float): Dropout probability.
-            retro_chunk_length (int): Retro chunk length (e.g., 64).
-            bias_dropout_add (Callable): Bias-dropout-add function.
-
-        Returns:
-            Output of bias-dropout-add.
-        """
-
-        # Extract input dict.
-        ns = x_with_bias["ns"]
-        bs = x_with_bias["bs"]
-        d = x_with_bias["d"]
-        l = x_with_bias["l"]
-        pad = x_with_bias["pad"]
-        attention_output = x_with_bias["attention_output"]  # [ m, bs*l, d ]
-        attention_bias = x_with_bias["attention_bias"]  # [ d ]
-
-        # Re-enable torch grad to enable fused optimization.
-        with torch.enable_grad():
-            # Bias-dropout-add.
-            x = bias_dropout_add(
-                (
-                    attention_output,
-                    None if attention_bias is None else attention_bias.expand_as(attention_output),
-                ),
-                torch.zeros_like(attention_output),
-                prob,
-            )
-
-            # Permute chunks back to sequence dimension.
-            # 1. [ m, bs*l, d ]
-            # 2. [ m, bs, l, d ]
-            # 3. [ l, m, bs, d ]
-            # 4. [ m*l, bs, d ] == [ ns, bs, d ]
-            x = (
-                x.reshape(retro_chunk_length, bs, l, d)
-                .permute(2, 0, 1, 3)
-                .reshape(retro_chunk_length * l, bs, d)
-            )
-
-            # Prepend zeros for non-attending tokens.
-            x = torch.nn.functional.pad(x, (0, 0, 0, 0, pad, 0), "constant", 0)[
-                :ns
-            ]  # [ ns, bs, d ]
-
-            # Add residual. [ ns, bs, d ]
-            x = x + residual
-
-        # Output. [ ns, bs, d ]
-        return x
-
-    def forward(self, training: bool, fused: bool) -> partial:
-        """Retro decoder bias-dropout-add.
-
-        Args:
-            training (bool): If training, then apply dropout.
-            fused (bool): Fuse bias-dropout-add.
-
-        Returns:
-            The partial function for performing bias-dropout-add.
-        """
-        return partial(
-            self._forward,
-            retro_chunk_length=self.retro_chunk_length,
-            bias_dropout_add=get_bias_dropout_add(training, fused),
-        )
diff --git a/megatron/core/models/retro/decoder_spec.py b/megatron/core/models/retro/decoder_spec.py
deleted file mode 100644
index c872a4f77e1..00000000000
--- a/megatron/core/models/retro/decoder_spec.py
+++ /dev/null
@@ -1,202 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-
-"""Specs for Retro decoder."""
-
-import typing
-from typing import Optional
-
-from megatron.core.models.gpt.gpt_layer_specs import (
-    get_gpt_layer_local_spec,
-    get_gpt_layer_with_transformer_engine_spec,
-)
-from megatron.core.models.retro.config import RetroConfig
-from megatron.core.models.retro.decoder_attention import (
-    RetroDecoderBiasDropoutAdd,
-    RetroDecoderCrossAttention,
-)
-from megatron.core.models.retro.encoder_spec import get_retro_encoder_block_spec
-from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
-from megatron.core.transformer import ModuleSpec
-from megatron.core.transformer.attention import CrossAttentionSubmodules
-from megatron.core.transformer.dot_product_attention import DotProductAttention
-from megatron.core.transformer.transformer_block import (
-    TransformerBlockSubmodules,
-    get_num_layers_to_build,
-)
-from megatron.core.typed_torch import not_none
-
-try:
-    import apex  # pylint: disable=unused-import
-
-    from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
-
-    HAVE_APEX = True
-    LNImpl = FusedLayerNorm
-except ImportError:
-    import warnings
-
-    from megatron.core.transformer.torch_norm import WrappedTorchNorm
-
-    warnings.warn(f"Apex is not installed. Falling back to Torch Norm")
-    LNImpl = WrappedTorchNorm
-    HAVE_APEX = False
-
-try:
-    import transformer_engine as te  # pylint: disable=unused-import
-
-    from megatron.core.extensions.transformer_engine import (
-        TEColumnParallelLinear,
-        TEDotProductAttention,
-        TENorm,
-        TERowParallelLinear,
-    )
-
-    HAVE_TE = True
-except ImportError:
-    (TEColumnParallelLinear, TEDotProductAttention, TENorm, TERowParallelLinear) = (
-        None,
-        None,
-        None,
-        None,
-    )
-    HAVE_TE = False
-
-
-def get_retro_decoder_layer_te_spec(
-    encoder_block_spec: typing.Union[ModuleSpec, TransformerBlockSubmodules, None] = None
-) -> ModuleSpec:
-    """Retro decoder TE spec (uses Transformer Engine components).
-
-    A Retro decoder layer uses custom attention and bias-dropout-add operators
-    to perform chunked-cross attention. Additionally, the first Retro decoder
-    layer instantiates an entire encoder transformer block. As such, the decoder
-    cross attention module takes an optional encoder block spec, which is only
-    provided for the first Retro decoder layer.
-
-    Args:
-        encoder_block_spec (ModuleSpec): Retro encoder block spec, to be provided for
-            the first Retro decoder layer.
-
-    Returns:
-        A module spec with Transformer Engine modules.
-    """
-    spec = get_gpt_layer_with_transformer_engine_spec()
-    spec.submodules.pre_cross_attn_layernorm = TENorm
-    spec.submodules.cross_attention = ModuleSpec(
-        module=RetroDecoderCrossAttention,
-        params={"encoder_block_spec": encoder_block_spec},
-        submodules=CrossAttentionSubmodules(
-            linear_q=not_none(TEColumnParallelLinear),
-            linear_kv=not_none(TEColumnParallelLinear),
-            core_attention=not_none(TEDotProductAttention),
-            linear_proj=TERowParallelLinear,
-        ),
-    )
-    spec.submodules.cross_attn_bda = ModuleSpec(module=RetroDecoderBiasDropoutAdd)
-    return spec
-
-
-def get_retro_decoder_layer_local_spec(
-    encoder_block_spec: typing.Optional[ModuleSpec] = None,
-) -> ModuleSpec:
-    """Retro decoder local spec (uses Megatron-Core components).
-
-    A Retro decoder layer uses custom attention and bias-dropout-add operators
-    to perform chunked-cross attention. Additionally, the first Retro decoder
-    layer instantiates an entire encoder transformer block. As such, the decoder
-    cross attention module takes an optional encoder block spec, which is only
-    provided for the first Retro decoder layer.
-
-    Args:
-        encoder_block_spec (ModuleSpec): Retro encoder block spec, to be provided
-            for the first Retro decoder layer.
-
-    Returns:
-        A module spec with local modules.
-    """
-    spec = get_gpt_layer_local_spec()
-    spec.submodules.pre_cross_attn_layernorm = LNImpl
-    spec.submodules.cross_attention = ModuleSpec(
-        module=RetroDecoderCrossAttention,
-        params={"encoder_block_spec": encoder_block_spec},
-        submodules=CrossAttentionSubmodules(
-            linear_q=ColumnParallelLinear,
-            linear_kv=ColumnParallelLinear,
-            core_attention=DotProductAttention,
-            linear_proj=RowParallelLinear,
-        ),
-    )
-    spec.submodules.cross_attn_bda = ModuleSpec(module=RetroDecoderBiasDropoutAdd)
-    return spec
-
-
-def get_retro_decoder_block_spec(
-    config: RetroConfig,
-    use_transformer_engine: bool,
-    vp_stage: Optional[int] = None,
-    pp_rank: Optional[int] = None,
-) -> TransformerBlockSubmodules:
-    """Retro decoder block spec.
-
-    Retro decoder block implementation details:
-    - The retro decoder block consists of interleaved GPT layers
-        and customized Retro decoder layers.
-    - The Retro decoder layers are spaced three layers apart,
-        and start on layer 6 or 9 (depending on the total number of layers).
-    - The first decoder layer instantiates an encoder block,
-        and it therefore passes in an encoder_block_spec.
-
-    Args:
-        config (RetroConfig): Retro config.
-        use_transformer_engine (bool): If True, use Transformer Engine (instead of local modules.
-        vp_stage (Optional[int]): Virtual pipeline stage number.
-        pp_rank (Optional[int]): Pipeline parallel rank.
-
-    Returns:
-        Transformer block submodules for the given spec.
-    """
-
-    assert (
-        config.pipeline_model_parallel_size == 1
-    ), "retro does not currently support pipeline parallelism."
-
-    assert (
-        config.virtual_pipeline_model_parallel_size is None
-    ), "retro does not currently support virtual pipeline parallelism."
-
-    # Num layers.
-    num_layers = get_num_layers_to_build(config, vp_stage=vp_stage, pp_rank=pp_rank)
-
-    # Retro layer numbers.
-    retro_layer_start = 6 if num_layers <= 15 else 9
-    retro_layer_numbers = list(range(retro_layer_start, num_layers + 1, 3))
-
-    # Layer specs.
-    gpt_layer_spec = (
-        get_gpt_layer_with_transformer_engine_spec()
-        if use_transformer_engine
-        else get_gpt_layer_local_spec()
-    )
-    get_retro_decoder_layer_spec = (
-        get_retro_decoder_layer_te_spec
-        if use_transformer_engine
-        else get_retro_decoder_layer_local_spec
-    )
-    retro_layer_spec = get_retro_decoder_layer_spec()
-    retro_layer_spec_with_retriever = get_retro_decoder_layer_spec(
-        get_retro_encoder_block_spec(config, use_transformer_engine)
-    )
-
-    layer_specs = []
-    for layer_number in range(1, num_layers + 1):
-        if layer_number == retro_layer_numbers[0]:
-            layer_specs.append(retro_layer_spec_with_retriever)
-        elif layer_number in retro_layer_numbers:
-            layer_specs.append(retro_layer_spec)
-        else:
-            layer_specs.append(gpt_layer_spec)
-
-    # Block spec.
-    block_spec = TransformerBlockSubmodules(layer_specs=layer_specs)
-
-    return block_spec
diff --git a/megatron/core/models/retro/encoder_attention.py b/megatron/core/models/retro/encoder_attention.py
deleted file mode 100644
index 19fdae5b250..00000000000
--- a/megatron/core/models/retro/encoder_attention.py
+++ /dev/null
@@ -1,231 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Retro's cross attention modules for the encoder block."""
-
-from functools import partial
-from typing import Callable, List, Optional, Tuple, Type
-
-import torch
-from torch import Tensor
-
-from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
-from megatron.core.inference.contexts import BaseInferenceContext
-from megatron.core.models.retro.base_attention import BaseRetroCrossAttention
-from megatron.core.models.retro.config import RetroConfig
-from megatron.core.models.retro.utils import get_all_true_mask
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.utils import deprecate_inference_params
-
-
-class RetroEncoderCrossAttention(BaseRetroCrossAttention):
-    """Retro encoder's cross attention operator.
-
-    See this paper for more details: https://arxiv.org/abs/2112.04426.
-    Neighboring chunks are retrieved from the chunk database, encoded, and
-    used by the decoder layers for chunked cross attention.
-
-    Args:
-        config (RetroConfig): Retro config.
-        submodules (CrossAttentionSubmodules): Cross attention submodules.
-        layer_number (int): Layer number within transformer block.
-        attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
-    """
-
-    def forward(
-        self,
-        hidden_states: Tensor,
-        attention_mask: Tensor,
-        key_value_states: Tensor = None,
-        inference_context: BaseInferenceContext = None,
-        # rotary_pos_emb: Tensor = None, # unsupported for retro.
-        *,
-        inference_params: Optional[BaseInferenceContext] = None,
-    ) -> List[Tuple[Tensor, Optional[Tensor], Tensor]]:
-        """Cross attention for Retro encoder.
-
-        Notation:
-            ns : Sequence length.
-            bs : Batch size.
-            d  : Hidden size.
-            l  : Number of chunks per sample (i.e., seq_length/chunk_length).
-            k  : Number of neighbors.
-            r  : Number of retrieved tokens (neighbors + continuation).
-
-        Args:
-            hidden_states (Tensor): Transformer layer hidden states.
-            attention_mask (Tensor): Attention mask.
-            key_value_states (Tensor): Neighbor embeddings.
-            inference_context (BaseInferenceContext): Inference context.
-
-        Returns:
-            List of tuples, where each tuple is (attention_output, attention_bias, residual).
-        """
-
-        inference_context = deprecate_inference_params(inference_context, inference_params)
-
-        # Input shape. [ r, bs*l*k, d ]
-        ns, bs, d = hidden_states.shape
-
-        # Reshape sequence into neighboring chunks.
-        # - hidden_states:   [ r, bs*l*k, d ]
-        # - chunked_outputs: [ r, bs*l, k, d ]
-        chunked_outputs = hidden_states.reshape(
-            self.retro_retrieved_length, -1, self.retro_num_neighbors, d
-        )
-
-        # flash attn: [ b, h, sq, sk ]
-        # fused attn: [ b, 1, 1, sq ]
-        chunked_output_mask = get_all_true_mask(
-            size=(1, 1, chunked_outputs.shape[0], key_value_states.shape[0]),
-            device=chunked_outputs.device,
-        )
-
-        # Per-chunk attention.
-        attention_output_tuples = []
-        for k in range(self.retro_num_neighbors):
-
-            # Attend to current neighboring chunks.
-            # - chunked_output:   [ r, bs*l, d ]
-            # - key_value_states: [ m, bs*l, d ]
-            # - attention_output: [ r, bs*l, d ]
-            # - attention_bias:   [ d ]
-            chunked_output = chunked_outputs[:, :, k].contiguous()
-            attention_output, attention_bias = self.attn(
-                hidden_states=chunked_output,  # Q (neighbor embedding)
-                attention_mask=chunked_output_mask,
-                key_value_states=key_value_states,  # K, V (hidden act)
-            )
-
-            # Residual connection. [ r, bs*l, d ]
-            residual = chunked_output
-
-            # Collect tensors.
-            attention_output_tuples.append((attention_output, attention_bias, residual))
-
-        # Output. (List[Tuple[( [ r, bs*l, d ], [ d ] )]])
-        return attention_output_tuples
-
-
-class RetroEncoderBiasDropoutAdd(MegatronModule):
-    """Retro encoder's bias-dropout-add operator.
-
-    This operator applies bias-dropout-add individually on each neighboring
-    chunk that is retrieved from the chunk database.
-
-    Args:
-        config (RetroConfig): Retro config.
-    """
-
-    def __init__(self, config: RetroConfig):
-        super().__init__(config=config)
-        self.retro_num_neighbors = config.retro_num_neighbors
-
-    @classmethod
-    def _forward(
-        cls,
-        x_with_bias: List[Tuple[Tensor, Optional[Tensor], Tensor]],
-        residual: Tensor,
-        prob: float,
-        retro_num_neighbors: int,
-        bias_dropout_add: Callable,
-    ) -> Tensor:
-        """Per-chunk bias-dropout-add.
-
-        Args:
-            x_with_bias (dict): Attention output and bias tuple.
-            residual (Tensor): Transformer layer residual.
-            prob (float): Dropout probability.
-            retro_num_neighbors (int): Number of retrieved neighbor chunks (e.g., 2).
-            bias_dropout_add (Callable): Bias-dropout-add function.
-
-        Returns:
-            Output of bias-dropout-add.
-        """
-
-        # Re-enable torch grad to enable fused optimization.
-        with torch.enable_grad():
-
-            # Per-neighbor bias-dropout-add.
-            # - attention_output: [ r, bs*l, d ]
-            # - attention_bias:   [ d ]
-            # - residual:         [ r, bs*l, d ]
-            # - output:           [ r, bs*l, d ]
-            outputs = [
-                bias_dropout_add(
-                    (
-                        attention_output,
-                        None if attention_bias is None else attention_bias.expand_as(residual),
-                    ),
-                    residual,
-                    prob,
-                )
-                for attention_output, attention_bias, residual in x_with_bias
-            ]
-
-        # Concatenate outputs (to shape [r, k*bs*l, d]; see notation above).
-        r, _, d = outputs[0].shape
-        output = torch.stack(outputs, dim=1).reshape(r, -1, d)
-
-        # Output. [ r, k*bs*l, d ]
-        return output
-
-    def forward(self, training: bool, fused: bool) -> partial:
-        """Retro decoder bias-dropout-add.
-
-        Args:
-            training (bool): If training, then apply dropout.
-            fused (bool): Fuse bias-dropout-add.
-
-        Returns:
-            A partial function for performing bias-dropout-add.
-        """
-        return partial(
-            self._forward,
-            retro_num_neighbors=self.retro_num_neighbors,
-            bias_dropout_add=get_bias_dropout_add(training, fused),
-        )
-
-
-class RetroEncoderLayerNorm(MegatronModule):
-    """Retro encoder's layernorm operator.
-
-    This operator applies layernorm individually on each neighboring chunk that
-    is retrieved from the chunk database, and then concatenates the chunks into
-    a single tensor.
-
-    Args:
-        config (RetroConfig): Retro config.
-        submodules (Type): Layer norm class. (Named 'submodules' to fit external interface.)
-    """
-
-    def __init__(self, config: RetroConfig, submodules: Type, **kwargs: dict):
-        super().__init__(config=config)
-        norm_class = submodules
-        self.norm = norm_class(config=config, **kwargs)
-        self.retro_num_neighbors = config.retro_num_neighbors
-
-    def forward(self, input: Tensor) -> Tensor:
-        """Per-chunk layer norm.
-
-        Args:
-            input (Tensor): Input chunks, concatenated into a single tensor.
-
-        Returns:
-            Output of the layer norm.
-        """
-
-        # Input shape: [ r, k*bs*l, d ]. (see notation above in attention module)
-
-        # Split input into 'num_neighbors' tensors.
-        chunk_size = input.shape[1] // self.retro_num_neighbors
-        inputs = torch.split(input, chunk_size, dim=1)
-
-        # Norm.
-        outputs = [self.norm(inp.contiguous()) for inp in inputs]
-
-        # Concatenate layer norms (to shape [r, k*bs*l, d]; see notation above).
-        r, _, d = inputs[0].shape
-        output = torch.stack(outputs, dim=1).reshape(r, -1, d)
-
-        # Output. [ r, k*bs*l, d ]
-        return output
diff --git a/megatron/core/models/retro/encoder_spec.py b/megatron/core/models/retro/encoder_spec.py
deleted file mode 100644
index 0b5b94409a2..00000000000
--- a/megatron/core/models/retro/encoder_spec.py
+++ /dev/null
@@ -1,178 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Specs for Retro encoder."""
-
-from megatron.core.models.gpt.gpt_layer_specs import (
-    get_gpt_layer_local_spec,
-    get_gpt_layer_with_transformer_engine_spec,
-)
-from megatron.core.models.retro.config import RetroConfig
-from megatron.core.models.retro.encoder_attention import (
-    RetroEncoderBiasDropoutAdd,
-    RetroEncoderCrossAttention,
-    RetroEncoderLayerNorm,
-)
-from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
-from megatron.core.transformer import ModuleSpec
-from megatron.core.transformer.attention import CrossAttentionSubmodules
-from megatron.core.transformer.dot_product_attention import DotProductAttention
-from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.transformer.mlp import MLP, MLPSubmodules
-from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
-from megatron.core.typed_torch import not_none
-
-try:
-    import transformer_engine as te  # pylint: disable=unused-import
-
-    from megatron.core.extensions.transformer_engine import (
-        TEColumnParallelLinear,
-        TEDotProductAttention,
-        TENorm,
-        TERowParallelLinear,
-    )
-
-    HAVE_TE = True
-except ImportError:
-    (TEColumnParallelLinear, TEDotProductAttention, TENorm, TERowParallelLinear) = (
-        None,
-        None,
-        None,
-        None,
-    )
-    HAVE_TE = False
-
-try:
-    import apex  # pylint: disable=unused-import
-
-    from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
-
-    HAVE_APEX = True
-    LNImpl = FusedLayerNorm
-except ImportError:
-    import warnings
-
-    from megatron.core.transformer.torch_norm import WrappedTorchNorm
-
-    warnings.warn(f'Apex is not installed. Falling back to Torch Norm')
-    LNImpl = WrappedTorchNorm
-    HAVE_APEX = False
-
-
-def get_retro_encoder_layer_te_spec() -> ModuleSpec:
-    """Retro encoder TE spec (uses Transformer Engine components).
-
-    A Retro encoder layer uses custom attention, bias-dropout-add, and layernorm
-    operators to encode neighboring chunks that are retrieved from the chunk
-    database. Each operator is responsible for iterating the retrieved chunks
-    and processing them individually.
-
-    Returns:
-        A module spec if Transformer Engine modules.
-    """
-    spec = get_gpt_layer_with_transformer_engine_spec()
-    spec.submodules.pre_cross_attn_layernorm = TENorm
-    spec.submodules.cross_attention = ModuleSpec(
-        module=RetroEncoderCrossAttention,
-        params={"attn_mask_type": AttnMaskType.padding},
-        submodules=CrossAttentionSubmodules(
-            linear_q=not_none(TEColumnParallelLinear),
-            linear_kv=not_none(TEColumnParallelLinear),
-            core_attention=not_none(TEDotProductAttention),
-            linear_proj=TERowParallelLinear,
-        ),
-    )
-    spec.submodules.cross_attn_bda = ModuleSpec(module=RetroEncoderBiasDropoutAdd)
-    spec.submodules.pre_mlp_layernorm = ModuleSpec(module=RetroEncoderLayerNorm, submodules=TENorm)
-    spec.submodules.mlp = ModuleSpec(
-        module=MLP,
-        submodules=MLPSubmodules(linear_fc1=TEColumnParallelLinear, linear_fc2=TERowParallelLinear),
-    )
-    return spec
-
-
-def get_retro_encoder_layer_local_spec() -> ModuleSpec:
-    """Retro encoder local spec (uses Megatron-Core components).
-
-    A Retro encoder layer uses custom attention, bias-dropout-add, and layernorm
-    operators to encode neighboring chunks that are retrieved from the chunk
-    database. Each operator is responsible for iterating the retrieved chunks
-    and processing them individually.
-
-    Returns:
-        A module spec if local modules.
-    """
-    spec = get_gpt_layer_local_spec()
-    spec.submodules.pre_cross_attn_layernorm = LNImpl
-    spec.submodules.cross_attention = ModuleSpec(
-        module=RetroEncoderCrossAttention,
-        params={"attn_mask_type": AttnMaskType.padding},
-        submodules=CrossAttentionSubmodules(
-            linear_q=ColumnParallelLinear,
-            linear_kv=ColumnParallelLinear,
-            core_attention=DotProductAttention,
-            linear_proj=RowParallelLinear,
-        ),
-    )
-    spec.submodules.cross_attn_bda = ModuleSpec(module=RetroEncoderBiasDropoutAdd)
-    spec.submodules.pre_mlp_layernorm = ModuleSpec(module=RetroEncoderLayerNorm, submodules=LNImpl)
-    spec.submodules.mlp = ModuleSpec(
-        module=MLP,
-        submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
-    )
-    spec.submodules.sharded_state_dict_keys_map = {
-        'input_layernorm.': 'self_attention.linear_qkv.layer_norm_'
-    }  # pre_mlp_layernorm doesn't need remapping
-    return spec
-
-
-def get_retro_encoder_block_spec(
-    config: RetroConfig, use_transformer_engine: bool
-) -> TransformerBlockSubmodules:
-    """Retro encoder block spec.
-
-    The retro encoder block consists of one customized Retro encoder layer
-    (layer 1), and all of the following layers are standard GPT layers.
-
-    Args:
-      config (RetroConfig): Retro config.
-      use_transformer_engine (bool): If True, use Transformer Engine (instead of local modules).
-
-    Returns:
-        Transformer block submodules for the given spec.
-    """
-
-    # Num layers.
-    num_layers = config.retro_encoder_num_layers
-    retro_layer_numbers = [1]
-
-    # Layer specs.
-    gpt_layer_spec = (
-        get_gpt_layer_with_transformer_engine_spec()
-        if use_transformer_engine
-        else get_gpt_layer_local_spec()
-    )
-    get_retro_encoder_layer_spec = (
-        get_retro_encoder_layer_te_spec
-        if use_transformer_engine
-        else get_retro_encoder_layer_local_spec
-    )
-    retro_layer_spec = get_retro_encoder_layer_spec()
-    for spec in (gpt_layer_spec, retro_layer_spec):
-        spec.params["hidden_dropout"] = config.retro_encoder_hidden_dropout
-        spec.submodules.self_attention.params["attn_mask_type"] = AttnMaskType.padding
-        spec.submodules.self_attention.submodules.core_attention = ModuleSpec(
-            module=TEDotProductAttention if use_transformer_engine else DotProductAttention,
-            params={"attention_dropout": config.retro_encoder_attention_dropout},
-        )
-
-    layer_specs = []
-    for layer_number in range(1, num_layers + 1):
-        if layer_number in retro_layer_numbers:
-            layer_specs.append(retro_layer_spec)
-        else:
-            layer_specs.append(gpt_layer_spec)
-
-    # Block spec.
-    block_spec = TransformerBlockSubmodules(layer_specs=layer_specs)
-
-    return block_spec
diff --git a/megatron/core/models/retro/model.py b/megatron/core/models/retro/model.py
deleted file mode 100644
index 35fe7e8c878..00000000000
--- a/megatron/core/models/retro/model.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Retro Model."""
-from typing import Dict, Optional
-
-from torch import Tensor
-
-from megatron.core.dist_checkpointing.mapping import ShardedStateDict
-from megatron.core.inference.contexts import BaseInferenceContext
-from megatron.core.models.gpt import GPTModel
-from megatron.core.utils import deprecate_inference_params
-
-
-class RetroModel(GPTModel):
-    """Retro Model.
-
-    A Retro model mostly re-uses the GPTModel interface, with the only difference
-    being the embedding of the 'context' this is used by Retro for processing
-    neighbor tokens. This embedded context is then forwarded to the Transformer
-    Block.
-    """
-
-    def forward(
-        self,
-        input_ids: Tensor,
-        position_ids: Tensor,
-        attention_mask: Tensor,
-        context_input_ids: Tensor = None,
-        context_position_ids: Tensor = None,
-        context_mask: Tensor = None,
-        decoder_input: Tensor = None,
-        labels: Tensor = None,
-        inference_context: BaseInferenceContext = None,
-        *,
-        inference_params: Optional[BaseInferenceContext] = None,
-    ) -> Tensor:
-        """RetroModel forward method.
-
-        Foward input tokens & mask, along with neighbor tokens & mask, through
-        the Retro model..
-
-        Args:
-            input_ids (Tensor): Input token IDs.
-            position_ids (Tensor): Input position IDs.
-            attention_mask (Tensor): Input attention mask.
-            context_input_ids (Tensor): Context (i.e., neighbor) token IDs.
-            context_position_ids (Tensor): Context (i.e., neighbor) position IDs.
-            context_mask (Tensor): Context (i.e., neighbor) attention mask.
-            decoder_input (Tensor): When using pipeline parallelism, input_ids
-                and position_ids will only be used on the first stage, and for
-                all other stages decoder_input will be provided via communication
-                from the previous stage.
-            labels (Tensor): The labels of dimension [batch size, seq length].
-            inference_context (BaseInferenceContext): Inference context.
-
-        Returns:
-            Output tensor of forward pass.
-        """
-
-        # Argument shapes:
-        #   Notation:
-        #     ns : Sequence length.
-        #     bs : Batch size.
-        #     d  : Hidden size.
-        #     l  : Number of chunks per sample (i.e., seq_length/chunk_length).
-        #     k  : Number of neighbors.
-        #     r  : Number of retrieved tokens (neighbors + continuation).
-        # - input_ids:   [ bs, ns ]
-        # - context_ids: [ k*bs*l, r ]
-        # - context:     [ r, k*bs*l, d ]
-        # - output:      [ ns, bs, d ]
-
-        inference_context = deprecate_inference_params(inference_context, inference_params)
-
-        # Context embedding (e.g., for Retro neighbor tokens).
-        if context_input_ids is not None:
-            context = self.embedding(context_input_ids, context_position_ids)
-        else:
-            context = None
-
-        # Call GPTModel.forward, and pass in embedded context.
-        return super().forward(
-            input_ids=input_ids,
-            position_ids=position_ids,
-            attention_mask=attention_mask,
-            decoder_input=decoder_input,
-            labels=labels,
-            inference_context=inference_context,
-            extra_block_kwargs={"context": context, "context_mask": context_mask},
-        )
-
-    def sharded_state_dict(
-        self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
-    ) -> ShardedStateDict:
-        """Get sharded state dict.
-
-        Args:
-            prefix (str): Module name prefix.
-            sharded_offsets (tuple): Offsets of local shard within global tensor.
-            metadata (Optional[Dict]): Shard metadata.
-
-        Returns:
-            A <ShardedStateDict> ?
-        """
-        metadata = metadata or {}
-        metadata['non_homogeneous_layers'] = True
-        return super().sharded_state_dict(prefix, sharded_offsets, metadata)
diff --git a/megatron/core/models/retro/utils.py b/megatron/core/models/retro/utils.py
deleted file mode 100644
index 7d83c5d306f..00000000000
--- a/megatron/core/models/retro/utils.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-import os
-
-import torch
-
-
-def get_config_path(project_dir: str) -> str:
-    """Config copy stored within retro project dir."""
-    return os.path.join(project_dir, "config.json")
-
-
-def get_gpt_data_dir(project_dir: str) -> str:
-    """Get project-relative directory of GPT bin/idx datasets."""
-    return os.path.join(project_dir, "data")
-
-
-# ** Note ** : Retro's compatibility between cross attention and Flash/Fused
-#   Attention is currently a work in progress. We default to returning None for
-#   now.
-# def get_all_true_mask(size, device):
-#     return torch.full(size=size, fill_value=True, dtype=torch.bool, device=device)
-def get_all_true_mask(size, device):
-    return None
diff --git a/megatron/core/tokenizers/megatron_tokenizer.py b/megatron/core/tokenizers/megatron_tokenizer.py
index be01d0e554f..14b273e909b 100644
--- a/megatron/core/tokenizers/megatron_tokenizer.py
+++ b/megatron/core/tokenizers/megatron_tokenizer.py
@@ -15,7 +15,6 @@
         ("mamba", "MambaTokenizer"),
         ("bert", "BertTokenizer"),
         ("t5", "T5Tokenizer"),
-        ("retro", "RetroTokenizer"),
     ]
 )
 
@@ -104,7 +103,7 @@ def write_metadata(
             tokenizer_path (str): path to tokenizer model.
             tokenizer_library (str): tokenizer model library.
             model_type (str): type of the model to be used with tokenizer.
-                list of available model types: [gpt, bert, t5, mamba, retro, default].
+                list of available model types: [gpt, bert, t5, mamba, default].
                 `DefaultTokenizerText` will be used if model_type is not specified.
             tokenizer_class (MegatronTokenizerBase): pre-defined tokenizer class.
             chat_template (str): tokenizer chat template in jinja format.
diff --git a/megatron/core/tokenizers/text/models/__init__.py b/megatron/core/tokenizers/text/models/__init__.py
index 3610a8b98e4..d1788adb417 100644
--- a/megatron/core/tokenizers/text/models/__init__.py
+++ b/megatron/core/tokenizers/text/models/__init__.py
@@ -4,5 +4,4 @@
 from megatron.core.tokenizers.text.models.default_tokenizer import DefaultTokenizerText
 from megatron.core.tokenizers.text.models.gpt_tokenizer import GPTTokenizer
 from megatron.core.tokenizers.text.models.mamba_tokenizer import MambaTokenizer
-from megatron.core.tokenizers.text.models.retro_tokenizer import RetroTokenizer
 from megatron.core.tokenizers.text.models.t5_tokenizer import T5Tokenizer
diff --git a/megatron/core/tokenizers/text/models/retro_tokenizer.py b/megatron/core/tokenizers/text/models/retro_tokenizer.py
deleted file mode 100644
index a81af0c00f7..00000000000
--- a/megatron/core/tokenizers/text/models/retro_tokenizer.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-from megatron.core.tokenizers.text.text_tokenizer import MegatronTokenizerText
-
-
-class RetroTokenizer(MegatronTokenizerText):
-    """Base class for Megatron Retro tokenizer."""
-
-    def __init__(self, path: str = None, config: dict = None, **kwargs) -> None:
-        config['class_name'] = self.__class__.__name__
-        config['class_path'] = self.__class__.__module__
-        super().__init__(path, config, **kwargs)
diff --git a/megatron/legacy/data/orqa_wiki_dataset.py b/megatron/legacy/data/orqa_wiki_dataset.py
index 57bcc5891d1..033b2351cee 100644
--- a/megatron/legacy/data/orqa_wiki_dataset.py
+++ b/megatron/legacy/data/orqa_wiki_dataset.py
@@ -20,8 +20,7 @@ def get_open_retrieval_wiki_dataset():
     dataset = OpenRetrievalEvidenceDataset('2018 Wikipedia from DPR codebase',
                                            'evidence',
                                            args.evidence_data_path,
-                                           tokenizer,
-                                           args.retriever_seq_length)
+                                           tokenizer)
     return dataset
 
 
diff --git a/megatron/legacy/model/enums.py b/megatron/legacy/model/enums.py
index bc4e4aa29a0..bab179d1a04 100644
--- a/megatron/legacy/model/enums.py
+++ b/megatron/legacy/model/enums.py
@@ -5,9 +5,6 @@
 class LayerType(enum.Enum):
     encoder = 1
     decoder = 2
-    retro_encoder = 3
-    retro_decoder = 4
-    retro_decoder_with_retriever = 5
  
 class AttnType(enum.Enum):
     self_attn = 1
diff --git a/megatron/legacy/model/gpt_model.py b/megatron/legacy/model/gpt_model.py
index 3a2b831ebe7..66fd0979c46 100644
--- a/megatron/legacy/model/gpt_model.py
+++ b/megatron/legacy/model/gpt_model.py
@@ -76,9 +76,6 @@ def set_input_tensor(self, input_tensor):
         self.language_model.set_input_tensor(input_tensor)
 
     def forward(self, input_ids, position_ids, attention_mask,
-                retriever_input_ids=None,
-                retriever_position_ids=None,
-                retriever_attn_mask=None,
                 labels=None, tokentype_ids=None, inference_context=None, *, inference_params=None):
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
@@ -87,9 +84,6 @@ def forward(self, input_ids, position_ids, attention_mask,
             input_ids,
             position_ids,
             attention_mask,
-            retriever_input_ids=retriever_input_ids,
-            retriever_position_ids=retriever_position_ids,
-            retriever_attn_mask=retriever_attn_mask,
             inference_context=inference_context)
 
         if self.post_process:
diff --git a/megatron/legacy/model/language_model.py b/megatron/legacy/model/language_model.py
index b4e3c87c5e5..383230edb7f 100644
--- a/megatron/legacy/model/language_model.py
+++ b/megatron/legacy/model/language_model.py
@@ -360,7 +360,6 @@ def __init__(
         self.decoder_attn_mask_type = decoder_attn_mask_type
         self.add_pooler = add_pooler
         self.encoder_hidden_state = None
-        self.add_retriever = args.retro_add_retriever
         self.untie_embeddings_and_output_weights = args.untie_embeddings_and_output_weights
 
         # Embeddings.
@@ -399,9 +398,7 @@ def __init__(
         if self.add_encoder:
             self.encoder = ParallelTransformer(
                 config,
-                model_type=(
-                    args.model_type if not args.retro_add_retriever else ModelType.retro_decoder
-                ),
+                model_type=args.model_type,
                 self_attn_mask_type=self.encoder_attn_mask_type,
                 pre_process=self.pre_process,
                 post_process=self.post_process,
@@ -479,9 +476,6 @@ def forward(
         dec_input_ids=None,
         dec_position_ids=None,
         dec_attn_mask=None,
-        retriever_input_ids=None,
-        retriever_position_ids=None,
-        retriever_attn_mask=None,
         enc_dec_attn_mask=None,
         tokentype_ids=None,
         inference_context=None,
@@ -502,14 +496,6 @@ def forward(
         else:
             encoder_input = None
 
-        # Retriever embedding.
-        if self.add_retriever and self.pre_process:
-            retriever_input = self.embedding(
-                retriever_input_ids, retriever_position_ids, tokentype_ids=tokentype_ids
-            )
-        else:
-            retriever_input = None
-
         # Rotary positional embeddings
         rotary_pos_emb = None
         if self.use_rotary_position_embeddings:
@@ -524,8 +510,6 @@ def forward(
                 encoder_output = self.encoder(
                     encoder_input,
                     enc_attn_mask,
-                    retriever_input=retriever_input,
-                    retriever_attn_mask=retriever_attn_mask,
                     inference_context=inference_context,
                     rotary_pos_emb=rotary_pos_emb,
                 )
diff --git a/megatron/legacy/model/transformer.py b/megatron/legacy/model/transformer.py
index 2a662a55b16..ca3414eecdd 100644
--- a/megatron/legacy/model/transformer.py
+++ b/megatron/legacy/model/transformer.py
@@ -911,18 +911,6 @@ def __init__(self, config,
         # Normalize the attention output
         self.post_attention_norm = get_norm(config)
 
-        # Cross attention.
-        if self.layer_type in (LayerType.decoder,
-                               LayerType.retro_decoder,
-                               LayerType.retro_decoder_with_retriever,
-                               LayerType.retro_encoder):
-            self.inter_attention = ParallelAttention(
-                config,
-                layer_number,
-                attention_type=AttnType.cross_attn)
-            # Normalize the attention output.
-            self.post_inter_attention_norm = get_norm(config)
-
         # MLP
         if args.num_experts is not None:
             self.mlp = SwitchMLP(config)
@@ -936,25 +924,6 @@ def __init__(self, config,
         self.bias_dropout_add_exec_handler = \
                 nullcontext if use_nvfuser else torch.enable_grad
 
-        if args.retro_add_retriever:
-            self.retro_num_neighbors = args.retro_num_neighbors
-            self.retro_chunk_length = args.retro_chunk_length
-            self.retro_retrieved_length = \
-                args.retro_num_retrieved_chunks * args.retro_chunk_length
-
-        # Retriever (bi-directional transformer with cross attention)
-        if layer_type == LayerType.retro_decoder_with_retriever:
-            self.retriever = ParallelTransformer(
-                config=config,
-                model_type=ModelType.retro_encoder,
-                self_attn_mask_type=AttnMaskType.padding,
-                pre_process=True,
-                post_process=False,
-            )
-            self._retriever_key = 'retriever'
-        else:
-            self.retriever = None
-
     def default_decoder_cross_attention(self,
                                         encoder_output,
                                         enc_dec_attn_mask,
@@ -991,185 +960,8 @@ def default_decoder_cross_attention(self,
 
         return norm_input, norm_output
 
-    def retro_encoder_cross_attention(self,
-                                      retriever_output,
-                                      norm_input,
-                                      norm_output,
-                                      bias_dropout_add_func):
-        """Cross attention for Retro encoder.
-
-        Notation:
-            ns : Sequence length.
-            bs : Batch size.
-            d  : Hidden size.
-            l  : Number of chunks per sample (i.e., seq_length/chunk_length).
-            k  : Number of neighbors.
-            r  : Number of retrieved tokens (neighbors + continuation).
-        """
-
-        ns, bs, d = norm_output.shape # [r, bs * l * k, d]
-
-        # Divide sequence dimension into chunks.
-        chunked_outputs = norm_output.reshape(self.retro_retrieved_length,
-                                              -1,
-                                              self.retro_num_neighbors,
-                                              d)
-        chunked_outputs_before_norm = \
-            norm_input.reshape(self.retro_retrieved_length, -1,
-                               self.retro_num_neighbors, d) # [r, bs*l, k, d]
-
-        # Per-chunk attention.
-        norm_inputs = []
-        norm_outputs = []
-        for k in range(self.retro_num_neighbors):
-
-            # Attention.
-            chunked_output = chunked_outputs[:,:,k].contiguous()
-            attention_output, attention_bias = \
-                self.inter_attention(
-                    chunked_output, # Q (neighbor embedding)
-                    None,
-                    encoder_output=retriever_output) # K, V (hidden act)
-
-            # Residual connection.
-            if self.apply_residual_connection_post_norm:
-                residual = chunked_output
-            else:
-                residual = chunked_outputs_before_norm[:,:,k]
-
-            # Re-enable torch grad to enable fused optimization.
-            with torch.enable_grad():
-                norm_input = bias_dropout_add_func(
-                    attention_output,
-                    None if attention_bias is None else attention_bias.expand_as(residual),
-                    residual,
-                    self.hidden_dropout)
-                norm_inputs.append(norm_input)
-
-            # Layer norm.
-            norm_output = self.post_inter_attention_norm(norm_input)
-            norm_outputs.append(norm_output)
-
-        # Concatenate layer norms.
-        # norm_input : [r, k * bs * l, d]
-        # norm_output : [r, k * bs * l, d]
-        norm_input = torch.stack(norm_inputs, dim=1).reshape(ns, bs, d)
-        norm_output = torch.stack(norm_outputs, dim=1).reshape(ns, bs, d)
-
-        return norm_input, norm_output
-
-    def retro_decoder_cross_attention(self,
-                                      retriever_input,
-                                      retriever_output,
-                                      retriever_attn_mask,
-                                      norm_input,
-                                      norm_output,
-                                      inference_context,
-                                      bias_dropout_add_func,
-                                      *,
-                                      inference_params=None):
-        """Cross attention for Retro decoder.
-
-        Notation:
-            ns : Sequence length.
-            bs : Batch size.
-            d  : Hidden size.
-            l  : Number of chunks per sample (i.e., seq_length/chunk_length).
-            m  : Number of tokens per chunk.
-            k  : Number of neighbors.
-            r  : Number of retrieved tokens (neighbors + continuation).
-        """
-
-        inference_context = deprecate_inference_params(inference_context, inference_params)
-
-        ns, bs, d = norm_output.shape
-        l = int(np.ceil(ns / self.retro_chunk_length))
-
-        # Retrieve neighbors.
-        if self.layer_type == LayerType.retro_decoder_with_retriever:
-            first_ns = ns % self.retro_chunk_length
-            if first_ns > 0:
-                first_chunk, rest_chunk = \
-                    norm_output[:first_ns], norm_output[first_ns:]
-                first_chunk = torch.nn.functional.pad(
-                    first_chunk,
-                    (0, 0, 0, 0, 0, self.retro_chunk_length - first_ns),
-                    'constant',
-                    0)
-                chunked_output = \
-                    torch.cat((first_chunk, rest_chunk), dim=0) # [l * m, bs, d]
-            else:
-                chunked_output = norm_output # [l * m, bs, d]
-            chunked_output = chunked_output \
-                .reshape(l, self.retro_chunk_length, bs, d) \
-                .permute(1, 2, 0, 3) \
-                .reshape(self.retro_chunk_length, bs * l, d) \
-                .contiguous()
-
-            # Get Encoder Output
-            retriever_output = self.retriever(
-                hidden_states=retriever_input,
-                attention_mask=retriever_attn_mask,
-                retriever_output=chunked_output,
-                retriever_attn_mask=retriever_attn_mask,
-                inference_context=inference_context) # [r, k * bs * l , d]
-            retriever_output = retriever_output.reshape(
-                self.retro_retrieved_length * self.retro_num_neighbors, bs * l, d) # [r * k, bs * l, d]
-
-        # Chunks.
-        pad = (ns - 1) % self.retro_chunk_length
-        attending_chunks = norm_output[pad:]
-        padded_chunks = torch.nn.functional.pad(
-            attending_chunks,
-            (0, 0, 0, 0, 0, self.retro_chunk_length - 1),
-            'constant', 0)
-        padded_chunked_output = padded_chunks \
-            .reshape(l, self.retro_chunk_length, bs, d) \
-            .permute(1, 2, 0, 3)
-        padded_chunked_output = padded_chunked_output.reshape(
-            self.retro_chunk_length, bs * l, d).contiguous()
-
-        # Encoder output.
-        attention_output, attention_bias = \
-            self.inter_attention(padded_chunked_output,
-                                 None,
-                                 encoder_output=retriever_output)
-
-        # Residual connection.
-        if self.apply_residual_connection_post_norm:
-            residual = norm_output
-        else:
-            residual = norm_input
-
-        # Re-enable torch grad to enable fused optimization.
-        with torch.enable_grad():
-            norm_input = bias_dropout_add_func(
-                attention_output,
-                None if attention_bias is None else attention_bias.expand_as(attention_output),
-                torch.zeros_like(attention_output),
-                self.hidden_dropout)
-            norm_input = norm_input \
-                .reshape(self.retro_chunk_length, bs, l, d) \
-                .permute(2, 0, 1, 3) # [l, m, bs, d]
-            norm_input = norm_input.reshape(self.retro_chunk_length * l, bs, d)
-            norm_input = torch.nn.functional.pad(
-                norm_input,
-                (0, 0, 0, 0, pad, 0),
-                'constant', 0)[:ns] # [ns, b, d]
-            # TODO: better redesign with inference param
-            args = get_args()
-            norm_input = args.retro_attention_gate * norm_input + residual
-
-        # Layer norm post the decoder attention
-        norm_output = self.post_inter_attention_norm(norm_input)
-
-        return retriever_output, norm_input, norm_output
-
     def forward(self, hidden_states, attention_mask,
                 encoder_output=None, enc_dec_attn_mask=None,
-                retriever_input=None,
-                retriever_output=None,
-                retriever_attn_mask=None,
                 inference_context=None,
                 rotary_pos_emb=None,
                 *,
@@ -1177,15 +969,6 @@ def forward(self, hidden_states, attention_mask,
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
-        # Update the params in case the retro param changes during inference
-        # TODO: better redesign with inference param
-        args = get_args()
-        if args.retro_add_retriever:
-            self.retro_num_neighbors = args.retro_num_neighbors
-            self.retro_chunk_length = args.retro_chunk_length
-            self.retro_retrieved_length = \
-                args.retro_num_retrieved_chunks * args.retro_chunk_length
-
         # hidden_states: [s, b, h]
 
         # Layer norm at the beginning of the transformer layer.
@@ -1246,24 +1029,6 @@ def forward(self, hidden_states, attention_mask,
                     norm_input,
                     norm_output,
                     bias_dropout_add_func)
-        elif self.layer_type == LayerType.retro_encoder:
-            norm_input, norm_output = \
-                self.retro_encoder_cross_attention(
-                    retriever_output,
-                    norm_input,
-                    norm_output,
-                    bias_dropout_add_func)
-        elif self.layer_type in (LayerType.retro_decoder,
-                                 LayerType.retro_decoder_with_retriever):
-            retriever_output, norm_input, norm_output = \
-                self.retro_decoder_cross_attention(
-                    retriever_input,
-                    retriever_output,
-                    retriever_attn_mask,
-                    norm_input,
-                    norm_output,
-                    inference_context,
-                    bias_dropout_add_func)
         else:
             raise Exception("Unsupported layer type, '%s'." %
                             self.layer_type.name)
@@ -1305,10 +1070,7 @@ def forward(self, hidden_states, attention_mask,
                                               training=self.training)
             output = residual + self.drop_path(out)
 
-        if self.layer_type == LayerType.retro_decoder_with_retriever:
-            return output, retriever_output
-        else:
-            return output
+        return output
 
 
 class NoopTransformerLayer(MegatronModule):
@@ -1339,9 +1101,7 @@ def forward(self, hidden_states, attention_mask,
 
 def _get_num_layers(args, model_type, is_decoder=False):
     """Compute the number of transformer layers resident on the current rank."""
-    if model_type == ModelType.retro_encoder:
-        num_layers = args.retro_encoder_layers
-    elif mpu.get_pipeline_model_parallel_world_size() > 1:
+    if mpu.get_pipeline_model_parallel_world_size() > 1:
         assert args.num_layers == args.encoder_num_layers
         assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \
             'num_layers must be divisible by transformer_pipeline_model_parallel_size'
@@ -1364,22 +1124,6 @@ def _get_num_layers(args, model_type, is_decoder=False):
     return num_layers
 
 
-def _get_layer_type(model_type, default_layer_type, retro_layer_numbers,
-                    layer_number):
-    args = get_args()
-    if args.retro_add_retriever and layer_number in retro_layer_numbers:
-        if model_type == ModelType.retro_decoder:
-            return LayerType.retro_decoder_with_retriever \
-                if layer_number == retro_layer_numbers[0] \
-                   else LayerType.retro_decoder
-        elif model_type == ModelType.retro_encoder:
-            return LayerType.retro_encoder
-        else:
-            raise Exception("Unsupported model type, '%s'." % model_type)
-    else:
-        return default_layer_type
-
-
 class ParallelTransformer(MegatronModule):
     """Transformer class."""
 
@@ -1403,7 +1147,6 @@ def __init__(self, config,
         self.input_tensor = None
         self.drop_path_rate = drop_path_rate
         self.transformer_impl = args.transformer_impl
-        self.retro_add_retriever = args.retro_add_retriever
 
         # Store activation checkpoiting flag.
         self.recompute_granularity = config.recompute_granularity
@@ -1469,29 +1212,12 @@ def __init__(self, config,
             rate.item() for rate in
             torch.linspace(0, self.drop_path_rate, config.num_layers)]
 
-        self.retro_layer_numbers = None
-        if model_type == ModelType.retro_decoder:
-            retro_layer_start = 6 if config.num_layers <= 15 else 9
-            self.retro_layer_numbers = \
-                np.arange(retro_layer_start, args.num_layers + 1, 3).tolist()
-        if model_type == ModelType.retro_encoder:
-            self.retro_layer_numbers = [1]
-
-        # Transformer layers.
-        if args.retro_add_retriever:
-            assert self.recompute_granularity != 'full', \
-                "Full recompute not supported for Retro."
-            assert args.transformer_impl == 'local', \
-                "Transformer engine does not support Retro layers."
         def build_layer(layer_number):
             if args.transformer_impl == 'local':
-                current_layer_type = _get_layer_type(
-                    model_type, layer_type, self.retro_layer_numbers,
-                    layer_number)
                 return ParallelTransformerLayer(
                     config,
                     layer_number,
-                    layer_type=current_layer_type,
+                    layer_type=layer_type,
                     self_attn_mask_type=self_attn_mask_type,
                     drop_path_rate=self.drop_path_rates[layer_number - 1])
             else:
@@ -1575,17 +1301,6 @@ def build_layer(layer_number):
             self.layers = torch.nn.ModuleList(
                 [build_layer(i + 1 + offset) for i in range(self.num_layers)])
 
-            # Update dropout rate for Retro encoder.
-            if model_type == ModelType.retro_encoder:
-                for layer in self.layers:
-                    if layer.self_attention.use_flash_attn:
-                        layer.self_attention.core_attention_flash.dropout_p = \
-                            torch.nn.Dropout(args.retro_encoder_attention_dropout)
-                    else:
-                        layer.self_attention.core_attention.attention_dropout.p =\
-                            args.retro_encoder_attention_dropout
-                    layer.hidden_dropout = args.retro_encoder_hidden_dropout
-
         if self.post_process and self.post_norm:
             # Final layer norm before output.
             self.final_norm = get_norm(config)
@@ -1684,9 +1399,6 @@ def set_input_tensor(self, input_tensor):
 
     def forward(self, hidden_states, attention_mask,
                 encoder_output=None, enc_dec_attn_mask=None,
-                retriever_input=None,
-                retriever_output=None,
-                retriever_attn_mask=None,
                 inference_context=None,
                 rotary_pos_emb=None,
                 *,
@@ -1771,9 +1483,6 @@ def forward(self, hidden_states, attention_mask,
                             forward_kwargs['rotary_pos_emb'] = rotary_pos_emb
                     else:
                         forward_kwargs['rotary_pos_emb'] = rotary_pos_emb
-                        forward_kwargs['retriever_input'] = retriever_input
-                        forward_kwargs['retriever_output'] = retriever_output
-                        forward_kwargs['retriever_attn_mask'] = retriever_attn_mask
 
                     for index in range(self.num_layers):
                         layer = self._get_layer(index)
@@ -1783,14 +1492,6 @@ def forward(self, hidden_states, attention_mask,
                             attention_mask,
                             **forward_kwargs)
 
-                        # First Retro decoder layer returns both hidden_states
-                        # and retriever_output. Make retriever_output available
-                        # to subsequence Retro layers.
-                        if isinstance(hidden_states, tuple):
-                            assert len(hidden_states) == 2
-                            hidden_states, retriever_output = hidden_states
-                            forward_kwargs["retriever_output"] = retriever_output
-
                 # Skip counter update for eval and activation checkpointing
                 if torch.is_grad_enabled() and self.training:
                     self.microbatch_count += 1
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 8827a7bdf55..007a9842610 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -15,10 +15,6 @@
 from packaging.version import Version as PkgVersion
 
 from megatron.core.dist_checkpointing.validation import StrictHandling
-from megatron.core.models.retro.utils import (
-    get_config_path as get_retro_config_path,
-    get_gpt_data_dir as get_retro_data_dir,
-)
 from megatron.core.rerun_state_machine import RerunStateMachine
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
 from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
@@ -77,7 +73,6 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
     parser = _add_workload_inspector_server_args(parser)
     parser = _add_inference_args(parser)
     parser = _add_transformer_engine_args(parser)
-    parser = _add_retro_args(parser)
     parser = _add_experimental_args(parser)
     parser = _add_one_logger_args(parser)
     parser = _add_inprocess_restart_args(parser)
@@ -201,81 +196,6 @@ def validate_model_config_args_from_heterogeneous_config(args):
             f"Arguments differ from heterogeneous config: {incompatible_args_str}"
         )
 
-
-def load_retro_config(retro_project_dir):
-    '''Load Retro's config.json.'''
-
-    # Retro config path.
-    retro_config_path = get_retro_config_path(retro_project_dir)
-    assert os.path.exists(retro_config_path), \
-        "Retro project dir missing config.json."
-
-    # Load retro config.
-    with open(retro_config_path) as f:
-        retro_config = types.SimpleNamespace(**json.load(f))
-
-    return retro_config
-
-
-def load_retro_args(args):
-    """Load predefined args from Retro config (if applicable).
-
-    When using Retro (or GPT for comparison purposes), data arguments are
-    overridden by the saved config.json within the Retro project directory. This
-    is to ensure that the data used for pretraining is consistent with the data
-    that was preprocessed using the Retro preprocessing pipeline (see
-    `tools/retro/preprocess_data.py`).
-    """
-
-    # Return if no project directory is specified.
-    if args.retro_project_dir is None:
-        return
-
-    # Load retro config.
-    retro_config = load_retro_config(args.retro_project_dir)
-
-    # Retro data path is relative to project dir (via hard or soft links).
-    data_dir = get_retro_data_dir(args.retro_project_dir)
-    data_path = list(retro_config.retro_gpt_data_path)
-    if len(data_path) % 2 == 0:
-        for i in range(len(data_path) - 1, -1, -2):
-            data_path[i] = os.path.join(data_dir, data_path[i])
-    else:
-        assert len(data_path) == 1
-        data_path[0] = os.path.join(data_dir, data_path[0])
-
-    # Update args.
-    args.data_cache_path = retro_config.retro_gpt_data_cache_path
-    args.data_path = data_path if args.data_path is None else args.data_path
-    args.eval_interval = retro_config.retro_gpt_eval_interval
-    args.eval_iters = retro_config.retro_gpt_eval_iters
-    args.global_batch_size = retro_config.retro_gpt_global_batch_size
-    args.max_position_embeddings = retro_config.retro_gpt_seq_length
-    args.merge_file = os.path.join(
-        args.retro_project_dir,
-        retro_config.retro_gpt_merge_file,
-    ) if retro_config.retro_gpt_merge_file is not None else None
-    args.seed = retro_config.retro_gpt_seed
-    args.seq_length = retro_config.retro_gpt_seq_length
-    args.tokenizer_model = os.path.join(
-        args.retro_project_dir,
-        retro_config.retro_gpt_tokenizer_model,
-    ) if retro_config.retro_gpt_tokenizer_model is not None else None
-    args.tokenizer_type = retro_config.retro_gpt_tokenizer_type
-    args.train_samples = retro_config.retro_gpt_train_samples
-    args.vocab_file = os.path.join(
-        args.retro_project_dir,
-        retro_config.retro_gpt_vocab_file,
-    ) if retro_config.retro_gpt_vocab_file is not None else None
-
-    # Retro-specific args.
-    args.retro_block_size = retro_config.retro_block_size
-    args.retro_chunk_length = retro_config.retro_gpt_chunk_length
-    args.retro_neighbor_dirs = retro_config.retro_neighbor_dirs
-    args.retro_split_preprocessing = retro_config.retro_gpt_split
-    args.retro_bert_tokenizer_type = retro_config.retro_bert_tokenizer_type
-    args.retro_bert_vocab_file = retro_config.retro_bert_vocab_file
-
 def _eval_pattern(pattern):
     """ Validate and evaluate a string containing a Python list expression """
     assert isinstance(pattern, str)
@@ -386,9 +306,6 @@ def validate_args(args, defaults={}):
     # validate model config args from heterogeneous config (if provided).
     validate_model_config_args_from_heterogeneous_config(args)
 
-    # Load saved args from Retro (if applicable).
-    load_retro_args(args)
-
     # Set args.use_dist_ckpt from args.ckpt_format.
     if args.use_legacy_models:
         assert args.ckpt_format == "torch", \
@@ -1058,21 +975,6 @@ def validate_args(args, defaults={}):
         assert is_te_min_version("2.9.0"), \
             '--log-max-attention-logit is only supported with TE >= 2.9.0.'
 
-    # Retro checks.
-    if args.retro_add_retriever:
-
-        # Train samples should be auto-loaded.
-        assert args.train_samples is not None, \
-            "args.train_samples should be auto-loaded from the retro config."
-
-        # Sequence parallelism unsupported.
-        assert not args.sequence_parallel, \
-            "retro currently does not support sequence parallelism."
-
-        # Pipeline parallelism unsupported.
-        assert args.pipeline_model_parallel_size == 1, \
-            "retro currently does not support pipeline parallelism."
-
     if args.decoupled_lr is not None or args.decoupled_min_lr is not None:
         assert not args.use_legacy_models, \
             '--decoupled-lr and --decoupled-min-lr is not supported in legacy models.'
@@ -1707,54 +1609,6 @@ def _add_inference_args(parser):
     return parser
 
 
-def _add_retro_args(parser):
-    group = parser.add_argument_group(title='retro')
-
-    group.add_argument('--retro-project-dir', default=None,
-                       help='Retro project directory, which contains the '
-                       'preprocessed data for pretraining. This directory '
-                       'is built during preprocessing (see '
-                       'tools/retro/README.md), and contains subdirectories '
-                       'for the chunk database and pretraining neighbors.')
-    group.add_argument('--retro-add-retriever',
-                       action='store_true', default=False,
-                       help='Add a retriever to the transformer, for use in '
-                       'pretraining a Retro model.')
-    group.add_argument('--retro-cyclic-train-iters', type=int, default=None,
-                       help='Set number of training iterations for cyclic '
-                       'Retro training.')
-    group.add_argument('--retro-encoder-layers', type=int, default=2,
-                       help='Number of layers to use for the retrieval '
-                       'encoder.')
-    group.add_argument('--retro-encoder-hidden-dropout',
-                       type=float, default=0.1, help='Hidden dropout for '
-                       'retrieval encoder.')
-    group.add_argument('--retro-encoder-attention-dropout',
-                       type=float, default=0.1, help='Attention dropout for '
-                       'retrieval encoder.')
-    group.add_argument("--retro-num-neighbors", type=int, default=2,
-                       help='Number of neighbors to retrieve during '
-                       'pretraining.')
-    group.add_argument("--retro-num-retrieved-chunks", type=int, default=2,
-                       help='Number of chunks to retrieve from the retrieval '
-                       'database.')
-    group.add_argument("--retro-attention-gate", type=float, default=1,
-                       help="Gated cross attention.")
-    group.add_argument("--retro-no-verify-neighbor-count", action="store_false",
-                       dest="retro_verify_neighbor_count",
-                       help="Skip verifying that len(GPT dataset) == len(saved "
-                       "neighbors).")
-
-    # Enforce argument naming convention.
-    for action in group._group_actions:
-        prefix = action.dest.split("_")[0]
-        assert prefix == "retro", \
-            "Retro args must be prefixed with '--retro-*', for consistent " \
-            "styling. Please fix '%s'." % ", ".join(action.option_strings)
-
-    return parser
-
-
 def _add_network_size_args(parser):
     group = parser.add_argument_group(title='network size')
 
@@ -1891,7 +1745,6 @@ def _add_network_size_args(parser):
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
     return parser
 
-
 def _add_straggler_detector_args(parser):
     from megatron.training.resilience_config import StragglerDetectionConfig
 
@@ -2895,9 +2748,6 @@ def _add_data_args(parser):
                        'This should be exclusive of --seq-length')
     group.add_argument('--decoder-seq-length', type=int, default=None,
                        help="Maximum decoder sequence length to process.")
-    group.add_argument('--retriever-seq-length', type=int, default=256,
-                       help='Maximum sequence length for the biencoder model '
-                       'for retriever')
     group.add_argument('--sample-rate', type=float, default=1.0,
                        help='sample rate for training data. Supposed to be 0 '
                             ' < sample_rate < 1')
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 4a0218c7106..a3d307f1e30 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1766,7 +1766,6 @@ def load_model_state_dict(module, state_dict, strict: bool):
                 load_return = module.load_state_dict(state_dict, strict=False)
                 print(f"load_return: {load_return}")
     # Model.
-    strict = False if args.retro_add_retriever else strict
     if not skip_load_to_model_and_opt:
         if len(ddp_model) == 1:
             load_model_state_dict(ddp_model[0], state_dict['model'], strict)
diff --git a/megatron/training/one_logger_utils.py b/megatron/training/one_logger_utils.py
index ea41ba18af0..fd18c569d1d 100644
--- a/megatron/training/one_logger_utils.py
+++ b/megatron/training/one_logger_utils.py
@@ -297,8 +297,7 @@ def on_pretrain_start():
                 'one_logger_utils_version': _one_logger_utils_version,
             })
 
-def track_config_flags(train_iters, skip_train, do_train, do_valid, do_test,
-                           dataloader_type, retro_project_dir, retro_cyclic_train_iters):
+def track_config_flags(train_iters, skip_train, do_train, do_valid, do_test, dataloader_type):
     """Track flags about train/validation/test enablement
 
     Args:
@@ -308,16 +307,10 @@ def track_config_flags(train_iters, skip_train, do_train, do_valid, do_test,
         do_valid (bool): flags to do validation
         do_test (bool): flags to do test
         dataloader_type (str): dataloader type
-        retro_project_dir (str): Retro project directory
-        retro_cyclic_train_iters (int): iteration number for cyclic retro training
     """
     one_logger = get_one_logger()
     if one_logger:
         with one_logger.get_context_manager():
-            # Update train_iters for cyclic loader
-            if dataloader_type == 'cyclic' and retro_project_dir:
-                assert retro_cyclic_train_iters is not None
-                train_iters = retro_cyclic_train_iters
             # Track if training is enabled. Can only be done once args.do_train is assigned after dataloader is built.
             train_enabled = train_iters and (not skip_train) and do_train and train_iters > 0
             one_logger.log_metrics({
diff --git a/megatron/training/training.py b/megatron/training/training.py
index fbc267fba82..500d30b9e73 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -1008,8 +1008,6 @@ def pretrain(
         args.do_valid,
         args.do_test,
         args.dataloader_type,
-        args.retro_project_dir,
-        args.retro_cyclic_train_iters,
     )
 
     # Print setup timing.
@@ -1027,11 +1025,6 @@ def pretrain(
     if not args.skip_train:
         print_rank_0('training ...')
 
-        if args.dataloader_type == 'cyclic' and args.retro_project_dir:
-            assert args.retro_cyclic_train_iters is not None
-            args.train_iters = args.retro_cyclic_train_iters
-            print_rank_0("retro cyclic train iters : %d" % args.train_iters)
-
         iteration = 0
         if args.do_train and args.train_iters > 0:
             iteration, num_floating_point_operations_so_far = train(
diff --git a/megatron/training/yaml_arguments.py b/megatron/training/yaml_arguments.py
index dfb5a8166ed..70ccac4402c 100644
--- a/megatron/training/yaml_arguments.py
+++ b/megatron/training/yaml_arguments.py
@@ -318,23 +318,6 @@ def validate_yaml(args, defaults={}):
             raise RuntimeError(
                 "Using async gradient all reduce requires setting the environment "
                 "variable CUDA_DEVICE_MAX_CONNECTIONS to 1")
-
-    # Retro checks.
-    if getattr(args, 'retro_add_retriever', False):
-        raise Exception("Retro untested for yaml args. See arguments.py.")
-
-        # Sequence parallelism unsupported.
-        assert not args.sequence_parallel, \
-            "retro currently does not support sequence parallelism."
-
-        # Pipeline parallelism unsupported.
-        assert args.pipeline_model_parallel_size == 1, \
-            "retro currently does not support pipeline parallelism."
-
-    #TODO: Retro args loading not tested
-    # Load retro args (used by both Retro & GPT).
-    if getattr(args, 'retro_project_dir', None) is not None:
-        raise Exception("Retro untested for yaml args. See arguments.py.")
     
     # MoE Spec check
     if args.language_model.num_moe_experts is not None:
diff --git a/pretrain_retro.py b/pretrain_retro.py
deleted file mode 100644
index a0a8fa4cae5..00000000000
--- a/pretrain_retro.py
+++ /dev/null
@@ -1,258 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Pretrain Retro."""
-
-from functools import partial
-import torch
-from importlib import import_module
-
-from megatron.training import get_args
-from megatron.training import get_tokenizer
-from megatron.training import get_timers
-from megatron.training import print_rank_0
-from megatron.training.arguments import core_transformer_config_from_args
-from megatron.core import tensor_parallel
-from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
-from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
-from megatron.core.datasets.utils import get_blend_from_list
-from megatron.core.datasets.retro.query.retro_dataset import get_retro_datasets
-from megatron.core.datasets.retro.query.multi_split_gpt_dataset import MultiSplitGPTDataset, MultiSplitGPTDatasetConfig
-from megatron.core.enums import ModelType
-from megatron.core.models.retro import get_retro_decoder_block_spec, RetroConfig, RetroModel
-from megatron.core.models.retro.utils import get_all_true_mask
-from megatron.training import pretrain
-from megatron.training.utils import get_ltor_masks_and_position_ids
-from pretrain_gpt import (
-    is_dataset_built_on_rank,
-    loss_func,
-    model_provider as default_model_provider,
-    train_valid_test_datasets_provider as gpt_train_valid_test_datasets_provider,
-)
-
-
-def get_retro_config():
-    return core_transformer_config_from_args(get_args(), RetroConfig)
-
-
-def core_model_provider(pre_process=True, post_process=True):
-    """Build the model using Megatron-Core."""
-
-    args = get_args()
-    config = get_retro_config()
-
-    # NOTE: Experimental customization feature
-    if args.spec is not None:
-        block_spec = import_module(args.spec)()
-    else:
-        block_spec = get_retro_decoder_block_spec(config, use_transformer_engine=True)
-
-    print_rank_0('building Retro model ...')
-    model = RetroModel(
-        config=config,
-        transformer_layer_spec=block_spec,
-        vocab_size=args.padded_vocab_size,
-        max_sequence_length=args.max_position_embeddings,
-        pre_process=pre_process,
-        post_process=post_process,
-        fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
-        parallel_output=True,
-        share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
-        position_embedding_type=args.position_embedding_type,
-        rotary_percent=args.rotary_percent
-    )
-    return model
-
-
-def model_provider(pre_process=True, post_process=True):
-    """Build the model.
-
-    Select between two different model classes:
-      1. Default model (uses megatron.legacy.models/gpt_model.py).
-      2. Core model (uses megatron/core/models/retro/model.py).
-    """
-
-    args = get_args()
-    if not args.use_legacy_models and args.retro_add_retriever:
-        provider = core_model_provider
-    else:
-        provider = default_model_provider
-    model = provider(pre_process=pre_process, post_process=post_process)
-    return model
-
-
-def get_batch(data_iterator):
-    """Generate a batch"""
-
-    args = get_args()
-
-    if args.legacy_tokenizer:
-        tokenizer = get_tokenizer()
-    else:
-        tokenizer = build_tokenizer(args)
-        
-    config = get_retro_config()
-
-    # Items and their type.
-    keys = ['text']
-    if args.retro_add_retriever:
-        keys.append('neighbor_tokens')
-    datatype = torch.int64
-
-    # Broadcast data.
-    if data_iterator is not None:
-        data = next(data_iterator)
-    else:
-        data = None
-
-    data_b = tensor_parallel.broadcast_data(keys, data, datatype)
-
-    # Unpack.
-    tokens_ = data_b['text'].long()
-    labels = tokens_[:, 1:].contiguous()
-    tokens = tokens_[:, :-1].contiguous()
-
-    # Get the masks and postition ids.
-    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
-        tokens,
-        tokenizer.eod,
-        args.reset_position_ids,
-        args.reset_attention_mask,
-        args.eod_mask_loss)
-
-    if args.retro_add_retriever:
-        # note: [bs * l * k, r]
-        # note: 2x == neighbor, continuation
-        neighbor_tokens = data_b['neighbor_tokens'] \
-            .view(-1, config.retro_retrieved_length).long()
-        _, _, neighbor_position_ids = get_ltor_masks_and_position_ids(
-            neighbor_tokens,
-            tokenizer.eod,
-            args.reset_position_ids,
-            args.reset_attention_mask,
-            args.eod_mask_loss)
-        neighbor_attention_mask = get_all_true_mask(
-            (1, 1, config.retro_retrieved_length, config.retro_retrieved_length),
-            neighbor_tokens.device)
-        return tokens, labels, loss_mask, attention_mask, position_ids, \
-               neighbor_tokens, neighbor_attention_mask, neighbor_position_ids
-
-    else:
-        return tokens, labels, loss_mask, attention_mask, position_ids
-
-
-def forward_step(data_iterator, model):
-    """Forward step."""
-    args = get_args()
-    timers = get_timers()
-
-    # Get the batch.
-    timers('batch-generator').start()
-    if args.retro_add_retriever:
-        tokens, labels, loss_mask, attention_mask, position_ids, \
-            neighbor_tokens, neighbor_attention_mask, neighbor_position_ids = \
-                get_batch(data_iterator)
-    else:
-        tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
-            data_iterator)
-        neighbor_tokens, neighbor_attention_mask, neighbor_position_ids = \
-            None, None, None
-    timers('batch-generator').stop()
-
-    # Model call.
-    if args.use_legacy_models:
-        forward_kwargs = {
-            "retriever_input_ids" : neighbor_tokens,
-            "retriever_position_ids" : neighbor_position_ids,
-            "retriever_attn_mask" : neighbor_attention_mask,
-        }
-    else:
-        if args.retro_add_retriever:
-            forward_kwargs = {
-                "context_input_ids" : neighbor_tokens,
-                "context_position_ids" : neighbor_position_ids,
-                "context_mask" : neighbor_attention_mask,
-            }
-        else:
-            forward_kwargs = {}
- 
-    output_tensor = model(tokens, position_ids, attention_mask,
-                          labels=labels, **forward_kwargs)
-
-    return output_tensor, partial(loss_func, loss_mask)
-
-
-def train_valid_test_datasets_provider(train_valid_test_num_samples):
-    """Build train, valid, and test datasets."""
-    args = get_args()
-
-    if args.legacy_tokenizer:
-        tokenizer = get_tokenizer()
-    else:
-        tokenizer = build_tokenizer(args)
-
-    # Dataset config.
-    retro_config = get_retro_config()
-    data_config = MultiSplitGPTDatasetConfig(
-        random_seed=args.seed,
-        sequence_length=args.seq_length,
-        blend=get_blend_from_list(args.data_path),
-        blend_per_split=[
-            get_blend_from_list(args.train_data_path),
-            get_blend_from_list(args.valid_data_path),
-            get_blend_from_list(args.test_data_path)
-        ],
-        split=args.split,
-        split_preprocessing=retro_config.retro_split_preprocessing,
-        path_to_cache=args.data_cache_path,
-        return_document_ids=False,
-        tokenizer=tokenizer,
-        reset_position_ids=args.reset_position_ids,
-        reset_attention_mask=args.reset_attention_mask,
-        eod_mask_loss=args.eod_mask_loss,
-        mid_level_dataset_surplus=args.mid_level_dataset_surplus,
-        allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
-    )
-
-    # GPT datasets.
-    print_rank_0(" > multi-split gpt datasets.")
-    train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
-        MultiSplitGPTDataset,
-        train_valid_test_num_samples,
-        is_dataset_built_on_rank,
-        data_config,
-    ).build()
-
-    gpt_datasets = {
-        "train" : (train_ds, train_valid_test_num_samples[0]),
-        "valid" : (valid_ds, train_valid_test_num_samples[1]),
-        "test"  : (test_ds, train_valid_test_num_samples[2]),
-    }
-
-    # Retro datasets.
-    if args.retro_add_retriever:
-        return get_retro_datasets(
-            config=retro_config,
-            gpt_datasets=gpt_datasets,
-            sample_length=args.seq_length,
-            eod_token_id=get_tokenizer().eod,
-        )
-
-    # Multi-split GPT datasets.
-    else:
-        return (
-            gpt_datasets["train"][0],
-            gpt_datasets["valid"][0],
-            gpt_datasets["test"][0],
-        )
-
-
-if __name__ == "__main__":
-
-    # Temporary for transition to core datasets.
-    train_valid_test_datasets_provider.is_distributed = True
-
-    pretrain(train_valid_test_datasets_provider,
-             model_provider,
-             ModelType.retro_decoder,
-             forward_step,
-             args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})
diff --git a/tests/unit_tests/dist_checkpointing/models/test_t5_model.py b/tests/unit_tests/dist_checkpointing/models/test_t5_model.py
index 1e44ee527ae..e393c806a94 100644
--- a/tests/unit_tests/dist_checkpointing/models/test_t5_model.py
+++ b/tests/unit_tests/dist_checkpointing/models/test_t5_model.py
@@ -6,14 +6,6 @@
 from megatron.core import parallel_state as ps
 from megatron.core.dist_checkpointing import load, save
 from megatron.core.dist_checkpointing.validation import StrictHandling
-from megatron.core.models.retro.decoder_spec import (
-    get_retro_decoder_layer_local_spec,
-    get_retro_decoder_layer_te_spec,
-)
-from megatron.core.models.retro.encoder_spec import (
-    get_retro_encoder_layer_local_spec,
-    get_retro_encoder_layer_te_spec,
-)
 from megatron.core.models.T5 import T5Model
 from megatron.core.models.T5.t5_spec import decoder_model_with_local_spec as t5_decoder_local_spec
 from megatron.core.models.T5.t5_spec import (
@@ -94,14 +86,8 @@ def test_sharded_state_dict_save_load(
         self, tmp_path_dist_ckpt, src_spec_type, dst_spec_type, model_type
     ):
         enc_dec_spec_fn = {
-            'te': {
-                't5': (t5_encoder_te_spec, t5_decoder_te_spec),
-                'retro': (get_retro_encoder_layer_te_spec, get_retro_decoder_layer_te_spec),
-            },
-            'local': {
-                't5': (t5_encoder_local_spec, t5_decoder_local_spec),
-                'retro': (get_retro_encoder_layer_local_spec, get_retro_decoder_layer_local_spec),
-            },
+            'te': {'t5': (t5_encoder_te_spec, t5_decoder_te_spec)},
+            'local': {'t5': (t5_encoder_local_spec, t5_decoder_local_spec)},
         }
         src_encoder_decoder_spec_fn = enc_dec_spec_fn[src_spec_type][model_type]
         dst_encoder_decoder_spec_fn = enc_dec_spec_fn[dst_spec_type][model_type]
@@ -155,14 +141,8 @@ def test_parallel_reconfiguration_e2e(
         *dest_tp_pp, dst_encpp = dest_tp_pp_encpp
 
         enc_dec_spec_fn = {
-            'te': {
-                't5': (t5_encoder_te_spec, t5_decoder_te_spec),
-                'retro': (get_retro_encoder_layer_te_spec, get_retro_decoder_layer_te_spec),
-            },
-            'local': {
-                't5': (t5_encoder_local_spec, t5_decoder_local_spec),
-                'retro': (get_retro_encoder_layer_local_spec, get_retro_decoder_layer_local_spec),
-            },
+            'te': {'t5': (t5_encoder_te_spec, t5_decoder_te_spec)},
+            'local': {'t5': (t5_encoder_local_spec, t5_decoder_local_spec)},
         }
 
         common_test_parallel_reconfiguration_e2e(
diff --git a/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py b/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py
index 927b51d5ddb..19658986d6e 100644
--- a/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py
+++ b/tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py
@@ -131,7 +131,6 @@ def create_args():
     args.ckpt_fully_parallel_save = False
     args.ckpt_fully_parallel_load = False
     args.auto_detect_ckpt_format = False
-    args.retro_add_retriever = False
     args.ckpt_convert_update_legacy_dist_opt_format = False
     args.ckpt_step = None
     args.use_dist_ckpt = True
diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py
index ce068ef3227..dd12ecd7684 100644
--- a/tests/unit_tests/dist_checkpointing/utils.py
+++ b/tests/unit_tests/dist_checkpointing/utils.py
@@ -154,7 +154,6 @@ def init_checkpointing_mock_args(args, ckpt_dir, fully_parallel=False):
     args.consumed_train_samples = 0
     args.skipped_train_samples = 0
     args.consumed_valid_samples = 0
-    args.retro_add_retriever = False
     args.no_load_optim = False
     args.no_load_rng = False
     args.dist_ckpt_strictness = 'assume_ok_unexpected'
diff --git a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py
index 5b01aac6b2e..a3990d25ecf 100644
--- a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py
+++ b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py
@@ -126,7 +126,6 @@ def create_args():
     args.ckpt_fully_parallel_save = False
     args.ckpt_fully_parallel_load = False
     args.auto_detect_ckpt_format = False
-    args.retro_add_retriever = False
     args.ckpt_convert_update_legacy_dist_opt_format = False
     args.ckpt_step = None
     args.use_dist_ckpt = True
diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py
index 2964c67c7ce..9a7a44939a3 100644
--- a/tests/unit_tests/test_checkpointing.py
+++ b/tests/unit_tests/test_checkpointing.py
@@ -113,7 +113,6 @@ def create_args():
     args.dist_ckpt_optim_fully_reshardable = False
     args.distrib_optim_fully_reshardable_mem_efficient = False
     args.auto_detect_ckpt_format = False
-    args.retro_add_retriever = False
     args.ckpt_convert_update_legacy_dist_opt_format = False
     args.ckpt_step = None
     args.swiglu = True
diff --git a/tests/unit_tests/transformer/test_retro_attention.py b/tests/unit_tests/transformer/test_retro_attention.py
deleted file mode 100644
index 85c5347c909..00000000000
--- a/tests/unit_tests/transformer/test_retro_attention.py
+++ /dev/null
@@ -1,204 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-import os
-import types
-
-import pytest
-import torch
-
-from megatron.core.models.retro import RetroConfig, get_retro_decoder_block_spec
-from megatron.core.models.retro.decoder_attention import (
-    RetroDecoderBiasDropoutAdd,
-    RetroDecoderCrossAttention,
-)
-from megatron.core.models.retro.encoder_attention import (
-    RetroEncoderBiasDropoutAdd,
-    RetroEncoderCrossAttention,
-    RetroEncoderLayerNorm,
-)
-from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer.transformer_block import TransformerBlock
-from tests.unit_tests.test_utilities import Utils
-
-
-class TestRetroAttention:
-
-    @classmethod
-    def get_config(cls):
-        return RetroConfig(
-            num_layers=12,
-            hidden_size=16,
-            num_attention_heads=4,
-            use_cpu_initialization=True,
-            retro_num_neighbors=2,
-            retro_chunk_length=4,
-            retro_retrieved_length=8,
-            retro_split_preprocessing="98,2,0",
-        )
-
-    @classmethod
-    def get_modules(cls, config, use_transformer_engine, use_gpu):
-
-        # Retro decoder layer.
-        decoder_block_spec = get_retro_decoder_block_spec(
-            config, use_transformer_engine=use_transformer_engine
-        )
-        decoder_block = TransformerBlock(config=config, spec=decoder_block_spec)
-        decoder_layers = [
-            layer
-            for layer in decoder_block.layers
-            if isinstance(layer.cross_attention, RetroDecoderCrossAttention)
-        ]
-        decoder_layer = decoder_layers[0]
-
-        # Retro encoder layer.
-        encoder_block = decoder_layer.cross_attention.encoder
-        encoder_layers = [
-            layer
-            for layer in encoder_block.layers
-            if isinstance(layer.cross_attention, RetroEncoderCrossAttention)
-        ]
-        encoder_layer = encoder_layers[0]
-
-        # Modules.
-        modules = types.SimpleNamespace(
-            decoder_attn=decoder_layer.cross_attention,
-            decoder_bda=decoder_layer.cross_attn_bda,
-            encoder_attn=encoder_layer.cross_attention,
-            encoder_bda=encoder_layer.cross_attn_bda,
-            encoder_norm=encoder_layer.pre_mlp_layernorm,
-        )
-
-        # GPU.
-        if use_gpu:
-            [m.cuda() for m in vars(modules).values()]
-
-        return modules
-
-    def setup_method(self, method):
-        Utils.initialize_model_parallel(1, 1)
-        os.environ['NVTE_FLASH_ATTN'] = "0"
-        os.environ['NVTE_FUSED_ATTN'] = "0"
-
-        model_parallel_cuda_manual_seed(123)
-
-    def teardown_method(self, method):
-        Utils.destroy_model_parallel()
-
-    def test_constructor(self):
-
-        config = self.get_config()
-        modules = self.get_modules(config, use_transformer_engine=True, use_gpu=False)
-
-        assert isinstance(modules.decoder_attn, RetroDecoderCrossAttention)
-        assert isinstance(modules.decoder_bda, RetroDecoderBiasDropoutAdd)
-        assert isinstance(modules.encoder_attn, RetroEncoderCrossAttention)
-        assert isinstance(modules.encoder_bda, RetroEncoderBiasDropoutAdd)
-        assert isinstance(modules.encoder_norm, RetroEncoderLayerNorm)
-
-        assert modules.decoder_attn.attn.layer_number == 6
-        assert modules.encoder_attn.attn.layer_number == 1
-
-        get_nparams = lambda m: sum(p.numel() for p in m.parameters())
-        assert get_nparams(modules.decoder_attn) == 8768
-        assert get_nparams(modules.decoder_bda) == 0
-        assert get_nparams(modules.encoder_attn) == 1088
-        assert get_nparams(modules.encoder_bda) == 0
-        assert get_nparams(modules.encoder_norm) == 32
-
-    def test_cpu_forward(self):
-        # we can't currently do this because the global memory buffer is on GPU
-        pass
-
-    def run_gpu_forward(self, recompute_granularity, use_transformer_engine):
-
-        config = self.get_config()
-        config.recompute_granularity = recompute_granularity
-        modules = self.get_modules(config, use_transformer_engine, use_gpu=True)
-
-        seq_length = 32
-        micro_batch_size = 2
-        n_chunks_per_sample = seq_length // config.retro_chunk_length
-
-        # Init tensors.
-        hidden_states = torch.ones((seq_length, micro_batch_size, config.hidden_size)).cuda()
-        attention_mask = None
-        decoder_context = torch.ones(
-            (
-                config.retro_retrieved_length,
-                config.retro_num_neighbors * micro_batch_size * n_chunks_per_sample,
-                config.hidden_size,
-            )
-        ).cuda()
-        encoder_context = torch.ones(
-            (config.retro_chunk_length, micro_batch_size * n_chunks_per_sample, config.hidden_size)
-        ).cuda()
-
-        # Forward decoder.
-        decoder_attn_output = modules.decoder_attn(hidden_states, attention_mask, decoder_context)
-        with torch.enable_grad():
-            decoder_bda_output = modules.decoder_bda(True, True)(
-                decoder_attn_output, hidden_states, config.hidden_dropout
-            )
-
-        # Forward encoder.
-        encoder_attn_output_tuples = modules.encoder_attn(decoder_context, None, encoder_context)
-        with torch.enable_grad():
-            encoder_bda_output = modules.encoder_bda(True, True)(
-                encoder_attn_output_tuples, decoder_context, config.retro_encoder_hidden_dropout
-            )
-        encoder_norm_output = modules.encoder_norm(encoder_bda_output)
-
-        # Verify decoder.
-        assert set(decoder_attn_output.keys()) == set(
-            ["ns", "bs", "d", "l", "pad", "attention_output", "attention_bias", "context"]
-        )
-        assert decoder_attn_output["ns"] == seq_length
-        assert decoder_attn_output["bs"] == micro_batch_size
-        assert decoder_attn_output["d"] == config.hidden_size
-        assert decoder_attn_output["l"] == n_chunks_per_sample
-        assert decoder_attn_output["pad"] == 3
-        assert tuple(decoder_attn_output["attention_output"].shape) == (
-            config.retro_chunk_length,
-            micro_batch_size * n_chunks_per_sample,
-            config.hidden_size,
-        )
-        assert tuple(decoder_attn_output["attention_bias"].shape) == (config.hidden_size,)
-        assert decoder_attn_output["context"].shape == (
-            config.retro_retrieved_length * config.retro_num_neighbors,
-            micro_batch_size * n_chunks_per_sample,
-            config.hidden_size,
-        )
-        assert decoder_bda_output.shape == hidden_states.shape
-
-        # Verify encoder.
-        assert len(encoder_attn_output_tuples) == config.retro_num_neighbors
-        for output, bias, residual in encoder_attn_output_tuples:
-            assert tuple(output.shape) == (
-                config.retro_retrieved_length,
-                micro_batch_size * n_chunks_per_sample,
-                config.hidden_size,
-            )
-            assert tuple(bias.shape) == (config.hidden_size,)
-            assert tuple(residual.shape) == (
-                config.retro_retrieved_length,
-                micro_batch_size * n_chunks_per_sample,
-                config.hidden_size,
-            )
-        assert encoder_bda_output.shape == (
-            config.retro_retrieved_length,
-            config.retro_num_neighbors * micro_batch_size * n_chunks_per_sample,
-            config.hidden_size,
-        )
-        assert encoder_norm_output.shape == (
-            config.retro_retrieved_length,
-            config.retro_num_neighbors * micro_batch_size * n_chunks_per_sample,
-            config.hidden_size,
-        )
-
-    @pytest.mark.flaky
-    @pytest.mark.flaky_in_dev
-    def test_gpu_forward(self):
-        for recompute_granularity in (None, 'selective'):
-            for use_transformer_engine in (True, False):
-                self.run_gpu_forward(recompute_granularity, use_transformer_engine)
diff --git a/tools/bert_embedding/embed.py b/tools/bert_embedding/embed.py
index 2236182a751..effc6f6d91e 100644
--- a/tools/bert_embedding/embed.py
+++ b/tools/bert_embedding/embed.py
@@ -1,18 +1,20 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
 
 from functools import partial
+from types import SimpleNamespace
+from typing import Callable, Dict, List, Optional, Tuple, TypedDict
 import numpy as np
 import os
 import time
 import torch
+from torch.distributed import ProcessGroup
 from torch.utils.data import BatchSampler, DataLoader, SequentialSampler, Subset
 from torch.utils.data._utils.collate import default_collate
-from tqdm import tqdm
 
 from megatron.training import get_args, get_tokenizer, print_rank_0
 from megatron import core
 from megatron.training.arguments import core_transformer_config_from_args
-from megatron.core.datasets.retro.utils import get_blocks_by_rank
+from megatron.core import parallel_state
 from megatron.core.enums import ModelType
 from megatron.core.pipeline_parallel import get_forward_backward_func
 from megatron.legacy.model import BertModel
@@ -23,6 +25,20 @@
 from .external_libs import h5py
 from .huggingface import HuggingfaceEmbedder
 
+try:
+    from tqdm import tqdm
+
+    HAVE_TQDM = True
+except ImportError:
+    HAVE_TQDM = False
+
+try:
+    import h5py
+
+    HAVE_H5PY = True
+except ImportError:
+    HAVE_H5PY = False
+
 
 def collate_batch(samples):
     """Collate samples of various lengths.
@@ -127,6 +143,211 @@ def embed_data_loader(models, data_loader, tag):
     return embeddings
 
 
+class Block(TypedDict):
+    """Specific block arg type to mute mypy."""
+
+    range: Tuple[int, int]
+    path: str
+
+
+def get_blocks(
+    dirname: str, n_samples: int, block_size: int, validate: Optional[Callable] = None
+) -> SimpleNamespace:
+    """Divide range [0, num_samples) to sequence of block ranges.
+
+    This is a core method within the concept of block processing. The idea
+    is to divide a range (size n_samples) into a sequence of blocks. Each
+    block corresponds to a file within 'dirname' with name
+    '{start_idx}-{end_idx}.hdf5'. This method checks for the existence of
+    these files, and returns two lists, one for existing blocks and one for
+    missing blocks.
+
+    Args:
+        dirname (str): Path to directory containing block files.
+        n_samples (int): Ideal number of samples.
+            The total number of saved block data is <=n_samples.
+        block_size (int): Max number of samples per block file (e.g., 100000).
+        validate (Callable): Method for validating each block file during load.
+
+    Returns:
+        A namespace consisting of 2 lists: existing blocks, and missing blocks.
+        The total number of samples between the existing and missing blocks should
+        equal n_samples above.
+    """
+
+    if not HAVE_TQDM:
+        raise ImportError("tqdm is required to use the BertDataset. Please install tqdm.")
+
+    if not HAVE_H5PY:
+        raise ImportError("h5py is required to use the BertDataset. Please install h5py.")
+
+    assert os.path.isdir(dirname), "missing directory '%s.'" % dirname
+
+    # Block ranges.
+    block_start_idxs = list(range(0, n_samples, block_size))
+    block_end_idxs = [min(n_samples, i + block_size) for i in block_start_idxs]
+    block_ranges = list(zip(block_start_idxs, block_end_idxs))
+
+    # All block files (existing + missing).
+    n_digits = int(np.ceil(np.log(n_samples) / np.log(10)) + 1)
+
+    all_blocks: List[Block] = [
+        {
+            "range": r,
+            "path": os.path.join(
+                dirname, "%s-%s.hdf5" % tuple([str(i).zfill(n_digits) for i in r])
+            ),
+        }
+        for r in block_ranges
+    ]
+    all_block_path_set = set(block["path"] for block in all_blocks)
+
+    # Validate function.
+    validate = (lambda f: None) if validate is None else validate
+
+    # Delete corrupt files.
+    if torch.distributed.get_rank() == 0:
+        existing_block_paths = [
+            block["path"] for block in all_blocks if os.path.exists(block["path"])
+        ]
+        for index, path in enumerate(tqdm(existing_block_paths, "validating block.")):
+            assert path in all_block_path_set, "unexpected filename, '%s'." % path
+
+            try:
+                f = h5py.File(path, "r")
+            except Exception:
+                os.remove(path)
+                continue
+
+            try:
+                validate(f)
+            except Exception:
+                os.remove(path)
+            finally:
+                f.close()
+
+    # Wait for files to be deleted.
+    torch.distributed.barrier()
+
+    # Collect blocks.
+    blocks = SimpleNamespace(
+        existing=[b for b in all_blocks if os.path.exists(b["path"])],
+        missing=[b for b in all_blocks if not os.path.exists(b["path"])],
+    )
+
+    return blocks
+
+
+def get_blocks_by_rank(
+    dirname: str,
+    n_samples: int,
+    block_size: int,
+    validate: Optional[Callable] = None,
+    sample: Optional[float] = None,
+    process_group: Optional[ProcessGroup] = None,
+) -> SimpleNamespace:
+    """Divide existing and missing blocks evenly across all ranks.
+
+    See 'get_blocks()' above for description. The returned lists of existing and
+    missing blocks are split evenly across ranks via interleaving. This way,
+    each rank has a roughly equal number of blocks to process for a
+    downstream operation.
+
+    Args:
+        dirname (str): Path to directory containing block files.
+        n_samples (int): Ideal number of samples. The total number of saved block data
+            is <=n_samples.
+        block_size (int): Max number of samples per block file (e.g., 100000).
+        validate (Callable): Method for validating each block file during load.
+        sample (Optional[float]): If provided, sample a random subset of the blocks.
+            Used for validating preprocessing correctness.
+        process_group (Optional[ProcessGroup]): Process group for distributed operations.
+            If None, uses data parallel group.
+
+    Returns:
+        A namespace consisting of 2 lists: existing blocks, and missing blocks.
+        Each of these two lists is potentially a sub-sample of the total set of
+        existing and missing blocks, depending on whether sampling is used.
+        Additionally, the attributes n_existing_world and n_missing_world are the
+        total number of existing and missing blocks, independent of samples.
+        Therefore, (n_existing_world + n_missing_world) * block_size == n_samples.
+    """
+
+    if process_group is None:
+        process_group = parallel_state.get_data_parallel_group()
+
+    # Get world blocks.
+    blocks = get_blocks(dirname, n_samples, block_size, validate)
+
+    # This rank's existing and missing files.
+    rank_existing_blocks = blocks.existing[
+        process_group.rank() : len(blocks.existing) : process_group.size()
+    ]
+    rank_missing_blocks = blocks.missing[
+        process_group.rank() : len(blocks.missing) : process_group.size()
+    ]
+
+    # Extend rank's existing and missing blocks (with None) such that all ranks
+    # have equal length lists. This allows for easier tracking of global progress.
+    def get_world_max(n: int) -> int:
+        """Get max value across ranks.
+
+        Args:
+            n (int): Value on this rank.
+
+        Returns:
+            Max value across all ranks.
+        """
+        n_tensor = torch.cuda.LongTensor([n])
+        torch.distributed.all_reduce(n_tensor, op=torch.distributed.ReduceOp.MAX)
+        return n_tensor.item()
+
+    max_n_existing = get_world_max(len(rank_existing_blocks))
+    max_n_missing = get_world_max(len(rank_missing_blocks))
+
+    rank_existing_blocks += [None] * (max_n_existing - len(rank_existing_blocks))
+    rank_missing_blocks += [None] * (max_n_missing - len(rank_missing_blocks))
+
+    # Collect blocks.
+    blocks = SimpleNamespace(
+        n_existing_world=len(blocks.existing),
+        n_missing_world=len(blocks.missing),
+        existing=rank_existing_blocks,
+        missing=rank_missing_blocks,
+    )
+
+    if sample is not None:
+        # Sample existing and missing blocks evenly across all ranks. The
+        # returned lists of blocks are randomly sampled (without replacement)
+        # to yield `sample * len(blocks)` number of blocks.
+
+        # Randomly sample blocks.
+        def sample_blocks(_blocks: List[Optional[Dict]]) -> List[Optional[Dict]]:
+            """Sample a random subset of all blocks.
+
+            Args:
+                _blocks (List[Optional[Dict]]): List of all blocks.
+
+            Returns:
+                A random subset of the blocks.
+            """
+            n_blocks_sample = int(np.ceil(sample * len(_blocks)))
+            sampled_blocks: List[Optional[Dict]] = [b for b in _blocks if b is not None]
+
+            np.random.seed(None)
+            np.random.shuffle(sampled_blocks)
+
+            sampled_blocks = sampled_blocks[:n_blocks_sample]
+            sampled_blocks += [None] * (n_blocks_sample - len(sampled_blocks))
+
+            return sampled_blocks
+
+        blocks.existing = sample_blocks(blocks.existing)
+        blocks.missing = sample_blocks(blocks.missing)
+
+    return blocks
+
+
 class TextDataset(torch.utils.data.Dataset):
     '''Dataset that holds a list of strings.'''
 
diff --git a/tools/retro/README.md b/tools/retro/README.md
deleted file mode 100644
index 395005e73bf..00000000000
--- a/tools/retro/README.md
+++ /dev/null
@@ -1,256 +0,0 @@
-# Retro and InstructRetro
-
-Retro [(Borgeaud et al., 2022)](https://arxiv.org/abs/2112.04426) is an autoregressive decoder-only language model (LM)
-pretrained with retrieval-augmentation.
-Retro features practical scalability to support large-scale pretraining from scratch by retrieving from trillions of
-tokens.
-Pretraining with retrieval provides a more efficient storage mechanism of factual knowledge, when compared to storing
-factual knowledge implicitly within the network's parameters, thus largely reducing model parameters while achieving
-lower perplexity than standard GPT.
-Retro also provides the flexibility to update the
-knowledge stored in LMs [(Wang et al., 2023a)](https://arxiv.org/abs/2304.06762)
-by updating the retrieval database without training LMs again.
-
-InstructRetro [(Wang et al., 2023b)](https://arxiv.org/abs/2310.07713) further scales up the size of Retro to 48B,
-featuring the largest LLM pretrained with retrieval (as of December 2023).
-The obtained foundation model, Retro 48B, largely outperforms the GPT counterpart in terms of perplexity.
-With instruction tuning on Retro, InstructRetro demonstrates significant improvement over the instruction tuned GPT on
-downstream tasks in the zero-shot setting. Specifically, the average improvement of InstructRetro is 7% over its GPT
-counterpart across 8 short-form QA tasks, 10% over GPT across 4 challenging long-form QA tasks, and 16% over GPT across
-3 summarization tasks. We also find that one can ablate the encoder from InstructRetro architecture and directly use the
-InstructRetro decoder backbone as GPT, while achieving comparable results.
-
-This README provides an end-to-end tutorial to reproduce Retro and InstructRetro.
-
-# Contents
-
-* [Checkpoints](#checkpoints)
-* [End-to-end Reproduction Guide](#end-to-end-reproduction-guide)
-    * [Step 0: Prepare the environment](#step-0-prepare-the-environment)
-        * [Docker image](#docker-image)
-        * [Install dependencies](#install-dependencies)
-    * [Step 1: Build retrieval database](#step-1-build-retrieval-database)
-    * [Step 2: Pretraining](#step-2-pretraining)
-    * [Step 3: Perplexity evaluation](#step-3-perplexity-evaluation)
-    * [Step 4: Instruction tuning](#step-4-instruction-tuning)
-    * [Step 5: Downstream task evaluation](#step-5-downstream-task-evaluation)
-* [Citations](#citations)
-
-# Checkpoints
-
-We provide the pretrained checkpoints of Retro and InstructRetro in the following table. The checkpoints are available
-to download through the following links:
-
-| Model                   | Size | Instruction Tuning | Download Link 1                                                    | Download Link 2                                                                | Download Link 3                                                                                      |
-|-------------------------|------|--------------------|--------------------------------------------------------------------|--------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------|
-| `retro-8b-base-4k`      | 8b   |                    | [Huggingface](https://huggingface.co/nvidia/retro-8b-base-4k)      | [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/models/retro-8b-base-4k)      | [Google Drive](https://drive.google.com/drive/folders/1uSQ5DAsuvx_8XcbtnVfs_MGvEOcx0uK_?usp=sharing) |
-| `retro-8b-instruct-4k`  | 8b   | ✅                  | [Huggingface](https://huggingface.co/nvidia/retro-8b-instruct-4k)  | [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/models/retro-8b-instruct-4k)  | [Google Drive](https://drive.google.com/drive/folders/1v5dKaSN0cm2lwyAWpFaJtlTrLhtMZXsI?usp=sharing) |
-| `retro-48b-base-4k`     | 48b  |                    | [Huggingface](https://huggingface.co/nvidia/retro-48b-base-4k)     | [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/models/retro-48b-base-4k)     | [Google Drive](https://drive.google.com/drive/folders/1rtNpf0CiLElSHQcr3aLI3zgfI3teGTP5?usp=sharing) |
-| `retro-48b-instruct-4k` | 48b  | ✅                  | [Huggingface](https://huggingface.co/nvidia/retro-48b-instruct-4k) | [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/models/retro-48b-instruct-4k) | [Google Drive](https://drive.google.com/drive/folders/1qdb0AQjSsAPGlWaIu3wgHPjf_nwLeY5h?usp=sharing) |
-
-# End-to-end Reproduction Guide
-
-In this README, we provide an end-to-end reproduction guide for InstructRetro, covering from large-scale retrieval
-construction, pretraining, perplexity evaluation, instruction tuning, to downstream task evaluation.
-
-If you are interested in evaluation only, we also [open-sourced our checkpoints](#checkpoints) and you can directly go
-to [Step 5](#step-5-downstream-task-evaluation) to evaluate the checkpoints on downstream tasks.
-
-## Step 0: Prepare the environment
-
-We recommend using docker environment to run the code.
-
-### Docker image
-
-We provide a docker build file in [tools/retro/examples/Dockerfile](examples/Dockerfile) for the reproduction. The
-docker image is based on the [NGC docker](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags) `nvcr.io/nvidia/pytorch:23.09-py3`.
-
-### Install dependencies
-
-Clone the Megatron repo:
-
-```bash
-git clone --branch InstructRetro https://github.com/NVIDIA/Megatron-LM.git
-```
-
-If docker is not available, we recommend starting from a clean conda environment with the following runtime
-dependencies:
-
-- Python 3.10
-- NVIDIA CUDA® 12.2.1
-- NVIDIA cuBLAS 12.2.5.6
-- NVIDIA cuDNN 8.9.5
-- NVIDIA NCCL 2.18.5
-- PyTorch 2.1.0a0+32f93b1
-
-Then install Retro-specific dependencies, including:
-
-```bash
-pip install -U faiss-gpu
-pip install -U transformers
-pip install -U sentencepiece
-pip install -U h5py
-pip install -U nltk
-pip install -U einops
-```
-
-## Step 1: Build retrieval database
-
-In this step, we build a large-scale retrieval database for InstructRetro
-through [Faiss](https://github.com/facebookresearch/faiss) to retrieve from trillions of tokens, and preprocess (and
-save) the retrieval neighbors for the pretraining step.
-
-Please refer to [tools/retro/build_db.md](build_db.md) for more details.
-
-## Step 2: Pretraining
-
-*Please strictly follow Step 1 to build the retrieval database before pretraining to make sure the preprocessed
-retrieval neighbors match the pretraining corpus.*
-
-In the pretraining step, we support both pretraining from scratch and continued pretraining from a pretrained GPT model.
-
-We provide a template pretraining script to pretrain 843M Retro from scratch. Prepare your own arguments and update our
-templates in [tools/retro/examples/pretrain_model.sh](examples/pretrain_model.sh). Please note that the data path should
-be exactly matching the one used in Step 1 to make sure the preprocessed retrieval neighbors match the pretraining
-corpus.
-
-[//]: # (Take the example of the Wikipedia corpus)
-
-```bash
-bash tools/retro/examples/pretrain_model.sh
-```
-
-After pretraining, the model checkpoints will be saved in the `--save` directory if you specified the arg
-in `pretrain_model.sh`.
-
-To continue pretraining with retrieval from a pretrained GPT model, please specify `--load` in `pretrain_model.sh` to
-load the pretrained GPT model checkpoint (the architecture of GPT, including hidden size, number of layers, and
-activation methods, should be exactly the same as the one used for Retro). You should also
-specify  `--no-load-optim --finetune` to make sure the optimizer state is not loaded from the pretrained GPT model and
-the continued pretraining with retrieval is from a clean start. After the first job / the first run, you will continue
-pretraining with retrieval from your last checkpoint. In the follow-up jobs, you should launch the pretraining without
-the flags `--no-load-optim --finetune` to make sure the optimizer state is correctly loaded from your last job.
-
-## Step 3: Perplexity evaluation
-
-During pretraining, we will automatically evaluate the model perplexity on the specified validation corpus
-every `--eval-interval` steps. The validation corpus should be exactly the same as the one used in Step 1 to make sure
-the preprocessed retrieval neighbors match the pretraining corpus.
-
-To evaluate the perplexity of a pretrained model, please add `--skip-train` in `pretrain_model.sh` to skip the
-pretraining step and only evaluate the perplexity of the model specified in `--load` on the validation corpus. Run the
-above command again to evaluate the perplexity of a pretrained model:
-
-```bash
-bash tools/retro/examples/pretrain_model.sh
-```
-
-## Step 4: Instruction tuning
-
-In this step, we fine-tune the pretrained model on the downstream task with instructions. We provide a template
-instruction tuning script to fine-tune 843M Retro.
-
-We also provide an open-source blend of instruction tuning datasets. The dataset is available to download
-through [here](https://drive.google.com/file/d/1nzKwwYf8lYb9gN3P4YO8pFNU_B2nMYe1/view?usp=sharing). The blendable
-dataset consists of the following open-source instruction tuning datasets:
-
-### Instruction Tuning Dataset Breakdown
-
-| Dataset                                                    | Samples | Epochs | Sampling Prob |
-|------------------------------------------------------------|--------:|-------:|--------------:|
-| [soda](https://arxiv.org/abs/2212.10465)                   |    2560 |  0.005 |         0.020 |
-| [eli5](https://arxiv.org/abs/1907.09190)                   |    2561 |  0.055 |         0.020 |
-| [self_instruct_short](https://arxiv.org/abs/2212.10560)    |    1280 |  0.043 |         0.010 |
-| [self_instruct_long](https://arxiv.org/abs/2212.10560)     |    2560 |  0.333 |         0.020 |
-| [unnatural-instructions](https://arxiv.org/abs/2212.09689) |    2560 |  0.024 |         0.020 |
-| [flan_cot](https://arxiv.org/abs/2210.11416)               |    1280 |  0.093 |         0.010 |
-| [dolly](https://arxiv.org/abs/2305.13735)                  |    6400 |  0.938 |         0.050 |
-| [oasst-skip-noncode](https://open-assistant.io/)           |  104558 |  1.839 |         0.817 |
-| [oasst-skip-code](https://open-assistant.io/)              |    4243 |  1.839 |         0.033 |
-
-Refer to the paper links above for more details about each instruction tuning dataset.
-
-*We note that the provided instruction tuning dataset is all from open-source instruction tuning datasets. It is
-slightly different from what we use in [InstructRetro](https://arxiv.org/abs/2310.07713), which contains private and
-proprietary datasets. Thus a 1-2% accuracy difference in downstream tasks may be expected.*
-
-### Instruction tuning script
-
-Download
-the [blended instruction tuning dataset](https://drive.google.com/file/d/1nzKwwYf8lYb9gN3P4YO8pFNU_B2nMYe1/view?usp=sharing)
-in your data home directory `$DATA_HOME` and update our templates
-in [tools/retro/sft/sft_retro_lm.sh](sft/sft_retro_lm.sh).
-
-An example command to run instruction tuning on 843M Retro is as follows:
-
-```bash
-                                      [blend-dataset-name] [model-size] [batch-size]  [lr]    [checkpoints]
-bash tools/retro/sft/sft_retro_lm.sh       open_inst               843m            128    5e-6  <path/to/pretrained/retro>
-```
-
-The `blend_dataset_name` argument will blend all the datasets within the `$DATA_HOME` following the weights and
-configurations specified in the `${blend_dataset_name}.sh` ([open_inst.sh](sft/open_inst.sh) in the example above).
-The checkpoints will be saved in the `--save` directory. For example, it will be saved to
-`<SFT_HOME>/checkpoints/applications/retro-sft_pp1_same_format_ctx1_843m_128_5e-6`.
-
-## Step 5: Downstream task evaluation
-
-In this step, we demonstrate how to run InstructRetro for zero-shot evaluation on downstream question answering (QA)
-tasks. We provide the pre-processed open-source evaluation datasets with a unified format for different tasks. The
-evaluation datasets used in our paper are available to download
-through [here](https://drive.google.com/drive/folders/1xw-N0LJR_lIWnH6BKzHIb49quVCS_V72?usp=sharing). Please stick to
-the same retro workdir used in Step 0-4 to make sure the preprocessed retrieval neighbors match the pretraining corpus.
-If you directly come to Step 5, an example retro workdir with `args.json` for 800M Retro is
-provided [here](https://drive.google.com/file/d/121GqAdMvf8bJEBZRt-SD4uhW-SRWgI3s/view?usp=sharing). Note that the args
-in the json can be overwritten through the command line.
-
-We present an example command to run retro generation given the InstructRetro checkpoints and the Natural Question (NQ)
-task. The example command is for the 843m InstructRetro obtained in Step 4. Please specify the directory for the NQ
-dataset and update the command accordingly for other checkpoints.
-
-```bash
-bash tools/retro/text_generation/retro_generate.sh nq 843m greedy test  0 20000 1000 5 pp1 <SFT_HOME>/checkpoints/applications/retro-sft_pp1_same_format_ctx1_843m_128_5e-6 2
-```
-
-The generated responses will be saved in the corresponding checkpoint directory. For example, for the 843m
-InstructRetro, it will be saved to
-`<SFT_HOME>/checkpoints/applications/retro-sft_pp1_same_format_ctx1_843m_128_5e-6/retro-generate-nq_5_2_843m_test_greedy_0_20000_1000.txt`.
-
-To evaluate the F1 / Exact Match (EM) scores of the generated responses, we provide an example script to run the
-evaluation on the NQ dataset. Please specify the directory for the NQ dataset and update the command accordingly for
-other checkpoints and downstream tasks.
-
-```bash
-python3 tools/retro/text_generation/evaluate.py
-```
-
-# Citations
-
-See more details from our papers:
-
-[Shall we Pretrain Autoregressive Language Models with Retrieval? A Comprehensive Study.](https://arxiv.org/abs/2304.06762)
-
-_Boxin Wang, Wei Ping, Peng Xu, Lawrence McAfee, Zihan Liu, Mohammad Shoeybi, Yi Dong, Oleksii Kuchaiev, Bo Li, Chaowei
-Xiao, Anima Anandkumar, Bryan Catanzaro._ (EMNLP 2023)
-
-[InstructRetro: Instruction Tuning post Retrieval-Augmented Pretraining.](https://arxiv.org/abs/2310.07713)
-
-_Boxin Wang, Wei Ping, Lawrence McAfee, Peng Xu, Bo Li, Mohammad Shoeybi, Bryan Catanzaro._
-
-Please cite the papers as follows if you use the data or code from this repo:
-
-```bibtex
-@inproceedings{wang2023shall,
-    title   = {Shall We Pretrain Autoregressive Language Models with Retrieval? A Comprehensive Study},
-    author  = {Boxin Wang and Wei Ping and Peng Xu and Lawrence McAfee and Zihan Liu and Mohammad Shoeybi and Yi Dong and Oleksii Kuchaiev and Bo Li and Chaowei Xiao and Anima Anandkumar and Bryan Catanzaro},
-    journal = {The 2023 Conference on Empirical Methods in Natural Language Processing},
-    year    = {2023}
-}
-
-@article{wang2023instructretro,
-    title   = {InstructRetro: Instruction Tuning post Retrieval-Augmented Pretraining},
-    author  = {Boxin Wang and Wei Ping and Lawrence McAfee and Peng Xu and Bo Li and Mohammad Shoeybi and Bryan Catanzaro},
-    year    = {2023},
-    journal = {arXiv preprint arXiv: 2310.07713}
-}
-```
diff --git a/tools/retro/build_db.md b/tools/retro/build_db.md
deleted file mode 100644
index c99952485ab..00000000000
--- a/tools/retro/build_db.md
+++ /dev/null
@@ -1,421 +0,0 @@
-This directory contains a collection of tools for building the retrieval database and pretraining neighbors for Retro. This preprocessing pipeline is broken into 3 main stages:
-
-1. **Build retrieval chunk database** : Used for retrieving neighbors and continuation chunks, which are then passed through the retrieval encoder.
-2. **Build index for similarity search** : Train and build a search index for querying chunk neighbors.
-3. **Query pretraining neighbors** : For matching pretraining samples to database chunks. Neighbors are generated separately for training, validation, and test datasets.
-
-The following overview goes into more detail on the pipeline, code structure, usage, and pretraining.
-
-<!-- ################ contents ################ -->
-# Contents
-
-  * [Quick start](#quick-start)
-  * [Tutorial](#tutorial)
-  * [Code structure](#code-structure)
-  * [Arguments](#arguments)
-  <!-- * [Pretraining](#pretraining) -->
-
-<!-- ################ quick start ################ -->
-
-# Quick Start
-Key files:
-
-- `main.py` : Entry point for processing.
-- `examples/preprocess_data.sh` : Example preprocessing launch (calls `main.py`).
-- `examples/pretrain_data.sh` : Example pretraining launch (calls `pretrain_retro.py`).
-
-Use `--retro-tasks` to move through the preprocessing pipeline.
-
-- Simplest setup (builds everything): `--retro-tasks build`
-- Alternatively, for tuning compute resources, run stages independently:
-  - Build retrieval database: `--retro-tasks db-build`
-  - Build search index: `--retro-tasks index-build`
-  - Query neighbors: `--retro-tasks pretraining-query-neighbors`
-
-Sample code flow:
-
-- `main.py` : Entry point (e.g., using `--retro-tasks X`).
-- `db/build.py` : Build retrieval database.
-- `index/build.py` : Build search index. Calls the following two files:
-  - `index/train.py` : Train index on subset of database.
-  - `index/add.py` : Add database chunks to index.
-- `pretraining/query.py` : Query pretraining samples for database neighbors (saved to disk and used during pretraining).
-
-<!-- ################ tutorial ################ -->
-
-# Tutorial
-
-In this tutorial example, we use the Wikipedia corpus to demonstrate how we build a retrieval database and index for this corpus, and then query the pretraining datasets for their neighbors.
-
-## Step 1: Prepare your retrieval text corpus
-
-The format of text corpus follows the same format as in Megatron training. See [data precessing](../../README.md#data-preprocessing) for more details on how to convert your json dataset into the mmap format.
-
-Assume we have the Wikipedia corpus in the following format:
-
-```
-<retrieval/db/path>/Wikipedia_shuf_text_document.bin
-<retrieval/db/path>/Wikipedia_shuf_text_document.idx
-```
-
-We note that the retrieval database can also be a blend of multiple text corpus.
-
-## Step 2: Build retrieval chunk database
-
-This *database* (stored as a 2-D array, NOT a relational database) consists of a list of chunks (traditionally length 64) extracted from the original GPT token dataset. This is simply a consecutive, non-overlapping chunking of the token dataset. Chunking only takes place within a document, and therefore the final chunk of each document has length: 1 <= chunk_length <= max_chunk_length.
-
-We discard chunks that would convert to an empty Bert sequence (rare case, happens ~1/100,000 chunks in our case), since we use Bert embeddings for building our index. Thus, the total number of chunks in the database will be slightly less than a naive calculation.
-
-Take the Wikipedia corpus as an example to build the retrieval chunk database:
-
-Prepare the following arguments and update our templates in [tools/retro/examples/preprocess_data.sh](examples/preprocess_data.sh):
-- `--retro-workdir`: The directory in which the preprocessing pipeline saves its datasets and configuration files. 
-  **This argument should remain consistent for a full pass through the pipeline, and for pretraining.**
-- `--data-path`: text corpus path to build retrieval database. In the case of Wikipedia corpus, it could be
-```bash
-WIK="${DATA_HOME}/Wikipedia_shuf_text_document"
-
-DATA_BLEND=" \
-  1 ${WIK} \
-"
-```
-- `--load`: bert path to load bert embedder
-- `--vocab-file` and `--retro-bert-vocab-file`: bert vocab file
-- `--retro-gpt-tokenizer-model`: gpt tokenizer model file
-
-Then launch the script:
-```bash
-bash tools/retro/examples/preprocess_data.sh db-build
-```
-
-After the `db-build` is finished, the output includes:
-- The launching args will be saved in your `<retro-workdir>/args.json` for the following steps. 
-- The retrieval chunk database will be saved in your `<retro-workdir>/db/` with your dataset information in `<retro-workdir>/db/indexed_dataset_infos.json`.  
-
-## Step 3: Build index for similarity search
-
-To match pretraining chunks to database chunks, a search index must be built to perform this querying. We use Faiss (https://github.com/facebookresearch/faiss) for training and building this index. Generally, the index is trained on a subset of all chunks in the database (specified via `--retro-index-ntrain`). After training, all chunks are added into the index, to be available during querying.
-
-Indexes only accept 1-D floating point vectors for training and adding, so each chunk must first be embedded before passing to the index for either training or adding. We use Bert embeddings for this purpose, and the embeddings are generated automatically within the pipeline.
-
-Take the Wikipedia corpus as an example to build the retrieval chunk database:
-
-```bash
-bash tools/retro/examples/preprocess_data.sh index-train
-```
-The `index-train` step is expected to take less than 4-hour on a single DGX-A100 node given the template index configuration. 
-To scale up for larger retrieval database, please carefully tune the faiss hyper-parameters specified in `--retro-index-str`. Please refer to [Faiss](https://github.com/facebookresearch/faiss/wiki/The-index-factory) to learn more about the index configuration.  
-
-After the index is trained, the centroids, HNSW graph, and product quantizer is determined. However, the index is still empty, as there is no chunk added.
-
-Take the example of the Wikipedia corpus, with the default template, the output of `index-train` includes:
-- The embedded Bert embeddings of the sampled chunks for `index-train` is saved in `<retro-workdir>/index/train_emb/`.  
-- The empty index is saved in `<retro-workdir>/index/faiss-par-add/OPQ32_64,IVF65536_HNSW8,PQ32/empty_0.970.faissindex`.
-
-Then we add all chunks in the retrieval database into the index so that we perform fast query over the whole retrieval database:
-```bash
-bash tools/retro/examples/preprocess_data.sh index-add
-```
-
-We note that this step can be time-consuming as it will go through the whole retrieval database, embed chunk tokens  to BERT embeddings, and add them into the index. Please make sure you successfully add the whole retrieval database before moving on to the next stage.
-
-*In case your job is interrupted in the middle, you can just run the script again, and it will automatically skip the chunks that have been added into the index and start from the chunk where it is interrupted.*
-
-
-Following the Wikipedia configuration, an example output of the step `index-add` includes:
-- The index with retrieval data chunks added is saved in `<retro-workdir>/index/faiss-par-add/OPQ32_64,IVF65536_HNSW8,PQ32/added_0.970_0.950.faissindex`, which can be used to query the neighbors for pretraining.
-
-## Step 4: Query pretraining neighbors
-
-To ensure fast Retro pretraining, the database neighbors for pretraining samples are pre-computed and saved to disk, for efficient access within the Retro dataset. In this stage, the pretraining datasets (training, validation, and test) are iterated, each sample is broken into chunks, and the chunks are used for querying the index. Similar to when building the index, each chunk is embedded (via Bert) before querying the index.
-
-The saved neighbors are labeled with unique dataset properties (i.e., seed, sequence length, number of samples, etc.) to ensure the neighbors generated during preprocessing match the neighbors requested during pretraining. Please also make sure the pretraining configuration is the same as this step so that the neighbors are aligned.
-
-There are query-time hyper-parameters that can be tuned to improve the quality of the neighbors. These are specified in `RETRO_QUERY_EF_SEARCH` and `RETRO_QUERY_NPROBE`. The most important parameter is `RETRO_QUERY_NPROBE`, which controls the number of clusters to search during querying. This parameter can be tuned to improve the quality of the neighbors, but will also increase the query time. 
-We recommend following the tutorial of [faiss](https://github.com/facebookresearch/faiss/wiki/Index-IO,-cloning-and-hyper-parameter-tuning) to tune the hyper-parameters for your own retrieval database. 
-
-Take the Wikipedia corpus as an example to query the neighbors in the retrieval database:
-
-```bash
-bash tools/retro/examples/preprocess_data.sh query-pretraining-neighbors
-```
-
-The output of `query-pretraining-neighbors` on the Wikipedia corpus includes:
-- `<retro-workdir>/wiki/query/train_855ab50e05151610301e2a74c4030fbc`, which contains the pre-retrieved neighbors for the pretraining dataset. 
-- `<retro-workdir>/wiki/query/valid_40bc7330318d64accec28e1e63c59bad`, which contains the pre-retrieved neighbors for the validation set of the pretraining corpus.
-
-## Step 5: Visualization of retrieval neighbors
-
-We also provide cli tools to help visualize and inspect the quality of your retrieved neighbors. 
-
-To use the CLI, open a Python terminal via the `python` command, and then load a Retro workdir with the following:
-
-```
-from tools.retro.cli import retro
-retro.init("/path/to/retro/workdir")
-```
-
-This initializes Megatron, and prepares the Retro data for inspection. We also print out some example commands to help you get familiar with the command lines.   
-
-An example output for the Wikipedia Corpus:
-
-```text
-setting number of micro-batches to constant 32
-> building BertWordPieceLowerCase tokenizer ...
-> initializing torch distributed ...
-> initialized tensor model parallel with size 1
-> initialized pipeline model parallel with size 1
-> compiling dataset index builder ...
-...
-...
- > sample ratios:
-   dataset 0, input: 1, achieved: 1
-> size of blendable dataset: 201000 samples
-> elapsed time for building blendable dataset indices: 0.00 (sec)
-> building indices for blendable datasets ...
- > sample ratios:
-   dataset 0, input: 1, achieved: 1
-> size of blendable dataset: 12864 samples
-> finished creating pretrained GPT datasets ...
-
-+++++++++++++++++++++++++++++++++++++++++++++++++++
-examples ... [ *note*: 'db' = chunk db; 'pt' = pretraining corpus. ]
-+++++++++++++++++++++++++++++++++++++++++++++++++++
-
-~~~~ indexed datasets ~~~~
-retro.get_db_num_indexed_datasets() : 1
-retro.get_db_indexed_dataset_infos() :
-  [(1.000000, Wikipedia_shuf_text_document)]
-
-~~~~ counts ~~~~
-retro.get_db_num_chunks : 68104992.
-
-retro.get_pt_num_samples('train') : 201000.
-retro.get_pt_num_samples('valid') : 12864.
-retro.get_pt_num_chunks('train') : 1608000.
-retro.get_pt_num_chunks('valid') : 102912.
-
-~~~~ tokens, text ~~~~
-retro.get_db_chunk_gpt(chunk_id) : [46809, 218340, 716, 647, ... , 251525, 872, 692, 4042]
-retro.get_db_chunk_bert(chunk_id) : [10680, 16216, 4313, 1745 ... , 8117, 1007, 1012, 1997]
-retro.get_db_chunk_text(chunk_id) : Jonas Geirnaert\n\nJonas  ... ort Flatlife (11 min). Of
-retro.get_db_chunk_and_continuation_text(chunk_id) :
-  ['Jonas Geirnaert  Jonas Ge ... ort Flatlife (11 min). Of',
-   'the copy he sent in for s ... abet, clearly has one. On']
-
-retro.get_pt_sample('train', sample_id) :
-  {
-    'dataset_idx' : 0
-    'text' : [   676     14  40656 184 ... 4\n    276  17361 251542]
-    'doc_ids' : [1246422 1596948 2403969]
-    'neighbor_chunks' : [[[  657380   657381]\n   ... \n  [34108760 34108761]]]
-    'neighbor_tokens' : [[[   276   9596 251511 . ... .    889    646   1723]]]
-  }
-
-(e.g., sample = retro.get_pt_sample(...))
-
-  sample['text'].shape : (513,)
-  sample['neighbor_tokens'].shape : (8, 20, 128)
-  sample['text'] : [   676     14  40656 184 ... 4\n    276  17361 251542]
-  sample['neighbor_tokens'][17][1] : [    14     14  30291   1 ... 682    328    379 251527]
-  retro.gpt_to_text(sample['text']) : also\nLatgalians (modern) ... ission criticised the AVN
-  retro.gpt_to_text(sample['neighbor_tokens']) : \n\nHis second marriage o ... Augusta Eardley-Wilmot (2
-+++++++++++++++++++++++++++++++++++++++++++++++++++
-```
-
-We can also directly call the function `retro.print_neighbor_texts(sample_id, chunk_id)` to inspect the retrieval neighbors for a specific sample and chunk within the pretraining corpus. For example,  
-
-```text
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-PRETRAINING CHUNK:
-  - also\nLatgalians (modern)\n\nReferences\n\nCategory:Defunct political parti ... e.\n\nAbout \nThe company was established established in 1997. It is listed
-NEIGHBOR_CHUNKS:
-  - the sides.\n\nNotes\n\nReferences\n\nCategory:Obaku Zen\n*\nCategory:Japane ... 2, 2008. It was founded by Anand Jagannathan, CEO of parent company Kriyari
-  - 2007).\n\nSee also\n Satellite Communications\n Tonga\n\nReferences\n\nExte ... y Procter & Gamble (P&G) in 1985 in order for P&G to compete in the "beauty
-  - Japan\nCategory:Fish of Russia\nCategory:Fish described in 1845 Mareco Inde ... lic Opinion (WAPOR)\n European Society for Opinion and Marketing Research (
-  - The current director of the company is Albert Bosch.\n\nSee also\n Coupon\n ...  some articles in Basque. Deia is the main product of the Editorial Iparrag
-  - A.Ş have been traded on the Istanbul Stock Exchange since 2000.\n\nReferenc ... with stores in California, New York City, and London.\n\nHistory \nSnapette
-  - \nCategory:Hawaiian mythology\nCategory:Hawaiian religion\nCategory:Religio ... crative state contracts. In 2008 Prokom became a part of the Asseco capital
-  - , and the Baltic countries, as well as an online store.\n\nReferences\n\nEx ... nd are involved in intracellular trafficking. This protein does not contain
-  - juice producer\nFood industry of Russia\n\nReferences\n\nExternal links\nWi ... panies formerly listed on the New York Stock Exchange General Grant's March
-  - is in private ownership.\n\nReferences\n\nExternal links\n\nCategory:Online ... ten and directed by Brent Hodge. The film stars Aubrey Plaza, Molly Hawkey,
-  - company's display technology to manufacture and sell display-only engines.\ ... for a group of naval vessels (a division in naval usage).\n\nUsage\n Russia
-  - .\n\nCarrols also operated a chain of outlets in neighbouring Estonia from  ... rama film directed by Raajeev Walia. It is produced by Aman Mehta and Bijal
-  - \n\nExternal links\nHightail website\nThe Next Web on YouSendIt rebrand to  ... eptember 2014, sitting mainly in the criminal division of that court.\n\nBe
-  - American television seasons\nCategory:2014 American television seasons\nCat ...  Canada and larger European cities.\n\nIn 2010, advertising in New Zealand,
-  - .\n\nNotes\n\nCategory:Trade unions\nCategory:Industrial Workers of the Wor ... x people, some of whom may have been working on a part-time basis. Its head
-  - \n List of podcasting companies\n\nReferences\n\nExternal links\n \n\nCateg ... ct.\n\nCategory:Populated places in the Ashanti Region Nkeirouka Ezekh\n\nN
-  - \n\nReferences\n\nExternal links\n ADESE official website\n\nCategory:Compa ...  State Street, and UBS Warburg. Its first CEO was Ian M. Drachman. The firm
-  - Hotel\n Sulake Corporation\n Sulake Press Room\n Habbo Hotel - Blog\n\nCate ... l: 김진태; born December 19, 1980), better known by his stage name Verbal Jint
-  - hockey player\n Ruutu.fi, a Finnish television streaming service operated b ...  from the bottom, a BDSM term\n Topping cycle, a cycle used in power plants
-  - of Surakarta\nCategory:Indonesian names\nCategory:Indonesian families\nCate ... mber 13, 2013 in Izhevsk on Universitetskaya Street (later it was given the
-  - facilities are also in Ankara and the company HQ is in Istanbul.\n\nReferen ... is currently a World Wide Web Consortium Working Draft.\n\nSee also\n Voice
-```
-
-The code snippet for the above example is also equivalent to
-```python
-tokens = retro.get_pt_sample('train', 0)
-for token_ids in tokens["neighbor_tokens"][0]:
-    print("- %s" % (retro.gpt_to_text(token_ids)))
-    print("-" * 20)
-```
-<!-- ################ code structure ################ -->
-# Code structure
-
-### `tools/retro/main.py`
-
-This is the main entry point for Retro preprocessing. Call `main.py --help` to see arguments. Additionally, some Retro arguments are in Megatron's core arguments, so also see `add_retro_args()` section of `megatron/arguments.py` for additional arguments. Two of the most important arguments to customize are `--retro-workdir` and `--retro-tasks`.
-
-- **`--retro-workdir`** : Set the directory in which the preprocessing pipeline saves its datasets and configuration files. This argument should remain consistent for a full pass through the pipeline, and for pretraining.
-
-- **`--retro-tasks`** : Set the stages of preprocessing to perform. As mentioned previously, the three high-level stages are: 1) build retrieval database, 2) build search index, and 3) query pretraining neighbors. `--retro-tasks` can be used to either run the full pipeline, or run each of these stages in isolation. The latter case is useful for tuning compute resources for each stage. For example, index training utilizes GPUs and requires relatively less time, while querying neighbors uses the CPU and is a relatively slow process. Example tasks include:
-
-  - **`--retro-tasks build`** : Run entire preprocessing pipeline.
-  - **`--retro-tasks db-build`** : Build retrieval database.
-  - **`--retro-tasks index-build`** : Train and build search index.
-  - **`--retro-tasks pretraining-query-neighbors`** : Query pretraining neighbors.
-
-Multiple tasks can be specified by separating with commas (e.g., `--retro-tasks db-build,index-build`). Additionally, various 'miscellaneous' tasks are currently including, primarily for validating data for each stage; these task names can be seen in `main.py`.
-
-### `tools/retro/examples`
-
-Example scripts for setting arguments and launch Retro preprocessing. The key files here are:
-
-- **`preprocess_data.sh`** : Example launch script for preprocessing retro data.
-- **`pretrain_model.sh`** : Example launch script for pretraining a retro model.
-
-### `tools/retro/db`
-
-Build the retrieval chunk database. The key files here are:
-
-- **`build.py`** : Entry point for building the database. This code is responsible for iterating the input datasets (i.e., `--data-path`), parsing each dataset into consecutive chunks, checking for empty Bert (Wordpiece) conversions, and storing this information to disk. Two databases are created: 1) the retrieval database, and 2) a sampled database used for training the search index.
-- **`dataset.py`** : Defines database class, for iterating or accessing chunks in the database. Each chunk contains its tokens, Bert conversion length, and dataset index.
-
-Input data:
-
-<!-- - Token datasets, as generated by `tools/preprocess_data.py`. Each dataset should include a `.bin` and `.idx` file. Multiple datasets can be specified by using a blended configuration (see `--data-path` in `megatron/arguments.py`). -->
-- Token datasets, as loaded by `gpt_dataset.py`. Multiple datasets can be specified by using a blended configuration (see `--data-path` in `megatron/arguments.py`).
-
-Output data:
-
-- **`<RETRO_WORKDIR>/db/merged/train.hdf5`** : The main retrieval database. (*Database* here is used to denote a list of indexed chunks, rather than a *relational database*.) The chunks in this database are added to the search index, and are used for retrieval during pretraining. This file contains a single dataset `'chunks'`, which contains 5 columns:
-
-  - `dataset_idx` : Dataset index, from list of blended indexed datasets.
-  - `document_idx` : Document index within dataset.
-  - `chunk_start_idx` : Chunk's starting token index within document.
-  - `chunk_end_idx` : Chunk's ending token index (exclusive) within document.
-  - `bert_chunk_length` : Length of Bert token sequence, after converting from GPT.
-
-- **`<RETRO_WORKDIR>/db/merged/sampled.hdf5`** : Subset of training database that is used for training the search index. This file has the same structure as detailed above. In general, this database is significanly smaller than the `train.hdf5` database, since the search index only needs a relatively small number of samples to understand the data's structure. After training, all chunks in the main database (`train.hdf5`) are *added* to the search index.
-
-### `tools/retro/index`
-
-Build the search index. The key files here are:
-
-- `build.py` : Entry point for building the search index. First, the index is trained on the sampled chunk database (see above) by calling `train.py`, and then all chunks for the full database are added to the index by calling `add.py`. Note that training requires first embedding (using Bert) all chunks (a parallel operation), and then loading these embeddings and training the index (a sequential operation), so it's best to change one's compute setup after all chunks have been embedded and saved to disk.
-- `indexes/faiss_base.py` : Wrapper class for building a Faiss index, following the standard `train()` and `add()` operations.
-- `indexes/faiss_par_add.py` : Similar to above, except it uses an embarrassingly parallel (multi-node, multi-process) `add()` operation. Vectors are first added to separate index copies, and then merged together.
-
-Input data:
-
-- **`<RETRO_WORKDIR>/db/merged/sampled.hdf5`** : Chunks used for training the search index.
-- **`<RETRO_WORKDIR>/db/merged/train.hdf5`** : Chunks used for adding to the *trained* search index.
-
-Output data:
-
-- **`<RETRO_WORKDIR>/index/<RETRO_INDEX_TYPE>/<RETRO_INDEX_STR>/added.faissindex`** : The final index, which has been trained and has had all database chunks added to it. This index is ready for querying neighbors. Here, `RETRO_INDEX_TYPE` and `RETRO_INDEX_STR` correspond to the same-name arguments `--retro-index-type` (e.g., `faiss-par-add`) and `--retro-index-str` (e.g., `OPQ32_256,IVF4194304_HNSW32,PQ32`).
-- **`<RETRO_WORKDIR>/index/<RETRO_INDEX_TYPE>/<RETRO_INDEX_STR>/empty.faissindex`** : Generally can be discarded once `added.faissindex` has been built, but this file contains the *post-training*, *pre-adding* index. Useful for debugging or building other indexes.
-
-### `tools/retro/pretraining`
-
-Query the pretraining datasets (training, validation, test) for their neighbors within the database. Neighbors are queried during preprocessing -- rather than during pretraining -- because querying is a fairly slow operation, so it would be a bottleneck if performed during pretraining. Queried neighbors are tagged with their unique identifying information (e.g., `train_indexmap_27662746ns_2048sl_1234s`), so as to avoid incorrect references during pretraining. The key files here are:
-
-- **`query.py`** : Entry point for querying. The pretraining datasets are iterated, and each chunk within each sample is queried using the search index. These neighbors are filtered by discarding any database chunks that fall within the same document as any chunk within a pretraining sample.
-- **`chunk_dataset.py`** : This creates an iterable 'chunk' dataset form of a pretraining dataset. This is just a light wrapper, but makes it easier to deterministically iterate and assign IDs to each chunk in a sample dataset.
-- **`retro_dataset.py`** : The Retro dataset used for pretraining (not used in preprocessing). Each sample returns the sample tokens, along with neighbor tokens for each chunk within the sample.
-
-Input data:
-
-- Token datasets, as loaded by `gpt_dataset.py`.
-- **`<RETRO_WORKDIR>/index/<RETRO_INDEX_TYPE>/<RETRO_INDEX_STR>/added.faissindex`** : The trained index, with all database chunks added to it (see previous section for details).
-
-Output data:
-
-- **`<RETRO_WORKDIR>/{train,valid,test}_XXns_YYsl_ZZs/WW.hdf5`** : These directories/files contain the indexes of neighbors for each chunk within each sample of the pretraining datasets. Each directory (e.g., `train_indexmap_2047435ns_2048sl_1234s`) contains a list of HDF5 files (e.g., one file might be called `0075700000-0075800000.hdf5`). Each HDF5 file contains a consecutive subset of neighbor IDs for a given chunk, for indexing into the main retrieval database. All HDF5 files taken together within a given directory, represent the entire set of neighbors for a dataset. The size of these HDF5 files is determined by the argument `--retro-block-size`. The `XX`, `YY`, `ZZ`, `WW` notation above denotes the dataset properties that are used for uniquely tagging the neighbor files, to ensure compatibility during model pretraining. These neighbor files are ultimated used by `retro_dataset.py` during pretraining, for building Retro samples.
-
-### `tools/retro/cli`
-
-Inspect preprocessed data. To use the CLI, open a Python terminal via the `python` command, and then load a Retro workdir with the following:
-
-```
-from tools.retro.cli import retro
-retro.init("/path/to/retro/workdir")
-```
-
-This initializes Megatron, and prepares the Retro data for inspection. See the printed usage for available functions. Several routines are included for viewing data in the retrieval database and viewing pretraining samples and neighbors. For example:
-
-```python
-retro.get_db_num_indexed_datasets() # 15
-retro.get_db_chunk_text(92874113) # 'research project at ...  and philosophy'
-retro.get_pt_sample('train', 62005) # '[16084, 26158, 25387 ..., 6898, 9568]'
-```
-
-Most methods within the CLI are prefixed to denote the data being inspected:
-
-- **'db'** : Retrieval database (i.e., chunk tokens, document IDs, and dataset IDs)
-- **'pt'** : Pretraining datasets (i.e., sample tokens and neighbor tokens)
-
-### `tools/retro/utils.py`
-
-A collection of utility methods. Most importantly, this contains:
-
-- **`def get_gpt_tokenizer()`** : Get the GPT tokenizer.
-- **`def get_bert_tokenizer()`** : Get the Bert tokenizer.
-- **`class GPTToTextDataset`** : Wrapper class that converts GPT (BPE) samples to raw text.
-
-### `tools/bert_embedding`
-
-Generate Bert embeddings. The main files here are:
-
-- **`embed.py`** : Entry point for generating embeddings, and contains the two main embedding classes, `BertEmbedder` and `DiskDataParallelBertEmbedder` (more below). This file contains code for generating Megatron embeddings, while the file below contains code for Huggingface embeddings.
-- **`huggingface.py`** : Used by `embed.py` when the embedder is configured (see below) to output Huggingface embeddings.
-- **`dataset.py`** : Wrapper class for converting a raw-text dataset to Bert (Wordpiece) tokens.
-
-The Bert embeddings can be configured along two axes. The first axis is the output type:
-
-- **`class BertEmbedder`** : This class takes a raw-text dataset as input, generates its embeddings, and returns a Numpy array. The main functions are `embed_text_dataset` (accepts a raw-text dataset) and `embed_text` (accepts a string).
-- **`class DiskDataParallelBertEmbedder`** : This class wraps `BertEmbedder`, and rather than returning a Numpy array, it saves the embeddings to disk. Additionally, this class automatically splits data across data parallel ranks (using interleaving), and also processes data in a specified `block_size` (e.g., 1,000,000).
-
-The second axis is the type of embedding model to use, controlled by the argument `--bert-embedder-type`:
-
-- **`--bert-embedder-type megatron`** : Use Megatron's Bert model. The specific model used is dependent on the loaded checkpoint, vocab file, and tokenizer.
-- **`--bert-embedder-type huggingface`** : Use Huggingface's `bert-large-cased`. (*Note*: Huggingface's inclusion is likely to be deprecated; and there is no ability to configure cased/uncased.)
-
-### Pretraining
-
-- **`pretrain_retro.py`** : Launch script for pretraining Retro. Similar to `pretrain_gpt.py`, except this script handles loading neighbor tokens and setting up the neighbor attention mask.
-<!-- - `megatron/data/gpt_dataset.py` : ? -->
-- **`megatron/model/retro_transformer.py`** : Implementation of Retro model, including the main transformer, the retrieval encoder, and chunked cross-attention layers. Note that currently, `retro_transformer.py` contains several classes that are nearly identical to `transformer.py`, except for 1 or 2 lines, due to code changes that are yet to be integrated.
-- **`tools/retro/pretraining/retro_dataset.py`** : The Retro dataset used for pretraining (not used in preprocessing). Each sample returns the sample tokens, along with neighbor tokens for each chunk within the sample.
-
-
-<!-- ################ arguments ################ -->
-# Arguments
-
-See `tools/retro/main.py`'s `add_retro_args()` and `megatron/arguments.py`'s `_add_retro_args()` for details and descriptions. Here we list some particularly important arguments:
-
-- `--retro-workdir` : Mentioned previously, this argument determines the directory in which a set of Retro data is stored (during preprocessing) and loaded (during pretraining). Any change in this directory during preprocessing may result in preprocessing starting over from scratch, and any change before pretraining will result in pretraining throwing an error.
-- Preprocessing
-  - `--retro-gpt-chunk-length` : Retro chunk length (e.g., 64 in original paper).
-  - `--retro-tasks` : Comma-separated list of preprocessing tasks. Generally, the `build` task is the simplest way to run the preprocessing pipeline. For finer control, individual stages can be run by using tasks (in order): `db-build`, `index-build`, and `pretraining-query-neighbors`.
-  - `--retro-index-str` : Faiss index string that defines the index configuration. This will vary based on data size, compute/disk setup, and user needs. For example, this string looks something like `IVF262144_HNSW32,Flat` or `OPQ32_256,IVF4194304_HNSW32,PQ32`.
-- Pretraining
-  - `--retro-add-retriever` : Must be used to select Retro model.
-  - `--retro-num-neighbors` : Number of neighbors to retrieve from the retrieval database (defaults to 2).
-  - `--retro-num-retrieved-chunks` : For each neighbor, the number consecutive chunks to retrieve, including the initial neighbor (defaults to 2).
-  - `--retro-attention-gate` : Gated mechanism to incorporate information of cross attention from retrieved neighbor  (defaults to 1 during pretraining).
-
-<!-- ################ pretraining ################ -->
-<!-- # Pretraining -->
-<!-- - New retro args in arguments.py (add_retro_args). -->
-<!-- - Most important arg is `--retro-add-retriever`. -->
diff --git a/tools/retro/cli/__init__.py b/tools/retro/cli/__init__.py
deleted file mode 100644
index 2531017a28b..00000000000
--- a/tools/retro/cli/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-from .cli import retro
diff --git a/tools/retro/cli/__main__.py b/tools/retro/cli/__main__.py
deleted file mode 100644
index 37d096a9538..00000000000
--- a/tools/retro/cli/__main__.py
+++ /dev/null
@@ -1,9 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-import os
-
-from . import retro
-
-
-if __name__ == "__main__":
-    retro.init(os.environ["RETRO_PROJECT_DIR"])
diff --git a/tools/retro/cli/cli.py b/tools/retro/cli/cli.py
deleted file mode 100644
index a5d953d2f7a..00000000000
--- a/tools/retro/cli/cli.py
+++ /dev/null
@@ -1,301 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-import json
-import numpy as np
-import os
-import typing as T
-from types import SimpleNamespace
-
-from megatron.training.arguments import load_retro_config, parse_args, validate_args
-from megatron.core.datasets.retro.db.dataset import DBDataset
-from megatron.core.datasets.retro.db.utils import (
-    get_indexed_dataset_infos as get_db_indexed_dataset_infos,
-    get_merged_train_dataset as get_db_dataset,
-)
-from megatron.core.datasets.retro.query.retro_dataset import get_retro_datasets, RetroDataset
-from megatron.training.global_vars import set_global_variables
-from megatron.training.training import build_train_valid_test_datasets, update_train_iters
-from pretrain_retro import train_valid_test_datasets_provider
-from tools.retro.preprocess_data import get_tokenizers
-
-
-def shorten_str(s: str, n: int) -> str:
-    s = "\\n".join(s.splitlines())
-    return s if len(s) <= n else "%s ... %s" % (s[: n // 2], s[-n // 2 :])
-
-
-class retro:
-
-    config = None
-
-    ##############################################
-    # initialize.
-    ##############################################
-
-    @classmethod
-    def init(cls, project_dir: str) -> None:
-        '''Initialize Megatron, tokenizers, and datasets.'''
-
-        # Megatron args.
-        args = parse_args(extra_args_provider=None, ignore_unknown_args=False)
-        args.retro_project_dir = project_dir
-        args.micro_batch_size = 1
-        args.num_layers = 1
-        args.hidden_size = 1
-        args.num_attention_heads = 1
-        args.async_tensor_model_parallel_allreduce = False
-        args.retro_add_retriever = True # for building RetroDataset
-        validate_args(args)
-        set_global_variables(args)
-        update_train_iters(args)
-
-        # Retro config.
-        cls.config = load_retro_config(project_dir)
-        cls.config.retro_project_dir = project_dir
-        cls.config.retro_tokenizers = get_tokenizers(cls.config)
-
-        # Chunk database dataset.
-        cls.db_indexed_dataset_infos = get_db_indexed_dataset_infos(project_dir)
-        cls.db_dataset = get_db_dataset(project_dir,
-                                        cls.config.retro_gpt_chunk_length,
-                                        cls.config.retro_tokenizers.gpt.eod)
-
-        # Pretraining datasets.
-        pt_train_ds, pt_valid_ds, pt_test_ds = build_train_valid_test_datasets(
-            train_valid_test_datasets_provider)
-        cls.pt_datasets = SimpleNamespace(
-            train=pt_train_ds,
-            valid=pt_valid_ds,
-            test=pt_test_ds,
-        )
-
-        # Print usage.
-        cls.print_usage()
-
-    ##############################################
-    # utils.
-    ##############################################
-
-    @classmethod
-    def gpt_to_text(cls, token_ids: np.ndarray) -> str:
-        '''GPT tokens to text.'''
-        return cls.config.retro_tokenizers.gpt.detokenize(
-            token_ids.tolist() if isinstance(token_ids, np.ndarray) else token_ids
-        )
-
-    @classmethod
-    def text_to_bert(cls, text: str) -> np.ndarray:
-        '''Text to Bert tokens.'''
-        return cls.config.retro_tokenizers.bert.tokenize(text)
-
-    ##############################################
-    # chunk db.
-    ##############################################
-
-    @classmethod
-    def get_db_num_indexed_datasets(cls) -> int:
-        '''Number of indexed datasets within blended dataset.'''
-        return len(cls.db_indexed_dataset_infos)
-
-    @classmethod
-    def get_db_indexed_dataset_infos(cls) -> T.List[T.Tuple[float, str]]:
-        '''Dataset infos, including number of training & sampled sets.'''
-        return [(info["ratio"], info["prefix"]) for info in cls.db_indexed_dataset_infos]
-
-    @classmethod
-    def get_db_dataset(cls) -> DBDataset:
-        return cls.db_dataset
-
-    @classmethod
-    def get_db_num_chunks(cls) -> int:
-        '''Number of DB chunks.'''
-        return len(cls.get_db_dataset())
-
-    @classmethod
-    def get_db_chunk_gpt(cls, idx: int) -> T.List[int]:
-        '''Get DB chunk as GPT token ids.'''
-        return cls.get_db_dataset()[idx]["text"].tolist()
-
-    @classmethod
-    def get_db_chunk_bert(cls, idx: int) -> T.List[int]:
-        '''Get DB chunk as Bert token ids.'''
-        return cls.text_to_bert(cls.get_db_chunk_text(idx))
-
-    @classmethod
-    def get_db_chunk_text(cls, idx: int) -> str:
-        '''Get DB chunk as text.'''
-        return cls.gpt_to_text(cls.get_db_chunk_gpt(idx))
-
-    @classmethod
-    def get_db_chunk_and_continuation_text(cls, idx: int) -> T.List[str]:
-        '''Get DB chunk along with continuation, as text.'''
-
-        # Modulus used here to match original implementation (i.e., last
-        # chunks continuation wraps around to first chunk).
-        return [
-            cls.get_db_chunk_text(idx),
-            cls.get_db_chunk_text((idx + 1) % len(cls.get_db_dataset())),
-        ]
-
-    ##############################################
-    # pretraining corpus.
-    ##############################################
-
-    @classmethod
-    def get_pt_num_samples_and_chunks(cls, data_key: str) -> T.Tuple[int, int]:
-        '''Number of samples & chunks (e.g., 32*n_samples) in corpus.'''
-        assert hasattr(cls.pt_datasets, data_key), (
-            "pretraining set '%s' not found (choices: %s)."
-            % (data_key, ", ".join(vars(cls.pt_datasets).keys()))
-        )
-        chunk_dataset = getattr(cls.pt_datasets, data_key).chunk_dataset
-        return (
-            len(chunk_dataset.sample_dataset),
-            len(chunk_dataset),
-        )
-
-    @classmethod
-    def get_pt_num_samples(cls, data_key: str) -> int:
-        '''Number of pretraining samples.'''
-        return cls.get_pt_num_samples_and_chunks(data_key)[0]
-
-    @classmethod
-    def get_pt_num_chunks(cls, data_key: str) -> int:
-        '''Number of pretraining chunks (e.g., 32*n_samples).'''
-        return cls.get_pt_num_samples_and_chunks(data_key)[1]
-
-    @classmethod
-    def get_pt_dataset(cls, data_key: str) -> RetroDataset:
-        return getattr(cls.pt_datasets, data_key)
-
-    @classmethod
-    def get_pt_sample(cls, data_key: str, idx: int) -> dict:
-        return getattr(cls.pt_datasets, data_key)[idx]
-
-    @classmethod
-    def get_neighbor_tokens(cls, sample_id: int, chunk_id: int, data_key: str="train") -> T.Optional[dict]:
-        try:
-            sample = cls.get_pt_sample(data_key, sample_id)
-            sample_token_ids = sample["text"]
-            chunk_length = cls.args.retro_gpt_chunk_length
-            chunk_start_idx = chunk_id * chunk_length
-            chunk_end_idx = min(sample_token_ids.shape[0], chunk_start_idx + chunk_length)
-            chunk_token_ids = sample_token_ids[chunk_start_idx:chunk_end_idx]
-            neighbor_token_ids = sample["neighbor_tokens"][chunk_id]
-            return {
-                "chunk_tokens": chunk_token_ids,
-                "neighbor_tokens": neighbor_token_ids,
-            }
-        except Exception:
-            return None
-
-    @classmethod
-    def print_neighbor_texts(cls, sample_id: int, chunk_id: int, data_key: str="train") -> None:
-        tokens: dict = cls.get_neighbor_tokens(sample_id, chunk_id, data_key)
-        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
-        try:
-            print("PRETRAINING CHUNK:")
-            print("  - %s" % shorten_str(cls.gpt_to_text(tokens["chunk_tokens"]), 150))
-            print("NEIGHBOR_CHUNKS:")
-            for token_ids in tokens["neighbor_tokens"]:
-                print("  - %s" % shorten_str(cls.gpt_to_text(token_ids), 150))
-        except Exception:
-            print("<no neighbors for sample %d>" % sample_id)
-
-    ##############################################
-    # usage.
-    ##############################################
-
-    @classmethod
-    def print_usage(cls) -> None:
-        '''Print usage.'''
-
-        print()
-        print("+++++++++++++++++++++++++++++++++++++++++++++++++++")
-        print("examples ... [ *note*: 'db' = chunk db; 'pt' = pretraining corpus. ]")
-        print("+++++++++++++++++++++++++++++++++++++++++++++++++++")
-
-        print()
-        print("~~~~ indexed datasets ~~~~")
-        print("retro.get_db_num_indexed_datasets() : %s" % cls.get_db_num_indexed_datasets())
-        print("retro.get_db_indexed_dataset_infos() :")
-        for i, (ratio, prefix) in enumerate(cls.get_db_indexed_dataset_infos()):
-            print(
-                "  %s(%f, %s)%s"
-                % (
-                    "[" if i == 0 else " ",
-                    ratio,
-                    prefix,
-                    "]" if i == len(cls.db_indexed_dataset_infos) - 1 else ",",
-                )
-            )
-
-        print()
-        print("~~~~ counts ~~~~")
-        print("retro.get_db_num_chunks : %d." % cls.get_db_num_chunks())
-
-        print()
-        for sq_key in ("sample", "chunk"):
-            for data_key in ("train", "valid"):  # test?
-                print(
-                    "retro.get_pt_num_%ss('%s') : %d."
-                    % (sq_key, data_key, getattr(cls, f"get_pt_num_{sq_key}s")(data_key))
-                )
-
-        print()
-        print("~~~~ tokens, text ~~~~")
-        print(
-            "retro.get_db_chunk_gpt(chunk_id) : %s"
-            % shorten_str(str(retro.get_db_chunk_gpt(0)), 50)
-        )
-        print(
-            "retro.get_db_chunk_bert(chunk_id) : %s"
-            % shorten_str(str(retro.get_db_chunk_bert(0)), 50)
-        )
-        print(
-            "retro.get_db_chunk_text(chunk_id) : %s"
-            % shorten_str(retro.get_db_chunk_text(0).strip(), 50)
-        )
-        print("retro.get_db_chunk_and_continuation_text(chunk_id) :")
-        for i, t in enumerate(retro.get_db_chunk_and_continuation_text(0)):
-            print(
-                "  %s'%s'%s"
-                % (
-                    "[" if i == 0 else " ",
-                    shorten_str(t.strip().replace("\n", " "), 50),
-                    "]" if i == 1 else ",",
-                )
-            )
-
-        sample = cls.get_pt_sample("train", 0)
-        sample_chunk_id = sample["neighbor_tokens"].shape[0] // 2
-        sample_neighbor_id = 0
-        print()
-        print("retro.get_pt_sample('train', sample_id) :")
-        print("  {")
-        for k, v in sample.items():
-            print("    '%s' : %s" % (k, shorten_str(str(v), 50)))
-        print("  }")
-
-        print()
-        print("(e.g., sample = retro.get_pt_sample(...))")
-        print()
-        print("  sample['text'].shape : %s" % str(sample["text"].shape))
-        print("  sample['neighbor_tokens'].shape : %s" % str(sample["neighbor_tokens"].shape))
-        print("  sample['text'] : %s" % shorten_str(str(sample["text"]), 50))
-        print(
-            "  sample['neighbor_tokens'][17][1] : %s"
-            % shorten_str(str(sample["neighbor_tokens"][sample_chunk_id][sample_neighbor_id]), 50)
-        )
-        print(
-            "  retro.gpt_to_text(sample['text']) : %s"
-            % shorten_str(cls.gpt_to_text(sample["text"]), 50)
-        )
-        print(
-            "  retro.gpt_to_text(sample['neighbor_tokens']) : %s"
-            % shorten_str(
-                cls.gpt_to_text(sample["neighbor_tokens"][sample_chunk_id][sample_neighbor_id]), 50
-            )
-        )
-
-        print("+++++++++++++++++++++++++++++++++++++++++++++++++++")
diff --git a/tools/retro/config_utils.py b/tools/retro/config_utils.py
deleted file mode 100644
index 00676c66ffd..00000000000
--- a/tools/retro/config_utils.py
+++ /dev/null
@@ -1,632 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-
-"""Config utils."""
-
-import argparse
-from collections import namedtuple, OrderedDict
-import dataclasses
-import enum
-import inspect
-import os
-import re
-import types
-import typing as T
-
-
-PARAM_KEYWORDS = {
-    "param",
-    "parameter",
-    "arg",
-    "argument",
-    "attribute",
-    "key",
-    "keyword",
-}
-RAISES_KEYWORDS = {"raises", "raise", "except", "exception"}
-DEPRECATION_KEYWORDS = {"deprecation", "deprecated"}
-RETURNS_KEYWORDS = {"return", "returns"}
-YIELDS_KEYWORDS = {"yield", "yields"}
-EXAMPLES_KEYWORDS = {"example", "examples"}
-
-
-class ParseError(RuntimeError):
-    """Base class for all parsing related errors."""
-
-
-class DocstringStyle(enum.Enum):
-    """Docstring style."""
-
-    REST = 1
-    GOOGLE = 2
-    NUMPYDOC = 3
-    EPYDOC = 4
-    AUTO = 255
-
-
-class RenderingStyle(enum.Enum):
-    """Rendering style when unparsing parsed docstrings."""
-
-    COMPACT = 1
-    CLEAN = 2
-    EXPANDED = 3
-
-
-class DocstringMeta:
-    """Docstring meta information.
-
-    Symbolizes lines in form of
-
-        :param arg: description
-        :raises ValueError: if something happens
-    """
-
-    def __init__(
-        self, args: T.List[str], description: T.Optional[str]
-    ) -> None:
-        """Initialize self.
-
-        :param args: list of arguments. The exact content of this variable is
-            dependent on the kind of docstring; it's used to distinguish
-            between custom docstring meta information items.
-        :param description: associated docstring description.
-        """
-        self.args = args
-        self.description = description
-
-
-class DocstringParam(DocstringMeta):
-    """DocstringMeta symbolizing :param metadata."""
-
-    def __init__(
-        self,
-        args: T.List[str],
-        description: T.Optional[str],
-        arg_name: str,
-        type_name: T.Optional[str],
-        is_optional: T.Optional[bool],
-        default: T.Optional[str],
-    ) -> None:
-        """Initialize self."""
-        super().__init__(args, description)
-        self.arg_name = arg_name
-        self.type_name = type_name
-        self.is_optional = is_optional
-        self.default = default
-
-
-class DocstringReturns(DocstringMeta):
-    """DocstringMeta symbolizing :returns or :yields metadata."""
-
-    def __init__(
-        self,
-        args: T.List[str],
-        description: T.Optional[str],
-        type_name: T.Optional[str],
-        is_generator: bool,
-        return_name: T.Optional[str] = None,
-    ) -> None:
-        """Initialize self."""
-        super().__init__(args, description)
-        self.type_name = type_name
-        self.is_generator = is_generator
-        self.return_name = return_name
-
-
-class DocstringRaises(DocstringMeta):
-    """DocstringMeta symbolizing :raises metadata."""
-
-    def __init__(
-        self,
-        args: T.List[str],
-        description: T.Optional[str],
-        type_name: T.Optional[str],
-    ) -> None:
-        """Initialize self."""
-        super().__init__(args, description)
-        self.type_name = type_name
-        self.description = description
-
-
-class DocstringDeprecated(DocstringMeta):
-    """DocstringMeta symbolizing deprecation metadata."""
-
-    def __init__(
-        self,
-        args: T.List[str],
-        description: T.Optional[str],
-        version: T.Optional[str],
-    ) -> None:
-        """Initialize self."""
-        super().__init__(args, description)
-        self.version = version
-        self.description = description
-
-
-class DocstringExample(DocstringMeta):
-    """DocstringMeta symbolizing example metadata."""
-
-    def __init__(
-        self,
-        args: T.List[str],
-        snippet: T.Optional[str],
-        description: T.Optional[str],
-    ) -> None:
-        """Initialize self."""
-        super().__init__(args, description)
-        self.snippet = snippet
-        self.description = description
-
-
-class Docstring:
-    """Docstring object representation."""
-
-    def __init__(
-        self,
-        style=None,  # type: T.Optional[DocstringStyle]
-    ) -> None:
-        """Initialize self."""
-        self.short_description = None  # type: T.Optional[str]
-        self.long_description = None  # type: T.Optional[str]
-        self.blank_after_short_description = False
-        self.blank_after_long_description = False
-        self.meta = []  # type: T.List[DocstringMeta]
-        self.style = style  # type: T.Optional[DocstringStyle]
-
-    @property
-    def params(self) -> T.List[DocstringParam]:
-        """Return a list of information on function params."""
-        return {m.arg_name:m for m in self.meta if isinstance(m, DocstringParam)}
-
-    @property
-    def raises(self) -> T.List[DocstringRaises]:
-        """Return a list of information on the exceptions that the function
-        may raise.
-        """
-        return [
-            item for item in self.meta if isinstance(item, DocstringRaises)
-        ]
-
-    @property
-    def returns(self) -> T.Optional[DocstringReturns]:
-        """Return a single information on function return.
-
-        Takes the first return information.
-        """
-        for item in self.meta:
-            if isinstance(item, DocstringReturns):
-                return item
-        return None
-
-    @property
-    def many_returns(self) -> T.List[DocstringReturns]:
-        """Return a list of information on function return."""
-        return [
-            item for item in self.meta if isinstance(item, DocstringReturns)
-        ]
-
-    @property
-    def deprecation(self) -> T.Optional[DocstringDeprecated]:
-        """Return a single information on function deprecation notes."""
-        for item in self.meta:
-            if isinstance(item, DocstringDeprecated):
-                return item
-        return None
-
-    @property
-    def examples(self) -> T.List[DocstringExample]:
-        """Return a list of information on function examples."""
-        return [
-            item for item in self.meta if isinstance(item, DocstringExample)
-        ]
-
-
-class SectionType(enum.IntEnum):
-    """Types of sections."""
-
-    SINGULAR = 0
-    """For sections like examples."""
-
-    MULTIPLE = 1
-    """For sections like params."""
-
-    SINGULAR_OR_MULTIPLE = 2
-    """For sections like returns or yields."""
-
-
-class Section(namedtuple("SectionBase", "title key type")):
-    """A docstring section."""
-
-
-GOOGLE_TYPED_ARG_REGEX = re.compile(r"\s*(.+?)\s*\(\s*(.*[^\s]+)\s*\)")
-GOOGLE_ARG_DESC_REGEX = re.compile(r".*\. Defaults to (.+)\.")
-MULTIPLE_PATTERN = re.compile(r"(\s*[^:\s]+:)|([^:]*\]:.*)")
-
-DEFAULT_SECTIONS = [
-    Section("Arguments", "param", SectionType.MULTIPLE),
-    Section("Args", "param", SectionType.MULTIPLE),
-    Section("Parameters", "param", SectionType.MULTIPLE),
-    Section("Params", "param", SectionType.MULTIPLE),
-    Section("Raises", "raises", SectionType.MULTIPLE),
-    Section("Exceptions", "raises", SectionType.MULTIPLE),
-    Section("Except", "raises", SectionType.MULTIPLE),
-    Section("Attributes", "attribute", SectionType.MULTIPLE),
-    Section("Example", "examples", SectionType.SINGULAR),
-    Section("Examples", "examples", SectionType.SINGULAR),
-    Section("Returns", "returns", SectionType.SINGULAR_OR_MULTIPLE),
-    Section("Yields", "yields", SectionType.SINGULAR_OR_MULTIPLE),
-]
-
-
-class GoogleDocstringParser:
-    """Parser for Google-style docstrings."""
-
-    def __init__(
-        self, sections: T.Optional[T.List[Section]] = None, title_colon=True
-    ):
-        """Setup sections.
-
-        :param sections: Recognized sections or None to defaults.
-        :param title_colon: require colon after section title.
-        """
-        if not sections:
-            sections = DEFAULT_SECTIONS
-        self.sections = {s.title: s for s in sections}
-        self.title_colon = title_colon
-        self._setup()
-
-    def _setup(self):
-        if self.title_colon:
-            colon = ":"
-        else:
-            colon = ""
-        self.titles_re = re.compile(
-            "^("
-            + "|".join(f"({t})" for t in self.sections)
-            + ")"
-            + colon
-            + "[ \t\r\f\v]*$",
-            flags=re.M,
-        )
-
-    def _build_meta(self, text: str, title: str) -> DocstringMeta:
-        """Build docstring element.
-
-        :param text: docstring element text
-        :param title: title of section containing element
-        :return:
-        """
-
-        section = self.sections[title]
-
-        if (
-            section.type == SectionType.SINGULAR_OR_MULTIPLE
-            and not MULTIPLE_PATTERN.match(text)
-        ) or section.type == SectionType.SINGULAR:
-            return self._build_single_meta(section, text)
-
-        if ":" not in text:
-            # raise ParseError(f"Expected a colon in {text!r}.")
-            return None
-
-        # Split spec and description
-        before, desc = text.split(":", 1)
-        if desc:
-            desc = desc[1:] if desc[0] == " " else desc
-            if "\n" in desc:
-                first_line, rest = desc.split("\n", 1)
-                desc = first_line + "\n" + inspect.cleandoc(rest)
-            desc = desc.strip("\n")
-
-        return self._build_multi_meta(section, before, desc)
-
-    @staticmethod
-    def _build_single_meta(section: Section, desc: str) -> DocstringMeta:
-        if section.key in RETURNS_KEYWORDS | YIELDS_KEYWORDS:
-            return DocstringReturns(
-                args=[section.key],
-                description=desc,
-                type_name=None,
-                is_generator=section.key in YIELDS_KEYWORDS,
-            )
-        if section.key in RAISES_KEYWORDS:
-            return DocstringRaises(
-                args=[section.key], description=desc, type_name=None
-            )
-        if section.key in EXAMPLES_KEYWORDS:
-            return DocstringExample(
-                args=[section.key], snippet=None, description=desc
-            )
-        if section.key in PARAM_KEYWORDS:
-            raise ParseError("Expected paramenter name.")
-        return DocstringMeta(args=[section.key], description=desc)
-
-    @staticmethod
-    def _build_multi_meta(
-        section: Section, before: str, desc: str
-    ) -> DocstringMeta:
-        if section.key in PARAM_KEYWORDS:
-            match = GOOGLE_TYPED_ARG_REGEX.match(before)
-            if match:
-                arg_name, type_name = match.group(1, 2)
-                if type_name.endswith(", optional"):
-                    is_optional = True
-                    type_name = type_name[:-10]
-                elif type_name.endswith("?"):
-                    is_optional = True
-                    type_name = type_name[:-1]
-                else:
-                    is_optional = False
-            else:
-                arg_name, type_name = before, None
-                is_optional = None
-
-            match = GOOGLE_ARG_DESC_REGEX.match(desc)
-            default = match.group(1) if match else None
-
-            return DocstringParam(
-                args=[section.key, before],
-                description=desc,
-                arg_name=arg_name,
-                type_name=type_name,
-                is_optional=is_optional,
-                default=default,
-            )
-        if section.key in RETURNS_KEYWORDS | YIELDS_KEYWORDS:
-            return DocstringReturns(
-                args=[section.key, before],
-                description=desc,
-                type_name=before,
-                is_generator=section.key in YIELDS_KEYWORDS,
-            )
-        if section.key in RAISES_KEYWORDS:
-            return DocstringRaises(
-                args=[section.key, before], description=desc, type_name=before
-            )
-        return DocstringMeta(args=[section.key, before], description=desc)
-
-    def add_section(self, section: Section):
-        """Add or replace a section.
-
-        :param section: The new section.
-        """
-
-        self.sections[section.title] = section
-        self._setup()
-
-    def parse(self, text: str) -> Docstring:
-        """Parse the Google-style docstring into its components.
-
-        :returns: parsed docstring
-        """
-        ret = Docstring(style=DocstringStyle.GOOGLE)
-        if not text:
-            return ret
-
-        # Clean according to PEP-0257
-        text = inspect.cleandoc(text)
-
-        # Find first title and split on its position
-        match = self.titles_re.search(text)
-        if match:
-            desc_chunk = text[: match.start()]
-            meta_chunk = text[match.start() :]
-        else:
-            desc_chunk = text
-            meta_chunk = ""
-
-        # Break description into short and long parts
-        parts = desc_chunk.split("\n", 1)
-        ret.short_description = parts[0] or None
-        if len(parts) > 1:
-            long_desc_chunk = parts[1] or ""
-            ret.blank_after_short_description = long_desc_chunk.startswith(
-                "\n"
-            )
-            ret.blank_after_long_description = long_desc_chunk.endswith("\n\n")
-            ret.long_description = long_desc_chunk.strip() or None
-
-        # Split by sections determined by titles
-        matches = list(self.titles_re.finditer(meta_chunk))
-        if not matches:
-            return ret
-        splits = []
-        for j in range(len(matches) - 1):
-            splits.append((matches[j].end(), matches[j + 1].start()))
-        splits.append((matches[-1].end(), len(meta_chunk)))
-
-        chunks = OrderedDict()  # type: T.Mapping[str,str]
-        for j, (start, end) in enumerate(splits):
-            title = matches[j].group(1)
-            if title not in self.sections:
-                continue
-
-            # Clear Any Unknown Meta
-            # Ref: https://github.com/rr-/docstring_parser/issues/29
-            meta_details = meta_chunk[start:end]
-            unknown_meta = re.search(r"\n\S", meta_details)
-            if unknown_meta is not None:
-                meta_details = meta_details[: unknown_meta.start()]
-
-            chunks[title] = meta_details.strip("\n")
-        if not chunks:
-            return ret
-
-        # Add elements from each chunk
-        for title, chunk in chunks.items():
-            # Determine indent
-            indent_match = re.search(r"^\s*", chunk)
-            if not indent_match:
-                raise ParseError(f'Can\'t infer indent from "{chunk}"')
-            indent = indent_match.group()
-
-            # Check for singular elements
-            if self.sections[title].type in [
-                SectionType.SINGULAR,
-                SectionType.SINGULAR_OR_MULTIPLE,
-            ]:
-                part = inspect.cleandoc(chunk)
-                ret.meta.append(self._build_meta(part, title))
-                continue
-
-            # Split based on lines which have exactly that indent
-            _re = "^" + indent + r"(?=\S)"
-            c_matches = list(re.finditer(_re, chunk, flags=re.M))
-            if not c_matches:
-                raise ParseError(f'No specification for "{title}": "{chunk}"')
-            c_splits = []
-            for j in range(len(c_matches) - 1):
-                c_splits.append((c_matches[j].end(), c_matches[j + 1].start()))
-            c_splits.append((c_matches[-1].end(), len(chunk)))
-            for j, (start, end) in enumerate(c_splits):
-                part = chunk[start:end].strip("\n")
-                ret.meta.append(self._build_meta(part, title))
-
-        return ret
-
-
-def verify_and_get_config_attr_descs(config_cls, strict_docstring_match=True):
-
-    assert dataclasses.is_dataclass(config_cls), f"uh oh <{config_cls.__name__}>."
-
-    # Parse docstring.
-    try:
-        docstring = GoogleDocstringParser().parse(config_cls.__doc__)
-    except Exception as e:
-        raise Exception(f"error parsing {config_cls.__name__} docstring.")
-    
-    # Get attributes and types.
-    config_attrs = docstring.params
-    config_types = config_cls.__annotations__
-
-    # Verify attribute names.
-    config_attr_keys = set(config_attrs.keys())
-    config_type_keys = set(config_types.keys())
-    missing_attr_keys = config_type_keys - config_attr_keys
-    extra_attr_keys = config_attr_keys - config_type_keys
-    if strict_docstring_match:
-        assert not missing_attr_keys and not extra_attr_keys, f"{config_cls.__name__} docstring is either missing attributes ({', '.join(missing_attr_keys) if missing_attr_keys else '--'}) or contains extra attributes ({', '.join(extra_attr_keys) if extra_attr_keys else '--'})."
-
-    # @todo
-    # Verify attribute type names.
-    # for key in config_attr_keys:
-    #     ... todo ...
-
-    # Verify base class attributes.
-    attrs = {k:v for base_cls in config_cls.__bases__ if dataclasses.is_dataclass(base_cls) for k,v in verify_and_get_config_attr_descs(base_cls, strict_docstring_match=strict_docstring_match).items()}
-    for key in config_attr_keys:
-        if key in config_types:
-            attrs[key] = {
-                "desc" : config_attrs[key].description,
-                "type" : config_types[key],
-            }
-
-    return attrs
-
-
-def add_config_args(parser, config_cls):
-    attrs = verify_and_get_config_attr_descs(config_cls, strict_docstring_match=False)
-    for key, attr in attrs.items():
-        _type = attr["type"]
-        if dataclasses.is_dataclass(_type):
-            group = parser.add_argument_group(title=attr["desc"])
-            add_config_args(group, _type)
-        else:
-
-            default_value = getattr(config_cls, key)
-            args = {
-                "help" : attr["desc"],
-                "default" : default_value,
-            }
-
-            if _type == bool:
-                assert isinstance(args["default"], (bool, type(None))), \
-                    f"boolean attribute '{key}' of {config_cls.__name__} " \
-                    "has non-boolean default value."
-
-                # When default=True, add 'no-{key}' arg.
-                if default_value:
-                    args["action"] = "store_false"
-                    args["dest"] = key
-                    key = "no-" + key
-                else:
-                    args["action"] = "store_true"
-
-            elif _type in (int, float):
-                args["type"] = _type
-
-            elif _type == list:
-                args["nargs"] = "*"
-
-            # else: ....... treat as string arg
-            #     raise Exception(f"specialize action for '{key}', type <{_type}>.")
-
-            try:
-                parser.add_argument(f"--{key.replace('_', '-')}", **args)
-            except argparse.ArgumentError as e:
-                pass
-
-
-def get_config_leaf_field_names(config_cls):
-    names = set()
-    for field in dataclasses.fields(config_cls):
-        if dataclasses.is_dataclass(field.type):
-            names.update(get_config_leaf_field_names(field.type))
-        else:
-            names.add(field.name)
-    return names
-
-
-def config_from_args(args, config_cls, add_custom_args=False):
-
-    # Collect config data in a dict.
-    data = {}
-    for field in dataclasses.fields(config_cls):
-        if dataclasses.is_dataclass(field.type):
-            data[field.name] = config_from_args(args, field.type)
-        else:
-            data[field.name] = getattr(args, field.name)
-
-    # Add custom args. (e.g., for tools, tasks)
-    if add_custom_args:
-
-        config_keys = get_config_leaf_field_names(config_cls)
-        arg_keys = set(vars(args).keys())
-        custom_keys = arg_keys - config_keys
-
-        custom_data = {k:v for k, v in vars(args).items() if k in custom_keys}
-        custom_config_cls = dataclasses.make_dataclass(
-            "CustomConfig",
-            [(k, type(v)) for k, v in custom_data.items()])
-        custom_config = custom_config_cls(**custom_data)
-        data["custom"] = custom_config
-
-    # Create config. [ todo: programmatically create dataclass that inherits
-    # TransformerConfig. ]
-    config = config_cls(**data)
-
-    return config
-
-
-def flatten_config(config, base_config_cls=None):
-
-    # Lift sub-config data.
-    flat_config = {}
-    for field in dataclasses.fields(config):
-        value = getattr(config, field.name)
-        if dataclasses.is_dataclass(value):
-            flat_config = { **flat_config, **flatten_config(value) }
-        else:
-            flat_config[field.name] = value
-
-    # Convert to dataclass.
-    if base_config_cls:
-        base_keys = set(field.name for field in dataclasses.fields(base_config_cls))
-        flat_config_cls = dataclasses.make_dataclass(
-            cls_name="FlatMegatronConfig",
-            fields=[(k, T.Any, dataclasses.field(default=None))
-                    for k, v in flat_config.items()
-                    if k not in base_keys],
-            bases=(base_config_cls,))
-        flat_config = flat_config_cls(**flat_config)
-
-    return flat_config
diff --git a/tools/retro/docker/Dockerfile b/tools/retro/docker/Dockerfile
deleted file mode 100644
index e8945b373a4..00000000000
--- a/tools/retro/docker/Dockerfile
+++ /dev/null
@@ -1,19 +0,0 @@
-FROM nvcr.io/nvidia/pytorch:23.09-py3
-
-RUN pip install -U faiss-gpu
-
-RUN apt update
-
-RUN apt install -qy htop
-
-RUN pip install -U transformers
-
-RUN pip install --upgrade google-api-python-client
-
-RUN pip install sentencepiece
-
-RUN pip install h5py
-
-RUN pip install nltk
-
-RUN pip install einops
diff --git a/tools/retro/preprocess_data.py b/tools/retro/preprocess_data.py
deleted file mode 100644
index 3a01d67cabb..00000000000
--- a/tools/retro/preprocess_data.py
+++ /dev/null
@@ -1,296 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
-
-"""Preprocess data for Retro.
-
-Stages (see argument '--retro-tasks'):
-- Build chunk database (DB).
-- Build index (train, add).
-- Query pretraining neighbors.
-"""
-
-import json
-import os
-import sys
-import torch
-
-from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
-from megatron.core.datasets.utils import get_blend_from_list
-from megatron.core.datasets.retro.db import build_db
-from megatron.core.datasets.retro.index import add_to_index, train_index
-from megatron.core.datasets.retro.config import (
-    RetroBertEmbedders,
-    RetroGPTChunkDatasets,
-    RetroPreprocessingConfig,
-    RetroTokenizers,
-)
-from megatron.core.datasets.retro.query.gpt_chunk_dataset import build_gpt_chunk_datasets_from_gpt_datasets
-from megatron.core.datasets.retro.query.multi_split_gpt_dataset import (
-    MultiSplitGPTDataset,
-    MultiSplitGPTDatasetConfig,
-)
-from megatron.core.datasets.retro.query.query import query_neighbors
-from megatron.core.datasets.retro.query.utils import get_query_dir
-from megatron.core.datasets.retro.utils import retro_makedir
-from megatron.core.models.retro.utils import (
-    get_config_path,
-    get_gpt_data_dir,
-)
-from megatron.training import get_args, initialize_megatron, print_rank_0
-from megatron.training.arguments import core_transformer_config_from_args
-from megatron.training.tokenizer.tokenizer import (
-    _BertWordPieceTokenizer,
-    _GPT2BPETokenizer,
-    _GPTSentencePieceTokenizer,
-)
-from megatron.training import get_train_valid_test_num_samples
-from pretrain_gpt import is_dataset_built_on_rank
-from tools.bert_embedding import BertEmbedder, DiskDataParallelBertEmbedder
-from tools.retro.config_utils import add_config_args
-
-
-def add_retro_args(parser):
-    group = parser.add_argument_group(title="Retro preprocessing")
-    add_config_args(group, RetroPreprocessingConfig)
-    return parser
-
-
-def initialize_megatron_retro():
-    '''Initialize megatron & save Retro config.'''
-
-    # Prevent arguments.py from overriding preprocessing args.
-    project_dir_idx = sys.argv.index("--retro-project-dir")
-    retro_project_dir = sys.argv[project_dir_idx + 1]
-    del sys.argv[project_dir_idx] # delete key
-    del sys.argv[project_dir_idx] # delete value
-
-    # Initialize.
-    initialize_megatron(extra_args_provider=add_retro_args)
-
-    args = get_args()
-    args.retro_project_dir = retro_project_dir
-
-    # Retro config.
-    config = get_retro_preprocessing_config()
-
-    # Save retro config.
-    if config.retro_task_validate is None:
-        retro_makedir(config, config.retro_project_dir)
-        save_config(config)
-
-    return config
-
-
-def get_bert_embedders(config):
-    mem_embedder = BertEmbedder(
-        batch_size = config.retro_bert_batch_size,
-        max_bert_seq_length = config.retro_bert_max_chunk_length,
-        embedder_type = "megatron",
-    )
-    return RetroBertEmbedders(
-        mem = mem_embedder,
-        disk = DiskDataParallelBertEmbedder(mem_embedder, config.retro_block_size),
-    )
-
-
-def get_gpt_chunk_datasets(config):
-
-    args = get_args()
-
-    # Dataset config.
-    data_dir = get_gpt_data_dir(config.retro_project_dir)
-    blend = list(config.retro_gpt_data_path)
-    for i in range(len(blend) - 1, -1, -2):
-        blend[i] = os.path.join(data_dir, blend[i])
-    data_config = MultiSplitGPTDatasetConfig(
-        random_seed=config.retro_gpt_seed,
-        sequence_length=config.retro_gpt_seq_length,
-        blend=get_blend_from_list(blend),
-        blend_per_split=[
-            get_blend_from_list(args.train_data_path),
-            get_blend_from_list(args.valid_data_path),
-            get_blend_from_list(args.test_data_path)
-        ],
-        split=config.retro_gpt_split,
-        split_preprocessing=config.retro_gpt_split,
-        path_to_cache=config.retro_gpt_data_cache_path,
-        return_document_ids=True,
-        tokenizer=config.retro_tokenizers.gpt,
-        reset_position_ids=args.reset_position_ids,
-        reset_attention_mask=args.reset_attention_mask,
-        eod_mask_loss=args.eod_mask_loss,
-        mid_level_dataset_surplus=args.mid_level_dataset_surplus,
-    )
-
-    # GPT datasets.
-    print_rank_0(" > multi-split gpt datasets.")
-    train_valid_test_num_samples = get_train_valid_test_num_samples()
-    train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
-        MultiSplitGPTDataset,
-        train_valid_test_num_samples,
-        is_dataset_built_on_rank,
-        data_config,
-    ).build()
-
-    gpt_datasets = {
-        "train" : (train_ds, train_valid_test_num_samples[0]),
-        "valid" : (valid_ds, train_valid_test_num_samples[1]),
-        "test"  : (test_ds, train_valid_test_num_samples[2]),
-    }
-
-    # Chunk datasets.
-    chunk_datasets = build_gpt_chunk_datasets_from_gpt_datasets(
-        project_dir=config.retro_project_dir,
-        gpt_datasets=gpt_datasets,
-        sample_length=config.retro_gpt_seq_length,
-        chunk_length=config.retro_gpt_chunk_length,
-    )
-    chunk_datasets = RetroGPTChunkDatasets(**chunk_datasets)
-
-    return chunk_datasets
-
-
-def get_gpt_tokenizer(config):
-    '''GPT (BPE) tokenizer.'''
-    tokenizer_type = config.retro_gpt_tokenizer_type
-    if tokenizer_type == "GPT2BPETokenizer":
-        assert config.retro_gpt_vocab_file and config.retro_gpt_merge_file
-        return _GPT2BPETokenizer(
-            vocab_file=os.path.join(
-                config.retro_project_dir,
-                config.retro_gpt_vocab_file,
-            ),
-            merge_file=os.path.join(
-                config.retro_project_dir,
-                config.retro_gpt_merge_file,
-            ),
-        )
-    elif tokenizer_type == 'GPTSentencePieceTokenizer':
-        assert config.retro_gpt_tokenizer_model is not None
-        return _GPTSentencePieceTokenizer(os.path.join(
-            config.retro_project_dir,
-            config.retro_gpt_tokenizer_model,
-        ))
-    else:
-        raise Exception("unrecognized gpt tokenizer, '%s'." % tokenizer_type)
-
-
-def get_bert_tokenizer(config):
-    '''Bert (Wordpiece) tokenizer.'''
-    lower_case = {
-        "BertWordPieceLowerCase" : True,
-        "BertWordPieceCase" : False,
-    }[config.retro_bert_tokenizer_type]
-    return _BertWordPieceTokenizer(
-        vocab_file=os.path.join(
-            config.retro_project_dir,
-            config.retro_bert_vocab_file,
-        ),
-        lower_case=lower_case,
-    )
-
-
-def get_tokenizers(config):
-    return RetroTokenizers(
-        gpt = get_gpt_tokenizer(config),
-        bert = get_bert_tokenizer(config),
-    )
-
-
-def get_retro_preprocessing_config():
-
-    # Arguments.
-    args = get_args()
-
-    # Retro config.
-    config = core_transformer_config_from_args(
-        args, config_class=RetroPreprocessingConfig)
-
-    # Add tools.
-    config.retro_tokenizers = get_tokenizers(config)
-    config.retro_bert_embedders = get_bert_embedders(config)
-    config.retro_gpt_chunk_datasets = get_gpt_chunk_datasets(config)
-
-    return config
-
-
-def save_config(config):
-    '''Save copy of config within retro project dir.'''
-
-    if torch.distributed.get_rank() == 0:
-
-        # GPT config + block size.
-        config_subset = {
-            k:v for k,v in vars(config).items()
-            if k.startswith("retro_gpt") and k != "retro_gpt_chunk_datasets"
-        }
-        config_subset["retro_block_size"] = config.retro_block_size
-
-        # Bert config.
-        config_subset["retro_bert_tokenizer_type"] = config.retro_bert_tokenizer_type
-        config_subset["retro_bert_vocab_file"] = config.retro_bert_vocab_file
-
-        # Neighbor directories.
-        query_dir = get_query_dir(config.retro_project_dir)
-        config_subset["retro_neighbor_dirs"] = {
-            k : (os.path.relpath(v["neighbor_dir"], query_dir) if v is not None else None)
-            for k, v in vars(config.retro_gpt_chunk_datasets).items()
-        }
-
-        # Save.
-        config_path = get_config_path(config.retro_project_dir)
-        with open(config_path, "w") as f:
-            json.dump(config_subset, f, indent=4, sort_keys=True)
-
-    torch.distributed.barrier()
-
-
-if __name__ == "__main__":
-
-    # Initalize Megatron.
-    config = initialize_megatron_retro()
-
-    # Expand tasks.
-    task_remap = {
-        "build" : [ "db-build", "index-train", "index-add", "query-neighbors" ],
-        "index-build" : [ "index-train", "index-add" ],
-        "db-build" : [ "db-build" ],
-        "index-train" : [ "index-train" ],
-        "index-add" : [ "index-add" ],
-        "query-neighbors" : [ "query-neighbors" ],
-    }
-    tasks = []
-    for task in config.retro_tasks:
-        tasks.extend(task_remap[task])
-    config.retro_tasks = tasks
-
-    # Select task to run.
-    for task in tasks:
-
-        print_rank_0("start '%s%s'." % (
-            "" if config.retro_task_validate is None else "[validate] ",
-            task,
-        ))
-
-        # DB (i.e., chunk db).
-        if task == "db-build":
-            build_db(config)
-
-        # Index.
-        elif task == "index-train":
-            train_index(config)
-        elif task == "index-add":
-            add_to_index(config)
-
-        # Query.
-        elif task == "query-neighbors":
-            query_neighbors(config)
-
-        else:
-            raise Exception("specialize for task '%s'." % task)
-
-        torch.distributed.barrier()
-
-        print_rank_0("end '%s%s'." % (
-            "" if config.retro_task_validate is None else "[validate] ",
-            task,
-        ))
diff --git a/tools/retro/sft/README.md b/tools/retro/sft/README.md
deleted file mode 100644
index e5898790383..00000000000
--- a/tools/retro/sft/README.md
+++ /dev/null
@@ -1,3 +0,0 @@
-## Note
-
-The content within this `sft` directory is still under active development and will be updated soon.
\ No newline at end of file
diff --git a/tools/retro/sft/dataset_conv.py b/tools/retro/sft/dataset_conv.py
deleted file mode 100644
index 3dd8fa9cd56..00000000000
--- a/tools/retro/sft/dataset_conv.py
+++ /dev/null
@@ -1,446 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-import re
-import json
-import os
-from typing import Any, Iterable, Dict, Optional
-
-from numpy import ndarray
-from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
-from megatron.core.datasets.utils import Split
-import torch
-import numpy
-import glob
-from collections import OrderedDict
-
-from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
-from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset
-from megatron.core.datasets.utils import Split
-from dataclasses import dataclass
-
-
-_DATASET_NAME_PATTERNS = {
-    Split.train: r"(?P<name>[^\0]+)\/(?P=name)\_QA\_train.json",
-    Split.valid: r"(?P<name>[^\0]+)\/(?P=name)\_QA\_dev.json",
-}
-
-
-@dataclass
-class JsonQADatasetConfig(BlendedMegatronDatasetConfig):
-    """Configuration object for the QA finetuning pipeline
-    """
-    ft_neighbours: int = 1
-
-    bert_retriever_neighbours: bool = False
-
-    longform_answer: bool = False
-
-    inference_only: bool = False
-
-    retrieved_neighbours: bool = False
-
-    fix_newsqa: bool = True
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        assert self.blend_per_split is not None
-
-
-@dataclass
-class RetroJsonQADatasetConfig(JsonQADatasetConfig):
-    """Configuration object for the Retro QA finetuning pipeline
-    """
-    retro_num_neighbors: int = None
-
-    retro_gpt_retrieved_length: int = None
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        assert self.retro_num_neighbors is not None
-        assert self.retro_gpt_retrieved_length is not None
-
-
-class JsonQADataset(MegatronDataset):
-
-    def __init__(self, dataset: Any, dataset_path: str, indices: ndarray, num_samples: Optional[int], index_split: Split, config: BlendedMegatronDatasetConfig) -> None:
-        super().__init__(dataset, dataset_path, indices, num_samples, index_split, config)
-        matches = re.findall(_DATASET_NAME_PATTERNS[index_split], dataset_path)
-        assert len(matches) == 1
-        assert len(matches[0]) > 0
-        self.dataset_name = matches[0]
-
-    @staticmethod
-    def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int:
-        return len(low_level_dataset)
-
-    @staticmethod
-    def build_low_level_dataset(dataset_path: str, config: JsonQADatasetConfig) -> Iterable:
-        assert os.path.isfile(dataset_path), f"{dataset_path} does not exist on disk"
-        return preprocess(dataset_path, config)
-
-    def __len__(self) -> int:
-        return len(self.dataset)
-
-    def __getitem__(self, idx: int) -> Dict[str, ndarray]:
-        sample = self.dataset[idx % len(self.dataset)]
-
-        # unpack tokens
-        query, answer, neighbours = sample
-
-        # tokenization
-        output_tokens = self.config.tokenizer.tokenize(answer)
-
-        input_tokens = reformat_prompt(
-            query,
-            neighbours,
-            self.dataset_name,
-            self.config.ft_neighbours,
-            len(output_tokens),
-            self.config.tokenizer,
-            self.config.sequence_length
-        )
-
-        # padding
-        tokens, answer_mask = pad_and_convert_to_numpy(
-            input_tokens, output_tokens, self.config.tokenizer.pad, self.config.sequence_length, self.config.tokenizer.eos
-        )
-
-        train_sample = {
-            'text': tokens,
-            'answer_mask': answer_mask,
-        }
-
-        return train_sample
-
-
-class RetroJsonQADataset(JsonQADataset):
-
-    def __getitem__(self, idx: int) -> Dict[str, ndarray]:
-
-        sample = self.dataset[idx % len(self.dataset)]
-
-        # unpack tokens
-        query, answer, neighbours = sample
-
-        # tokenization
-        output_tokens = self.config.tokenizer.tokenize(answer)
-
-        input_tokens = reformat_prompt_retro(
-            query,
-            neighbours,
-            self.dataset_name,
-            self.config.ft_neighbours,
-            len(output_tokens),
-            self.config.tokenizer,
-            self.config.sequence_length
-        )
-
-        # padding
-        tokens, answer_mask = pad_and_convert_to_numpy(
-            input_tokens,
-            output_tokens,
-            self.config.tokenizer.pad,
-            self.config.sequence_length,
-            self.config.tokenizer.eos
-        )
-
-        # get retro neighbors
-        # context chunk and answer chunk
-        n_chunks_per_sample = 2
-        num_neighbors = self.config.retro_num_neighbors
-        # disable retro encoder
-        neighbor_tokens = numpy.zeros(
-            [n_chunks_per_sample, num_neighbors, self.config.retro_gpt_retrieved_length],
-            dtype=numpy.int64
-        )
-
-        train_sample = {
-            'text': tokens,
-            'answer_mask': answer_mask,
-            'neighbor_tokens': neighbor_tokens,
-            'context_len': len(input_tokens)
-        }
-
-        return train_sample
-
-
-def format_multichoice(multichoice_options):
-    options_text = ["({}) {}".format(chr(ord('A') + i), option) for i, option in
-                    zip(range(len(multichoice_options)), multichoice_options)]
-    return "Choose one based on the following options: {}".format(" ".join(options_text))
-
-
-def format_multichoice_question(question, multichoice_options):
-    return "{}\n{}".format(question, format_multichoice(multichoice_options))
-
-
-def format_answer(answer):
-    return " {}".format(answer)
-
-
-def preprocess(dataset_path: str, config: JsonQADatasetConfig):
-    assert config.ft_neighbours > 0
-    if config.longform_answer:
-        nq_examples = []
-        with open(dataset_path, "r") as f:
-            for fn in f:
-                nq_examples.append(json.loads(fn))
-    else:
-        nq_examples = []
-        for my_data_file in sorted(glob.glob(dataset_path)):
-            with open(my_data_file, "r", encoding='utf-8') as f:
-                nq_examples.extend(json.load(f))
-
-    data = []
-    for instance in nq_examples:
-        question = instance["question"]
-        if 'qa_type' in instance and instance['qa_type'] == "multi_choice_qa":
-            question = format_multichoice_question(question, instance["multichoice_options"])
-        if config.bert_retriever_neighbours:
-            contexts = instance["bert_pretrain_corpus_neighbours"]
-            neighbours = ["source: " + ctx for ctx in contexts]
-        else:
-            if config.retrieved_neighbours:
-                contexts = instance["ctxs"]
-                neighbours = ["title: " + ctx["title"] + ", source: " + ctx["text"] for ctx in contexts]
-            else:
-                if "sub-paragraphs" in instance:
-                    if type(instance["sub-paragraphs"]) == list:  # doc2dial:
-                        neighbours = [
-                            "title: " + instance["sub-paragraphs"][0] + ", source: " + instance["sub-paragraphs"][1]]
-                    else:
-                        neighbours = ["title: , source: " + instance["sub-paragraphs"]]
-                elif config.fix_newsqa and "sub_paragraph" in instance:
-                    neighbours = ["title: , source: " + instance["sub_paragraph"]]
-                else:
-                    neighbours = ["title: , source: "]
-
-        if config.inference_only:
-            data.append((question, None, neighbours))
-        else:
-            if config.longform_answer:
-                if "longform_answer" in instance:
-                    answers = [instance["longform_answer"]]
-                else:
-                    continue
-            else:
-                if "answers" in instance:
-                    answers = instance["answers"]
-                elif "answer" in instance:
-                    if type(instance["answer"]) is str:
-                        answers = [instance["answer"]]
-                    elif type(instance["answer"]) is list:
-                        answers = instance["answer"]
-                    else:
-                        answers = [str(instance["answer"])]
-                else:
-                    raise ValueError("need to have answer or answers")
-            if len(answers) < 1:
-                continue
-            else:
-                if type(answers[0]) is dict:
-                    answers = [answers[0]["text"].strip()]
-                elif type(answers[0]) is str:
-                    answers = [answers[0]]
-                else:
-                    raise ValueError("unsupported type for answer(s)")
-
-            for answer in answers:
-                answer = format_answer(answer)
-                data.append((question, answer, neighbours))
-
-    return data
-
-
-def count_stat(dataset, tokenizer, k):
-    nb_lens = []
-    for i, d in enumerate(dataset):
-        query, answer, neighbours = d
-        nb_lens.extend([len(tokenizer.tokenize(neighbour)) for neighbour in neighbours[:k]])
-
-    print("len of nb", len(nb_lens))
-    print("max of len nb", max(nb_lens))
-    print("num of cut ", sum([l > 128 for l in nb_lens]), sum([l > 128 for l in nb_lens]) // len(nb_lens))
-    print("last max", sorted(nb_lens)[-10:])
-
-
-def reformat_prompt_retro(query, neighbours, dataset_name, ft_neighbours, \
-                          max_output_len, tokenizer, max_seq_length):
-    system = ("System: This is a chat between a user and an artificial intelligence assistant. The assistant gives "
-              "helpful, detailed, and polite answers to the user's questions.\n\n")
-
-    if dataset_name in ["oasst", "quiet_cockatoo", "open_inst", "quiet-cockatoo_commercial"]:
-        input_tokens = tokenizer.tokenize(system + query)
-        return input_tokens
-
-    short_span_with_context = ["drop", "NarrativeQA", "QASC", "Quoref", "ROPES", "squad1.1", "squad2.0", "newsqa", "nq",
-                               "tqa", "quac"]
-    yes_no_without_context = ["BoolQ"]
-    multichoices = [""]
-    formatted_dataset_name = ["doc2dial", "quac", "qrecc", "sharc"]
-
-    if dataset_name in formatted_dataset_name:
-        dialogue_turn = query
-    else:
-        if dataset_name in short_span_with_context:
-            user = "{} Answer the above question with a short phrase.".format(query)
-        elif dataset_name in yes_no_without_context:
-            user = "{} Answer the above question with True or False.".format(query)
-        else:
-            user = "{} Answer the above question with a long complete answer.".format(query)
-
-        if dataset_name in short_span_with_context:
-            dialogue_format = "User: {}\n\nAssistant: The answer is"
-            dialogue_turn = dialogue_format.format(user)
-        else:
-            dialogue_format = "User: {}\n\nAssistant:"
-            dialogue_turn = dialogue_format.format(user)
-
-    if ft_neighbours > 0:
-        context = "\n\n".join(neighbours[0:ft_neighbours]) + "\n\n"
-        context_tokens = tokenizer.tokenize(context)
-        dialogue_tokens = tokenizer.tokenize(dialogue_turn)
-        system_tokens = tokenizer.tokenize(system)
-        context_tokens = context_tokens[:max_seq_length - max_output_len - len(dialogue_tokens) - len(system_tokens)]
-        context = tokenizer.detokenize(context_tokens)
-
-        all_input = system + context + dialogue_turn
-        print(all_input)
-        input_tokens = tokenizer.tokenize(all_input)
-    else:
-        all_input = system + dialogue_turn
-        input_tokens = tokenizer.tokenize(all_input)
-
-    return input_tokens
-
-
-def flan_format(system, context, dialogue_turn, template_id=0):
-    templates = [
-        "{}User: Answer based on context:\n\n{}{}",
-        "{}User: {}Answer this question based on the article: {}",
-        "{}User: {}{}",
-        "{}User: {}Answer this question: {}",
-        "{}User: Read this article and answer this question {}{}",
-        "{}User: {}Based on the above article, answer a question. {}",
-        "{}User: Context: {}Question: {}"
-    ]
-    template = templates[template_id - 1].format(system, context, dialogue_turn)
-    return template
-
-
-def reformat_prompt(query, neighbours, dataset_name, ft_neighbours, \
-                    max_output_len, tokenizer, max_seq_length, template_id=0):
-    system = ("System: This is a chat between a user and an artificial intelligence assistant. The assistant gives "
-              "helpful, detailed, and polite answers to the user's questions based on the context. The assistant "
-              "should also indicate when the answer cannot be found in the context.\n\n")
-
-    if dataset_name in ["oasst", "quiet_cockatoo", "open_inst", "quiet-cockatoo_commercial"]:
-        input_tokens = tokenizer.tokenize(system + query)
-        return input_tokens
-
-    short_span_with_context = ["drop", "NarrativeQA", "QASC", "Quoref", "ROPES", "squad1.1", "squad2.0", "newsqa", "nq",
-                               "BioASQ", "DuoRC_ParaphraseRC", "TextbookQA", "tqa"]
-    yes_no_without_context = ["boolq", "multirc"]
-    multichoices = ["race"]
-    # multi-turn qa datasets
-    formatted_dataset_name = ["convqa", "chatgptgen", "doc2dial", "quac", "qrecc", "sharc"]
-
-    if dataset_name in formatted_dataset_name:
-        dialogue_turn = query
-    else:
-        if dataset_name in short_span_with_context:
-            if template_id == 0:
-                user = "Answer the following question with a short span. {}".format(query)
-            else:
-                user = query
-        elif dataset_name in yes_no_without_context:
-            user = "Answer the following question with True or False. {}".format(query)
-        elif dataset_name in multichoices:
-            user = "Answer the following question by selecting one of the provided options. {}".format(query)
-        else:
-            if template_id == 0:
-                user = "Please give a full and complete answer for the question. {}".format(query)
-            else:
-                user = query
-
-        if dataset_name in short_span_with_context:
-            if template_id == 0:
-                dialogue_format = "User: {}\n\nAssistant: The answer is"
-            else:
-                dialogue_format = "{}\n\nAssistant: The answer is"
-            dialogue_turn = dialogue_format.format(user)
-        else:
-            if template_id == 0:
-                dialogue_format = "User: {}\n\nAssistant:"
-            else:
-                dialogue_format = "{}\n\nAssistant:"
-            dialogue_turn = dialogue_format.format(user)
-
-    if ft_neighbours > 0:
-        context = "\n\n".join(neighbours[0:ft_neighbours]) + "\n\n"
-        context_tokens = tokenizer.tokenize(context)
-        dialogue_tokens = tokenizer.tokenize(dialogue_turn)
-        system_tokens = tokenizer.tokenize(system)
-        context_tokens = context_tokens[:max_seq_length - max_output_len - len(dialogue_tokens) - len(system_tokens)]
-        context = tokenizer.detokenize(context_tokens)
-
-        if template_id == 0:
-            all_input = system + context + dialogue_turn
-        else:
-            all_input = flan_format(system, context, dialogue_turn, template_id=template_id)
-        input_tokens = tokenizer.tokenize(all_input)
-    else:
-        all_input = system + dialogue_turn
-        input_tokens = tokenizer.tokenize(all_input)
-
-    return input_tokens
-
-
-def reformat_prompt_short(query, neighbours, dataset_name, ft_neighbours, \
-                          max_output_len, tokenizer, max_seq_length):
-    if not query.endswith("?"):
-        query = query + "?"
-    query = "Question: {} Answer: The answer is".format(query)
-
-    if ft_neighbours > 0:
-        context = "\n\n".join(neighbours[0:ft_neighbours]) + "\n\n"
-        context_tokens = tokenizer.tokenize(context)
-        dialogue_tokens = tokenizer.tokenize(query)
-        context_tokens = context_tokens[:max_seq_length - max_output_len - len(dialogue_tokens)]
-        context = tokenizer.detokenize(context_tokens)
-        all_input = context + query
-        input_tokens = tokenizer.tokenize(all_input)
-    else:
-        all_input = query
-        input_tokens = tokenizer.tokenize(all_input)
-
-    return input_tokens
-
-
-def pad_and_convert_to_numpy(input_ids, output_ids,
-                             pad_id, max_seq_length,
-                             eos_id):
-    """Pad sequences and convert them to numpy."""
-    if len(input_ids) > max_seq_length:
-        input_ids = input_ids[:max_seq_length - 1]
-
-    if len(input_ids + output_ids) > max_seq_length:
-        output_ids = output_ids[:max_seq_length - len(input_ids)]
-
-    tokens = input_ids + output_ids
-    answer_mask = [0] * len(input_ids) + [1] * len(output_ids)
-
-    # padding
-    num_tokens = len(tokens)
-    padding_length = max_seq_length - num_tokens
-    assert padding_length >= 0
-
-    # Tokens.
-    filler = [pad_id] * padding_length
-    tokens = numpy.array(tokens + [eos_id] + filler, dtype=numpy.int64)
-
-    # answer mask
-    answer_mask = answer_mask + [1] + [0] * padding_length
-    answer_mask = numpy.array(answer_mask, dtype=numpy.int64)
-
-    return tokens, answer_mask
diff --git a/tools/retro/sft/open_inst.sh b/tools/retro/sft/open_inst.sh
deleted file mode 100644
index 9ebe063b810..00000000000
--- a/tools/retro/sft/open_inst.sh
+++ /dev/null
@@ -1 +0,0 @@
-DATA_BLEND="1.0 open_inst"
diff --git a/tools/retro/sft/sft_retro.py b/tools/retro/sft/sft_retro.py
deleted file mode 100644
index e71d841f4df..00000000000
--- a/tools/retro/sft/sft_retro.py
+++ /dev/null
@@ -1,278 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-"""Pretrain GPT"""
-
-import torch
-from functools import partial, reduce
-import sys, os
-
-sys.path.append(os.path.abspath(os.path.join(
-    os.path.join(os.path.dirname(__file__), "../../../"))))
-from megatron.training import get_args, get_retro_args
-from megatron.training import print_rank_0
-from megatron.training import get_timers
-from megatron.training import get_tokenizer
-from megatron.core import tensor_parallel
-from megatron.core.enums import ModelType
-from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
-from megatron.core.datasets.utils import get_blend_from_list
-from megatron.training import pretrain
-from megatron.training.utils import get_ltor_masks_and_position_ids
-from megatron.training.utils import average_losses_across_data_parallel_group
-from pretrain_gpt import is_dataset_built_on_rank
-from model_provider import model_provider
-from gpt_builders import gpt_builder
-from tools.retro.sft.dataset_conv import JsonQADataset, JsonQADatasetConfig, RetroJsonQADataset, RetroJsonQADatasetConfig
-
-
-def get_tasks_args(parser):
-    """Provide extra arguments required for tasks."""
-    group = parser.add_argument_group(title='tasks')
-
-    # parameters for the knowledgeable dialogue generation
-    group.add_argument('--task', type=str, default=None,
-                       help='Task name.')
-    group.add_argument('--epochs', type=int, default=None,
-                       help='Number of finetunning epochs. Zero results in '
-                            'evaluation only.')
-    group.add_argument('--keep-last', action='store_true',
-                       help='Keep the last batch (maybe incomplete) in'
-                            'the data loader')
-    group.add_argument('--pretrained-checkpoint', type=str, default=None,
-                       help='Pretrained checkpoint used for finetunning.')
-    group.add_argument('--data-folder', type=str, default=None,
-                       help='dataset folder')
-    group.add_argument('--answer-loss-only', action='store_true', default=False,
-                       help='take the loss from answer part, ignore the context')
-    group.add_argument('--weight', type=float, default=1)
-    group.add_argument('--adaptor', action='store_true', default=False)
-    group.add_argument('--project-size', type=int, default=256)
-    group.add_argument('--cyclic-train-iters', type=int, default=None)
-    group.add_argument('--stored_params', type=dict, default=dict())
-    group.add_argument('--eval_ppl', action='store_true', default=False)
-    group.add_argument('--debug', action='store_true', default=False)
-    group.add_argument('--add_retriever', action='store_true', default=False)
-    group.add_argument('--return_doc_ids', action='store_true', default=False)
-    group.add_argument('--return_neighbor_ids', action='store_true', default=False)
-    group.add_argument('--add_offset_doc_ids', action='store_true', default=False)
-    group.add_argument('--offset_dict_path', type=str, default='')
-    group.add_argument('--neighbors_path', type=str, default='')
-    group.add_argument('--valid_neighbors_path', type=str, default='')
-    group.add_argument('--database_path', type=str, default='')
-    group.add_argument('--valid_database_path', type=str, default='')
-    group.add_argument('--encoder-layers', type=int, default=12)
-    group.add_argument('--encoder-hidden-dropout', type=float, default=0.1)
-    group.add_argument('--encoder-attention-dropout', type=float, default=0.1)
-    group.add_argument('--k', type=int, default=2)
-    group.add_argument('--r', type=int, default=128)
-    group.add_argument('--m', type=int, default=64)
-    group.add_argument('--dpr-mode', type=str, default="multi")
-    group.add_argument('--faiss-ckpt', type=str, default='')
-    group.add_argument('--original-db-file', type=str, default="")
-    group.add_argument('--ft_neighbours', type=int, default=1)
-    group.add_argument('--reuse-top', action='store_true', default=False)
-    group.add_argument('--shuffle_topn', action='store_true', default=False)
-    group.add_argument('--chunk0', action='store_true', default=False)
-    group.add_argument('--disable-encoder', action='store_true', default=False)
-    group.add_argument('--qa-space-pad', action='store_true', default=False)
-    group.add_argument('--retro-mask-encoder', action='store_true', default=False)
-    group.add_argument('--without-title', action='store_true', default=False)
-    group.add_argument('--longform-answer', action='store_true', default=False)
-    group.add_argument('--bert-retriever-neighbours', action='store_true', default=False)
-    group.add_argument('--prefix', action='store_true', default=False)
-    group.add_argument('--question-in-encoder', action='store_true', default=False)
-    group.add_argument('--reset_eval', type=bool, default=True)  ## by default reset eval for each eval
-    return parser
-
-
-def get_batch(data_iterator):
-    """Generate a batch"""
-    args = get_args()
-    tokenizer = get_tokenizer()
-
-    # Items and their type.
-    keys = ['text', 'answer_mask']
-    datatype = torch.int64
-
-    if args.retro_add_retriever:
-        keys += 'neighbor_tokens', 'context_len'
-
-    # Broadcast data.
-    if data_iterator is not None:
-        try:
-            data = next(data_iterator)
-
-        except Exception:
-            data = data_iterator
-            raise ValueError("error with data_iterator")
-    else:
-        data = None
-
-    data_b = tensor_parallel.broadcast_data(keys, data, datatype)
-    chunk_size = torch.min(data_b['context_len'])
-    retro_args = get_retro_args()
-    # two chunk retro has at least seq_len / 2 of chunk size
-    retro_args.retro_gpt_chunk_length = max(args.seq_length // 2, args.seq_length - chunk_size.item())
-
-    # Unpack.
-    tokens_ = data_b['text'].long()
-    labels = tokens_[:, 1:].contiguous()
-    tokens = tokens_[:, :-1].contiguous()
-
-    answer_mask = data_b["answer_mask"].float()[:, 1:].contiguous()
-
-    if args.retro_add_retriever:
-        neighbor_tokens = data_b['neighbor_tokens'].view(-1,
-                                                         retro_args.retro_gpt_retrieved_length).long()  # [bs * l * k, r]
-
-    # Get the masks and postition ids.
-    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
-        tokens,
-        tokenizer.eod,
-        args.reset_position_ids,
-        args.reset_attention_mask,
-        args.eod_mask_loss)
-
-    if args.answer_loss_only:
-        loss_mask = loss_mask * answer_mask
-
-    if args.retro_add_retriever:
-        _, _, neighbor_position_ids = get_ltor_masks_and_position_ids(
-            neighbor_tokens,
-            tokenizer.eod,
-            args.reset_position_ids,
-            args.reset_attention_mask,
-            args.eod_mask_loss)
-        neighbor_attention_mask = None
-        return tokens, labels, loss_mask, attention_mask, position_ids, \
-            neighbor_tokens, neighbor_attention_mask, neighbor_position_ids
-    else:
-        return tokens, labels, loss_mask, attention_mask, position_ids
-
-
-def loss_func(loss_mask, output_tensor):
-    losses = output_tensor.float()
-    loss_mask = loss_mask.view(-1).float()
-    loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
-
-    # Reduce loss for logging.
-    averaged_loss = average_losses_across_data_parallel_group([loss])
-
-    return loss, {'lm loss': averaged_loss[0]}
-
-
-def forward_step(data_iterator, model):
-    """Forward step."""
-    args = get_args()
-    timers = get_timers()
-
-    if args.retro_add_retriever:
-        timers('batch-generator', log_level=2).start()
-        tokens, labels, loss_mask, attention_mask, position_ids, \
-            neighbor_tokens, neighbor_attention_mask, neighbor_position_ids = get_batch(
-            data_iterator)
-        timers('batch-generator').stop()
-        output_tensor = model(tokens, position_ids, attention_mask,
-                              retriever_input_ids=neighbor_tokens,
-                              retriever_position_ids=neighbor_position_ids,
-                              retriever_attn_mask=neighbor_attention_mask,
-                              labels=labels)
-    else:
-        timers('batch-generator', log_level=2).start()
-        tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
-            data_iterator)
-        timers('batch-generator').stop()
-        output_tensor = model(tokens, position_ids, attention_mask,
-                              labels=labels)
-
-    return output_tensor, partial(loss_func, loss_mask)
-
-
-def train_valid_test_datasets_provider(train_val_test_num_samples):
-    """Build train, valid, and test datasets."""
-    args = get_args()
-    retro_args = get_retro_args()
-
-    tokenizer = get_tokenizer()
-
-    def fix_and_split_blend_pair(pair):
-        weight, name = pair
-        return [
-            [weight, os.path.join(args.data_folder, name, f"{name}_QA_train.json")],
-            [weight, os.path.join(args.data_folder, name, f"{name}_QA_dev.json")],
-            None,
-        ]
-
-    blend = [args.data_path[i:i+2] for i in range(0, len(args.data_path), 2)]
-
-    if len(blend) == 1:
-        blend_per_split =  [
-            os.path.join(args.data_folder, blend[0], f"{blend[0]}_QA_train.json"),
-            os.path.join(args.data_folder, blend[0], f"{blend[0]}_QA_dev.json"),
-            None,
-        ]
-    else:
-        blend_per_split = [
-            list(
-                reduce(
-                    lambda x, y: x + y,
-                    list(zip(*map(fix_and_split_blend_pair, blend)))[0]
-                )
-            ),
-            None,
-            None,
-        ]
-
-    blend_per_split = [get_blend_from_list(blend) for blend in blend_per_split]
-
-    extra_kwargs = {}
-
-    if args.retro_add_retriever:
-        dataset_cls = RetroJsonQADataset
-        config_cls = RetroJsonQADatasetConfig
-        extra_kwargs["retro_num_neighbors"] = args.retro_num_neighbors
-        extra_kwargs["retro_gpt_retrieved_length"] = retro_args.retro_gpt_retrieved_length
-    else:
-        dataset_cls = JsonQADataset
-        config_cls = JsonQADatasetConfig
-
-    config = config_cls(
-        random_seed=args.seed,
-        sequence_length=args.seq_length,
-        blend_per_split=blend_per_split,
-        split=args.split,
-        path_to_cache=args.data_cache_path,
-        tokenizer=tokenizer,
-        ft_neighbours=args.ft_neighbours,
-        bert_retriever_neighbours=args.bert_retriever_neighbours,
-        longform_answer=args.longform_answer,
-        inference_only=False,
-        retrieved_neighbours=False,
-        fix_newsqa=True,
-        mid_level_dataset_surplus=args.mid_level_dataset_surplus,
-        **extra_kwargs
-    )
-
-    print_rank_0('> building train, validation, and test datasets '
-                 'for GPT ...')
-    train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
-        dataset_cls,
-        train_val_test_num_samples,
-        is_dataset_built_on_rank,
-        config
-    ).build()
-    print_rank_0("> finished creating GPT datasets ...")
-
-    return train_ds, valid_ds, test_ds
-
-
-if __name__ == "__main__":
-
-    # Temporary for transition to core datasets
-    train_valid_test_datasets_provider.is_distributed = True
-
-    pretrain(train_valid_test_datasets_provider, partial(model_provider, gpt_builder),
-        ModelType.retro_decoder,  # ModelType.encoder_or_decoder,
-        forward_step,
-        extra_args_provider=get_tasks_args
-    )
diff --git a/tools/retro/sft/sft_retro_lm.sh b/tools/retro/sft/sft_retro_lm.sh
deleted file mode 100644
index 8c13f1052c1..00000000000
--- a/tools/retro/sft/sft_retro_lm.sh
+++ /dev/null
@@ -1,150 +0,0 @@
-#!/bin/bash
-# bash examples/qa/finetune_normal_lm.sh landrover_tasb_retrieved 843m 1 3e-6 1
-
-blend_name=$1
-model_size=$2
-global_bsz=$3
-lr=$4
-ft_neighbours=1
-model_card=pp1
-ckpt=$5
-TASK=none
-
-train_iters=1000
-
-
-DATA_HOME="<path/to/instruction/tuning/data/directory>"
-data_folder="$DATA_HOME"
-
-SFT_HOME="<path/to/megatron/repo>"
-
-TOKENIZER_MODEL="<path/to/gpt/tokenizer/model>"
-
-RETRO_WORKDIR="<path/to/retro/workdir>"
-
-K=2
-
-PRETRAINED_CHECKPOINT=${ckpt}
-
-SAVENAME="retro-${blend_name}_${model_card}_same_format_ctx${ft_neighbours}_${model_size}_${global_bsz}_${lr}"
-CHECKPOINT_PATH="${SFT_HOME}/checkpoints/applications/${SAVENAME}"
-TENSORBOARD_DIR="${SFT_HOME}/tensorboard/${SAVENAME}"
-mkdir -p ${TENSORBOARD_DIR}
-
-. ./tools/retro/sft/"${blend_name}".sh
-
-
-if [[ $model_size == "843m" ]]; then
-    # model param
-    mod_par=1
-    layers=24
-    hid_dim=1024
-    heads=16
-    pip_par=1
-
-    # node param
-    num_nodes=1
-    lr=5e-6
-    min_lr=5e-6
-fi
-
-
-GPT_ARGS="--apply-layernorm-1p \
-        --untie-embeddings-and-output-weights \
-        --disable-bias-linear \
-        --no-position-embedding \
-        --use-rotary-position-embeddings \
-        --rotary-percent 0.5 \
-        --swiglu \
-        --attention-dropout 0.0 \
-        --hidden-dropout 0.0 \
-        --pipeline-model-parallel-size $pip_par \
-        --tensor-model-parallel-size $mod_par \
-        --num-layers $layers \
-        --hidden-size $hid_dim \
-        --num-attention-heads $heads \
-        --seq-length 4096 \
-        --max-position-embeddings 4096 \
-        --lr-decay-style cosine \
-        --tokenizer-type GPTSentencePieceTokenizer \
-        --tokenizer-model ${TOKENIZER_MODEL} \
-        --clip-grad 1.0 \
-        --weight-decay 0.01 \
-        --adam-beta1 0.9 \
-        --adam-beta2 0.98 \
-        --log-params-norm \
-        --log-num-zeros-in-grad \
-        --bf16 \
-        --use-distributed-optimizer \
-"
-
-FT_ARGS="--eod-mask-loss \
-    --answer-loss-only \
-    --ft_neighbours ${ft_neighbours} \
-    --task $TASK"
-
-
-OUTPUT_ARGS="--log-interval 10 \
-             --save-interval 500 \
-             --eval-interval 200 \
-             --tensorboard-dir ${TENSORBOARD_DIR} \
-             --log-validation-ppl-to-tensorboard \
-             --eval-iters 100"
-
-options=" \
-    $GPT_ARGS \
-    --retro-workdir ${RETRO_WORKDIR} \
-    --retro-add-retriever \
-    --retro-num-neighbors ${K} \
-    --retro-attention-gate 0 \
-    --data-path ${DATA_BLEND} \
-    --data-folder ${data_folder} \
-    --recompute-activations \
-    --lr $lr \
-    --micro-batch-size 1 \
-    --global-batch-size ${global_bsz} \
-    --min-lr ${min_lr} \
-    --retro-cyclic-train-iters ${train_iters} \
-    --train-iters ${train_iters} \
-    --dataloader-type cyclic \
-    --save $CHECKPOINT_PATH \
-    $OUTPUT_ARGS \
-    $FT_ARGS"
-
-if [[ -d "$CHECKPOINT_PATH" ]]; then
-  options="$options \
-      --load $CHECKPOINT_PATH "
-else
-  echo $PRETRAINED_CHECKPOINT
-  options="$options \
-      --load $PRETRAINED_CHECKPOINT \
-      --finetune \
-      --no-load-rng \
-      --no-load-optim "
-fi
-
-######## Command. ########
-
-run_cmd="python -u ${SFT_HOME}/tools/retro/sft/sft_retro.py ${options}"
-
-export NCCL_DEBUG=INFO
-export NCCL_IB_TIMEOUT=19
-export NCCL_IB_SL=1
-export CUDA_DEVICE_MAX_CONNECTIONS=1
-
-NPROCS=8
-CMD="\
-    pwd && cd ${SFT_HOME} && pwd && \
-    export PYTHONPATH=$PYTHONPATH:${SFT_HOME} && \
-    python -m torch.distributed.run \
-    --nproc_per_node ${NPROCS} \
-    --nnodes 1 \
-    --node_rank 0 \
-    --master_port 6000 \
-    ${run_cmd} \
-"
-echo "~~~~~~~~~~~~~~~~~~~~~~~~~~"
-echo "CMD = '$CMD'."
-echo "~~~~~~~~~~~~~~~~~~~~~~~~~~"
-eval $CMD
-
diff --git a/tools/retro/text_generation/evaluate.py b/tools/retro/text_generation/evaluate.py
deleted file mode 100755
index 2031118cdc3..00000000000
--- a/tools/retro/text_generation/evaluate.py
+++ /dev/null
@@ -1,200 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-
-import sys
-import os
-from tqdm import tqdm
-import string
-import json
-import regex
-import numpy as np
-
-sys.path.append(os.path.abspath(os.path.join(
-    os.path.join(os.path.dirname(__file__), "../../../"))))
-from tools.retro.text_generation.metrics import F1Metric
-
-
-def normalize_answer(s):
-    def remove_articles(text):
-        return regex.sub(r'\b(a|an|the)\b', ' ', text)
-
-    def white_space_fix(text):
-        return ' '.join(text.split())
-
-    def remove_punc(text):
-        exclude = set(string.punctuation)
-        return ''.join(ch for ch in text if ch not in exclude)
-
-    def lower(text):
-        return text.lower()
-
-    return white_space_fix(remove_articles(remove_punc(lower(s))))
-
-
-def compute_f1_score(predicted_answers, groundtruth_answer, exp_name="default"):
-    """Evaluating F1 Score"""
-    print(len(predicted_answers), len(groundtruth_answer))
-    if len(predicted_answers) != len(groundtruth_answer):
-        groundtruth_answer = groundtruth_answer[:len(predicted_answers)]
-
-    guess_list = []
-    answer_list = []
-
-    assert len(guess_list) == len(answer_list), \
-        "lengths of guess and answer are different!"
-
-    for pred, ans in zip(predicted_answers, groundtruth_answer):
-        pred = pred.strip()
-        if type(ans) == str:
-            ans = ans.strip()
-        elif type(ans) == dict:
-            ans = ans['text'].strip()
-        elif ans == None:
-            continue
-        if "<|endoftext|>" in pred:
-            pred = pred.replace("<|endoftext|>", "")
-        if ans == "no_passages_used":
-            ans = ""
-        guess_list.append(pred)
-        answer_list.append(ans)
-
-    precision, recall, f1 = F1Metric.compute_all_pairs(guess_list, answer_list)
-    print('Method: %s; Precision: %.4f; recall: %.4f; f1: %.4f' % ( \
-        exp_name, precision, recall, f1))
-
-
-def load_groundtruth_file(data_file):
-    with open(data_file, "r") as f:
-        nq_examples = json.load(f)
-
-    data = []
-    for instance in nq_examples:
-        if "answers" in instance:
-            answers = instance["answers"]
-            if len(answers) < 1:
-                answers = [None]
-        elif "answer" in instance:
-            if type(instance["answer"]) is str:
-                answers = [instance["answer"]]
-            elif type(instance["answer"]) is list:
-                answers = instance["answer"]
-            else:
-                answers = [str(instance["answer"])]
-        else:
-            raise ValueError("need to have answer or answers")
-        data.append(answers[0])
-
-    return data
-
-
-def read_prediction(prediction_file):
-    prediction_list = []
-    print('reading %s' % prediction_file)
-    with open(prediction_file, "r") as f:
-        for i, line in enumerate(tqdm(f)):
-            if prediction_file.endswith("jsonl"):
-                line = json.loads(line)["pred"]
-                # print(line)
-            line = line.replace("Answer:", "")
-            line = line.replace("Answer: ", "")
-            line = line.replace('????  ', "")
-            line = line.replace('A: ', "")
-            line = line.replace("A:", "")
-
-            line = line.strip()
-
-            if "<|endoftext|>" in line:
-                line = line.replace("<|endoftext|>", "")
-            line = normalize_answer(line)  # normalize the answer
-            prediction_list.append(line)
-
-    return prediction_list
-
-
-def exact_match_score(prediction, ground_truth):
-    return normalize_answer(prediction) == normalize_answer(ground_truth)
-
-
-def ems(prediction, ground_truths):
-    return max([exact_match_score(prediction, gt) for gt in ground_truths])
-
-
-def evaluate_ems(prediction_file, ground_truth_file, dev_num=3000):
-    prediction_list = read_prediction(prediction_file)
-    ground_truths_list = []
-
-    if ground_truth_file.endswith(('txt', 'lst')):
-        raw_data = open(ground_truth_file, 'r')
-    else:
-        with open(ground_truth_file, 'r') as f:
-            raw_data = json.load(f)
-    if "dev" in ground_truth_file:
-        raw_data = raw_data[:dev_num]
-        prediction_list = prediction_list[:dev_num]
-
-    for each in raw_data:
-        if ground_truth_file.endswith('txt'):
-            each = json.loads(each)
-
-        if 'answers' in each:
-            ground_truths_list.append(each['answers'])
-        elif 'answer' in each:
-            ground_truths_list.append(each['answer'])
-        else:
-            ground_truths_list.append([each])
-
-    exactmatch = []
-
-    good_example_list = []
-    for i, each in enumerate(prediction_list):
-        score = ems(each, ground_truths_list[i])
-        exactmatch.append(score)
-        if score:
-            good_example_list.append(i)
-
-    final_em_score = np.mean(exactmatch)
-
-    print('Exact Match: %.4f;' % final_em_score)
-
-    print('done :-)')
-
-    return final_em_score, exactmatch
-
-
-def load_prediction(data_file):
-    data = []
-    with open(data_file, "r") as f:
-        for line in f.readlines():
-            data.append(line.strip())
-
-    return data
-
-
-def evaluate_f1(ground_truth_file, prediction_file, reduced_test_only=False):
-    groundtruth_answer = load_groundtruth_file(ground_truth_file)
-    predicted_answers = load_prediction(prediction_file)
-    if not reduced_test_only:
-        compute_f1_score(predicted_answers, groundtruth_answer)
-
-
-if __name__ == "__main__":
-    model_names = []
-    model_names += "retro-open_inst_pp1_same_format_ctx1_843m_128_5e-6",
-
-    for model_name in model_names:
-        ckpt_path = "/path/to/checkpoints/{}/".format(model_name)
-
-        n_ctx = 5
-        n_enc = 2
-        iter = 1000
-        model_param = "843m"
-
-        prediction_file = ckpt_path + "/retro-generate-nq_{}_{}_{}_test_greedy_0_20000_{}.txt".format(
-            n_ctx, n_enc, model_param, iter)
-        ground_truth_file = "/path/to/NQ/test.json"
-        print(prediction_file)
-        print(ground_truth_file)
-        evaluate_f1(ground_truth_file, prediction_file)
-        evaluate_ems(prediction_file, ground_truth_file)
-
-        print("=====================================")
diff --git a/tools/retro/text_generation/metrics.py b/tools/retro/text_generation/metrics.py
deleted file mode 100755
index bd0b5fe6b32..00000000000
--- a/tools/retro/text_generation/metrics.py
+++ /dev/null
@@ -1,80 +0,0 @@
-
-# The following code is adapted from
-# https://github.com/facebookresearch/ParlAI/blob/master/parlai/core/metrics.py, 
-# which is licensed under the MIT license. More details on the license can be 
-# found at https://github.com/facebookresearch/ParlAI/blob/master/LICENSE.
-
-"""Provides standard metric evaluations for dialog."""
-
-from collections import Counter
-from typing import List
-import numpy as np
-import re
-from nltk import ngrams
-
-re_art = re.compile(r'\b(a|an|the)\b')
-re_punc = re.compile(r'[!"#$%&()*+,-./:;<=>?@\[\]\\^`{|}~_\']')
-
-
-def normalize_answer(s):
-    """
-    Lower text and remove punctuation, articles and extra whitespace.
-    """
-    s = s.lower()
-    s = re_punc.sub(' ', s)
-    s = re_art.sub(' ', s)
-    s = ' '.join(s.split())
-    return s
-
-
-class F1Metric:
-    """
-    Helper class which computes token-level F1.
-    """
-
-    @staticmethod
-    def _prec_recall_f1_score(pred_items, gold_items):
-        """
-        Compute precision, recall and f1 given a set of gold and prediction items.
-        :param pred_items: iterable of predicted values
-        :param gold_items: iterable of gold values
-        :return: tuple (p, r, f1) for precision, recall, f1
-        """
-        common = Counter(gold_items) & Counter(pred_items)
-        num_same = sum(common.values())
-        if num_same == 0:
-            return 0, 0, 0
-        precision = 1.0 * num_same / len(pred_items)
-        recall = 1.0 * num_same / len(gold_items)
-        f1 = (2 * precision * recall) / (precision + recall)
-        return precision, recall, f1
-
-    @staticmethod
-    def compute_each_pair(guess: str, answer: str, n=1):
-        if answer == "":
-            return None, None, None
-        if guess == "":
-            return 0, 0, 0
-        g_tokens = normalize_answer(guess).split()
-        a_tokens = normalize_answer(answer).split()
-        g_tokens = list(ngrams(g_tokens, n))
-        a_tokens = list(ngrams(a_tokens, n))
-        precision, recall, f1 = F1Metric._prec_recall_f1_score(g_tokens, a_tokens)
-        return precision, recall, f1
-
-    @staticmethod
-    def compute_all_pairs(guesses: List[str], answers: List[str], n=1):
-        # additional augment:
-        print("guess:", len(guesses), ", answers:", len(answers))
-        assert len(guesses) == len(answers)
-
-        precision_list, recall_list, f1_list = [], [], []
-        for guess, answer in zip(guesses, answers):
-            precision, recall, f1 = F1Metric.compute_each_pair(guess, answer, n)
-            if precision is None or recall is None or f1 is None:
-                continue
-            precision_list.append(precision)
-            recall_list.append(recall)
-            f1_list.append(f1)
-
-        return np.mean(precision_list), np.mean(recall_list), np.mean(f1_list)
diff --git a/tools/retro/text_generation/retro_api.py b/tools/retro/text_generation/retro_api.py
deleted file mode 100644
index b70677485d4..00000000000
--- a/tools/retro/text_generation/retro_api.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-
-"""Inference API."""
-import numpy as np
-import torch
-from megatron.core import mpu
-from megatron.training import print_rank_0, get_retro_args, get_args, get_tokenizer
-from megatron.inference.text_generation.communication import broadcast_float_list, broadcast_tensor, broadcast_int_list
-from megatron.inference.text_generation.generation import (
-    score_and_return_on_first_stage)
-from tools.retro.text_generation.retro_generation import (
-    retro_generate_tokens_probs_and_return_on_first_stage)
-from megatron.inference.text_generation.tokenization import (
-    detokenize_generations)
-
-
-def tokenize_prompts(prompts=None, tokens_to_generate=None,
-                     add_BOS=None, rank=0):
-    """Tokenize prompts and make them avaiable on all ranks."""
-
-    # On all ranks set to None so we can pass them to functions
-    sizes_list = None
-    prompts_tokens_cuda_long_tensor = None
-    prompts_length_cuda_long_tensor = None
-
-    # On the specified rank, build the above.
-    if torch.distributed.get_rank() == rank:
-        assert prompts is not None
-        assert tokens_to_generate is not None
-        # Tensor of tokens padded and their unpadded length.
-        prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor = \
-            _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS)
-        # We need the sizes of these tensors for the boradcast
-        sizes_list = [prompts_tokens_cuda_long_tensor.size(0), # Batch size
-                      prompts_tokens_cuda_long_tensor.size(1)] # Sequence lenght
-
-    # First, broadcast the sizes.
-    sizes_tensor = broadcast_int_list(2, int_list=sizes_list, rank=rank)
-
-    # Now that we have the sizes, we can boradcast the tokens
-    # and length tensors.
-    sizes = sizes_tensor.tolist()
-    prompts_tokens_cuda_long_tensor = broadcast_tensor(
-        sizes, torch.int64, tensor=prompts_tokens_cuda_long_tensor, rank=rank)
-    prompts_length_cuda_long_tensor = broadcast_tensor(
-        sizes[0], torch.int64, tensor=prompts_length_cuda_long_tensor,
-        rank=rank)
-
-    return prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor
-
-
-def _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS):
-    """Given a set of prompts and number of tokens to generate:
-        - tokenize prompts
-        - set the sequence length to be the max of length of prompts
-          plus the number of tokens we would like to generate
-        - pad all the sequences to this length so we can convert them
-          into a 2D tensor.
-    """
-
-    # Tokenize all the prompts.
-    tokenizer = get_tokenizer()
-    if add_BOS:
-        prompts_tokens = [[tokenizer.eod] + tokenizer.tokenize(prompt)
-                          for prompt in prompts]
-    else:
-        prompts_tokens = [tokenizer.tokenize(prompt) for prompt in prompts]
-
-    # Now we have a list of list of tokens which each list has a different
-    # size. We want to extend this list to:
-    #   - incorporate the tokens that need to be generated
-    #   - make all the sequences equal length.
-    # Get the prompts length.
-    prompts_length = [len(prompt_tokens) for prompt_tokens in prompts_tokens]
-    # Get the max prompts length.
-    max_prompt_len = max(prompts_length)
-    # Set the tokens to generate to the max prompts length for Retro
-    args = get_args()
-    if args.retro_add_retriever:
-        tokens_to_generate = max_prompt_len
-    # Number of tokens in the each sample of the batch.
-    samples_length = max_prompt_len + tokens_to_generate
-    # Now update the list of list to be of the same size: samples_length.
-    for prompt_tokens, prompt_length in zip(prompts_tokens, prompts_length):
-        padding_size = samples_length - prompt_length
-        prompt_tokens.extend([tokenizer.eod] * padding_size)
-
-    # Now we are in a structured format, we can convert to tensors.
-    prompts_tokens_tensor = torch.cuda.LongTensor(prompts_tokens)
-    prompts_length_tensor = torch.cuda.LongTensor(prompts_length)
-
-    return prompts_tokens_tensor, prompts_length_tensor
-
-
-def retro_generate_and_post_process(model,
-                              prompts=None,
-                              neighbours_array=None,
-                              tokens_to_generate=0,
-                              return_output_log_probs=False,
-                              top_k_sampling=0,
-                              top_p_sampling=0.0,
-                              temperature=1.0,
-                              add_BOS=False,
-                              use_eod_token_for_early_termination=True,
-                              random_seed=-1,
-                              logits_mask=None):
-    """Run inference and post-process outputs, i.e., detokenize,
-    move to cpu and convert to list."""
-
-    # Main inference.
-    tokens, lengths, output_log_probs = retro_generate(
-        model,
-        prompts=prompts,
-        neighbours_array=neighbours_array,
-        tokens_to_generate=tokens_to_generate,
-        return_output_log_probs=return_output_log_probs,
-        top_k_sampling=top_k_sampling,
-        top_p_sampling=top_p_sampling,
-        temperature=temperature,
-        add_BOS=add_BOS,
-        use_eod_token_for_early_termination=use_eod_token_for_early_termination,
-        random_seed=random_seed,
-        logits_mask=logits_mask)
-
-    # Only post-process on first stage.
-    if mpu.is_pipeline_first_stage():
-        tokens, prompts_plus_generations, prompts_plus_generations_segments = \
-            detokenize_generations(tokens, lengths, True)
-
-        if return_output_log_probs:
-            output_log_probs = output_log_probs.cpu().numpy().tolist()
-            for i, (prob, seg) in enumerate(zip(output_log_probs, prompts_plus_generations_segments)):
-                output_log_probs[i] = prob[:len(seg) - 1]
-
-        return prompts_plus_generations, prompts_plus_generations_segments, \
-               output_log_probs, tokens
-
-    return None
-
-
-def retro_generate(model,
-             prompts=None,
-             neighbours_array=None,
-             tokens_to_generate=0,
-             return_output_log_probs=False,
-             top_k_sampling=0,
-             top_p_sampling=0.0,
-             temperature=1.0,
-             add_BOS=False,
-             use_eod_token_for_early_termination=True,
-             stop_on_double_eol=False,
-             stop_on_eol=False,
-             random_seed=-1,
-             logits_mask=None):
-    """Given prompts and input parameters, run inference and return:
-       tokens: prompts plus the generated tokens.
-       lengths: length of the prompt + generations. Note that we can
-           discard tokens in the tokens tensor that are after the
-           corresponding length.
-       output_log_probs: log probs of the tokens.
-    """
-
-    # Make sure input params are avaialble to all ranks.
-    values = [tokens_to_generate,
-              return_output_log_probs,
-              top_k_sampling, top_p_sampling,
-              temperature, add_BOS, use_eod_token_for_early_termination,
-              stop_on_double_eol,
-              stop_on_eol,
-              random_seed]
-    values_float_tensor = broadcast_float_list(10, float_list=values)
-    tokens_to_generate = int(values_float_tensor[0].item())
-    return_output_log_probs = bool(values_float_tensor[1].item())
-    top_k_sampling = int(values_float_tensor[2].item())
-    top_p_sampling = values_float_tensor[3].item()
-    temperature = values_float_tensor[4].item()
-    add_BOS = bool(values_float_tensor[5].item())
-    use_eod_token_for_early_termination = bool(values_float_tensor[6].item())
-    stop_on_double_eol = bool(values_float_tensor[7].item())
-    stop_on_eol = bool(values_float_tensor[8].item())
-    random_seed = int(values_float_tensor[9].item())
-
-    if random_seed != -1:
-        torch.random.manual_seed(random_seed)
-
-    # Tokenize prompts and get the batch.
-    # Note that these tensors are broadcaseted to all ranks.
-    if torch.distributed.get_rank() == 0:
-        assert prompts is not None
-
-    context_tokens_tensor, context_length_tensor = tokenize_prompts(
-        prompts=prompts, tokens_to_generate=tokens_to_generate, add_BOS=add_BOS)
-
-    retro_args = get_retro_args()
-    retro_args.retro_gpt_chunk_length = context_length_tensor.item()
-
-    retro_args = get_retro_args()
-    args = get_args()
-    r = retro_args.retro_gpt_retrieved_length
-    l = int(np.ceil(min(args.max_position_embeddings, context_tokens_tensor.size(1)) / retro_args.retro_gpt_chunk_length))
-    if torch.distributed.get_rank() == 0:
-        neighbours_array = neighbours_array.reshape(1, args.retro_num_neighbors, r).repeat(l, axis=0)  ## dim (l, k, r)
-
-    if tokens_to_generate == 0:
-        return score_and_return_on_first_stage(
-            model, context_tokens_tensor, context_length_tensor)
-
-    # Main inference function.
-    # Note that the outputs are available on the first stage.
-    return retro_generate_tokens_probs_and_return_on_first_stage(
-        model, context_tokens_tensor, context_length_tensor,
-        neighbours_array=neighbours_array,
-        return_output_log_probs=return_output_log_probs,
-        top_k=top_k_sampling,
-        top_p=top_p_sampling,
-        temperature=temperature,
-        use_eod_token_for_early_termination=use_eod_token_for_early_termination,
-        stop_on_double_eol=stop_on_double_eol,
-        stop_on_eol=stop_on_eol,
-        logits_mask=logits_mask)
\ No newline at end of file
diff --git a/tools/retro/text_generation/retro_generate.sh b/tools/retro/text_generation/retro_generate.sh
deleted file mode 100755
index 53f7d76476f..00000000000
--- a/tools/retro/text_generation/retro_generate.sh
+++ /dev/null
@@ -1,125 +0,0 @@
-#!/bin/bash
-
-TASK=$1
-model_size=$2
-sampling=$3
-split=$4
-gen_start=$5
-num_gen=$6
-ckpt_step=${7}
-ft_neighbours=${8}
-model_card=${9}
-ckpt=${10}
-K=${11}
-retrieve=${12}
-
-QA_HOME="<path/to/megatron/repo>"
-
-TOKENIZER_MODEL="<path/to/gpt/tokenizer/model>"
-
-RETRO_WORKDIR="<path/to/retro/workdir>"
-
-
-if [[ $model_size == "843m" ]]; then
-    mod_par=1
-    layers=24
-    hid_dim=1024
-    heads=16
-    pip_par=1
-fi
-
-GPT_ARGS="--apply-layernorm-1p \
-        --untie-embeddings-and-output-weights \
-        --disable-bias-linear \
-        --no-position-embedding \
-        --use-rotary-position-embeddings \
-        --rotary-percent 0.5 \
-        --swiglu \
-        --attention-dropout 0.0 \
-        --hidden-dropout 0.0 \
-        --pipeline-model-parallel-size $pip_par \
-        --tensor-model-parallel-size $mod_par \
-        --num-layers $layers \
-        --hidden-size $hid_dim \
-        --num-attention-heads $heads \
-        --seq-length 4096 \
-        --max-position-embeddings 4096 \
-        --lr-decay-style cosine \
-        --tokenizer-type GPTSentencePieceTokenizer \
-        --tokenizer-model ${TOKENIZER_MODEL} \
-        --clip-grad 1.0 \
-        --weight-decay 0.01 \
-        --adam-beta1 0.9 \
-        --adam-beta2 0.98 \
-        --log-params-norm \
-        --log-num-zeros-in-grad \
-        --bf16 \
-"
-
-
-sample_input_file="/path/to/instruct_tuning/data/$TASK/${split}.json"
-
-top_k=1
-micro_bsz=1
-SAMPLE_ARGS="--top_k $top_k"
-
-CHECKPOINT_PATH=${ckpt}
-sample_output_file="${CHECKPOINT_PATH}/retro-generate-${TASK}_${ft_neighbours}_${K}_${model_size}_${split}_${sampling}_${gen_start}_${num_gen}_${ckpt_step}.txt"
-
-DIR=`pwd`
-
-echo $sample_input_file
-echo $sample_output_file
-
-
-GEN_ARGS="$SAMPLE_ARGS \
-          --gen-start-idx $gen_start \
-          --num-gen $num_gen \
-          --ckpt-step ${ckpt_step} \
-          --sample-input-file $sample_input_file \
-          --sample-output-file $sample_output_file \
-          --retro-workdir ${RETRO_WORKDIR} \
-          --retro-add-retriever \
-          --retro-num-neighbors ${K} \
-          --reuse-top \
-          --retro-attention-gate 0 \
-          "
-
-if [[ $retrieve == 1 ]]; then
-    GEN_ARGS="$GEN_ARGS \
-          --use-retrieved-neighbours \
-          "
-fi
-
-FT_ARGS="--eod-mask-loss \
-    --answer-loss-only \
-    --ft_neighbours ${ft_neighbours} \
-    --task $TASK"
-
-DISTRIBUTED_ARGS="--nproc_per_node ${mod_par} \
-                  --nnodes ${pip_par} \
-                  --node_rank 0 \
-                  --master_port 8889"
-
-######## Command. ########
-
-COMMAND="python -m torch.distributed.run $DISTRIBUTED_ARGS ${DIR}/tools/retro/text_generation/retro_text_generation.py"
-
-COMMAND="$COMMAND \
-       $GPT_ARGS \
-       $GEN_ARGS \
-       --load $CHECKPOINT_PATH \
-       --micro-batch-size $micro_bsz \
-       $FT_ARGS"
-
-export NCCL_DEBUG=INFO
-export NCCL_IB_TIMEOUT=19
-export NCCL_IB_SL=1
-export CUDA_DEVICE_MAX_CONNECTIONS=1
-
-
-echo "~~~~~~~~~~~~~~~~~~~~~~~~~~"
-echo "CMD = '$CMD'."
-echo "~~~~~~~~~~~~~~~~~~~~~~~~~~"
-eval $COMMAND
-
diff --git a/tools/retro/text_generation/retro_generation.py b/tools/retro/text_generation/retro_generation.py
deleted file mode 100644
index f69103de772..00000000000
--- a/tools/retro/text_generation/retro_generation.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-
-"""Generation utilities."""
-import torch
-import torch.nn.functional as F
-from megatron.training import get_args, get_tokenizer
-from megatron.training import get_retro_args
-from megatron.core import mpu
-from megatron.training.utils import get_ltor_masks_and_position_ids, unwrap_model
-from megatron.inference.text_generation.communication import (
-    copy_from_last_to_first_pipeline_stage,
-    broadcast_from_last_pipeline_stage,
-    broadcast_from_last_to_first_pipeline_stage, broadcast_int_list, broadcast_tensor)
-from megatron.inference.text_generation.generation import _build_attention_mask_and_position_ids
-from megatron.inference.text_generation.sampling import sample
-
-
-
-def retro_generate_tokens_probs_and_return_on_first_stage(
-        model, tokens, lengths, neighbours_array=None,
-        return_output_log_probs=False,
-        top_k=0, top_p=0.0,
-        temperature=1.0,
-        use_eod_token_for_early_termination=True,
-        stop_on_double_eol=False,
-        stop_on_eol=False,
-        logits_mask=None):
-    """Main token generation function.
-
-    Args:
-        model: no interleaving is supported.
-        tokens: prompt tokens extended to be of size [b, max-sequence-length]
-        lengths: original prompt length, size: [b]
-        neighbours_array: neighbours array of size [b, l, k, r]
-        return_output_log_probs: flag to calculate the log probability of
-            the generated tokens. Note that the log probability is the one
-            from the original logit.
-        top_k, top_p: top-k and top-p sampling parameters.
-            Note that top-k = 1 is gready. Also, these paramters are
-            exclusive meaning that:
-                if top-k > 0 then we expect top-p=0.
-                if top-p > 0 then we check for top-k=0.
-        temperature: sampling temperature.
-        use_eod_token_for_early_termination: if True, do early termination if
-            all the sequences have reached this token.
-    Note: Outside of model, other parameters only need to be available on
-          rank 0.
-
-    Returns: Note that is size is adjusted to a lower value than
-             max-sequence-length if generation is terminated early.
-        tokens: prompt and generated tokens. size: [b, :]
-        generated_sequence_lengths: total length (including prompt) of
-            the generated sequence. size: [b]
-        output_log_probs: log probability of the selected tokens. size: [b, s]
-    """
-
-    args = get_args()
-    retro_args = get_retro_args()
-
-    tokenizer = get_tokenizer()
-
-    batch_size = tokens.size(0)
-    min_prompt_length = lengths.min().item()
-    max_sequence_length = tokens.size(1)
-    print("max_sequence_length", max_sequence_length)
-    print("min_prompt_length", min_prompt_length)
-    max_sequence_length = min(max_sequence_length, args.max_position_embeddings)
-
-    # If the context is too big, this happens
-    if min_prompt_length >= max_sequence_length:
-        raise ValueError("context length + tokens_to_generate too large")
-
-    # forward step.
-    unwrapped_model = unwrap_model(
-        model)
-    unwrapped_model.language_model.seq_length = max_sequence_length
-
-    # Added termination_id to support the case that we want to terminate the
-    # generation once that id is generated.
-    if hasattr(args, 'eos_id'):
-        termination_id = args.eos_id
-    else:
-        termination_id = tokenizer.eod
-
-    # ===================
-    # Pre-allocate memory
-    # ===================
-
-    # Log probability of the sequence (prompt + generated tokens).
-    output_log_probs = None
-    output_log_probs_size = (batch_size, max_sequence_length - 1)
-    # Lengths of generated seuquence including including prompts.
-    generated_sequence_lengths = None
-    if mpu.is_pipeline_last_stage():
-        if return_output_log_probs:
-            output_log_probs = torch.empty(output_log_probs_size,
-                                           dtype=torch.float32,
-                                           device=torch.cuda.current_device())
-        generated_sequence_lengths = torch.ones(
-            batch_size, dtype=torch.int64,
-            device=torch.cuda.current_device()) * max_sequence_length
-
-    # Whether we have reached a termination id.
-    is_generation_done = torch.zeros(batch_size, dtype=torch.uint8,
-                                     device=torch.cuda.current_device())
-
-    # =============
-    # Run infernece
-    # =============
-
-    with torch.no_grad():
-        attention_mask, position_ids = _build_attention_mask_and_position_ids(
-            tokens)
-        for context_length in range(min_prompt_length, max_sequence_length):
-            prev_context_length = 0
-            sizes_list = None
-            neighbor_tokens_cuda_long_tensor = None
-
-            # get the chunks for retrieval
-            if torch.distributed.get_rank() == 0:
-                neighbor_tokens = neighbours_array
-                neighbor_tokens_cuda_long_tensor = torch.cuda.LongTensor(
-                    neighbor_tokens.reshape((-1, retro_args.retro_gpt_retrieved_length)))
-                sizes_list = [neighbor_tokens_cuda_long_tensor.size(0),  # Batch size
-                              neighbor_tokens_cuda_long_tensor.size(1)]  # Sequence lenght
-            sizes_tensor = broadcast_int_list(2, int_list=sizes_list)
-            sizes = sizes_tensor.tolist()
-            neighbor_tokens_cuda_long_tensor = broadcast_tensor(
-                sizes, torch.int64, tensor=neighbor_tokens_cuda_long_tensor)
-
-            _, _, neighbor_position_ids = get_ltor_masks_and_position_ids(
-                neighbor_tokens_cuda_long_tensor,
-                tokenizer.eod,
-                args.reset_position_ids,
-                args.reset_attention_mask,
-                args.eod_mask_loss)
-            neighbor_attention_mask = None
-
-            # Pick the slice that we need to pass through the network.
-            tokens2use = tokens[:, prev_context_length:4096]
-            positions2use = position_ids[:, prev_context_length:4096]
-            attention_mask2use = attention_mask[
-                                 ..., prev_context_length:4096, :4096]
-
-            logits = model(tokens2use, positions2use, attention_mask2use,
-                           retriever_input_ids=neighbor_tokens_cuda_long_tensor,
-                           retriever_position_ids=neighbor_position_ids, retriever_attn_mask=neighbor_attention_mask,
-                           )
-
-            if mpu.is_pipeline_last_stage():
-                # Always the last stage should have an output.
-                assert logits is not None
-
-                # Sample.
-                last_token_logits = logits[:, context_length - 1, :]
-                # last_token_logits = logits[:, -1, :]
-
-                # word banning
-                if logits_mask is not None:
-                    last_token_logits[:, logits_mask] = float('-Inf')
-
-                new_sample = sample(last_token_logits,
-                                    top_k=top_k,
-                                    top_p=top_p,
-                                    temperature=temperature,
-                                    vocab_size=tokenizer.vocab_size)
-
-                # If a prompt length is smaller or equal th current context
-                # length, it means we have started generating tokens
-                started = lengths <= context_length
-                # Update the tokens.
-                tokens[started, context_length] = new_sample[started]
-
-                # Calculate the log probabilities.
-                if return_output_log_probs:
-                    log_probs = F.log_softmax(logits, dim=2)
-                    if return_output_log_probs:
-                        # Pick the tokens that we need to get the log
-                        # probabilities for. Note that next input token is
-                        # the token which we selected in the current logits,
-                        # so shift by 1.
-                        indices = torch.unsqueeze(
-                            tokens[
-                            :,
-                            (prev_context_length + 1):(context_length + 1)],
-                            2)
-                        output_log_probs[:,
-                        prev_context_length:context_length] = \
-                            torch.gather(log_probs, 2, indices).squeeze(2)
-
-            # Update the tokens on the first stage so the next input to
-            # the network is correct.
-            copy_from_last_to_first_pipeline_stage(batch_size, torch.int64,
-                                                   tokens[:, context_length])
-
-            # Update the context length for the next token generation.
-            prev_context_length = context_length
-
-            # Check if all the sequences have hit the termination_id.
-            done = None
-            if mpu.is_pipeline_last_stage():
-                # TODO(rprenger) These stopping methods are tokenizer dependent
-                # instead tokenization should be in the inference loop so stop sequences can be used
-                if stop_on_double_eol:
-                    hit_double_eol = (new_sample == 628).byte() & started.byte()
-                    hit_two_eols = (new_sample == 198).byte() & (
-                            tokens[:, context_length - 1] == 198).byte() & started.byte()
-                    done_token = hit_double_eol | hit_two_eols
-                elif stop_on_eol:
-                    hit_double_eol = (new_sample == 628).byte() & started.byte()
-                    hit_eol = (new_sample == 198).byte() & started.byte()
-                    done_token = hit_double_eol | hit_eol
-                elif context_length > min_prompt_length + 64:  # previous retrov1 limitations
-                    done_token = 1
-                else:
-                    done_token = (new_sample == termination_id).byte() & \
-                                 started.byte()
-
-                just_finished = (done_token & ~is_generation_done).bool()
-                generated_sequence_lengths[just_finished.view(-1)] = \
-                    context_length + 1
-                is_generation_done = is_generation_done | done_token
-                done = torch.all(is_generation_done)
-            done = broadcast_from_last_pipeline_stage(1, torch.uint8,
-                                                      tensor=done)
-            if use_eod_token_for_early_termination and done:
-                break
-
-    # ===================================================
-    # Update the length of based on max generated length.
-    # ===================================================
-
-    tokens = tokens[:, :(context_length + 1)]
-    if mpu.is_pipeline_last_stage():
-        if return_output_log_probs:
-            output_log_probs = output_log_probs[:, :context_length]
-
-    # ======================================
-    # Broadcast to the first pipeline stage.
-    # ======================================
-
-    generated_sequence_lengths = broadcast_from_last_to_first_pipeline_stage(
-        batch_size, torch.int64, generated_sequence_lengths)
-    if return_output_log_probs:
-        output_log_probs_size = (batch_size, context_length)
-        output_log_probs = broadcast_from_last_to_first_pipeline_stage(
-            output_log_probs_size, torch.float32, output_log_probs)
-
-    return tokens, generated_sequence_lengths, output_log_probs
diff --git a/tools/retro/text_generation/retro_text_generation.py b/tools/retro/text_generation/retro_text_generation.py
deleted file mode 100755
index 27050090446..00000000000
--- a/tools/retro/text_generation/retro_text_generation.py
+++ /dev/null
@@ -1,263 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
-"""Sample Generate GPT"""
-import torch
-import os
-import sys
-from typing import Union
-
-sys.path.append(os.path.abspath(os.path.join(
-    os.path.join(os.path.dirname(__file__), "../../../"))))
-from megatron.training import get_args, get_retro_args
-from megatron.training import print_rank_0
-from megatron.training import get_tokenizer
-from megatron.training.checkpointing import load_checkpoint
-from megatron.training.initialize import initialize_megatron
-from megatron.core.models.gpt import GPTModel
-from megatron.training import get_model
-from tools.retro.text_generation.retro_api import retro_generate_and_post_process
-from tools.retro.sft.sft_retro import get_tasks_args
-from tools.retro.sft.dataset_conv import reformat_prompt, preprocess, reformat_prompt_short
-import numpy as np
-import time
-import megatron.legacy.model
-from megatron.training.arguments import core_transformer_config_from_args
-
-
-
-def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megatron.legacy.model.GPTModel]:
-    """Builds the model.
-
-    Args:
-        pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
-        post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
-
-
-    Returns:
-        Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model
-    """
-    print_rank_0('building GPT model ...')
-    args = get_args()
-    config = core_transformer_config_from_args(args)
-
-    assert args.use_legacy_models, 'retro text generation only implemented for legacy models'
-    
-    # not support core model yet
-    model = megatron.legacy.model.GPTModel(
-        config,
-        num_tokentypes=0,
-        parallel_output=False,
-        pre_process=pre_process,
-        post_process=post_process
-    )
-
-    return model
-
-
-def pad_neighbours_for_query_only(args, nb_tokens, pad_id, ft_neighbours):
-    # take top k neighbours and padding
-    neighbours_tokens = []
-    retro_args = get_retro_args()
-    r = retro_args.retro_gpt_retrieved_length
-
-    if args.reuse_top:
-        valid_nb_tokens = nb_tokens[:args.retro_num_neighbors]
-    else:
-        valid_nb_tokens = nb_tokens[ft_neighbours:args.retro_num_neighbors + ft_neighbours]
-
-    for nb_token in valid_nb_tokens:
-        if len(nb_token) >= r:
-            nb_token = nb_token[:r]
-        else:
-            nb_token = nb_token + [pad_id] * (r - len(nb_token))
-        neighbours_tokens.append(nb_token)
-    print("len(nb_tokens)", len(nb_tokens))
-    print("len(neighbours_tokens)", len(neighbours_tokens))
-    print("args.retro_num_neighbors", args.retro_num_neighbors)
-
-    if len(neighbours_tokens) < args.retro_num_neighbors:
-        assert ValueError("neighbours are not enough, add empty ones and create mask for those empty ones")
-    neighbours_tokens = np.array(neighbours_tokens)
-    return neighbours_tokens
-
-
-def add_text_generate_args(parser):
-    """Text generation arguments."""
-
-    parser = get_tasks_args(parser)
-    group = parser.add_argument_group(title='text generation')
-
-    group.add_argument("--temperature", type=float, default=1.0,
-                       help='Sampling temperature.')
-    group.add_argument("--greedy", action='store_true', default=False,
-                       help='Use greedy sampling.')
-    group.add_argument("--top_p", type=float, default=0.0,
-                       help='Top p sampling.')
-    group.add_argument("--top_k", type=int, default=0,
-                       help='Top k sampling.')
-    group.add_argument("--out-seq-length", type=int, default=256,
-                       help='Size of the output generated text.')
-    group.add_argument("--sample-input-file", type=str, default=None,
-                       help='Get input from file instead of interactive mode, '
-                            'each line is an input.')
-    group.add_argument("--sample-output-file", type=str, default=None,
-                       help='Output file got from --sample-input-file')
-    group.add_argument("--num-samples", type=int, default=0,
-                       help='Number of samples to generate unconditionally, '
-                            'defaults to 0 and interactive conditional sampling')
-    group.add_argument("--genfile", type=str,
-                       help='Output file when generating unconditionally')
-    group.add_argument("--recompute", action='store_true',
-                       help='During generation recompute all attention '
-                            'instead of using previously computed keys/values.')
-    group.add_argument("--epsilon", type=float, default=0.01,
-                       help="Minimum factor by which each probability is multiplied")
-    group.add_argument("--debug-gen", action='store_true',
-                       help="If set, additional debugging output is printed to stdout")
-    group.add_argument('--length-penalty', type=float, default=1.0,
-                       help='length penalty')
-    group.add_argument('--gen-start-idx', type=int, default=0,
-                       help='project size for adapters')
-    group.add_argument('--num-gen', type=int, default=-1,
-                       help='project size for adapters')
-    group.add_argument('--ckpt-step', type=int, default=None,
-                       help='setting ckpt step manually')
-    group.add_argument("--short-format", action='store_true',
-                       help='Use short format QA')
-    group.add_argument("--use-retrieved-neighbours", action='store_true', default=False,
-                       help='Use retrieved neighbours')
-    group.add_argument('--template-id', type=int, default=0,
-                       help='template id for generation,')
-    return parser
-
-
-def generate_samples_conditional(model):
-    args = get_args()
-    start = time.time()
-    avg_time = []
-    tokenizer = get_tokenizer()
-    model.eval()
-    if torch.distributed.get_rank() == 0:
-
-        data = preprocess(args.sample_input_file, inference_only=True,
-                          retrieved_neighbours=args.use_retrieved_neighbours)
-        print("total rows {}".format(len(data)))
-        all_data = data[args.gen_start_idx:]  # start from gen_start_idx
-        if args.num_gen > 0:
-            all_data = all_data[:args.num_gen]
-        input_count = len(all_data)
-        input_pos = 0
-
-    terminate_runs = 0
-    while True:
-        torch.distributed.barrier()
-        if torch.distributed.get_rank() == 0:
-            sentences = []
-            n_arrays = []
-            print("global batch size", args.global_batch_size)
-            for _ in range(args.global_batch_size):
-                print(input_pos)
-                if input_pos >= input_count:
-                    print("reach the last row")
-                    break
-                else:
-                    sample = all_data[input_pos]
-                input_pos += 1
-
-                if True:
-                    max_target_len = args.out_seq_length
-                    query, _, neighbours = sample
-
-                    neighbours_array = pad_neighbours_for_query_only(args,
-                                                                     [tokenizer.tokenize(neighbour) for neighbour in
-                                                                      neighbours], tokenizer.eod, args.ft_neighbours)
-                    print("neighbours_array.shape", neighbours_array.shape)
-                    tokenizer = get_tokenizer()
-
-                    if args.short_format:
-                        input_tokens = reformat_prompt_short(query, neighbours, args.task, args.ft_neighbours,
-                                                             max_target_len,
-                                                             tokenizer, args.seq_length)
-                    else:
-                        input_tokens = reformat_prompt(query, neighbours, args.task, args.ft_neighbours, max_target_len,
-                                                       tokenizer, args.seq_length, template_id=args.template_id)
-                    raw_text = tokenizer.detokenize(input_tokens)
-                    print(raw_text)
-                else:
-                    raise ValueError("invalid arg for task")
-                sentences.append(raw_text)
-            retro_args = get_retro_args()
-
-            resp_sentences, resp_sentences_seg, scores, \
-                tokens = retro_generate_and_post_process(model, prompts=sentences,
-                                                         neighbours_array=neighbours_array,
-                                                         tokens_to_generate=args.seq_length - retro_args.retro_gpt_chunk_length,
-                                                         return_output_log_probs=False,
-                                                         top_k_sampling=args.top_k,
-                                                         top_p_sampling=args.top_p,
-                                                         add_BOS=False,
-                                                         temperature=1.0)
-            print("len of resp_sentences", len(resp_sentences))
-            for prompt, generation in zip(sentences, resp_sentences):
-                datum = generation[len(prompt):]
-                print("prompt:", generation[:len(prompt)])
-                if "<|endoftext|>" in datum:
-                    datum = datum[:datum.find("<|endoftext|>")].strip()
-                datum = datum.replace("\n", " ")
-                print("cont:", datum)
-                yield datum
-            avg_time.append((time.time() - start) / args.global_batch_size)
-            print("avg time for each sample: ", sum(avg_time) / len(avg_time))
-            start = time.time()
-            if input_pos >= input_count:
-                print("finish all lines")
-                terminate_runs = 1
-        else:
-            retro_generate_and_post_process(model)
-
-        terminate_runs_tensor = torch.cuda.LongTensor([terminate_runs])
-        torch.distributed.broadcast(terminate_runs_tensor, 0)
-        terminate_runs = terminate_runs_tensor[0].item()
-
-        if terminate_runs == 1:
-            return
-
-
-def generate_and_write_samples_conditional(model):
-    args = get_args()
-    if args.sample_output_file is None:
-        sample_output_file = args.sample_input_file + ".out"
-        print('`sample-output-file` not specified, setting '
-              'it to {}'.format(sample_output_file))
-    else:
-        sample_output_file = args.sample_output_file
-    with open(sample_output_file, 'w') as f:
-        for datum in generate_samples_conditional(model):
-            if torch.distributed.get_rank() == 0:
-                f.write(datum + '\n')
-
-
-def main():
-    """Main program."""
-
-    initialize_megatron(extra_args_provider=add_text_generate_args,
-                        args_defaults={'no_load_rng': True,
-                                       'no_load_optim': True})
-
-    # Set up model and load checkpoint
-    model = get_model(model_provider, wrap_with_ddp=False)
-    print(model)
-    args = get_args()
-
-    if args.load is not None:
-        _ = load_checkpoint(model, None, None)
-    model = model[0]
-
-    # Generate samples.
-    if args.sample_input_file is not None:
-        print(f"{args.sample_input_file}")
-        generate_and_write_samples_conditional(model)
-
-
-if __name__ == "__main__":
-    main()

From c22615e437f537613ec9f55c683139e798a16560 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Wed, 28 Jan 2026 20:18:48 +0100
Subject: [PATCH 55/79] ci: Mark test_compatible_with_nd_parallel as flaky
 (#3122)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .../megatron_fsdp/test_mcore_fully_sharded_data_parallel.py      | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py
index 77274ec4d50..d4c664cda9c 100644
--- a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py
+++ b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py
@@ -623,6 +623,7 @@ def _training_loop(seed=42, **kwargs):
 
         return outputs
 
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(
         not is_torch_min_version("2.4.0"), reason="Test needs to be updated for torch >= 2.4.0"
     )

From a883e96bf4727f3c79b717af98dfac9161bc3df7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Wed, 28 Jan 2026 20:27:43 +0100
Subject: [PATCH 56/79] build: Use merge-commit-sha for container (#3123)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/workflows/cicd-main.yml | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
index 4bca7f7882e..bf8c7b85941 100644
--- a/.github/workflows/cicd-main.yml
+++ b/.github/workflows/cicd-main.yml
@@ -202,8 +202,28 @@ jobs:
       && needs.pre-flight.outputs.is_merge_group == 'false'
       && !cancelled()
     steps:
+      - name: Get PR info
+        id: get-pr-info
+        if: startsWith(github.ref, 'refs/heads/pull-request/')
+        uses: nv-gha-runners/get-pr-info@main
+
+      - name: Get merge commit sha
+        shell: bash -x -e -u -o pipefail {0}
+        id: sha
+        env:
+          IS_PR: ${{ startsWith(github.ref, 'refs/heads/pull-request/') }}
+        run: |
+          if [[ "$IS_PR" == "true" ]]; then
+            SHA=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').merge_commit_sha }}
+          else
+            SHA=${GITHUB_SHA}
+          fi
+          echo "main=${SHA}" | tee -a "$GITHUB_OUTPUT"
+
       - name: Checkout
         uses: actions/checkout@v4
+        with:
+          ref: ${{ steps.sha.outputs.main }}
 
       - name: Setup python
         uses: actions/setup-python@v5
@@ -216,11 +236,6 @@ jobs:
           apt-get update
           apt-get install -y gh
 
-      - name: Get PR info
-        id: get-pr-info
-        if: startsWith(github.ref, 'refs/heads/pull-request/')
-        uses: nv-gha-runners/get-pr-info@main
-
       - name: Has lts label
         id: has-lts-label
         env:

From e2ff203a3b7a11f6eceb0889e25b7320b9c4cfe7 Mon Sep 17 00:00:00 2001
From: GitHub Actions <github-actions[bot]@users.noreply.github.com>
Date: Thu, 29 Jan 2026 00:13:53 +0000
Subject: [PATCH 57/79] Update copy-pr-bot.yaml [skip ci]

---
 .github/copy-pr-bot.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml
index e8e5bcd69cb..477f87d90f5 100644
--- a/.github/copy-pr-bot.yaml
+++ b/.github/copy-pr-bot.yaml
@@ -1,4 +1,4 @@
 enabled: true
 auto_sync_draft: false
 auto_sync_ready: true
-trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]
+trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "prajwal1210", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]

From 50132f25315a1b6d7b179af3c9c752e86dd8d0d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Wed, 28 Jan 2026 21:29:49 +0100
Subject: [PATCH 58/79] ci: Add unit tests to merge queue (#3125)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/workflows/cicd-main.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
index bf8c7b85941..d69985bff40 100644
--- a/.github/workflows/cicd-main.yml
+++ b/.github/workflows/cicd-main.yml
@@ -399,8 +399,8 @@ jobs:
         success()
         || needs.pre-flight.outputs.is_ci_workload == 'true'
         || needs.pre-flight.outputs.force_run_all == 'true'
+        || needs.pre-flight.outputs.is_merge_group == 'true'
       )
-      && needs.pre-flight.outputs.is_merge_group == 'false'
       && !cancelled()
     outputs:
       integration-tests: ${{ steps.main.outputs.integration-tests }}

From 42986ace1a07da57d6de926cf9a248e5e55bfd6a Mon Sep 17 00:00:00 2001
From: helen ngo <helenn@nvidia.com>
Date: Wed, 28 Jan 2026 16:13:52 -0500
Subject: [PATCH 59/79] Refactor `rl_offload_kv_cache_during_training` to
 offload KV cache to CPU while retaining fixed virtual address  (#3048)

---
 .../inference/gpt/gpt_dynamic_inference.py    |  1 +
 .../inference/contexts/dynamic_context.py     | 44 ++++++++++++++-----
 .../core/inference/engines/dynamic_engine.py  |  8 ++++
 megatron/rl/inference/megatron.py             |  3 +-
 megatron/rl/rl_utils.py                       | 12 +++--
 megatron/training/arguments.py                | 11 ++++-
 6 files changed, 61 insertions(+), 18 deletions(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index 679dd78b42b..88b744b3ac0 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -191,6 +191,7 @@ def get_inference_context(
         cuda_graph_max_tokens=args.inference_dynamic_batching_cuda_graph_max_tokens,
         cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
         metrics_writer=metrics_writer,
+        offload_kv_cache=args.rl_offload_kv_cache_during_training
     )
 
     return context
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index c6b30f47f78..5dc2d503097 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -58,6 +58,14 @@
 except ImportError:
     HAVE_FLASHINFER = False
 
+try:
+    from torch_memory_saver import torch_memory_saver
+
+    torch_memory_saver.hook_mode = "torch"
+    HAVE_TORCH_MEMORY_SAVER = True
+except ImportError:
+    HAVE_TORCH_MEMORY_SAVER = False
+
 try:
     import wandb  # pylint: disable=unused-import
 
@@ -286,6 +294,7 @@ def __init__(
         metrics_writer: Optional['WandbModule'] = None,
         request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None,
         persist_cuda_graphs: Optional[bool] = False,
+        offload_kv_cache: Optional[bool] = False,
     ):
         super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
 
@@ -542,6 +551,12 @@ def __init__(
             )
         )
 
+        # Whether to offload the KV cache. Determines where the KV cache is allocated within memory.
+        self.offload_kv_cache = offload_kv_cache
+        assert not (
+            self.offload_kv_cache and self.unified_memory_level
+        ), "The KV cache should not be instantiated in unified memory when it is offloaded during training."
+
         self._using_cuda_graph_this_step = False
         self.use_cuda_graphs_for_non_decode_steps = use_cuda_graphs_for_non_decode_steps
         # Deal with chunked prefill
@@ -657,19 +672,26 @@ def allocate_memory_buffer():
                     device=torch.cuda.current_device(),
                 )
             else:
-                self.memory_buffer = torch.empty(
-                    (
-                        2,  # key and value
-                        self.num_attention_layers,
-                        self.block_allocator.total_count,
-                        self.block_size_tokens,
-                        self.num_attention_heads_per_partition,
-                        self.hidden_size_per_attention_head,
-                    ),
-                    dtype=self.params_dtype,
-                    device=torch.cuda.current_device(),
+                ctx = (
+                    torch_memory_saver.region(tag="kv_cache", enable_cpu_backup=True)
+                    if HAVE_TORCH_MEMORY_SAVER and self.offload_kv_cache
+                    else nullcontext()
                 )
 
+                with ctx:
+                    self.memory_buffer = torch.empty(
+                        (
+                            2,  # key and value
+                            self.num_attention_layers,
+                            self.block_allocator.total_count,
+                            self.block_size_tokens,
+                            self.num_attention_heads_per_partition,
+                            self.hidden_size_per_attention_head,
+                        ),
+                        dtype=self.params_dtype,
+                        device=torch.cuda.current_device(),
+                    )
+
         # Optional state tensors for hybrid models
         def allocate_mamba_states():
             """Allocate Mamba states. This function is called below within
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index c42246d5624..e23520cf65f 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -90,6 +90,11 @@
 except ImportError:
     HAVE_PSUTIL = False
 
+from megatron.core.inference.contexts.dynamic_context import HAVE_TORCH_MEMORY_SAVER
+
+if HAVE_TORCH_MEMORY_SAVER:
+    from torch_memory_saver import torch_memory_saver
+
 
 class EngineSuspendedError(Exception):
     """Engine is currently suspended and not performing steps."""
@@ -338,6 +343,9 @@ def create_cuda_graphs(self, reset_context: bool = True):
 
         self.capture_stats = capture_stats
 
+        if HAVE_TORCH_MEMORY_SAVER:
+            torch_memory_saver.pause("kv_cache")
+
     @internal_api
     async def start_listening_to_data_parallel_coordinator(
         self,
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 73ab5024a64..48c02774fca 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -166,7 +166,8 @@ def get_dynamic_inference_engine(
         cuda_graph_max_tokens=args.inference_dynamic_batching_cuda_graph_max_tokens,
         cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
         metrics_writer=metrics_writer,
-        persist_cuda_graphs=args.rl_training_cuda_graphs
+        persist_cuda_graphs=args.rl_training_cuda_graphs,
+        offload_kv_cache=args.rl_offload_kv_cache_during_training
     )
 
     inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context, pg_collection=pg_collection)
diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index 5d6b3b77653..26b990236ba 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -84,6 +84,9 @@
     is_batch_invariant_mode_enabled,
 )
 
+from megatron.core.inference.contexts.dynamic_context import HAVE_TORCH_MEMORY_SAVER
+if HAVE_TORCH_MEMORY_SAVER:
+    from torch_memory_saver import torch_memory_saver
 
 logger = logging.getLogger(__name__)
 
@@ -1617,9 +1620,9 @@ def megatron_rl_inference_mode(
 
         with nvtx_range("onload-kv-cache-before-inference"):
             if offload_kv_cache_during_training:
-                assert (
-                    reset_cuda_graphs
-                ), "reset_cuda_graphs must be True when offloading kv cache during training"
+                # Restore the KV cache by re-binding physical pages to a consistent virtual address
+                torch_memory_saver.resume("kv_cache")
+
                 logger.debug(
                     f"[{dist.get_rank()}] Restoring kv cache ({inference_interface._inference_engine.context.memory_buffer.numel() / 1024**3:.2f} GB) to GPU"
                 )
@@ -1654,7 +1657,8 @@ def megatron_rl_inference_mode(
                 logger.debug(
                     f"[{dist.get_rank()}] Offloading kv cache ({kv_cache.numel() * kv_cache.element_size() / 1024**3:.2f} GB) to CPU"
                 )
-                inference_interface._inference_engine.context.memory_buffer = kv_cache.cpu()
+                torch_memory_saver.pause("kv_cache")
+
             elif remove_kv_cache_during_training:
                 inference_interface._inference_engine.context.memory_buffer = None
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 007a9842610..dd4cf5b5c89 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -328,6 +328,14 @@ def validate_args(args, defaults={}):
     total_model_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size
     args.data_parallel_size = args.world_size // total_model_size
 
+    # Assert that `torch_memory_saver` is installed if offloading KV cache during RL.
+    if args.rl_offload_kv_cache_during_training:
+        try:
+            from torch_memory_saver import torch_memory_saver
+        except ImportError:
+            raise AssertionError("To use offload-kv-cache-during-training, `torch_memory_saver` must be installed. See https://github.com/fzyzcjy/torch_memory_saver.")
+        assert not args.inference_dynamic_batching_unified_memory_level, "The KV cache should not be instantiated in unified memory when it is offloaded during training."
+
     # Batch size checks if running RL.
     if args.perform_rl_step:
         assert not (args.rl_remove_kv_cache_during_training and args.rl_offload_kv_cache_during_training), \
@@ -1992,8 +2000,7 @@ def _add_rl_args(parser):
                             'round-robin: distribute bins cyclically across ranks for better load balancing')
     group.add_argument('--rl-training-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool,
                        default=False,
-                       help='If set, do not call `delete_cuda_graphs` or `toggle_cuda_graphs` when the inference engine is suspended. '
-                            'Use only when all training and inference cudagraphs and the KV cache fit on device.')
+                       help='If set, do not call `delete_cuda_graphs` or `toggle_cuda_graphs` when the inference engine is suspended.')
     group.add_argument('--rl-inference-tensor-model-parallel-size', type=int, default=None,
                        help='Degree of tensor model parallelism for inference for RL.')     
     group.add_argument(

From d41bf66191622d508f327c7c7fc539604e97ec4e Mon Sep 17 00:00:00 2001
From: Philip Petrakian <ppetrakian@nvidia.com>
Date: Wed, 28 Jan 2026 14:56:33 -0800
Subject: [PATCH 60/79] Disable Greptile status comments (#3127)

---
 greptile.json | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/greptile.json b/greptile.json
index e08c2def387..38013ea8869 100644
--- a/greptile.json
+++ b/greptile.json
@@ -35,5 +35,6 @@
       "included": false,
       "collapsible": false,
       "defaultOpen": false
-    }
+    },
+    "statusCommentsEnabled": false
   }
\ No newline at end of file

From 9f05aac7dcd52ced62d72dd9acc933bf23770660 Mon Sep 17 00:00:00 2001
From: Philip Petrakian <ppetrakian@nvidia.com>
Date: Wed, 28 Jan 2026 16:59:44 -0800
Subject: [PATCH 61/79] Create CodeRabbit config (#3131)

---
 .coderabbit.yaml | 62 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100644 .coderabbit.yaml

diff --git a/.coderabbit.yaml b/.coderabbit.yaml
new file mode 100644
index 00000000000..f20c23e1fe6
--- /dev/null
+++ b/.coderabbit.yaml
@@ -0,0 +1,62 @@
+# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json
+language: "en-US"
+
+# Instruct the AI to be highly selective - only critical issues
+tone_instructions: |
+  IMPORTANT: Be EXTREMELY restrictive with comments. You should almost never comment.
+
+  ONLY comment if the issue is Critical (🔴) or Major (🟠) severity.
+  NEVER make Minor (🟡), Trivial (🔵), or Info (⚪) severity comments.
+
+  The ONLY acceptable comments are:
+  - Critical bugs: null pointer dereferences, infinite loops, security vulnerabilities, data corruption
+  - Major bugs: logic errors that break functionality, race conditions, resource leaks
+  - Obvious typos in error messages or user-facing strings that would confuse users
+  - Forgotten updates: docstrings that directly contradict the code behavior
+
+  Do NOT comment on (even if you think it's helpful):
+  - Style, formatting, or naming conventions
+  - Refactoring opportunities
+  - Performance suggestions
+  - Missing docstrings or comments
+  - Code organization
+  - Type hints or annotations
+  - Import ordering
+  - Any "minor" improvements
+
+  If you would label a comment as "Minor" severity or lower, DO NOT POST IT.
+  When in doubt, stay silent. Aim for 0-2 comments per PR maximum.
+
+reviews:
+  # Use chill profile - filters out nitpicks automatically
+  profile: "chill"
+
+  # Disable all summary features
+  high_level_summary: false
+  high_level_summary_in_walkthrough: false
+
+  # Disable walkthrough comment entirely
+  collapse_walkthrough: true
+  changed_files_summary: false
+  sequence_diagrams: false
+
+  # Disable status/effort estimates
+  review_status: false
+  commit_status: false
+  estimate_code_review_effort: false
+
+  # Disable auto-suggestions for labels/reviewers
+  suggested_labels: false
+  suggested_reviewers: false
+
+  # Disable related issues/PRs lookup
+  assess_linked_issues: false
+  related_issues: false
+  related_prs: false
+
+  # Auto-review disabled - only review when explicitly requested via @coderabbitai review
+  auto_review:
+    enabled: false
+
+chat:
+  auto_reply: true

From f0b1cb2544dbe2d264253f272599fb87732ea7ef Mon Sep 17 00:00:00 2001
From: Charlie Truong <chtruong@nvidia.com>
Date: Wed, 28 Jan 2026 21:55:51 -0600
Subject: [PATCH 62/79] build: Explicitly set minimum torch version to >= 2.6.0
 (#3085)

Signed-off-by: Charlie Truong <chtruong@nvidia.com>
---
 pyproject.toml | 2 +-
 uv.lock        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7f5ceac6203..1e3ed2c76be 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ dynamic = ["version", "readme"]
 description = "Megatron Core - a library for efficient and scalable training of transformer based models"
 requires-python = ">=3.10"
 license = { text = "Apache 2.0" }
-dependencies = ["torch", "numpy", "packaging>=24.2"]
+dependencies = ["torch>=2.6.0", "numpy", "packaging>=24.2"]
 authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
 maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
 keywords = [
diff --git a/uv.lock b/uv.lock
index 705340af107..bfb4a0ee734 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2373,7 +2373,7 @@ requires-dist = [
     { name = "tensorstore", marker = "extra == 'dev'", specifier = "~=0.1,!=0.1.46,!=0.1.72" },
     { name = "tensorstore", marker = "extra == 'lts'", specifier = "~=0.1,!=0.1.46,!=0.1.72" },
     { name = "tiktoken", marker = "extra == 'mlm'" },
-    { name = "torch" },
+    { name = "torch", specifier = ">=2.6.0" },
     { name = "tqdm", marker = "extra == 'dev'" },
     { name = "tqdm", marker = "extra == 'lts'" },
     { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=d9b7fc5770a88af06e2e9c2bd97b550614c3a69f" },

From 190f5b66325162def96f4ed8c43ed0c705dc2133 Mon Sep 17 00:00:00 2001
From: kwyss-nvidia <kwyss@nvidia.com>
Date: Wed, 28 Jan 2026 21:07:05 -0800
Subject: [PATCH 63/79] Move kitchen extension file to private kitchen
 repository (#2779)

Co-authored-by: Jared Casper <155158+jaredcasper@users.noreply.github.com>
---
 megatron/core/extensions/kitchen.py           | 1836 +----------------
 megatron/core/models/gpt/gpt_layer_specs.py   |    5 +-
 megatron/core/transformer/mlp.py              |    4 +
 megatron/training/arguments.py                |    7 +-
 .../unit_tests/extension/test_kitchen_sdpa.py |   11 +-
 .../models/test_gpt_model_quantization.py     |    4 +-
 .../transformer/test_quantization_config.py   |   22 +-
 7 files changed, 51 insertions(+), 1838 deletions(-)

diff --git a/megatron/core/extensions/kitchen.py b/megatron/core/extensions/kitchen.py
index ad9be01fb60..a8a83fb341c 100644
--- a/megatron/core/extensions/kitchen.py
+++ b/megatron/core/extensions/kitchen.py
@@ -1,1810 +1,30 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
-import logging
-import math
-import warnings
-from dataclasses import dataclass, fields
-from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple
-
-import torch
-from torch import Tensor
-
-from megatron.core import tensor_parallel
-from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
-from megatron.core.fusions.fused_softmax import FusedScaleMaskSoftmax
-from megatron.core.model_parallel_config import ModelParallelConfig
-from megatron.core.models.backends import BackendSpecProvider
-from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.parallel_state import (
-    get_expert_data_parallel_rank,
-    get_expert_model_parallel_rank,
-    get_expert_model_parallel_world_size,
-)
-from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.quantization.quant_config import MatchContext, QuantizationConfig
-from megatron.core.tensor_parallel.random import (
-    get_cuda_rng_tracker,
-    get_data_parallel_rng_tracker_name,
-    get_expert_parallel_rng_tracker_name,
-)
-from megatron.core.tensor_parallel.utils import divide
-from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.transformer.mlp import MLPSubmodules
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP, TEGroupedMLP
-from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.transformer.utils import attention_mask_func, make_sharded_tensors_for_checkpoint
-from megatron.core.utils import get_tensor_model_parallel_group_if_none, log_single_rank
-
-logger = logging.getLogger(__name__)
-
-
-# Parsing constant
-_KITCHEN_CONFIG_TYPE_KEY = "kitchen_config_type"
-try:
-    import nvidia_kitchen  # type: ignore[import-not-found]
-
-    # Kitchen imports for SDPA
-    from nvidia_kitchen.attention import (  # type: ignore[import-not-found]
-        QAttentionParams,
-        QuantizedBMM,
-    )
-    from nvidia_kitchen.config import (  # type: ignore[import-not-found]
-        QLinearParams,
-        get_qlinear_params_from_qat_params,
-    )
-    from nvidia_kitchen.config_attention_recipe import (  # type: ignore[import-not-found]
-        get_qattention_params_from_qat_params,
-    )
-    from nvidia_kitchen.config_fa_recipe import (  # type: ignore[import-not-found]
-        get_qfa_params_from_recipe_name,
-    )
-
-    # Kitchen imports for FA
-    from nvidia_kitchen.fa import KitchenFlashAttentionModule  # type: ignore[import-not-found]
-    from nvidia_kitchen.fa_params import QFlashAttentionParams  # type: ignore[import-not-found]
-
-    HAVE_KITCHEN = True
-except ImportError:
-    from unittest.mock import MagicMock
-
-    HAVE_KITCHEN = False
-    nvidia_kitchen = MagicMock()
-    QuantizedBMM = MagicMock()
-    QAttentionParams = MagicMock()
-    QLinearParams = MagicMock()
-    get_qlinear_params_from_qat_params = MagicMock()
-    get_qattention_params_from_qat_params = MagicMock()
-    KitchenFlashAttentionModule = MagicMock()
-    QFlashAttentionParams = MagicMock()
-    get_qfa_params_from_recipe_name = MagicMock()
-
-
-class KitchenConfigType(Enum):
-    """Configuration object types in config dictionary"""
-
-    QLINEAR_PARAMS = "QLinearParams"
-    QATTENTION_PARAMS = "QAttentionParams"
-    QFLASHATTENTION_PARAMS = "QFlashAttentionParams"
-
-    COMPOUND_PARAMS = "CompoundParams"
-
-
-@dataclass
-class QFlashAttentionParamsConfigSchema:
-    """Dataclass to parse values from config dict of 'QFlashAttentionParams' type"""
-
-    kitchen_config_type: KitchenConfigType
-    recipe_name: str
-
-    @classmethod
-    def parse_config_dict(cls, config_dict: Dict[Any, Any]) -> 'QFlashAttentionParamsConfigSchema':
-        """
-        Parse config dictionary and return a schema instance.
-
-
-        Expected config format:
-            {"kitchen_config_type": "QFlashAttentionParams", "recipe_name": <str>}
-        """
-        expected_keys = cls.get_expected_keys()
-        actual_keys = set(config_dict.keys())
-
-        # Check for missing keys
-        missing = expected_keys - actual_keys
-        if missing:
-            raise KeyError(f"Missing required keys: {missing}")
-
-        # Check for unexpected keys
-        unexpected = actual_keys - expected_keys
-        if unexpected:
-            raise KeyError(f"Unexpected keys in config: {unexpected}")
-
-        try:
-            config_type = KitchenConfigType(config_dict[_KITCHEN_CONFIG_TYPE_KEY])
-        except ValueError:
-            raise ValueError(f"Unsupported config type '{config_dict['kitchen_config_type']}'.")
-
-        if config_type != KitchenConfigType.QFLASHATTENTION_PARAMS:
-            raise ValueError(f"Parsing config dict of incorrect type '{config_type}'")
-
-        # Create instance with converted enum
-        return cls(kitchen_config_type=config_type, recipe_name=config_dict["recipe_name"])
-
-    @classmethod
-    def get_expected_keys(cls) -> Set[str]:
-        """Get expected keys from the dataclass fields."""
-        return {field.name for field in fields(cls)}
-
-    def __post_init__(self):
-        # config type check
-        if not isinstance(self.kitchen_config_type, KitchenConfigType):
-            raise TypeError(
-                "kitchen_config_type must be KitchenConfigType, "
-                f"got {type(self.kitchen_config_type)}"
-            )
-
-        if self.kitchen_config_type != KitchenConfigType.QFLASHATTENTION_PARAMS:
-            raise TypeError(
-                f"kitchen_config_type must be QFlashAttentionParams got {self.kitchen_config_type}"
-            )
-        # recipe_name check
-        if not isinstance(self.recipe_name, str):
-            raise ValueError(f"recipe_name must be a string, got {self.recipe_name}")
-
-    def to_kitchen_qfa(self) -> QFlashAttentionParams:
-        """Converts to kitchen library's QFlashAttentionParams object."""
-        return get_qfa_params_from_recipe_name(self.recipe_name)
-
-
-@dataclass
-class QAttentionParamsConfigSchema:
-    """Dataclass to parse values from config dict of 'QAttentionParams' type"""
-
-    kitchen_config_type: KitchenConfigType
-    recipe_idx: int
-
-    @classmethod
-    def parse_config_dict(cls, config_dict: Dict[Any, Any]) -> 'QAttentionParamsConfigSchema':
-        """
-        Parse config dictionary and return a schema instance.
-
-
-        Expected config format: {"kitchen_config_type": "QLinearParams", "recipe_idx": <int>}
-        """
-        expected_keys = cls.get_expected_keys()
-        actual_keys = set(config_dict.keys())
-
-        # Check for missing keys
-        missing = expected_keys - actual_keys
-        if missing:
-            raise KeyError(f"Missing required keys: {missing}")
-
-        # Check for unexpected keys
-        unexpected = actual_keys - expected_keys
-        if unexpected:
-            raise KeyError(f"Unexpected keys in config: {unexpected}")
-
-        try:
-            config_type = KitchenConfigType(config_dict[_KITCHEN_CONFIG_TYPE_KEY])
-        except ValueError:
-            raise ValueError(f"Unsupported config type '{config_dict['kitchen_config_type']}'.")
-
-        if config_type != KitchenConfigType.QATTENTION_PARAMS:
-            raise ValueError(f"Parsing config dict of incorrect type '{config_type}'")
-
-        # Create instance with converted enum
-        return cls(kitchen_config_type=config_type, recipe_idx=config_dict["recipe_idx"])
-
-    @classmethod
-    def get_expected_keys(cls) -> Set[str]:
-        """Get expected keys from the dataclass fields."""
-        return {field.name for field in fields(cls)}
-
-    def __post_init__(self):
-        # config type check
-        if not isinstance(self.kitchen_config_type, KitchenConfigType):
-            raise TypeError(
-                "kitchen_config_type must be KitchenConfigType, "
-                f"got {type(self.kitchen_config_type)}"
-            )
-
-        if self.kitchen_config_type != KitchenConfigType.QATTENTION_PARAMS:
-            raise TypeError(
-                f"kitchen_config_type must be QAttentionParams got {self.kitchen_config_type}"
-            )
-        # recipe_idx check
-        if not isinstance(self.recipe_idx, int) or self.recipe_idx <= 0:
-            raise ValueError(f"recipe_idx must be a positive integer, got {self.recipe_idx}")
-
-    def to_kitchen_qattention(self) -> QAttentionParams:
-        """Converts to kitchen library's QAttentionParams object."""
-        return get_qattention_params_from_qat_params(self.recipe_idx)
-
-
-@dataclass
-class QLinearParamsConfigSchema:
-    """Dataclass to parse values from config dict of 'QLinearParams' type"""
-
-    kitchen_config_type: KitchenConfigType
-    recipe_idx: int
-
-    @classmethod
-    def parse_config_dict(cls, config_dict: Dict[Any, Any]) -> "QLinearParamsConfigSchema":
-        """
-        Parse config dictionary and return a schema instance.
-
-
-        Expected config format: {"kitchen_config_type": "QLinearParams", "recipe_idx": <int>}
-        """
-        expected_keys = cls.get_expected_keys()
-        actual_keys = set(config_dict.keys())
-
-        # Check for missing keys
-        missing = expected_keys - actual_keys
-        if missing:
-            raise KeyError(f"Missing required keys: {missing}")
-
-        # Check for unexpected keys
-        unexpected = actual_keys - expected_keys
-        if unexpected:
-            raise KeyError(f"Unexpected keys in config: {unexpected}")
-
-        try:
-            config_type = KitchenConfigType(config_dict[_KITCHEN_CONFIG_TYPE_KEY])
-        except ValueError:
-            raise ValueError(f"Unsupported config type '{config_dict['kitchen_config_type']}'.")
-
-        if config_type != KitchenConfigType.QLINEAR_PARAMS:
-            raise ValueError(f"Parsing config dict of incorrect type '{config_type}'")
-
-        # Create instance with converted enum
-        return cls(kitchen_config_type=config_type, recipe_idx=config_dict["recipe_idx"])
-
-    @classmethod
-    def get_expected_keys(cls) -> Set[str]:
-        """Get expected keys from the dataclass fields."""
-        return {field.name for field in fields(cls)}
-
-    def __post_init__(self):
-        # config type check
-        if not isinstance(self.kitchen_config_type, KitchenConfigType):
-            raise TypeError(
-                "kitchen_config_type must be KitchenConfigType, "
-                f"got {type(self.kitchen_config_type)}"
-            )
-
-        if self.kitchen_config_type != KitchenConfigType.QLINEAR_PARAMS:
-            raise TypeError(
-                f"kitchen_config_type must be QLinearParams got {self.kitchen_config_type}"
-            )
-        # recipe_idx check
-        if not isinstance(self.recipe_idx, int) or self.recipe_idx <= 0:
-            raise ValueError(f"recipe_idx must be a positive integer, got {self.recipe_idx}")
-
-    def to_kitchen_qlinear(self) -> QLinearParams:
-        """Converts to kitchen library's QLinearParams object."""
-        return get_qlinear_params_from_qat_params(self.recipe_idx)
-
-
-@dataclass
-class CompoundParamsConfigSchema:
-    """Dataclass to parse values from config dict of 'CompoundParams' type"""
-
-    kitchen_config_type: KitchenConfigType
-    configs: Dict[Any, Any]
-
-    q_linear_params: Optional[QLinearParamsConfigSchema] = None
-    q_attention_params: Optional[QAttentionParamsConfigSchema] = None
-    q_fa_params: Optional[QFlashAttentionParamsConfigSchema] = None
-
-    @classmethod
-    def parse_config_dict(cls, config_dict: Dict[Any, Any]) -> 'CompoundParamsConfigSchema':
-        """
-        Parse config dictionary and return a schema instance.
-
-        Expected config format: {
-            "kitchen_config_type": "CompoundParams",
-            "configs": [
-                {"kitchen_config_type": "QLinearParams", "recipe_idx": <int>},
-                {"kitchen_config_type": "QAttentionParams", "recipe_idx": <int>},
-            ]
-        }
-
-        or {
-            "kitchen_config_type": "CompoundParams",
-            "configs": [
-                {"kitchen_config_type": "QLinearParams", "recipe_idx": <int>},
-                {"kitchen_config_type": "QFlashAttentionParams", "recipe_name": <str>},
-            ]
-        }
-        """
-        expected_keys = cls.get_expected_keys()
-        actual_keys = set(config_dict.keys())
-
-        # Check for missing keys
-        missing = expected_keys - actual_keys
-        if missing:
-            raise KeyError(f"Missing required keys: {missing}")
-
-        # Check for unexpected keys
-        unexpected = actual_keys - expected_keys
-        if unexpected:
-            raise KeyError(f"Unexpected keys in config: {unexpected}")
-
-        try:
-            config_type = KitchenConfigType(config_dict[_KITCHEN_CONFIG_TYPE_KEY])
-        except ValueError:
-            raise ValueError(f"Unsupported config type '{config_dict['kitchen_config_type']}'.")
-
-        if config_type != KitchenConfigType.COMPOUND_PARAMS:
-            raise ValueError(f"Parsing config dict of incorrect type '{config_type}'")
-
-        # Create instance with converted enum
-        return cls(kitchen_config_type=config_type, configs=config_dict["configs"])
-
-    @classmethod
-    def get_expected_keys(cls) -> Set[str]:
-        """Get expected keys from the dataclass fields."""
-        return {
-            field.name
-            for field in fields(cls)
-            if field.name not in ["q_linear_params", "q_attention_params", "q_fa_params"]
-        }
-
-    def __post_init__(self):
-        if not isinstance(self.kitchen_config_type, KitchenConfigType):
-            raise TypeError(
-                "kitchen_config_type must be KitchenConfigType, "
-                f"got {type(self.kitchen_config_type)}"
-            )
-
-        if self.kitchen_config_type != KitchenConfigType.COMPOUND_PARAMS:
-            raise TypeError(
-                f"kitchen_config_type must be CompoundParams got {self.kitchen_config_type}"
-            )
-
-        for config in self.configs:
-            if config["kitchen_config_type"] == KitchenConfigType.QLINEAR_PARAMS.value:
-                self.q_linear_params = QLinearParamsConfigSchema.parse_config_dict(config)
-            elif config["kitchen_config_type"] == KitchenConfigType.QATTENTION_PARAMS.value:
-                self.q_attention_params = QAttentionParamsConfigSchema.parse_config_dict(config)
-            elif config["kitchen_config_type"] == KitchenConfigType.QFLASHATTENTION_PARAMS.value:
-                self.q_fa_params = QFlashAttentionParamsConfigSchema.parse_config_dict(config)
-            else:
-                raise ValueError(f"Unsupported config type '{config['kitchen_config_type']}'.")
-
-    def get_qlinear_params(self) -> Optional[QLinearParams]:
-        """
-        Returns the QLinearParams object for the compound params.
-        """
-        return self.q_linear_params.to_kitchen_qlinear() if self.q_linear_params else None
-
-    def get_qattention_params(self) -> Optional[QAttentionParams]:
-        """
-        Returns the QAttentionParams object for the compound params.
-        """
-        return self.q_attention_params.to_kitchen_qattention() if self.q_attention_params else None
-
-    def get_qfa_params(self) -> Optional[QFlashAttentionParams]:
-        """
-        Returns the QFlashAttentionParams object for the compound params.
-        """
-        return self.q_fa_params.to_kitchen_qfa() if self.q_fa_params else None
-
-
-@dataclass
-class KitchenQuantizationParams:
-    """Quantization parameters used for kitchen extensions"""
-
-    # Could be extended with sparsity, etc.
-    # match_input is what selected the config.
-    qlinear_params: Optional[QLinearParams]
-
-    match_input: MatchContext
-    params_config_key: str
-
-    qattention_params: Optional[QAttentionParams] = None
-    qfa_params: Optional[QFlashAttentionParams] = None
-
-    @staticmethod
-    def parse_from_config(quant_config: QuantizationConfig) -> "KitchenQuantizationParams":
-        """Parses quantization config for a layer or throw an error."""
-        if not HAVE_KITCHEN:
-            raise ImportError("Kitchen not available. Kitchen is not released publicly.")
-
-        assert (
-            quant_config is not None
-        ), "Kitchen extension expects a quantization config for linear layers."
-        config = quant_config.config
-        try:
-            config_type = KitchenConfigType(config[_KITCHEN_CONFIG_TYPE_KEY])
-        except KeyError:
-            raise ValueError(
-                f"Kitchen config dictionary must have '{_KITCHEN_CONFIG_TYPE_KEY}' key."
-            )
-        except ValueError:
-            raise ValueError(f"Unsupported config type '{config['kitchen_config_type']}'.")
-
-        if config_type == KitchenConfigType.QLINEAR_PARAMS:
-            return KitchenQuantizationParams(
-                qlinear_params=QLinearParamsConfigSchema.parse_config_dict(
-                    config
-                ).to_kitchen_qlinear(),
-                qattention_params=None,
-                qfa_params=None,
-                match_input=quant_config.match_input,
-                params_config_key=quant_config.config_key,
-            )
-        elif config_type == KitchenConfigType.QATTENTION_PARAMS:
-            return KitchenQuantizationParams(
-                qlinear_params=None,
-                qattention_params=QAttentionParamsConfigSchema.parse_config_dict(
-                    config
-                ).to_kitchen_qattention(),
-                qfa_params=None,
-                match_input=quant_config.match_input,
-                params_config_key=quant_config.config_key,
-            )
-        elif config_type == KitchenConfigType.QFLASHATTENTION_PARAMS:
-            return KitchenQuantizationParams(
-                qlinear_params=None,
-                qattention_params=None,
-                qfa_params=QFlashAttentionParamsConfigSchema.parse_config_dict(
-                    config
-                ).to_kitchen_qfa(),
-                match_input=quant_config.match_input,
-                params_config_key=quant_config.config_key,
-            )
-        elif config_type == KitchenConfigType.COMPOUND_PARAMS:
-            compound_params = CompoundParamsConfigSchema.parse_config_dict(config)
-            return KitchenQuantizationParams(
-                qlinear_params=compound_params.get_qlinear_params(),
-                qattention_params=compound_params.get_qattention_params(),
-                qfa_params=compound_params.get_qfa_params(),
-                match_input=quant_config.match_input,
-                params_config_key=quant_config.config_key,
-            )
-        else:
-            raise NotImplementedError(f"Unhandled configuration type {config_type}")
-
-
-def _get_extra_kitchen_kwargs(config: TransformerConfig):
-    extra_kitchen_kwargs = {"params_dtype": config.params_dtype}
-
-    if config.use_cpu_initialization:
-        raise ValueError("Kitchen backend does not support use_cpu_initialization.")
-    elif config.init_model_with_meta_device:
-        extra_kitchen_kwargs["device"] = "meta"
-    else:
-        extra_kitchen_kwargs["device"] = torch.cuda.current_device()
-    return extra_kitchen_kwargs
-
-
-class KitchenLinear(nvidia_kitchen.Linear):
-    """
-    Wrapper for Kitchen's `Linear` layer.
-
-    Note that if Megatron's parallel_state has not been initialized
-    yet, the tp_group passed to Kitchen will be None and must be set later
-    via set_tensor_parallel_group().
-
-    parallel_mode currently supports 3 different values:
-        - "column": Split the weight matrix along output dimension (for KitchenColumnParallelLinear)
-        - "row": Split the weight matrix along input dimension (for KitchenRowParallelLinear)
-        - "duplicated": No tensor parallelism and weight is duplicated across TP ranks
-        - Note: For expert linear layers, we will disable communication logic here
-                as TP communication is handled in token_dispatcher.
-    """
-
-    def __init__(
-        self,
-        input_size: int,
-        output_size: int,
-        *,
-        parallel_mode: Optional[str],
-        config: ModelParallelConfig,
-        init_method: Callable,
-        bias: bool,
-        skip_bias_add: bool,
-        skip_weight_param_allocation: bool,
-        tp_comm_buffer_name: Optional[str] = None,
-        layer_number: Optional[int] = None,
-        is_expert: bool = False,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-        self.config = config
-
-        # Kitchen returns a zero length Tensor when bias=False and
-        # return_bias=True, but we prefer None.  So in that case we
-        # tell TE to not return the bias, and return None
-        # ourselves. This way our forward always returns two values
-        # and we don't have to deal with the zero length Tensor.
-        self.kitchen_return_bias = skip_bias_add and bias
-        self.is_first_microbatch = True
-        self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
-        if skip_weight_param_allocation:
-            raise ValueError("Kitchen linear layers do not support skip_weight_param_allocation")
-
-        # Save params for finish_init
-        self.stashed_input_size = input_size
-        self.stashed_output_size = output_size
-        self.stashed_parallel_mode = parallel_mode
-        self.stashed_init_method = init_method
-        self.stashed_bias = bias
-        self.stashed_tp_comm_buffer_name = tp_comm_buffer_name
-        self.stashed_layer_number = layer_number
-        self.stashed_is_expert = is_expert
-        self.stashed_tp_group = tp_group
-
-        self.init_finished = False
-
-    def finish_init(self, quantization_config: QuantizationConfig):
-        """Required post-init of quantization configuration."""
-        extra_kwargs = _get_extra_kitchen_kwargs(self.config)
-
-        # Restore args from stash
-        input_size = self.stashed_input_size
-        output_size = self.stashed_output_size
-        parallel_mode = self.stashed_parallel_mode
-        init_method = self.stashed_init_method
-        bias = self.stashed_bias
-        tp_comm_buffer_name = self.stashed_tp_comm_buffer_name
-        layer_number = self.stashed_layer_number
-        is_expert = self.stashed_is_expert
-        tp_group = self.stashed_tp_group
-
-        self.kitchen_quant_params = KitchenQuantizationParams.parse_from_config(quantization_config)
-        assert self.kitchen_quant_params.qlinear_params is not None
-        extra_kwargs["qlinear_params"] = self.kitchen_quant_params.qlinear_params
-
-        if tp_comm_buffer_name:
-            self.config.tp_comm_overlap = False
-            warnings.warn(
-                f"The user buffer name {tp_comm_buffer_name} is not supported in "
-                "Kitchen. Disabling TP communication overlap for this layer."
-            )
-            extra_kwargs["ub_name"] = tp_comm_buffer_name
-
-        extra_kwargs["layer_number"] = layer_number
-
-        if parallel_mode == "duplicated":
-            assert tp_group is None, "duplicated linear should not have tp_group set"
-            tp_size = 1
-        else:
-            assert tp_group is not None, "Parallel linear should always have tp_group set"
-            tp_size = tp_group.size()
-
-        self.expert_parallel = self.config.expert_model_parallel_size > 1
-        if is_expert:
-            rng_tracker_name = get_expert_parallel_rng_tracker_name()
-        else:
-            if parallel_mode == "duplicated":
-                rng_tracker_name = get_data_parallel_rng_tracker_name()
-            else:
-                rng_tracker_name = None
-        extra_kwargs["rng_tracker_name"] = rng_tracker_name
-
-        kitchen_parallel_mode = parallel_mode
-        if parallel_mode == "duplicated":
-            # Handle non-parallel case
-            tp_group = None
-            tp_size = 1
-            explicit_expert_comm = False
-            kitchen_parallel_mode = None
-        else:
-            # Disable communications in kitchen when using TP or EP by megatron
-            explicit_expert_comm = is_expert and (tp_size > 1 or self.expert_parallel)
-
-            if explicit_expert_comm:
-                if parallel_mode == "column":
-                    output_size = divide(output_size, tp_size)
-                elif parallel_mode == "row":
-                    input_size = divide(input_size, tp_size)
-                kitchen_parallel_mode = None
-                tp_size = 1
-                tp_group = None
-
-        super().__init__(
-            in_features=input_size,
-            out_features=output_size,
-            sequence_parallel=self.config.sequence_parallel,
-            fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
-            # Pass None if not initialized for backward compatibility with the ckpt converter.
-            tp_group=tp_group if torch.distributed.is_initialized() else None,
-            tp_size=tp_size,
-            get_rng_state_tracker=(
-                get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
-            ),
-            init_method=(init_method if self.config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            return_bias=self.kitchen_return_bias,
-            parallel_mode=kitchen_parallel_mode,
-            **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:
-                # Reduce the gradient on DP group
-                setattr(param, "allreduce", True)
-                if parallel_mode == "duplicated":
-                    # Reduce the gradient further on the TP group since the weight is
-                    # duplicated across TP ranks
-                    setattr(param, "sequence_parallel", self.config.sequence_parallel)
-
-        del self.stashed_input_size
-        del self.stashed_output_size
-        del self.stashed_parallel_mode
-        del self.stashed_init_method
-        del self.stashed_bias
-        del self.stashed_tp_comm_buffer_name
-        del self.stashed_layer_number
-        del self.stashed_is_expert
-        del self.stashed_tp_group
-        self.init_finished = True
-
-    def forward(self, x):
-        """Forward."""
-        assert self.init_finished
-        _is_first_microbatch = (
-            None if self.disable_parameter_transpose_cache else self.is_first_microbatch
-        )
-        out = super().forward(x, is_first_microbatch=_is_first_microbatch)
-        self.is_first_microbatch = False
-
-        # Kitchen only returns a tuple when return_bias is True, otherwise
-        # it returns a single Tensor, we always want to return two
-        # values regardless of the arguments.
-        if self.kitchen_return_bias:
-            return out
-        return out, None
-
-    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
-        """Replicate cross TP/DP."""
-
-        # Provide the dist-ckpt support when KitchenLinear is directly used
-        # It can only happen with duplicated parallel mode
-        assert (
-            self.parallel_mode is None
-        ), "KitchenLinear sharded_state_dict can only be used with duplicated parallel mode"
-        state_dict = self.state_dict(prefix="", keep_vars=True)
-        return make_sharded_tensors_for_checkpoint(state_dict, prefix, None, sharded_offsets)
-
-
-class KitchenColumnParallelLinear(KitchenLinear):
-    """
-    Wrapper for the Kitchen's `Linear` layer but specialized similar
-    to megatron's `ColumnParallelLinear` layer.
-    """
-
-    def __init__(
-        self,
-        input_size: int,
-        output_size: int,
-        *,
-        config: ModelParallelConfig,
-        init_method: Callable,
-        gather_output: bool,
-        bias: bool,
-        skip_bias_add: bool,
-        is_expert: bool,
-        skip_weight_param_allocation: bool = False,
-        tp_comm_buffer_name: Optional[str] = None,
-        layer_number: Optional[int] = None,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-
-        if gather_output:
-            raise ValueError("Kitchen linear layers do not support gather_output = True")
-        tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
-        world_size = tp_group.size()
-        rank = tp_group.rank()
-
-        super().__init__(
-            input_size=input_size,
-            output_size=output_size,
-            parallel_mode="column",
-            config=config,
-            init_method=(init_method if config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            skip_bias_add=skip_bias_add,
-            is_expert=is_expert,
-            skip_weight_param_allocation=skip_weight_param_allocation,
-            tp_comm_buffer_name=tp_comm_buffer_name,
-            layer_number=layer_number,
-            tp_group=tp_group,
-        )
-
-        if config.use_cpu_initialization:
-            raise ValueError("Kitchen extension doesn't support use_cpu_initialization.")
-
-    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)
-        return make_sharded_tensors_for_checkpoint(
-            state_dict, prefix, {"weight": 0, "bias": 0}, sharded_offsets
-        )
-
-    def __repr__(self):
-        return (
-            f"{type(self).__name__}(in_features={self.in_features}, "
-            f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
-        )
-
-
-class KitchenRowParallelLinear(KitchenLinear):
-    """
-    Wrapper for Kitchen's `Linear` layer but specialized similar
-    to megatron's `RowParallelLinear` layer.
-    """
-
-    def __init__(
-        self,
-        input_size: int,
-        output_size: int,
-        *,
-        config: ModelParallelConfig,
-        init_method: Callable,
-        bias: bool,
-        input_is_parallel: bool,
-        skip_bias_add: bool,
-        is_expert: bool,
-        tp_comm_buffer_name: Optional[str] = None,
-        layer_number: Optional[int] = None,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-
-        if not input_is_parallel:
-            raise ValueError("Kitchen linear layers do not support input_is_parallel = False")
-        tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
-
-        super().__init__(
-            input_size=input_size,
-            output_size=output_size,
-            parallel_mode="row",
-            config=config,
-            init_method=(init_method if config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            skip_bias_add=skip_bias_add,
-            skip_weight_param_allocation=False,
-            # We don't currently use this for row parallel layers # pylint: disable=line-too-long
-            is_expert=is_expert,
-            tp_comm_buffer_name=tp_comm_buffer_name,
-            layer_number=layer_number,
-            tp_group=tp_group,
-        )
-        if config.use_cpu_initialization:
-            raise ValueError("Kitchen extension does not support use_cpu_initialization.")
-
-    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)
-        return make_sharded_tensors_for_checkpoint(
-            state_dict, prefix, {"weight": 1}, sharded_offsets
-        )
-
-    def __repr__(self):
-        return (
-            f"{type(self).__name__}(in_features={self.in_features}, "
-            f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
-        )
-
-
-class KitchenGroupedLinear(nvidia_kitchen.GroupedLinear):
-    """
-    Wrapper for Kitchen's `GroupedLinear` layer.
-
-    Note that if Megatron's parallel_state has not been initialized
-    yet, the tp_group passed to TE will be None and must be set later
-    via set_tensor_parallel_group().
-    """
-
-    def __init__(
-        self,
-        num_gemms: int,
-        input_size: int,
-        output_size: int,
-        *,
-        parallel_mode: Optional[str],
-        config: ModelParallelConfig,
-        init_method: Callable,
-        bias: bool,
-        skip_bias_add: bool,
-        is_expert: bool = False,
-        tp_comm_buffer_name: Optional[str] = None,
-        layer_number: Optional[int] = None,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-
-        self.config = config
-
-        # Kitchen returns a zero length Tensor when bias=False and
-        # return_bias=True, but we prefer None.  So in that case we
-        # tell TE to not return the bias, and return None
-        # ourselves. This way our forward always returns two values
-        # and we don't have to deal with the zero length Tensor.
-        self.kitchen_return_bias = skip_bias_add and bias
-        self.is_first_microbatch = True
-        self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
-
-        # Stash parameters for finish_init
-        self.stashed_num_gemms = num_gemms
-        self.stashed_input_size = input_size
-        self.stashed_output_size = output_size
-        self.stashed_parallel_mode = parallel_mode
-        self.stashed_init_method = init_method
-        self.stashed_bias = bias
-        self.stashed_is_expert = is_expert
-        self.stashed_tp_comm_buffer_name = tp_comm_buffer_name
-        self.stashed_layer_number = layer_number
-        self.stashed_tp_group = tp_group
-        self.init_finished = False
-
-    def finish_init(self, quantization_config: QuantizationConfig) -> None:
-        """Required post-init of quantization configuration."""
-        # Restore parameters from stash
-        num_gemms = self.stashed_num_gemms
-        input_size = self.stashed_input_size
-        output_size = self.stashed_output_size
-        parallel_mode = self.stashed_parallel_mode
-        init_method = self.stashed_init_method
-        bias = self.stashed_bias
-        is_expert = self.stashed_is_expert
-        tp_comm_buffer_name = self.stashed_tp_comm_buffer_name
-        layer_number = self.stashed_layer_number
-        tp_group = self.stashed_tp_group
-
-        extra_kwargs = _get_extra_kitchen_kwargs(self.config)
-        extra_kwargs["ub_name"] = tp_comm_buffer_name
-        extra_kwargs["layer_number"] = layer_number
-
-        self.kitchen_quant_params = KitchenQuantizationParams.parse_from_config(quantization_config)
-        assert self.kitchen_quant_params.qlinear_params is not None
-        extra_kwargs["qlinear_params"] = self.kitchen_quant_params.qlinear_params
-
-        self.expert_parallel = self.config.expert_model_parallel_size > 1
-        if is_expert:
-            extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name()
-
-        # The comms between TP and EP group is explicitly handled by MoE token dispatcher.
-        # So we disable comms by making Kitchen agnostic of model parallel.
-        tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
-        tp_size = tp_group.size()
-
-        self.explicit_expert_comm = is_expert and (tp_size > 1 or self.expert_parallel)
-
-        if self.explicit_expert_comm:
-            if parallel_mode == "column":
-                output_size = divide(output_size, tp_size)
-            elif parallel_mode == "row":
-                input_size = divide(input_size, tp_size)
-            parallel_mode = None
-            tp_size = 1
-            tp_group = None
-
-        super().__init__(
-            num_gemms=num_gemms,
-            in_features=input_size,
-            out_features=output_size,
-            sequence_parallel=self.config.sequence_parallel,
-            fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
-            tp_group=tp_group if torch.distributed.is_initialized() else None,
-            tp_size=tp_size,
-            get_rng_state_tracker=(
-                get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
-            ),
-            init_method=(init_method if self.config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            return_bias=self.kitchen_return_bias,
-            parallel_mode=parallel_mode,
-            **extra_kwargs,
-        )
-
-        for param in self.parameters():
-            setattr(param, "allreduce", not (is_expert and self.expert_parallel))
-
-        def merge_extra_states(
-            self,
-            state_dict,
-            prefix,
-            local_metadata,
-            strict,
-            missing_keys,
-            unexpected_keys,
-            error_msgs,
-        ):
-            """
-            Merge multiple "_extra_state" into one.
-            """
-            self.init_fp8_metadata(num_gemms=self.num_gemms)
-            fp8_checkpoint = self.fp8_meta["fp8_checkpoint"] or self.fp8 or self.fp8_calibration
-
-            try:
-                state_list = [
-                    state_dict.pop(f"{prefix}_extra_state{i}") for i in range(1, self.num_gemms)
-                ]
-            except KeyError:
-                # "_extra_state{i}" only exists for dist-ckpt. Return for torch native ckpt.
-                return
-
-            if not fp8_checkpoint:
-                return
-            state_list = [state_dict.pop(f"{prefix}_extra_state")] + state_list
-            state_list = [self._decode_extra_state(state) for state in state_list]
-            extra_fp8_variables = state_list[0]["extra_fp8_variables"]
-            extra_fp8_variables["num_gemms"] = self.num_gemms
-            extra_state = {"extra_fp8_variables": extra_fp8_variables}
-            state_dict[f"{prefix}_extra_state"] = self._encode_extra_state(extra_state)
-
-        self._register_load_state_dict_pre_hook(merge_extra_states, with_module=True)
-        del self.stashed_num_gemms
-        del self.stashed_input_size
-        del self.stashed_output_size
-        del self.stashed_parallel_mode
-        del self.stashed_init_method
-        del self.stashed_bias
-        del self.stashed_is_expert
-        del self.stashed_tp_comm_buffer_name
-        del self.stashed_layer_number
-        del self.stashed_tp_group
-        self.init_finished = True
-
-    def forward(self, x, m_splits):
-        """Forward."""
-        assert self.init_finished
-        _is_first_microbatch = (
-            None if self.disable_parameter_transpose_cache else self.is_first_microbatch
-        )
-        out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch)
-        self.is_first_microbatch = False
-
-        # Kitchen only returns a tuple when return_bias is True, otherwise
-        # it returns a single Tensor, we always want to return two
-        # values regardless of the arguments.
-        if self.kitchen_return_bias:
-            return out
-        return out, None
-
-    def _encode_extra_state(self, state):
-        torch.cuda.synchronize()
-        state_serialized = bytearray(pickle.dumps(state))
-        state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
-        return state_serialized
-
-    def _decode_extra_state(self, state):
-        if isinstance(state, torch.Tensor):
-            return pickle.loads(state.detach().cpu().numpy().tobytes())
-        elif isinstance(state, io.BytesIO):
-            state.seek(0)
-            return torch.load(state, map_location="cuda")
-        else:
-            raise RuntimeError("Unsupported checkpoint format.")
-
-    def _split_extra_state(self, state):
-        fp8_checkpoint = self.fp8_meta["fp8_checkpoint"]
-        # Kitchen is compatible with TE checkpoint format, but never
-        # uses fp8_checkpoints.
-        assert not fp8_checkpoint
-        return [state] * self.num_gemms
-
-    def _sharded_state_dict_grouped(
-        self, tp_axis_map, prefix="", sharded_offsets=(), metadata=None
-    ):
-        """
-        prefix should be module_name to make keys identical to sequetial ones.
-        """
-        assert self.init_finished
-        sharded_state_dict = {}
-        full_state_dict = self.state_dict(prefix="", keep_vars=True)
-        num_global_experts = get_expert_model_parallel_world_size() * self.num_gemms
-        local_expert_indices_offset = get_expert_model_parallel_rank() * self.num_gemms
-        ep_axis = len(sharded_offsets)
-        extra_states = self._split_extra_state(full_state_dict["_extra_state"])
-        for gemm_idx in range(self.num_gemms):
-            state_dict = {
-                f"{gemm_idx}.weight": full_state_dict[f"weight{gemm_idx}"],
-                f"{gemm_idx}._extra_state": extra_states[gemm_idx],
-            }
-            if self.use_bias:
-                state_dict[f"{gemm_idx}.bias"] = full_state_dict[f"bias{gemm_idx}"]
-            sub_sd = make_sharded_tensors_for_checkpoint(
-                state_dict,
-                "",
-                tp_axis_map,
-                (
-                    *sharded_offsets,
-                    (ep_axis, local_expert_indices_offset + gemm_idx, num_global_experts),
-                ),
-            )
-            # Remove expert layers indexing from sharded keys
-            replace_prefix_for_sharding(sub_sd, f"{gemm_idx}.", prefix)
-            sharded_state_dict.update(
-                {
-                    f"{prefix}weight{gemm_idx}": sub_sd[f"{gemm_idx}.weight"],
-                    f"{prefix}_extra_state{'' if gemm_idx == 0 else gemm_idx}": sub_sd[
-                        f"{gemm_idx}._extra_state"
-                    ],
-                }
-            )
-            if self.use_bias:
-                sharded_state_dict[f"{prefix}bias{gemm_idx}"] = sub_sd[f"{gemm_idx}.bias"]
-        # Adjust replica ids - replication along DP modulo EP
-        for k, sh_ten in sharded_state_dict.items():
-            replica_id = sh_ten.replica_id
-            assert (
-                len(replica_id) == 3
-            ), f"Expected replica_id for {k} to be in (PP, TP, DP) format, got: {replica_id}"
-            if getattr(sh_ten, "is_data_parallel_fully_shard", False):
-                edp_replica_id = 0
-            else:
-                edp_replica_id = get_expert_data_parallel_rank()
-            sh_ten.replica_id = (*replica_id[:2], edp_replica_id)
-        return sharded_state_dict
-
-
-class KitchenColumnParallelGroupedLinear(KitchenGroupedLinear):
-    """
-    Wrapper for Kitchen's `GroupedLinear` layer but specialized
-    to column-parallel style.
-    """
-
-    def __init__(
-        self,
-        num_gemms: int,
-        input_size: int,
-        output_size: int,
-        *,
-        config: ModelParallelConfig,
-        init_method: Callable,
-        bias: bool,
-        skip_bias_add: bool,
-        is_expert: bool,
-        tp_comm_buffer_name: Optional[str] = None,
-        layer_number: Optional[int] = None,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-
-        super().__init__(
-            num_gemms=num_gemms,
-            input_size=input_size,
-            output_size=output_size,
-            parallel_mode="column",
-            config=config,
-            init_method=(init_method if config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            skip_bias_add=skip_bias_add,
-            is_expert=is_expert,
-            tp_comm_buffer_name=tp_comm_buffer_name,
-            layer_number=layer_number,
-            tp_group=tp_group,
-        )
-
-    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
-        """
-        For each gemm, sharding along axis 0, bias sharded.
-        Assume sharded_offsets[-1] is the expert parallel offset.
-        """
-        tp_axis_map = {}
-        for gemm_idx in range(self.num_gemms):
-            tp_axis_map.update({f"{gemm_idx}.weight": 0, f"{gemm_idx}.bias": 0})
-        return super()._sharded_state_dict_grouped(tp_axis_map, prefix, sharded_offsets, metadata)
-
-
-class KitchenRowParallelGroupedLinear(KitchenGroupedLinear):
-    """
-    Wrapper for Kitchen's `GroupedLinear` layer but specialized
-    to row-parallel style.
-    """
-
-    def __init__(
-        self,
-        num_gemms: int,
-        input_size: int,
-        output_size: int,
-        *,
-        config: ModelParallelConfig,
-        init_method: Callable,
-        bias: bool,
-        skip_bias_add: bool,
-        is_expert: bool,
-        tp_comm_buffer_name: Optional[str] = None,
-        layer_number: Optional[int] = None,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-
-        super().__init__(
-            num_gemms=num_gemms,
-            input_size=input_size,
-            output_size=output_size,
-            parallel_mode="row",
-            config=config,
-            init_method=(init_method if config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            skip_bias_add=skip_bias_add,
-            is_expert=is_expert,
-            tp_comm_buffer_name=tp_comm_buffer_name,
-            layer_number=layer_number,
-            tp_group=tp_group,
-        )
-
-    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
-        """
-        For each gemm, sharding along axis 1, bias not sharded.
-        Assume sharded_offsets[-1] is the expert parallel offset.
-        """
-        tp_axis_map = {f"{gemm_idx}.weight": 1 for gemm_idx in range(self.num_gemms)}
-        return super()._sharded_state_dict_grouped(tp_axis_map, prefix, sharded_offsets, metadata)
-
-
-class KitchenLayerNormColumnParallelLinear(nvidia_kitchen.LayerNormLinear):
-    """
-    Wrapper for Kitchen's `LayerNormLinear` layer that combines
-    layernorm and linear layers
-    """
-
-    def __init__(
-        self,
-        input_size: int,
-        output_size: int,
-        *,
-        config: TransformerConfig,
-        init_method: Callable,
-        gather_output: bool,
-        bias: bool,
-        skip_bias_add: bool,
-        is_expert: bool,
-        skip_weight_param_allocation: bool = False,
-        layer_number: Optional[int] = None,
-        tp_comm_buffer_name: Optional[str] = None,
-        tp_group: Optional[torch.distributed.ProcessGroup] = None,
-    ):
-        if not HAVE_KITCHEN:
-            raise ImportError(
-                "Kitchen extension requires the nvidia_kitchen package. "
-                "Please install it with `pip install nvidia-kitchen`."
-            )
-
-        self.config = config
-
-        if gather_output:
-            raise ValueError("Kitchen linear layers do not support gather_output = True")
-
-        if is_expert:
-            raise ValueError("Kitchen linear layers do not yet support MoE")
-
-        if skip_weight_param_allocation:
-            raise ValueError("Kitchen linear layers do not support skip_weight_param_allocation")
-
-        tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
-        # Kitchen returns a zero length Tensor when bias=False and
-        # return_bias=True, but we prefer None.  So in that case we
-        # tell Kitchen to not return the bias, and return None
-        # ourselves. This way our forward always returns two values
-        # and we don't have to deal with the zero length Tensor.
-        self.kitchen_return_bias = skip_bias_add and bias
-        self.is_first_microbatch = True
-        self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
-        self.tp_size = tp_group.size()
-        self.tp_rank = tp_group.rank()
-
-        if self.config.tp_comm_overlap:
-            raise ValueError("Kitchen LayerNormLinear does not support tp_comm_overlap")
-
-        if self.config.symmetric_ar_type is not None:
-            raise ValueError("Kitchen LayerNormLinear does not support symmetric all-reduce")
-
-        if config.use_cpu_initialization:
-            raise ValueError("Kitchen extension does not support use_cpu_initialization")
-
-        # Stash parameters for finish_init.
-        self.stashed_input_size = input_size
-        self.stashed_output_size = output_size
-        self.stashed_init_method = init_method
-        self.stashed_gather_output = gather_output
-        self.stashed_bias = bias
-        self.stashed_skip_bias_add = skip_bias_add
-        self.stashed_is_expert = is_expert
-        self.stashed_skip_weight_param_allocation = skip_weight_param_allocation
-        self.stashed_layer_number = layer_number
-        self.stashed_tp_comm_buffer_name = tp_comm_buffer_name
-        self.stashed_tp_group = tp_group
-        self.init_finished = False
-
-    def finish_init(self, quantization_config: QuantizationConfig) -> None:
-        """Required post-init of quantization configuration."""
-        # Restore parameters from stash
-        input_size = self.stashed_input_size
-        output_size = self.stashed_output_size
-        init_method = self.stashed_init_method
-        gather_output = self.stashed_gather_output
-        bias = self.stashed_bias
-        skip_bias_add = self.stashed_skip_bias_add
-        is_expert = self.stashed_is_expert
-        skip_weight_param_allocation = self.stashed_skip_weight_param_allocation
-        layer_number = self.stashed_layer_number
-        tp_comm_buffer_name = self.stashed_tp_comm_buffer_name
-        tp_group = self.stashed_tp_group
-
-        extra_kwargs = _get_extra_kitchen_kwargs(self.config)
-        extra_kwargs["normalization"] = self.config.normalization
-        self.kitchen_quant_params = KitchenQuantizationParams.parse_from_config(quantization_config)
-        assert self.kitchen_quant_params.qlinear_params is not None
-        extra_kwargs["qlinear_params"] = self.kitchen_quant_params.qlinear_params
-        extra_kwargs["ub_name"] = tp_comm_buffer_name
-
-        super().__init__(
-            in_features=input_size,
-            out_features=output_size,
-            eps=self.config.layernorm_epsilon,
-            sequence_parallel=self.config.sequence_parallel,
-            fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
-            tp_group=tp_group if torch.distributed.is_initialized() else None,
-            tp_size=self.config.tensor_model_parallel_size,
-            get_rng_state_tracker=(
-                get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
-            ),
-            init_method=(init_method if self.config.perform_initialization else (lambda w: None)),
-            bias=bias,
-            return_bias=self.kitchen_return_bias,
-            parallel_mode="column",
-            return_layernorm_output=False,
-            zero_centered_gamma=self.config.layernorm_zero_centered_gamma,
-            layer_number=layer_number,
-            **extra_kwargs,
-        )
-        del self.stashed_input_size
-        del self.stashed_output_size
-        del self.stashed_init_method
-        del self.stashed_gather_output
-        del self.stashed_bias
-        del self.stashed_skip_bias_add
-        del self.stashed_is_expert
-        del self.stashed_skip_weight_param_allocation
-        del self.stashed_layer_number
-        del self.stashed_tp_comm_buffer_name
-        del self.stashed_tp_group
-        self.init_finished = True
-
-    def forward(self, x):
-        """Forward."""
-        assert self.init_finished
-        _is_first_microbatch = (
-            None if self.disable_parameter_transpose_cache else self.is_first_microbatch
-        )
-        out = super().forward(x, is_first_microbatch=_is_first_microbatch)
-        self.is_first_microbatch = False
-
-        # Kitchen only returns a tuple when return_bias is True, otherwise
-        # it returns a single Tensor, we always want to return two
-        # values regardless of the arguments.
-        if self.kitchen_return_bias:
-            return out
-        return out, None
-
-    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
-        """Sharding along axis 0, bias sharded"""
-        assert self.init_finished
-        state_dict = self.state_dict(prefix="", keep_vars=True)
-        return make_sharded_tensors_for_checkpoint(
-            state_dict, prefix, {"weight": 0, "bias": 0}, sharded_offsets
-        )
-
-    def __repr__(self):
-        return (
-            f"{type(self).__name__}(in_features={self.in_features}, "
-            f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
-        )
-
-
-class KitchenFlashAttention(MegatronModule):
-    """
-    Flash Attention implementation for Kitchen.
-    """
-
-    def __init__(
-        self,
-        config: TransformerConfig,
-        layer_number: int,
-        attn_mask_type: AttnMaskType,
-        attention_type: str,
-        attention_dropout: Optional[float] = None,
-        softmax_scale: Optional[float] = None,
-        cp_comm_type: Optional[str] = None,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-    ):
-        super().__init__(config=config)
-
-        self.config = config
-
-        self.is_first_microbatch = True
-
-        assert (
-            self.config.context_parallel_size == 1
-        ), "Context parallelism is not supported by KitchenDotProductAttention!"
-
-        assert (
-            self.config.window_size is None
-        ), "Sliding Window Attention is not supported by KitchenDotProductAttention!"
-
-        self.layer_number = max(1, layer_number)
-        self.attn_mask_type = attn_mask_type
-        self.attention_type = attention_type  # unused for now
-
-        kv_channels = self.config.kv_channels
-        assert kv_channels is not None, "kv_channels must be set for KitchenFlashAttention"
-        projection_size = kv_channels * self.config.num_attention_heads
-
-        # Per attention head and per partition values.
-        if pg_collection is None:
-            raise ValueError("DotProductAttention was called without ProcessGroupCollection")
-        else:
-            assert hasattr(
-                pg_collection, 'tp'
-            ), "DotProductAttention pg_collection must have tp process group"
-
-        world_size = pg_collection.tp.size()
-        self.hidden_size_per_partition = divide(projection_size, world_size)
-        self.hidden_size_per_attention_head = divide(projection_size, config.num_attention_heads)
-        self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, world_size)
-        self.num_query_groups_per_partition = divide(self.config.num_query_groups, world_size)
-
-        coeff = None
-        if softmax_scale is None:
-            self.softmax_scale = 1.0 / math.sqrt(self.hidden_size_per_attention_head)
-        else:
-            self.softmax_scale = softmax_scale
-
-        if self.config.apply_query_key_layer_scaling:
-            coeff = self.layer_number
-            self.softmax_scale /= coeff
-
-        self.attention_dropout = (
-            self.config.attention_dropout if attention_dropout is None else attention_dropout
-        )
-
-        self.init_finished = False
-
-    def finish_init(self, quantization_config: QuantizationConfig):
-        """
-        Finishes the initialization of the KitchenFlashAttention module.
-        """
-        extra_kwargs = _get_extra_kitchen_kwargs(self.config)
-        self.kitchen_quant_params = KitchenQuantizationParams.parse_from_config(quantization_config)
-        extra_kwargs["qfa_params"] = self.kitchen_quant_params.qfa_params
-
-        self.flash_attention_module = KitchenFlashAttentionModule(
-            num_attention_heads=self.config.num_attention_heads,
-            kv_channels=self.config.kv_channels,
-            num_gqa_groups=self.config.num_query_groups,
-            attention_dropout=self.attention_dropout,
-            qkv_format='sbhd',
-            attn_mask_type=self.attn_mask_type.name,
-            window_size=self.config.window_size,
-            sequence_parallel=self.config.sequence_parallel,
-            # tp_size=self.config.tensor_model_parallel_size,
-            get_rng_state_tracker=(
-                None  # TODO(Frank): Support cuda rng tracker for dropout.
-                # get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
-            ),
-            layer_number=self.layer_number,
-            attention_type=self.attention_type,
-            softmax_scale=self.softmax_scale,
-            qfa_params=self.kitchen_quant_params.qfa_params,
-        )
-        self.init_finished = True
-
-    def forward(
-        self,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        attention_mask: Optional[Tensor],
-        attn_mask_type: Optional[AttnMaskType] = None,
-        attention_bias: Optional[Tensor] = None,
-        packed_seq_params: Optional[PackedSeqParams] = None,
-    ):
-        """Forward."""
-        assert self.init_finished, "Must call finish_init before forward."
-        assert packed_seq_params is None, (
-            "Packed sequence is not supported by KitchenDotProductAttention."
-            "Please use TEDotProductAttention instead."
-        )
-        assert (
-            attention_bias is None
-        ), "Attention bias is not supported for KitchenDotProductAttention."
-        _is_first_microbatch = self.is_first_microbatch
-
-        # TODO(Frank): Handles group query attention internally to kitchen backend.
-        if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
-            key = key.repeat_interleave(
-                self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
-            )
-            value = value.repeat_interleave(
-                self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
-            )
-
-        if attn_mask_type is not None:
-            assert (
-                attn_mask_type == AttnMaskType.causal
-            ), "Only causal mask is supported for KitchenFlashAttention."
-        attn_mask_type_str = "causal"
-
-        # TODO(Frank): Figure out the story for qkv_layout
-        core_attn_out = self.flash_attention_module(
-            query,
-            key,
-            value,
-            attn_mask_type=attn_mask_type_str,
-            window_size=self.config.window_size,
-            is_first_microbatch=_is_first_microbatch,
-        )
-        self.is_first_microbatch = False
-
-        return core_attn_out
-
-
-class KitchenDotProductAttention(MegatronModule):
-    """
-    Region where selective activation recomputation is applied.
-    This region is memory intensive but less compute intensive which
-    makes activation checkpointing more efficient for LLMs (20B+).
-    See Reducing Activation Recomputation in Large Transformer Models:
-    https://arxiv.org/abs/2205.05198 for more details.
-
-    We use the following notation:
-     h: hidden size
-     n: number of attention heads
-     p: number of tensor model parallel partitions
-     b: batch size
-     s: sequence length
-    """
-
-    def __init__(
-        self,
-        config: TransformerConfig,
-        layer_number: int,
-        attn_mask_type: AttnMaskType,
-        attention_type: str,
-        attention_dropout: Optional[float] = None,
-        softmax_scale: Optional[float] = None,
-        cp_comm_type: Optional[str] = None,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-    ):
-        super().__init__(config=config)
-
-        self.config: TransformerConfig = config
-
-        assert (
-            self.config.context_parallel_size == 1
-        ), "Context parallelism is not supported by KitchenDotProductAttention!"
-
-        assert (
-            self.config.window_size is None
-        ), "Sliding Window Attention is not supported by KitchenDotProductAttention!"
-
-        self.layer_number = max(1, layer_number)
-        self.attn_mask_type = attn_mask_type
-        self.attention_type = attention_type  # unused for now
-
-        kv_channels = self.config.kv_channels
-        assert kv_channels is not None, "kv_channels must be set for KitchenFlashAttention"
-        projection_size = kv_channels * self.config.num_attention_heads
-
-        # Per attention head and per partition values.
-        if pg_collection is None:
-            raise ValueError("DotProductAttention was called without ProcessGroupCollection")
-        else:
-            assert hasattr(
-                pg_collection, 'tp'
-            ), "DotProductAttention pg_collection must have tp process group"
-
-        world_size = pg_collection.tp.size()
-        self.hidden_size_per_partition = divide(projection_size, world_size)
-        self.hidden_size_per_attention_head = divide(projection_size, config.num_attention_heads)
-        self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, world_size)
-        self.num_query_groups_per_partition = divide(self.config.num_query_groups, world_size)
-
-        coeff = None
-        if softmax_scale is None:
-            self.softmax_scale = 1.0 / math.sqrt(self.hidden_size_per_attention_head)
-        else:
-            self.softmax_scale = softmax_scale
-
-        if self.config.apply_query_key_layer_scaling:
-            coeff = self.layer_number
-            self.softmax_scale /= coeff
-
-        self.scale_mask_softmax = FusedScaleMaskSoftmax(
-            input_in_fp16=self.config.fp16,
-            input_in_bf16=self.config.bf16,
-            attn_mask_type=self.attn_mask_type,
-            scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
-            mask_func=attention_mask_func,
-            softmax_in_fp32=self.config.attention_softmax_in_fp32,
-            scale=coeff,
-        )
-
-        # Dropout. Note that for a single iteration, this layer will generate
-        # different outputs on different number of parallel partitions but
-        # on average it should not be partition dependent.
-        self.attention_dropout = torch.nn.Dropout(
-            self.config.attention_dropout if attention_dropout is None else attention_dropout
-        )
-        self.auto_grad_qbmm = QuantizedBMM
-        self.init_finished = False
-
-    def finish_init(self, quantization_config: QuantizationConfig):
-        """
-        Finishes the initialization of the KitchenDotProductAttention module.
-        """
-        extra_kwargs = _get_extra_kitchen_kwargs(self.config)
-        self.kitchen_quant_params = KitchenQuantizationParams.parse_from_config(quantization_config)
-        extra_kwargs["qattention_params"] = self.kitchen_quant_params.qattention_params
-        self.qattention_params = self.kitchen_quant_params.qattention_params
-        self.init_finished = True
-
-    def forward(
-        self,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        attention_mask: Optional[Tensor],
-        attn_mask_type: Optional[AttnMaskType] = None,
-        attention_bias: Optional[Tensor] = None,
-        packed_seq_params: Optional[PackedSeqParams] = None,
-    ) -> Tensor:
-        """Forward."""
-        assert self.init_finished, "Must call finish_init before forward."
-        assert packed_seq_params is None, (
-            "Packed sequence is not supported by KitchenDotProductAttention."
-            "Please use TEDotProductAttention instead."
-        )
-        assert (
-            attention_bias is None
-        ), "Attention bias is not supported for KitchenDotProductAttention."
-
-        # ===================================
-        # Raw attention scores. [b, n/p, s, s]
-        # ===================================
-
-        # expand the key and value [sk, b, ng, hn] -> [sk, b, np, hn]
-        # This is a noop for normal attention where ng == np. When using group query attention this
-        # creates a view that has the keys and values virtually repeated along their dimension to
-        # match the number of queries.
-
-        # attn_mask_type is not used.
-        if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
-            key = key.repeat_interleave(
-                self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
-            )
-            value = value.repeat_interleave(
-                self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
-            )
-
-        # [b, np, sq, sk]
-        output_size = (query.size(1), query.size(2), query.size(0), key.size(0))
-
-        # [sq, b, np, hn] -> [sq, b * np, hn]
-        # This owill be a simple view when doing normal attention, but in group query attention
-        # the key and value tensors are repeated to match the queries so you can't use
-        # simple strides to extract the queries.
-        query = query.reshape(output_size[2], output_size[0] * output_size[1], -1)
-        # [sk, b, np, hn] -> [sk, b * np, hn]
-        key = key.view(output_size[3], output_size[0] * output_size[1], -1)
-
-        bmm_args: List[Any] = []
-        if torch.is_grad_enabled():
-            bmm_fn = self.auto_grad_qbmm.apply
-        else:
-            bmm_fn = self.auto_grad_qbmm.forward
-            bmm_args.append(None)
-        bmm_args.extend(
-            [
-                query,
-                key,
-                self.softmax_scale,
-                False,
-                torch.is_grad_enabled(),
-                self.layer_number,
-                self.qattention_params,
-            ]
-        )
-        # Raw attention scores. [b * np, sq, sk]
-        matmul_result = bmm_fn(*bmm_args)
-
-        # change view to [b, np, sq, sk]
-        attention_scores = matmul_result.view(*output_size)
-
-        # ===========================
-        # Attention probs and dropout
-        # ===========================
-
-        # attention scores and attention mask [b, np, sq, sk]
-        attention_probs: Tensor = self.scale_mask_softmax(attention_scores, attention_mask)
-
-        # This is actually dropping out entire tokens to attend to, which might
-        # seem a bit unusual, but is taken from the original Transformer paper.
-
-        if not self.config.sequence_parallel:
-            with tensor_parallel.get_cuda_rng_tracker().fork():
-                attention_probs = self.attention_dropout(attention_probs)
-        else:
-            attention_probs = self.attention_dropout(attention_probs)
-
-        # =========================
-        # Context layer. [sq, b, hp]
-        # =========================
-
-        # value -> context layer.
-        # [sk, b, np, hn] --> [b, np, sq, hn]
-
-        # context layer shape: [b, np, sq, hn]
-        output_size = (value.size(1), value.size(2), query.size(0), value.size(3))
-
-        # change view [sk, b * np, hn]
-        value = value.view(value.size(0), output_size[0] * output_size[1], -1)
-
-        # change view [b * np, sq, sk]
-        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
-
-        # matmul: [b * np, sq, hn]
-        bmm_args = []
-        if torch.is_grad_enabled():
-            bmm_fn = self.auto_grad_qbmm.apply
-        else:
-            bmm_fn = self.auto_grad_qbmm.forward
-            bmm_args.append(None)
-        bmm_args.extend(
-            [
-                attention_probs,
-                value,
-                1.0,
-                True,
-                torch.is_grad_enabled(),
-                self.layer_number,
-                self.qattention_params,
-            ]
-        )
-        context = bmm_fn(*bmm_args)
-        # change view [b, np, sq, hn]
-        context = context.view(*output_size)
-
-        # [b, np, sq, hn] --> [sq, b, np, hn]
-        context = context.permute(2, 0, 1, 3).contiguous()
-
-        # [sq, b, np, hn] --> [sq, b, hp]
-        new_context_shape = context.size()[:-2] + (self.hidden_size_per_partition,)
-        context = context.view(*new_context_shape)
-
-        return context
-
-
-class KitchenSpecProvider(BackendSpecProvider):
-    """A protocol for providing the submodules used in Spec building."""
-
-    def __init__(
-        self,
-        fallback: BackendSpecProvider,
-        use_kitchen_attention: bool = False,
-        kitchen_attention_backend: str = "sdpa",
-    ):
-        self.fallback = fallback
-        self.use_kitchen_attention = use_kitchen_attention
-        self.kitchen_attention_backend = kitchen_attention_backend
-
-    def column_parallel_linear(self) -> type[KitchenColumnParallelLinear]:
-        """Which column parallel linear module kitchen backend uses"""
-        return KitchenColumnParallelLinear
-
-    def row_parallel_linear(self) -> type:
-        """Which row parallel linear module kitchen backend uses"""
-        return KitchenRowParallelLinear
-
-    def fuse_layernorm_and_linear(self) -> bool:
-        """Does kitchen backend support a single module for layernorm and linear"""
-        # NOTE(kwyss): This is coupled with get_mlp_module_spec_for_backend and
-        # the initialization of TransformerLayerSubmodules such as in
-        # get_gpt_layer_local_spec or get_gpt_layer_with_transformer_engine_spec
-        # where an explicit norm may be provided. Kitchen extension chooses to
-        # match the topology of the fallback with this code.
-        # Arguably, we should pass the info down to get_mlp_module_spec_for_backend
-        # explicitly about whether to include a norm.
-        return self.fallback.fuse_layernorm_and_linear()
-
-    def column_parallel_layer_norm_linear(self) -> type[KitchenLayerNormColumnParallelLinear]:
-        """Which module for sequential layernorm and linear"""
-        return KitchenLayerNormColumnParallelLinear
-
-    def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type:
-        """Which module to use for layer norm"""
-        return self.fallback.layer_norm(rms_norm=rms_norm, for_qk=for_qk)
-
-    def core_attention(
-        self,
-    ) -> type[KitchenDotProductAttention] | type[KitchenFlashAttention] | type:
-        """Which module to use for attention"""
-        if not self.use_kitchen_attention:
-            log_single_rank(
-                logger,
-                logging.WARNING,
-                "KitchenSpecProvider: Using fallback (likely TE) as core attention.",
-            )
-            return self.fallback.core_attention()
-
-        if self.kitchen_attention_backend == "sdpa":
-            log_single_rank(
-                logger,
-                logging.WARNING,
-                "KitchenSpecProvider: Using Kitchen SDPA as core attention.",
-            )
-            return KitchenDotProductAttention
-        elif self.kitchen_attention_backend == "fa":
-            log_single_rank(
-                logger, logging.WARNING, "KitchenSpecProvider: Using Kitchen FA as core attention."
-            )
-            return KitchenFlashAttention
-        else:
-            raise ValueError(
-                f"Invalid kitchen_attention_backend: {self.kitchen_attention_backend}. "
-                "Must be 'sdpa' or 'fa'."
-            )
-
-    def grouped_mlp_modules(
-        self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: bool
-    ) -> Tuple[type, Optional[MLPSubmodules]]:
-        """Which module and submodules to use for grouped mlp"""
-        if moe_use_grouped_gemm and not moe_use_legacy_grouped_gemm:
-            # NOTE: TEGroupedMLP is a bit of a misnomer.
-            # It doesn't strictly require TE except for the GroupedLinear,
-            # which Kitchen also provides an implementation of.
-            return TEGroupedMLP, MLPSubmodules(
-                linear_fc1=KitchenColumnParallelGroupedLinear,
-                linear_fc2=KitchenRowParallelGroupedLinear,
-            )
-        elif moe_use_grouped_gemm:
-            warnings.warn(
-                "The legacy GroupedMLP will be deprecated in Megatron-Core v0.12.0. "
-                "Please update the TransformerEngine to version>=1.7.0 and use TEGroupedMLP."
-            )
-            return GroupedMLP, None
-        else:
-            return SequentialMLP, MLPSubmodules(
-                linear_fc1=KitchenColumnParallelLinear, linear_fc2=KitchenRowParallelLinear
-            )
-
-    def activation_func(self) -> type:
-        """Which module to use for activation function"""
-        return self.fallback.activation_func()
+HAVE_KITCHEN = False
+
+from unittest.mock import MagicMock
+
+AutogradFunctionImplementation = MagicMock()
+KitchenSpecProvider = MagicMock()
+
+QAttentionParamsConfigSchema = MagicMock()
+QFlashAttentionParamsConfigSchema = MagicMock()
+QLinearParamsConfigSchema = MagicMock()
+QLinearParams = MagicMock()
+QuantizeRecipe = MagicMock()
+QuantizeRecipeAttnBMM = MagicMock()
+get_qattention_params_from_predefined = MagicMock()
+get_qfa_params_from_recipe_name = MagicMock()
+get_qlinear_params_from_predefined = MagicMock()
+get_qlinear_params_from_qat_params = MagicMock()
+
+KitchenColumnParallelGroupedLinear = MagicMock()
+KitchenColumnParallelLinear = MagicMock()
+KitchenDotProductAttention = MagicMock()
+KitchenFlashAttention = MagicMock()
+KitchenLayerNormColumnParallelLinear = MagicMock()
+KitchenRowParallelGroupedLinear = MagicMock()
+KitchenRowParallelLinear = MagicMock()
+
+# N.B. Kitchen extension is not released publicly.
+# This extension is just a stub.
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index 974e33f88e8..49501ee54eb 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -50,11 +50,8 @@
     HAVE_TE = False
 
 try:
-    import nvidia_kitchen  # type: ignore[import-not-found]  # pylint: disable=unused-import
+    from megatron.core.extensions.kitchen import HAVE_KITCHEN, KitchenSpecProvider
 
-    from megatron.core.extensions.kitchen import KitchenSpecProvider
-
-    HAVE_KITCHEN = True
 except ImportError:
     HAVE_KITCHEN = False
 
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index 2bc3949a421..ad80bcfe4e4 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -107,6 +107,10 @@ def __init__(
         if self.config.gated_linear_unit:
             ffn_hidden_size *= 2
             fc1_stride = 2
+            if self.config.use_kitchen:
+                # Kitchen Linear doesn't support stride != 1.
+                # Weight resharding across TP sizes will have aforementioned problems.
+                fc1_stride = 1
         else:
             fc1_stride = 1
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index dd4cf5b5c89..f269ad02879 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -3298,13 +3298,12 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser):
     If kitchen isn't available, nothing to do here, return unchanged parser
     """
     try:
-        from megatron.core.extensions.kitchen import KitchenSpecProvider
+        from megatron.core.extensions.kitchen import KitchenSpecProvider, HAVE_KITCHEN
 
-        have_kitchen = True
     except (ImportError, ModuleNotFoundError):
-        have_kitchen = False
+        HAVE_KITCHEN = False
 
-    if have_kitchen:
+    if HAVE_KITCHEN:
         group = parser.add_argument_group(title="kitchen")
         recipe_or_config_group = group.add_mutually_exclusive_group(required=False)
         recipe_or_config_group.add_argument(
diff --git a/tests/unit_tests/extension/test_kitchen_sdpa.py b/tests/unit_tests/extension/test_kitchen_sdpa.py
index c0503ffae83..6875b005c0b 100644
--- a/tests/unit_tests/extension/test_kitchen_sdpa.py
+++ b/tests/unit_tests/extension/test_kitchen_sdpa.py
@@ -8,7 +8,6 @@
 import torch
 
 from megatron.core import parallel_state
-from megatron.core.extensions.kitchen import KitchenDotProductAttention, KitchenFlashAttention
 from megatron.core.extensions.transformer_engine import TEDotProductAttention
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.quantization.quant_config import RecipeConfig
@@ -21,14 +20,18 @@
 from tests.unit_tests.test_utilities import Utils
 
 try:
-    import nvidia_kitchen  # type: ignore[import-not-found]
+    from megatron.core.extensions.kitchen import (
+        HAVE_KITCHEN,
+        KitchenDotProductAttention,
+        KitchenFlashAttention,
+    )
 
-    HAVE_KITCHEN = True
 except ImportError:
     from unittest.mock import MagicMock
 
     HAVE_KITCHEN = False
-    nvidia_kitchen = MagicMock()
+    KitchenDotProductAttention = MagicMock()
+    KitchenFlashAttention = MagicMock()
 
 try:
     import transformer_engine  # type: ignore[import-untyped]
diff --git a/tests/unit_tests/models/test_gpt_model_quantization.py b/tests/unit_tests/models/test_gpt_model_quantization.py
index b1fbf9f54ef..e993c9be8d2 100644
--- a/tests/unit_tests/models/test_gpt_model_quantization.py
+++ b/tests/unit_tests/models/test_gpt_model_quantization.py
@@ -16,10 +16,8 @@
     HAVE_TE = False
 
 try:
-    import nvidia_kitchen  # type: ignore[import-not-found]
-
-    HAVE_KITCHEN = True
     from megatron.core.extensions.kitchen import (
+        HAVE_KITCHEN,
         KitchenColumnParallelGroupedLinear,
         KitchenColumnParallelLinear,
         KitchenDotProductAttention,
diff --git a/tests/unit_tests/transformer/test_quantization_config.py b/tests/unit_tests/transformer/test_quantization_config.py
index 03556a1c638..6d92e59b9a9 100644
--- a/tests/unit_tests/transformer/test_quantization_config.py
+++ b/tests/unit_tests/transformer/test_quantization_config.py
@@ -7,28 +7,20 @@
 from megatron.core.quantization.quant_config import GlobMatcher, MatchContext, RecipeConfig
 
 try:
-    import nvidia_kitchen  # type: ignore[import-not-found]
-    from nvidia_kitchen.config import (  # type: ignore[import-not-found]
+    from megatron.core.extensions.kitchen import (
+        HAVE_KITCHEN,
         AutogradFunctionImplementation,
+        QAttentionParamsConfigSchema,
+        QFlashAttentionParamsConfigSchema,
+        QLinearParamsConfigSchema,
         QuantizeRecipe,
-        get_qlinear_params_from_predefined,
-    )
-    from nvidia_kitchen.config_attention_recipe import (  # type: ignore[import-not-found]
         QuantizeRecipeAttnBMM,
         get_qattention_params_from_predefined,
-    )
-    from nvidia_kitchen.config_fa_recipe import (  # type: ignore[import-not-found]
         get_qfa_params_from_recipe_name,
+        get_qlinear_params_from_predefined,
     )
 
-    from megatron.core.extensions.kitchen import (  # type: ignore[import-not-found]
-        QAttentionParamsConfigSchema,
-        QFlashAttentionParamsConfigSchema,
-        QLinearParamsConfigSchema,
-    )
-
-    HAVE_KITCHEN = True
-except ImportError as e:
+except ImportError:
     HAVE_KITCHEN = False
 
 

From 287d2f47c3ee724fc5dfda6e798e17156da1f33f Mon Sep 17 00:00:00 2001
From: Jon Barker <jbarker@nvidia.com>
Date: Wed, 28 Jan 2026 23:50:48 -0700
Subject: [PATCH 64/79] Fix RL optimizer offload (#3112)

Co-authored-by: Jon Barker <jbarker@oci-hsg-cs-001-vscode-01.cm.cluster>
---
 examples/rl/model_configs/nemotron6_3b_moe.sh |   6 -
 .../distributed/distributed_data_parallel.py  |  38 ++++
 .../core/distributed/param_and_grad_buffer.py |  36 ++++
 megatron/rl/rl_utils.py                       |  24 ++-
 tests/unit_tests/rl/test_rl_utils.py          | 175 ++++++++++++++++++
 5 files changed, 264 insertions(+), 15 deletions(-)

diff --git a/examples/rl/model_configs/nemotron6_3b_moe.sh b/examples/rl/model_configs/nemotron6_3b_moe.sh
index 8efe0b2debb..eff4f6cf0b3 100644
--- a/examples/rl/model_configs/nemotron6_3b_moe.sh
+++ b/examples/rl/model_configs/nemotron6_3b_moe.sh
@@ -5,12 +5,6 @@ EP=${EP:-32}
 NODES_REQUIRED=${NODES_REQUIRED:-4}
 LLM="nemotron6_3b_moe"
 
-ROOT_DIR="/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/nemotron6"
-
-CHECKPOINT="${ROOT_DIR}/3b_hybrid_moe/checkpoints/phase2_lc_reinit_emb/"
-
-TOKENIZER_MODEL="${ROOT_DIR}/tokenizers/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json"
-
 echo "Using Nemotron6 3B MOE model checkpoint"
 SCRIPT_PATH="${BASH_SOURCE[0]}"
 source $(dirname $SCRIPT_PATH)/common.sh
diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index 421b279b17d..55179ff3024 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -571,3 +571,41 @@ def broadcast_params(self):
                 src=torch.distributed.get_global_rank(data_parallel_group, 0),
                 group=data_parallel_group,
             )
+
+    def offload_grad_buffers(self, synchronize: bool = True, empty_cache: bool = True) -> None:
+        """
+        Free all grad_data tensors to release GPU memory.
+
+        Uses storage().resize_(0) to release memory while keeping tensor views intact.
+        All bucket.grad_data and param.main_grad views remain valid tensor objects
+        (though accessing them during offload is undefined behavior).
+
+        Args:
+            synchronize: Whether to call torch.cuda.synchronize() before freeing.
+            empty_cache: Whether to call torch.cuda.empty_cache() after freeing.
+        """
+        if synchronize:
+            torch.cuda.synchronize()
+
+        for buffer in self.buffers + self.expert_parallel_buffers:
+            buffer.offload_to_cpu(move_params=False, move_grads=True)
+
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+    def restore_grad_buffers(self, synchronize: bool = True) -> None:
+        """
+        Reallocate grad_data tensors on GPU.
+
+        All existing views (bucket.grad_data, param.main_grad) automatically
+        become valid again since they share the same storage. The grad_data
+        is zeroed after reallocation.
+
+        Args:
+            synchronize: Whether to call torch.cuda.synchronize() after allocation.
+        """
+        for buffer in self.buffers + self.expert_parallel_buffers:
+            buffer.reload_from_cpu(move_params=False, move_grads=True)
+
+        if synchronize:
+            torch.cuda.synchronize()
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index 7abdaab103b..b9480533d7a 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -801,6 +801,10 @@ def _does_param_require_new_bucket(param):
                     requires_grad=False,
                 )
 
+        self.grad_data_size = 0
+        self.param_data_size = 0
+        self.param_data_cpu = None
+
         # Finally, map param.data and param.main_grad fields to buffers.
         bucket_params = []
         bucket_start_index = 0
@@ -951,6 +955,38 @@ def reset(self):
         """
         self.grad_data.zero_()
 
+    def offload_to_cpu(self, move_params: bool = True, move_grads: bool = True) -> None:
+        """
+        Offload the buffers to CPU.
+        """
+        if move_grads and self.grad_data is not None and self.grad_data.storage().size() > 0:
+            self.grad_data_size = self.grad_data.storage().size()
+            self.grad_data.storage().resize_(0)
+        if move_params and self.param_data is not None and self.param_data.storage().size() > 0:
+            self.param_data_size = self.param_data.storage().size()
+            if self.param_data_cpu is not None:
+                self.param_data_cpu.copy_(self.param_data, non_blocking=True)
+            else:
+                self.param_data_cpu = self.param_data.cpu().pin_memory()
+            self.param_data.storage().resize_(0)
+
+    def reload_from_cpu(self, move_params: bool = True, move_grads: bool = True):
+        """
+        Reload the buffers from CPU.
+        """
+        if (
+            move_params
+            and self.param_data is not None
+            and self.param_data_cpu is not None
+            and self.param_data.storage().size() == 0
+        ):
+            self.param_data.storage().resize_(self.param_data_size)
+            self.param_data.copy_(self.param_data_cpu, non_blocking=True)
+        if move_grads and self.grad_data is not None and self.grad_data_size > 0:
+            self.grad_data.storage().resize_(self.grad_data_size)
+            self.grad_data.zero_()
+            self.grad_data_size = 0
+
 
 def partition_buckets(
     buffers: List[_ParamAndGradBuffer], force_single_bucket_group: bool = False
diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index 26b990236ba..973a396b909 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -104,7 +104,6 @@ def _maybe_prefetch_separate_inference_model_weights(model_core, *, to_cpu: bool
         return
     if args.rl_inference_model_unified_memory_level != 1:
         return
-
     device = -1 if to_cpu else int(torch.cuda.current_device())
     # Note: include_buffers=False because buffers created with explicit device= in register_buffer()
     # are not allocated via the UVM mempool and will fail UVM operations. Only parameters are UVM-allocated.
@@ -461,13 +460,13 @@ def get_environment_rollouts(
     args = get_args()
     nvtx_range = get_nvtx_range()
 
+    if args.rl_offload_optimizer_during_inference:
+        with nvtx_range("offload-optimizer-state-and-grad-buffers-during-inference"):
+            model[0].offload_grad_buffers()
+            optimizer.offload_to_cpu()
+             
     # If we have seperate training and inference models we to refit weights from the training model to the inference model.
     if inference_model is not None:
-        if args.rl_offload_optimizer_during_inference:
-            with nvtx_range("offload-optimizer-before-refit"):
-                optimizer.offload_to_cpu()
-                torch.cuda.empty_cache()
-
         # If the separate inference model weights were prefetched to CPU while idle, bring them
         # back to GPU before refit/copy and before any CUDA-graph'd inference.
         with nvtx_range("prefetch-inference-model-weights-to-gpu"):
@@ -496,7 +495,7 @@ def get_environment_rollouts(
             optimizer,
             args.cuda_graph_impl,
             args.rl_reset_cuda_graphs,
-            args.rl_offload_optimizer_during_inference,
+            False, # offload optimizer during rollout collection is handled above
             args.rl_offload_kv_cache_during_training,
             args.rl_remove_kv_cache_during_training,
         ) as inference_interface:
@@ -536,6 +535,11 @@ def get_environment_rollouts(
             torch.distributed.broadcast_object_list(rollouts, src=0)
         logger.debug(f"Got rollouts on rank {rank}")
 
+    if args.rl_offload_optimizer_during_inference:
+        with nvtx_range("restore-optimizer-state-and-grad-buffers-after-inference"):
+            model[0].restore_grad_buffers()
+            optimizer.restore_from_cpu()
+
     if lang_rl_log_dir and rank == get_pg_rank(inference_pg_collection.tp):
         with open(
             lang_rl_log_dir
@@ -1609,7 +1613,8 @@ def megatron_rl_inference_mode(
     with torch.no_grad():
 
         if offload_optimizer_during_inference:
-            with nvtx_range("offload-optimizer-before-inference"):
+            with nvtx_range("offload-optimizer-state-and-grad-buffers-before-inference"):
+                model[0].offload_grad_buffers()
                 optimizer.offload_to_cpu()
 
         # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to.
@@ -1672,7 +1677,8 @@ def megatron_rl_inference_mode(
             _maybe_prefetch_separate_inference_model_weights(model_core, to_cpu=True)
 
         if offload_optimizer_during_inference:
-            with nvtx_range("onload-optimizer-after-inference"):
+            with nvtx_range("onload-optimizer-state-and-grad-buffers-after-inference"):
+                model[0].restore_grad_buffers()
                 optimizer.restore_from_cpu()
 
         lang_module.train()
diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py
index d3570bee108..8747f2e8c35 100644
--- a/tests/unit_tests/rl/test_rl_utils.py
+++ b/tests/unit_tests/rl/test_rl_utils.py
@@ -7,11 +7,16 @@
 import pytest
 import torch
 
+from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
 from megatron.core.enums import ModelType
 from megatron.core.models.common.language_module.language_module import LanguageModule
+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.num_microbatches_calculator import destroy_num_microbatches_calculator
+from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
 from megatron.core.pipeline_parallel.utils import is_pp_last_stage
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.rl import rl_utils
 from megatron.rl.agent.api import TokenRollout
@@ -416,3 +421,173 @@ def test_prepare_trajectories(self, use_sequence_packing):
                 got_t = got if torch.is_tensor(got) else torch.tensor(got, dtype=torch.float32)
                 exp_t = torch.tensor(exp, dtype=torch.float32, device=got_t.device)
                 torch.testing.assert_close(got_t, exp_t, rtol=0, atol=0)
+
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
+            for tp, pp in itertools.product([1, 2], [1, 2])
+            if tp * pp <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    def test_grad_buffer_offload(self, initialize_model_parallel):
+        """Test that grad buffer offload/restore correctly frees and restores GPU memory."""
+        world_size, dp, tp, pp = initialize_model_parallel
+        self.create_test_args(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
+
+        model_parallel_cuda_manual_seed(123)
+
+        # Create a realistic GPTModel as used in RL training
+        transformer_config = TransformerConfig(
+            num_layers=2, hidden_size=64, num_attention_heads=4, use_cpu_initialization=True
+        )
+        gpt_model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
+            vocab_size=256,
+            max_sequence_length=32,
+        ).cuda()
+
+        ddp_config = DistributedDataParallelConfig(
+            grad_reduce_in_fp32=True,
+            use_distributed_optimizer=True,
+            overlap_grad_reduce=False,
+            bucket_size=None,  # Single bucket for simplicity
+        )
+
+        ddp_model = DistributedDataParallel(
+            transformer_config, ddp_config=ddp_config, module=gpt_model
+        )
+
+        all_buffers = ddp_model.buffers + ddp_model.expert_parallel_buffers
+
+        # Verify initial storage is allocated
+        initial_sizes = [buf.grad_data.storage().size() for buf in all_buffers]
+        assert all(size > 0 for size in initial_sizes), "Expected non-zero initial storage"
+
+        # Offload grad buffers to CPU
+        ddp_model.offload_grad_buffers()
+
+        # Verify storage is released
+        for buf in all_buffers:
+            assert buf.grad_data.storage().size() == 0, "Expected zero storage after offload"
+
+        # Restore grad buffers to GPU
+        ddp_model.restore_grad_buffers()
+
+        # Verify storage is restored
+        restored_sizes = [buf.grad_data.storage().size() for buf in all_buffers]
+        assert (
+            initial_sizes == restored_sizes
+        ), f"Expected restored sizes {restored_sizes} to match initial {initial_sizes}"
+
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
+            for tp, pp in itertools.product([1, 2], [1, 2])
+            if tp * pp <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    def test_optimizer_offload(self, initialize_model_parallel):
+        """Test that optimizer offload_to_cpu/restore_from_cpu correctly moves state to/from CPU."""
+        world_size, dp, tp, pp = initialize_model_parallel
+        self.create_test_args(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
+        model_parallel_cuda_manual_seed(123)
+
+        # Create a realistic GPTModel as used in RL training
+        transformer_config = TransformerConfig(
+            num_layers=2, hidden_size=64, num_attention_heads=4, use_cpu_initialization=True
+        )
+        gpt_model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
+            vocab_size=256,
+            max_sequence_length=32,
+        ).cuda()
+
+        ddp_config = DistributedDataParallelConfig(
+            grad_reduce_in_fp32=True,
+            use_distributed_optimizer=True,
+            overlap_grad_reduce=False,
+            bucket_size=None,  # Single bucket for simplicity
+        )
+
+        ddp_model = DistributedDataParallel(
+            transformer_config, ddp_config=ddp_config, module=gpt_model
+        )
+
+        # Create optimizer
+        optimizer_config = OptimizerConfig(
+            optimizer='adam', bf16=True, use_distributed_optimizer=True
+        )
+        optimizer = get_megatron_optimizer(optimizer_config, [ddp_model])
+
+        # Manually initialize optimizer state (simulating what happens after first step)
+        # This avoids needing to run a full forward/backward/step cycle
+        for opt in optimizer.chained_optimizers:
+            if hasattr(opt, 'optimizer') and opt.optimizer is not None:
+                for group in opt.optimizer.param_groups:
+                    for p in group['params']:
+                        if len(opt.optimizer.state[p]) == 0:
+                            # Initialize Adam state (exp_avg and exp_avg_sq) on GPU
+                            opt.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data)
+                            opt.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data)
+                            opt.optimizer.state[p]['step'] = torch.tensor(1)
+
+        # Helper to check if optimizer state tensors are on GPU or CPU
+        def get_optimizer_state_devices():
+            devices = set()
+            for opt in optimizer.chained_optimizers:
+                if hasattr(opt, 'optimizer') and opt.optimizer is not None:
+                    for state_dict in opt.optimizer.state.values():
+                        for v in state_dict.values():
+                            if isinstance(v, torch.Tensor):
+                                devices.add(str(v.device))
+            return devices
+
+        # Verify optimizer state is initially on GPU
+        initial_devices = get_optimizer_state_devices()
+        assert any(
+            'cuda' in d for d in initial_devices
+        ), f"Expected optimizer state on GPU initially, got devices: {initial_devices}"
+
+        # Record GPU memory before offload
+        torch.cuda.synchronize()
+        memory_before_offload = torch.cuda.memory_allocated()
+
+        # Offload optimizer state to CPU
+        optimizer.offload_to_cpu()
+
+        # Verify GPU memory decreased (optimizer state should be freed)
+        torch.cuda.synchronize()
+        memory_after_offload = torch.cuda.memory_allocated()
+        assert memory_after_offload < memory_before_offload, (
+            f"Expected GPU memory to decrease after offload. "
+            f"Before: {memory_before_offload}, After: {memory_after_offload}"
+        )
+
+        # Verify optimizer state is now on CPU
+        offloaded_devices = get_optimizer_state_devices()
+        assert all(
+            'cpu' in d for d in offloaded_devices
+        ), f"Expected all optimizer state on CPU after offload, got devices: {offloaded_devices}"
+
+        # Restore optimizer state to GPU
+        optimizer.restore_from_cpu()
+
+        # Verify optimizer state is back on GPU
+        restored_devices = get_optimizer_state_devices()
+        assert any(
+            'cuda' in d for d in restored_devices
+        ), f"Expected optimizer state on GPU after restore, got devices: {restored_devices}"
+
+        # Verify GPU memory increased after restore (optimizer state reallocated)
+        torch.cuda.synchronize()
+        memory_after_restore = torch.cuda.memory_allocated()
+        assert memory_after_restore > memory_after_offload, (
+            f"Expected GPU memory to increase after restore. "
+            f"After offload: {memory_after_offload}, After restore: {memory_after_restore}"
+        )

From 3955c49ed9af5e5b38dccdd30c1323c00b9bcd29 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Thu, 29 Jan 2026 11:42:52 +0100
Subject: [PATCH 65/79] Revert "Fix RL optimizer offload (#3112)" (#3141)

---
 examples/rl/model_configs/nemotron6_3b_moe.sh |   6 +
 .../distributed/distributed_data_parallel.py  |  38 ----
 .../core/distributed/param_and_grad_buffer.py |  36 ----
 megatron/rl/rl_utils.py                       |  24 +--
 tests/unit_tests/rl/test_rl_utils.py          | 175 ------------------
 5 files changed, 15 insertions(+), 264 deletions(-)

diff --git a/examples/rl/model_configs/nemotron6_3b_moe.sh b/examples/rl/model_configs/nemotron6_3b_moe.sh
index eff4f6cf0b3..8efe0b2debb 100644
--- a/examples/rl/model_configs/nemotron6_3b_moe.sh
+++ b/examples/rl/model_configs/nemotron6_3b_moe.sh
@@ -5,6 +5,12 @@ EP=${EP:-32}
 NODES_REQUIRED=${NODES_REQUIRED:-4}
 LLM="nemotron6_3b_moe"
 
+ROOT_DIR="/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/nemotron6"
+
+CHECKPOINT="${ROOT_DIR}/3b_hybrid_moe/checkpoints/phase2_lc_reinit_emb/"
+
+TOKENIZER_MODEL="${ROOT_DIR}/tokenizers/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json"
+
 echo "Using Nemotron6 3B MOE model checkpoint"
 SCRIPT_PATH="${BASH_SOURCE[0]}"
 source $(dirname $SCRIPT_PATH)/common.sh
diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index 55179ff3024..421b279b17d 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -571,41 +571,3 @@ def broadcast_params(self):
                 src=torch.distributed.get_global_rank(data_parallel_group, 0),
                 group=data_parallel_group,
             )
-
-    def offload_grad_buffers(self, synchronize: bool = True, empty_cache: bool = True) -> None:
-        """
-        Free all grad_data tensors to release GPU memory.
-
-        Uses storage().resize_(0) to release memory while keeping tensor views intact.
-        All bucket.grad_data and param.main_grad views remain valid tensor objects
-        (though accessing them during offload is undefined behavior).
-
-        Args:
-            synchronize: Whether to call torch.cuda.synchronize() before freeing.
-            empty_cache: Whether to call torch.cuda.empty_cache() after freeing.
-        """
-        if synchronize:
-            torch.cuda.synchronize()
-
-        for buffer in self.buffers + self.expert_parallel_buffers:
-            buffer.offload_to_cpu(move_params=False, move_grads=True)
-
-        if empty_cache:
-            torch.cuda.empty_cache()
-
-    def restore_grad_buffers(self, synchronize: bool = True) -> None:
-        """
-        Reallocate grad_data tensors on GPU.
-
-        All existing views (bucket.grad_data, param.main_grad) automatically
-        become valid again since they share the same storage. The grad_data
-        is zeroed after reallocation.
-
-        Args:
-            synchronize: Whether to call torch.cuda.synchronize() after allocation.
-        """
-        for buffer in self.buffers + self.expert_parallel_buffers:
-            buffer.reload_from_cpu(move_params=False, move_grads=True)
-
-        if synchronize:
-            torch.cuda.synchronize()
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index b9480533d7a..7abdaab103b 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -801,10 +801,6 @@ def _does_param_require_new_bucket(param):
                     requires_grad=False,
                 )
 
-        self.grad_data_size = 0
-        self.param_data_size = 0
-        self.param_data_cpu = None
-
         # Finally, map param.data and param.main_grad fields to buffers.
         bucket_params = []
         bucket_start_index = 0
@@ -955,38 +951,6 @@ def reset(self):
         """
         self.grad_data.zero_()
 
-    def offload_to_cpu(self, move_params: bool = True, move_grads: bool = True) -> None:
-        """
-        Offload the buffers to CPU.
-        """
-        if move_grads and self.grad_data is not None and self.grad_data.storage().size() > 0:
-            self.grad_data_size = self.grad_data.storage().size()
-            self.grad_data.storage().resize_(0)
-        if move_params and self.param_data is not None and self.param_data.storage().size() > 0:
-            self.param_data_size = self.param_data.storage().size()
-            if self.param_data_cpu is not None:
-                self.param_data_cpu.copy_(self.param_data, non_blocking=True)
-            else:
-                self.param_data_cpu = self.param_data.cpu().pin_memory()
-            self.param_data.storage().resize_(0)
-
-    def reload_from_cpu(self, move_params: bool = True, move_grads: bool = True):
-        """
-        Reload the buffers from CPU.
-        """
-        if (
-            move_params
-            and self.param_data is not None
-            and self.param_data_cpu is not None
-            and self.param_data.storage().size() == 0
-        ):
-            self.param_data.storage().resize_(self.param_data_size)
-            self.param_data.copy_(self.param_data_cpu, non_blocking=True)
-        if move_grads and self.grad_data is not None and self.grad_data_size > 0:
-            self.grad_data.storage().resize_(self.grad_data_size)
-            self.grad_data.zero_()
-            self.grad_data_size = 0
-
 
 def partition_buckets(
     buffers: List[_ParamAndGradBuffer], force_single_bucket_group: bool = False
diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index 973a396b909..26b990236ba 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -104,6 +104,7 @@ def _maybe_prefetch_separate_inference_model_weights(model_core, *, to_cpu: bool
         return
     if args.rl_inference_model_unified_memory_level != 1:
         return
+
     device = -1 if to_cpu else int(torch.cuda.current_device())
     # Note: include_buffers=False because buffers created with explicit device= in register_buffer()
     # are not allocated via the UVM mempool and will fail UVM operations. Only parameters are UVM-allocated.
@@ -460,13 +461,13 @@ def get_environment_rollouts(
     args = get_args()
     nvtx_range = get_nvtx_range()
 
-    if args.rl_offload_optimizer_during_inference:
-        with nvtx_range("offload-optimizer-state-and-grad-buffers-during-inference"):
-            model[0].offload_grad_buffers()
-            optimizer.offload_to_cpu()
-             
     # If we have seperate training and inference models we to refit weights from the training model to the inference model.
     if inference_model is not None:
+        if args.rl_offload_optimizer_during_inference:
+            with nvtx_range("offload-optimizer-before-refit"):
+                optimizer.offload_to_cpu()
+                torch.cuda.empty_cache()
+
         # If the separate inference model weights were prefetched to CPU while idle, bring them
         # back to GPU before refit/copy and before any CUDA-graph'd inference.
         with nvtx_range("prefetch-inference-model-weights-to-gpu"):
@@ -495,7 +496,7 @@ def get_environment_rollouts(
             optimizer,
             args.cuda_graph_impl,
             args.rl_reset_cuda_graphs,
-            False, # offload optimizer during rollout collection is handled above
+            args.rl_offload_optimizer_during_inference,
             args.rl_offload_kv_cache_during_training,
             args.rl_remove_kv_cache_during_training,
         ) as inference_interface:
@@ -535,11 +536,6 @@ def get_environment_rollouts(
             torch.distributed.broadcast_object_list(rollouts, src=0)
         logger.debug(f"Got rollouts on rank {rank}")
 
-    if args.rl_offload_optimizer_during_inference:
-        with nvtx_range("restore-optimizer-state-and-grad-buffers-after-inference"):
-            model[0].restore_grad_buffers()
-            optimizer.restore_from_cpu()
-
     if lang_rl_log_dir and rank == get_pg_rank(inference_pg_collection.tp):
         with open(
             lang_rl_log_dir
@@ -1613,8 +1609,7 @@ def megatron_rl_inference_mode(
     with torch.no_grad():
 
         if offload_optimizer_during_inference:
-            with nvtx_range("offload-optimizer-state-and-grad-buffers-before-inference"):
-                model[0].offload_grad_buffers()
+            with nvtx_range("offload-optimizer-before-inference"):
                 optimizer.offload_to_cpu()
 
         # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to.
@@ -1677,8 +1672,7 @@ def megatron_rl_inference_mode(
             _maybe_prefetch_separate_inference_model_weights(model_core, to_cpu=True)
 
         if offload_optimizer_during_inference:
-            with nvtx_range("onload-optimizer-state-and-grad-buffers-after-inference"):
-                model[0].restore_grad_buffers()
+            with nvtx_range("onload-optimizer-after-inference"):
                 optimizer.restore_from_cpu()
 
         lang_module.train()
diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py
index 8747f2e8c35..d3570bee108 100644
--- a/tests/unit_tests/rl/test_rl_utils.py
+++ b/tests/unit_tests/rl/test_rl_utils.py
@@ -7,16 +7,11 @@
 import pytest
 import torch
 
-from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
 from megatron.core.enums import ModelType
 from megatron.core.models.common.language_module.language_module import LanguageModule
-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.num_microbatches_calculator import destroy_num_microbatches_calculator
-from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
 from megatron.core.pipeline_parallel.utils import is_pp_last_stage
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.rl import rl_utils
 from megatron.rl.agent.api import TokenRollout
@@ -421,173 +416,3 @@ def test_prepare_trajectories(self, use_sequence_packing):
                 got_t = got if torch.is_tensor(got) else torch.tensor(got, dtype=torch.float32)
                 exp_t = torch.tensor(exp, dtype=torch.float32, device=got_t.device)
                 torch.testing.assert_close(got_t, exp_t, rtol=0, atol=0)
-
-    @pytest.mark.parametrize(
-        "initialize_model_parallel",
-        [
-            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
-            for tp, pp in itertools.product([1, 2], [1, 2])
-            if tp * pp <= Utils.world_size
-        ],
-        indirect=["initialize_model_parallel"],
-    )
-    def test_grad_buffer_offload(self, initialize_model_parallel):
-        """Test that grad buffer offload/restore correctly frees and restores GPU memory."""
-        world_size, dp, tp, pp = initialize_model_parallel
-        self.create_test_args(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
-
-        model_parallel_cuda_manual_seed(123)
-
-        # Create a realistic GPTModel as used in RL training
-        transformer_config = TransformerConfig(
-            num_layers=2, hidden_size=64, num_attention_heads=4, use_cpu_initialization=True
-        )
-        gpt_model = GPTModel(
-            config=transformer_config,
-            transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
-            vocab_size=256,
-            max_sequence_length=32,
-        ).cuda()
-
-        ddp_config = DistributedDataParallelConfig(
-            grad_reduce_in_fp32=True,
-            use_distributed_optimizer=True,
-            overlap_grad_reduce=False,
-            bucket_size=None,  # Single bucket for simplicity
-        )
-
-        ddp_model = DistributedDataParallel(
-            transformer_config, ddp_config=ddp_config, module=gpt_model
-        )
-
-        all_buffers = ddp_model.buffers + ddp_model.expert_parallel_buffers
-
-        # Verify initial storage is allocated
-        initial_sizes = [buf.grad_data.storage().size() for buf in all_buffers]
-        assert all(size > 0 for size in initial_sizes), "Expected non-zero initial storage"
-
-        # Offload grad buffers to CPU
-        ddp_model.offload_grad_buffers()
-
-        # Verify storage is released
-        for buf in all_buffers:
-            assert buf.grad_data.storage().size() == 0, "Expected zero storage after offload"
-
-        # Restore grad buffers to GPU
-        ddp_model.restore_grad_buffers()
-
-        # Verify storage is restored
-        restored_sizes = [buf.grad_data.storage().size() for buf in all_buffers]
-        assert (
-            initial_sizes == restored_sizes
-        ), f"Expected restored sizes {restored_sizes} to match initial {initial_sizes}"
-
-    @pytest.mark.parametrize(
-        "initialize_model_parallel",
-        [
-            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
-            for tp, pp in itertools.product([1, 2], [1, 2])
-            if tp * pp <= Utils.world_size
-        ],
-        indirect=["initialize_model_parallel"],
-    )
-    def test_optimizer_offload(self, initialize_model_parallel):
-        """Test that optimizer offload_to_cpu/restore_from_cpu correctly moves state to/from CPU."""
-        world_size, dp, tp, pp = initialize_model_parallel
-        self.create_test_args(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
-        model_parallel_cuda_manual_seed(123)
-
-        # Create a realistic GPTModel as used in RL training
-        transformer_config = TransformerConfig(
-            num_layers=2, hidden_size=64, num_attention_heads=4, use_cpu_initialization=True
-        )
-        gpt_model = GPTModel(
-            config=transformer_config,
-            transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
-            vocab_size=256,
-            max_sequence_length=32,
-        ).cuda()
-
-        ddp_config = DistributedDataParallelConfig(
-            grad_reduce_in_fp32=True,
-            use_distributed_optimizer=True,
-            overlap_grad_reduce=False,
-            bucket_size=None,  # Single bucket for simplicity
-        )
-
-        ddp_model = DistributedDataParallel(
-            transformer_config, ddp_config=ddp_config, module=gpt_model
-        )
-
-        # Create optimizer
-        optimizer_config = OptimizerConfig(
-            optimizer='adam', bf16=True, use_distributed_optimizer=True
-        )
-        optimizer = get_megatron_optimizer(optimizer_config, [ddp_model])
-
-        # Manually initialize optimizer state (simulating what happens after first step)
-        # This avoids needing to run a full forward/backward/step cycle
-        for opt in optimizer.chained_optimizers:
-            if hasattr(opt, 'optimizer') and opt.optimizer is not None:
-                for group in opt.optimizer.param_groups:
-                    for p in group['params']:
-                        if len(opt.optimizer.state[p]) == 0:
-                            # Initialize Adam state (exp_avg and exp_avg_sq) on GPU
-                            opt.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data)
-                            opt.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data)
-                            opt.optimizer.state[p]['step'] = torch.tensor(1)
-
-        # Helper to check if optimizer state tensors are on GPU or CPU
-        def get_optimizer_state_devices():
-            devices = set()
-            for opt in optimizer.chained_optimizers:
-                if hasattr(opt, 'optimizer') and opt.optimizer is not None:
-                    for state_dict in opt.optimizer.state.values():
-                        for v in state_dict.values():
-                            if isinstance(v, torch.Tensor):
-                                devices.add(str(v.device))
-            return devices
-
-        # Verify optimizer state is initially on GPU
-        initial_devices = get_optimizer_state_devices()
-        assert any(
-            'cuda' in d for d in initial_devices
-        ), f"Expected optimizer state on GPU initially, got devices: {initial_devices}"
-
-        # Record GPU memory before offload
-        torch.cuda.synchronize()
-        memory_before_offload = torch.cuda.memory_allocated()
-
-        # Offload optimizer state to CPU
-        optimizer.offload_to_cpu()
-
-        # Verify GPU memory decreased (optimizer state should be freed)
-        torch.cuda.synchronize()
-        memory_after_offload = torch.cuda.memory_allocated()
-        assert memory_after_offload < memory_before_offload, (
-            f"Expected GPU memory to decrease after offload. "
-            f"Before: {memory_before_offload}, After: {memory_after_offload}"
-        )
-
-        # Verify optimizer state is now on CPU
-        offloaded_devices = get_optimizer_state_devices()
-        assert all(
-            'cpu' in d for d in offloaded_devices
-        ), f"Expected all optimizer state on CPU after offload, got devices: {offloaded_devices}"
-
-        # Restore optimizer state to GPU
-        optimizer.restore_from_cpu()
-
-        # Verify optimizer state is back on GPU
-        restored_devices = get_optimizer_state_devices()
-        assert any(
-            'cuda' in d for d in restored_devices
-        ), f"Expected optimizer state on GPU after restore, got devices: {restored_devices}"
-
-        # Verify GPU memory increased after restore (optimizer state reallocated)
-        torch.cuda.synchronize()
-        memory_after_restore = torch.cuda.memory_allocated()
-        assert memory_after_restore > memory_after_offload, (
-            f"Expected GPU memory to increase after restore. "
-            f"After offload: {memory_after_offload}, After restore: {memory_after_restore}"
-        )

From 4913c46c5af2ce584b12855b8f121b0370e212a1 Mon Sep 17 00:00:00 2001
From: Asha Anoosheh <aanoosheh@nvidia.com>
Date: Thu, 29 Jan 2026 15:29:13 +0100
Subject: [PATCH 66/79] Revise and move KD docs (#3108)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
---
 examples/post_training/modelopt/README.md     |  1 +
 .../post_training/modelopt}/distillation.md   | 38 ++++++++++---------
 megatron/post_training/arguments.py           | 20 +++++-----
 megatron/post_training/model_builder.py       | 20 ++++++----
 4 files changed, 43 insertions(+), 36 deletions(-)
 rename {megatron/post_training/docs => examples/post_training/modelopt}/distillation.md (76%)

diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md
index 6ebdd8ac5d6..93b5022b2aa 100644
--- a/examples/post_training/modelopt/README.md
+++ b/examples/post_training/modelopt/README.md
@@ -8,6 +8,7 @@
 [Configuration](./ADVANCED.md#advanced-configuration) |
 [Slurm Examples](./ADVANCED.md#slurm-examples) |
 [Speculative Decoding](./speculative.md) |
+[Knowledge Distillation](./distillation.md) |
 [Advanced Topics](./ADVANCED.md)
 
 </div>
diff --git a/megatron/post_training/docs/distillation.md b/examples/post_training/modelopt/distillation.md
similarity index 76%
rename from megatron/post_training/docs/distillation.md
rename to examples/post_training/modelopt/distillation.md
index f071cb0efba..49f73c4edde 100644
--- a/megatron/post_training/docs/distillation.md
+++ b/examples/post_training/modelopt/distillation.md
@@ -1,9 +1,5 @@
 # Megatron-LM ModelOpt Distillation Integration
 
-## Table of Contents
-
-[[_TOC_]]
-
 ## How To
 
 ### Prerequisites
@@ -16,22 +12,22 @@ We require the following pieces of data:
 * Teacher model weights
 * Student model weights (unless starting from scratch)
 * NeMo-format config file for teacher model
-* Distillation run config file
 * Tokenizer
 * Dataset
 
-It also requires the installation of the [NVIDIA Model Optimizer library](https://github.com/NVIDIA/Model-Optimizer)
+And optionally:
+* Distillation run config file
 
 ### Teacher checkpoint format
 
 We enforce the use of a config yaml in [NeMo](https://github.com/NVIDIA/NeMo) checkpoint-format style to define the arguments to the teacher model.
 The normal command-line arguments go toward constructing the student, thus the values in this file
-override the student arguments before being handed to the teacher constructor. This file must be
-named `model_config.yaml` and be placed in the root of the teacher model checkpoint folder.
-Unlike NeMo-generated checkpoints, Megatron-LM checkpoints do not contain these files by default and must be manually created.
+override the student arguments before being handed to the teacher constructor. This file must be either passed in via
+`--export-kd-teacher-model-config` or be named `model_config.yaml` in the root of the teacher model checkpoint folder.
+Unlike NeMo-generated checkpoints, Megatron-LM checkpoints do not contain this file by default and must be manually created.
 
-> NOTE: Not all keys in the NEMO-style yaml correspond 1:1 to the argument names for Megatron-LM. These
-are converted in `megatron/inference/gpt/model_provider.py`.
+> NOTE: Not all keys in the NeMo-style yaml correspond 1:1 to the argument names for Megatron-LM. These
+are converted in `megatron/post_training/model_builder.py`.
 
 ### Distillation config format
 
@@ -44,23 +40,33 @@ intermediate_layer_pairs:
   - ["decoder.final_layernorm", "decoder.layers.30.input_layernorm"]
 skip_lm_loss: true
 kd_loss_scale: 10.0
+logit_kl_temperature: 1.0
 ```
 
 * `logit_layers` defines the names of the student and teacher submodules, respectively, whose outputs are the logits.
 * `intermediate_layer_pairs` defines the potentially multiple – or zero – pairs of intermediate activation layers to also perform loss on.
-* `skip_lm_loss` decides whether or not to compute and combine the original training LM loss with the KD loss
+* `skip_lm_loss` decides whether or not to compute and combine the original training LM loss with the KD loss.
 * `kd_loss_scale` will scale the KD loss before adding it to the LM loss, if `skip_lm_loss` is `False`.
+* `logit_kl_temperature` is the temperature smoothing factor to multiply the logits by prior to softmax and loss.
+
+Without this configuration file, the default logits-only distillation with scale and temperatures of 1.0 will be performed.
 
 ### Training
 
-Distillation is triggered by calling `pretrain_gpt.py` with the additional following arguments:
+Distillation is triggered by calling `pretrain_gpt.py` or `pretrain_mamba.py` with the following arguments:
 
 ```bash
---kd-teacher-load <path-to-teacher-checkpoint>
---kd-distill-cfg <path-to-distill-config-yaml-file>
+--export-kd-teacher-load <path-to-teacher-checkpoint>
 --export-te-mcore-model
 ```
 
+optionally alongside the additional following arguments:
+
+```bash
+--export-kd-distill-cfg <path-to-distill-config-yaml-file>
+--export-kd-teacher-model-config <path-to-teacher-model-config-file>
+```
+
 > NOTE: If the teacher checkpoint happens to be in a different format from the student's (whose format is specified via `--ckpt-format`), it can
 be distinguished separately using the additional flag `--export-kd-teacher-ckpt-format`.
 
@@ -81,8 +87,6 @@ both defined in `modelopt.torch.distill.plugins.megatron`.
 
 * Interleaved Pipeline Parallel is unsupported for Distillation.
 
-* Only Megatron-Core models (not legacy Megatron) are supported for Distillation.
-
 ## Known Issues
 
 * An unknown memory allocation (a few megabytes per microbatch) takes place when the model is converted to a
diff --git a/megatron/post_training/arguments.py b/megatron/post_training/arguments.py
index 845fe9f17c3..1e988680142 100644
--- a/megatron/post_training/arguments.py
+++ b/megatron/post_training/arguments.py
@@ -51,24 +51,16 @@ def add_modelopt_args(parser):
     )
     # Knowledge Distillation
     group.add_argument(
-        '--export-kd-cfg',
+        '--export-kd-teacher-load',
         type=str,
-        default=None,
-        help='Path to distillation configuration yaml file.',
+        help='Path to checkpoint to load as distillation teacher. (Enables distillation mode automatically)',
     )
-
     group.add_argument(
-        '--teacher-model-config',
+        '--export-kd-teacher-model-config',
         type=str,
         default=None,
         help='Path to teacher model config for distillation. If not provided, defaults to ${export_kd_teacher_load}/model_config.yaml.',
     )
-
-    group.add_argument(
-        '--export-kd-teacher-load',
-        type=str,
-        help='Path to checkpoint to load as distillation teacher.',
-    )
     group.add_argument(
         '--export-kd-teacher-ckpt-format',
         type=str,
@@ -76,6 +68,12 @@ def add_modelopt_args(parser):
         choices=['torch', 'torch_dist', 'torch_dcp'],
         help="Checkpoint format of teacher model, if different from student's.",
     )
+    group.add_argument(
+        '--export-kd-cfg',
+        type=str,
+        default=None,
+        help='Path to distillation configuration yaml file, in order to use non-default settings.',
+    )
 
     # Finetuning
     group.add_argument(
diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py
index 71111ced069..fea837c96c3 100644
--- a/megatron/post_training/model_builder.py
+++ b/megatron/post_training/model_builder.py
@@ -47,9 +47,9 @@ def _add_load_convert_hooks(model: MCoreGPTModel):
 def _load_teacher_model_config(checkpoint_path: str) -> Namespace:
     """Reads teacher config from a file.
 
-    The config provided via --teacher-model-config should specify
-    (in NEMO format) any model architecture settings which differ from the main student model's.
-    This function will translate NEMO field names to MCore as needed.
+    The config provided, either in the teacher checkpoint dir or via `--export-kd-teacher-model-config`,
+    should specify (in NeMo yaml config format) any model architecture settings which differ from the main student model's.
+    This function will translate NeMo field names to MCore as needed.
     """
     required_teacher_fields = (
         "num_layers",
@@ -59,18 +59,22 @@ def _load_teacher_model_config(checkpoint_path: str) -> Namespace:
     )
 
     args = get_args()
-    config_path = os.path.join(checkpoint_path, "model_config.yaml") if args.teacher_model_config is None else args.teacher_model_config
+    if args.export_kd_teacher_model_config is not None:
+        config_path = args.export_kd_teacher_model_config
+    else:
+        config_path = os.path.join(checkpoint_path, "model_config.yaml")
     if not os.path.exists(config_path):
         raise FileNotFoundError(
-            "Teacher checkpoint dir must contain a NEMO-format yaml config named 'model_config.yaml'"
+            f"Teacher model-config file {config_path} not found.\n"
+            "Teacher checkpoint dir must contain a NeMo-format config named 'model_config.yaml'"
+            " or provide it via --export-kd-teacher-model-config."
         )
     with open(config_path) as f:
         config = yaml.safe_load(f)
 
-    missing_keys = [k for k in required_teacher_fields if k not in config]
-    if missing_keys:
+    if missing_keys := [k for k in required_teacher_fields if k not in config]:
         raise ValueError(
-            f"Teacher `model_config.yaml` file missing the following fields: {missing_keys}"
+            f"Teacher model config file ({config_path}) missing the following required fields: {missing_keys}"
         )
 
     if "encoder_seq_length" in config:

From 558fdafab3a86147efff3de99e15720724037a06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Thu, 29 Jan 2026 16:02:10 +0100
Subject: [PATCH 67/79] build: Bump FLA (#3139)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 pyproject.toml |  2 +-
 uv.lock        | 22 ++++++++++------------
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1e3ed2c76be..567954ca4a1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -78,7 +78,7 @@ dev = [
     "opentelemetry-api~=1.33.1",
     "mamba-ssm~=2.2",
     "causal-conv1d~=1.5",
-    "flash-linear-attention~=0.3.2",
+    "flash-linear-attention~=0.4.0",
     "nv-grouped-gemm~=1.1",
     "megatron-energon[av_decode]~=6.0",
     "av",
diff --git a/uv.lock b/uv.lock
index bfb4a0ee734..1867e8aaddf 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1333,15 +1333,15 @@ wheels = [
 
 [[package]]
 name = "fla-core"
-version = "0.3.2"
+version = "0.4.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "einops" },
     { name = "torch", marker = "sys_platform == 'never'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/67/c6/10a1149b07e6bab45b2cb2d07f6b827716c2baf5f3404161753f25c6389b/fla_core-0.3.2.tar.gz", hash = "sha256:d38db16bc4e1c6fa8c04df442f246da1e6926a209426bc6ef703d41bfbc37c92", size = 296725, upload-time = "2025-09-10T07:43:40.155Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/f1/de/0d6bd5664ba2e711cabdde11ccb41ddcdd866c531e40900af3601bd7b8c6/fla_core-0.4.1.tar.gz", hash = "sha256:38ab28966eeadc2141b29e87c2bf72a8a4851e00af9d25bbbc3596b1fb53450d", size = 319608, upload-time = "2025-12-24T18:07:37.669Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/7e/f5/74947b33c07682280e65adbdf17c4ee94b30232df2f728bafecf13d1d820/fla_core-0.3.2-py3-none-any.whl", hash = "sha256:e751d5a41e33eee721a6fb6588bd857f6f36e0d14719a23b1ebdbd617d307209", size = 413594, upload-time = "2025-09-10T07:43:37.786Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/43/945ef69eb48a14c30fd7323d3e0b560c821ae71e6d3ef979e06a901bc3b9/fla_core-0.4.1-py3-none-any.whl", hash = "sha256:93c6afe4c80fc7bc705fa8aeea6a46d2cf2d77383f9619a41863c7114c801bab", size = 437282, upload-time = "2025-12-24T18:07:34.41Z" },
 ]
 
 [[package]]
@@ -1360,17 +1360,15 @@ wheels = [
 
 [[package]]
 name = "flash-linear-attention"
-version = "0.3.2"
+version = "0.4.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "datasets" },
     { name = "fla-core" },
-    { name = "pytest" },
     { name = "transformers" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/84/f6/e62c1e562a288557eba7f06f168a7615813d1a227327b8beb8ba426da2c5/flash_linear_attention-0.3.2.tar.gz", hash = "sha256:9147747316c2951fed4ebeb4fa87977c05d807dc70c93b46250b68a6eb1183e2", size = 150880, upload-time = "2025-09-10T07:43:41.37Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/46/83/7d8ec7ffb5229080b1c9b772338ff588cbd63282ac355ede2a12a6e174a8/flash_linear_attention-0.4.1.tar.gz", hash = "sha256:127ee7273ed15ac17f72bcf4c75e1051719d8fbe0a2d1d047e59406f36d81ee2", size = 158280, upload-time = "2025-12-24T18:07:38.812Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/a0/d0/35ce9eac5f52c72005095aaa12a393d2656ed7ffedf925b2381a6b76d10c/flash_linear_attention-0.3.2-py3-none-any.whl", hash = "sha256:604e73361437ba786420ab195e2caa3fd19280503761e703fa353c5ce5c65376", size = 274592, upload-time = "2025-09-10T07:43:39.107Z" },
+    { url = "https://files.pythonhosted.org/packages/63/d5/6327559a9d5b9243b10c3984f1bcef256ed2ad06d105a3bb8f7b2979659c/flash_linear_attention-0.4.1-py3-none-any.whl", hash = "sha256:d18bdfe9d1f4b424676444eac9d50fb8433b70e5d4e0e0878b20bcbcdbea57ce", size = 287415, upload-time = "2025-12-24T18:07:35.815Z" },
 ]
 
 [[package]]
@@ -2347,7 +2345,7 @@ requires-dist = [
     { name = "emerging-optimizers", marker = "extra == 'lts'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.1.0" },
     { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" },
     { name = "fastapi", marker = "extra == 'lts'", specifier = "~=0.50" },
-    { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "~=0.3.2" },
+    { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "~=0.4.0" },
     { name = "flashinfer-python", marker = "extra == 'dev'", specifier = "~=0.5.0" },
     { name = "flashinfer-python", marker = "extra == 'lts'", specifier = "~=0.5.0" },
     { name = "flask-restful", marker = "extra == 'mlm'" },
@@ -4149,12 +4147,12 @@ name = "pytest"
 version = "8.3.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-    { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
+    { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" },
+    { name = "exceptiongroup", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" },
     { name = "iniconfig" },
     { name = "packaging" },
     { name = "pluggy" },
-    { name = "tomli", marker = "python_full_version < '3.11'" },
+    { name = "tomli", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/ae/3c/c9d525a414d506893f0cd8a8d0de7706446213181570cdbd766691164e40/pytest-8.3.5.tar.gz", hash = "sha256:f4efe70cc14e511565ac476b57c279e12a855b11f48f212af1080ef2263d3845", size = 1450891, upload-time = "2025-03-02T12:54:54.503Z" }
 wheels = [

From 409af92e88039ade6e9b790c40545e4759a4f65f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Thu, 29 Jan 2026 16:04:35 +0100
Subject: [PATCH 68/79] ci: Add job timeouts (#3142)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .github/workflows/cicd-main.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
index d69985bff40..ad26a5ba0f6 100644
--- a/.github/workflows/cicd-main.yml
+++ b/.github/workflows/cicd-main.yml
@@ -360,6 +360,7 @@ jobs:
       - cicd-container-build
       - cicd-parse-unit-tests
     runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner }}
+    timeout-minutes: 60
     name: "${{ matrix.bucket }} - latest"
     if: |
       (
@@ -389,6 +390,7 @@ jobs:
 
   cicd-parse-integration-tests:
     runs-on: ubuntu-latest
+    timeout-minutes: 60
     needs:
       - pre-flight
       - cicd-wait-in-queue

From f4af1bfcfa124bdfc1398986b300a115ebb53608 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Thu, 29 Jan 2026 16:07:12 +0100
Subject: [PATCH 69/79] ci: Set NODE_RANK (#3143)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 tests/functional_tests/shell_test_utils/run_ci_test.sh | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/functional_tests/shell_test_utils/run_ci_test.sh b/tests/functional_tests/shell_test_utils/run_ci_test.sh
index 3f1500b502c..3d47e591749 100644
--- a/tests/functional_tests/shell_test_utils/run_ci_test.sh
+++ b/tests/functional_tests/shell_test_utils/run_ci_test.sh
@@ -132,6 +132,8 @@ SKIP_PYTEST=$(cat $TRAINING_PARAMS_PATH |
 
 export RECORD_CHECKPOINTS=${RECORD_CHECKPOINTS:-"false"}
 
+NODE_RANK=${SLURM_NODEID:-${SLURM_NODEID:-0}}
+
 for i in $(seq 1 $N_REPEAT); do
     # Move TB logs into a repeat-specific directory
     DIR=$(dirname "$_TENSORBOARD_PATH")

From 0b619c2090b153ac599fbf4cc330421bc8864a84 Mon Sep 17 00:00:00 2001
From: yobi byte <yobibyte@users.noreply.github.com>
Date: Thu, 29 Jan 2026 15:22:55 +0000
Subject: [PATCH 70/79] Multiturn rollout support prep (#2966)

---
 megatron/rl/agent/api.py                      |  14 +-
 megatron/rl/agent/reward_only_agent.py        |   8 +-
 megatron/rl/rl_utils.py                       | 552 ++++++++++--------
 megatron/rl/sequence_packing_utils.py         | 128 ++--
 megatron/training/arguments.py                |   2 -
 megatron/training/training.py                 |  16 +-
 tests/unit_tests/rl/test_rl_utils.py          | 121 +++-
 .../rl/test_sequence_packing_utils.py         |  56 ++
 train_rl.py                                   |   1 +
 9 files changed, 544 insertions(+), 354 deletions(-)

diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py
index 34efa68d85a..9568db3a54d 100644
--- a/megatron/rl/agent/api.py
+++ b/megatron/rl/agent/api.py
@@ -46,8 +46,8 @@ class GroupedRolloutRequest(Request):
 class Rollout(AgentBaseModel):
     """Data for language-based Rollout."""
 
-    trajectory: str
-    prompt_length: int | None = None
+    trajectory: list[str]
+    prompt_length: list[int] | None = None
     reward: float = None
     env_id: str | None = None
     problem_id: str | None = None
@@ -56,10 +56,10 @@ class Rollout(AgentBaseModel):
 class TokenRollout(AgentBaseModel):
     """Tokenized representation of a language-based Rollout."""
 
-    trajectory: list[int]
+    trajectory: list[list[int]]
     reward: list[float] | float
-    generation_mask: list[list[int]] | list[bool] | None = None
-    logprobs: list[float] | None = None
+    generation_mask: list[list[bool]] | None = None
+    logprobs: list[list[float]] | None = None
     env_id: str | None = None
     problem_id: str | None = None
 
@@ -67,8 +67,8 @@ class TokenRollout(AgentBaseModel):
 class ContrastiveRollout(AgentBaseModel):
     """Contrastive/Preference data for language-based Rollout."""
 
-    chosen_trajectory: str
-    rejected_trajectory: str
+    chosen_trajectory: list[str]
+    rejected_trajectory: list[str]
 
 
 class Head2HeadRolloutRequest(Request):
diff --git a/megatron/rl/agent/reward_only_agent.py b/megatron/rl/agent/reward_only_agent.py
index 2e81674c74d..53b1f7407b2 100644
--- a/megatron/rl/agent/reward_only_agent.py
+++ b/megatron/rl/agent/reward_only_agent.py
@@ -104,16 +104,16 @@ async def rollout_from_response(
                 for x in range(len(response.token_ids))
             ]
             rollout = TokenRollout(
-                trajectory=response.token_ids,
+                trajectory=[response.token_ids],
                 reward=await self.get_reward(response_text, golden),
-                logprobs=logprobs,
-                generation_mask=generation_mask,
+                logprobs=[logprobs],
+                generation_mask=[generation_mask],
                 env_id=self.env_id,
                 problem_id=golden['problem_id'] if 'problem_id' in golden else None,
             )
         else:
             rollout = Rollout(
-                trajectory=raw_text,
+                trajectory=[raw_text],
                 reward=await self.get_reward(response_text, golden),
                 env_id=self.env_id,
                 problem_id=golden['problem_id'] if 'problem_id' in golden else None,
diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index 26b990236ba..f5ba6d35d8c 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -2,18 +2,18 @@
 
 import gc
 
+import copy
+from functools import partial
 # Keep this to make the env registered.
 import itertools
-import json
-import logging
 import math
+import logging
 import pickle
 from collections import Counter, defaultdict
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
-from difflib import SequenceMatcher
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Optional 
 
 import numpy as np
 import torch
@@ -26,6 +26,7 @@
 from megatron.core.datasets.megatron_tokenizer import MegatronLegacyTokenizer
 from megatron.core.full_cuda_graph import FullCudaGraphWrapper
 from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.num_microbatches_calculator import reconfigure_num_microbatches_calculator
 from megatron.core.optimizer import MegatronOptimizer
 from megatron.core.pipeline_parallel import get_forward_backward_func
 from megatron.core.pipeline_parallel.utils import is_pp_last_stage, get_pp_last_rank
@@ -45,10 +46,10 @@
     compute_packed_inference_logprobs_stats,
     pack_all_trajectories,
     load_packed_data_by_index,
-    update_sequence_packing_metrics,
     get_sequence_packing_tensorboard_metrics,
     get_sequence_packing_log_info,
     get_default_packed_seq_params,
+    update_microbatch_calculator,
 )
 from megatron.rl.agent.api import (
     EvaluationRequest,
@@ -66,7 +67,6 @@
 from megatron.training.global_vars import (
     get_args,
     get_tensorboard_writer,
-    get_timers,
     get_tokenizer,
     get_wandb_writer,
 )
@@ -210,13 +210,13 @@ def verify_model_weights_swap(
         if inf_was_training:
             inf_core.train()
 
-GroupedRollouts = list[list[TokenRollout | Rollout]]
+Rollouts = list[TokenRollout | Rollout]
+GroupedRollouts = list[Rollouts]
 
 
 @dataclass(slots=True)
 class RolloutStats:
     mean_reward: float
-    mean_sim: None | float
     mean_length: float
     mean_length_std: float
     max_length: float
@@ -234,6 +234,7 @@ class RolloutStats:
     min_inf_prob: None | float
     max_inf_prob: None | float
     mean_inf_prob: None | float
+    num_turns: list[int] # num_turns per traj
 
 
 # Runtime state container for RL-specific data that shouldn't be checkpointed
@@ -243,7 +244,6 @@ class RLRuntimeState:
     def __init__(self):
         self.packing_context = None
         self.last_collection_iteration = 0
-        self.global_batches_per_collection = 0
         self.sequences_this_iteration_on_rank = 0
         self.latest_batch_num_sequences = 0
 
@@ -609,20 +609,20 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa
 
     """
 
+    args = get_args()
     # Ensure packed_seq_params is always provided for CUDA graph signature consistency
     if packed_seq_params is None and sequence_packing:
         packed_seq_params = get_default_packed_seq_params(
             seq_length=tokens.shape[1],
+            max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin,
             device=tokens.device,
         )
 
     nvtx_range = get_nvtx_range()
 
     with nvtx_range("get-logprobs", time=False):
-
         with nvtx_range("forward-pass", time=False):
             # TODO(vitalyk): use fp16/bf16 as a function argument. Do not use args.
-            args = get_args()
 
             attention_mask_for_forward = None
 
@@ -654,19 +654,46 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa
             return logprobs
 
 
+def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[int]]) -> np.ndarray:
+    """Calculate GRPO advantages from rewards/num_turns.
+
+    For multiturn rollouts, the logic is a bit more involved.
+    # For training, we'll be turning each turn into a trajectory with the same reward
+    # within a trajectory, e.g. if [[a,b],[c,d,e]] trajectory has reward 1.0, we will
+    # get [a,b] with 1.0 and [c,d,e] with 1.0 when doing updates.
+    """
+
+    rewards = np.array(rewards)
+
+    num_turns = np.array(num_turns)
+    # Each outer dimension of num_turns is a group. Sum of those gives total num_turns per group.
+    # Let's use this to calculate advantage.
+    # mean/std should be repeated based on group lens
+    group_turns = num_turns.sum(axis=-1)
+    reward_means = rewards.mean(axis=1, keepdims=True).repeat(group_turns)
+    reward_stds = rewards.std(axis=1, keepdims=True).repeat(group_turns)
+
+    # rewards are originally [g, group_size]
+    # Making an assumption that all groups are of the same size!
+    # @vitalyk: this will go away when we start sending env-based sample reqs.
+    rewards = rewards.flatten().repeat(num_turns.flatten())
+
+    return ((rewards - reward_means) / (1e-4 + reward_stds)).tolist()
+
+
 def compute_group_stats(
-    rollouts: GroupedRollouts, tokenizer: MegatronLegacyTokenizer
+    rollouts: GroupedRollouts, tokenizer: MegatronLegacyTokenizer, seq_len: int,
 ) -> RolloutStats:
     """Add group-based rollout stats for logging.
 
     Args:
         rollouts: Rollouts to generate the stats for. Each inner list is a group (as in GRPO group), i.e. all rollouts are for the same prompt.
         tokenizer: Tokenizer to tokenize the rollouts in case they are raw strings.
+        seq_len: Maximum sequence length.
 
     Returns:
        RolloutStats object containing all the stats.
     """
-    args = get_args()
     # TODO (rkirby) Maybe do some of this after the tensor building
     group_reward_means = []
     group_reward_stds = []
@@ -674,54 +701,45 @@ def compute_group_stats(
     group_length_stds = []
     group_length_maxs = []
     group_length_mins = []
-    group_rollout_similarities = []
+    rewards = []
+    num_turns = [] # num_turns per traj
     for group in rollouts:
         group_rewards = []
         group_lengths = []
+        group_num_turns = []
         for rollout in group:
+            group_num_turns.append(len(rollout.trajectory))
             if isinstance(rollout, TokenRollout):
-                lang_rl_log(
-                    f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} tokens] {tokenizer.detokenize(rollout.trajectory)}"
-                )
-                assert (len(rollout.trajectory) == args.seq_length) or (
-                    rollout.trajectory[-1] == tokenizer.eod
-                ), f"Rollout is not the correct length: {len(rollout.trajectory)} {rollout.trajectory[-1]}\n{tokenizer.detokenize(rollout.trajectory)}"
+                for turn_traj in rollout.trajectory:
+                    detokenized_traj = tokenizer.detokenize(turn_traj)
+                    lang_rl_log(
+                        f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} tokens] {detokenized_traj}"
+                    )
+                    # TODO(vitalyk): how does multiturn change EOD/EOT?
+                    assert (len(turn_traj) == seq_len) or (
+                        turn_traj[-1] == tokenizer.eod
+                    ), f"Rollout is not the correct length: {len(turn_traj)} {turn_traj[-1]}\n{detokenized_traj}"
             else:
                 lang_rl_log(
                     f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} chars] {rollout.trajectory}"
                 )
             group_rewards.append(rollout.reward)
-            group_lengths.append(len(rollout.trajectory))
-        if args.rl_calculate_intra_group_similarity:
-            # We can probably compute this outside, but in case we switch to different group sizes for different envs, let's keep it here.
-            combos = itertools.combinations(range(len(group)), 2)
-            # For every pair (excluding ourselves), check the sequence similarity and log.
-            # Use this to track the diversity of generated rollouts within a group.
-            intra_group_sim = np.mean(
-                list(
-                    map(
-                        lambda idx_pair: SequenceMatcher(
-                            None, group[idx_pair[0]].trajectory, group[idx_pair[1]].trajectory
-                        ).ratio(),
-                        combos,
-                    )
-                )
-            )
-            group_rollout_similarities.append(intra_group_sim)
-        else:
-            group_rollout_similarities = None
+            #TODO(vitalyk): What is the semantics behind traj length in multiturn? Should we take the last only? Average them instead of extending?
+            group_lengths.extend(len(t) for t in rollout.trajectory)
 
         group_length_maxs.append(max(group_lengths))
         group_length_mins.append(min(group_lengths))
         group_reward_means.append(np.mean(group_rewards))
         group_reward_stds.append(np.std(group_rewards))
+        rewards.append(group_rewards)
         group_length_means.append(np.mean(group_lengths))
         # https://arxiv.org/abs/2504.21233 reports that lens variants hurts.
         # Let's track this.
         group_length_stds.append(np.std(group_lengths))
+        num_turns.append(group_num_turns)
+
     stats = RolloutStats(
         mean_reward=np.mean(group_reward_means),
-        mean_sim=np.mean(group_rollout_similarities) if group_rollout_similarities else None,
         mean_length=np.mean(group_length_means),
         mean_length_std=np.mean(group_length_stds),
         max_length=np.max(group_length_maxs),
@@ -737,8 +755,9 @@ def compute_group_stats(
         min_inf_prob=None,
         max_inf_prob=None,
         mean_inf_prob=None,
-        rewards=None,  # We will fill those in later in prepare_data_for_update.
-        advantages=None,  # We will fill those in later in prepare_data_for_update.
+        rewards=[r for group in rewards for r in group],
+        advantages=calculate_grpo_advantages(rewards, num_turns),
+        num_turns=[nt for group in num_turns for nt in group],
     )
     return stats
 
@@ -809,11 +828,10 @@ def maybe_log_training_metrics(
                         columns=['Trajectories', 'Tokens', 'Rewards'],
                         rows=[
                             [
-                                (
-                                    tokenizer.detokenize(r.trajectory)
+                                [(tokenizer.detokenize(turn)
                                     if isinstance(r, TokenRollout)
-                                    else r.trajectory
-                                ),
+                                    else turn) for turn in r.trajectory
+                                ],
                                 r.trajectory,
                                 r.reward,
                             ]
@@ -821,11 +839,6 @@ def maybe_log_training_metrics(
                         ],
                     ),
                 },
-                **(
-                    {'mean_intra_group_similarity': group_stats.mean_sim}
-                    if group_stats.mean_sim
-                    else {}
-                ),
             },
             step=current_iteration,
         )
@@ -834,10 +847,9 @@ def maybe_log_training_metrics(
 
 
 def prepare_trajectories(
-    rollouts: GroupedRollouts, tokenizer: MegatronLegacyTokenizer, seq_length: int
+    rollouts: Rollouts, tokenizer: MegatronLegacyTokenizer, seq_length: int, sequence_packing: bool, skip_bos_token: bool
 ):
     """Pad trajectories and extract the generation masks.
-
     Args:
         rollouts: Rollouts to extract trajectories from.
         tokenizer: Tokenizer to get the padding token and potentially tokenize.
@@ -854,6 +866,7 @@ def prepare_trajectories(
 
     DEFAULT_PAD_TOKENS = ['<|finetune_right_pad_id|>']
 
+
     if isinstance(tokenizer, _HuggingFaceTokenizer):
         if not tokenizer.pad:
             for pad_token in DEFAULT_PAD_TOKENS:
@@ -884,17 +897,19 @@ def prepare_trajectories(
     trajs = []
     generation_masks = []
     inference_logprobs = []
-    for group in rollouts:
-        for rollout in group:
-            generation_mask = rollout.generation_mask if isinstance(rollout, TokenRollout) else None
-
-            trajectory = (
-                rollout.trajectory.copy()
-                if isinstance(rollout, TokenRollout)
-                else tokenizer.tokenize(rollout.trajectory)
-            )
-            inf_logprobs = rollout.logprobs
-
+    for rollout in rollouts:
+        # traj, gen mask and logprobs are lists now.
+        # each list entry is a turn, single-turn environments just have a single-element list.
+        # We assume that all lengths of the structs above have the same lengths (number of turns).
+
+        all_turns_trajectories = (
+            copy.deepcopy(rollout.trajectory)
+            if isinstance(rollout, TokenRollout)
+            else tokenizer.tokenize(rollout.trajectory)
+        )
+        for turn_idx, trajectory in enumerate(all_turns_trajectories):
+            inf_logprobs = rollout.logprobs[turn_idx]
+            generation_mask = rollout.generation_mask[turn_idx] if isinstance(rollout, TokenRollout) else None
             length = len(trajectory)
             assert length <= seq_length, "Rollout too long, how did this happen?"
             if len(trajectory) < seq_length:
@@ -916,8 +931,7 @@ def prepare_trajectories(
             else:
                 inference_logprobs.append(None)
 
-            env_id = rollout.env_id
-            env_id_counts[env_id] += 1
+        env_id_counts[rollout.env_id] += 1
 
     if torch.distributed.is_initialized():
         logger.info(f"[{dist.get_rank()}] Rollout counts:")
@@ -927,34 +941,18 @@ def prepare_trajectories(
     generation_masks = torch.tensor(generation_masks, dtype=torch.bool, device='cpu')
     trajs = torch.tensor(trajs, device='cpu')
 
-    args = get_args()
     # Only process if we have inference_logprobs
     if inference_logprobs and any(lp is not None for lp in inference_logprobs):
-        if args.rl_use_sequence_packing:
-            # For sequence packing, we need to pad all logprobs to the same size
-            padded_logprobs = []
-            for logprobs in inference_logprobs:
-                if logprobs is not None:
-                    if len(logprobs) < seq_length:
-                        # Pad with zeros (these positions will be masked anyway)
-                        padding_size = seq_length - len(logprobs)
-                        padded = torch.nn.functional.pad(logprobs, (0, padding_size), value=0.0)
-                        padded_logprobs.append(padded)
-                    else:
-                        padded_logprobs.append(logprobs)
-                else:
-                    # Create zero tensor for None logprobs
-                    padded_logprobs.append(torch.zeros(seq_length))
-            inference_logprobs = torch.stack(padded_logprobs)
-        else:
-            # For non-packing mode, keep as list of tensors (unpadded)
-            # This preserves the original behavior where each sequence can have different lengths
-            pass
+        # We need to pad all logprobs to the same size for sequence packing.
+        # For non-packing mode, keep as list of tensors (unpadded)
+        # This preserves the original behavior where each sequence can have different lengths
+        if sequence_packing:
+            inference_logprobs = _pad_nonnull_with_zeros(inference_logprobs, seq_length)
     else:
         inference_logprobs = None
 
     # Some sanity checks regarding the tokenization
-    if not args.rl_skip_bos_token:
+    if not skip_bos_token:
         assert (
             tokenizer.bos is None or (trajs[:, 0] == tokenizer.bos).all()
         ), "First token should be bos"
@@ -975,11 +973,92 @@ def prepare_trajectories(
     return trajs, generation_masks, inference_logprobs
 
 
+def logprobs_forward_step(data_iterator, model, is_correction, packing_context=None):
+    # Avoid self.training checks which will trigger cudagraph capture; this path reuses
+    # the forward pass from training after it has been captured on the 1st iteration.
+    model.eval()
+
+    if packing_context is not None:
+        # When using sequence packing, the data iterator returns a tuple with a single element, the bin index.
+        bin_tensor = next(data_iterator)[0]
+        #TODO(jalbericiola): change for named tuple
+        (b_trajs, _, _, _, b_posids, _, _, _, _, _, b_packed_seq_params) = (
+            load_packed_data_by_index(bin_tensor.item(), packing_context, is_correction)
+        )
+    else:
+        b_trajs, b_posids = next(data_iterator)
+        b_packed_seq_params = None
+
+    logprobs = (
+        get_logprobs(
+            model,
+            b_trajs.cuda(),
+            b_posids.cuda(),
+            no_grad=True,
+            sequence_packing=b_packed_seq_params is not None, 
+            packed_seq_params=b_packed_seq_params,
+        ),
+        None,
+    )
+    model.train()
+    return logprobs
+
+
+def _compute_logprobs_batch(
+    model,
+    data_loader,
+    forward_backward_func,
+    packing_context,
+    trajs_batch_size, # n_bins for seq packing, and batch_size for non seq packing
+    seq_length,
+    logprobs_batch_size,
+    decoder_seq_length,
+    dtype,
+    pp_group,
+    is_correction,
+):
+    """Compute logprobs for all batches in the data loader."""
+    logprobs_list = []
+    data_iterator = iter(data_loader)
+    for i in range(len(data_loader)):
+        output_tensor = forward_backward_func(
+            forward_step_func=partial(logprobs_forward_step, is_correction=is_correction, packing_context=packing_context),
+            data_iterator=data_iterator,
+            model=model,
+            num_microbatches=1,
+            seq_length=seq_length,
+            micro_batch_size=logprobs_batch_size,
+            decoder_seq_length=decoder_seq_length,
+            forward_only=True,
+            adjust_tensor_shapes_fn=None,
+        )
+        if is_pp_last_stage(pp_group):
+            logprobs_list.append(output_tensor[0].detach())
+
+    if is_pp_last_stage(pp_group):
+        logprobs = torch.concat(logprobs_list, dim=0)
+        assert logprobs.dtype == dtype
+    else:
+        logprobs = torch.empty(
+            trajs_batch_size,
+            seq_length-1,
+            dtype=dtype,
+            device=torch.cuda.current_device(),
+        )
+
+    # Only PP>1 needs a broadcast from the last stage; for PP=1 the output is already local.
+    if get_pg_size(pp_group) > 1:
+        dist.broadcast(logprobs, src=get_pp_last_rank(pp_group), group=pp_group)
+    return logprobs.cpu()
+
+
 def prepare_data_for_update(
     model: list[LanguageModule],
     ref_state_dict: Dict[str, Any],
     rollouts: GroupedRollouts,
     tokenizer: MegatronLegacyTokenizer,
+    sequence_packing: bool,
+    is_correction: bool,
 ) -> RerunDataIterator:
     """Extract data for the update from raw rollouts.
 
@@ -988,6 +1067,8 @@ def prepare_data_for_update(
         ref_state_dict: Reference policy state dict.
         rollouts: Rollouts to extract the data from.
         tokenizer: Tokenizer to pad/tokenize data.
+        sequence_packing: Use sequence packing if True.
+        is_correction: Prepare data for IS correction if True.
 
     Returns:
         Cycled iterator over dataset batches. In GRPO we might want to go over the same data multiple times.
@@ -1009,59 +1090,50 @@ def prepare_data_for_update(
 
     with nvtx_range("prepare-data-for-update"):
         with nvtx_range("compute-group-stats"):
-            # These are computed on all rollouts for reporting purposes
-            group_stats = compute_group_stats(rollouts, tokenizer)
-            rewards = np.array([[rollout.reward for rollout in group] for group in rollouts])
-            group_stats.rewards = rewards.flatten().tolist()
-            group_stats.advantages = (
-                (
-                    (rewards - rewards.mean(axis=1, keepdims=True))
-                    / (1e-4 + rewards.std(axis=1, keepdims=True))
-                )
-                .flatten()
-                .tolist()
-            )
-            global_rollout_count = len(group_stats.rewards)
-
-        with nvtx_range("prepare_advantages", time=True):        
-            # [g, group_size]
-            # Making an assumption that all groups are of the same size!
-            rewards = torch.tensor(rewards, device='cpu')
-            advantages = (rewards - rewards.mean(axis=1, keepdim=True)) / (
-                1e-4 + rewards.std(axis=1, keepdim=True)
-            )
-
-            # Flatten advantages for training and move to GPU
-            advantages = global_advantages = advantages.view(-1).cuda()
+            group_stats = compute_group_stats(rollouts, tokenizer, args.seq_length)
+            # TODO(vitalyk): why do we need global_advantages here? go inside packing
+            advantages = global_advantages = torch.tensor(group_stats.advantages, dtype=dtype).cuda()
 
         # Now split the rollouts across the data parallel ranks for training
         # This needs to be done at this point because we are about to calculate logprobs
         # Note :- For EP, do not use the expert data parallel group here. Always 
         # use the regular data parallel group. 
+
+        # Use one group as an exampling for logging later.
+        example_group = rollouts[0]
+
+        # Let's expand rollouts getting rid of the groups.
+        # We need this to correctly split the rollouts across dp groups.
+        # And we do not actually need them grouped in anything below anyways.
+        rollouts = [r for g in rollouts for r in g]
+        total_turns_sampled = len(rollouts)
+
+        # We might sample more than we consume in one step.
+        samples_ratio_per_step = args.global_batch_size / (args.grpo_prompts_per_step * args.grpo_group_size)
+        assert samples_ratio_per_step <= 1, "You cannot use more data than you sampled."
+
         if (data_parallel_world_size := mpu.get_data_parallel_world_size()) > 0:
             data_split_size = len(rollouts) // data_parallel_world_size
             data_split_range = (
                 mpu.get_data_parallel_rank() * data_split_size,
                 (mpu.get_data_parallel_rank() + 1) * data_split_size,
             )
+            # TODO(vitalyk): This has to be rewritten assuming we are multiturn now.
             rollouts = rollouts[data_split_range[0] : data_split_range[1]]
+            local_num_turns = sum(group_stats.num_turns[data_split_range[0] : data_split_range[1]])
+            steps_before = sum(group_stats.num_turns[:data_split_range[0]])
+            advantages = advantages[steps_before:steps_before+local_num_turns]
             # First we calculate them on a global level and then we split and recalculate on a local level.
             # Sequence packing and reporting needs it global but non-packing wants it local.
-            rewards = torch.tensor([[r.reward for r in group] for group in rollouts], device='cpu')
-            advantages = (rewards - rewards.mean(axis=1, keepdim=True)) / (
-                1e-4 + rewards.std(axis=1, keepdim=True)
-            )
-
-            # Flatten advantages for training and move to GPU
-            advantages = advantages.view(-1).cuda()
 
         with nvtx_range("prepare_trajectories"):
             trajs, generation_masks, inference_logprobs = prepare_trajectories(
-                rollouts, tokenizer, args.seq_length
+                rollouts, tokenizer, args.seq_length, sequence_packing, args.rl_skip_bos_token
             )
 
+        packing_context = None
         # Build trajectories based on sequence packing or standard processing
-        if args.rl_use_sequence_packing:
+        if sequence_packing:
             with nvtx_range("sequence_packing", time=True):
                 runtime_state.packing_context = packing_context = pack_all_trajectories(
                     trajs, 
@@ -1101,7 +1173,6 @@ def prepare_data_for_update(
                 )
                 logprobs_batch_size = args.micro_batch_size
 
-
         with torch.no_grad(), nvtx_range("compute_logprobs", time=True):
             # Before we can update the model, we need to get the logprobs for the \pi_{old} model.
 
@@ -1112,41 +1183,6 @@ def prepare_data_for_update(
                     forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps
                 )
 
-            def logprobs_forward_step(data_iterator, model):
-
-                # Avoid self.training checks which will trigger cudagraph capture; this path reuses
-                # the forward pass from training after it has been captured on the 1st iteration.
-                model.eval()
-
-                if args.rl_use_sequence_packing:
-                    # When using sequence packing, the data iterator returns a tuple with a single element, the bin index.
-                    bin_tensor = next(data_iterator)[0]
-                    #TODO(jalbericiola): change for named tuple
-                    (b_trajs, _, _, _, b_posids, _, _, _, _, _, b_packed_seq_params) = (
-                        load_packed_data_by_index(bin_tensor.item(), packing_context, args.rl_inference_logprobs_is_correction)
-                    )
-                else:
-                    batch_data = next(data_iterator)
-                    b_trajs, b_posids = batch_data
-                    b_packed_seq_params = None
-
-                b_trajs = b_trajs.cuda()
-                b_posids = b_posids.cuda()
-
-                logprobs = (
-                    get_logprobs(
-                        model,
-                        b_trajs,
-                        b_posids,
-                        no_grad=True,
-                        sequence_packing=args.rl_use_sequence_packing,                       
-                        packed_seq_params=b_packed_seq_params,
-                    ),
-                    None,
-                )
-
-                model.train()
-                return logprobs
 
             dtype = (
                 torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32)
@@ -1155,43 +1191,20 @@ def logprobs_forward_step(data_iterator, model):
             pg_collection = get_attr_wrapped_model(model, "pg_collection")
             pp_group = pg_collection.pp
 
-            def _compute_logprobs_batch():
-                """Compute logprobs for all batches in the data loader."""
-                logprobs_list = []
-                data_iterator = iter(data_loader)
-                for i in range(len(data_loader)):
-                    output_tensor = forward_backward_func(
-                        forward_step_func=logprobs_forward_step,
-                        data_iterator=data_iterator,
-                        model=model,
-                        num_microbatches=1,
-                        seq_length=args.seq_length,
-                        micro_batch_size=logprobs_batch_size,
-                        decoder_seq_length=args.decoder_seq_length,
-                        forward_only=True,
-                        adjust_tensor_shapes_fn=None,
-                    )
-                    if is_pp_last_stage(pp_group):
-                        logprobs_list.append(output_tensor[0].detach())
-
-                if is_pp_last_stage(pp_group):
-                    logprobs = torch.concat(logprobs_list, dim=0)
-                    assert logprobs.dtype == dtype
-                else:
-                    logprobs = torch.empty(
-                        len(compute_trajs),
-                        args.seq_length - 1,
-                        dtype=dtype,
-                        device=torch.cuda.current_device(),
-                    )
-
-                # Only PP>1 needs a broadcast from the last stage; for PP=1 the output is already local.
-                if get_pg_size(pp_group) > 1:
-                    dist.broadcast(logprobs, src=get_pp_last_rank(pp_group), group=pp_group)
-                return logprobs.cpu()
-
             with torch.no_grad(), nvtx_range("compute_old_logprobs", time=True):
-                old_logprobs = _compute_logprobs_batch()
+                old_logprobs = _compute_logprobs_batch(
+                    model=model,
+                    data_loader=data_loader,
+                    forward_backward_func=forward_backward_func,
+                    packing_context=packing_context,
+                    trajs_batch_size=len(compute_trajs),
+                    seq_length=args.seq_length,
+                    logprobs_batch_size=logprobs_batch_size,
+                    decoder_seq_length=args.decoder_seq_length,
+                    dtype=dtype,
+                    pp_group=pp_group,
+                    is_correction=args.rl_inference_logprobs_is_correction,
+                )
 
             with torch.no_grad(), nvtx_range("compute_ref_logprobs", time=True):
                 # We need to load the ref model state dict and compute the logprobs for the ref model
@@ -1199,8 +1212,19 @@ def _compute_logprobs_batch():
                     k: (v.cpu() if v is not None else v) for k, v in model.state_dict().items()
                 }
                 model.load_state_dict(ref_state_dict)
-
-                ref_logprobs = _compute_logprobs_batch()
+                ref_logprobs = _compute_logprobs_batch(
+                    model=model,
+                    data_loader=data_loader,
+                    forward_backward_func=forward_backward_func,
+                    packing_context=packing_context,
+                    trajs_batch_size=len(compute_trajs),
+                    seq_length=args.seq_length,
+                    logprobs_batch_size=logprobs_batch_size,
+                    decoder_seq_length=args.decoder_seq_length,
+                    dtype=dtype,
+                    pp_group=pp_group,
+                    is_correction=args.rl_inference_logprobs_is_correction,
+                )
 
                 # logprobs are [b, seq, h] now.
                 model.load_state_dict(cur_st_dict)
@@ -1210,7 +1234,7 @@ def _compute_logprobs_batch():
             torch.cuda.empty_cache()
 
 
-        if args.rl_use_sequence_packing:
+        if sequence_packing:
             with nvtx_range("pack_logprobs", time=True):
                 # Store logprobs on gpu in packing context
                 # Since PackingContext is a dataclass, we add these as new attributes
@@ -1239,6 +1263,22 @@ def _compute_logprobs_batch():
                     packing_context.packed_inference_logprobs = packed_inference_logprobs.cuda()
                     # Only mark as having inference logprobs for IS correction if enabled
                     packing_context.has_inference_logprobs = args.rl_inference_logprobs_is_correction
+            with nvtx_range("create_dataloader"):
+                # @vitalyk: This function also reconfigures the data loader to count the
+                # global_batch_size in the bins frame of reference.
+                # I think it will be a better design if we split the data loader creating and logic
+                # that reconfigures the microbatch calculator.
+
+                update_microbatch_calculator(
+                    samples_ratio_per_step=samples_ratio_per_step,
+                    num_bins_this_rank = len(packing_context.packed_trajs),
+                    bin_seq_indices = packing_context.packing_info.bin_seq_indices,
+                    global_batch_size=args.global_batch_size, 
+                    rampup_batch_size=args.rampup_batch_size, 
+                    micro_batch_size=args.micro_batch_size, 
+                    decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
+               )
+                loader = get_microbatch_dataloader(len(packing_context.packed_trajs), args.micro_batch_size)
         else:
             with nvtx_range("align_inference_logprobs", time=True):
                 if inference_logprobs is not None:
@@ -1253,14 +1293,20 @@ def _compute_logprobs_batch():
                     # Nullify logprobs if not used in IS correction,
                     if not args.rl_inference_logprobs_is_correction:
                         inference_logprobs = None
-
-        with nvtx_range("create_dataloader"):
-            if args.rl_use_sequence_packing:
-               loader, optimizer_steps = get_microbatch_dataloader(packing_context)
-               runtime_state.global_batches_per_collection = optimizer_steps
-            else:
+            with nvtx_range("create_dataloader"):
+                # Because of multiturn, our batch sizes for non-sequence packed trajectories are not fixed anymore.
+                # As in sequence packing above, we need to reconfigure it too.
                 runtime_state.packing_context = None
-                runtime_state.global_batches_per_collection = global_rollout_count / args.global_batch_size
+
+                reconfigure_num_microbatches_calculator(
+                    rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0,
+                    global_batch_size=math.ceil(samples_ratio_per_step*total_turns_sampled), 
+                    rampup_batch_size=args.rampup_batch_size, 
+                    micro_batch_size=args.micro_batch_size, 
+                    decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
+                    data_parallel_size=mpu.get_data_parallel_world_size(),
+                )
+
                 dataset_tensors = [
                     compute_trajs,
                     advantages,
@@ -1269,20 +1315,20 @@ def _compute_logprobs_batch():
                     original_position_ids,
                     ref_logprobs,
                 ]
-                if args.rl_inference_logprobs_is_correction and inference_logprobs is not None:
+                if is_correction and inference_logprobs is not None:
                     dataset_tensors.append(inference_logprobs)
                 else:
                     dataset_tensors.append(torch.zeros_like(old_logprobs))
-
                 data = TensorDataset(*dataset_tensors)
                 loader = DataLoader(data, batch_size=args.micro_batch_size)
 
+
         with nvtx_range("log-wandb-tb"):
             maybe_log_training_metrics(
                 group_stats=group_stats,
                 current_iteration=args.curr_iteration,
                 tokenizer=tokenizer,
-                example_group=rollouts[0],
+                example_group=example_group,
                 wandb_writer=wandb_writer,
                 tb_writer=tb_writer,
             )
@@ -1290,66 +1336,65 @@ def _compute_logprobs_batch():
     return RerunDataIterator(itertools.cycle(loader))
 
 
-def get_rollout_data_iterator(
-    model: LanguageModule,
-    inference_model: LanguageModule | None,
-    optimizer: MegatronOptimizer,
-    iteration: int,
-    ref_state_dict: Dict[str, torch.Tensor],
-) -> RerunDataIterator:
-
-    args = get_args()
-    tokenizer = get_tokenizer()
-
-    buffered_rollouts = get_environment_rollouts(
-        model, inference_model, optimizer, args.grpo_prompts_per_step, args.grpo_group_size
-    )
-    buffered_rollouts = prepare_data_for_update(model, ref_state_dict, buffered_rollouts, tokenizer)
-
-    return buffered_rollouts
-
-
-def setup_grpo_data_iterator(
+def get_grpo_data_iterator(
     model: LanguageModule,
     inference_model: LanguageModule | None,
     optimizer: MegatronOptimizer,
     iteration: int,
     ref_state_dict: Dict[str, torch.Tensor],
+    grpo_iterations: int,
+    grpo_prompts_per_step: int,
+    grpo_group_size: int,
+    global_batch_size: int,
+    sequence_packing: bool,
+    is_correction: bool,
     buffered_rollouts: RerunDataIterator | None = None,
 ) -> RerunDataIterator:
     """
-    Set up the data iterator for GRPO training.
+    Get the data iterator for GRPO training.
+
+    Depending on the sampling parameters either performs data collections or returns
+    the buffered_rollouts as is.
 
     Args:
         model: The language model
         optimizer: The Megatron optimizer
         iteration: Current training iteration
         ref_state_dict: Reference model state dict for GRPO
+        grpo_iterations: How many steps we reuse the sampled data for.
+        grpo_prompts_per_step: How many prompts we sample per data collection.
+        grpo_group_size: How many samples we do per prompt.
+        global_batch_size: Global batch size.
+        sequence_packing: Use sequence packing if True.
+        is_correction: Use IS correction if True.
         buffered_rollouts: Previously collected rollouts (if any)
 
     Returns:
         RerunDataIterator for the current training step
     """
-    args = get_args()
     runtime_state = get_rl_runtime_state()
 
-    if inference_model is not None:
-        inference_pg_collection = unwrap_model(inference_model[0]).pg_collection
-    else:
-        inference_pg_collection = ProcessGroupCollection.use_mpu_process_groups()
-
     # We collect new rollouts when we've gone over the collected data 'grpo_iterations' times.
+    global_batches_per_collection = (grpo_prompts_per_step * grpo_group_size) // global_batch_size 
     if (
         buffered_rollouts is None or
         iteration == runtime_state.last_collection_iteration + 
-        (args.grpo_iterations * runtime_state.global_batches_per_collection)
+        (grpo_iterations * global_batches_per_collection)
     ):
-        train_data_iterator = get_rollout_data_iterator(model,inference_model, optimizer, iteration, ref_state_dict)
+
+        buffered_rollouts = get_environment_rollouts(
+            model, inference_model, optimizer, grpo_prompts_per_step, grpo_group_size
+        )
+        buffered_rollouts = prepare_data_for_update(model=model, 
+            ref_state_dict=ref_state_dict, 
+            rollouts=buffered_rollouts,
+            tokenizer=get_tokenizer(),
+            sequence_packing=sequence_packing,
+            is_correction=is_correction,
+            )
         runtime_state.reset_iteration_counters(iteration)
-    else:
-        train_data_iterator = buffered_rollouts
 
-    return train_data_iterator
+    return buffered_rollouts
 
 
 def evaluate_and_print_results_rl(
@@ -1389,7 +1434,7 @@ def evaluate_and_print_results_rl(
 
             rank = torch.distributed.get_rank()
             if rank == 0:
-                logger.info(f"Collecting evaluation results...")
+                logger.info("Collecting evaluation results...")
                 agent = get_agent(args)
                 request = EvaluationRequest(
                     inference_interface=inference_interface,
@@ -1702,3 +1747,32 @@ def get_iteration_sequence_count(args):
     if torch.distributed.is_initialized():
         torch.distributed.all_reduce(sequences_tensor, group=mpu.get_data_parallel_group())
     return int(sequences_tensor.item())
+    
+def _pad_nonnull_with_zeros(data: list[Optional[torch.Tensor]], max_len: int) -> torch.Tensor:
+    """Pad each element of a list of tensors to the length required.
+    Args:
+        data: List of tensors to pad.
+        max_len: Maximum length to pad to. Must be higher or equal than the max len of the data tensors.
+    Returns:
+        A padded tensor which is a stacked list of padded input tensors.
+
+    """
+    if all([el is None for el in data]):
+        raise ValueError("At least one element of the data list should be not None.")
+    padded_data = []
+    for chunk in data:
+        if chunk is not None:
+            padding_size = max_len - len(chunk)
+            if padding_size > 0:
+                # Pad with zeros (these positions will be masked anyway)
+                padded = torch.nn.functional.pad(chunk, (0, padding_size), value=0.0)
+                padded_data.append(padded)
+            elif padding_size == 0:
+                padded_data.append(chunk)
+            else:
+                raise ValueError("One of the input tensors has larger length than padding max len.")
+        else:
+            # Create zero tensor for None logprobs
+            padded_data.append(torch.zeros(max_len))
+    return torch.stack(padded_data)
+
diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py
index a5703a4580c..4d983764f77 100644
--- a/megatron/rl/sequence_packing_utils.py
+++ b/megatron/rl/sequence_packing_utils.py
@@ -10,7 +10,6 @@
 from megatron.training.global_vars import get_args, get_tokenizer
 from megatron.training.utils import get_nvtx_range
 from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core import mpu
 import logging
 import typing
@@ -78,7 +77,6 @@ def load_packed_data_by_index(bin_idx: int, packing_context: PackingContext, log
     Args:
         bin_idx: Index of the bin to load.
     """
-    args = get_args()
     # Get packing context (should always be available in packed mode)
     idx = slice(bin_idx, bin_idx + 1)
 
@@ -156,9 +154,8 @@ def log_packing_efficiency(packing_context: PackingContext):
     packing_efficiency = my_tokens / total_capacity if total_capacity > 0 else 0
     avg_seq_length = total_tokens / len(packing_info.seq_lengths)
     rank = mpu.get_data_parallel_rank()
-    data_parallel_world_size = mpu.get_data_parallel_world_size()
 
-    log_single_rank(logger, logging.INFO, f"[Sequence Packing] Statistics:")
+    log_single_rank(logger, logging.INFO, "[Sequence Packing] Statistics:")
     log_single_rank(
         logger,
         logging.INFO,
@@ -269,7 +266,7 @@ def log_packing_efficiency(packing_context: PackingContext):
             log_single_rank(
                 logger,
                 logging.INFO,
-                f"[Sequence Packing]  Round-robin distribution quality:",
+                "[Sequence Packing]  Round-robin distribution quality:",
             )
             log_single_rank(
                 logger,
@@ -398,7 +395,7 @@ def create_empty_bins(
         empty_packing_info_entries,
     )
 
-def get_default_packed_seq_params(seq_length: int, device: torch.device) -> PackedSeqParams:
+def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, device: torch.device) -> PackedSeqParams:
     """Create a default PackedSeqParams that acts as no-op for a single sequence.
 
     This ensures CUDA graph signature consistency when packed_seq_params
@@ -407,6 +404,7 @@ def get_default_packed_seq_params(seq_length: int, device: torch.device) -> Pack
 
     Args:
         seq_length: The sequence length 
+        max_sequences_per_bin: Max sequences to pack in a bin.
         device: Device to create tensors on.
 
     Returns:
@@ -416,7 +414,7 @@ def get_default_packed_seq_params(seq_length: int, device: torch.device) -> Pack
     args = get_args()
 
     # Pad to the maximum number of sequences in the bin for the attention kernel.
-    cu_seqlens = torch.full((args.rl_sequence_packing_max_sequences_per_bin,), seq_length, dtype=torch.int32, device=device)
+    cu_seqlens = torch.full((max_sequences_per_bin,), seq_length, dtype=torch.int32, device=device)
     cu_seqlens[0] = 0
 
     return PackedSeqParams(
@@ -774,7 +772,7 @@ def pack_sequences(
 
         seq_per_bin = [len(indices) for indices in packing_info.bin_seq_indices]
         log_single_rank(
-            logger, logging.DEBUG, (f"Initial packing output (before distribution):")
+            logger, logging.DEBUG, ("Initial packing output (before distribution):")
         )
         log_single_rank(
             logger,
@@ -969,33 +967,20 @@ def distribute_packed_bins(
 def pack_all_trajectories(trajs, generation_masks, inference_logprobs, global_advantages, bin_size, max_sequences_per_bin, packing_algo):
     tokenizer = get_tokenizer()
     data_parallel_world_size = mpu.get_data_parallel_world_size()
+    data_parallel_group = mpu.get_data_parallel_group()
     nvtx_range = get_nvtx_range()
 
     with nvtx_range("regather_trajectories", time=True):
-        # Regather trajectories from all ranks for packing
-        trajs = trajs.cuda()
-        trajs_list = [torch.empty_like(trajs) for _ in range(data_parallel_world_size)]
-        torch.distributed.all_gather(
-            trajs_list, trajs, group=mpu.get_data_parallel_group()
-        )
-        trajs = torch.cat(trajs_list, dim=0)
-
-        # Gather all generation masks
-        generation_masks = generation_masks.cuda()
-        masks_list = [torch.empty_like(generation_masks) for _ in range(data_parallel_world_size)]
-        torch.distributed.all_gather(
-            masks_list, generation_masks, group=mpu.get_data_parallel_group()
-        )
-        generation_masks = torch.cat(masks_list, dim=0)
-
-        # Gather inference logprobs if present
+        def _gather(data):
+            data = data.cuda()
+            data_list = [torch.empty_like(data) for _ in range(data_parallel_world_size)]
+            torch.distributed.all_gather(data_list, data, group=data_parallel_group)
+            return torch.cat(data_list, dim=0)
+
+        trajs = _gather(trajs)    
+        generation_masks = _gather(generation_masks) 
         if inference_logprobs is not None:
-            inference_logprobs = inference_logprobs.cuda()
-            logprobs_list = [torch.empty_like(inference_logprobs) for _ in range(data_parallel_world_size)]
-            torch.distributed.all_gather(
-                logprobs_list, inference_logprobs, group=mpu.get_data_parallel_group()
-            )
-            inference_logprobs = torch.cat(logprobs_list, dim=0)
+            inference_logprobs = _gather(inference_logprobs)
 
     with nvtx_range("pack_sequences", time=True):
         # Create packer with max sequences per bin limit to prevent extreme imbalance
@@ -1073,53 +1058,63 @@ def pack_all_trajectories(trajs, generation_masks, inference_logprobs, global_ad
 
     return packing_context
 
+def update_microbatch_calculator(
+    samples_ratio_per_step: float,
+    num_bins_this_rank: int,
+    bin_seq_indices: List[List[int]],
+    global_batch_size: int, 
+    rampup_batch_size: int, 
+    micro_batch_size: int, 
+    decrease_batch_size_if_needed: bool,
+):
+    """Return a data loader with seqpacked indices with microbatches in bins frame of reference.
+    Args:
+        samples_ratio_per_step: Fraction of sampled trajectories to use per iteration.
+        num_bins_this_rank: Amount of packing bins that belongs to current rank.
+        bin_seq_indices: Global seq indices in the bin, see PackingInfo.
+        global_batch_size: Current global batch size.
+        rampup_batch_size: Rampup batch size. See num_microbatches_calculator.py for more.
+        micro_batch_size: Micro batch size at init.
+        decrease_batch_size_if_needed: Scale down batch size. See num_microbatches_calculator.py for more.
+
+    As a side effect, we calculate the global batch size in the bins frame of reference.
+    In sequence packing, our batch dimension shrinks as we move some trajs onto free
+    space in sequence dimension. The resulting batch size is what we return here.
+    """
 
-def get_microbatch_dataloader(packing_context: PackingContext) -> Tuple[DataLoader, int]:
-    args = get_args()
-    num_bins_this_rank = len(packing_context.packed_trajs)
     dp_world_size = mpu.get_data_parallel_world_size()
 
-    # Ratio of collected sequences to the global batch size
-    pct_of_sequences_per_batch = len(packing_context.packing_info.seq_lengths) / args.global_batch_size
-
     # Ceiling division means we will reuse some bins
     # If we did floor we would leave some behind
-    local_bins_per_step = math.ceil(pct_of_sequences_per_batch * num_bins_this_rank)
-    effective_global_batch_size = local_bins_per_step * dp_world_size
+    local_bins_per_step = math.ceil(samples_ratio_per_step * num_bins_this_rank)
 
-    # Store packing plan in runtime state for the training loop to use
-    optimizer_steps = -(-num_bins_this_rank // local_bins_per_step)
+    bins_bs = local_bins_per_step * dp_world_size
 
     old_num_microbatches = get_num_microbatches()
-
     reconfigure_num_microbatches_calculator(
         rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0,
-        rampup_batch_size=args.rampup_batch_size,
-        global_batch_size=effective_global_batch_size,
-        micro_batch_size=args.micro_batch_size,
+        rampup_batch_size=rampup_batch_size,
+        global_batch_size=bins_bs,
+        micro_batch_size=micro_batch_size,
         data_parallel_size=dp_world_size,
-        decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
+        decrease_batch_size_if_needed=decrease_batch_size_if_needed,
     )
-
     new_num_microbatches = get_num_microbatches()
 
     log_single_rank(
-        logger, logging.INFO, f"[Sequence Packing] Multi-step training plan:"
-    )
-    log_single_rank(
-        logger,
-        logging.INFO,
-        f"[Sequence Packing]  - Target sequences per step: {args.global_batch_size}",
+        logger, logging.INFO, "[Sequence Packing] Multi-step training plan:"
     )
+
     log_single_rank(
         logger,
         logging.INFO,
-        f"[Sequence Packing]  - Bins per rank per step: {pct_of_sequences_per_batch}*{num_bins_this_rank}={local_bins_per_step}",
+        f"[Sequence Packing]  - Bins per rank per step: {samples_ratio_per_step}*{num_bins_this_rank}={local_bins_per_step}",
     )
+
     log_single_rank(
         logger,
         logging.INFO,
-        f"[Sequence Packing]  - Total optimizer steps: {optimizer_steps}",
+        f"[Sequence Packing]  - Target sequences per step: {global_batch_size}",
     )
     log_single_rank(
         logger,
@@ -1127,8 +1122,10 @@ def get_microbatch_dataloader(packing_context: PackingContext) -> Tuple[DataLoad
         f"[Sequence Packing]  - Microbatches per step: {new_num_microbatches} (was {old_num_microbatches})",
     )
 
-    bin_seq_indices = packing_context.packing_info.bin_seq_indices
-    for step in range(min(3, optimizer_steps)):
+    # Opt steps only depends on how much we sample and how much we consume.
+    # We make sure this is an integer division, check validate_args in arguments.py for details.
+    opt_steps = int(1 / samples_ratio_per_step)
+    for step in range(min(3, opt_steps)):
         start_idx = step * local_bins_per_step
         end_idx = min(start_idx + local_bins_per_step, num_bins_this_rank)
         step_bins = end_idx - start_idx
@@ -1145,22 +1142,13 @@ def get_microbatch_dataloader(packing_context: PackingContext) -> Tuple[DataLoad
             f"[Sequence Packing]  - Step {step + 1}: {step_bins} bins, ~{est_global_seqs} sequences globally",
         )
 
-    if optimizer_steps > 3:
-        log_single_rank(logger, logging.INFO, f"  - ... ({optimizer_steps - 3} more steps)")
+    if opt_steps > 3:
+        log_single_rank(logger, logging.INFO, f"  - ... ({opt_steps - 3} more steps)")
 
+def get_microbatch_dataloader(num_bins_this_rank, micro_batch_size):
     bin_indices = torch.arange(num_bins_this_rank)
     dataset = TensorDataset(bin_indices)
-    loader = DataLoader(dataset, batch_size=args.micro_batch_size, shuffle=False, collate_fn=lambda x: x[0], drop_last=True)
-    return loader, optimizer_steps
-
-def update_sequence_packing_metrics(args):
-    """Update bin tracking for sequence packing mode."""
-    if args.rl_use_sequence_packing:
-        bin_count = (
-            mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
-        )
-        args.consumed_train_bins += bin_count
-
+    return DataLoader(dataset, batch_size=micro_batch_size, shuffle=False, collate_fn=lambda x: x[0])
 
 def get_sequence_packing_log_info(args):
     """Get logging information for sequence packing mode."""
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index f269ad02879..4ed4eb1cb73 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1987,8 +1987,6 @@ def _add_rl_args(parser):
                        help='If set, use inference logprobs in importance sampling correction of the loss.')
     group.add_argument('--rl-importance-sampling-truncation-coef', type=float, default=None,
                        help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.")
-    group.add_argument('--rl-calculate-intra-group-similarity', action=argparse.BooleanOptionalAction, default=False,
-                       help='If set, calculate the intra-group similarity of rollouts.')
     group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False,
                        help='Enable sequence packing')
     group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50,
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 500d30b9e73..87d9fe8b841 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -2766,8 +2766,15 @@ def get_e2e_base_metrics():
 
         if getattr(args, 'perform_rl_step', False):
             with torch.no_grad():
-                train_data_iterator = rl_utils.setup_grpo_data_iterator(
-                    model, inference_model, optimizer, iteration, ref_state_dict, buffered_rollouts
+                train_data_iterator = rl_utils.get_grpo_data_iterator(
+                    model, inference_model, optimizer, iteration, ref_state_dict,
+                    grpo_iterations=args.grpo_iterations,
+                    grpo_prompts_per_step=args.grpo_prompts_per_step,
+                    grpo_group_size=args.grpo_group_size,
+                    global_batch_size=args.global_batch_size,
+                    sequence_packing=args.rl_use_sequence_packing,
+                    buffered_rollouts=buffered_rollouts,
+                    is_correction=args.rl_inference_logprobs_is_correction,
                 )
                 # Buffered rollouts are used as a state container for setups when
                 # we use previously-generated data for an update.
@@ -2846,7 +2853,10 @@ def get_e2e_base_metrics():
         if getattr(args, 'perform_rl_step', False) and args.rl_use_sequence_packing:
             iteration_sequences = rl_utils.get_iteration_sequence_count(args)
             # Track bins separately for packed mode
-            rl_utils.update_sequence_packing_metrics(args)
+            bin_count = (
+                mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
+            )
+            args.consumed_train_bins += bin_count
         else:
             batch_size = (
                 mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py
index d3570bee108..fccdb832a07 100644
--- a/tests/unit_tests/rl/test_rl_utils.py
+++ b/tests/unit_tests/rl/test_rl_utils.py
@@ -74,6 +74,8 @@ def initialize_model_parallel(request, monkeypatch):
     Skips if world_size < tp * pp.
     """
     monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
+    monkeypatch.setenv("WANDB_MODE", "disabled")
+    monkeypatch.setenv("LOG_TO_WANDB", "false")
 
     tp, pp = request.param
     world_size = Utils.world_size
@@ -107,6 +109,7 @@ def create_test_args(self, **kwargs):
         args.hidden_size = 128
         args.max_position_embeddings = 256
         args.seq_length = 256
+        args.wandb_project = None
 
         args.micro_batch_size = 1
 
@@ -269,58 +272,68 @@ def test_grpo_loss_truncation(self):
     def test_prepare_data_for_update(self, initialize_model_parallel):
         """Test that getting logprobs at least does not crash."""
         world_size, dp, tp, pp = initialize_model_parallel
+        # Here I assume that we will be consuming all data in one step.
+        group_size = 2
         self.create_test_args(
             micro_batch_size=2,
             seq_length=4,
             curr_iteration=1,
             tensor_model_parallel_size=tp,
             pipeline_model_parallel_size=pp,
+            global_batch_size=dp * 2,
+            grpo_prompts_per_step=dp,
+            grpo_group_size=group_size,
         )
 
         model = MockModel()
         tokenizer = MockTokenizer()
 
         r1 = TokenRollout(
-            trajectory=[1, 2, 3],
+            trajectory=[[1, 2, 3]],
             reward=3.14,
-            generation_mask=[False, True, True],
-            logprobs=[0.1, 0.2, 0.3],
+            generation_mask=[[False, True, True]],
+            logprobs=[[0.1, 0.2, 0.3]],
             env_id='MEGAENV',
             problem_id="2",
         )
         r2 = TokenRollout(
-            trajectory=[1, 2, 3, 4],
+            trajectory=[[1, 2, 3, 4]],
             reward=0.14,
-            generation_mask=[False, True, True, True],
-            logprobs=[0.1, 0.2, 0.3, -1.2],
+            generation_mask=[[False, True, True, True]],
+            logprobs=[[0.1, 0.2, 0.3, -1.2]],
             env_id='MEGAENV',
             problem_id="2",
         )
+
         rollouts = [[r1, r2] for _ in range(dp)]
         try:
-            rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
+            rl_utils.prepare_data_for_update(
+                [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False
+            )
         except AssertionError as e:
             # We expect trajectories to come padded there.
             assert str(e).startswith('Rollout is not the correct length')
 
         r1 = TokenRollout(
-            trajectory=torch.tensor([1, 2, 3, tokenizer.eod], dtype=torch.float).cuda(),
+            trajectory=torch.tensor([[1, 2, 3, tokenizer.eod]], dtype=torch.float).cuda(),
             reward=3.14,
-            generation_mask=torch.tensor([False, True, True, True], dtype=torch.float).cuda(),
-            logprobs=torch.tensor([-0.2, -0.3, -3.2]).cuda(),
+            generation_mask=torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(),
+            logprobs=torch.tensor([[-0.2, -0.3, -3.2]]).cuda(),
             env_id='MEGAENV',
             problem_id="2",
         )
         r2 = TokenRollout(
-            trajectory=torch.tensor([1, 2, 234, tokenizer.eod], dtype=torch.float).cuda(),
+            trajectory=torch.tensor([[1, 2, 234, tokenizer.eod]], dtype=torch.float).cuda(),
             reward=0.14,
-            generation_mask=torch.tensor([False, True, True, True], dtype=torch.float).cuda(),
-            logprobs=torch.tensor([-0.2, -0.3, -1.2]),
+            generation_mask=torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(),
+            logprobs=torch.tensor([[-0.2, -0.3, -1.2]]),
             env_id='MEGAENV',
             problem_id="2",
         )
         rollouts = [[r1, r2] for _ in range(dp)]
-        data_iter = rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
+        data_iter = rl_utils.prepare_data_for_update(
+            [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False
+        )
 
         _, _, old_logprobs, _, _, _, _ = next(data_iter)
         # All logits are ones in the MockModel.
@@ -328,7 +341,8 @@ def test_prepare_data_for_update(self, initialize_model_parallel):
         torch.testing.assert_close(old_logprobs.exp(), torch.ones_like(old_logprobs) / VOCAB)
 
     @pytest.mark.parametrize("use_sequence_packing", [True, False])
-    def test_prepare_trajectories(self, use_sequence_packing):
+    @pytest.mark.parametrize("num_turns", [1, 2])
+    def test_prepare_trajectories(self, use_sequence_packing, num_turns):
         """Test that rollouts are properly prepared for training."""
         seq_length = 8
         self.create_test_args(
@@ -342,34 +356,38 @@ def test_prepare_trajectories(self, use_sequence_packing):
 
         # Create rollouts of varying lengths
         r1 = TokenRollout(
-            trajectory=[1, 2, 3, tokenizer.eod],
+            trajectory=[[1, 2, 3, tokenizer.eod]] * num_turns,
             reward=3.14,
-            generation_mask=[False, True, True, True],
-            logprobs=[0.1, 0.2, 0.3, 0.35],
+            generation_mask=[[False, True, True, True]] * num_turns,
+            logprobs=[[0.1, 0.2, 0.3, 0.35]] * num_turns,
             env_id='MEGAENV',
             problem_id="1",
         )
         r2 = TokenRollout(
-            trajectory=[4, 5, 6, 7, tokenizer.eod],
+            trajectory=[[4, 5, 6, 7, tokenizer.eod]] * num_turns,
             reward=0.14,
-            generation_mask=[False, True, True, True, True],
-            logprobs=[0.4, 0.5, 0.6, 0.7, 0.75],
+            generation_mask=[[False, True, True, True, True]] * num_turns,
+            logprobs=[[0.4, 0.5, 0.6, 0.7, 0.75]] * num_turns,
             env_id='MEGAENV',
             problem_id="2",
         )
         r3 = TokenRollout(
-            trajectory=[8, 9, tokenizer.eod],
+            trajectory=[[8, 9, tokenizer.eod]] * num_turns,
             reward=2.71,
-            generation_mask=[False, True, True],
-            logprobs=[0.8, 0.9, 0.95],
+            generation_mask=[[False, True, True]] * num_turns,
+            logprobs=[[0.8, 0.9, 0.95]] * num_turns,
             env_id='MEGAENV',
             problem_id="3",
         )
 
-        rollouts = [[r1, r2, r3]]
+        rollouts = [r1, r2, r3]
 
         trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories(
-            rollouts, tokenizer, seq_length
+            rollouts,
+            tokenizer,
+            seq_length,
+            sequence_packing=use_sequence_packing,
+            skip_bos_token=False,
         )
 
         expected_trajs = torch.tensor(
@@ -380,7 +398,7 @@ def test_prepare_trajectories(self, use_sequence_packing):
             ],
             dtype=torch.long,
             device=trajs.device,
-        )
+        ).repeat_interleave(num_turns, dim=0)
         assert torch.equal(trajs, expected_trajs)
 
         expected_genmask = torch.tensor(
@@ -391,7 +409,7 @@ def test_prepare_trajectories(self, use_sequence_packing):
             ],
             dtype=torch.bool,
             device=genmask.device,
-        )
+        ).repeat_interleave(num_turns, dim=0)
         assert torch.equal(genmask, expected_genmask)
 
         if use_sequence_packing:
@@ -403,7 +421,7 @@ def test_prepare_trajectories(self, use_sequence_packing):
                 ],
                 dtype=torch.float32,
                 device=inference_logprobs.device,
-            )
+            ).repeat_interleave(num_turns, dim=0)
             torch.testing.assert_close(inference_logprobs, expected_logprobs, rtol=0, atol=0)
         else:
             expected_logprobs = [
@@ -411,8 +429,53 @@ def test_prepare_trajectories(self, use_sequence_packing):
                 [0.4, 0.5, 0.6, 0.7, 0.75],
                 [0.8, 0.9, 0.95],
             ]
+            expected_logprobs = [el for el in expected_logprobs for _ in range(num_turns)]
             assert len(inference_logprobs) == len(expected_logprobs)
             for got, exp in zip(inference_logprobs, expected_logprobs):
                 got_t = got if torch.is_tensor(got) else torch.tensor(got, dtype=torch.float32)
                 exp_t = torch.tensor(exp, dtype=torch.float32, device=got_t.device)
                 torch.testing.assert_close(got_t, exp_t, rtol=0, atol=0)
+
+    def test_single_turn_advantage_calculation(self):
+        rewards = [[-1, 1], [4, 4]]
+        num_turns = [[1, 1], [1, 1]]
+        advs = rl_utils.calculate_grpo_advantages(rewards, num_turns)
+        torch.testing.assert_close(
+            torch.tensor(advs), torch.tensor([-1, 1.0, 0.0, 0.0]), atol=1e-4, rtol=1e-5
+        )
+
+    def test_multi_turn_advantage_calculation(self):
+        rewards = [[-1, 1], [4, 4]]
+        num_turns = [[2, 1], [1, 3]]
+        advs = rl_utils.calculate_grpo_advantages(rewards, num_turns)
+        torch.testing.assert_close(
+            torch.tensor(advs),
+            torch.tensor([-1, -1, 1.0, 0.0, 0.0, 0.0, 0.0]),
+            atol=1e-4,
+            rtol=1e-5,
+        )
+
+    def test_pad_list_of_nones(self):
+        with pytest.raises(ValueError) as e_info:
+            rl_utils._pad_nonnull_with_zeros([None] * 3, 42)
+        assert "At least one" in str(e_info)
+
+    def test_pad_with_wrong_params(self):
+        with pytest.raises(ValueError) as e_info:
+            rl_utils._pad_nonnull_with_zeros([torch.zeros(5)], 4)
+        assert "larger length" in str(e_info)
+
+    def test_pad_full_size(self):
+        padded = rl_utils._pad_nonnull_with_zeros([torch.zeros(5), torch.zeros(5)], 5)
+        assert padded.shape == (2, 5)
+
+    def test_pad_some_nones(self):
+        padded = rl_utils._pad_nonnull_with_zeros([None, torch.zeros(5)], 5)
+        assert padded.shape == (2, 5)
+        assert (padded[0] == 0).all()
+
+    def test_pad_normal(self):
+        padded = rl_utils._pad_nonnull_with_zeros(
+            [torch.zeros(2), torch.zeros(3), torch.zeros(4)], 5
+        )
+        assert padded.shape == (3, 5)
diff --git a/tests/unit_tests/rl/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py
index 548aedf55fd..44a3de762f0 100644
--- a/tests/unit_tests/rl/test_sequence_packing_utils.py
+++ b/tests/unit_tests/rl/test_sequence_packing_utils.py
@@ -1,5 +1,8 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+from unittest.mock import patch
+
+import pytest
 import torch
 
 from megatron.rl import rl_utils, sequence_packing_utils
@@ -407,3 +410,56 @@ def test_compute_packed_inference_logprobs_stats_shape_mismatch():
 
     # Stats should remain None due to shape mismatch
     assert group_stats.mean_piold_to_inf_prob is None
+
+
+@pytest.mark.parametrize(
+    "ratio,local_bins,world,expected_bs",
+    [
+        (1.0, 1, 8, 8),  # no stale data (ratio 1.), everything divides perfectly.
+        (1.0, 42, 8, 42 * 8),  # no stale data (ratio 1.), everything divides perfectly, more bins
+        (
+            0.5,
+            1,
+            8,
+            8,
+        ),  # 0.5 means we use half of all seqs per step, they all fit 1 bin -> we should reuse
+        (1 / 3, 4, 8, 16),  # third of the data per step, nonint division
+    ],
+)
+def test_get_bins_bs_and_steps(ratio, local_bins, world, expected_bs):
+    # Make a dummy struct to check only the required fields.
+    # Divide by ratio to make sure the samples are divisible by global_bs in the test.
+    n_seqs = int(world * 7 / ratio)
+    global_bs_in_seq = int(n_seqs * ratio)
+
+    def side_eff(
+        rank,
+        rampup_batch_size,
+        global_batch_size,
+        micro_batch_size,
+        data_parallel_size,
+        decrease_batch_size_if_needed,
+    ):
+        # Inside of the get_microbatch_dataloader, we compute the batch size in bins.
+        # We want to test this variable.
+        global actual_bs
+        actual_bs = global_batch_size
+
+    with patch('megatron.rl.sequence_packing_utils.get_num_microbatches', return_value=1):
+        with patch(
+            'megatron.rl.sequence_packing_utils.reconfigure_num_microbatches_calculator',
+            side_effect=side_eff,
+        ):
+            with patch('megatron.core.mpu.get_data_parallel_world_size', return_value=world):
+                sequence_packing_utils.update_microbatch_calculator(
+                    samples_ratio_per_step=ratio,
+                    num_bins_this_rank=local_bins,
+                    bin_seq_indices=[],
+                    global_batch_size=global_bs_in_seq,
+                    rampup_batch_size=1,
+                    micro_batch_size=1,
+                    decrease_batch_size_if_needed=False,
+                )
+
+    # Iterator is local, batch size is global
+    assert expected_bs == actual_bs
diff --git a/train_rl.py b/train_rl.py
index 299843bcff3..cfc010b3c04 100644
--- a/train_rl.py
+++ b/train_rl.py
@@ -260,6 +260,7 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False):
     if packed_seq_params is None:
         packed_seq_params = get_default_packed_seq_params(
             seq_length=tokens.shape[1],
+            max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin,
             device=tokens.device,
         )
 

From 36411ddff10623eeeb184192738e9b787352073d Mon Sep 17 00:00:00 2001
From: Jon Barker <jbarker@nvidia.com>
Date: Thu, 29 Jan 2026 11:40:44 -0700
Subject: [PATCH 71/79] Reapply 3955c49ed9af5e5b38dccdd30c1323c00b9bcd29
 (#3146)

---
 examples/rl/model_configs/nemotron6_3b_moe.sh |   6 -
 .../distributed/distributed_data_parallel.py  |  38 ++++
 .../core/distributed/param_and_grad_buffer.py |  36 ++++
 megatron/rl/rl_utils.py                       |  24 ++-
 .../golden_values_dev_dgx_h100.json           |  28 +--
 tests/unit_tests/rl/test_rl_utils.py          | 175 ++++++++++++++++++
 6 files changed, 278 insertions(+), 29 deletions(-)

diff --git a/examples/rl/model_configs/nemotron6_3b_moe.sh b/examples/rl/model_configs/nemotron6_3b_moe.sh
index 8efe0b2debb..eff4f6cf0b3 100644
--- a/examples/rl/model_configs/nemotron6_3b_moe.sh
+++ b/examples/rl/model_configs/nemotron6_3b_moe.sh
@@ -5,12 +5,6 @@ EP=${EP:-32}
 NODES_REQUIRED=${NODES_REQUIRED:-4}
 LLM="nemotron6_3b_moe"
 
-ROOT_DIR="/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/nemotron6"
-
-CHECKPOINT="${ROOT_DIR}/3b_hybrid_moe/checkpoints/phase2_lc_reinit_emb/"
-
-TOKENIZER_MODEL="${ROOT_DIR}/tokenizers/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json"
-
 echo "Using Nemotron6 3B MOE model checkpoint"
 SCRIPT_PATH="${BASH_SOURCE[0]}"
 source $(dirname $SCRIPT_PATH)/common.sh
diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index 421b279b17d..55179ff3024 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -571,3 +571,41 @@ def broadcast_params(self):
                 src=torch.distributed.get_global_rank(data_parallel_group, 0),
                 group=data_parallel_group,
             )
+
+    def offload_grad_buffers(self, synchronize: bool = True, empty_cache: bool = True) -> None:
+        """
+        Free all grad_data tensors to release GPU memory.
+
+        Uses storage().resize_(0) to release memory while keeping tensor views intact.
+        All bucket.grad_data and param.main_grad views remain valid tensor objects
+        (though accessing them during offload is undefined behavior).
+
+        Args:
+            synchronize: Whether to call torch.cuda.synchronize() before freeing.
+            empty_cache: Whether to call torch.cuda.empty_cache() after freeing.
+        """
+        if synchronize:
+            torch.cuda.synchronize()
+
+        for buffer in self.buffers + self.expert_parallel_buffers:
+            buffer.offload_to_cpu(move_params=False, move_grads=True)
+
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+    def restore_grad_buffers(self, synchronize: bool = True) -> None:
+        """
+        Reallocate grad_data tensors on GPU.
+
+        All existing views (bucket.grad_data, param.main_grad) automatically
+        become valid again since they share the same storage. The grad_data
+        is zeroed after reallocation.
+
+        Args:
+            synchronize: Whether to call torch.cuda.synchronize() after allocation.
+        """
+        for buffer in self.buffers + self.expert_parallel_buffers:
+            buffer.reload_from_cpu(move_params=False, move_grads=True)
+
+        if synchronize:
+            torch.cuda.synchronize()
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index 7abdaab103b..b9480533d7a 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -801,6 +801,10 @@ def _does_param_require_new_bucket(param):
                     requires_grad=False,
                 )
 
+        self.grad_data_size = 0
+        self.param_data_size = 0
+        self.param_data_cpu = None
+
         # Finally, map param.data and param.main_grad fields to buffers.
         bucket_params = []
         bucket_start_index = 0
@@ -951,6 +955,38 @@ def reset(self):
         """
         self.grad_data.zero_()
 
+    def offload_to_cpu(self, move_params: bool = True, move_grads: bool = True) -> None:
+        """
+        Offload the buffers to CPU.
+        """
+        if move_grads and self.grad_data is not None and self.grad_data.storage().size() > 0:
+            self.grad_data_size = self.grad_data.storage().size()
+            self.grad_data.storage().resize_(0)
+        if move_params and self.param_data is not None and self.param_data.storage().size() > 0:
+            self.param_data_size = self.param_data.storage().size()
+            if self.param_data_cpu is not None:
+                self.param_data_cpu.copy_(self.param_data, non_blocking=True)
+            else:
+                self.param_data_cpu = self.param_data.cpu().pin_memory()
+            self.param_data.storage().resize_(0)
+
+    def reload_from_cpu(self, move_params: bool = True, move_grads: bool = True):
+        """
+        Reload the buffers from CPU.
+        """
+        if (
+            move_params
+            and self.param_data is not None
+            and self.param_data_cpu is not None
+            and self.param_data.storage().size() == 0
+        ):
+            self.param_data.storage().resize_(self.param_data_size)
+            self.param_data.copy_(self.param_data_cpu, non_blocking=True)
+        if move_grads and self.grad_data is not None and self.grad_data_size > 0:
+            self.grad_data.storage().resize_(self.grad_data_size)
+            self.grad_data.zero_()
+            self.grad_data_size = 0
+
 
 def partition_buckets(
     buffers: List[_ParamAndGradBuffer], force_single_bucket_group: bool = False
diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index f5ba6d35d8c..364a80db81e 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -104,7 +104,6 @@ def _maybe_prefetch_separate_inference_model_weights(model_core, *, to_cpu: bool
         return
     if args.rl_inference_model_unified_memory_level != 1:
         return
-
     device = -1 if to_cpu else int(torch.cuda.current_device())
     # Note: include_buffers=False because buffers created with explicit device= in register_buffer()
     # are not allocated via the UVM mempool and will fail UVM operations. Only parameters are UVM-allocated.
@@ -461,13 +460,13 @@ def get_environment_rollouts(
     args = get_args()
     nvtx_range = get_nvtx_range()
 
+    if args.rl_offload_optimizer_during_inference:
+        with nvtx_range("offload-optimizer-state-and-grad-buffers-during-inference"):
+            model[0].offload_grad_buffers()
+            optimizer.offload_to_cpu()
+             
     # If we have seperate training and inference models we to refit weights from the training model to the inference model.
     if inference_model is not None:
-        if args.rl_offload_optimizer_during_inference:
-            with nvtx_range("offload-optimizer-before-refit"):
-                optimizer.offload_to_cpu()
-                torch.cuda.empty_cache()
-
         # If the separate inference model weights were prefetched to CPU while idle, bring them
         # back to GPU before refit/copy and before any CUDA-graph'd inference.
         with nvtx_range("prefetch-inference-model-weights-to-gpu"):
@@ -496,7 +495,7 @@ def get_environment_rollouts(
             optimizer,
             args.cuda_graph_impl,
             args.rl_reset_cuda_graphs,
-            args.rl_offload_optimizer_during_inference,
+            False, # offload optimizer during rollout collection is handled above
             args.rl_offload_kv_cache_during_training,
             args.rl_remove_kv_cache_during_training,
         ) as inference_interface:
@@ -536,6 +535,11 @@ def get_environment_rollouts(
             torch.distributed.broadcast_object_list(rollouts, src=0)
         logger.debug(f"Got rollouts on rank {rank}")
 
+    if args.rl_offload_optimizer_during_inference:
+        with nvtx_range("restore-optimizer-state-and-grad-buffers-after-inference"):
+            model[0].restore_grad_buffers()
+            optimizer.restore_from_cpu()
+
     if lang_rl_log_dir and rank == get_pg_rank(inference_pg_collection.tp):
         with open(
             lang_rl_log_dir
@@ -1654,7 +1658,8 @@ def megatron_rl_inference_mode(
     with torch.no_grad():
 
         if offload_optimizer_during_inference:
-            with nvtx_range("offload-optimizer-before-inference"):
+            with nvtx_range("offload-optimizer-state-and-grad-buffers-before-inference"):
+                model[0].offload_grad_buffers()
                 optimizer.offload_to_cpu()
 
         # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to.
@@ -1717,7 +1722,8 @@ def megatron_rl_inference_mode(
             _maybe_prefetch_separate_inference_model_weights(model_core, to_cpu=True)
 
         if offload_optimizer_during_inference:
-            with nvtx_range("onload-optimizer-after-inference"):
+            with nvtx_range("onload-optimizer-state-and-grad-buffers-after-inference"):
+                model[0].restore_grad_buffers()
                 optimizer.restore_from_cpu()
 
         lang_module.train()
diff --git a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
index 339197f8f0f..1c13e432979 100644
--- a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/golden_values_dev_dgx_h100.json
@@ -28,11 +28,11 @@
         "end_step": 5,
         "step_interval": 1,
         "values": {
-            "1": 76691947520.0,
-            "2": 76708724736.0,
-            "3": 76708724736.0,
-            "4": 76708724736.0,
-            "5": 76708724736.0
+            "1": 76683395072.0,
+            "2": 76694667264.0,
+            "3": 76694667264.0,
+            "4": 76694667264.0,
+            "5": 76694667264.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -40,11 +40,11 @@
         "end_step": 5,
         "step_interval": 1,
         "values": {
-            "1": 76691955712.0,
-            "2": 77045972992.0,
-            "3": 77046243328.0,
-            "4": 77047095296.0,
-            "5": 77047095296.0
+            "1": 76683403264.0,
+            "2": 77029359616.0,
+            "3": 77029900288.0,
+            "4": 77030817792.0,
+            "5": 77030817792.0
         }
     },
     "iteration-time": {
@@ -53,10 +53,10 @@
         "step_interval": 1,
         "values": {
             "1": "nan",
-            "2": 135.80412,
-            "3": 81.98981,
-            "4": 82.00576,
-            "5": 82.33207
+            "2": 152.30721,
+            "3": 105.10506,
+            "4": 104.09995,
+            "5": 102.75745
         }
     }
 }
\ No newline at end of file
diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py
index fccdb832a07..cff62d40f0e 100644
--- a/tests/unit_tests/rl/test_rl_utils.py
+++ b/tests/unit_tests/rl/test_rl_utils.py
@@ -7,11 +7,16 @@
 import pytest
 import torch
 
+from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
 from megatron.core.enums import ModelType
 from megatron.core.models.common.language_module.language_module import LanguageModule
+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.num_microbatches_calculator import destroy_num_microbatches_calculator
+from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
 from megatron.core.pipeline_parallel.utils import is_pp_last_stage
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.rl import rl_utils
 from megatron.rl.agent.api import TokenRollout
@@ -479,3 +484,173 @@ def test_pad_normal(self):
             [torch.zeros(2), torch.zeros(3), torch.zeros(4)], 5
         )
         assert padded.shape == (3, 5)
+
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
+            for tp, pp in itertools.product([1, 2], [1, 2])
+            if tp * pp <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    def test_grad_buffer_offload(self, initialize_model_parallel):
+        """Test that grad buffer offload/restore correctly frees and restores GPU memory."""
+        world_size, dp, tp, pp = initialize_model_parallel
+        self.create_test_args(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
+
+        model_parallel_cuda_manual_seed(123)
+
+        # Create a realistic GPTModel as used in RL training
+        transformer_config = TransformerConfig(
+            num_layers=2, hidden_size=64, num_attention_heads=4, use_cpu_initialization=True
+        )
+        gpt_model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
+            vocab_size=256,
+            max_sequence_length=32,
+        ).cuda()
+
+        ddp_config = DistributedDataParallelConfig(
+            grad_reduce_in_fp32=True,
+            use_distributed_optimizer=True,
+            overlap_grad_reduce=False,
+            bucket_size=None,  # Single bucket for simplicity
+        )
+
+        ddp_model = DistributedDataParallel(
+            transformer_config, ddp_config=ddp_config, module=gpt_model
+        )
+
+        all_buffers = ddp_model.buffers + ddp_model.expert_parallel_buffers
+
+        # Verify initial storage is allocated
+        initial_sizes = [buf.grad_data.storage().size() for buf in all_buffers]
+        assert all(size > 0 for size in initial_sizes), "Expected non-zero initial storage"
+
+        # Offload grad buffers to CPU
+        ddp_model.offload_grad_buffers()
+
+        # Verify storage is released
+        for buf in all_buffers:
+            assert buf.grad_data.storage().size() == 0, "Expected zero storage after offload"
+
+        # Restore grad buffers to GPU
+        ddp_model.restore_grad_buffers()
+
+        # Verify storage is restored
+        restored_sizes = [buf.grad_data.storage().size() for buf in all_buffers]
+        assert (
+            initial_sizes == restored_sizes
+        ), f"Expected restored sizes {restored_sizes} to match initial {initial_sizes}"
+
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp), id=f"tp{tp}-pp{pp}")
+            for tp, pp in itertools.product([1, 2], [1, 2])
+            if tp * pp <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    def test_optimizer_offload(self, initialize_model_parallel):
+        """Test that optimizer offload_to_cpu/restore_from_cpu correctly moves state to/from CPU."""
+        world_size, dp, tp, pp = initialize_model_parallel
+        self.create_test_args(tensor_model_parallel_size=tp, pipeline_model_parallel_size=pp)
+        model_parallel_cuda_manual_seed(123)
+
+        # Create a realistic GPTModel as used in RL training
+        transformer_config = TransformerConfig(
+            num_layers=2, hidden_size=64, num_attention_heads=4, use_cpu_initialization=True
+        )
+        gpt_model = GPTModel(
+            config=transformer_config,
+            transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
+            vocab_size=256,
+            max_sequence_length=32,
+        ).cuda()
+
+        ddp_config = DistributedDataParallelConfig(
+            grad_reduce_in_fp32=True,
+            use_distributed_optimizer=True,
+            overlap_grad_reduce=False,
+            bucket_size=None,  # Single bucket for simplicity
+        )
+
+        ddp_model = DistributedDataParallel(
+            transformer_config, ddp_config=ddp_config, module=gpt_model
+        )
+
+        # Create optimizer
+        optimizer_config = OptimizerConfig(
+            optimizer='adam', bf16=True, use_distributed_optimizer=True
+        )
+        optimizer = get_megatron_optimizer(optimizer_config, [ddp_model])
+
+        # Manually initialize optimizer state (simulating what happens after first step)
+        # This avoids needing to run a full forward/backward/step cycle
+        for opt in optimizer.chained_optimizers:
+            if hasattr(opt, 'optimizer') and opt.optimizer is not None:
+                for group in opt.optimizer.param_groups:
+                    for p in group['params']:
+                        if len(opt.optimizer.state[p]) == 0:
+                            # Initialize Adam state (exp_avg and exp_avg_sq) on GPU
+                            opt.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data)
+                            opt.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data)
+                            opt.optimizer.state[p]['step'] = torch.tensor(1)
+
+        # Helper to check if optimizer state tensors are on GPU or CPU
+        def get_optimizer_state_devices():
+            devices = set()
+            for opt in optimizer.chained_optimizers:
+                if hasattr(opt, 'optimizer') and opt.optimizer is not None:
+                    for state_dict in opt.optimizer.state.values():
+                        for v in state_dict.values():
+                            if isinstance(v, torch.Tensor):
+                                devices.add(str(v.device))
+            return devices
+
+        # Verify optimizer state is initially on GPU
+        initial_devices = get_optimizer_state_devices()
+        assert any(
+            'cuda' in d for d in initial_devices
+        ), f"Expected optimizer state on GPU initially, got devices: {initial_devices}"
+
+        # Record GPU memory before offload
+        torch.cuda.synchronize()
+        memory_before_offload = torch.cuda.memory_allocated()
+
+        # Offload optimizer state to CPU
+        optimizer.offload_to_cpu()
+
+        # Verify GPU memory decreased (optimizer state should be freed)
+        torch.cuda.synchronize()
+        memory_after_offload = torch.cuda.memory_allocated()
+        assert memory_after_offload < memory_before_offload, (
+            f"Expected GPU memory to decrease after offload. "
+            f"Before: {memory_before_offload}, After: {memory_after_offload}"
+        )
+
+        # Verify optimizer state is now on CPU
+        offloaded_devices = get_optimizer_state_devices()
+        assert all(
+            'cpu' in d for d in offloaded_devices
+        ), f"Expected all optimizer state on CPU after offload, got devices: {offloaded_devices}"
+
+        # Restore optimizer state to GPU
+        optimizer.restore_from_cpu()
+
+        # Verify optimizer state is back on GPU
+        restored_devices = get_optimizer_state_devices()
+        assert any(
+            'cuda' in d for d in restored_devices
+        ), f"Expected optimizer state on GPU after restore, got devices: {restored_devices}"
+
+        # Verify GPU memory increased after restore (optimizer state reallocated)
+        torch.cuda.synchronize()
+        memory_after_restore = torch.cuda.memory_allocated()
+        assert memory_after_restore > memory_after_offload, (
+            f"Expected GPU memory to increase after restore. "
+            f"After offload: {memory_after_offload}, After restore: {memory_after_restore}"
+        )

From dbd8dda552e39d6ca0628662b760b1c2fae4e305 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Thu, 29 Jan 2026 22:14:38 +0100
Subject: [PATCH 72/79] Revert "Multiturn rollout support prep (#2966)" (#3153)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 megatron/rl/agent/api.py                      |  14 +-
 megatron/rl/agent/reward_only_agent.py        |   8 +-
 megatron/rl/rl_utils.py                       | 552 ++++++++----------
 megatron/rl/sequence_packing_utils.py         | 128 ++--
 megatron/training/arguments.py                |   2 +
 megatron/training/training.py                 |  16 +-
 tests/unit_tests/rl/test_rl_utils.py          | 121 +---
 .../rl/test_sequence_packing_utils.py         |  56 --
 train_rl.py                                   |   1 -
 9 files changed, 354 insertions(+), 544 deletions(-)

diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py
index 9568db3a54d..34efa68d85a 100644
--- a/megatron/rl/agent/api.py
+++ b/megatron/rl/agent/api.py
@@ -46,8 +46,8 @@ class GroupedRolloutRequest(Request):
 class Rollout(AgentBaseModel):
     """Data for language-based Rollout."""
 
-    trajectory: list[str]
-    prompt_length: list[int] | None = None
+    trajectory: str
+    prompt_length: int | None = None
     reward: float = None
     env_id: str | None = None
     problem_id: str | None = None
@@ -56,10 +56,10 @@ class Rollout(AgentBaseModel):
 class TokenRollout(AgentBaseModel):
     """Tokenized representation of a language-based Rollout."""
 
-    trajectory: list[list[int]]
+    trajectory: list[int]
     reward: list[float] | float
-    generation_mask: list[list[bool]] | None = None
-    logprobs: list[list[float]] | None = None
+    generation_mask: list[list[int]] | list[bool] | None = None
+    logprobs: list[float] | None = None
     env_id: str | None = None
     problem_id: str | None = None
 
@@ -67,8 +67,8 @@ class TokenRollout(AgentBaseModel):
 class ContrastiveRollout(AgentBaseModel):
     """Contrastive/Preference data for language-based Rollout."""
 
-    chosen_trajectory: list[str]
-    rejected_trajectory: list[str]
+    chosen_trajectory: str
+    rejected_trajectory: str
 
 
 class Head2HeadRolloutRequest(Request):
diff --git a/megatron/rl/agent/reward_only_agent.py b/megatron/rl/agent/reward_only_agent.py
index 53b1f7407b2..2e81674c74d 100644
--- a/megatron/rl/agent/reward_only_agent.py
+++ b/megatron/rl/agent/reward_only_agent.py
@@ -104,16 +104,16 @@ async def rollout_from_response(
                 for x in range(len(response.token_ids))
             ]
             rollout = TokenRollout(
-                trajectory=[response.token_ids],
+                trajectory=response.token_ids,
                 reward=await self.get_reward(response_text, golden),
-                logprobs=[logprobs],
-                generation_mask=[generation_mask],
+                logprobs=logprobs,
+                generation_mask=generation_mask,
                 env_id=self.env_id,
                 problem_id=golden['problem_id'] if 'problem_id' in golden else None,
             )
         else:
             rollout = Rollout(
-                trajectory=[raw_text],
+                trajectory=raw_text,
                 reward=await self.get_reward(response_text, golden),
                 env_id=self.env_id,
                 problem_id=golden['problem_id'] if 'problem_id' in golden else None,
diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
index 364a80db81e..973a396b909 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -2,18 +2,18 @@
 
 import gc
 
-import copy
-from functools import partial
 # Keep this to make the env registered.
 import itertools
-import math
+import json
 import logging
+import math
 import pickle
 from collections import Counter, defaultdict
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
+from difflib import SequenceMatcher
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional 
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -26,7 +26,6 @@
 from megatron.core.datasets.megatron_tokenizer import MegatronLegacyTokenizer
 from megatron.core.full_cuda_graph import FullCudaGraphWrapper
 from megatron.core.models.common.language_module.language_module import LanguageModule
-from megatron.core.num_microbatches_calculator import reconfigure_num_microbatches_calculator
 from megatron.core.optimizer import MegatronOptimizer
 from megatron.core.pipeline_parallel import get_forward_backward_func
 from megatron.core.pipeline_parallel.utils import is_pp_last_stage, get_pp_last_rank
@@ -46,10 +45,10 @@
     compute_packed_inference_logprobs_stats,
     pack_all_trajectories,
     load_packed_data_by_index,
+    update_sequence_packing_metrics,
     get_sequence_packing_tensorboard_metrics,
     get_sequence_packing_log_info,
     get_default_packed_seq_params,
-    update_microbatch_calculator,
 )
 from megatron.rl.agent.api import (
     EvaluationRequest,
@@ -67,6 +66,7 @@
 from megatron.training.global_vars import (
     get_args,
     get_tensorboard_writer,
+    get_timers,
     get_tokenizer,
     get_wandb_writer,
 )
@@ -209,13 +209,13 @@ def verify_model_weights_swap(
         if inf_was_training:
             inf_core.train()
 
-Rollouts = list[TokenRollout | Rollout]
-GroupedRollouts = list[Rollouts]
+GroupedRollouts = list[list[TokenRollout | Rollout]]
 
 
 @dataclass(slots=True)
 class RolloutStats:
     mean_reward: float
+    mean_sim: None | float
     mean_length: float
     mean_length_std: float
     max_length: float
@@ -233,7 +233,6 @@ class RolloutStats:
     min_inf_prob: None | float
     max_inf_prob: None | float
     mean_inf_prob: None | float
-    num_turns: list[int] # num_turns per traj
 
 
 # Runtime state container for RL-specific data that shouldn't be checkpointed
@@ -243,6 +242,7 @@ class RLRuntimeState:
     def __init__(self):
         self.packing_context = None
         self.last_collection_iteration = 0
+        self.global_batches_per_collection = 0
         self.sequences_this_iteration_on_rank = 0
         self.latest_batch_num_sequences = 0
 
@@ -613,20 +613,20 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa
 
     """
 
-    args = get_args()
     # Ensure packed_seq_params is always provided for CUDA graph signature consistency
     if packed_seq_params is None and sequence_packing:
         packed_seq_params = get_default_packed_seq_params(
             seq_length=tokens.shape[1],
-            max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin,
             device=tokens.device,
         )
 
     nvtx_range = get_nvtx_range()
 
     with nvtx_range("get-logprobs", time=False):
+
         with nvtx_range("forward-pass", time=False):
             # TODO(vitalyk): use fp16/bf16 as a function argument. Do not use args.
+            args = get_args()
 
             attention_mask_for_forward = None
 
@@ -658,46 +658,19 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa
             return logprobs
 
 
-def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[int]]) -> np.ndarray:
-    """Calculate GRPO advantages from rewards/num_turns.
-
-    For multiturn rollouts, the logic is a bit more involved.
-    # For training, we'll be turning each turn into a trajectory with the same reward
-    # within a trajectory, e.g. if [[a,b],[c,d,e]] trajectory has reward 1.0, we will
-    # get [a,b] with 1.0 and [c,d,e] with 1.0 when doing updates.
-    """
-
-    rewards = np.array(rewards)
-
-    num_turns = np.array(num_turns)
-    # Each outer dimension of num_turns is a group. Sum of those gives total num_turns per group.
-    # Let's use this to calculate advantage.
-    # mean/std should be repeated based on group lens
-    group_turns = num_turns.sum(axis=-1)
-    reward_means = rewards.mean(axis=1, keepdims=True).repeat(group_turns)
-    reward_stds = rewards.std(axis=1, keepdims=True).repeat(group_turns)
-
-    # rewards are originally [g, group_size]
-    # Making an assumption that all groups are of the same size!
-    # @vitalyk: this will go away when we start sending env-based sample reqs.
-    rewards = rewards.flatten().repeat(num_turns.flatten())
-
-    return ((rewards - reward_means) / (1e-4 + reward_stds)).tolist()
-
-
 def compute_group_stats(
-    rollouts: GroupedRollouts, tokenizer: MegatronLegacyTokenizer, seq_len: int,
+    rollouts: GroupedRollouts, tokenizer: MegatronLegacyTokenizer
 ) -> RolloutStats:
     """Add group-based rollout stats for logging.
 
     Args:
         rollouts: Rollouts to generate the stats for. Each inner list is a group (as in GRPO group), i.e. all rollouts are for the same prompt.
         tokenizer: Tokenizer to tokenize the rollouts in case they are raw strings.
-        seq_len: Maximum sequence length.
 
     Returns:
        RolloutStats object containing all the stats.
     """
+    args = get_args()
     # TODO (rkirby) Maybe do some of this after the tensor building
     group_reward_means = []
     group_reward_stds = []
@@ -705,45 +678,54 @@ def compute_group_stats(
     group_length_stds = []
     group_length_maxs = []
     group_length_mins = []
-    rewards = []
-    num_turns = [] # num_turns per traj
+    group_rollout_similarities = []
     for group in rollouts:
         group_rewards = []
         group_lengths = []
-        group_num_turns = []
         for rollout in group:
-            group_num_turns.append(len(rollout.trajectory))
             if isinstance(rollout, TokenRollout):
-                for turn_traj in rollout.trajectory:
-                    detokenized_traj = tokenizer.detokenize(turn_traj)
-                    lang_rl_log(
-                        f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} tokens] {detokenized_traj}"
-                    )
-                    # TODO(vitalyk): how does multiturn change EOD/EOT?
-                    assert (len(turn_traj) == seq_len) or (
-                        turn_traj[-1] == tokenizer.eod
-                    ), f"Rollout is not the correct length: {len(turn_traj)} {turn_traj[-1]}\n{detokenized_traj}"
+                lang_rl_log(
+                    f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} tokens] {tokenizer.detokenize(rollout.trajectory)}"
+                )
+                assert (len(rollout.trajectory) == args.seq_length) or (
+                    rollout.trajectory[-1] == tokenizer.eod
+                ), f"Rollout is not the correct length: {len(rollout.trajectory)} {rollout.trajectory[-1]}\n{tokenizer.detokenize(rollout.trajectory)}"
             else:
                 lang_rl_log(
                     f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} chars] {rollout.trajectory}"
                 )
             group_rewards.append(rollout.reward)
-            #TODO(vitalyk): What is the semantics behind traj length in multiturn? Should we take the last only? Average them instead of extending?
-            group_lengths.extend(len(t) for t in rollout.trajectory)
+            group_lengths.append(len(rollout.trajectory))
+        if args.rl_calculate_intra_group_similarity:
+            # We can probably compute this outside, but in case we switch to different group sizes for different envs, let's keep it here.
+            combos = itertools.combinations(range(len(group)), 2)
+            # For every pair (excluding ourselves), check the sequence similarity and log.
+            # Use this to track the diversity of generated rollouts within a group.
+            intra_group_sim = np.mean(
+                list(
+                    map(
+                        lambda idx_pair: SequenceMatcher(
+                            None, group[idx_pair[0]].trajectory, group[idx_pair[1]].trajectory
+                        ).ratio(),
+                        combos,
+                    )
+                )
+            )
+            group_rollout_similarities.append(intra_group_sim)
+        else:
+            group_rollout_similarities = None
 
         group_length_maxs.append(max(group_lengths))
         group_length_mins.append(min(group_lengths))
         group_reward_means.append(np.mean(group_rewards))
         group_reward_stds.append(np.std(group_rewards))
-        rewards.append(group_rewards)
         group_length_means.append(np.mean(group_lengths))
         # https://arxiv.org/abs/2504.21233 reports that lens variants hurts.
         # Let's track this.
         group_length_stds.append(np.std(group_lengths))
-        num_turns.append(group_num_turns)
-
     stats = RolloutStats(
         mean_reward=np.mean(group_reward_means),
+        mean_sim=np.mean(group_rollout_similarities) if group_rollout_similarities else None,
         mean_length=np.mean(group_length_means),
         mean_length_std=np.mean(group_length_stds),
         max_length=np.max(group_length_maxs),
@@ -759,9 +741,8 @@ def compute_group_stats(
         min_inf_prob=None,
         max_inf_prob=None,
         mean_inf_prob=None,
-        rewards=[r for group in rewards for r in group],
-        advantages=calculate_grpo_advantages(rewards, num_turns),
-        num_turns=[nt for group in num_turns for nt in group],
+        rewards=None,  # We will fill those in later in prepare_data_for_update.
+        advantages=None,  # We will fill those in later in prepare_data_for_update.
     )
     return stats
 
@@ -832,10 +813,11 @@ def maybe_log_training_metrics(
                         columns=['Trajectories', 'Tokens', 'Rewards'],
                         rows=[
                             [
-                                [(tokenizer.detokenize(turn)
+                                (
+                                    tokenizer.detokenize(r.trajectory)
                                     if isinstance(r, TokenRollout)
-                                    else turn) for turn in r.trajectory
-                                ],
+                                    else r.trajectory
+                                ),
                                 r.trajectory,
                                 r.reward,
                             ]
@@ -843,6 +825,11 @@ def maybe_log_training_metrics(
                         ],
                     ),
                 },
+                **(
+                    {'mean_intra_group_similarity': group_stats.mean_sim}
+                    if group_stats.mean_sim
+                    else {}
+                ),
             },
             step=current_iteration,
         )
@@ -851,9 +838,10 @@ def maybe_log_training_metrics(
 
 
 def prepare_trajectories(
-    rollouts: Rollouts, tokenizer: MegatronLegacyTokenizer, seq_length: int, sequence_packing: bool, skip_bos_token: bool
+    rollouts: GroupedRollouts, tokenizer: MegatronLegacyTokenizer, seq_length: int
 ):
     """Pad trajectories and extract the generation masks.
+
     Args:
         rollouts: Rollouts to extract trajectories from.
         tokenizer: Tokenizer to get the padding token and potentially tokenize.
@@ -870,7 +858,6 @@ def prepare_trajectories(
 
     DEFAULT_PAD_TOKENS = ['<|finetune_right_pad_id|>']
 
-
     if isinstance(tokenizer, _HuggingFaceTokenizer):
         if not tokenizer.pad:
             for pad_token in DEFAULT_PAD_TOKENS:
@@ -901,19 +888,17 @@ def prepare_trajectories(
     trajs = []
     generation_masks = []
     inference_logprobs = []
-    for rollout in rollouts:
-        # traj, gen mask and logprobs are lists now.
-        # each list entry is a turn, single-turn environments just have a single-element list.
-        # We assume that all lengths of the structs above have the same lengths (number of turns).
-
-        all_turns_trajectories = (
-            copy.deepcopy(rollout.trajectory)
-            if isinstance(rollout, TokenRollout)
-            else tokenizer.tokenize(rollout.trajectory)
-        )
-        for turn_idx, trajectory in enumerate(all_turns_trajectories):
-            inf_logprobs = rollout.logprobs[turn_idx]
-            generation_mask = rollout.generation_mask[turn_idx] if isinstance(rollout, TokenRollout) else None
+    for group in rollouts:
+        for rollout in group:
+            generation_mask = rollout.generation_mask if isinstance(rollout, TokenRollout) else None
+
+            trajectory = (
+                rollout.trajectory.copy()
+                if isinstance(rollout, TokenRollout)
+                else tokenizer.tokenize(rollout.trajectory)
+            )
+            inf_logprobs = rollout.logprobs
+
             length = len(trajectory)
             assert length <= seq_length, "Rollout too long, how did this happen?"
             if len(trajectory) < seq_length:
@@ -935,7 +920,8 @@ def prepare_trajectories(
             else:
                 inference_logprobs.append(None)
 
-        env_id_counts[rollout.env_id] += 1
+            env_id = rollout.env_id
+            env_id_counts[env_id] += 1
 
     if torch.distributed.is_initialized():
         logger.info(f"[{dist.get_rank()}] Rollout counts:")
@@ -945,18 +931,34 @@ def prepare_trajectories(
     generation_masks = torch.tensor(generation_masks, dtype=torch.bool, device='cpu')
     trajs = torch.tensor(trajs, device='cpu')
 
+    args = get_args()
     # Only process if we have inference_logprobs
     if inference_logprobs and any(lp is not None for lp in inference_logprobs):
-        # We need to pad all logprobs to the same size for sequence packing.
-        # For non-packing mode, keep as list of tensors (unpadded)
-        # This preserves the original behavior where each sequence can have different lengths
-        if sequence_packing:
-            inference_logprobs = _pad_nonnull_with_zeros(inference_logprobs, seq_length)
+        if args.rl_use_sequence_packing:
+            # For sequence packing, we need to pad all logprobs to the same size
+            padded_logprobs = []
+            for logprobs in inference_logprobs:
+                if logprobs is not None:
+                    if len(logprobs) < seq_length:
+                        # Pad with zeros (these positions will be masked anyway)
+                        padding_size = seq_length - len(logprobs)
+                        padded = torch.nn.functional.pad(logprobs, (0, padding_size), value=0.0)
+                        padded_logprobs.append(padded)
+                    else:
+                        padded_logprobs.append(logprobs)
+                else:
+                    # Create zero tensor for None logprobs
+                    padded_logprobs.append(torch.zeros(seq_length))
+            inference_logprobs = torch.stack(padded_logprobs)
+        else:
+            # For non-packing mode, keep as list of tensors (unpadded)
+            # This preserves the original behavior where each sequence can have different lengths
+            pass
     else:
         inference_logprobs = None
 
     # Some sanity checks regarding the tokenization
-    if not skip_bos_token:
+    if not args.rl_skip_bos_token:
         assert (
             tokenizer.bos is None or (trajs[:, 0] == tokenizer.bos).all()
         ), "First token should be bos"
@@ -977,92 +979,11 @@ def prepare_trajectories(
     return trajs, generation_masks, inference_logprobs
 
 
-def logprobs_forward_step(data_iterator, model, is_correction, packing_context=None):
-    # Avoid self.training checks which will trigger cudagraph capture; this path reuses
-    # the forward pass from training after it has been captured on the 1st iteration.
-    model.eval()
-
-    if packing_context is not None:
-        # When using sequence packing, the data iterator returns a tuple with a single element, the bin index.
-        bin_tensor = next(data_iterator)[0]
-        #TODO(jalbericiola): change for named tuple
-        (b_trajs, _, _, _, b_posids, _, _, _, _, _, b_packed_seq_params) = (
-            load_packed_data_by_index(bin_tensor.item(), packing_context, is_correction)
-        )
-    else:
-        b_trajs, b_posids = next(data_iterator)
-        b_packed_seq_params = None
-
-    logprobs = (
-        get_logprobs(
-            model,
-            b_trajs.cuda(),
-            b_posids.cuda(),
-            no_grad=True,
-            sequence_packing=b_packed_seq_params is not None, 
-            packed_seq_params=b_packed_seq_params,
-        ),
-        None,
-    )
-    model.train()
-    return logprobs
-
-
-def _compute_logprobs_batch(
-    model,
-    data_loader,
-    forward_backward_func,
-    packing_context,
-    trajs_batch_size, # n_bins for seq packing, and batch_size for non seq packing
-    seq_length,
-    logprobs_batch_size,
-    decoder_seq_length,
-    dtype,
-    pp_group,
-    is_correction,
-):
-    """Compute logprobs for all batches in the data loader."""
-    logprobs_list = []
-    data_iterator = iter(data_loader)
-    for i in range(len(data_loader)):
-        output_tensor = forward_backward_func(
-            forward_step_func=partial(logprobs_forward_step, is_correction=is_correction, packing_context=packing_context),
-            data_iterator=data_iterator,
-            model=model,
-            num_microbatches=1,
-            seq_length=seq_length,
-            micro_batch_size=logprobs_batch_size,
-            decoder_seq_length=decoder_seq_length,
-            forward_only=True,
-            adjust_tensor_shapes_fn=None,
-        )
-        if is_pp_last_stage(pp_group):
-            logprobs_list.append(output_tensor[0].detach())
-
-    if is_pp_last_stage(pp_group):
-        logprobs = torch.concat(logprobs_list, dim=0)
-        assert logprobs.dtype == dtype
-    else:
-        logprobs = torch.empty(
-            trajs_batch_size,
-            seq_length-1,
-            dtype=dtype,
-            device=torch.cuda.current_device(),
-        )
-
-    # Only PP>1 needs a broadcast from the last stage; for PP=1 the output is already local.
-    if get_pg_size(pp_group) > 1:
-        dist.broadcast(logprobs, src=get_pp_last_rank(pp_group), group=pp_group)
-    return logprobs.cpu()
-
-
 def prepare_data_for_update(
     model: list[LanguageModule],
     ref_state_dict: Dict[str, Any],
     rollouts: GroupedRollouts,
     tokenizer: MegatronLegacyTokenizer,
-    sequence_packing: bool,
-    is_correction: bool,
 ) -> RerunDataIterator:
     """Extract data for the update from raw rollouts.
 
@@ -1071,8 +992,6 @@ def prepare_data_for_update(
         ref_state_dict: Reference policy state dict.
         rollouts: Rollouts to extract the data from.
         tokenizer: Tokenizer to pad/tokenize data.
-        sequence_packing: Use sequence packing if True.
-        is_correction: Prepare data for IS correction if True.
 
     Returns:
         Cycled iterator over dataset batches. In GRPO we might want to go over the same data multiple times.
@@ -1094,50 +1013,59 @@ def prepare_data_for_update(
 
     with nvtx_range("prepare-data-for-update"):
         with nvtx_range("compute-group-stats"):
-            group_stats = compute_group_stats(rollouts, tokenizer, args.seq_length)
-            # TODO(vitalyk): why do we need global_advantages here? go inside packing
-            advantages = global_advantages = torch.tensor(group_stats.advantages, dtype=dtype).cuda()
+            # These are computed on all rollouts for reporting purposes
+            group_stats = compute_group_stats(rollouts, tokenizer)
+            rewards = np.array([[rollout.reward for rollout in group] for group in rollouts])
+            group_stats.rewards = rewards.flatten().tolist()
+            group_stats.advantages = (
+                (
+                    (rewards - rewards.mean(axis=1, keepdims=True))
+                    / (1e-4 + rewards.std(axis=1, keepdims=True))
+                )
+                .flatten()
+                .tolist()
+            )
+            global_rollout_count = len(group_stats.rewards)
+
+        with nvtx_range("prepare_advantages", time=True):        
+            # [g, group_size]
+            # Making an assumption that all groups are of the same size!
+            rewards = torch.tensor(rewards, device='cpu')
+            advantages = (rewards - rewards.mean(axis=1, keepdim=True)) / (
+                1e-4 + rewards.std(axis=1, keepdim=True)
+            )
+
+            # Flatten advantages for training and move to GPU
+            advantages = global_advantages = advantages.view(-1).cuda()
 
         # Now split the rollouts across the data parallel ranks for training
         # This needs to be done at this point because we are about to calculate logprobs
         # Note :- For EP, do not use the expert data parallel group here. Always 
         # use the regular data parallel group. 
-
-        # Use one group as an exampling for logging later.
-        example_group = rollouts[0]
-
-        # Let's expand rollouts getting rid of the groups.
-        # We need this to correctly split the rollouts across dp groups.
-        # And we do not actually need them grouped in anything below anyways.
-        rollouts = [r for g in rollouts for r in g]
-        total_turns_sampled = len(rollouts)
-
-        # We might sample more than we consume in one step.
-        samples_ratio_per_step = args.global_batch_size / (args.grpo_prompts_per_step * args.grpo_group_size)
-        assert samples_ratio_per_step <= 1, "You cannot use more data than you sampled."
-
         if (data_parallel_world_size := mpu.get_data_parallel_world_size()) > 0:
             data_split_size = len(rollouts) // data_parallel_world_size
             data_split_range = (
                 mpu.get_data_parallel_rank() * data_split_size,
                 (mpu.get_data_parallel_rank() + 1) * data_split_size,
             )
-            # TODO(vitalyk): This has to be rewritten assuming we are multiturn now.
             rollouts = rollouts[data_split_range[0] : data_split_range[1]]
-            local_num_turns = sum(group_stats.num_turns[data_split_range[0] : data_split_range[1]])
-            steps_before = sum(group_stats.num_turns[:data_split_range[0]])
-            advantages = advantages[steps_before:steps_before+local_num_turns]
             # First we calculate them on a global level and then we split and recalculate on a local level.
             # Sequence packing and reporting needs it global but non-packing wants it local.
+            rewards = torch.tensor([[r.reward for r in group] for group in rollouts], device='cpu')
+            advantages = (rewards - rewards.mean(axis=1, keepdim=True)) / (
+                1e-4 + rewards.std(axis=1, keepdim=True)
+            )
+
+            # Flatten advantages for training and move to GPU
+            advantages = advantages.view(-1).cuda()
 
         with nvtx_range("prepare_trajectories"):
             trajs, generation_masks, inference_logprobs = prepare_trajectories(
-                rollouts, tokenizer, args.seq_length, sequence_packing, args.rl_skip_bos_token
+                rollouts, tokenizer, args.seq_length
             )
 
-        packing_context = None
         # Build trajectories based on sequence packing or standard processing
-        if sequence_packing:
+        if args.rl_use_sequence_packing:
             with nvtx_range("sequence_packing", time=True):
                 runtime_state.packing_context = packing_context = pack_all_trajectories(
                     trajs, 
@@ -1177,6 +1105,7 @@ def prepare_data_for_update(
                 )
                 logprobs_batch_size = args.micro_batch_size
 
+
         with torch.no_grad(), nvtx_range("compute_logprobs", time=True):
             # Before we can update the model, we need to get the logprobs for the \pi_{old} model.
 
@@ -1187,6 +1116,41 @@ def prepare_data_for_update(
                     forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps
                 )
 
+            def logprobs_forward_step(data_iterator, model):
+
+                # Avoid self.training checks which will trigger cudagraph capture; this path reuses
+                # the forward pass from training after it has been captured on the 1st iteration.
+                model.eval()
+
+                if args.rl_use_sequence_packing:
+                    # When using sequence packing, the data iterator returns a tuple with a single element, the bin index.
+                    bin_tensor = next(data_iterator)[0]
+                    #TODO(jalbericiola): change for named tuple
+                    (b_trajs, _, _, _, b_posids, _, _, _, _, _, b_packed_seq_params) = (
+                        load_packed_data_by_index(bin_tensor.item(), packing_context, args.rl_inference_logprobs_is_correction)
+                    )
+                else:
+                    batch_data = next(data_iterator)
+                    b_trajs, b_posids = batch_data
+                    b_packed_seq_params = None
+
+                b_trajs = b_trajs.cuda()
+                b_posids = b_posids.cuda()
+
+                logprobs = (
+                    get_logprobs(
+                        model,
+                        b_trajs,
+                        b_posids,
+                        no_grad=True,
+                        sequence_packing=args.rl_use_sequence_packing,                       
+                        packed_seq_params=b_packed_seq_params,
+                    ),
+                    None,
+                )
+
+                model.train()
+                return logprobs
 
             dtype = (
                 torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32)
@@ -1195,20 +1159,43 @@ def prepare_data_for_update(
             pg_collection = get_attr_wrapped_model(model, "pg_collection")
             pp_group = pg_collection.pp
 
+            def _compute_logprobs_batch():
+                """Compute logprobs for all batches in the data loader."""
+                logprobs_list = []
+                data_iterator = iter(data_loader)
+                for i in range(len(data_loader)):
+                    output_tensor = forward_backward_func(
+                        forward_step_func=logprobs_forward_step,
+                        data_iterator=data_iterator,
+                        model=model,
+                        num_microbatches=1,
+                        seq_length=args.seq_length,
+                        micro_batch_size=logprobs_batch_size,
+                        decoder_seq_length=args.decoder_seq_length,
+                        forward_only=True,
+                        adjust_tensor_shapes_fn=None,
+                    )
+                    if is_pp_last_stage(pp_group):
+                        logprobs_list.append(output_tensor[0].detach())
+
+                if is_pp_last_stage(pp_group):
+                    logprobs = torch.concat(logprobs_list, dim=0)
+                    assert logprobs.dtype == dtype
+                else:
+                    logprobs = torch.empty(
+                        len(compute_trajs),
+                        args.seq_length - 1,
+                        dtype=dtype,
+                        device=torch.cuda.current_device(),
+                    )
+
+                # Only PP>1 needs a broadcast from the last stage; for PP=1 the output is already local.
+                if get_pg_size(pp_group) > 1:
+                    dist.broadcast(logprobs, src=get_pp_last_rank(pp_group), group=pp_group)
+                return logprobs.cpu()
+
             with torch.no_grad(), nvtx_range("compute_old_logprobs", time=True):
-                old_logprobs = _compute_logprobs_batch(
-                    model=model,
-                    data_loader=data_loader,
-                    forward_backward_func=forward_backward_func,
-                    packing_context=packing_context,
-                    trajs_batch_size=len(compute_trajs),
-                    seq_length=args.seq_length,
-                    logprobs_batch_size=logprobs_batch_size,
-                    decoder_seq_length=args.decoder_seq_length,
-                    dtype=dtype,
-                    pp_group=pp_group,
-                    is_correction=args.rl_inference_logprobs_is_correction,
-                )
+                old_logprobs = _compute_logprobs_batch()
 
             with torch.no_grad(), nvtx_range("compute_ref_logprobs", time=True):
                 # We need to load the ref model state dict and compute the logprobs for the ref model
@@ -1216,19 +1203,8 @@ def prepare_data_for_update(
                     k: (v.cpu() if v is not None else v) for k, v in model.state_dict().items()
                 }
                 model.load_state_dict(ref_state_dict)
-                ref_logprobs = _compute_logprobs_batch(
-                    model=model,
-                    data_loader=data_loader,
-                    forward_backward_func=forward_backward_func,
-                    packing_context=packing_context,
-                    trajs_batch_size=len(compute_trajs),
-                    seq_length=args.seq_length,
-                    logprobs_batch_size=logprobs_batch_size,
-                    decoder_seq_length=args.decoder_seq_length,
-                    dtype=dtype,
-                    pp_group=pp_group,
-                    is_correction=args.rl_inference_logprobs_is_correction,
-                )
+
+                ref_logprobs = _compute_logprobs_batch()
 
                 # logprobs are [b, seq, h] now.
                 model.load_state_dict(cur_st_dict)
@@ -1238,7 +1214,7 @@ def prepare_data_for_update(
             torch.cuda.empty_cache()
 
 
-        if sequence_packing:
+        if args.rl_use_sequence_packing:
             with nvtx_range("pack_logprobs", time=True):
                 # Store logprobs on gpu in packing context
                 # Since PackingContext is a dataclass, we add these as new attributes
@@ -1267,22 +1243,6 @@ def prepare_data_for_update(
                     packing_context.packed_inference_logprobs = packed_inference_logprobs.cuda()
                     # Only mark as having inference logprobs for IS correction if enabled
                     packing_context.has_inference_logprobs = args.rl_inference_logprobs_is_correction
-            with nvtx_range("create_dataloader"):
-                # @vitalyk: This function also reconfigures the data loader to count the
-                # global_batch_size in the bins frame of reference.
-                # I think it will be a better design if we split the data loader creating and logic
-                # that reconfigures the microbatch calculator.
-
-                update_microbatch_calculator(
-                    samples_ratio_per_step=samples_ratio_per_step,
-                    num_bins_this_rank = len(packing_context.packed_trajs),
-                    bin_seq_indices = packing_context.packing_info.bin_seq_indices,
-                    global_batch_size=args.global_batch_size, 
-                    rampup_batch_size=args.rampup_batch_size, 
-                    micro_batch_size=args.micro_batch_size, 
-                    decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
-               )
-                loader = get_microbatch_dataloader(len(packing_context.packed_trajs), args.micro_batch_size)
         else:
             with nvtx_range("align_inference_logprobs", time=True):
                 if inference_logprobs is not None:
@@ -1297,20 +1257,14 @@ def prepare_data_for_update(
                     # Nullify logprobs if not used in IS correction,
                     if not args.rl_inference_logprobs_is_correction:
                         inference_logprobs = None
-            with nvtx_range("create_dataloader"):
-                # Because of multiturn, our batch sizes for non-sequence packed trajectories are not fixed anymore.
-                # As in sequence packing above, we need to reconfigure it too.
-                runtime_state.packing_context = None
-
-                reconfigure_num_microbatches_calculator(
-                    rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0,
-                    global_batch_size=math.ceil(samples_ratio_per_step*total_turns_sampled), 
-                    rampup_batch_size=args.rampup_batch_size, 
-                    micro_batch_size=args.micro_batch_size, 
-                    decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
-                    data_parallel_size=mpu.get_data_parallel_world_size(),
-                )
 
+        with nvtx_range("create_dataloader"):
+            if args.rl_use_sequence_packing:
+               loader, optimizer_steps = get_microbatch_dataloader(packing_context)
+               runtime_state.global_batches_per_collection = optimizer_steps
+            else:
+                runtime_state.packing_context = None
+                runtime_state.global_batches_per_collection = global_rollout_count / args.global_batch_size
                 dataset_tensors = [
                     compute_trajs,
                     advantages,
@@ -1319,20 +1273,20 @@ def prepare_data_for_update(
                     original_position_ids,
                     ref_logprobs,
                 ]
-                if is_correction and inference_logprobs is not None:
+                if args.rl_inference_logprobs_is_correction and inference_logprobs is not None:
                     dataset_tensors.append(inference_logprobs)
                 else:
                     dataset_tensors.append(torch.zeros_like(old_logprobs))
+
                 data = TensorDataset(*dataset_tensors)
                 loader = DataLoader(data, batch_size=args.micro_batch_size)
 
-
         with nvtx_range("log-wandb-tb"):
             maybe_log_training_metrics(
                 group_stats=group_stats,
                 current_iteration=args.curr_iteration,
                 tokenizer=tokenizer,
-                example_group=example_group,
+                example_group=rollouts[0],
                 wandb_writer=wandb_writer,
                 tb_writer=tb_writer,
             )
@@ -1340,65 +1294,66 @@ def prepare_data_for_update(
     return RerunDataIterator(itertools.cycle(loader))
 
 
-def get_grpo_data_iterator(
+def get_rollout_data_iterator(
+    model: LanguageModule,
+    inference_model: LanguageModule | None,
+    optimizer: MegatronOptimizer,
+    iteration: int,
+    ref_state_dict: Dict[str, torch.Tensor],
+) -> RerunDataIterator:
+
+    args = get_args()
+    tokenizer = get_tokenizer()
+
+    buffered_rollouts = get_environment_rollouts(
+        model, inference_model, optimizer, args.grpo_prompts_per_step, args.grpo_group_size
+    )
+    buffered_rollouts = prepare_data_for_update(model, ref_state_dict, buffered_rollouts, tokenizer)
+
+    return buffered_rollouts
+
+
+def setup_grpo_data_iterator(
     model: LanguageModule,
     inference_model: LanguageModule | None,
     optimizer: MegatronOptimizer,
     iteration: int,
     ref_state_dict: Dict[str, torch.Tensor],
-    grpo_iterations: int,
-    grpo_prompts_per_step: int,
-    grpo_group_size: int,
-    global_batch_size: int,
-    sequence_packing: bool,
-    is_correction: bool,
     buffered_rollouts: RerunDataIterator | None = None,
 ) -> RerunDataIterator:
     """
-    Get the data iterator for GRPO training.
-
-    Depending on the sampling parameters either performs data collections or returns
-    the buffered_rollouts as is.
+    Set up the data iterator for GRPO training.
 
     Args:
         model: The language model
         optimizer: The Megatron optimizer
         iteration: Current training iteration
         ref_state_dict: Reference model state dict for GRPO
-        grpo_iterations: How many steps we reuse the sampled data for.
-        grpo_prompts_per_step: How many prompts we sample per data collection.
-        grpo_group_size: How many samples we do per prompt.
-        global_batch_size: Global batch size.
-        sequence_packing: Use sequence packing if True.
-        is_correction: Use IS correction if True.
         buffered_rollouts: Previously collected rollouts (if any)
 
     Returns:
         RerunDataIterator for the current training step
     """
+    args = get_args()
     runtime_state = get_rl_runtime_state()
 
+    if inference_model is not None:
+        inference_pg_collection = unwrap_model(inference_model[0]).pg_collection
+    else:
+        inference_pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+
     # We collect new rollouts when we've gone over the collected data 'grpo_iterations' times.
-    global_batches_per_collection = (grpo_prompts_per_step * grpo_group_size) // global_batch_size 
     if (
         buffered_rollouts is None or
         iteration == runtime_state.last_collection_iteration + 
-        (grpo_iterations * global_batches_per_collection)
+        (args.grpo_iterations * runtime_state.global_batches_per_collection)
     ):
-
-        buffered_rollouts = get_environment_rollouts(
-            model, inference_model, optimizer, grpo_prompts_per_step, grpo_group_size
-        )
-        buffered_rollouts = prepare_data_for_update(model=model, 
-            ref_state_dict=ref_state_dict, 
-            rollouts=buffered_rollouts,
-            tokenizer=get_tokenizer(),
-            sequence_packing=sequence_packing,
-            is_correction=is_correction,
-            )
+        train_data_iterator = get_rollout_data_iterator(model,inference_model, optimizer, iteration, ref_state_dict)
         runtime_state.reset_iteration_counters(iteration)
+    else:
+        train_data_iterator = buffered_rollouts
 
-    return buffered_rollouts
+    return train_data_iterator
 
 
 def evaluate_and_print_results_rl(
@@ -1438,7 +1393,7 @@ def evaluate_and_print_results_rl(
 
             rank = torch.distributed.get_rank()
             if rank == 0:
-                logger.info("Collecting evaluation results...")
+                logger.info(f"Collecting evaluation results...")
                 agent = get_agent(args)
                 request = EvaluationRequest(
                     inference_interface=inference_interface,
@@ -1753,32 +1708,3 @@ def get_iteration_sequence_count(args):
     if torch.distributed.is_initialized():
         torch.distributed.all_reduce(sequences_tensor, group=mpu.get_data_parallel_group())
     return int(sequences_tensor.item())
-    
-def _pad_nonnull_with_zeros(data: list[Optional[torch.Tensor]], max_len: int) -> torch.Tensor:
-    """Pad each element of a list of tensors to the length required.
-    Args:
-        data: List of tensors to pad.
-        max_len: Maximum length to pad to. Must be higher or equal than the max len of the data tensors.
-    Returns:
-        A padded tensor which is a stacked list of padded input tensors.
-
-    """
-    if all([el is None for el in data]):
-        raise ValueError("At least one element of the data list should be not None.")
-    padded_data = []
-    for chunk in data:
-        if chunk is not None:
-            padding_size = max_len - len(chunk)
-            if padding_size > 0:
-                # Pad with zeros (these positions will be masked anyway)
-                padded = torch.nn.functional.pad(chunk, (0, padding_size), value=0.0)
-                padded_data.append(padded)
-            elif padding_size == 0:
-                padded_data.append(chunk)
-            else:
-                raise ValueError("One of the input tensors has larger length than padding max len.")
-        else:
-            # Create zero tensor for None logprobs
-            padded_data.append(torch.zeros(max_len))
-    return torch.stack(padded_data)
-
diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py
index 4d983764f77..a5703a4580c 100644
--- a/megatron/rl/sequence_packing_utils.py
+++ b/megatron/rl/sequence_packing_utils.py
@@ -10,6 +10,7 @@
 from megatron.training.global_vars import get_args, get_tokenizer
 from megatron.training.utils import get_nvtx_range
 from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core import mpu
 import logging
 import typing
@@ -77,6 +78,7 @@ def load_packed_data_by_index(bin_idx: int, packing_context: PackingContext, log
     Args:
         bin_idx: Index of the bin to load.
     """
+    args = get_args()
     # Get packing context (should always be available in packed mode)
     idx = slice(bin_idx, bin_idx + 1)
 
@@ -154,8 +156,9 @@ def log_packing_efficiency(packing_context: PackingContext):
     packing_efficiency = my_tokens / total_capacity if total_capacity > 0 else 0
     avg_seq_length = total_tokens / len(packing_info.seq_lengths)
     rank = mpu.get_data_parallel_rank()
+    data_parallel_world_size = mpu.get_data_parallel_world_size()
 
-    log_single_rank(logger, logging.INFO, "[Sequence Packing] Statistics:")
+    log_single_rank(logger, logging.INFO, f"[Sequence Packing] Statistics:")
     log_single_rank(
         logger,
         logging.INFO,
@@ -266,7 +269,7 @@ def log_packing_efficiency(packing_context: PackingContext):
             log_single_rank(
                 logger,
                 logging.INFO,
-                "[Sequence Packing]  Round-robin distribution quality:",
+                f"[Sequence Packing]  Round-robin distribution quality:",
             )
             log_single_rank(
                 logger,
@@ -395,7 +398,7 @@ def create_empty_bins(
         empty_packing_info_entries,
     )
 
-def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, device: torch.device) -> PackedSeqParams:
+def get_default_packed_seq_params(seq_length: int, device: torch.device) -> PackedSeqParams:
     """Create a default PackedSeqParams that acts as no-op for a single sequence.
 
     This ensures CUDA graph signature consistency when packed_seq_params
@@ -404,7 +407,6 @@ def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, d
 
     Args:
         seq_length: The sequence length 
-        max_sequences_per_bin: Max sequences to pack in a bin.
         device: Device to create tensors on.
 
     Returns:
@@ -414,7 +416,7 @@ def get_default_packed_seq_params(seq_length: int, max_sequences_per_bin: int, d
     args = get_args()
 
     # Pad to the maximum number of sequences in the bin for the attention kernel.
-    cu_seqlens = torch.full((max_sequences_per_bin,), seq_length, dtype=torch.int32, device=device)
+    cu_seqlens = torch.full((args.rl_sequence_packing_max_sequences_per_bin,), seq_length, dtype=torch.int32, device=device)
     cu_seqlens[0] = 0
 
     return PackedSeqParams(
@@ -772,7 +774,7 @@ def pack_sequences(
 
         seq_per_bin = [len(indices) for indices in packing_info.bin_seq_indices]
         log_single_rank(
-            logger, logging.DEBUG, ("Initial packing output (before distribution):")
+            logger, logging.DEBUG, (f"Initial packing output (before distribution):")
         )
         log_single_rank(
             logger,
@@ -967,20 +969,33 @@ def distribute_packed_bins(
 def pack_all_trajectories(trajs, generation_masks, inference_logprobs, global_advantages, bin_size, max_sequences_per_bin, packing_algo):
     tokenizer = get_tokenizer()
     data_parallel_world_size = mpu.get_data_parallel_world_size()
-    data_parallel_group = mpu.get_data_parallel_group()
     nvtx_range = get_nvtx_range()
 
     with nvtx_range("regather_trajectories", time=True):
-        def _gather(data):
-            data = data.cuda()
-            data_list = [torch.empty_like(data) for _ in range(data_parallel_world_size)]
-            torch.distributed.all_gather(data_list, data, group=data_parallel_group)
-            return torch.cat(data_list, dim=0)
-
-        trajs = _gather(trajs)    
-        generation_masks = _gather(generation_masks) 
+        # Regather trajectories from all ranks for packing
+        trajs = trajs.cuda()
+        trajs_list = [torch.empty_like(trajs) for _ in range(data_parallel_world_size)]
+        torch.distributed.all_gather(
+            trajs_list, trajs, group=mpu.get_data_parallel_group()
+        )
+        trajs = torch.cat(trajs_list, dim=0)
+
+        # Gather all generation masks
+        generation_masks = generation_masks.cuda()
+        masks_list = [torch.empty_like(generation_masks) for _ in range(data_parallel_world_size)]
+        torch.distributed.all_gather(
+            masks_list, generation_masks, group=mpu.get_data_parallel_group()
+        )
+        generation_masks = torch.cat(masks_list, dim=0)
+
+        # Gather inference logprobs if present
         if inference_logprobs is not None:
-            inference_logprobs = _gather(inference_logprobs)
+            inference_logprobs = inference_logprobs.cuda()
+            logprobs_list = [torch.empty_like(inference_logprobs) for _ in range(data_parallel_world_size)]
+            torch.distributed.all_gather(
+                logprobs_list, inference_logprobs, group=mpu.get_data_parallel_group()
+            )
+            inference_logprobs = torch.cat(logprobs_list, dim=0)
 
     with nvtx_range("pack_sequences", time=True):
         # Create packer with max sequences per bin limit to prevent extreme imbalance
@@ -1058,63 +1073,53 @@ def _gather(data):
 
     return packing_context
 
-def update_microbatch_calculator(
-    samples_ratio_per_step: float,
-    num_bins_this_rank: int,
-    bin_seq_indices: List[List[int]],
-    global_batch_size: int, 
-    rampup_batch_size: int, 
-    micro_batch_size: int, 
-    decrease_batch_size_if_needed: bool,
-):
-    """Return a data loader with seqpacked indices with microbatches in bins frame of reference.
-    Args:
-        samples_ratio_per_step: Fraction of sampled trajectories to use per iteration.
-        num_bins_this_rank: Amount of packing bins that belongs to current rank.
-        bin_seq_indices: Global seq indices in the bin, see PackingInfo.
-        global_batch_size: Current global batch size.
-        rampup_batch_size: Rampup batch size. See num_microbatches_calculator.py for more.
-        micro_batch_size: Micro batch size at init.
-        decrease_batch_size_if_needed: Scale down batch size. See num_microbatches_calculator.py for more.
-
-    As a side effect, we calculate the global batch size in the bins frame of reference.
-    In sequence packing, our batch dimension shrinks as we move some trajs onto free
-    space in sequence dimension. The resulting batch size is what we return here.
-    """
 
+def get_microbatch_dataloader(packing_context: PackingContext) -> Tuple[DataLoader, int]:
+    args = get_args()
+    num_bins_this_rank = len(packing_context.packed_trajs)
     dp_world_size = mpu.get_data_parallel_world_size()
 
+    # Ratio of collected sequences to the global batch size
+    pct_of_sequences_per_batch = len(packing_context.packing_info.seq_lengths) / args.global_batch_size
+
     # Ceiling division means we will reuse some bins
     # If we did floor we would leave some behind
-    local_bins_per_step = math.ceil(samples_ratio_per_step * num_bins_this_rank)
+    local_bins_per_step = math.ceil(pct_of_sequences_per_batch * num_bins_this_rank)
+    effective_global_batch_size = local_bins_per_step * dp_world_size
 
-    bins_bs = local_bins_per_step * dp_world_size
+    # Store packing plan in runtime state for the training loop to use
+    optimizer_steps = -(-num_bins_this_rank // local_bins_per_step)
 
     old_num_microbatches = get_num_microbatches()
+
     reconfigure_num_microbatches_calculator(
         rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0,
-        rampup_batch_size=rampup_batch_size,
-        global_batch_size=bins_bs,
-        micro_batch_size=micro_batch_size,
+        rampup_batch_size=args.rampup_batch_size,
+        global_batch_size=effective_global_batch_size,
+        micro_batch_size=args.micro_batch_size,
         data_parallel_size=dp_world_size,
-        decrease_batch_size_if_needed=decrease_batch_size_if_needed,
+        decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
     )
+
     new_num_microbatches = get_num_microbatches()
 
     log_single_rank(
-        logger, logging.INFO, "[Sequence Packing] Multi-step training plan:"
+        logger, logging.INFO, f"[Sequence Packing] Multi-step training plan:"
     )
-
     log_single_rank(
         logger,
         logging.INFO,
-        f"[Sequence Packing]  - Bins per rank per step: {samples_ratio_per_step}*{num_bins_this_rank}={local_bins_per_step}",
+        f"[Sequence Packing]  - Target sequences per step: {args.global_batch_size}",
+    )
+    log_single_rank(
+        logger,
+        logging.INFO,
+        f"[Sequence Packing]  - Bins per rank per step: {pct_of_sequences_per_batch}*{num_bins_this_rank}={local_bins_per_step}",
     )
-
     log_single_rank(
         logger,
         logging.INFO,
-        f"[Sequence Packing]  - Target sequences per step: {global_batch_size}",
+        f"[Sequence Packing]  - Total optimizer steps: {optimizer_steps}",
     )
     log_single_rank(
         logger,
@@ -1122,10 +1127,8 @@ def update_microbatch_calculator(
         f"[Sequence Packing]  - Microbatches per step: {new_num_microbatches} (was {old_num_microbatches})",
     )
 
-    # Opt steps only depends on how much we sample and how much we consume.
-    # We make sure this is an integer division, check validate_args in arguments.py for details.
-    opt_steps = int(1 / samples_ratio_per_step)
-    for step in range(min(3, opt_steps)):
+    bin_seq_indices = packing_context.packing_info.bin_seq_indices
+    for step in range(min(3, optimizer_steps)):
         start_idx = step * local_bins_per_step
         end_idx = min(start_idx + local_bins_per_step, num_bins_this_rank)
         step_bins = end_idx - start_idx
@@ -1142,13 +1145,22 @@ def update_microbatch_calculator(
             f"[Sequence Packing]  - Step {step + 1}: {step_bins} bins, ~{est_global_seqs} sequences globally",
         )
 
-    if opt_steps > 3:
-        log_single_rank(logger, logging.INFO, f"  - ... ({opt_steps - 3} more steps)")
+    if optimizer_steps > 3:
+        log_single_rank(logger, logging.INFO, f"  - ... ({optimizer_steps - 3} more steps)")
 
-def get_microbatch_dataloader(num_bins_this_rank, micro_batch_size):
     bin_indices = torch.arange(num_bins_this_rank)
     dataset = TensorDataset(bin_indices)
-    return DataLoader(dataset, batch_size=micro_batch_size, shuffle=False, collate_fn=lambda x: x[0])
+    loader = DataLoader(dataset, batch_size=args.micro_batch_size, shuffle=False, collate_fn=lambda x: x[0], drop_last=True)
+    return loader, optimizer_steps
+
+def update_sequence_packing_metrics(args):
+    """Update bin tracking for sequence packing mode."""
+    if args.rl_use_sequence_packing:
+        bin_count = (
+            mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
+        )
+        args.consumed_train_bins += bin_count
+
 
 def get_sequence_packing_log_info(args):
     """Get logging information for sequence packing mode."""
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 4ed4eb1cb73..f269ad02879 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1987,6 +1987,8 @@ def _add_rl_args(parser):
                        help='If set, use inference logprobs in importance sampling correction of the loss.')
     group.add_argument('--rl-importance-sampling-truncation-coef', type=float, default=None,
                        help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.")
+    group.add_argument('--rl-calculate-intra-group-similarity', action=argparse.BooleanOptionalAction, default=False,
+                       help='If set, calculate the intra-group similarity of rollouts.')
     group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False,
                        help='Enable sequence packing')
     group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50,
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 87d9fe8b841..500d30b9e73 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -2766,15 +2766,8 @@ def get_e2e_base_metrics():
 
         if getattr(args, 'perform_rl_step', False):
             with torch.no_grad():
-                train_data_iterator = rl_utils.get_grpo_data_iterator(
-                    model, inference_model, optimizer, iteration, ref_state_dict,
-                    grpo_iterations=args.grpo_iterations,
-                    grpo_prompts_per_step=args.grpo_prompts_per_step,
-                    grpo_group_size=args.grpo_group_size,
-                    global_batch_size=args.global_batch_size,
-                    sequence_packing=args.rl_use_sequence_packing,
-                    buffered_rollouts=buffered_rollouts,
-                    is_correction=args.rl_inference_logprobs_is_correction,
+                train_data_iterator = rl_utils.setup_grpo_data_iterator(
+                    model, inference_model, optimizer, iteration, ref_state_dict, buffered_rollouts
                 )
                 # Buffered rollouts are used as a state container for setups when
                 # we use previously-generated data for an update.
@@ -2853,10 +2846,7 @@ def get_e2e_base_metrics():
         if getattr(args, 'perform_rl_step', False) and args.rl_use_sequence_packing:
             iteration_sequences = rl_utils.get_iteration_sequence_count(args)
             # Track bins separately for packed mode
-            bin_count = (
-                mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
-            )
-            args.consumed_train_bins += bin_count
+            rl_utils.update_sequence_packing_metrics(args)
         else:
             batch_size = (
                 mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py
index cff62d40f0e..8747f2e8c35 100644
--- a/tests/unit_tests/rl/test_rl_utils.py
+++ b/tests/unit_tests/rl/test_rl_utils.py
@@ -79,8 +79,6 @@ def initialize_model_parallel(request, monkeypatch):
     Skips if world_size < tp * pp.
     """
     monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
-    monkeypatch.setenv("WANDB_MODE", "disabled")
-    monkeypatch.setenv("LOG_TO_WANDB", "false")
 
     tp, pp = request.param
     world_size = Utils.world_size
@@ -114,7 +112,6 @@ def create_test_args(self, **kwargs):
         args.hidden_size = 128
         args.max_position_embeddings = 256
         args.seq_length = 256
-        args.wandb_project = None
 
         args.micro_batch_size = 1
 
@@ -277,68 +274,58 @@ def test_grpo_loss_truncation(self):
     def test_prepare_data_for_update(self, initialize_model_parallel):
         """Test that getting logprobs at least does not crash."""
         world_size, dp, tp, pp = initialize_model_parallel
-        # Here I assume that we will be consuming all data in one step.
-        group_size = 2
         self.create_test_args(
             micro_batch_size=2,
             seq_length=4,
             curr_iteration=1,
             tensor_model_parallel_size=tp,
             pipeline_model_parallel_size=pp,
-            global_batch_size=dp * 2,
-            grpo_prompts_per_step=dp,
-            grpo_group_size=group_size,
         )
 
         model = MockModel()
         tokenizer = MockTokenizer()
 
         r1 = TokenRollout(
-            trajectory=[[1, 2, 3]],
+            trajectory=[1, 2, 3],
             reward=3.14,
-            generation_mask=[[False, True, True]],
-            logprobs=[[0.1, 0.2, 0.3]],
+            generation_mask=[False, True, True],
+            logprobs=[0.1, 0.2, 0.3],
             env_id='MEGAENV',
             problem_id="2",
         )
         r2 = TokenRollout(
-            trajectory=[[1, 2, 3, 4]],
+            trajectory=[1, 2, 3, 4],
             reward=0.14,
-            generation_mask=[[False, True, True, True]],
-            logprobs=[[0.1, 0.2, 0.3, -1.2]],
+            generation_mask=[False, True, True, True],
+            logprobs=[0.1, 0.2, 0.3, -1.2],
             env_id='MEGAENV',
             problem_id="2",
         )
-
         rollouts = [[r1, r2] for _ in range(dp)]
         try:
-            rl_utils.prepare_data_for_update(
-                [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False
-            )
+            rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
         except AssertionError as e:
             # We expect trajectories to come padded there.
             assert str(e).startswith('Rollout is not the correct length')
 
         r1 = TokenRollout(
-            trajectory=torch.tensor([[1, 2, 3, tokenizer.eod]], dtype=torch.float).cuda(),
+            trajectory=torch.tensor([1, 2, 3, tokenizer.eod], dtype=torch.float).cuda(),
             reward=3.14,
-            generation_mask=torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(),
-            logprobs=torch.tensor([[-0.2, -0.3, -3.2]]).cuda(),
+            generation_mask=torch.tensor([False, True, True, True], dtype=torch.float).cuda(),
+            logprobs=torch.tensor([-0.2, -0.3, -3.2]).cuda(),
             env_id='MEGAENV',
             problem_id="2",
         )
         r2 = TokenRollout(
-            trajectory=torch.tensor([[1, 2, 234, tokenizer.eod]], dtype=torch.float).cuda(),
+            trajectory=torch.tensor([1, 2, 234, tokenizer.eod], dtype=torch.float).cuda(),
             reward=0.14,
-            generation_mask=torch.tensor([[False, True, True, True]], dtype=torch.float).cuda(),
-            logprobs=torch.tensor([[-0.2, -0.3, -1.2]]),
+            generation_mask=torch.tensor([False, True, True, True], dtype=torch.float).cuda(),
+            logprobs=torch.tensor([-0.2, -0.3, -1.2]),
             env_id='MEGAENV',
             problem_id="2",
         )
         rollouts = [[r1, r2] for _ in range(dp)]
-        data_iter = rl_utils.prepare_data_for_update(
-            [model], {}, rollouts, tokenizer, sequence_packing=False, is_correction=False
-        )
+        data_iter = rl_utils.prepare_data_for_update([model], {}, rollouts, tokenizer)
 
         _, _, old_logprobs, _, _, _, _ = next(data_iter)
         # All logits are ones in the MockModel.
@@ -346,8 +333,7 @@ def test_prepare_data_for_update(self, initialize_model_parallel):
         torch.testing.assert_close(old_logprobs.exp(), torch.ones_like(old_logprobs) / VOCAB)
 
     @pytest.mark.parametrize("use_sequence_packing", [True, False])
-    @pytest.mark.parametrize("num_turns", [1, 2])
-    def test_prepare_trajectories(self, use_sequence_packing, num_turns):
+    def test_prepare_trajectories(self, use_sequence_packing):
         """Test that rollouts are properly prepared for training."""
         seq_length = 8
         self.create_test_args(
@@ -361,38 +347,34 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns):
 
         # Create rollouts of varying lengths
         r1 = TokenRollout(
-            trajectory=[[1, 2, 3, tokenizer.eod]] * num_turns,
+            trajectory=[1, 2, 3, tokenizer.eod],
             reward=3.14,
-            generation_mask=[[False, True, True, True]] * num_turns,
-            logprobs=[[0.1, 0.2, 0.3, 0.35]] * num_turns,
+            generation_mask=[False, True, True, True],
+            logprobs=[0.1, 0.2, 0.3, 0.35],
             env_id='MEGAENV',
             problem_id="1",
         )
         r2 = TokenRollout(
-            trajectory=[[4, 5, 6, 7, tokenizer.eod]] * num_turns,
+            trajectory=[4, 5, 6, 7, tokenizer.eod],
             reward=0.14,
-            generation_mask=[[False, True, True, True, True]] * num_turns,
-            logprobs=[[0.4, 0.5, 0.6, 0.7, 0.75]] * num_turns,
+            generation_mask=[False, True, True, True, True],
+            logprobs=[0.4, 0.5, 0.6, 0.7, 0.75],
             env_id='MEGAENV',
             problem_id="2",
         )
         r3 = TokenRollout(
-            trajectory=[[8, 9, tokenizer.eod]] * num_turns,
+            trajectory=[8, 9, tokenizer.eod],
             reward=2.71,
-            generation_mask=[[False, True, True]] * num_turns,
-            logprobs=[[0.8, 0.9, 0.95]] * num_turns,
+            generation_mask=[False, True, True],
+            logprobs=[0.8, 0.9, 0.95],
             env_id='MEGAENV',
             problem_id="3",
         )
 
-        rollouts = [r1, r2, r3]
+        rollouts = [[r1, r2, r3]]
 
         trajs, genmask, inference_logprobs = rl_utils.prepare_trajectories(
-            rollouts,
-            tokenizer,
-            seq_length,
-            sequence_packing=use_sequence_packing,
-            skip_bos_token=False,
+            rollouts, tokenizer, seq_length
         )
 
         expected_trajs = torch.tensor(
@@ -403,7 +385,7 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns):
             ],
             dtype=torch.long,
             device=trajs.device,
-        ).repeat_interleave(num_turns, dim=0)
+        )
         assert torch.equal(trajs, expected_trajs)
 
         expected_genmask = torch.tensor(
@@ -414,7 +396,7 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns):
             ],
             dtype=torch.bool,
             device=genmask.device,
-        ).repeat_interleave(num_turns, dim=0)
+        )
         assert torch.equal(genmask, expected_genmask)
 
         if use_sequence_packing:
@@ -426,7 +408,7 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns):
                 ],
                 dtype=torch.float32,
                 device=inference_logprobs.device,
-            ).repeat_interleave(num_turns, dim=0)
+            )
             torch.testing.assert_close(inference_logprobs, expected_logprobs, rtol=0, atol=0)
         else:
             expected_logprobs = [
@@ -434,57 +416,12 @@ def test_prepare_trajectories(self, use_sequence_packing, num_turns):
                 [0.4, 0.5, 0.6, 0.7, 0.75],
                 [0.8, 0.9, 0.95],
             ]
-            expected_logprobs = [el for el in expected_logprobs for _ in range(num_turns)]
             assert len(inference_logprobs) == len(expected_logprobs)
             for got, exp in zip(inference_logprobs, expected_logprobs):
                 got_t = got if torch.is_tensor(got) else torch.tensor(got, dtype=torch.float32)
                 exp_t = torch.tensor(exp, dtype=torch.float32, device=got_t.device)
                 torch.testing.assert_close(got_t, exp_t, rtol=0, atol=0)
 
-    def test_single_turn_advantage_calculation(self):
-        rewards = [[-1, 1], [4, 4]]
-        num_turns = [[1, 1], [1, 1]]
-        advs = rl_utils.calculate_grpo_advantages(rewards, num_turns)
-        torch.testing.assert_close(
-            torch.tensor(advs), torch.tensor([-1, 1.0, 0.0, 0.0]), atol=1e-4, rtol=1e-5
-        )
-
-    def test_multi_turn_advantage_calculation(self):
-        rewards = [[-1, 1], [4, 4]]
-        num_turns = [[2, 1], [1, 3]]
-        advs = rl_utils.calculate_grpo_advantages(rewards, num_turns)
-        torch.testing.assert_close(
-            torch.tensor(advs),
-            torch.tensor([-1, -1, 1.0, 0.0, 0.0, 0.0, 0.0]),
-            atol=1e-4,
-            rtol=1e-5,
-        )
-
-    def test_pad_list_of_nones(self):
-        with pytest.raises(ValueError) as e_info:
-            rl_utils._pad_nonnull_with_zeros([None] * 3, 42)
-        assert "At least one" in str(e_info)
-
-    def test_pad_with_wrong_params(self):
-        with pytest.raises(ValueError) as e_info:
-            rl_utils._pad_nonnull_with_zeros([torch.zeros(5)], 4)
-        assert "larger length" in str(e_info)
-
-    def test_pad_full_size(self):
-        padded = rl_utils._pad_nonnull_with_zeros([torch.zeros(5), torch.zeros(5)], 5)
-        assert padded.shape == (2, 5)
-
-    def test_pad_some_nones(self):
-        padded = rl_utils._pad_nonnull_with_zeros([None, torch.zeros(5)], 5)
-        assert padded.shape == (2, 5)
-        assert (padded[0] == 0).all()
-
-    def test_pad_normal(self):
-        padded = rl_utils._pad_nonnull_with_zeros(
-            [torch.zeros(2), torch.zeros(3), torch.zeros(4)], 5
-        )
-        assert padded.shape == (3, 5)
-
     @pytest.mark.parametrize(
         "initialize_model_parallel",
         [
diff --git a/tests/unit_tests/rl/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py
index 44a3de762f0..548aedf55fd 100644
--- a/tests/unit_tests/rl/test_sequence_packing_utils.py
+++ b/tests/unit_tests/rl/test_sequence_packing_utils.py
@@ -1,8 +1,5 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from unittest.mock import patch
-
-import pytest
 import torch
 
 from megatron.rl import rl_utils, sequence_packing_utils
@@ -410,56 +407,3 @@ def test_compute_packed_inference_logprobs_stats_shape_mismatch():
 
     # Stats should remain None due to shape mismatch
     assert group_stats.mean_piold_to_inf_prob is None
-
-
-@pytest.mark.parametrize(
-    "ratio,local_bins,world,expected_bs",
-    [
-        (1.0, 1, 8, 8),  # no stale data (ratio 1.), everything divides perfectly.
-        (1.0, 42, 8, 42 * 8),  # no stale data (ratio 1.), everything divides perfectly, more bins
-        (
-            0.5,
-            1,
-            8,
-            8,
-        ),  # 0.5 means we use half of all seqs per step, they all fit 1 bin -> we should reuse
-        (1 / 3, 4, 8, 16),  # third of the data per step, nonint division
-    ],
-)
-def test_get_bins_bs_and_steps(ratio, local_bins, world, expected_bs):
-    # Make a dummy struct to check only the required fields.
-    # Divide by ratio to make sure the samples are divisible by global_bs in the test.
-    n_seqs = int(world * 7 / ratio)
-    global_bs_in_seq = int(n_seqs * ratio)
-
-    def side_eff(
-        rank,
-        rampup_batch_size,
-        global_batch_size,
-        micro_batch_size,
-        data_parallel_size,
-        decrease_batch_size_if_needed,
-    ):
-        # Inside of the get_microbatch_dataloader, we compute the batch size in bins.
-        # We want to test this variable.
-        global actual_bs
-        actual_bs = global_batch_size
-
-    with patch('megatron.rl.sequence_packing_utils.get_num_microbatches', return_value=1):
-        with patch(
-            'megatron.rl.sequence_packing_utils.reconfigure_num_microbatches_calculator',
-            side_effect=side_eff,
-        ):
-            with patch('megatron.core.mpu.get_data_parallel_world_size', return_value=world):
-                sequence_packing_utils.update_microbatch_calculator(
-                    samples_ratio_per_step=ratio,
-                    num_bins_this_rank=local_bins,
-                    bin_seq_indices=[],
-                    global_batch_size=global_bs_in_seq,
-                    rampup_batch_size=1,
-                    micro_batch_size=1,
-                    decrease_batch_size_if_needed=False,
-                )
-
-    # Iterator is local, batch size is global
-    assert expected_bs == actual_bs
diff --git a/train_rl.py b/train_rl.py
index cfc010b3c04..299843bcff3 100644
--- a/train_rl.py
+++ b/train_rl.py
@@ -260,7 +260,6 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False):
     if packed_seq_params is None:
         packed_seq_params = get_default_packed_seq_params(
             seq_length=tokens.shape[1],
-            max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin,
             device=tokens.device,
         )
 

From f58b6d6f3947d6131b5a094f0ba775b760b9cd88 Mon Sep 17 00:00:00 2001
From: Philip Petrakian <ppetrakian@nvidia.com>
Date: Thu, 29 Jan 2026 13:33:18 -0800
Subject: [PATCH 73/79] Fix coderabbit instructions error (#3150)

---
 .coderabbit.yaml | 27 ++-------------------------
 1 file changed, 2 insertions(+), 25 deletions(-)

diff --git a/.coderabbit.yaml b/.coderabbit.yaml
index f20c23e1fe6..160bda5f0f6 100644
--- a/.coderabbit.yaml
+++ b/.coderabbit.yaml
@@ -1,31 +1,8 @@
 # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json
 language: "en-US"
 
-# Instruct the AI to be highly selective - only critical issues
-tone_instructions: |
-  IMPORTANT: Be EXTREMELY restrictive with comments. You should almost never comment.
-
-  ONLY comment if the issue is Critical (🔴) or Major (🟠) severity.
-  NEVER make Minor (🟡), Trivial (🔵), or Info (⚪) severity comments.
-
-  The ONLY acceptable comments are:
-  - Critical bugs: null pointer dereferences, infinite loops, security vulnerabilities, data corruption
-  - Major bugs: logic errors that break functionality, race conditions, resource leaks
-  - Obvious typos in error messages or user-facing strings that would confuse users
-  - Forgotten updates: docstrings that directly contradict the code behavior
-
-  Do NOT comment on (even if you think it's helpful):
-  - Style, formatting, or naming conventions
-  - Refactoring opportunities
-  - Performance suggestions
-  - Missing docstrings or comments
-  - Code organization
-  - Type hints or annotations
-  - Import ordering
-  - Any "minor" improvements
-
-  If you would label a comment as "Minor" severity or lower, DO NOT POST IT.
-  When in doubt, stay silent. Aim for 0-2 comments per PR maximum.
+# Only comment on Critical/Major bugs. No Minor, Trivial, or style comments.
+tone_instructions: "Only comment on Critical or Major bugs. Never comment on Minor issues, style, refactoring, or suggestions. When in doubt, stay silent."
 
 reviews:
   # Use chill profile - filters out nitpicks automatically

From 063624b259f293e736a9c2368941538872d4969b Mon Sep 17 00:00:00 2001
From: Antoni-Joan Solergibert <asolergibert@nvidia.com>
Date: Thu, 29 Jan 2026 22:57:13 +0100
Subject: [PATCH 74/79] Force input ids generated by mock dataset are <
 vocab_size (#2945)

---
 megatron/core/datasets/gpt_dataset.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py
index 3549db88001..d895eac87d3 100644
--- a/megatron/core/datasets/gpt_dataset.py
+++ b/megatron/core/datasets/gpt_dataset.py
@@ -801,10 +801,11 @@ class MockGPTLowLevelDataset:
     """The hard-coded number of samples to generate"""
 
     max_sequence_length: int = 4096
-    """The hard-coded max sequence length to generate"""
+    """The hard-coded max sequence length of the random generated sequences"""
 
     def __init__(self, tokenizer: MegatronTokenizerBase) -> None:
-        self.tokenizer = tokenizer
+        self.vocab_size = tokenizer.vocab_size
+        self.eod_token = tokenizer.eod
         rng = numpy.random.default_rng(seed=self.seed)
         self.sequence_lengths = rng.integers(
             low=1, high=self.max_sequence_length, size=self.size, dtype=numpy.int32
@@ -816,7 +817,7 @@ def __len__(self) -> int:
     def __getitem__(self, idx: int) -> numpy.number:
         length = self.sequence_lengths[idx]
         sample = numpy.int64(
-            numpy.concatenate([numpy.arange(length - 1) + 1, [self.tokenizer.eod]])
+            numpy.concatenate([(numpy.arange(length - 1) + 1) % self.vocab_size, [self.eod_token]])
         )
         return sample
 

From 4652e7b95b95db96b1c796ae42850174ffb68823 Mon Sep 17 00:00:00 2001
From: Antoni-Joan Solergibert <asolergibert@nvidia.com>
Date: Thu, 29 Jan 2026 23:47:43 +0100
Subject: [PATCH 75/79] Add a check to make sure we are distributing all the
 layers when using `--decoder-first-pipeline-num-layers` &
 `--decoder-last-pipeline-num-layers`  (#2947)

---
 megatron/core/transformer/transformer_config.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 633c51e789e..f50bbed0a44 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -1365,6 +1365,13 @@ def __post_init__(self):
                 num_layers -= self.num_layers_in_last_pipeline_stage
                 pipeline_parallel_size -= 1
 
+            # Ensure you either have middle pp stages and layers or none of them.
+            if bool(num_layers) != bool(pipeline_parallel_size):
+                raise ValueError(
+                    f"Mismatch: {num_layers} middle layers remaining but {pipeline_parallel_size} "
+                    f"middle PP stages available."
+                )
+
             # Here pipeline_parallel_size is the number of middle PP stages. If there are middle
             # PP stages, check number of layers at middle stage is divisible by middle PP size.
             if pipeline_parallel_size and not num_layers % pipeline_parallel_size == 0:

From 18deeff12eb69979bae2fc3da19e95ca4cdecdbc Mon Sep 17 00:00:00 2001
From: GitHub Actions <github-actions[bot]@users.noreply.github.com>
Date: Fri, 30 Jan 2026 00:14:09 +0000
Subject: [PATCH 76/79] Update copy-pr-bot.yaml [skip ci]

---
 .github/copy-pr-bot.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml
index 477f87d90f5..f43437d19c0 100644
--- a/.github/copy-pr-bot.yaml
+++ b/.github/copy-pr-bot.yaml
@@ -1,4 +1,4 @@
 enabled: true
 auto_sync_draft: false
 auto_sync_ready: true
-trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "prajwal1210", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]
+trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "kwyss-nvidia", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "prajwal1210", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]

From 67f3515c8a0539d38c23f59dd8b85efdcbb3af3d Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com>
Date: Thu, 29 Jan 2026 15:42:44 -0800
Subject: [PATCH 77/79] Automatically choose available ports in ZMQ (#2278)

---
 .../gpt_dynamic_inference_with_coordinator.py |   9 +-
 .../data_parallel_inference_coordinator.py    |  57 +-
 .../core/inference/engines/dynamic_engine.py  |  51 +-
 megatron/core/inference/inference_client.py   |   8 +-
 megatron/core/inference/utils.py              |   9 +-
 megatron/rl/inference/megatron.py             |   4 +-
 ...est_data_parallel_inference_coordinator.py | 532 +++++++++---------
 7 files changed, 381 insertions(+), 289 deletions(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
index 18191fd38af..cbb7a1aa745 100644
--- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
+++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
@@ -42,7 +42,7 @@
 async def main(
     engine: DynamicInferenceEngine,
     requests: List[Request],
-    port: int,
+    port: int | None = None,
     sampling_params: SamplingParams | None = None,
 ):
     if sampling_params is not None:
@@ -55,8 +55,9 @@ async def main(
     # once you call engine.start_listening_to_data_parallel_coordinator,
     # the engine will start accepting requests from the data parallel coordinator.
     # and processing them in an asyncio coroutine.
+    # leaving inference_coordinator_port as None will find a free port automatically.
     
-    await engine.start_listening_to_data_parallel_coordinator(
+    dp_addr = await engine.start_listening_to_data_parallel_coordinator(
         inference_coordinator_port=port,
         launch_inference_coordinator=True,
     )
@@ -83,7 +84,7 @@ async def main(
 
     # Create client and run example.
     if dist.get_rank() == 0:
-        client = InferenceClient(port)  # submits requests to the inference coordinator
+        client = InferenceClient(dp_addr)  # submits requests to the inference coordinator
         await client.start()
         base_arrival_time = time.time_ns() / 10**9
         for request in requests:
@@ -258,4 +259,4 @@ async def main(
 
         # Stop Nsight profiler.
         if os.environ.get("NSIGHT_PREFIX"):
-            torch.cuda.cudart().cudaProfilerStop()
\ No newline at end of file
+            torch.cuda.cudart().cudaProfilerStop()
diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py
index 3a1747facb2..9a1a11a8c2b 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator.py
@@ -1,11 +1,14 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import errno
 import faulthandler
 import logging
 import signal
+import socket
 from collections import deque
 from itertools import cycle
 from multiprocessing import Event
+from multiprocessing.connection import Connection
 
 import torch
 
@@ -65,7 +68,13 @@ class DataParallelInferenceCoordinator:
         next_request_id (int): A counter for generating unique server-side request IDs.
     """
 
-    def __init__(self, inference_coordinator_port: int, data_parallel_size: int, tokenizer):
+    def __init__(
+        self,
+        pipe_connection: Connection,
+        data_parallel_size: int,
+        tokenizer,
+        inference_coordinator_port: int | None = None,
+    ):
         """
         Initializes the inference coordinator.
 
@@ -74,9 +83,11 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int, tok
         ranks to connect before proceeding.
 
         Args:
-            inference_coordinator_port (int): The TCP port number to bind the server to.
+            pipe_connection (Connection): A connecting pipe to the parent process.
             data_parallel_size (int): The number of TP-coordinator workers that are
                 expected to connect.
+            tokenizer: The tokenizer to use for prompt tokenization and detokenization.
+            inference_coordinator_port (Optional[int]): The TCP port number to bind the server to.
         """
         assert HAVE_ZMQ, (
             "please install the pyzmq library to use DataParallelInferenceCoordinator\n"
@@ -86,6 +97,8 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int, tok
             "please install the messagepack library to use DataParallelInferenceCoordinator\n"
             "pip install msgpack"
         )
+        self.pipe_connection = pipe_connection
+        self.data_parallel_size = data_parallel_size
         self.context = zmq.Context()
 
         # This is the central router socket
@@ -95,9 +108,33 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int, tok
         # 3. data parallel ranks return completed requests to this socket. We route them back to
         #    the user that had submitted the request originally.
 
+        # Get local IP.
+        local_ip = socket.gethostname()
+
         self.router_socket = self.context.socket(zmq.ROUTER)
-        self.router_socket.bind(f"tcp://0.0.0.0:{inference_coordinator_port}")
-        self.data_parallel_size = data_parallel_size
+        is_bound = False
+        if inference_coordinator_port is not None:
+            try:
+                self.router_socket.bind(f"tcp://{local_ip}:{inference_coordinator_port}")
+                is_bound = True
+            except zmq.error.ZMQError as e:
+                if e.errno == errno.EADDRINUSE:
+                    logging.warning(
+                        f"Port {inference_coordinator_port} is already in use. "
+                        "Binding to a random available port instead."
+                    )
+            except Exception:
+                logging.warning(
+                    f"Unknown error when binding to port {inference_coordinator_port}. "
+                    "Attempting to bind to a random available port instead."
+                )
+        if not is_bound:
+            self.router_socket.bind_to_random_port(f"tcp://{local_ip}")
+        self.addr = self.router_socket.getsockopt_string(zmq.LAST_ENDPOINT)
+
+        # Send the address to the parent process.
+        self.pipe_connection.send(self.addr)
+        self.pipe_connection.close()
 
         logging.info("Inference Coordinator: waiting for connections from data parallel ranks...")
         # First wait for all data parallel ranks to establish connections.
@@ -300,7 +337,12 @@ def detokenize(self, finished_request_record):
 
     @classmethod
     def entrypoint(
-        cls, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int, tokenizer
+        cls,
+        pipe_connection: Connection,
+        ready_event: Event,
+        data_parallel_size: int,
+        tokenizer,
+        inference_coordinator_port: int | None = None,
     ):
         """
         Class method to instantiate and run the coordinator, for use in a separate process.
@@ -309,12 +351,15 @@ def entrypoint(
         that it is fully initialized and listening, and then starts the main event loop.
 
         Args:
+            pipe_connection (Connection): A connecting pipe to the parent process.
             ready_event (Event): A threading or multiprocessing event object that is set()
                 once the coordinator is ready to accept connections.
             inference_coordinator_port (int): The port to bind to.
             data_parallel_size (int): The number of expected TP-coordinators.
         """
-        coordinator = cls(inference_coordinator_port, data_parallel_size, tokenizer)
+        coordinator = cls(
+            pipe_connection, data_parallel_size, tokenizer, inference_coordinator_port
+        )
         ready_event.set()
         try:
             coordinator.start()
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index e23520cf65f..0a95e8f4a53 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -4,7 +4,6 @@
 import concurrent.futures
 import logging
 import multiprocessing
-import os
 import socket
 import struct
 import time
@@ -39,7 +38,7 @@
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
-from megatron.core.inference.utils import Counter, await_process_event
+from megatron.core.inference.utils import Counter, await_process_call
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
 from megatron.core.utils import (
@@ -349,7 +348,7 @@ def create_cuda_graphs(self, reset_context: bool = True):
     @internal_api
     async def start_listening_to_data_parallel_coordinator(
         self,
-        inference_coordinator_port: int,
+        inference_coordinator_port: int | None = None,
         launch_inference_coordinator: bool = True,
         *,
         loop: Optional[asyncio.AbstractEventLoop] = None,
@@ -379,11 +378,20 @@ async def start_listening_to_data_parallel_coordinator(
         (`self.run_engine`) as a background asyncio task.
 
         Args:
-            inference_coordinator_port (int): The network port where the central
+            inference_coordinator_port (int | None): The network port where the central
                 `InferenceCoordinator` is or will be listening.
+                If None, a random available port will be selected.
+                If not None, the coordinator will attempt to bind to this port, but should it
+                not succeed (e.g., if the port is already in use), it may bind to a different port.
+                The actual port used is returned by this method.
             launch_inference_coordinator (bool, optional): If True, the global rank 0
                 process will spawn and manage the `InferenceCoordinator`
                 process. Defaults to True.
+
+        Returns:
+            inference_coordinator_addresss (str): The network address of the central
+                `InferenceCoordinator`, which may not have the same port as what the user requested
+                with `inference_coordinator_port`.
         """
 
         assert HAVE_ZMQ, (
@@ -411,24 +419,43 @@ async def start_listening_to_data_parallel_coordinator(
         self.is_mp_coordinator = tp_rank == 0 and pp_rank == 0
         self.is_dp_coordinator = (dp_rank == 0) and self.is_mp_coordinator
 
+        local_ip = socket.gethostname()
+
         # Spawn a DP coordinator process and get the connection info.
         if launch_inference_coordinator and self.is_dp_coordinator:
             spawn_context = multiprocessing.get_context('spawn')
+            dp_pipe, dp_process_pipe = spawn_context.Pipe()
             coordinator_ready_event = spawn_context.Event()
             self.inference_coordinator_process = spawn_context.Process(
                 target=DataParallelInferenceCoordinator.entrypoint,
                 args=(
+                    dp_process_pipe,
                     coordinator_ready_event,
-                    inference_coordinator_port,
                     get_pg_size(self.pg_collection.dp),
                     self.controller.tokenizer,
+                    inference_coordinator_port,
                 ),
             )
             self.inference_coordinator_process.start()
+            await await_process_call(dp_pipe.poll, self.inference_coordinator_process)
+            dp_addr = dp_pipe.recv()
+            dp_pipe.close()
+
+            # Check if the port number is not inference_coordinator_port
+            actual_port = int(dp_addr.rsplit(":", 1)[-1])
+            if inference_coordinator_port != None and actual_port != inference_coordinator_port:
+                logging.warning(
+                    f"Requested InferenceCoordinator port {inference_coordinator_port} "
+                    f"but got port {actual_port} instead. This happens if the request port "
+                    f"is already in use."
+                )
+        elif not launch_inference_coordinator:
+            dp_addr = f"tcp://{local_ip}:{inference_coordinator_port}"
+        else:
+            dp_addr = None
 
         # Find available ports for MP and bind to them.
         if self.is_mp_coordinator:
-            local_ip = socket.gethostname()
             mp_req_sock = self.zmq_context.socket(zmq.PUB)
             mp_req_sock.bind_to_random_port(f"tcp://{local_ip}")
             mp_req_addr = mp_req_sock.getsockopt_string(zmq.LAST_ENDPOINT)
@@ -441,12 +468,13 @@ async def start_listening_to_data_parallel_coordinator(
             mp_len_addr = None
 
         # Broadcast addresses to respective ranks.
+        bcast = [dp_addr]
+        torch.distributed.broadcast_object_list(bcast, src=dp_src, group=dp_group)
+        [dp_addr] = bcast
         bcast = [mp_req_addr, mp_len_addr]
         torch.distributed.broadcast_object_list(bcast, src=mp_src, group=mp_group)
         [mp_req_addr, mp_len_addr] = bcast
 
-        ip_address_of_dp_coordinator = os.getenv('MASTER_ADDR', '127.0.0.1')
-        dp_addr = f"tcp://{ip_address_of_dp_coordinator}:{inference_coordinator_port}"
         identity = f'mp-coord-{dp_rank}'
         if self.is_mp_coordinator:
             # 1. Create dealer sockets where tp_rank = 0 and pp_rank = 0
@@ -495,13 +523,18 @@ async def start_listening_to_data_parallel_coordinator(
             )
 
         if launch_inference_coordinator and self.is_dp_coordinator:
-            await await_process_event(coordinator_ready_event, self.inference_coordinator_process)
+            await await_process_call(
+                coordinator_ready_event.wait, self.inference_coordinator_process
+            )
             logging.info("Inference co-ordinator is ready to receive requests!")
+            logging.info(f"Data parallel coordinator can be found at {dp_addr}")
 
         # Finally run the engine infinite loop
         loop = get_asyncio_loop(loop)
         self.engine_loop_task = loop.create_task(self.run_engine_with_coordinator(loop=loop))
 
+        return dp_addr
+
     @contextmanager
     @staticmethod
     def suspend_resume_ctx(key: str, *, unified_memory_level: int) -> None:
diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py
index 8659368b9fa..a927a393b8c 100644
--- a/megatron/core/inference/inference_client.py
+++ b/megatron/core/inference/inference_client.py
@@ -2,7 +2,6 @@
 
 import asyncio
 import logging
-import os
 import time
 from typing import Awaitable, List, Optional, Union
 
@@ -54,12 +53,12 @@ class InferenceClient:
             completed requests.
     """
 
-    def __init__(self, inference_coordinator_port: int):
+    def __init__(self, inference_coordinator_address: str):
         """
         Initializes the InferenceClient.
 
         Args:
-            inference_coordinator_port (int): The port number on which the
+            inference_coordinator_address (str): The address on which the
                 inference coordinator is listening.
         """
         assert (
@@ -70,8 +69,7 @@ def __init__(self, inference_coordinator_port: int):
         ), "please install the messagepack library to use InferenceClient - pip install msgpack"
         self.context = zmq.Context()
         socket = self.context.socket(zmq.DEALER)
-        inference_coordinator_address = os.getenv('MASTER_ADDR', '127.0.0.1')
-        socket.connect(f"tcp://{inference_coordinator_address}:{inference_coordinator_port}")
+        socket.connect(inference_coordinator_address)
 
         self._loop = None
         self.running = asyncio.Event()
diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index 55536a52088..0bdaff64be1 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -139,10 +139,8 @@ def tensor_swap(x, src_idxs, dst_idxs):
     x[dst_idxs], x[src_idxs] = x[src_idxs], x[dst_idxs]
 
 
-async def await_process_event(
-    event: multiprocessing.Event, process: multiprocessing.Process, timeout: float = 1.0
-) -> None:
-    """Repeatedly wait for a multiprocessing event to be set, aborting upon process failure.
+async def await_process_call(call, process: multiprocessing.Process, timeout: float = 1.0):
+    """Repeatedly wait for a multiprocessing callable to resolve, aborting upon process failure.
 
     Note that the timeout in this function is only for checking process liveness.
     Its value should be set to a relatively high number. The only problem a high timeout
@@ -155,8 +153,7 @@ async def await_process_event(
         timeout: The timeout for each wait iteration in seconds.
     """
     while True:
-        signal = await asyncio.to_thread(event.wait, timeout)
-        if signal:
+        if await asyncio.to_thread(call, timeout):
             return
         if not process.is_alive():
             raise RuntimeError(
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 48c02774fca..4e9364b3ae9 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -280,12 +280,12 @@ async def launch(cls, model: GPTModel, **kwargs):
         inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(
             args, model, inference_logging_step_interval, metrics_writer
         )
-        await inference_engine.start_listening_to_data_parallel_coordinator(
+        dp_addr = await inference_engine.start_listening_to_data_parallel_coordinator(
             inference_coordinator_port=41521, launch_inference_coordinator=True,
         )
         if dist.get_rank() == 0:
             # TODO: We have to do this only on the rank 0 process, should be fixed in the future when we have support for multiple inference clients. !2278
-            client = InferenceClient(inference_coordinator_port=41521)
+            client = InferenceClient(inference_coordinator_address=dp_addr)
             await client.start()
         else:
             client = None
diff --git a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py
index 7b4fb4b4250..57326291a73 100644
--- a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py
+++ b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py
@@ -1,20 +1,18 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import asyncio
-import random
+import itertools
 import time
 from collections import deque
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, Optional
 
+import msgpack
 import pytest
-import torch.distributed as dist
+import torch
 from tqdm import tqdm
 
-from megatron.core.inference.data_parallel_inference_coordinator import (
-    DataParallelInferenceCoordinator,
-)
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine, RequestEntry
+from megatron.core.inference.headers import Headers
 from megatron.core.inference.inference_client import InferenceClient
 from megatron.core.inference.inference_request import (
     DynamicInferenceRequest,
@@ -22,6 +20,7 @@
     Status,
 )
 from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.utils import get_asyncio_loop
 from tests.unit_tests.test_utilities import Utils
 
@@ -29,10 +28,35 @@
     import zmq
 
     HAVE_ZMQ = True
-except Exception:
+except ImportError:
     HAVE_ZMQ = False
 
-IS_ZMQ_FLAKY = True
+NUM_REQUESTS = 10
+NUM_TOKENS = 2
+DEFAULT_PORT = 46581
+ZMQ_FLAKY_SHUTDOWN = True
+
+
+class DummyTokenizer:
+    """Dummy tokenizer."""
+
+    def __init__(self, vocab_size: int = 10, bos: int | None = None, eod: int = 0, pad: int = 0):
+        self.vocab_size = vocab_size
+        self.bos = bos
+        self.eod = eod
+        self.pad = pad
+
+    def tokenize(self, prompt):
+        if isinstance(prompt, str):
+            return [int(tok) % self.vocab_size for tok in prompt.strip().split()]
+        return list(prompt)
+
+    def detokenize(self, tokens, skip_special_tokens: bool = False):
+        if isinstance(tokens, torch.Tensor):
+            tokens = tokens.tolist()
+        if skip_special_tokens and self.eod in tokens:
+            tokens = [tok for tok in tokens if tok != self.eod]
+        return " ".join(str(tok) for tok in tokens)
 
 
 class DummyContext:
@@ -45,6 +69,16 @@ def get_active_request_count(self) -> int:
         return self.active_cnt
 
 
+class DummyController:
+    """Dummy inference controller."""
+
+    def __init__(self):
+        self.tokenizer = DummyTokenizer()
+
+    def dummy_forward(self):
+        pass
+
+
 class DummyEngine(DynamicInferenceEngine):
     """Dummy inference engine that only implements coordinator-related methods."""
 
@@ -56,12 +90,15 @@ def __init__(self):
         self.is_suspended = False
         self._loop = get_asyncio_loop()
         self.context = DummyContext()
+        self.controller = DummyController()
         self.running = asyncio.Event()
         self.paused = asyncio.Event()
         self.stopped = asyncio.Event()
         self.pending_microbatch = deque()
         self.received_pause: bool = False
         self.received_stop: bool = False
+        self.pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        self.rank = torch.distributed.get_rank()
 
     def add_request(
         self, request_id: int, prompt: str, sampling_params: Optional[SamplingParams] = None
@@ -99,6 +136,13 @@ async def async_step(self, *, verbose: Optional[bool] = False) -> Dict:
                 finished_request_records.append(entry.record)
                 entry.future.set_result(entry.record)
                 to_remove.append(request_id)
+                # Send signal to coordinator.
+                if self.is_mp_coordinator:
+                    payload = msgpack.packb(
+                        [Headers.ENGINE_REPLY.value, [entry.record.serialize()]], use_bin_type=True
+                    )
+                    self.socket_for_receiving_requests.send(payload)
+
         for request_id in to_remove:
             del self.requests[request_id]
 
@@ -119,301 +163,289 @@ async def async_step(self, *, verbose: Optional[bool] = False) -> Dict:
         }
 
 
-@dataclass
-class CoordinatorTestConfig:
-    """Test configuration args."""
-
-    port: int = 46581
-    mp_port: int = 49581
-    launch_inference_coordinator: bool = True
-    stop_engines: bool = True
-    verify_results: bool = True
-
-    num_requests: int = 10**1
-    min_time_offset: float = 10 ** (-4)
-    max_time_offset: float = 10 ** (-3)
-    num_steps_to_finish: int = 1
-    num_iterations: int = 1
-
-    tensor_model_parallel_size: int = 1
-    pipeline_model_parallel_size: int = 1
-
-
-@dataclass
-class CoordinatorTestEnv:
-    """Test environment, including requests."""
-
-    config: CoordinatorTestConfig
-    requests: List[Tuple]
-    engine: DummyEngine
-    responses: List[List[DynamicInferenceRequest]] = field(default_factory=list)
-    timing_data: Dict[str, Optional[float]] = field(
-        default_factory=lambda: {
-            "start_time": None,
-            "init_time": None,
-            "done_time": None,
-            "stop_time": None,
-        }
-    )
+@pytest.fixture
+def initialize_model_parallel(request, monkeypatch):
+    """Fixture to initialize and destroy model parallel.
 
+    Parameters are passed via request.param as a tuple: (tp, pp, ep).
+    Defaults to (1, 1, 1) if not parametrized.
+    """
+    monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
 
-class TestCoordinator:
-
-    @classmethod
-    def _build_requests(cls, test_config: CoordinatorTestConfig) -> List[Tuple]:
-        ret = []
+    tp, pp, ep = getattr(request, "param", (1, 1, 1))
+    world_size = Utils.world_size
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=tp,
+        pipeline_model_parallel_size=pp,
+        expert_model_parallel_size=ep,
+    )
+    dp = world_size // (tp * pp * ep)
+    yield world_size, dp, tp, pp, ep
+    Utils.destroy_model_parallel()
 
-        for _ in range(test_config.num_requests):
-            arrival_delta = random.uniform(test_config.min_time_offset, test_config.max_time_offset)
-            num_tokens = test_config.num_steps_to_finish
-            ret.append(
-                ("Hello world!", SamplingParams(num_tokens_to_generate=num_tokens), arrival_delta)
-            )
-        return ret
 
-    @classmethod
-    def _build_test_env(cls, test_config):
-        Utils.initialize_model_parallel(
-            tensor_model_parallel_size=test_config.tensor_model_parallel_size,
-            pipeline_model_parallel_size=test_config.pipeline_model_parallel_size,
-        )
-        requests = cls._build_requests(test_config)
+@pytest.mark.skipif(ZMQ_FLAKY_SHUTDOWN, reason="ZMQ shutdown is flaky")
+class TestCoordinator:
+    """Test class for Data Parallel Inference Coordinator."""
+
+    def build_requests(self, num_requests=NUM_REQUESTS, num_tokens=NUM_TOKENS):
+        """Build a list of test requests."""
+        return [
+            ("Hello world!", SamplingParams(num_tokens_to_generate=num_tokens))
+            for _ in range(num_requests)
+        ]
+
+    async def run_coordinator_test(
+        self,
+        *,
+        launch_coordinator=True,
+        stop_engines=True,
+        num_requests=NUM_REQUESTS,
+        num_tokens=NUM_TOKENS,
+    ):
+        """Run a coordinator test. Model parallel must already be initialized."""
         engine = DummyEngine()
-        engine.num_steps_to_finish = test_config.num_steps_to_finish
-        return CoordinatorTestEnv(config=test_config, requests=requests, engine=engine)
-
-    @classmethod
-    async def _run_test(cls, **test_config_kwargs):
-        # Test environment.
-        test_config = CoordinatorTestConfig(**test_config_kwargs)
-        env = cls._build_test_env(test_config)
-
-        # Connect each engine to their respective processes.
-        env.timing_data["start_time"] = time.time()
-        await env.engine.start_listening_to_data_parallel_coordinator(
-            inference_coordinator_port=test_config.port,
-            launch_inference_coordinator=test_config.launch_inference_coordinator,
+        requests = self.build_requests(num_requests, num_tokens)
+
+        dp_addr = await engine.start_listening_to_data_parallel_coordinator(
+            inference_coordinator_port=DEFAULT_PORT, launch_inference_coordinator=launch_coordinator
         )
 
-        results_success = False
-        shutdown_success = False
         try:
-            if dist.get_rank() == 0:
-                client = InferenceClient(test_config.port)
+            if torch.distributed.get_rank() == 0:
+                client = InferenceClient(dp_addr)
                 await client.start()
-                env.timing_data["init_time"] = time.time()
 
-                all_results = []
-                for _ in range(test_config.num_iterations):
-                    futures = []
-                    for request in tqdm(env.requests, "add_requests"):
-                        prompt, sampling_params, arrival_delta = request
-                        await asyncio.sleep(arrival_delta)
-                        fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
-                        futures.append(fut)
-                    results = await asyncio.wait_for(asyncio.gather(*futures), timeout=10.0)
-                    all_results.append(results)
-                env.timing_data["done_time"] = time.time()
-            results_success = True
-        finally:
-            try:
-                if dist.get_rank() == 0:
-                    if test_config.stop_engines:
-                        await asyncio.wait_for(client.stop_engines(), timeout=10.0)
-                    client.stop()
-                if test_config.stop_engines:
-                    await asyncio.wait_for(env.engine.engine_loop_task, timeout=10.0)
-                shutdown_success = True
-            except:
-                env.engine.engine_loop_task.cancel()
-
-        env.timing_data["stop_time"] = time.time()
-
-        assert results_success, "Did not receive all results successfully."
-        assert shutdown_success, "Did not shutdown successfully."
-        if dist.get_rank() == 0:
-            env.responses = all_results
-            if test_config.verify_results:
-                for batch in all_results:
-                    for record in batch:
-                        request = record[-1]
-                        assert request.status == Status.COMPLETED
+                futures = [
+                    client.add_request(prompt=prompt, sampling_params=params)
+                    for prompt, params in requests
+                ]
+                results = await asyncio.wait_for(asyncio.gather(*futures), timeout=10.0)
 
-        return env
+                for record in results:
+                    assert record[-1].status == Status.COMPLETED
+        finally:
+            if torch.distributed.get_rank() == 0:
+                if stop_engines:
+                    await asyncio.wait_for(client.stop_engines(), timeout=10.0)
+                client.stop()
+            if stop_engines:
+                try:
+                    await asyncio.wait_for(engine.engine_loop_task, timeout=30.0)
+                except asyncio.TimeoutError:
+                    engine.engine_loop_task.cancel()
 
-    def teardown_method(self, method):
-        Utils.destroy_model_parallel()
+        return dp_addr
 
     @pytest.mark.internal
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
     @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
     @pytest.mark.asyncio
-    async def test_simple(self):
-        """Simple test with no TP or PP."""
-        env = await self._run_test(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+    @pytest.mark.parametrize(
+        "initialize_model_parallel",
+        [
+            pytest.param((tp, pp, ep), id=f"tp{tp}-pp{pp}-ep{ep}")
+            for tp, pp, ep in itertools.product([1, 2], [1, 2], [1, 2])
+            if tp * pp * ep <= Utils.world_size
+        ],
+        indirect=["initialize_model_parallel"],
+    )
+    async def test_parallel_configs(self, initialize_model_parallel):
+        """Test coordinator with various TP, PP, and EP configurations."""
+        await self.run_coordinator_test()
 
     @pytest.mark.internal
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
     @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
     @pytest.mark.asyncio
-    async def test_tp(self):
-        """Simple test with TP, but no PP."""
-        env = await self._run_test(tensor_model_parallel_size=2, pipeline_model_parallel_size=1)
+    async def test_coordinator_lifecycle(self, initialize_model_parallel):
+        """Test coordinator connection and port conflict behavior."""
+        engine1 = DummyEngine()
+        engine2 = None
+        engine3 = None
+        third_addr = None
+
+        # Launch first coordinator - binds to DEFAULT_PORT
+        first_addr = await engine1.start_listening_to_data_parallel_coordinator(
+            inference_coordinator_port=DEFAULT_PORT, launch_inference_coordinator=True
+        )
 
-    @pytest.mark.internal
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
-    @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
-    @pytest.mark.asyncio
-    async def test_pp(self):
-        """Simple test with no TP, but PP."""
-        env = await self._run_test(tensor_model_parallel_size=1, pipeline_model_parallel_size=2)
+        try:
+            # Cancel engine1 loop without sending stop to coordinator
+            # This keeps coordinator process alive and holding the port
+            engine1.engine_loop_task.cancel()
+            try:
+                await engine1.engine_loop_task
+            except asyncio.CancelledError:
+                pass
+
+            # Connect engine2 to existing coordinator (don't launch new one)
+            engine2 = DummyEngine()
+            second_addr = await engine2.start_listening_to_data_parallel_coordinator(
+                inference_coordinator_port=DEFAULT_PORT, launch_inference_coordinator=False
+            )
 
-    @pytest.mark.internal
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
-    @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
-    @pytest.mark.asyncio
-    async def test_tp_pp(self):
-        """Simple test with both TP and PP."""
-        env = await self._run_test(tensor_model_parallel_size=2, pipeline_model_parallel_size=2)
+            # Should connect to same port, but will not always in CI due to port conflicts.
+            first_port = int(first_addr.rsplit(":", 1)[-1])
+            second_port = int(second_addr.rsplit(":", 1)[-1])
+            # assert second_port == first_port
 
-    @pytest.mark.internal
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
-    @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
-    @pytest.mark.asyncio
-    async def test_pp(self):
-        """Simple test with no TP, but PP."""
-        env = await self._run_test(tensor_model_parallel_size=1, pipeline_model_parallel_size=2)
+            # Cancel engine2
+            engine2.engine_loop_task.cancel()
+            try:
+                await engine2.engine_loop_task
+            except asyncio.CancelledError:
+                pass
+
+            # Launch new coordinator - should get different port since first is holding it
+            engine3 = DummyEngine()
+            third_addr = await engine3.start_listening_to_data_parallel_coordinator(
+                inference_coordinator_port=DEFAULT_PORT, launch_inference_coordinator=True
+            )
 
-    @pytest.mark.internal
-    @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
-    @pytest.mark.asyncio
-    async def test_tp_pp(self):
-        """Simple test with both TP and PP."""
-        env = await self._run_test(tensor_model_parallel_size=2, pipeline_model_parallel_size=2)
+            # Verify we got a different port due to conflict
+            third_port = int(third_addr.rsplit(":", 1)[-1])
+            assert (
+                third_port != first_port
+            ), f"Expected different port due to conflict, but got same: {third_port}"
+
+        finally:
+            # Clean up engine3's coordinator
+            if engine3 is not None and third_addr is not None:
+                client3 = InferenceClient(third_addr)
+                await client3.start()
+                await asyncio.wait_for(client3.stop_engines(), timeout=10.0)
+                client3.stop()
+                try:
+                    await asyncio.wait_for(engine3.engine_loop_task, timeout=30.0)
+                except asyncio.TimeoutError:
+                    engine3.engine_loop_task.cancel()
+
+            # Rebuild engine and reconnect to engine1's coordinator
+            first_port = int(first_addr.rsplit(":", 1)[-1])
+            engine1 = DummyEngine()
+            await engine1.start_listening_to_data_parallel_coordinator(
+                inference_coordinator_port=first_port, launch_inference_coordinator=False
+            )
+            client1 = InferenceClient(first_addr)
+            await client1.start()
+            await asyncio.wait_for(client1.stop_engines(), timeout=10.0)
+            client1.stop()
+            try:
+                await asyncio.wait_for(engine1.engine_loop_task, timeout=30.0)
+            except asyncio.TimeoutError:
+                engine1.engine_loop_task.cancel()
 
     @pytest.mark.internal
     @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
     @pytest.mark.asyncio
-    async def test_pause(self):
-        """Pause/resume test."""
-        test_config = CoordinatorTestConfig(
-            tensor_model_parallel_size=2, pipeline_model_parallel_size=1, num_requests=32
-        )
-        env = self._build_test_env(test_config)
+    async def test_pause(self, initialize_model_parallel):
+        """Test pause and resume functionality."""
+        engine = DummyEngine()
+        requests = self.build_requests(num_requests=32)
 
-        await env.engine.start_listening_to_data_parallel_coordinator(
-            inference_coordinator_port=test_config.port, launch_inference_coordinator=True
+        dp_addr = await engine.start_listening_to_data_parallel_coordinator(
+            inference_coordinator_port=DEFAULT_PORT, launch_inference_coordinator=True
         )
 
-        success = False
+        success = True
         try:
-            if dist.get_rank() == 0:
-                # Start client as usual.
-                client = InferenceClient(test_config.port)
+            if torch.distributed.get_rank() == 0:
+                client = InferenceClient(dp_addr)
                 await client.start()
 
-                ### TEST 1: Pause after all requests have finished.
-                futures = []
-                for i, request in enumerate(env.requests[:2]):
-                    prompt, sampling_params, _ = request
-                    fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
-                    futures.append(fut)
-                # Wait a sufficient time for the requests to complete.
+                # Submit requests and pause after completion.
+                futures = [client.add_request(prompt=p, sampling_params=s) for p, s in requests[:2]]
                 await asyncio.sleep(0.1)
-                # Get a pause awaitable.
-                to_pause = client.pause_engines()
-                awaitables = futures + [to_pause]
-                # Gather all awaitables; assert that the requests actually complete.
+                awaitables = futures + [client.pause_engines()]
                 try:
-                    await asyncio.wait_for(asyncio.gather(*awaitables), timeout=0.1)
+                    await asyncio.wait_for(asyncio.gather(*awaitables), timeout=0.5)
                 except asyncio.TimeoutError:
-                    pytest.fail("Simple pause did not succeed.")
+                    pytest.fail("Pause operation timed out.")
 
-                ### TEST 2: Ensure that requests can be added while paused.
-                prompt, sampling_params, _ = env.requests[2]
-                paused_fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
+                # Ensure that requests can be added while paused.
+                prompt, params = requests[2]
+                future = client.add_request(prompt=prompt, sampling_params=params)
                 with pytest.raises(asyncio.TimeoutError):
-                    await asyncio.wait_for(paused_fut, timeout=0.1)
+                    await asyncio.wait_for(future, timeout=0.1)
 
-                ### TEST 3: Resume after pause and drain the queued requests.
+                # Resume and verify new requests complete.
                 client.unpause_engines()
                 # TODO: The system should not be incorrectly raising a cancelled error here.
                 with pytest.raises(asyncio.CancelledError):
-                    await paused_fut
-
-                ### TEST 4: Add new requests after resume.
-                futures = []
-                for i, request in enumerate(env.requests[3:4]):
-                    prompt, sampling_params, _ = request
-                    fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
-                    futures.append(fut)
-                # Wait a sufficient time for the requests to complete.
+                    await future
+
+                futures = [
+                    client.add_request(prompt=p, sampling_params=s) for p, s in requests[3:4]
+                ]
                 await asyncio.sleep(0.1)
-                # Gather all awaitables; assert that the requests actually complete.
                 try:
-                    await asyncio.wait_for(asyncio.gather(*futures), timeout=0.1)
+                    await asyncio.wait_for(asyncio.gather(*futures), timeout=0.5)
                 except asyncio.TimeoutError:
-                    pytest.fail("Simple resume did not succeed.")
-
-                ### TEST 5: Pause while requests are being processed.
-                ### Note: this situation cannot occur in a synchronous system.
-                if False:
-                    for request in env.engine.requests[4:6]:
-                        request.sampling_params.num_tokens_to_generate = 100
-                    futures = []
-                    for i, request in enumerate(env.requests[4:6]):
-                        prompt, sampling_params, _ = request
-                        fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
-                        futures.append(fut)
-                    # Do not wait for the requests to complete.
-                    await client.pause_engines()
-                    # Gather all awaitables; assert that the requests do not complete.
-                    with pytest.raises(asyncio.TimeoutError):
-                        await asyncio.wait_for(asyncio.gather(*futures), timeout=0.1)
-            success = True
+                    pytest.fail("Resumed requests did not complete in time.")
+        except:
+            success = False
         finally:
             try:
-                if dist.get_rank() == 0:
+                if torch.distributed.get_rank() == 0:
                     await asyncio.wait_for(client.stop_engines(), timeout=5.0)
                     client.stop()
-                await asyncio.wait_for(env.engine.engine_loop_task, timeout=5.0)
+                await asyncio.wait_for(engine.engine_loop_task, timeout=30.0)
             except asyncio.TimeoutError:
-                env.engine.engine_loop_task.cancel()
-        assert success, "Pause/resume test did not complete successfully."
+                engine.engine_loop_task.cancel()
+        assert success, "Pause/resume test failed."
 
     @pytest.mark.internal
     @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test")
-    @pytest.mark.skipif(IS_ZMQ_FLAKY, reason="pyzmq is flaky in CI")
     @pytest.mark.asyncio
-    async def test_throughput(self):
+    async def test_throughput(self, initialize_model_parallel):
         """Throughput test with no TP or PP."""
-        import torch
-        import torch.distributed as dist
-
-        env = await self._run_test(
-            tensor_model_parallel_size=1,
-            pipeline_model_parallel_size=1,
-            num_requests=10**4,
-            num_iterations=10,
-            min_time_offset=0.0,
-            max_time_offset=0.0,
+        num_requests = 10**4
+        num_iterations = 10
+
+        engine = DummyEngine()
+        requests = self.build_requests(num_requests=num_requests)
+
+        start_time = time.time()
+        dp_addr = await engine.start_listening_to_data_parallel_coordinator(
+            inference_coordinator_port=DEFAULT_PORT, launch_inference_coordinator=True
         )
 
+        try:
+            if torch.distributed.get_rank() == 0:
+                client = InferenceClient(dp_addr)
+                await client.start()
+                init_time = time.time()
+
+                for _ in range(num_iterations):
+                    futures = []
+                    for prompt, sampling_params in tqdm(requests, "add_requests"):
+                        fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
+                        futures.append(fut)
+                    await asyncio.wait_for(asyncio.gather(*futures), timeout=10.0)
+                done_time = time.time()
+        finally:
+            if torch.distributed.get_rank() == 0:
+                await asyncio.wait_for(client.stop_engines(), timeout=10.0)
+                client.stop()
+            try:
+                await asyncio.wait_for(engine.engine_loop_task, timeout=30.0)
+            except asyncio.TimeoutError:
+                engine.engine_loop_task.cancel()
+
+        stop_time = time.time()
+
         flags = torch.tensor([1, 1, 1], dtype=torch.int, device=torch.cuda.current_device())
 
         init_duration = golden_init_duration = None
         run_duration = golden_run_duration = None
         stop_duration = golden_stop_duration = None
 
-        if dist.get_rank() == 0:
-            init_duration = (env.timing_data["init_time"] - env.timing_data["start_time"]) * 10**3
-            golden_init_duration = 4445.64  # ms
-            run_duration = (env.timing_data["done_time"] - env.timing_data["init_time"]) * 10**3
-            golden_run_duration = 2906.29  # ms
-            stop_duration = (env.timing_data["stop_time"] - env.timing_data["done_time"]) * 10**3
-            golden_stop_duration = 33.17  # ms
+        if torch.distributed.get_rank() == 0:
+            init_duration = (init_time - start_time) * 10**3
+            golden_init_duration = 6974.43  # ms
+            run_duration = (done_time - init_time) * 10**3
+            golden_run_duration = 4392.63  # ms
+            stop_duration = (stop_time - done_time) * 10**3
+            golden_stop_duration = 931.49  # ms
 
             def clamp_to_golden_value(value, golden_value, delta=0.1):
                 return value > golden_value * (1 - delta) and value < golden_value * (1 + delta)
@@ -426,10 +458,9 @@ def clamp_to_golden_value(value, golden_value, delta=0.1):
                 flags[2] = 0
 
         # Synchronize results
-        dist.broadcast(flags, src=0)
+        torch.distributed.broadcast(flags, src=0)
 
-        if dist.get_rank() == 0:
-            # Print current results.
+        if torch.distributed.get_rank() == 0:
             print(f"Initialization time: {init_duration:.2f} ms")
             print(f"Run time: {run_duration:.2f} ms")
             print(f"Stop time: {stop_duration:.2f} ms")
@@ -449,23 +480,10 @@ def clamp_to_golden_value(value, golden_value, delta=0.1):
 
             print(
                 f"ZMQ throughput is approximately "
-                f"{env.config.num_requests * env.config.num_iterations / (run_duration):.2f} "
+                f"{num_requests * num_iterations / run_duration:.2f} "
                 f"requests/ms"
             )
         else:
             assert flags[0].item() == 1
             assert flags[1].item() == 1
             assert flags[2].item() == 1
-
-
-if __name__ == "__main__":
-    test = TestCoordinator()
-    asyncio.run(test.test_simple())
-    asyncio.run(test.test_tp())
-    asyncio.run(test.test_pp())
-    asyncio.run(test.test_tp_pp())
-    asyncio.run(test.test_pause())
-    asyncio.run(test.test_throughput())
-    test.teardown_method(None)
-    print("~~~")
-    print("success.")

From 639c08ab828408e2e563ea310d82b5fb977edb77 Mon Sep 17 00:00:00 2001
From: Maanu Grover <109391026+maanug-nv@users.noreply.github.com>
Date: Thu, 29 Jan 2026 16:26:31 -0800
Subject: [PATCH 78/79] Generate arguments from TransformerConfig (#2896)

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/core/model_parallel_config.py        |  66 +-
 .../core/transformer/transformer_config.py    | 131 ++--
 megatron/training/arguments.py                | 612 +++---------------
 3 files changed, 213 insertions(+), 596 deletions(-)

diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
index a995c360eb6..3c6ff04d3b0 100644
--- a/megatron/core/model_parallel_config.py
+++ b/megatron/core/model_parallel_config.py
@@ -1,8 +1,8 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import warnings
-from dataclasses import dataclass
-from typing import Callable, ContextManager, Optional
+from dataclasses import dataclass, field
+from typing import Callable, ContextManager, Literal, Optional
 
 import torch
 
@@ -20,7 +20,7 @@ class ModelParallelConfig:
     tensor_model_parallel_size: int = 1
     """Intra-layer model parallelism. Splits tensors across GPU ranks."""
 
-    pipeline_model_parallel_comm_backend: Optional[str] = None
+    pipeline_model_parallel_comm_backend: Optional[Literal["nccl", "ucc"]] = None
     """Configuring backend option of pipeline parallel communication (e.g., nccl, ucc)
        If None, the default backend will be used.
     """
@@ -73,7 +73,9 @@ class ModelParallelConfig:
     """Distributes Moe Experts across sub data parallel dimension."""
 
     expert_tensor_parallel_size: Optional[int] = None
-    """Intra-layer tensor model parallelsm for expert layer. Splits tensors across GPU ranks."""
+    """Intra-layer tensor model parallelism for expert layer. Splits tensors across GPU ranks.
+       Default is None, which will be set to the value of tensor_model_parallel_size.
+    """
 
     moe_extended_tp: bool = False
     """NOTE: Deprecated from MCore v0.10. This flag is ignored.
@@ -83,12 +85,16 @@ class ModelParallelConfig:
     ###################
     # Initialization
     ###################
-    perform_initialization: bool = True
-    """If true, weights are initialized. This option can be useful when you know you are going to
+    perform_initialization: bool = field(
+        default=True, metadata={"argparse_meta": {"arg_names": ["--no-initialization"]}}
+    )
+    """Controls weights initialization. This option can be useful when you know you are going to
        load values from a checkpoint.
     """
 
-    use_cpu_initialization: bool = False
+    use_cpu_initialization: bool = field(
+        default=False, metadata={"argparse_meta": {"default": None}}
+    )
     """When set to False, we initialize the weights directly on the GPU. CPU initialization is the
        same regardless of tensor model parallelism, but GPU initialization is not. Transferring
        weights from CPU to GPU can take a significant amount of time for large models.
@@ -167,11 +173,14 @@ class ModelParallelConfig:
        must turn off gradient accumulation fusion.
     """
 
-    async_tensor_model_parallel_allreduce: bool = False
+    async_tensor_model_parallel_allreduce: bool = True
     """NOTE: Deprecated. This flag is ignored."""
 
-    use_te_rng_tracker: bool = False
+    use_te_rng_tracker: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--te-rng-tracker"]}}
+    )
     """If true, uses RNG state tracker in TransformerEngine if exists.
+    Required for CUDA graphs support.
     """
 
     tp_comm_overlap: bool = False
@@ -181,22 +190,22 @@ class ModelParallelConfig:
     """
 
     tp_comm_bulk_wgrad: bool = True
-    """If true, allows All-Gather overlap with Bprop activation gradient GEMM. Don't care if
+    """Controls All-Gather overlap with Bprop activation gradient GEMM. Don't care if
        tp_comm_overlap is False.
     """
 
     tp_comm_bulk_dgrad: bool = True
-    """If true, allows Reduce-Scatter overlap with Bprop weight gradient GEMM. Don't care if
+    """Controls Reduce-Scatter overlap with Bprop weight gradient GEMM. Don't care if
        tp_comm_overlap is False.
     """
 
     tp_comm_overlap_ag: bool = True
-    """If true, allows All-Gather overlap with GEMM by pipelining the GEMM and All-Gather.
+    """Controls All-Gather overlap with GEMM by pipelining the GEMM and All-Gather.
        Don't care if tp_comm_overlap is False.
     """
 
     tp_comm_overlap_rs: bool = True
-    """If true, allows Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter.
+    """Controls Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter.
        Don't care if tp_comm_overlap is False.
     """
 
@@ -207,7 +216,7 @@ class ModelParallelConfig:
 
     tp_comm_split_ag: bool = True
     """Deprecated from TransformerEngine v1.6.0.
-       If true, allows All-Gather overlap with Fprop GEMM by pipelining the GEMM and All-Gather
+       Controls All-Gather overlap with Fprop GEMM by pipelining the GEMM and All-Gather
        splits. Don't care if tp_comm_overlap is False.
     """
 
@@ -219,7 +228,7 @@ class ModelParallelConfig:
 
     tp_comm_split_rs: bool = True
     """Deprecated from TransformerEngine v1.6.0.
-       If true, allows Reduce-Scatter overlap with Fprop GEMM by pipelining the GEMM and
+       Controls Reduce-Scatter overlap with Fprop GEMM by pipelining the GEMM and
        Reduce-Scatter splits. Don't care if tp_comm_overlap is False.
     """
 
@@ -234,7 +243,7 @@ class ModelParallelConfig:
        Defaults to False.
     """
 
-    cross_entropy_fusion_impl: str = 'native'
+    cross_entropy_fusion_impl: Literal['native', 'te'] = 'native'
     """If 'native', MCore based CE loss fusion is used, if 'te', Parallel CE loss
        from Transformer Engine library is used. Defaults to 'native'.
     """
@@ -249,10 +258,8 @@ class ModelParallelConfig:
        If true, the AllGather -> Gemm overlap for FC1 layer of MLP gets disabled
     """
 
-    tp_comm_bootstrap_backend: str = 'nccl'
-    """
-       Set the bootstrapping backend out of 'nccl', 'mpi', and 'gloo'
-    """
+    tp_comm_bootstrap_backend: Literal['nccl', 'mpi', 'gloo'] = 'nccl'
+    """Set the bootstrapping backend of Tensor parallel communications."""
 
     overlap_moe_expert_parallel_comm: bool = False
     """Overlap EP A2A communications with independent computations of different micro-batches
@@ -324,13 +331,21 @@ class ModelParallelConfig:
        Defaults to 0, which means all micro-batches are deferred.
     """
 
-    overlap_p2p_comm_warmup_flush: bool = False
+    overlap_p2p_comm_warmup_flush: bool = field(
+        default=False,
+        metadata={"argparse_meta": {"arg_names": ["--overlap-p2p-communication-warmup-flush"]}},
+    )
     """If true, overlap communication and computation in warm up and flush phase.
        Only valid when overlap_p2p_comm is True and batch_p2p_comm is False. 
        Defaults to False.
     """
 
-    microbatch_group_size_per_vp_stage: Optional[int] = None
+    microbatch_group_size_per_vp_stage: Optional[int] = field(
+        default=None,
+        metadata={
+            "argparse_meta": {"arg_names": ["--microbatch-group-size-per-virtual-pipeline-stage"]}
+        },
+    )
     """This value specifies the number of micro-batches that are executed 
        at a time for a given virtual stage (both forward and backward).
        Default (in __post_init__() method below) to pipeline_parallel_size 
@@ -377,8 +392,11 @@ class ModelParallelConfig:
     ###################
     # Timing
     ###################
-    barrier_with_L1_time: bool = True
-    """If true, use barrier with level 1 time measurements. It is up to the user to make sure
+    barrier_with_L1_time: bool = field(
+        default=True,
+        metadata={"argparse_meta": {"arg_names": ["--no-barrier-with-level-1-timing"]}},
+    )
+    """Controls barrier with level 1 time measurements. It is up to the user to make sure
        calling barrier with their timers will not result in hangs. This can happen if for example
        the user adds a level 1 timer that is not called by all ranks.
     """
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index f50bbed0a44..eaae585905e 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -1,7 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Callable, List, Literal, Optional, Tuple, Union
 
 import torch
@@ -42,14 +42,22 @@ class TransformerConfig(ModelParallelConfig):
     # model architecture
     ####################
 
-    num_layers: int = 0
+    num_layers: int = field(default=0, metadata={"argparse_meta": {"default": None}})
     """Number of transformer layers in a transformer block."""
 
     mtp_num_layers: Optional[int] = None
-    """Number of Multi-Token Prediction (MTP) Layers."""
+    """Number of Multi-Token Prediction (MTP) Layers.
+    MTP extends the prediction scope to multiple future tokens at each position.
+    This MTP implementation sequentially predict additional tokens
+    by using D sequential modules to predict D additional tokens.
+    """
 
-    mtp_loss_scaling_factor: Optional[float] = None
-    """Weighting factor of Multi-Token Prediction (MTP) loss."""
+    mtp_loss_scaling_factor: Optional[float] = 0.1
+    """Weighting factor of Multi-Token Prediction (MTP) loss.
+    We compute the average of the MTP losses across all depths, 
+    and multiply it the scaling factor to obtain the overall MTP loss, 
+    which serves as an additional training objective.
+    """
 
     num_layers_in_first_pipeline_stage: Optional[int] = None
     """Number of transformer layers on first pipeline stage.
@@ -93,10 +101,10 @@ class TransformerConfig(ModelParallelConfig):
     """If set, the loss layer will be treated as a standard transformer
     layer in the context of partition and placement for pipeline parallelism."""
 
-    hidden_size: int = 0
+    hidden_size: int = field(default=0, metadata={"argparse_meta": {"default": None}})
     """Transformer hidden size."""
 
-    num_attention_heads: int = 0
+    num_attention_heads: int = field(default=0, metadata={"argparse_meta": {"default": None}})
     """Number of transformer attention heads."""
 
     attention_backend: AttnBackend = AttnBackend.auto
@@ -113,7 +121,9 @@ class TransformerConfig(ModelParallelConfig):
        Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and 
        and learnable offset."""
 
-    num_query_groups: Optional[int] = None
+    num_query_groups: Optional[int] = field(
+        default=None, metadata={"argparse_meta": {"default": 1}}
+    )
     """Number of query groups for group query attention. If None, normal attention is used."""
 
     ffn_hidden_size: Optional[int] = None
@@ -137,16 +147,22 @@ class TransformerConfig(ModelParallelConfig):
     apply_residual_connection_post_layernorm: bool = False
     """If True, uses the original BERT residule connection ordering."""
 
-    layernorm_epsilon: float = 1e-5
-    """Epsilon value for any LayerNorm operations."""
+    layernorm_epsilon: float = field(
+        default=1e-5, metadata={"argparse_meta": {"arg_names": ["--norm-epsilon"]}}
+    )
+    """Epsilon value for any LayerNorm/RMSNorm operations."""
 
-    layernorm_zero_centered_gamma: bool = False
+    layernorm_zero_centered_gamma: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--apply-layernorm-1p"]}}
+    )
     """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves
     numerical stability."""
 
-    add_bias_linear: bool = True
-    """Include a bias term in all linear layers (QKV projections, after core attention, and two in
-    MLP layer)."""
+    add_bias_linear: bool = field(
+        default=True, metadata={"argparse_meta": {"arg_names": ["--disable-bias-linear"]}}
+    )
+    """Include/exclude a bias term in all linear layers (QKV projections, after core attention,
+    and two in MLP layer)."""
 
     add_qkv_bias: bool = False
     """Add a bias term only for QKV projections."""
@@ -186,7 +202,7 @@ class TransformerConfig(ModelParallelConfig):
     - An integer N: Represents a (N-1):1 ratio, one full attention layer after (N-1) SWA layers.
     - A list that defines a custom pattern, e.g.: [1,1,1,1,0,0,0,0], where 1 represents SWA. """
 
-    normalization: str = "LayerNorm"
+    normalization: Literal['LayerNorm', 'RMSNorm'] = "LayerNorm"
     """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`."""
 
     qk_layernorm: bool = False
@@ -231,7 +247,7 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # attention variant
     ####################
-    experimental_attention_variant: Optional[str] = None
+    experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None
     """Type of attention variant to use. Currently support gated_delta_net and dsa."""
 
     ####################
@@ -263,19 +279,19 @@ class TransformerConfig(ModelParallelConfig):
     - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer
     - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]"""
 
-    linear_conv_kernel_dim: Optional[int] = None
+    linear_conv_kernel_dim: Optional[int] = 4
     """Conv kernel dimension for the gated delta net."""
 
-    linear_key_head_dim: Optional[int] = None
+    linear_key_head_dim: Optional[int] = 128
     """Query and key head dimension for the gated delta net."""
 
-    linear_value_head_dim: Optional[int] = None
+    linear_value_head_dim: Optional[int] = 128
     """Value and gate head dimension for the gated delta net."""
 
-    linear_num_key_heads: Optional[int] = None
+    linear_num_key_heads: Optional[int] = 16
     """Number of query and key heads for the gated delta net."""
 
-    linear_num_value_heads: Optional[int] = None
+    linear_num_value_heads: Optional[int] = 32
     """Number of value and gate heads for the gated delta net."""
 
     ####################
@@ -305,7 +321,10 @@ class TransformerConfig(ModelParallelConfig):
     embedding_init_method_std: Optional[float] = None
     """
     Standard deviation of the zero mean normal for the default initialization method for the 
-    embedding layer. If None, will be set to init_method_std.
+    embedding layer. If None, will be set to init_method_std. Setting this to a value around
+    1.0 may avoid loss spikes in training. Setting this to any value will also skip applying
+    weight decay on embedding weights to avoid shrinkage towards zero.
+    See https://arxiv.org/abs/2312.16903 for more details.
     """
 
     init_model_with_meta_device: bool = False
@@ -319,7 +338,7 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     apply_query_key_layer_scaling: bool = False
     """If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with
-    fp16."""
+    fp16. Also sets `attention_softmax_in_fp32` to True."""
 
     attention_softmax_in_fp32: bool = True
     """If True, run attention masking and softmax in fp32. This should be True if
@@ -361,7 +380,7 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # activation recomputation
     ####################
-    recompute_granularity: Optional[str] = None
+    recompute_granularity: Optional[Literal['full', 'selective']] = None
     """Determines which type of activation recompute to use.  Megatron-core supports 'selective'
     activation checkpointing where the submodules set in --recompute-modules is checkpointed.
     The default is "core_attn" which is the memory intensive part of attention.
@@ -372,7 +391,7 @@ class TransformerConfig(ModelParallelConfig):
     If set, must be 'selective' or 'full'. 'selective' always uses all layers.
     """
 
-    recompute_method: Optional[str] = None
+    recompute_method: Optional[Literal['uniform', 'block']] = None
     """Determines which transformer layers will be recomputed. uniform will uniformly divide the
     total number of transformer layers in a transformer block and recompute the input activation of
     each divided chunk at the specified granularity.  block will recompute the input activations for
@@ -386,7 +405,7 @@ class TransformerConfig(ModelParallelConfig):
     the number of transformer layers to recompute within each pipeline stage.  Must be None for
     'selective' activation checkpointing."""
 
-    distribute_saved_activations: Optional[bool] = None
+    distribute_saved_activations: Optional[bool] = False
     """If True, distribute recomputed activations across the model parallel group."""
 
     recompute_modules: Optional[List[str]] = None
@@ -407,12 +426,16 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # fp8 related
     ####################
-    fp8: Optional[str] = None
+    fp8: Optional[Literal['e4m3', 'hybrid']] = field(
+        default=None, metadata={"argparse_meta": {"arg_names": ["--fp8-format"]}}
+    )
     """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined
     choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8
     activation and weight tensors and e5m2 for all FP8 output activation gradient tensors."""
 
-    fp8_recipe: Optional[str] = "delayed"
+    fp8_recipe: Optional[Literal['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom']] = (
+        "delayed"
+    )
     """If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined
     choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed'
     uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only,
@@ -440,7 +463,7 @@ class TransformerConfig(ModelParallelConfig):
     fp8_amax_history_len: int = 1
     """The length of the amax history window used for scaling factor computation."""
 
-    fp8_amax_compute_algo: str = "most_recent"
+    fp8_amax_compute_algo: Literal['most_recent', 'max'] = "most_recent"
     """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2
     predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent`
     always chooses the most recently seen value.
@@ -484,15 +507,19 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # fp4 related
     ####################
-    fp4: Optional[str] = None
+    fp4: Optional[Literal['e2m1']] = field(
+        default=None, metadata={"argparse_meta": {"arg_names": ["--fp4-format"]}}
+    )
     """If set, enables the use of FP4 precision through Transformer Engine. Currently only 
     supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0)."""
 
-    fp4_recipe: Optional[str] = "nvfp4"
+    fp4_recipe: Optional[Literal['nvfp4', 'custom']] = "nvfp4"
     """If set, enables the use of FP4 precision through Transformer Engine. Currently only
     'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture."""
 
-    fp4_param: bool = False
+    fp4_param: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--fp4-param-gather"]}}
+    )
     """If set, keep the parameters in fp4 precision to save memory. This option must be used
     together with fp4 mode (i.e., TransformerConfig.fp4 is not None). Note that not all parameters
     will be converted to fp4; for example, biases will remain unchanged."""
@@ -522,7 +549,9 @@ class TransformerConfig(ModelParallelConfig):
 
     moe_shared_expert_overlap: bool = False
     """Enable overlapping between shared expert computations and dispatcher communications.
-    Without this, the shared experts execute before the router."""
+    Without this, the shared experts execute before the router. 
+    Only effective when moe-shared-expert-intermediate-size is set.
+    """
 
     moe_layer_freq: Union[int, List[int]] = 1
     """Frequency between MoE layers and Dense layers. Accepts either:
@@ -530,7 +559,7 @@ class TransformerConfig(ModelParallelConfig):
     - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]"""
 
     moe_ffn_hidden_size: Optional[int] = None
-    """MoE Feed-Forward Network hidden size"""
+    """MoE Feed-Forward Network hidden size. If not specified, defaults to the ffn_hidden_size."""
 
     moe_router_load_balancing_type: Union[str, List[str]] = "aux_loss"
     """The load balancing strategy for the router.
@@ -594,10 +623,10 @@ class TransformerConfig(ModelParallelConfig):
     """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax
     enabled. Defaults to None, which means no scaling."""
 
-    moe_router_score_function: str = "softmax"
+    moe_router_score_function: Literal['softmax', 'sigmoid'] = "softmax"
     """Score function for MoE routing. Can be "softmax" or "sigmoid"."""
 
-    moe_router_dtype: Optional[str] = None
+    moe_router_dtype: Optional[Literal['fp32', 'fp64']] = None
     """Data type for routing and expert output weighted averaging. Using fp32 or fp64 can
     improve stability especially when the number of experts is large (e.g. finegrained-moe).
     None means no changes for dtype."""
@@ -643,14 +672,14 @@ class TransformerConfig(ModelParallelConfig):
     specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is
     currently unsupported so should remain False."""
 
-    moe_token_dispatcher_type: str = "allgather"
+    moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex'] = "allgather"
     """The type of token dispatcher to use. The default is 'allgather'.
     Options are 'allgather','alltoall' and 'flex'."""
 
     moe_enable_deepep: bool = False
     """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models."""
 
-    moe_flex_dispatcher_backend: str = "deepep"
+    moe_flex_dispatcher_backend: Literal['deepep', 'hybridep'] = "deepep"
     """[Experimental] The backend to use for flex token dispatcher. The default is "deepep".
     Options are "deepep" and "hybridep". Currently only "hybridep" backend supports 
     the MNNVL case."""
@@ -667,7 +696,7 @@ class TransformerConfig(ModelParallelConfig):
     the expert capacity length, effective only after the moe_expert_capacity_factor is set. The
     default setting is False."""
 
-    moe_token_drop_policy: str = "probs"
+    moe_token_drop_policy: Literal['probs', 'position'] = "probs"
     """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with
     the lowest probabilities will be dropped. If "position", tokens at the end of each batch will
     be dropped.
@@ -680,7 +709,9 @@ class TransformerConfig(ModelParallelConfig):
     """Fuse token rearrangement ops during token dispatching."""
 
     moe_router_fusion: bool = False
-    """Fuse ops in routing and aux loss calculation."""
+    """Enable fusion for MoE TopK routing and aux-loss computation. This is only
+    supported in TransformerEngine 2.7.0 and above.
+    """
 
     moe_apply_probs_on_input: bool = False
     """Apply probs on input of experts instead of applying after activation and glu."""
@@ -742,7 +773,7 @@ class TransformerConfig(ModelParallelConfig):
     """DEPRECATED and replaced by cuda_graph_impl.
     When set to true, TransformerLayer layers are swapped with user provided CUDA graphs."""
 
-    cuda_graph_impl: str = "none"
+    cuda_graph_impl: Literal['none', 'local', 'transformer_engine'] = "none"
     """Determines the CUDA graph capture implementation.
     "none": no CUDA graph.
     "local": capture the CUDA graph using MCore local implementation. Either partial CUDA graph
@@ -794,8 +825,10 @@ class TransformerConfig(ModelParallelConfig):
     inference_sampling_seed: int = 42
     """ Random seed to use for sampling during inference. """
 
-    symmetric_ar_type: Optional[str] = None
-    """Type of symmetric all reduce to use"""
+    symmetric_ar_type: Optional[Literal['two_shot', "one_shot", "multimem_all_reduce"]] = None
+    """What type of symmetric all reduce to use. The default is None
+    which is no use of symmetric memory.
+    """
 
     use_inference_optimized_layers: bool = False
     """If True, use inference optimized transformer layers during inference."""
@@ -823,8 +856,10 @@ class TransformerConfig(ModelParallelConfig):
     """The number of heads used in Mamba layers.
     If None, the number of heads will be hidden_size * expand // mamba_head_dim."""
 
-    use_mamba_mem_eff_path: bool = True
-    """If True, use the memory efficient path for Mamba layers."""
+    use_mamba_mem_eff_path: bool = field(
+        default=True, metadata={"argparse_meta": {"arg_names": ["--disable-mamba-mem-eff-path"]}}
+    )
+    """Controls usage of the memory efficient path for Mamba layers."""
 
     mlp_chunks_for_prefill: int = 1
     """The number of chunks along the sequence dimension to use for MLP computation
@@ -842,7 +877,9 @@ class TransformerConfig(ModelParallelConfig):
     quant_recipe: Optional[RecipeConfig] = None
     """Configuration of any per-module quantization settings to be applied to the model"""
 
-    transformer_impl: str = "transformer_engine"
+    transformer_impl: Literal['local', 'transformer_engine', 'inference_optimized'] = (
+        "transformer_engine"
+    )
     """Transformer implementation to use.
     Options are 'transformer_engine' for Transformer Engine and 'local' for MCore."""
 
@@ -854,7 +891,7 @@ class TransformerConfig(ModelParallelConfig):
     Fine-grained activation offloading is a module-level offloading method
     instead of a layer-level offloading method like cpu_offloading."""
 
-    offload_modules: Optional[list[str]] = None
+    offload_modules: Optional[list[str]] = field(default_factory=list)
     """The submodules to offload its input.
     choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj",
              "mlp_norm", "expert_fc1", "moe_act".
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index f269ad02879..5246f44d206 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -77,7 +77,6 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
     parser = _add_one_logger_args(parser)
     parser = _add_inprocess_restart_args(parser)
     parser = _add_ft_package_args(parser)
-    parser = _add_config_logger_args(parser)
     parser = _add_rerun_machine_args(parser)
     parser = _add_msc_args(parser)
     parser = _add_kitchen_quantization_arguments(parser)
@@ -1322,8 +1321,6 @@ def core_transformer_config_from_args(args, config_class=None):
         if hasattr(args, f.name):
             kw_args[f.name] = getattr(args, f.name)
     kw_args['persist_layer_norm'] = not args.no_persist_layer_norm
-    kw_args['layernorm_zero_centered_gamma'] = args.apply_layernorm_1p
-    kw_args['layernorm_epsilon'] = args.norm_epsilon
     kw_args['deallocate_pipeline_outputs'] = True
     kw_args['pipeline_dtype'] = args.params_dtype
     kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm
@@ -1397,66 +1394,12 @@ def core_transformer_config_from_args(args, config_class=None):
 def _add_transformer_engine_args(parser):
     group = parser.add_argument_group(title='Transformer-Engine')
 
-    group.add_argument('--fp8-format', default=None,
-                       choices=['e4m3', 'hybrid'],
-                       help='Which fp8 format scheme to use for FP8 tensors in the forward and backward pass',
-                       dest='fp8')
-    # per tensor current scaling recipe selection
-    group.add_argument('--fp8-recipe', default='delayed',
-                       choices=['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom'],
-                       help='Which fp8 recipe to use for FP8 tensors in the forward and backward pass',
-                       dest='fp8_recipe')
-    group.add_argument('--fp8-quantizer-factory', default=None,
-                       help='Python import path to a callable quantizer factory, '
-                            'e.g., package.module.quantizer_factory.',
-                       dest='fp8_quantizer_factory')
     # delayed scaling only configs
-    group.add_argument('--fp8-margin', type=int, default=0,
-                       help='Scaling margin for fp8',
-                       dest='fp8_margin')
-    group.add_argument('--fp8-interval', type=int, default=1,
-                       help='DEPRECATED. This flag is ignored. Scaling update interval for fp8',
-                       dest='fp8_interval')
-    group.add_argument('--fp8-amax-history-len', type=int, default=1,
-                       help='Number of steps for which amax history is recorded per tensor',
-                       dest='fp8_amax_history_len')
-    group.add_argument('--fp8-amax-compute-algo', default='most_recent',
-                       choices=['most_recent', 'max'],
-                       help='Algorithm for computing amax from history',
-                       dest='fp8_amax_compute_algo')
-    group.add_argument('--no-fp8-wgrad', action='store_false',
-                       help='Execute wgrad in higher precision even for FP8 runs',
-                       dest='fp8_wgrad')
-    group.add_argument('--transformer-impl', default='transformer_engine',
-                       choices=['local', 'transformer_engine', 'inference_optimized'],
-                       help='Which Transformer implementation to use.')
     group.add_argument('--fp8-param-gather', action='store_true',
                        help='Keep the compute param in fp8 (do not use any other intermediate '
                             'dtype) and perform the param all-gather in fp8.')
-    group.add_argument('--first-last-layers-bf16', action='store_true',
-                       help='Construct first and last layers in bf16 when doing FP8 training.')
-    group.add_argument('--num-layers-at-start-in-bf16', type=int, default=1,
-                       help='Number of layers at start to construct in bf16 when --first-last-layers-bf16 is enabled.')
-    group.add_argument('--num-layers-at-end-in-bf16', type=int, default=1,
-                       help='Number of layers at end to construct in bf16 when --first-last-layers-bf16 is enabled.')
 
     # FP4 related arguments
-    group.add_argument('--fp4-format', default=None,
-                       choices=['e2m1'],
-                       help='Which nvfp4 format scheme to use for FP4 tensors in the forward and backward pass',
-                       dest='fp4')
-    group.add_argument('--fp4-recipe', default='nvfp4',
-                       choices=['nvfp4', 'custom'],
-                       help='Which fp4 recipe to use for FP4 tensors in the forward and backward pass',
-                       dest='fp4_recipe')
-    group.add_argument('--fp4-quantizer-factory', default=None,
-                       help='Python import path to a callable quantizer factory, '
-                            'e.g., package.module.quantizer_factory.',
-                       dest='fp4_quantizer_factory')
-    group.add_argument('--fp4-param-gather', action='store_true',
-                       help='Keep the compute param in fp4 (do not use any other intermediate '
-                            'dtype) and perform the param all-gather in fp4.',
-                       dest='fp4_param')
     group.add_argument('--te-precision-config-file', default=None,
                        help='Configuration file to select per-module precision overrides. '
                        'See TransformerEngineMixedPrecision.md')
@@ -1484,24 +1427,6 @@ def _add_inference_args(parser):
                        choices=["megatron", "huggingface"],
                        help='Select either Megatron or Huggingface as the '
                        'Bert embedder.')
-    group.add_argument('--flash-decode', default=False, action="store_true",
-                       help='Whether to use the flash decoding kernel.')
-    group.add_argument('--enable-cuda-graph', default=False, action="store_true",
-                       help='Deprecated. Use --cuda-graph-impl=local instead. '
-                       'Use local implementation of CUDA graph capture and replay. '
-                       '--cuda-graph-scope=\"full_iteration\" enables whole iteration CUDA graph. ')
-    group.add_argument("--cuda-graph-warmup-steps", type=int, default=3,
-                       help="Number of CUDA graph warmup steps")
-    group.add_argument('--external-cuda-graph', action='store_true',
-                       help='Deprecated. Use --cuda-graph-impl=transformer_engine instead. '
-                       'Use TE make_graphed_callables() to capture the CUDA graph. '
-                       'Use --cuda-graph-scope=\"attn\", \"mlp\", \"moe\", \"moe_router\", \"moe_preprocess\", \"mamba\" for partial capture. ')
-    group.add_argument('--cuda-graph-impl', type=str, default='none',
-                       choices=['none', 'local', 'transformer_engine'],
-                       help='Determines the CUDA graph capture implementation. '
-                       '"none": no CUDA graph. '
-                       '"local": capture the CUDA graph using MCore local implementation. --cuda-graph-scope=\"full_iteration\" enables whole iteration CUDA graph. '
-                       '"transformer_engine": capture the CUDA graph using TE make_graphed_callables().')
     group.add_argument('--cuda-graph-scope', nargs='+', type=lambda scope: CudaGraphScope[scope] if scope != "full" else scope, default=[],
                        help='Determines the CUDA graphs capturing scope. '
                        'choices: "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba", "full_iteration". '
@@ -1583,15 +1508,9 @@ def _add_inference_args(parser):
                        '1) allocate `memory_buffer` in unified memory. '
                        'Eventually, additional levels will be included to '
                        'control other tensors within the context.')
-    group.add_argument('--symmetric-ar-type', type=str, default=None,
-                       choices=['two_shot', "one_shot", "multimem_all_reduce", None],
-                       help='What type of symmetric all reduce to use. The default is none which is no use of symetric memory')
     group.add_argument('--nccl-all-reduce-for-prefill',
                        action='store_true', default=False,
                        help='When using symmeric all reduce kernels this will use regular nccl kernels for prefill. This can be more effecient when prefill is large as the nccl kernels can be more bandwith optimized')
-    group.add_argument('--mlp-chunks-for-prefill', type=int, default=1,
-                       help='Number of chunks along sequence dimension for MLP '
-                       'computation during prefill')
     # TODO(ksanthanam): Clean this up in future PR
     group.add_argument('--enable-chunked-prefill', dest='disable_chunked_prefill',
                        action='store_false', default=True,
@@ -1611,43 +1530,105 @@ def _add_inference_args(parser):
                        required=False, default=False, help='Enable inference wandb logging.')
     group.add_argument("--inference-coordinator-port", type=int, default=12346,
                        help="This port will be used to setup the inference coordinator on node-0")
-    group.add_argument("--inference-fuse-tp-communication", action="store_true", default=False,
-                       help="Use the fused communication kernel for tensor parallelism during inference. This "
-                       "kernel fuses reduce-scatter + residual-add + rms-norm + all-gather into one operation.")
     return parser
 
 
 def _add_network_size_args(parser):
+    exclude = [
+        # cannot provide callables over CLI
+        "timers",
+        "finalize_model_grads_func",
+        "grad_scale_func",
+        "no_sync_func",
+        "grad_sync_func",
+        "param_sync_func",
+        "_cpu_offloading_context",
+        "init_method",
+        "output_layer_init_method",
+        "embedding_init_method",
+        "activation_func",
+        # types affect docstring
+        "pipeline_model_parallel_layout",
+        "window_size",
+        "window_attn_skip_freq",
+        "no_rope_freq",
+        "moe_layer_freq",
+        "linear_attention_freq",
+        "moe_router_load_balancing_type",
+        "moe_aux_loss_coeff",
+        "cp_comm_type",
+        "cuda_graph_scope",
+        # no CLI argument exists for these
+        "virtual_pipeline_model_parallel_size",
+        "params_dtype",
+        "enable_autocast",
+        "autocast_dtype",
+        "num_microbatches_with_partial_activation_checkpoints",
+        "tp_comm_overlap_disable_qkv",
+        "tp_comm_overlap_disable_fc1",
+        "pipeline_dtype",
+        "variable_seq_lengths",
+        "batch_p2p_comm",
+        "batch_p2p_sync",
+        "deallocate_pipeline_outputs",
+        "cpu_offloading",
+        "cpu_offloading_activations",
+        "cpu_offloading_weights",
+        "cpu_offloading_double_buffering",
+        "num_layers_in_first_pipeline_stage",
+        "num_layers_in_last_pipeline_stage",
+        "softmax_scale",
+        "gated_linear_unit",
+        "bias_activation_fusion",
+        "activation_func_fp8_input_store",
+        "test_mode",
+        "memory_efficient_layer_norm",
+        "fused_single_qkv_rope",
+        "fp8_dot_product_attention",
+        "fp8_multi_head_attention",
+        "tp_only_amax_red",
+        "use_kitchen",
+        "moe_token_dropping",
+        "cuda_graph_use_single_mempool",
+        "cuda_graph_retain_backward_graph",
+        "disable_parameter_transpose_cache",
+        "inference_sampling_seed",
+        "use_inference_optimized_layers",
+        "heterogeneous_block_specs",
+        "hetereogenous_dist_checkpoint",
+        "quant_recipe",
+        # deprecated and no CLI arg exists
+        "tp_comm_atomic_ag",
+        "tp_comm_atomic_rs",
+        "moe_router_topk_limited_devices",
+        # already generated by another config
+        "inference_rng_tracker",
+        "use_te_rng_tracker",
+        "log_max_attention_logit",
+        "barrier_with_L1_time",
+        # args uses same var with a different name
+        "num_moe_experts",
+        "fp8_param",
+        # incompatible defaults in dataclass
+        "gradient_accumulation_fusion",
+        "overlap_p2p_comm",
+        "attention_softmax_in_fp32",
+        "masked_softmax_fusion",
+        "persist_layer_norm",
+        "bias_dropout_fusion",
+        "apply_rope_fusion",
+    ]
+    transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude)
+    transformer_group = transformer_factory.build_group(parser, "transformer configuration")
+
     group = parser.add_argument_group(title='network size')
 
-    group.add_argument('--num-layers', type=int, default=None,
-                       help='Number of transformer layers.')
     group.add_argument('--encoder-num-layers', type=int, default=None,
                        help='Number of encoder transformer layers.')
     group.add_argument('--decoder-num-layers', type=int, default=None,
                        help='Number of decoder transformer layers.')
-    group.add_argument('--hidden-size', type=int, default=None,
-                       help='Transformer hidden size.')
-    group.add_argument('--ffn-hidden-size', type=int, default=None,
-                       help='Transformer Feed-Forward Network hidden size. '
-                       'This is set to 4*hidden-size if not provided')
-    group.add_argument('--num-attention-heads', type=int, default=None,
-                       help='Number of transformer attention heads.')
-    group.add_argument('--attention-backend', type=lambda attn_backend: AttnBackend[attn_backend], default=AttnBackend.auto, choices = list(AttnBackend), help='Attention backend to use (flash,fused,unfused,local,auto). Defaults to auto')
-    group.add_argument('--kv-channels', type=int, default=None,
-                       help='Projection weights dimension in multi-head '
-                       'attention. This is set to '
-                       '   args.hidden_size // args.num_attention_heads '
-                       'if not provided.')
     group.add_argument('--group-query-attention', action='store_true',
                           help='Use group-query attention.')
-    group.add_argument('--num-query-groups', type=int, default=1)
-    group.add_argument('--attention-output-gate', action='store_true',
-                       help='Whether to apply output gate to the attention.')
-    group.add_argument('--softmax-type', type=str, default='vanilla',
-                       choices=['learnable', 'vanilla', 'off-by-one'],
-                       help='Type of softmax to use for the attention. Supports both a fixed offset and '
-                       'learnable offset.')
     group.add_argument('--window-size', type=tuple_type, default=None,
                        help='Window size for window attention. If not provided, '
                             'window attention will be disabled.')
@@ -1675,8 +1656,6 @@ def _add_network_size_args(parser):
                        help='Base to use for rotary positional embeddings, default 10000')
     group.add_argument('--rotary-percent', type=float, default=1.0,
                        help='Percent of rotary dimension to use, default 100%%')
-    group.add_argument('--rotary-interleaved', action='store_true',
-                          help='Use interleaved rotary embedding.')
     group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None,
                        help='Sequence length interpolation factor for rotary embeddings.')
     group.add_argument('--use-rope-scaling', action='store_true',
@@ -1696,23 +1675,9 @@ def _add_network_size_args(parser):
                        action='store_false',
                        help='Disable position embedding. Deprecated: use --position-embedding-type',
                        dest='add_position_embedding')
-    group.add_argument('--mrope-section', nargs='+', type=int, default=None,
-                       help='Multimodal rope section is for channel dimension, empty by default.')
     group.add_argument('--make-vocab-size-divisible-by', type=int, default=128,
                        help='Pad the vocab size to be divisible by this value.'
                        'This is added for computational efficieny reasons.')
-    group.add_argument('--normalization', default='LayerNorm',
-                       choices=['LayerNorm', 'RMSNorm'],
-                       help='Which normalization technique to use.')
-    group.add_argument('--norm-epsilon', type=float, default=1e-5,
-                       help='Epsilon for layer norm and RMS norm.')
-    group.add_argument('--apply-layernorm-1p', action='store_true',
-                       help='Adjust LayerNorm weights such that they are centered '
-                       'around zero. This improves numerical stability.')
-    group.add_argument('--apply-residual-connection-post-layernorm',
-                       action='store_true',
-                       help='If set, use original BERT residula connection '
-                       'ordering.')
     group.add_argument('--openai-gelu', action='store_true',
                        help='Use OpenAIs GeLU implementation. This option'
                        'should not be used unless for backward compatibility'
@@ -1723,12 +1688,6 @@ def _add_network_size_args(parser):
                        help='Use gated linear units and SiLU activation instead of default gelu')
     group.add_argument('--quick-geglu', action='store_true',
                        help='Use quick geglu activation instead of default gelu')
-    group.add_argument('--activation-func-clamp-value', type=float, default=None,
-                       help='Clamp the output of the linear_fc1 in the activation function. Only used when '
-                            'activation_func is quick_gelu.')
-    group.add_argument('--glu-linear-offset', type=float, default=0.0,
-                       help='Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). '
-                            'Only used when gated_linear_unit is True')
     group.add_argument('--onnx-safe', type=bool, required=False,
                        help='Use workarounds for known problems with '
                        'Torch ONNX exporter')
@@ -1737,20 +1696,6 @@ def _add_network_size_args(parser):
                        dest='bert_binary_head')
     group.add_argument('--untie-embeddings-and-output-weights', action='store_true',
                        help='Untie embeddings and output weights.')
-    group.add_argument('--multi-latent-attention', action='store_true',
-                       help='Use multi-latent attention for model.')
-    group.add_argument('--mtp-num-layers', type=int, default=None,
-                       help='Number of Multi-Token Prediction (MTP) Layers.'
-                       'MTP extends the prediction scope to multiple future tokens at each position.'
-                       'This MTP implementation sequentially predict additional tokens '
-                       'by using D sequential modules to predict D additional tokens.')
-    group.add_argument('--mtp-loss-scaling-factor', type=float, default=0.1,
-                       help='Scaling factor of Multi-Token Prediction (MTP) loss. '
-                       'We compute the average of the MTP losses across all depths, '
-                       'and multiply it the scaling factor to obtain the overall MTP loss, '
-                       'which serves as an additional training objective.')
-    group.add_argument('--moe-latent-size', type=int, default=None,
-                       help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
     return parser
 
 def _add_straggler_detector_args(parser):
@@ -1861,14 +1806,6 @@ def _add_ft_package_args(parser):
     return parser
 
 
-def _add_config_logger_args(parser):
-    group = parser.add_argument_group(title='config logger')
-    group.add_argument('--config-logger-dir', type=str, default='',
-                       help='If set, will dump all configs to --config-logger-dir',
-                       dest='config_logger_dir')
-    return parser
-
-
 def _add_logging_args(parser):
     from megatron.training.training_config import LoggerConfig
 
@@ -1881,10 +1818,6 @@ def _add_logging_args(parser):
 def _add_regularization_args(parser):
     group = parser.add_argument_group(title='regularization')
 
-    group.add_argument('--attention-dropout', type=float, default=0.1,
-                       help='Post attention dropout probability.')
-    group.add_argument('--hidden-dropout', type=float, default=0.1,
-                       help='Dropout probability for hidden state transformer.')
     group.add_argument('--weight-decay', type=float, default=0.01,
                        help='Weight decay coefficient for L2 regularization.')
     group.add_argument('--apply-wd-to-qk-layernorm', action='store_true',
@@ -2078,95 +2011,19 @@ def _add_training_args(parser):
     group.add_argument('--recompute-activations', action='store_true',
                        help='recompute activation to allow for training '
                        'with larger models, sequences, and batch sizes.')
-    group.add_argument('--recompute-granularity', type=str, default=None,
-                       choices=['full', 'selective'],
-                       help='Checkpoint activations to allow for training '
-                       'with larger models, sequences, and batch sizes. '
-                       'It is supported at two granularities 1) full: '
-                       'whole transformer layer is recomputed, '
-                       '2) selective: submodules set in --recompute-modules '
-                       'are recomputed, default is core_attn.')
     group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false',
                        help='Check for NaNs in loss and grad',
                        dest='check_for_nan_in_loss_and_grad')
     group.add_argument('--check-for-large-grads', action='store_true',
                        help='Check for unexpectedly large grads',
                        dest='check_for_large_grads')
-    group.add_argument('--distribute-saved-activations',
-                       action='store_true',
-                       help='If set, distribute recomputed activations '
-                       'across model parallel group.')
-    group.add_argument('--recompute-method', type=str, default=None,
-                       choices=['uniform', 'block'],
-                       help='1) uniform: uniformly divide the total number of '
-                       'Transformer layers and recompute the input activation of '
-                       'each divided chunk at specified granularity, '
-                       '2) recompute the input activations of only a set number of '
-                       'individual Transformer layers per pipeline stage and do the '
-                       'rest without any recomputing at specified granularity'
-                       'default) do not apply activations recompute to any layers')
-    group.add_argument('--recompute-num-layers', type=int, default=None,
-                       help='1) uniform: the number of Transformer layers in each '
-                       'uniformly divided recompute unit, '
-                       '2) block: the number of individual Transformer layers '
-                       'to recompute within each pipeline stage.')
-    group.add_argument('--recompute-modules', nargs='*', type=str, default=None,
-                       help='The submodules to recompute. '
-                       'choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", '
-                       '         "mlp", "moe", "shared_experts". '
-                       'default: ["core_attn"].'
-                       '"core_attn": recompute the core attention part of the transformer layer. '
-                       '"moe_act": recompute the MoE MLP activation function. '
-                       '"layernorm": recompute the input_layernorm and pre_mlp_layernorm. '
-                       '"mla_up_proj": recompute the MLA up projection and RoPE applying parts.'
-                       '"mlp": recompute the dense MLP layer.'
-                       '"moe": recompute the MoE layer.'
-                       '"shared_experts": recompute the shared experts in the MoE layer.'
-                       '"moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, '
-                       '"core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing.')
-    group.add_argument('--cpu-offloading-num-layers', type=int, default=0,
-                       help='The number of Transformer layers to offload to CPU.')
-    group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false',
-                       help='If not set, clone the output of the scatter in embedding layer to GC original tensor.',
-                       dest='clone_scatter_output_in_embedding')
     group.add_argument('--result-rejected-tracker-filename', type=str, default=None,
                        help='Optional name of file tracking `result_rejected` events.')
     group.add_argument('--disable-gloo-process-groups', action='store_false',
                        dest='enable_gloo_process_groups',
                        help='Disables creation and usage of Gloo process groups.')
-    group.add_argument('--tp-comm-overlap', action='store_true', help='Enables the '
-                       ' overlap of Tensor parallel communication and GEMM kernels.')
     group.add_argument('--tp-comm-overlap-cfg', type=str, default=None,
                        help='Config file when tp_comm_overlap is enabled.')
-    group.add_argument('--disable-tp-comm-overlap-ag', action='store_false',
-                       help=('Disables the All-Gather overlap with GEMM by '
-                             'pipelining the GEMM and All-Gather.'),
-                       dest='tp_comm_overlap_ag')
-    group.add_argument('--disable-tp-comm-overlap-rs', action='store_false',
-                       help=('Disables the Reduce-Scatter overlap with GEMM by '
-                             'pipelining the GEMM and Reduce-Scatter.'),
-                       dest='tp_comm_overlap_rs')
-    group.add_argument('--tp-comm-overlap-rs-dgrad', action='store_true',
-                       help = 'Enables the Reduce-Scatter overlap with dgrad GEMM.',
-                       dest='tp_comm_overlap_rs_dgrad')
-    group.add_argument('--disable-tp-comm-bulk-dgrad', action='store_false',
-                       help='Disables the All-Gather overlap with bprop activation gradient GEMM.',
-                       dest='tp_comm_bulk_dgrad')
-    group.add_argument('--disable-tp-comm-bulk-wgrad', action='store_false',
-                       help='Disables the Reduce-Scatter overlap with bprop weight gradient GEMM.',
-                       dest='tp_comm_bulk_wgrad')
-    group.add_argument('--tp-comm-bootstrap-backend', default='nccl', type=str,
-                       choices=['nccl', 'mpi', 'gloo'],
-                       help='Set the bootstrapping backend of Tensor parallel communications.')
-    group.add_argument('--use-cpu-initialization', action='store_true',
-                       default=None,
-                       help='If set, initialize weights on the CPU. This eliminates init differences based on tensor parallelism.')
-    group.add_argument('--deterministic-mode', action='store_true',
-                       help='Choose code that has deterministic execution. This usually '
-                       'means slower execution, but is good for debugging and testing.')
-    group.add_argument('--calculate-per-token-loss', action='store_true',
-                       help=('Scale cross entropy loss by the number of non-padded tokens in the '
-                             'global batch, versus the default behavior of assuming all tokens are non-padded.'))
 
     # deprecated
     group.add_argument('--checkpoint-activations', action='store_true',
@@ -2184,8 +2041,6 @@ def _add_training_args(parser):
                        help='Disable bias and swiglu fusion, the fusion is '
                        'available only when using megatron-core.',
                        dest='bias_swiglu_fusion')
-    group.add_argument('--use-fused-weighted-squared-relu', action='store_true',
-                       help='Use fused weighted squared relu when using MoE.')
     group.add_argument('--no-bias-dropout-fusion', action='store_false',
                        help='Disable bias and dropout fusion.',
                        dest='bias_dropout_fusion')
@@ -2197,27 +2052,9 @@ def _add_training_args(parser):
                       choices=['rope', 'yarn'],
                       help='Type of rope to use. Note that MLA takes yarn by default, '
                       'and common attention takes rope by default.')
-    group.add_argument('--cross-entropy-loss-fusion', action='store_true',
-                       help='Enabled fusion of cross entropy loss calculation.',
-                       dest='cross_entropy_loss_fusion')
-    group.add_argument('--cross-entropy-fusion-impl', type=str, default='native',
-                       choices=['native', 'te'],
-                       help='Implementation of cross entropy loss calculation.')
     group.add_argument('--use-flash-attn', action='store_true',
                        help='use FlashAttention implementation of attention. '
                        'https://arxiv.org/abs/2205.14135')
-    group.add_argument('--disable-bias-linear', action='store_false',
-                       help='Disable bias in the linear layers',
-                       dest='add_bias_linear')
-    group.add_argument('--add-qkv-bias', action='store_true',
-                       help='Enable bias only in the QKV linear layers',
-                       dest='add_qkv_bias')
-    group.add_argument('--qk-clip', action='store_true',
-                       help='Whether to use qk-clip for training stabilization, strongly recommended for Muon.')
-    group.add_argument('--qk-clip-alpha', type=float, default=0.5,
-                       help='The balancing alpha for qk-clip.')
-    group.add_argument('--qk-clip-threshold', type=float, default=100,
-                       help='The balancing threshold for qk-clip.')
     group.add_argument('--optimizer', type=str, default='adam',
                        choices=['adam', 'sgd', 'muon', 'dist_muon'],
                        help='Optimizer function')
@@ -2240,17 +2077,11 @@ def _add_training_args(parser):
     group.add_argument('--dataloader-type', type=str, default=None,
                        choices=['single', 'cyclic', 'external'],
                        help='Single pass vs multiple pass data loader')
-    group.add_argument('--no-async-tensor-model-parallel-allreduce',
-                       action='store_false',
-                       help='DEPRECATED. This flag is ignored.',
-                       dest='async_tensor_model_parallel_allreduce')
     group.add_argument('--no-persist-layer-norm', action='store_true',
                        help='Disable using persistent fused layer norm kernel. '
                        'This kernel supports only a set of hidden sizes. Please '
                        'check persist_ln_hidden_sizes if your hidden '
                        'size is supported.')
-    group.add_argument('--sequence-parallel', action='store_true',
-                       help='Enable sequence parallel optimization.')
     group.add_argument('--no-gradient-accumulation-fusion',
                        action='store_false',
                        help='Disable fusing gradient accumulation to weight '
@@ -2263,32 +2094,8 @@ def _add_training_args(parser):
                        '--use-legacy-models to not use core models.')
     group.add_argument('--use-legacy-models', action='store_true',
                        help='Use the legacy Megatron models, not Megatron-Core models.')
-    group.add_argument('--disable-tp-comm-split-ag', action='store_false',
-                       help='Disables the All-Gather overlap with fprop GEMM.',
-                       dest='tp_comm_split_ag')
-    group.add_argument('--disable-tp-comm-split-rs', action='store_false',
-                       help='Disables the Reduce-Scatter overlap with fprop GEMM.',
-                       dest='tp_comm_split_rs')
-    group.add_argument('--pipeline-model-parallel-comm-backend', type=str, default=None,
-                       choices=['nccl', 'ucc'],
-                       help='Select a communicator backend for pipeline parallel communication. '
-                       'If None, the default backend will be used.')
     group.add_argument('--high-priority-stream-groups', nargs='*', type=str, default=[],
                        help='The communicator group names to use high priority streams.')
-    group.add_argument('--use-te-activation-func', action='store_true',
-                       help='Use activation function kernel from Transformer Engine in MLP module.')
-    group.add_argument('--fine-grained-activation-offloading', action='store_true',
-                       help='Enable fine-grained activation offloading.')
-    group.add_argument('--offload-modules', nargs='*', type=str, default=[],
-                       help='The submodules to offload its input. Choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act".')
-    group.add_argument('--min-offloaded-tensor-size', type=int, default=1024*1024,
-                       help='The minimum size of the tensor to be offloaded.')
-    group.add_argument('--batch-invariant-mode', action='store_true',
-                       help='Use batch-invariant kernels for deterministic forward execution regardless '
-                       'of batch size. Ensures bitwise identical results when the same inputs are '
-                       'processed in different batch configurations. This is more strict than deterministic-mode '
-                       'which only ensures bitwise identical results when the same inputs are processed in the same batch configuration. '
-                       'This will significantly affect speed of training and inference as the kernels are not full optimized.')
     group.add_argument('--disable-jit-fuser', action='store_true',
                        help='Disable the JIT fuser.')
 
@@ -2310,19 +2117,6 @@ def _add_initialization_args(parser):
     rng_factory = ArgumentGroupFactory(RNGConfig)
     group = rng_factory.build_group(parser, "RNG and initialization")
 
-    group.add_argument('--init-method-std', type=float, default=0.02,
-                       help='Standard deviation of the zero mean normal '
-                       'distribution used for weight initialization.')
-    group.add_argument('--embedding-init-method-std', type=float, default=None,
-                       help='Standard deviation of the zero mean normal '
-                       'distribution used for embedding weight initialization. '
-                       'If unset, embeddings will be initialized the same way '
-                       'as other weights. Setting this to a value around 1.0 '
-                       'may avoid loss spikes in training. Setting this to any '
-                       'value will also skip applying weight decay on embedding '
-                       'weights to avoid shrinkage towards zero. See '
-                       'https://arxiv.org/abs/2312.16903 for more details.'
-                       )
     group.add_argument('--init-method-xavier-uniform', action='store_true',
                        help='Enable Xavier uniform parameter initialization')
 
@@ -2368,15 +2162,9 @@ def _add_checkpointing_args(parser):
                        help='Do not load optimizer when loading checkpoint.')
     group.add_argument('--no-load-rng', action='store_true', default=None,
                        help='Do not load rng state when loading checkpoint.')
-    group.add_argument('--no-initialization', action='store_false',
-                       help='Do not perform initialization when building model, '
-                       'can reduce startup time when definitely loading from a '
-                       'checkpoint',
-                       dest='perform_initialization')
     group.add_argument('--use-dist-ckpt', action='store_true',
                        dest='use_dist_ckpt_deprecated',
                        help='Deprecated: see --ckpt-format.')
-
     group.add_argument('--dist-ckpt-format',
                        dest='dist_ckpt_format_deprecated',
                        help='Deprecated: see --ckpt-format.')
@@ -2389,10 +2177,6 @@ def _add_checkpointing_args(parser):
 def _add_mixed_precision_args(parser):
     group = parser.add_argument_group(title='mixed precision')
 
-    group.add_argument('--fp16', action='store_true',
-                       help='Run model in fp16 mode.')
-    group.add_argument('--bf16', action='store_true',
-                       help='Run model in bfloat16 mode.')
     group.add_argument('--grad-reduce-in-bf16', action='store_true',
                        help='Reduce gradients in bfloat16.')
     group.add_argument('--loss-scale', type=float, default=None,
@@ -2407,11 +2191,6 @@ def _add_mixed_precision_args(parser):
                        help='Window over which to raise/lower dynamic scale.')
     group.add_argument('--hysteresis', type=int, default=2,
                        help='hysteresis for dynamic loss scaling')
-    group.add_argument('--fp32-residual-connection', action='store_true',
-                       help='Move residual connections to fp32.')
-    group.add_argument('--apply-query-key-layer-scaling', action='store_true',
-                       help='Scale Q * K^T by 1 / layer-number. '
-                       'Useful for fp16 training. Also sets `attention_softmax_in_fp32` to True.')
     group.add_argument('--attention-softmax-in-fp32', action='store_true',
                        help='Run attention masking and softmax in fp32.')
     group.add_argument('--accumulate-allreduce-grads-in-fp32',
@@ -2420,9 +2199,6 @@ def _add_mixed_precision_args(parser):
     group.add_argument('--fp16-lm-cross-entropy', action='store_true',
                        help='Move the cross entropy unreduced loss calculation'
                        'for lm head to fp16.')
-    group.add_argument('--disable-bf16-reduced-precision-matmul', action='store_true',
-                       help='If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to '
-                       'prevent matmul from using reduced precision accumulation when using BF16.')
     group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true',
                        help='If True, reuse the grad buffer for MXFP8 parameter all-gather.')
 
@@ -2432,10 +2208,6 @@ def _add_mixed_precision_args(parser):
 def _add_distributed_args(parser):
     group = parser.add_argument_group(title='distributed')
 
-    group.add_argument('--tensor-model-parallel-size', type=int, default=1,
-                       help='Degree of tensor model parallelism.')
-    group.add_argument('--pipeline-model-parallel-size', type=int, default=1,
-                       help='Degree of pipeline model parallelism.')
     group.add_argument('--decoder-first-pipeline-num-layers',
                        type=int, default=None,
                        help=('The number of transformer layers on the first pipeline stage of the decoder. '
@@ -2459,15 +2231,9 @@ def _add_distributed_args(parser):
                        help='Number of layers per virtual pipeline stage')
     group.add_argument('--num-virtual-stages-per-pipeline-rank', type=int, default=None,
                        help='Number of virtual pipeline stages per pipeline parallelism rank')
-    group.add_argument('--microbatch-group-size-per-virtual-pipeline-stage', type=int, default=None,
-                       help='Number of contiguous microbatches per virtual pipeline stage',
-                       dest='microbatch_group_size_per_vp_stage')
     group.add_argument('--no-overlap-p2p-communication', action='store_false',
                        help='overlap pipeline parallel communication with forward and backward chunks in 1F1B',
                        dest='overlap_p2p_comm')
-    group.add_argument('--overlap-p2p-communication-warmup-flush', action='store_true',
-                       default=False, help='if set, overlap pipeline parallel communication in warmup and flush',
-                       dest='overlap_p2p_comm_warmup_flush')
     group.add_argument('--distributed-backend', default='nccl',
                        choices=['nccl', 'gloo'],
                        help='Which backend to use for distributed training.')
@@ -2478,13 +2244,6 @@ def _add_distributed_args(parser):
                             'This timeout is applied to all process groups after initialization.')
     group.add_argument('--overlap-grad-reduce', action='store_true',
                        default=False, help='If set, overlap DDP grad reduce.')
-    group.add_argument('--defer-embedding-wgrad-compute', action='store_true',
-                       default=False, help='If set, defers the vocabulary projection linear layer weight'
-                       'gradient compute to pipeline flush.', dest='defer_embedding_wgrad_compute')
-    group.add_argument('--wgrad-deferral-limit', type=int, default=0, help='Number of micro-batches for which'
-                       'weight gradient computation of vocabulary projection is deferred, defaults to 0 which'
-                       'means all the micro-batches are deferred. Invalid if `defer-embedding-wgrad-compute`'
-                       'is not set')
     group.add_argument('--no-align-grad-reduce', action='store_false',
                        help='If not set, all PP stages will launch gradient reduces simultaneously. '
                        'Otherwise, each PP stage will independently launch as needed.',
@@ -2515,10 +2274,6 @@ def _add_distributed_args(parser):
     group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false',
                        help='If not set, use scatter/gather to optimize communication of tensors in pipeline.',
                        dest='scatter_gather_tensors_in_pipeline')
-    group.add_argument('--use-ring-exchange-p2p', action='store_true',
-                       default=False, help='If set, use custom-built ring exchange '
-                       'for p2p communications. Note that this option will require '
-                       'a custom built image that support ring-exchange p2p.')
     group.add_argument('--local-rank', type=int, default=int(os.getenv('LOCAL_RANK', '0')),
                        help='local rank passed from distributed launcher.')
     group.add_argument('--lazy-mpu-init', type=bool, required=False,
@@ -2527,12 +2282,6 @@ def _add_distributed_args(parser):
                        'complete it instead. Also turns on '
                        '--use-cpu-initialization flag. This is for '
                        'external DDP manager.' )
-    group.add_argument('--account-for-embedding-in-pipeline-split', action='store_true',
-                       default=False, help='If set, *input* embedding layer will be treated as a standard transformer'
-                       'layer in the context of partition and placement for pipeline parallelism.')
-    group.add_argument('--account-for-loss-in-pipeline-split', action='store_true',
-                       default=False, help='If set, loss layer will be treated as a standard transformer'
-                       'layer in the context of partition and placement for pipeline parallelism.')
     group.add_argument('--use-distributed-optimizer', action='store_true',
                        help='Use distributed optimizer.')
     group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub',
@@ -2557,7 +2306,6 @@ def _add_distributed_args(parser):
                    'to overlap reduce-scatter and all-gather operations.')
     group.add_argument('--use-megatron-fsdp', action='store_true',
                        help='Use the Megatron FSDP code path in DDP.')
-    group.add_argument('--init-model-with-meta-device', action='store_true')
     group.add_argument('--data-parallel-sharding-strategy', type=str, default='no_shard',
                        choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'],
                        help='Sharding strategy of data parallelism.')
@@ -2592,27 +2340,12 @@ def _add_distributed_args(parser):
     group.add_argument('--torch-fsdp2-no-reshard-after-forward', action='store_false', dest='torch_fsdp2_reshard_after_forward',
                        help='Whether to reshard weights after forward pass when using PyTorch FSDP2. '
                        'Set to enable FSDP ZeRO-2.')
-    group.add_argument('--context-parallel-size', type=int, default=1,
-                       help='Degree of context parallelism.')
     group.add_argument('--cp-comm-type', nargs='+', type=str, default=["p2p"],
                        help='Inter-gpu communication type for context parallelism: '
                        'p2p, a2a, allgather or a2a+p2p. If a single string is provided, '
                        'all layers will share the same communication type. Users can also '
                        'specify separated types for each layer like '
                        '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p')
-    group.add_argument('--hierarchical-context-parallel-sizes', nargs='+', type=int, default=None,
-                       help='Degrees of the hierarchical context parallelism. Users should '
-                       'provide a list to specify the sizes for different levels. '
-                       '--hierarchical-context-parallel-sizes 2 4 indicates every two adjacent gpus '
-                       'forms the first level of cp groups and the cp ranks with the same odevity '
-                       'forms the second level of cp groups.')
-    group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None,
-                       help='Maximum sequence length per DPxCP rank. This is used to calculate the '
-                       'number of sub-samples assigned to each DPxCP rank when using Hybrid Context Parallel.')
-    group.add_argument('--hybrid-context-parallel', action='store_true', default=False,
-                       help='Enables hybrid context parallel. This is used to balance the workload '
-                       'of each CP rank when we use packed samples with variable sequence lengths. '
-                       'Requires --max-seqlen-per-dp-cp-rank to be set.')
     group.add_argument('--nccl-communicator-config-path', type=str, default=None,
                        help='Path to the yaml file with NCCL communicator '
                        'configurations. The number of min/max thread groups and thread '
@@ -2952,21 +2685,11 @@ def _add_vision_args(parser):
     group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30,
                        help='warmup teacher temperaure epochs')
 
-    # regularization arguments
-    group.add_argument('--qk-layernorm', action='store_true',
-                       help='Whether to layer normalize the q and k attention embeddings.')
-    group.add_argument('--qk-l2-norm', action='store_true',
-                       help='Use llama 4 qk l2 norm')
-
     return parser
 
 def _add_moe_args(parser):
     group = parser.add_argument_group(title="moe")
     # General arguments
-    group.add_argument('--expert-model-parallel-size', type=int, default=1,
-                       help='Degree of expert model parallelism.')
-    group.add_argument('--expert-tensor-parallel-size', type=int, default=None,
-                       help='Degree of expert model parallelism. Default is None, which will be set to the value of --tensor-model-paralle-size.')
     group.add_argument('--num-experts', type=int, default=None,
                        help='Number of Experts in MoE (None means no MoE)')
     group.add_argument('--moe-layer-freq', type=moe_freq_type, default=1,
@@ -2977,34 +2700,6 @@ def _add_moe_args(parser):
                             'where 1 indicates an expert layer and 0 indicates a dense layer. '
                             'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, '
                             '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.')
-    group.add_argument('--moe-ffn-hidden-size', type=int, default=None,
-                       help='The hidden size of each expert\'s feed-forward network (ffn). '
-                       'If not specified, defaults to the ffn_hidden_size.')
-    group.add_argument('--moe-shared-expert-intermediate-size', type=int, default=None,
-                       help='Shared expert total ffn hidden size. '
-                       'It should be equal to "num_shared_experts * ffn_size_of_each_shared_expert" if there are multiple shared experts. '
-                       'None means no shared expert. '
-                       'By default, the shared experts execute before the router. However, when '
-                       '--moe-shared-expert-overlap or --overlap-moe-expert-parallel-comm is set, '
-                       'the shared experts execute after the router, before the routed experts. '
-                       'This makes the gradients from the router and the shared experts added in '
-                       'different orders to the hidden_states, causing minor numerical differences '
-                       'in the hidden_states gradient.')
-    group.add_argument('--moe-shared-expert-gate', action='store_true',
-                       help='Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set.')
-    group.add_argument('--moe-shared-expert-overlap', action='store_true',
-                       help='Enable overlapping between shared expert computations and dispatcher communications. '
-                       'Without this, the shared experts execute before the router. '
-                       'Only effective when moe-shared-expert-intermediate-size is set.')
-    group.add_argument('--moe-grouped-gemm', action='store_true',
-                       help='When there are multiple experts per rank, launch multiple local GEMM kernels in multiple streams to improve the utilization and performance with GroupedLinear in TransformerEngine.')
-    group.add_argument('--moe-use-legacy-grouped-gemm', action='store_true',
-                       help='Use legacy GroupedMLP rather than TEGroupedMLP. Note: The legacy one will be deprecated soon.')
-    group.add_argument('--moe-layer-recompute', action='store_true',
-                       help='Enable checkpointing for moe_layer, should be used when memory is not sufficient. '
-                       'Deprecated. Use "--recompute-granularity selective --recompute-modules moe" instead.')
-    group.add_argument('--moe-extended-tp', action='store_true',
-                       help='Deprecated. Use --expert-tensor-parallel-size instead.')
     group.add_argument('--moe-use-upcycling', action='store_true',
                        help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. '
                        'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.')
@@ -3013,94 +2708,10 @@ def _add_moe_args(parser):
                        choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'],
                        default='aux_loss',
                        help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".')
-    group.add_argument('--moe-router-dtype', type=str,
-                       choices=['fp32', 'fp64'],
-                       default=None,
-                       help='Data type for routing computation and expert output weighted averaging. '
-                            'Fp32/fp64 enhances numerical stability, especially with numerous experts. '
-                            'The perf impact should be negligible when used with permute fusion. '
-                            'None means no changes for dtype.')
-    group.add_argument('--moe-router-fusion', action='store_true',
-                       help='Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in TransformerEngine 2.7.0 and above.')
-    group.add_argument('--moe-router-score-function', type=str,
-                       choices=['softmax', 'sigmoid'],
-                       default='softmax',
-                       help='Score function for MoE TopK routing. Can be "softmax" or "sigmoid".')
-    group.add_argument('--moe-router-topk', type=int, default=2,
-                       help='Number of experts to route to for each token. The default is 2.')
-    group.add_argument('--moe-enable-routing-replay', action='store_true',
-                       help='Enable routing replay for MoE routers. When enabled, the router will '
-                            'use a pre-defined routing table instead of computing it on the fly.')
-    group.add_argument('--moe-router-pre-softmax', action='store_true',
-                       help='Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k.')
-    group.add_argument('--moe-router-num-groups', type=int, default=None,
-                       help='Number of groups to divide experts into for group-limited routing. When using group-limited routing: 1) Experts are divided into equal-sized groups, 2) For each token, a subset of groups are selected based on routing scores (sum of top-2 expert scores within each group), 3) From these selected groups, moe_router_topk experts are chosen.'
-                       'Two common use cases: 1) Device-limited routing: Set equal to expert parallel size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) 2) Node-limited routing: Set equal to number of nodes in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437)')
-    group.add_argument('--moe-router-group-topk', type=int, default=None,
-                       help='Number of selected groups for group-limited routing.')
-    group.add_argument('--moe-router-topk-scaling-factor', type=float, default=None,
-                       help='Scaling factor for routing score in top-k selection, only works when --moe-router-pre-softmax enabled. Defaults to None, which means no scaling.')
-    group.add_argument('--moe-router-enable-expert-bias', action='store_true',
-                       help='TopK routing with dynamic expert bias in the aux-loss-free load balancing strategy. '
-                       'The routing decision is based on the sum of the routing scores and the expert bias. '
-                       'See https://arxiv.org/abs/2408.15664 for details.')
-    group.add_argument('--moe-router-bias-update-rate', type=float, default=1e-3,
-                       help='Expert bias update rate in the aux-loss-free load balancing strategy. '
-                       'The expert bias is updated based on the number of assigned tokens to each expert in a global batch, '
-                       'where the bias is increased for the experts with less assigned tokens and decreased for the experts with more assigned tokens. '
-                       'The default value 1e-3 is same as that used in DeepSeekV3.')
-    group.add_argument('--moe-router-force-load-balancing', action='store_true',
-                       help='[Experimental] Force override routing to balance token distribution using random logits for MoE routers, supporting naive top-k and group-limited top-k. This experimental feature is for benchmarking purposes only!')
-    group.add_argument('--moe-router-padding-for-quantization', action='store_true',
-                       help='Pad the routing_map to make sure the number of tokens each expert received '
-                       'is a multiple of 16/32 for FP8/FP4 precision. It is suggested to enable this for '
-                       'dropless training with FP8/FP4 precision when num_local_experts > 1. This is a more '
-                       'efficient way to pad for FP8/FP4 which eliminates the explicit padding in the '
-                       'GroupedMLP layer.')
-    group.add_argument('--moe-router-padding-for-fp8', action='store_true',
-                       help='[Compatibility alias for --moe-router-padding-for-quantization] '
-                       'Enabling this will also enable --moe-router-padding-for-quantization.')
     group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0,
                        help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.')
-    group.add_argument('--moe-z-loss-coeff', type=float, default=None,
-                       help='Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended.')
-    group.add_argument('--moe-input-jitter-eps', type=float, default=None,
-                       help='Add noise to the input tensor by applying jitter with a specified epsilon value.')
-    group.add_argument('--moe-per-layer-logging', action='store_true',
-                       help='Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.')
     # Token dispatcher arguments
-    group.add_argument('--moe-token-dispatcher-type', type=str,
-                       choices=['allgather', 'alltoall', 'flex'],
-                       default='allgather',
-                       help="The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather', 'alltoall'. We recommend using 'alltoall' when applying expert parallelism. For more information, please refer to the documentation in core/moe/README.")
-    group.add_argument('--moe-enable-deepep', action='store_true',
-                       help='DEPRECATED: Please use --moe-flex-dispatcher-backend=deepep instead.')
-    group.add_argument('--moe-flex-dispatcher-backend', type=str,
-                       choices=['deepep', 'hybridep'],
-                       default='deepep',
-                       help='The backend to use for flex token dispatcher. The default is "deepep". Options are "deepep" and "hybridep".')
-    group.add_argument('--moe-deepep-num-sms', type=int, default=20,
-                       help='Number of SMs to use for DeepEP.')
-    group.add_argument('--moe-hybridep-num-sms', type=int, default=16,
-                       help='Number of SMs to use for HybridEP.')
-    group.add_argument('--moe-permute-fusion', action='store_true',
-                       help='Fuse token rearrangement ops during token dispatching.')
-    # Token dropping arguments
-    group.add_argument('--moe-expert-capacity-factor', type=float, default=None,
-                       help='The capacity factor for each expert, None means no token will be dropped.')
-    group.add_argument('--moe-pad-expert-input-to-capacity', action='store_true',
-                       help='Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set.')
-    group.add_argument('--moe-token-drop-policy', type=str, default='probs', choices=['probs', 'position'],
-                       help='The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with the lowest probabilities will be dropped. If "position", tokens at the end of each batch will be dropped.')
-    group.add_argument('--moe-apply-probs-on-input', action='store_true',
-                       help='Apply probs before mlp activation for moe routing.')
     # MoE communication overlap arguments
-    group.add_argument('--overlap-moe-expert-parallel-comm', action='store_true',
-                       help='Overlap the EP A2A communication by batch-level overlapping in 1f1b stage.')
-    group.add_argument('--delay-wgrad-compute', action='store_true',
-                       help='Delay the wgrad compute for batch-level overlapping')
-    group.add_argument('--ep-overlap-early-attn-memory-release', action='store_true',
-                       help='Release the memory of the attention module early in EP overlap.')
 
     group.add_argument('--moe-upcycling-granularity', type=int, default=1,
                        help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. '
@@ -3136,19 +2747,6 @@ def _add_mla_args(parser):
 
 def _add_experimental_attention_variant_args(parser):
     group = parser.add_argument_group(title="experimental_attention_variant")
-    group.add_argument('--experimental-attention-variant', default=None, choices=['gated_delta_net', 'dsa'], type=str,
-                       help='Type of attention variant to use. Currently support gated_delta_net and dsa.')
-    # DSA
-    group.add_argument('--dsa-indexer-n-heads', default=None, type=int,
-                       help='Number of indexer heads for sparse attention. If not set, defaults to num-attention-heads.')
-    group.add_argument('--dsa-indexer-head-dim', default=None, type=int,
-                       help='Dimension per indexer head for sparse attention. If not set, defaults to kv-channels.')
-    group.add_argument('--dsa-indexer-topk', default=None, type=int,
-                       help='Number of top-k tokens to select in sparse attention indexer.')
-    group.add_argument('--dsa-indexer-loss-coeff', default=None, type=float,
-                       help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.')
-    group.add_argument('--dsa-indexer-use-sparse-loss', action='store_true',
-                       help='Use sparse indexer loss. If set, the indexer loss will be computed using the top-k indices.')
     # Linear attention
     group.add_argument('--linear-attention-freq', type=la_freq_type, default=None,
                        help='Frequency between LA (linear attention) layers and'
@@ -3159,16 +2757,6 @@ def _add_experimental_attention_variant_args(parser):
                             'where 1 indicates an LA layer and 0 indicates a SDPA layer. '
                             'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, '
                             '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.')
-    group.add_argument('--linear-conv-kernel-dim', default=4, type=int,
-                       help='Conv kernel dimension for the gated delta net.')
-    group.add_argument('--linear-key-head-dim', default=128, type=int,
-                       help='Query and key head dimension for the gated delta net.')
-    group.add_argument('--linear-value-head-dim', default=128, type=int,
-                       help='Value and gate head dimension for the gated delta net.')
-    group.add_argument('--linear-num-key-heads', default=16, type=int,
-                       help='Number of query and key heads for the gated delta net.')
-    group.add_argument('--linear-num-value-heads', default=32, type=int,
-                       help='Number of value and gate heads for the gated delta net.')
     return parser
 
 def _add_heterogeneous_args(parser):
@@ -3247,20 +2835,6 @@ def _add_experimental_args(parser):
                             'hybrid ratio arguments, then the number of each type'
                             'of layer in the override pattern must match number in'
                             'the overidden pattern')
-    group.add_argument('--mamba-state-dim', type=int, default=128,
-                       help='State dimension for Mamba layers.')
-    group.add_argument('--mamba-head-dim', type=int, default=64,
-                       help='Head dimension for Mamba layers.')
-    group.add_argument('--mamba-num-groups', type=int, default=8,
-                       help='Number of groups for Mamba layers.')
-    group.add_argument('--mamba-num-heads', type=int, default=None,
-                       help='Number of heads for Mamba layers.'
-                       'If not set, then the number of heads will be '
-                       '--hidden-size * expand // --mamba-head-dim')
-    group.add_argument('--is-hybrid-model', default=False, action="store_true",
-                       help='Indicates whether the model is a hybrid model.')
-    group.add_argument('--disable-mamba-mem-eff-path', default=False, action="store_true",
-                       help='Disable Mamba efficient path.')
     group.add_argument('--yaml-cfg', type=str, default=None,
                        help = 'Config file to add additional arguments')
 
@@ -3320,18 +2894,6 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser):
             help="Use a default kitchen recipe for all linear layers as defined by QAT_PARAMS index. "
             "The argument has no effect on attention layers.",
         )
-        group.add_argument(
-            '--use-kitchen-attention',
-            action="store_true",
-            help="Have kitchen use its own attention with attention quantization recipe support.",
-        )
-        group.add_argument(
-            '--kitchen-attention-backend',
-            type=str,
-            default='sdpa',
-            choices=['fa', 'sdpa'],
-            help="The backend to use for kitchen attention. The default is 'fa'.",
-        )
     return parser
 
 def _add_sft_args(parser):

From 20e8ac8ff04f51d72c256fd9f247d5bbac71b4b8 Mon Sep 17 00:00:00 2001
From: Deyu Fu <deyuf@nvidia.com>
Date: Fri, 30 Jan 2026 16:42:28 +0800
Subject: [PATCH 79/79] fix merge main issues

Signed-off-by: Deyu Fu <deyuf@nvidia.com>
---
 megatron/training/arguments.py                  | 11 -----------
 tests/unit_tests/models/test_mamba_moe_model.py |  1 -
 2 files changed, 12 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 97338f1f528..1af066a8207 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1173,15 +1173,6 @@ def validate_args(args, defaults={}):
         args.no_load_optim = True
         warn_rank_0('enabling --no-load-optim when skipping training.')
 
-    # Experimental attention variant check
-    if args.linear_attention_type is not None:
-        print_rank_0(
-            '--linear-attention-type is deprecated, use --experimental-attention-variant instead.',
-            args.rank,
-        )
-        args.experimental_attention_variant = args.linear_attention_type
-        del args.linear_attention_type
-
     # Muon optimizer check
     if 'muon' in args.optimizer:
         # TODO: remove these checks once we support them
@@ -2788,8 +2779,6 @@ def _add_mla_args(parser):
 def _add_experimental_attention_variant_args(parser):
     group = parser.add_argument_group(title="experimental_attention_variant")
     # Linear attention
-    group.add_argument('--linear-attention-type', default=None, choices=['gated_delta_net'], type=str,
-                       help='(Deprecated, use --experimental-attention-variant instead) Type of linear attention to use. Currently support gated_delta_net.')
     group.add_argument('--linear-attention-freq', type=la_freq_type, default=None,
                        help='Frequency between LA (linear attention) layers and'
                             ' SDPA (scaled dot-product attention) layers. Accepts either: '
diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py
index aeedc96dfc7..2481649bc3f 100644
--- a/tests/unit_tests/models/test_mamba_moe_model.py
+++ b/tests/unit_tests/models/test_mamba_moe_model.py
@@ -273,7 +273,6 @@
     "offload_modules": [],
     "hybrid_context_parallel": False,
     "max_seqlen_per_dp_cp_rank": None,
-    "enable_routing_replay": False,
     "fallback_to_eager_attn": False,
     "linear_attention_type": None,
     "moe_router_force_biased": None,