diff --git a/megatron/core/models/common/utils.py b/megatron/core/models/common/utils.py
index 186a0c882fc..e5376128369 100644
--- a/megatron/core/models/common/utils.py
+++ b/megatron/core/models/common/utils.py
@@ -17,6 +17,7 @@
 import torch
 
 from megatron.core.pipeline_parallel.utils import ScheduleNode, make_viewless
+from megatron.core.transformer.cuda_graphs import is_cuda_graph_replay_suspended
 from megatron.core.transformer.enums import CudaGraphModule
 from megatron.core.transformer.module import GraphableMegatronModule, float16_to_fp32
 from megatron.core.transformer.transformer_layer import TransformerLayer, make_viewless_tensor
@@ -409,7 +410,11 @@ def __init__(self, layer):
 
     def backward_dw(self):
         """Run eager or graphed backward wgrad callables for the wrapped layer."""
-        is_replay = hasattr(self.layer, 'cuda_graphs') and self.layer.cuda_graphs
+        is_replay = (
+            hasattr(self.layer, 'cuda_graphs')
+            and self.layer.cuda_graphs
+            and not is_cuda_graph_replay_suspended()
+        )
         if self.shared_expert_dw_callable is not None and (
             not is_replay or CudaGraphModule.moe_router not in self.cuda_graph_modules
         ):
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 61648d7b602..e66eb047367 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -11,6 +11,7 @@
     FineGrainedActivationOffloadingInterface as off_interface,
 )
 from megatron.core.pipeline_parallel.utils import ScheduleNode
+from megatron.core.transformer.cuda_graphs import is_cuda_graph_replay_suspended
 from megatron.core.transformer.module import GraphableMegatronModule
 from megatron.core.transformer.moe.moe_layer import MoELayer
 from megatron.core.transformer.transformer_layer import TransformerLayer, make_viewless_tensor
@@ -75,6 +76,7 @@ def submodule_pre_dispatch_forward(node: ScheduleNode, hidden_states: torch.Tens
             isinstance(layer, GraphableMegatronModule)
             and hasattr(layer, 'cuda_graphs')
             and layer.cuda_graphs
+            and not is_cuda_graph_replay_suspended()
         ):
             layer.set_te_cuda_graph_backward_dw_wrapper()
             forward_func = layer._te_cuda_graph_replay
@@ -215,7 +217,11 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor):
         output = layer.mlp.postprocess(output, shared_expert_output)
 
         mlp_output_with_bias = (output, None)
-        if hasattr(layer, 'cuda_graphs') and layer.cuda_graphs:
+        if (
+            hasattr(layer, 'cuda_graphs')
+            and layer.cuda_graphs
+            and not is_cuda_graph_replay_suspended()
+        ):
             layer.mlp.cudagraph_tensor_store.clear()
         with layer.bias_dropout_add_exec_handler():
             hidden_states = layer.mlp_bda(layer.training, layer.config.bias_dropout_fusion)(
diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py
index 22322e1b346..75ac4edd280 100644
--- a/megatron/core/models/hybrid/hybrid_block.py
+++ b/megatron/core/models/hybrid/hybrid_block.py
@@ -27,6 +27,10 @@
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.cuda_graphs import annotate_first_last_layer
 from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.layer_boundary_observer import (
+    observe_transformer_layer_input,
+    observe_transformer_layer_output,
+)
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_layer import TransformerLayer
@@ -328,6 +332,7 @@ def get_inner_quant_context(config, layer_number):
                     inner_quant_context = get_inner_quant_context(
                         self.config, layer.layer_number - 1
                     )
+                    observe_transformer_layer_input(self, layer, hidden_states)
                     with inner_quant_context:
                         if isinstance(layer, TransformerLayer):
                             hidden_states, _ = layer(
@@ -352,6 +357,7 @@ def get_inner_quant_context(config, layer_number):
                     # for cross-attention, and is not needed in our model.
                     if isinstance(hidden_states, tuple):
                         hidden_states = hidden_states[0]
+                    observe_transformer_layer_output(self, layer, hidden_states)
 
         # Final layer norm.
         if self.post_process and self.post_layer_norm:
diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 27b675d1b8d..6c4f88c0799 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -691,6 +691,8 @@ def init_state_fn(opt, config=None):
         tp_group = pg_collection.tp
     # TODO(M4): plumb tp_group through optimizer constructors so this setattr disappears.
     setattr(optimizer, 'tp_group', tp_group)
+    if not hasattr(optimizer, 'model_chunks'):
+        setattr(optimizer, 'model_chunks', model_chunks)
 
     return optimizer
 
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index 9e030a6b17f..d9586680c42 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -54,6 +54,7 @@
 )
 from ..fp4_utils import is_nvfp4tensor, quantize_nvfp4_param_shard
 from ..fp8_utils import dequantize_fp8_tensor, is_float8tensor, quantize_param_shard
+from ..per_parameter_stats import PerParameterStatRegistry
 from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict
 from ..transformer.module import MegatronModule
 from .grad_scaler import MegatronGradScaler
@@ -778,6 +779,25 @@ def get_grad_stats_parallel_group(self) -> torch.distributed.ProcessGroup:
         """
         return getattr(self, 'grad_stats_parallel_group', None)
 
+    def _get_param_to_name_for_per_param_stats(
+        self, registry: PerParameterStatRegistry
+    ) -> Dict[torch.nn.Parameter, str]:
+        param_to_name = super()._get_param_to_name_for_per_param_stats(registry)
+
+        def add_shard_names(model_groups, shard_groups):
+            for model_group, shard_group in zip(model_groups, shard_groups):
+                for model_param, shard_param in zip(model_group, shard_group):
+                    if shard_param is not None and model_param in registry.param_to_name:
+                        param_to_name[shard_param] = registry.name_for_param(model_param)
+
+        add_shard_names(self.model_fp32_groups, self.shard_fp32_groups)
+        if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8:
+            add_shard_names(self.model_float16_groups, self.shard_float16_groups)
+        else:
+            add_shard_names(self.model_float16_groups, self.shard_fp32_from_float16_groups)
+
+        return param_to_name
+
     def state_dict(self):
         """
         The state dict contains all non-DP-rank-dependent (i.e., non-parameter-
diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py
index 606525f8097..8d841bf683c 100644
--- a/megatron/core/optimizer/layer_wise_optimizer.py
+++ b/megatron/core/optimizer/layer_wise_optimizer.py
@@ -10,6 +10,7 @@
 from megatron.core.dist_checkpointing.dict_utils import nested_values
 from megatron.core.dist_checkpointing.mapping import LocalNonpersistentObject, ShardedStateDict
 from megatron.core.distributed.param_and_grad_buffer import group_params_for_buffers
+from megatron.core.per_parameter_stats import NamedTensorBucket, PerParameterStatRegistry
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.utils import get_pg_rank, get_pg_size, log_single_rank
 
@@ -744,6 +745,21 @@ def _get_grad_norm_for_group(self, grad_norm_group: str):
         grad_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None)
         return grad_norm
 
+    def get_raw_moment_buckets_for_grad_norm(
+        self, registry: PerParameterStatRegistry
+    ) -> list[NamedTensorBucket]:
+        names = []
+        grads = []
+        for optimizer in self.chained_optimizers:
+            for name, param in optimizer.get_named_parameters_for_grad_norm(registry):
+                grad = optimizer._get_grad_for_grad_norm(param)
+                if not optimizer._include_param_in_grad_norm(param, grad):
+                    continue
+                names.append(name)
+                grads.append(grad.detach())
+
+        return [NamedTensorBucket(names, grads, (None,))]
+
     @torch.no_grad()
     def count_zeros(self):
         params = []
diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py
index b40b2cb2dd5..28b279e723f 100644
--- a/megatron/core/optimizer/optimizer.py
+++ b/megatron/core/optimizer/optimizer.py
@@ -46,6 +46,12 @@
     optim_state_to_sharding_state,
 )
 from ..dist_checkpointing.utils import add_prefix_for_sharding
+from ..per_parameter_stats import (
+    NamedTensorBucket,
+    PerParameterStatRegistry,
+    get_or_create_per_parameter_stat_registry,
+    reduce_raw_moments_by_param,
+)
 from ..transformer.module import param_is_not_shared
 from ..utils import log_single_rank
 from .clip_grads import clip_grad_by_total_norm_fp32, count_zeros_fp32, get_grad_norm_fp32
@@ -54,6 +60,8 @@
 
 logger = getLogger(__name__)
 
+_GRAD_RAW_MOMENTS_BY_PARAM_NORM_RTOL = 1e-2
+
 
 def _zero_grad_group_helper(
     group: List[torch.nn.Parameter], set_to_none: bool, use_decoupled_grad: bool = False
@@ -158,6 +166,9 @@ def __init__(
             )
         self.config = config
         self.init_state_fn = init_state_fn
+        self._per_param_grad_raw_moments_requested = False
+        self._per_param_stat_registry = None
+        self._latest_grad_raw_moments_by_param = None
 
     def get_parameters(self) -> List[torch.nn.Parameter]:
         """
@@ -170,6 +181,38 @@ def get_parameters(self) -> List[torch.nn.Parameter]:
                     params.append(param)
         return params
 
+    def _get_grad_for_grad_norm(self, param: torch.nn.Parameter) -> torch.Tensor | None:
+        if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8 or (
+            # Megatron-FSDP always uses decoupled_grad with FusedAdam.
+            self.config.use_precision_aware_optimizer
+            and getattr(param, "__fsdp_param__", False)
+        ):
+            grad = param.decoupled_grad if hasattr(param, "decoupled_grad") else None
+            if (
+                getattr(param, "__fsdp_param__", False)
+                and grad is not None
+                and hasattr(grad, "_local_tensor")
+            ):
+                # Megatron-FSDP gradients are DTensors.
+                grad = grad._local_tensor
+        elif getattr(param, "__fsdp_param__", False):
+            # Megatron-FSDP gradients are DTensors.
+            grad = param.grad._local_tensor if param.grad is not None else None
+        else:
+            grad = param.grad
+        return grad
+
+    def _include_param_in_grad_norm(
+        self, param: torch.nn.Parameter, grad: torch.Tensor | None
+    ) -> bool:
+        return (
+            grad is not None
+            and param_is_not_shared(param)
+            and tensor_parallel.param_is_not_tensor_parallel_duplicate(
+                param, getattr(self, 'tp_group', None)
+            )
+        )
+
     def prepare_model_params_for_param_sync(self) -> None:
         """Stage optimizer-owned model params before an explicit DDP param sync."""
         return
@@ -192,30 +235,8 @@ def _filter_grads_for_norm(
         for param in params:
             if param_filter is not None and not param_filter(param):
                 continue
-            if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8 or (
-                # Megatron-FSDP always uses decoupled_grad with FusedAdam.
-                self.config.use_precision_aware_optimizer
-                and getattr(param, "__fsdp_param__", False)
-            ):
-                grad = param.decoupled_grad if hasattr(param, "decoupled_grad") else None
-                if (
-                    getattr(param, "__fsdp_param__", False)
-                    and grad is not None
-                    and hasattr(grad, "_local_tensor")
-                ):
-                    # Megatron-FSDP gradients are DTensors.
-                    grad = grad._local_tensor
-            elif getattr(param, "__fsdp_param__", False):
-                # Megatron-FSDP gradients are DTensors.
-                grad = param.grad._local_tensor if param.grad is not None else None
-            else:
-                grad = param.grad
-            grad_not_none = grad is not None
-            is_not_shared = param_is_not_shared(param)
-            is_not_tp_duplicate = tensor_parallel.param_is_not_tensor_parallel_duplicate(
-                param, getattr(self, 'tp_group', None)
-            )
-            if grad_not_none and is_not_shared and is_not_tp_duplicate:
+            grad = self._get_grad_for_grad_norm(param)
+            if self._include_param_in_grad_norm(param, grad):
                 grads_for_norm.append(grad)
         return grads_for_norm
 
@@ -261,6 +282,112 @@ def has_grad_norm_group(self, grad_norm_group: str) -> bool:
             cache[grad_norm_group] = bool(flag.item() > 0)
         return cache[grad_norm_group]
 
+    def _get_param_to_name_for_per_param_stats(
+        self, registry: PerParameterStatRegistry
+    ) -> dict[torch.nn.Parameter, str]:
+        param_to_name = {}
+        for model_param, name in registry.param_to_name.items():
+            param_to_name[model_param] = name
+            main_param = getattr(model_param, 'main_param', None)
+            if main_param is not None:
+                param_to_name[main_param] = name
+        return param_to_name
+
+    def get_named_parameters_for_grad_norm(
+        self, registry: PerParameterStatRegistry
+    ) -> list[tuple[str, torch.nn.Parameter]]:
+        """Return named optimizer parameters that are present in the model registry."""
+        param_to_name = self._get_param_to_name_for_per_param_stats(registry)
+        return [
+            (param_to_name[param], param)
+            for param in self.get_parameters()
+            if param in param_to_name
+        ]
+
+    def get_raw_moment_buckets_for_grad_norm(
+        self, registry: PerParameterStatRegistry
+    ) -> list[NamedTensorBucket]:
+        """Build gradient buckets for per-parameter raw-moment reductions."""
+        names = []
+        grads = []
+        for name, param in self.get_named_parameters_for_grad_norm(registry):
+            grad = self._get_grad_for_grad_norm(param)
+            if not self._include_param_in_grad_norm(param, grad):
+                continue
+            names.append(name)
+            grads.append(grad.detach())
+
+        return [NamedTensorBucket(names, grads, (self.get_grad_stats_parallel_group(),))]
+
+    def get_grad_raw_moments_by_param(
+        self,
+        registry: PerParameterStatRegistry | None = None,
+        expert_model_parallel_group: torch.distributed.ProcessGroup | None = None,
+    ) -> tuple[list[tuple[str, dict[str, float]]], dict[str, float]]:
+        """Compute per-parameter gradient raw moments and aggregate moments."""
+        if registry is None:
+            registry = get_or_create_per_parameter_stat_registry(
+                self.model_chunks, expert_model_parallel_group=expert_model_parallel_group
+            )
+        return reduce_raw_moments_by_param(
+            registry, self.get_raw_moment_buckets_for_grad_norm(registry)
+        )
+
+    def request_grad_raw_moments_by_param(
+        self,
+        model_chunks: Any,
+        expert_model_parallel_group: torch.distributed.ProcessGroup | None = None,
+    ) -> None:
+        """Request per-parameter gradient raw moments for the next optimizer step."""
+        self._per_param_stat_registry = get_or_create_per_parameter_stat_registry(
+            model_chunks, expert_model_parallel_group=expert_model_parallel_group
+        )
+        self._per_param_grad_raw_moments_requested = True
+        self._latest_grad_raw_moments_by_param = None
+
+    def consume_grad_raw_moments_by_param(self) -> list[tuple[str, dict[str, float]]] | None:
+        """Return and clear the most recently recorded gradient raw moments."""
+        grad_raw_moments_by_param = self._latest_grad_raw_moments_by_param
+        self._latest_grad_raw_moments_by_param = None
+        return grad_raw_moments_by_param
+
+    def _clear_grad_raw_moments_by_param_request(self) -> None:
+        self._per_param_grad_raw_moments_requested = False
+        self._latest_grad_raw_moments_by_param = None
+
+    def _maybe_record_grad_raw_moments_by_param(
+        self, scalar_grad_norm: float | torch.Tensor | None = None
+    ) -> None:
+        if not self._per_param_grad_raw_moments_requested:
+            return
+
+        grad_raw_moments_by_param, aggregate_moments = self.get_grad_raw_moments_by_param(
+            self._per_param_stat_registry
+        )
+        self._latest_grad_raw_moments_by_param = grad_raw_moments_by_param
+        self._per_param_grad_raw_moments_requested = False
+
+        if scalar_grad_norm is None:
+            return
+        if any(self.has_grad_norm_group(group) for group in SEPARATE_GRAD_NORM_GROUPS):
+            return
+        if isinstance(scalar_grad_norm, torch.Tensor):
+            scalar_grad_norm = scalar_grad_norm.item()
+        scalar_grad_norm = float(scalar_grad_norm)
+        reconstructed_norm = aggregate_moments["sum_2"] ** 0.5
+        rel_diff = (
+            abs(reconstructed_norm - scalar_grad_norm) / scalar_grad_norm
+            if scalar_grad_norm > 0
+            else 0.0
+        )
+        if rel_diff > _GRAD_RAW_MOMENTS_BY_PARAM_NORM_RTOL:
+            warnings.warn(
+                "per-parameter gradient raw moments recombine to an l2 norm of "
+                f"{reconstructed_norm:.6e}, but the directly-computed gradient norm is "
+                f"{scalar_grad_norm:.6e} (relative difference {rel_diff:.2e} > "
+                f"{_GRAD_RAW_MOMENTS_BY_PARAM_NORM_RTOL:.0e})."
+            )
+
     def get_grad_stats_parallel_group(self) -> torch.distributed.ProcessGroup:
         """Process group for reducing gradient statistics (num_zeros & norm).
 
@@ -328,6 +455,7 @@ def clip_grad_norm(self, clip_grad: float) -> float:
         grad_norm = get_grad_norm_fp32(
             grads_for_norm, grad_stats_parallel_group=self.get_grad_stats_parallel_group()
         )
+        self._maybe_record_grad_raw_moments_by_param(grad_norm)
 
         if clip_grad > 0.0 and params:
             # Only reduce group grad norms when clipping can use them.
@@ -748,6 +876,7 @@ def step(self):
 
         found_inf_flag = self.prepare_grads()
         if found_inf_flag:
+            self._clear_grad_raw_moments_by_param_request()
             return False, None, None
 
         # Clip the main gradients.
@@ -758,6 +887,9 @@ def step(self):
         grad_norm = 0.0
         if self.config.clip_grad > 0.0:
             grad_norm = self.clip_grad_norm(self.config.clip_grad)
+        elif self._per_param_grad_raw_moments_requested:
+            grad_norm = self.get_grad_norm()
+            self._maybe_record_grad_raw_moments_by_param(grad_norm)
         if timers is not None:
             timers('optimizer-clip-main-grad').stop()
 
@@ -1118,6 +1250,7 @@ def step(self):
 
         found_inf_flag = self.prepare_grads()
         if found_inf_flag:
+            self._clear_grad_raw_moments_by_param_request()
             return False, None, None
 
         # Clip gradients.
@@ -1128,6 +1261,9 @@ def step(self):
         grad_norm = None
         if self.config.clip_grad > 0.0:
             grad_norm = self.clip_grad_norm(self.config.clip_grad)
+        elif self._per_param_grad_raw_moments_requested:
+            grad_norm = self.get_grad_norm()
+            self._maybe_record_grad_raw_moments_by_param(grad_norm)
         if timers is not None:
             timers('optimizer-clip-main-grad').stop()
 
@@ -1238,6 +1374,9 @@ class ChainedOptimizer(MegatronOptimizer):
 
     def __init__(self, chained_optimizers: List[MegatronOptimizer]):
         self.model_chunks = []
+        self._per_param_grad_raw_moments_requested = False
+        self._per_param_stat_registry = None
+        self._latest_grad_raw_moments_by_param = None
         # chained_optimizers would be empty in the case that a rank
         # has no trainable parameters
         if chained_optimizers:
@@ -1566,6 +1705,14 @@ def get_grad_stats_parallel_group(self) -> torch.distributed.ProcessGroup:
         )
         return self.chained_optimizers[0].get_grad_stats_parallel_group()
 
+    def get_raw_moment_buckets_for_grad_norm(
+        self, registry: PerParameterStatRegistry
+    ) -> list[NamedTensorBucket]:
+        buckets = []
+        for optimizer in self.chained_optimizers:
+            buckets.extend(optimizer.get_raw_moment_buckets_for_grad_norm(registry))
+        return buckets
+
     @torch.no_grad()
     def get_grad_norm(self):
         if len(self.chained_optimizers) == 1:
@@ -1669,6 +1816,7 @@ def step(self):
         self.grad_norms_by_group = {}
         found_inf_flag = self.prepare_grads()
         if found_inf_flag:
+            self._clear_grad_raw_moments_by_param_request()
             return False, None, None
 
         grad_norm = self.get_grad_norm()
@@ -1681,6 +1829,7 @@ def step(self):
         )
         if should_clip:
             self._compute_grad_norms_by_group()
+        self._maybe_record_grad_raw_moments_by_param(grad_norm)
 
         # Clip gradients.
         for optimizer in self.chained_optimizers:
diff --git a/megatron/core/parameter_names.py b/megatron/core/parameter_names.py
new file mode 100644
index 00000000000..6078bc1c191
--- /dev/null
+++ b/megatron/core/parameter_names.py
@@ -0,0 +1,257 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Canonical logical names and deterministic indices for model parameters."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import dataclass, field
+
+import torch
+
+from megatron.core.utils import unwrap_model
+
+_GROUPED_EXPERT_PATTERN = re.compile(
+    r"^((?:.*\.)?mlp\.experts\.linear_fc\d\.(?:weight|bias))(\d+)(.*)$"
+)
+_SEQUENTIAL_EXPERT_PATTERN = re.compile(r"^((?:.*\.)?mlp\.experts\.local_experts\.)(\d+)(\..*)$")
+
+
+@dataclass(frozen=True, init=False)
+class CanonicalParameterNameIndex(Mapping[str, int]):
+    """Deterministic indices for a set of canonical parameter names.
+
+    Duplicate names are collapsed and the remaining names are ordered
+    lexicographically. The resulting mapping is immutable after construction.
+
+    Args:
+        names: Canonical parameter names to index.
+    """
+
+    names: tuple[str, ...]
+    _name_to_index: dict[str, int] = field(repr=False, compare=False, hash=False)
+
+    def __init__(self, names: Iterable[str]) -> None:
+        ordered_names = tuple(sorted(set(names)))
+        object.__setattr__(self, "names", ordered_names)
+        object.__setattr__(
+            self, "_name_to_index", {name: index for index, name in enumerate(ordered_names)}
+        )
+
+    def __getitem__(self, name: str) -> int:
+        return self._name_to_index[name]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.names)
+
+    def __len__(self) -> int:
+        return len(self.names)
+
+
+class CanonicalParameterNameMap(Mapping[torch.nn.Parameter, str]):
+    """Map model parameters to topology-independent logical names.
+
+    Pipeline-local layer indices are replaced with the global ``layer_number``
+    assigned to their owning layer module. Expert-local indices are replaced
+    with global expert indices when an expert-parallel rank and size are
+    supplied. The map only contains original model parameters; consumers are
+    responsible for mapping optimizer copies or shards back to those parameters.
+
+    Construction performs no distributed collectives and does not read global
+    process-group state. Call :meth:`all_gather_index` explicitly when every
+    rank in a process group needs the same global name index.
+
+    Args:
+        model_chunks: A model module or iterable of model chunks.
+        expert_parallel_rank: Rank within the expert-model-parallel group.
+        expert_parallel_size: Size of the expert-model-parallel group.
+
+    Raises:
+        ValueError: If the model list or expert topology is invalid, or if two
+            distinct local parameters resolve to the same canonical name.
+    """
+
+    def __init__(
+        self,
+        model_chunks: Iterable[torch.nn.Module] | torch.nn.Module,
+        *,
+        expert_parallel_rank: int = 0,
+        expert_parallel_size: int = 1,
+    ) -> None:
+        if expert_parallel_size < 1:
+            raise ValueError("expert_parallel_size must be at least 1")
+        if not 0 <= expert_parallel_rank < expert_parallel_size:
+            raise ValueError(
+                f"expert_parallel_rank must be in [0, {expert_parallel_size}), "
+                f"got {expert_parallel_rank}"
+            )
+
+        normalized_chunks = _normalize_model_chunks(model_chunks)
+        if not normalized_chunks:
+            raise ValueError("Cannot build canonical parameter names for an empty model list.")
+
+        self.model_chunks = tuple(unwrap_model(normalized_chunks))
+        self.expert_parallel_rank = expert_parallel_rank
+        self.expert_parallel_size = expert_parallel_size
+        self._param_to_name = self._build_param_to_name()
+        self._local_index = CanonicalParameterNameIndex(self._param_to_name.values())
+
+    def __getitem__(self, param: torch.nn.Parameter) -> str:
+        return self._param_to_name[param]
+
+    def __iter__(self) -> Iterator[torch.nn.Parameter]:
+        return iter(self._param_to_name)
+
+    def __len__(self) -> int:
+        return len(self._param_to_name)
+
+    @property
+    def local_index(self) -> CanonicalParameterNameIndex:
+        """Return deterministic indices for locally present parameter names."""
+        return self._local_index
+
+    def name_for_param(self, param: torch.nn.Parameter) -> str:
+        """Return the canonical logical name for a model parameter."""
+        return self[param]
+
+    def all_gather_index(
+        self, group: torch.distributed.ProcessGroup | None = None
+    ) -> CanonicalParameterNameIndex:
+        """Collect and deterministically index canonical names from a process group.
+
+        This is a collective operation. Every rank in ``group`` must call it in
+        the same collective order. If distributed communication is unavailable
+        or uninitialized, the local index is returned without communication.
+
+        Args:
+            group: Process group whose canonical names should be collected. If
+                ``None``, use the default world process group.
+
+        Returns:
+            A global, deterministic parameter-name index shared by the group.
+        """
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            return self.local_index
+
+        world_size = torch.distributed.get_world_size(group=group)
+        if world_size == 1:
+            return self.local_index
+
+        gathered_names: list[tuple[str, ...] | None] = [None] * world_size
+        torch.distributed.all_gather_object(gathered_names, self.local_index.names, group=group)
+        return CanonicalParameterNameIndex(
+            name for rank_names in gathered_names if rank_names is not None for name in rank_names
+        )
+
+    def _build_param_to_name(self) -> dict[torch.nn.Parameter, str]:
+        param_to_name: dict[torch.nn.Parameter, str] = {}
+        name_to_param: dict[str, torch.nn.Parameter] = {}
+
+        for model_chunk in self.model_chunks:
+            layer_prefixes = _build_global_layer_prefixes(model_chunk)
+            num_experts = _get_num_moe_experts(model_chunk)
+            expert_offset = _get_local_expert_offset(
+                num_experts, self.expert_parallel_rank, self.expert_parallel_size
+            )
+
+            for local_name, param in model_chunk.named_parameters():
+                canonical_name = _canonical_parameter_name(
+                    local_name, layer_prefixes, num_experts, expert_offset
+                )
+
+                previous_name = param_to_name.get(param)
+                if previous_name is not None and previous_name != canonical_name:
+                    raise ValueError(
+                        "A shared parameter resolved to multiple canonical names: "
+                        f"{previous_name!r} and {canonical_name!r}."
+                    )
+
+                previous_param = name_to_param.get(canonical_name)
+                if previous_param is not None and previous_param is not param:
+                    raise ValueError(
+                        f"Canonical parameter name {canonical_name!r} refers to multiple "
+                        "distinct local parameters."
+                    )
+
+                param_to_name[param] = canonical_name
+                name_to_param[canonical_name] = param
+
+        return param_to_name
+
+
+def _normalize_model_chunks(
+    model_chunks: Iterable[torch.nn.Module] | torch.nn.Module,
+) -> list[torch.nn.Module]:
+    if isinstance(model_chunks, torch.nn.Module):
+        return [model_chunks]
+    return list(model_chunks)
+
+
+def _build_global_layer_prefixes(model_chunk: torch.nn.Module) -> dict[str, str]:
+    """Build local-to-global prefixes for numbered layer modules."""
+    prefixes = {}
+    for module_name, module in model_chunk.named_modules():
+        if not module_name or "mtp" in module_name:
+            continue
+        layer_number = getattr(module, "layer_number", None)
+        if not isinstance(layer_number, int):
+            continue
+        parts = module_name.split(".")
+        if not parts[-1].isdigit():
+            continue
+        parts[-1] = str(layer_number - 1)
+        prefixes[module_name] = ".".join(parts)
+    return prefixes
+
+
+def _canonical_parameter_name(
+    local_name: str, layer_prefixes: Mapping[str, str], num_experts: int | None, expert_offset: int
+) -> str:
+    name = _replace_longest_prefix(local_name, layer_prefixes)
+    return _global_expert_parameter_name(name, num_experts, expert_offset)
+
+
+def _replace_longest_prefix(name: str, replacements: Mapping[str, str]) -> str:
+    parts = name.split(".")
+    for index in range(len(parts), 0, -1):
+        prefix = ".".join(parts[:index])
+        replacement = replacements.get(prefix)
+        if replacement is None:
+            continue
+        suffix = ".".join(parts[index:])
+        return replacement if not suffix else f"{replacement}.{suffix}"
+    return name
+
+
+def _global_expert_parameter_name(
+    local_name: str, num_experts: int | None, expert_offset: int
+) -> str:
+    if not num_experts or expert_offset == 0:
+        return local_name
+
+    for pattern in (_GROUPED_EXPERT_PATTERN, _SEQUENTIAL_EXPERT_PATTERN):
+        match = pattern.match(local_name)
+        if match is not None:
+            prefix, local_expert_index, suffix = match.groups()
+            return f"{prefix}{int(local_expert_index) + expert_offset}{suffix}"
+
+    return local_name
+
+
+def _get_local_expert_offset(
+    num_experts: int | None, expert_parallel_rank: int, expert_parallel_size: int
+) -> int:
+    if not num_experts or expert_parallel_size == 1:
+        return 0
+    if num_experts % expert_parallel_size != 0:
+        raise ValueError(
+            f"num_moe_experts ({num_experts}) must be divisible by "
+            f"expert_parallel_size ({expert_parallel_size})"
+        )
+    return expert_parallel_rank * (num_experts // expert_parallel_size)
+
+
+def _get_num_moe_experts(model_chunk: torch.nn.Module) -> int | None:
+    config = getattr(model_chunk, "config", None)
+    return getattr(config, "num_moe_experts", None)
diff --git a/megatron/core/per_parameter_stats.py b/megatron/core/per_parameter_stats.py
new file mode 100644
index 00000000000..62c74743b98
--- /dev/null
+++ b/megatron/core/per_parameter_stats.py
@@ -0,0 +1,314 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Shared helpers for high-cardinality per-parameter statistics."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Iterable, Sequence
+
+import torch
+
+try:
+    from transformer_engine.pytorch.optimizers import multi_tensor_applier, multi_tensor_raw_moments
+except ImportError:
+    multi_tensor_applier = None
+    multi_tensor_raw_moments = None
+
+from megatron.core.parameter_names import CanonicalParameterNameMap
+from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model
+
+RAW_MOMENT_FIELDS = ("count", "sum_1", "sum_2", "sum_3", "sum_4")
+_RAW_MOMENTS_DTYPE = torch.float32
+_MULTI_TENSOR_RAW_MOMENTS_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_MULTI_TENSOR_RAW_MOMENTS_SPLIT_ALIGNMENT = 4
+_MAX_MULTI_TENSOR_RAW_MOMENTS_NUMEL = torch.iinfo(torch.int32).max - (
+    torch.iinfo(torch.int32).max % _MULTI_TENSOR_RAW_MOMENTS_SPLIT_ALIGNMENT
+)
+_DISABLE_MULTI_TENSOR_RAW_MOMENTS_ENV = "MEGATRON_DISABLE_MULTI_TENSOR_RAW_MOMENTS"
+_TRUTHY_ENV_VALUES = {"1", "true", "yes", "on"}
+
+
+@dataclass(frozen=True)
+class NamedTensorBucket:
+    """Named tensors that should be reduced over the same process groups."""
+
+    names: Sequence[str]
+    tensors: Sequence[torch.Tensor]
+    reduce_groups: tuple[torch.distributed.ProcessGroup | None, ...] = ()
+
+
+class PerParameterStatRegistry:
+    """Canonical parameter-name registry for per-parameter statistics."""
+
+    def __init__(
+        self,
+        model_chunks: Iterable[torch.nn.Module] | torch.nn.Module,
+        expert_model_parallel_group: torch.distributed.ProcessGroup | None = None,
+    ):
+        self.model_chunks = unwrap_model(_normalize_model_chunks(model_chunks))
+        self.expert_model_parallel_group = expert_model_parallel_group
+        self.expert_model_parallel_rank, self.expert_model_parallel_size = (
+            _get_expert_model_parallel_rank_size(expert_model_parallel_group)
+        )
+        self.cache_key = _registry_cache_key(
+            self.model_chunks,
+            expert_model_parallel_group,
+            self.expert_model_parallel_rank,
+            self.expert_model_parallel_size,
+        )
+        self.parameter_names = CanonicalParameterNameMap(
+            self.model_chunks,
+            expert_parallel_rank=self.expert_model_parallel_rank,
+            expert_parallel_size=self.expert_model_parallel_size,
+        )
+        self.parameter_name_index = self.parameter_names.all_gather_index()
+        self.param_to_name = dict(self.parameter_names.items())
+        self.name_to_index = self.parameter_name_index
+        self.index_to_name = self.parameter_name_index.names
+
+    def name_for_param(self, param: torch.nn.Parameter) -> str:
+        """Return the canonical name for ``param``."""
+        return self.param_to_name[param]
+
+    @property
+    def num_params(self) -> int:
+        """Number of globally known parameters."""
+        return len(self.name_to_index)
+
+
+def get_or_create_per_parameter_stat_registry(
+    model_chunks: Iterable[torch.nn.Module] | torch.nn.Module,
+    expert_model_parallel_group: torch.distributed.ProcessGroup | None = None,
+) -> PerParameterStatRegistry:
+    """Return a per-model cached parameter-stat registry."""
+    unwrapped_model_chunks = unwrap_model(_normalize_model_chunks(model_chunks))
+    if not unwrapped_model_chunks:
+        raise ValueError("Cannot build a per-parameter stat registry for an empty model list.")
+
+    expert_model_parallel_rank, expert_model_parallel_size = _get_expert_model_parallel_rank_size(
+        expert_model_parallel_group
+    )
+    cache_key = _registry_cache_key(
+        unwrapped_model_chunks,
+        expert_model_parallel_group,
+        expert_model_parallel_rank,
+        expert_model_parallel_size,
+    )
+    cache_owner = unwrapped_model_chunks[0]
+    registry = getattr(cache_owner, "_per_parameter_stat_registry", None)
+    if registry is None or registry.cache_key != cache_key:
+        registry = PerParameterStatRegistry(
+            unwrapped_model_chunks, expert_model_parallel_group=expert_model_parallel_group
+        )
+        cache_owner._per_parameter_stat_registry = registry
+    return registry
+
+
+def reduce_raw_moments_by_param(
+    registry: PerParameterStatRegistry, buckets: Sequence[NamedTensorBucket]
+) -> tuple[list[tuple[str, dict[str, float]]], dict[str, float]]:
+    """Reduce named tensor raw moments by parameter name.
+
+    Args:
+        registry: Canonical parameter-name registry.
+        buckets: Named tensor buckets with the process groups needed to assemble each bucket's
+            local raw moments into global per-parameter raw moments.
+
+    Returns:
+        A ``(values, aggregate_moments)`` tuple. ``values`` is a list of
+        ``(name, raw_moment_dict)`` tuples ordered by canonical parameter index.
+    """
+    device = _select_device(buckets)
+    moments = torch.zeros(
+        (registry.num_params, len(RAW_MOMENT_FIELDS)), dtype=_RAW_MOMENTS_DTYPE, device=device
+    )
+
+    for bucket in buckets:
+        if len(bucket.names) != len(bucket.tensors):
+            raise ValueError(
+                f"NamedTensorBucket has {len(bucket.names)} names but "
+                f"{len(bucket.tensors)} tensors."
+            )
+
+        bucket_moments = torch.zeros_like(moments)
+        if bucket.names:
+            indices = torch.tensor(
+                [registry.name_to_index[name] for name in bucket.names],
+                dtype=torch.long,
+                device=device,
+            )
+            bucket_moments.index_add_(0, indices, _local_raw_moments(bucket.tensors, device))
+
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            for group in bucket.reduce_groups:
+                torch.distributed.all_reduce(
+                    bucket_moments, op=torch.distributed.ReduceOp.SUM, group=group
+                )
+
+        moments += bucket_moments
+
+    rows = moments.tolist()
+    aggregate_moments = raw_moment_row_to_dict(moments.sum(dim=0).tolist())
+    return [
+        (name, raw_moment_row_to_dict(rows[idx])) for idx, name in enumerate(registry.index_to_name)
+    ], aggregate_moments
+
+
+def _select_device(buckets: Sequence[NamedTensorBucket]) -> torch.device:
+    for bucket in buckets:
+        if bucket.tensors:
+            return bucket.tensors[0].device
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{torch.cuda.current_device()}")
+    return torch.device("cpu")
+
+
+def _normalize_model_chunks(
+    model_chunks: Iterable[torch.nn.Module] | torch.nn.Module,
+) -> list[torch.nn.Module]:
+    if isinstance(model_chunks, (list, tuple)):
+        return list(model_chunks)
+    return [model_chunks]
+
+
+def _local_raw_moments(tensors: Sequence[torch.Tensor], device: torch.device) -> torch.Tensor:
+    if not tensors:
+        return torch.zeros((0, len(RAW_MOMENT_FIELDS)), dtype=_RAW_MOMENTS_DTYPE, device=device)
+
+    if _can_use_multi_tensor_raw_moments(tensors, device):
+        return _multi_tensor_raw_moments(tensors, device)
+
+    rows = [_torch_raw_moment_row(tensor, device=device) for tensor in tensors]
+    return torch.stack(rows)
+
+
+def raw_moment_row(tensor: torch.Tensor, device: torch.device | None = None) -> torch.Tensor:
+    """Return count and raw sums of powers 1-4 for ``tensor`` as an fp32 row."""
+    device = tensor.device if device is None else device
+    if _can_use_multi_tensor_raw_moments([tensor], device):
+        return _multi_tensor_raw_moments([tensor], device)[0]
+    return _torch_raw_moment_row(tensor, device=device)
+
+
+def _torch_raw_moment_row(tensor: torch.Tensor, device: torch.device | None = None) -> torch.Tensor:
+    """Torch fallback for count and raw sums of powers 1-4."""
+    device = tensor.device if device is None else device
+    values = tensor.detach().to(device=device, dtype=_RAW_MOMENTS_DTYPE)
+    values_2 = values * values
+    return torch.stack(
+        [
+            torch.tensor(float(values.numel()), dtype=_RAW_MOMENTS_DTYPE, device=device),
+            values.sum(),
+            values_2.sum(),
+            (values_2 * values).sum(),
+            (values_2 * values_2).sum(),
+        ]
+    )
+
+
+def _can_use_multi_tensor_raw_moments(
+    tensors: Sequence[torch.Tensor], device: torch.device
+) -> bool:
+    disabled = os.getenv(_DISABLE_MULTI_TENSOR_RAW_MOMENTS_ENV, "").lower() in _TRUTHY_ENV_VALUES
+    return (
+        not disabled
+        and multi_tensor_applier is not None
+        and multi_tensor_raw_moments is not None
+        and device.type == "cuda"
+        and all(
+            tensor.device == device
+            and tensor.dtype in _MULTI_TENSOR_RAW_MOMENTS_DTYPES
+            and tensor.is_contiguous()
+            for tensor in tensors
+        )
+    )
+
+
+def _multi_tensor_raw_moments(
+    tensors: Sequence[torch.Tensor], device: torch.device
+) -> torch.Tensor:
+    grouped_indices = _group_tensor_indices_by_device_and_dtype(tensors)
+    if len(grouped_indices) == 1:
+        return _multi_tensor_raw_moments_for_group(tensors, device)
+
+    rows = torch.empty(
+        (len(tensors), len(RAW_MOMENT_FIELDS)), dtype=_RAW_MOMENTS_DTYPE, device=device
+    )
+    for indices in grouped_indices.values():
+        group_tensors = [tensors[index] for index in indices]
+        group_rows = _multi_tensor_raw_moments_for_group(group_tensors, group_tensors[0].device)
+        rows[torch.tensor(indices, dtype=torch.long, device=device)] = group_rows.to(device=device)
+    return rows
+
+
+def _multi_tensor_raw_moments_for_group(
+    tensors: Sequence[torch.Tensor], device: torch.device
+) -> torch.Tensor:
+    split_tensors, source_indices = _split_tensors_for_multi_tensor_raw_moments(tensors)
+    device = split_tensors[0].device
+    dummy_overflow_buf = torch.zeros(1, dtype=torch.int, device=device)
+    split_rows = multi_tensor_applier(multi_tensor_raw_moments, dummy_overflow_buf, [split_tensors])
+    if len(split_tensors) == len(tensors):
+        return split_rows
+
+    rows = torch.zeros(
+        (len(tensors), len(RAW_MOMENT_FIELDS)), dtype=_RAW_MOMENTS_DTYPE, device=device
+    )
+    rows.index_add_(0, torch.tensor(source_indices, dtype=torch.long, device=device), split_rows)
+    return rows
+
+
+def _split_tensors_for_multi_tensor_raw_moments(
+    tensors: Sequence[torch.Tensor],
+) -> tuple[list[torch.Tensor], list[int]]:
+    split_tensors = []
+    source_indices = []
+    for index, tensor in enumerate(tensors):
+        flat_tensor = tensor.detach().view(-1)
+        if flat_tensor.numel() == 0 or flat_tensor.numel() <= _MAX_MULTI_TENSOR_RAW_MOMENTS_NUMEL:
+            split_tensors.append(flat_tensor)
+            source_indices.append(index)
+            continue
+
+        for start in range(0, flat_tensor.numel(), _MAX_MULTI_TENSOR_RAW_MOMENTS_NUMEL):
+            length = min(_MAX_MULTI_TENSOR_RAW_MOMENTS_NUMEL, flat_tensor.numel() - start)
+            split_tensors.append(flat_tensor.narrow(0, start, length))
+            source_indices.append(index)
+    return split_tensors, source_indices
+
+
+def _group_tensor_indices_by_device_and_dtype(
+    tensors: Sequence[torch.Tensor],
+) -> dict[tuple[torch.device, torch.dtype], list[int]]:
+    groups = {}
+    for index, tensor in enumerate(tensors):
+        groups.setdefault((tensor.device, tensor.dtype), []).append(index)
+    return groups
+
+
+def raw_moment_row_to_dict(row: Sequence[float]) -> dict[str, float]:
+    """Convert a raw-moment row to a JSON-serializable mapping."""
+    return {field: float(row[idx]) for idx, field in enumerate(RAW_MOMENT_FIELDS)}
+
+
+def _get_expert_model_parallel_rank_size(
+    expert_model_parallel_group: torch.distributed.ProcessGroup | None,
+) -> tuple[int, int]:
+    return get_pg_rank(expert_model_parallel_group), get_pg_size(expert_model_parallel_group)
+
+
+def _registry_cache_key(
+    model_chunks: Sequence[torch.nn.Module],
+    expert_model_parallel_group: torch.distributed.ProcessGroup | None,
+    expert_model_parallel_rank: int,
+    expert_model_parallel_size: int,
+) -> tuple[tuple[int, ...], int | None, int, int]:
+    group_id = id(expert_model_parallel_group) if expert_model_parallel_group is not None else None
+    return (
+        tuple(id(model_chunk) for model_chunk in model_chunks),
+        group_id,
+        expert_model_parallel_rank,
+        expert_model_parallel_size,
+    )
diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py
index d852afe5d59..8cfa2c18594 100644
--- a/megatron/core/recompute.py
+++ b/megatron/core/recompute.py
@@ -9,6 +9,10 @@
 from megatron.core.fp4_utils import get_fp4_context
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.layer_boundary_observer import (
+    observe_transformer_layer_input,
+    observe_transformer_layer_output,
+)
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_layer import TransformerLayer
 
@@ -58,6 +62,7 @@ def custom_forward(
                 # Use self.layers[index] (not self._get_layer) so this
                 # function works for both TransformerBlock and HybridStack.
                 layer = self.layers[index]
+                observe_transformer_layer_input(self, layer, hidden_states)
 
                 # Get appropriate inner quantization context
                 if use_inner_quantization_context:
@@ -101,6 +106,7 @@ def custom_forward(
                 # Some layer paths may still return a tuple (defensive).
                 if isinstance(hidden_states, tuple):
                     hidden_states = hidden_states[0]
+                observe_transformer_layer_output(self, layer, hidden_states)
             return hidden_states, context
 
         return custom_forward
diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py
index 210f39fa217..50ae0421d79 100644
--- a/megatron/core/transformer/cuda_graphs.py
+++ b/megatron/core/transformer/cuda_graphs.py
@@ -8,14 +8,15 @@
 import os
 import time
 from collections import defaultdict
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
+from contextvars import ContextVar
 from copy import deepcopy
 from dataclasses import dataclass, is_dataclass
 from enum import Enum
 from functools import partial
 from itertools import chain, zip_longest
 from math import ceil
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List
 
 import torch
 from torch.utils._pytree import tree_map as tree_map_pyt
@@ -66,6 +67,9 @@
 
 _IS_GRAPH_CAPTURING = False
 _IS_GRAPH_WARMUP = False
+_CUDA_GRAPH_REPLAY_SUSPENDED: ContextVar[bool] = ContextVar(
+    "cuda_graph_replay_suspended", default=False
+)
 logger = logging.getLogger(__name__)
 
 
@@ -111,6 +115,21 @@ def is_graph_capturing():
     return _IS_GRAPH_CAPTURING
 
 
+def is_cuda_graph_replay_suspended() -> bool:
+    """Return whether CUDA graph dispatch is temporarily suspended."""
+    return _CUDA_GRAPH_REPLAY_SUSPENDED.get()
+
+
+@contextmanager
+def suspend_cuda_graph_replay() -> Iterator[None]:
+    """Temporarily run graphable modules through their eager paths."""
+    token = _CUDA_GRAPH_REPLAY_SUSPENDED.set(True)
+    try:
+        yield
+    finally:
+        _CUDA_GRAPH_REPLAY_SUSPENDED.reset(token)
+
+
 def _set_capture_start():
     """Set graph capture has started."""
     global _IS_GRAPH_CAPTURING
@@ -1393,7 +1412,7 @@ def __init__(
             func = getattr(base_module, function_name)
 
             def wrapped_func(*args, eager=False, cache_key=None, **kwargs):
-                if eager:
+                if eager or is_cuda_graph_replay_suspended():
                     return func(*args, **kwargs)
                 out = self(base_module, args, kwargs, cache_key=cache_key)
                 # Unwrap single-element tuple to match the original function's return type.
diff --git a/megatron/core/transformer/layer_boundary_observer.py b/megatron/core/transformer/layer_boundary_observer.py
new file mode 100644
index 00000000000..50fec7fe489
--- /dev/null
+++ b/megatron/core/transformer/layer_boundary_observer.py
@@ -0,0 +1,51 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Scoped observation of transformer-layer input and output boundaries."""
+
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Literal
+
+from torch import Tensor, nn
+
+LayerBoundary = Literal["input", "output"]
+LayerBoundaryObserver = Callable[[nn.Module, nn.Module, LayerBoundary, Tensor], None]
+
+_LAYER_BOUNDARY_OBSERVER: ContextVar[LayerBoundaryObserver | None] = ContextVar(
+    "layer_boundary_observer", default=None
+)
+
+
+@contextmanager
+def observe_transformer_layer_boundaries(observer: LayerBoundaryObserver) -> Iterator[None]:
+    """Make ``observer`` active for transformer-layer boundary notifications."""
+    token = _LAYER_BOUNDARY_OBSERVER.set(observer)
+    try:
+        yield
+    finally:
+        _LAYER_BOUNDARY_OBSERVER.reset(token)
+
+
+def observe_transformer_layer_input(
+    stack: nn.Module, layer: nn.Module, hidden_states: Tensor
+) -> None:
+    """Notify the active observer of the decoder's initial residual stream."""
+    if getattr(layer, "layer_number", None) != 1:
+        return
+    _observe_transformer_layer_boundary(stack, layer, "input", hidden_states)
+
+
+def observe_transformer_layer_output(
+    stack: nn.Module, layer: nn.Module, hidden_states: Tensor
+) -> None:
+    """Notify the active observer of a post-layer residual stream."""
+    _observe_transformer_layer_boundary(stack, layer, "output", hidden_states)
+
+
+def _observe_transformer_layer_boundary(
+    stack: nn.Module, layer: nn.Module, boundary: LayerBoundary, hidden_states: Tensor
+) -> None:
+    observer = _LAYER_BOUNDARY_OBSERVER.get()
+    if observer is not None:
+        observer(stack, layer, boundary, hidden_states)
diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py
index 558b1b07a15..fb432c1a5de 100644
--- a/megatron/core/transformer/module.py
+++ b/megatron/core/transformer/module.py
@@ -363,6 +363,10 @@ def _should_call_te_cudagraph(self, *args, **kwargs):
         )
 
     def __call__(self, *args, **kwargs):
+        from megatron.core.transformer.cuda_graphs import is_cuda_graph_replay_suspended
+
+        if is_cuda_graph_replay_suspended():
+            return super().__call__(*args, **kwargs)
         if self._should_call_local_cudagraph(*args, **kwargs):
             return self.cudagraph_manager(self, args, kwargs)
         elif self._should_call_te_cudagraph(*args, **kwargs):
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index 0415035ffbe..db6f2b68549 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -23,6 +23,10 @@
 from megatron.core.recompute import checkpointed_forward
 from megatron.core.transformer.cuda_graphs import annotate_first_last_layer
 from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType
+from megatron.core.transformer.layer_boundary_observer import (
+    observe_transformer_layer_input,
+    observe_transformer_layer_output,
+)
 from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.torch_norm import LayerNormBuilder
@@ -660,6 +664,7 @@ def forward(
                     else:
                         inner_quantization_context = nullcontext()
 
+                    observe_transformer_layer_input(self, layer, hidden_states)
                     with self.offload_context, inner_quantization_context:
                         hidden_states, context = layer(
                             hidden_states=hidden_states,
@@ -676,6 +681,7 @@ def forward(
                             sequence_len_offset=sequence_len_offset,
                             padding_mask=padding_mask,
                         )
+                    observe_transformer_layer_output(self, layer, hidden_states)
 
                     if (
                         torch.is_grad_enabled()
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index f6ea382077e..b1317ffe9f2 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -18,7 +18,12 @@
 from megatron.core.inference.utils import InferenceMode
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup, make_weakref
+from megatron.core.transformer.cuda_graphs import (
+    is_cuda_graph_replay_suspended,
+    is_graph_capturing,
+    is_graph_warmup,
+    make_weakref,
+)
 from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType
 from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
 from megatron.core.transformer.mlp import MLP
@@ -1716,6 +1721,15 @@ def _forward_mlp(
                 "alongside inference."
             )
 
+        if self.use_partial_cudagraphs and is_cuda_graph_replay_suspended():
+            # The partial-graph path restores graph-owned dispatcher metadata before expert
+            # dispatch. An eager router produces its own metadata, so run the complete eager
+            # MoE path instead when graph replay is suspended.
+            self.mlp.fwd_execution_map = ["route", "expert_compute", "postprocess"]
+            return super()._forward_mlp(
+                hidden_states, padding_mask=padding_mask, packed_seq_params=packed_seq_params
+            )
+
         def _forward_mlp_partial_cudagraphs(
             hidden_states, inference_context=None, padding_mask=None
         ):
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 8174fbd5eb3..87e6664a95a 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -385,6 +385,37 @@ def tuple_type(x):
     assert isinstance(x, str)
     return tuple(int(i) for i in x.strip('()').split(','))
 
+
+def _validate_raw_moment_logging_args(args):
+    stats_logging_enabled = any(
+        interval > 0
+        for interval in (
+            args.log_param_stats_interval,
+            args.log_wgrad_stats_interval,
+            args.log_activation_stats_interval,
+            args.log_dgrad_stats_interval,
+            args.log_residual_stats_interval,
+            args.log_residual_grad_stats_interval,
+        )
+    )
+    if stats_logging_enabled and (args.use_megatron_fsdp or args.use_torch_fsdp2):
+        raise ValueError(
+            'Raw-moment statistics logging is not supported with '
+            '--use-megatron-fsdp or --use-torch-fsdp2.'
+        )
+    if (
+        (
+            args.log_residual_stats_interval > 0
+            or args.log_residual_grad_stats_interval > 0
+        )
+        and args.overlap_moe_expert_parallel_comm
+    ):
+        raise ValueError(
+            'Residual and residual-gradient statistics logging is not supported with '
+            '--overlap-moe-expert-parallel-comm.'
+        )
+
+
 def validate_args(args, defaults={}):
 
     # Prep for checkpoint conversion.
@@ -626,6 +657,8 @@ def validate_args(args, defaults={}):
         else:
             setattr(args, key, defaults[key])
 
+    _validate_raw_moment_logging_args(args)
+
     if args.data_path is not None and args.split is None:
         legacy_default_split_value = '969, 30, 1'
         warn_rank_0('Please specify --split when using --data-path. Using legacy default value '
diff --git a/megatron/training/config/training_config.py b/megatron/training/config/training_config.py
index fb5598b8d42..f017a00bb88 100644
--- a/megatron/training/config/training_config.py
+++ b/megatron/training/config/training_config.py
@@ -250,6 +250,39 @@ class LoggerConfig:
     log_params_norm: bool = False
     """If set, calculate and log parameters norm."""
 
+    log_param_stats_interval: int = 0
+    """Training-step interval for logging count and raw sums of powers 1-4 for each parameter
+    separately (keyed by parameter name) to JSONL statistics files. Values of 0 or less disable
+    logging."""
+
+    log_wgrad_stats_interval: int = 0
+    """Training-step interval for logging count and raw sums of powers 1-4 for each parameter's
+    pre-clipping gradient separately (keyed by parameter name) to JSONL statistics files. Values
+    of 0 or less disable logging."""
+
+    log_activation_stats_interval: int = 0
+    """Training-step interval for logging count and raw sums of powers 1-4 for activations keyed
+    by module site to JSONL statistics files. Will suspend CUDA graphs for the logged steps.
+    Values of 0 or less disable logging."""
+
+    log_dgrad_stats_interval: int = 0
+    """Training-step interval for logging count and raw sums of powers 1-4 for backward data
+    gradients keyed by module site to JSONL statistics files. Will suspend CUDA graphs for the
+    logged steps. Values of 0 or less disable logging."""
+
+    log_residual_stats_interval: int = 0
+    """Training-step interval for logging count and raw sums of powers 1-4 for the initial decoder
+    residual stream and each post-layer residual stream. Values of 0 or less disable logging."""
+
+    log_residual_grad_stats_interval: int = 0
+    """Training-step interval for logging count and raw sums of powers 1-4 for the backward
+    gradients of the initial decoder residual stream and each post-layer residual stream. Values
+    of 0 or less disable logging."""
+
+    statistics_log_dir: str | None = None
+    """Directory for high-cardinality JSONL statistics. If unset, statistics use
+    tensorboard_dir when available, then the --save directory as a fallback."""
+
     log_throughput: bool = False
     """If set, calculate and log throughput per GPU."""
 
diff --git a/megatron/training/raw_moment_logging.py b/megatron/training/raw_moment_logging.py
new file mode 100644
index 00000000000..530beb55513
--- /dev/null
+++ b/megatron/training/raw_moment_logging.py
@@ -0,0 +1,613 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Activation, dgrad, residual-stream, and residual-dgrad raw-moment logging."""
+
+from __future__ import annotations
+
+import re
+from collections import OrderedDict
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Iterable, Iterator
+
+import torch
+import torch.nn as nn
+
+from megatron.core import parallel_state
+from megatron.core.per_parameter_stats import (
+    RAW_MOMENT_FIELDS,
+    raw_moment_row,
+    raw_moment_row_to_dict,
+)
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.layer_boundary_observer import (
+    LayerBoundary,
+    observe_transformer_layer_boundaries,
+)
+from megatron.core.transformer.moe.router import Router
+from megatron.core.utils import unwrap_model
+
+from .activation_logging import LINEAR_TYPES
+
+_LAYER_NAME_PATTERN = re.compile(r"layers\.(\d+)")
+_COLUMN_PARALLEL_CLASS_NAMES = {
+    "TEColumnParallelLinear",
+    "TELayerNormColumnParallelLinear",
+    "TEColumnParallelGroupedLinear",
+}
+_ROW_PARALLEL_CLASS_NAMES = {"TERowParallelLinear", "TERowParallelGroupedLinear"}
+_SEQUENCE_PARALLEL_CLASS_NAMES = {"TENorm"}
+
+
+@dataclass(frozen=True)
+class _SitePolicy:
+    reduce_groups: tuple[torch.distributed.ProcessGroup, ...]
+    owner_groups: tuple[torch.distributed.ProcessGroup, ...]
+
+
+@dataclass
+class _RawMomentSite:
+    name: str
+    policy: _SitePolicy
+    moments: torch.Tensor | None = None
+
+
+@dataclass
+class _RawMomentCollection:
+    sites: OrderedDict[str | tuple[int, int, LayerBoundary], _RawMomentSite] = field(
+        default_factory=OrderedDict
+    )
+    latest: list[tuple[str, dict[str, float]]] | None = None
+
+
+class RawMomentLogger:
+    """Collect activation, dgrad, residual-stream, and residual-dgrad raw moments."""
+
+    def __init__(self):
+        self._activation_hooks: list[torch.utils.hooks.RemovableHandle] = []
+        self._dgrad_hooks: list[torch.utils.hooks.RemovableHandle] = []
+        self._residual_dgrad_hooks: list[
+            tuple[torch.utils.hooks.RemovableHandle, torch.autograd.graph.GradientEdge]
+        ] = []
+        self._activations = _RawMomentCollection()
+        self._dgrads = _RawMomentCollection()
+        self._residuals = _RawMomentCollection()
+        self._residual_dgrads = _RawMomentCollection()
+
+    def register_activation_hooks(self, model: Iterable[nn.Module] | nn.Module) -> None:
+        """Register forward hooks for activation raw moments."""
+        assert not self._activation_hooks
+        self._activations.sites.clear()
+        for module, input_site, output_site in _iter_hook_modules(model):
+            input_site = self._activations.sites.setdefault(input_site.name, input_site)
+            if output_site is not None:
+                output_site = self._activations.sites.setdefault(output_site.name, output_site)
+
+            def hook(_, args, __kwargs, output, input_site=input_site, output_site=output_site):
+                if not torch.is_grad_enabled():
+                    return
+                self._add_tensor(input_site, _first_item(args))
+                if output_site is not None:
+                    self._add_tensor(output_site, _first_item(output))
+
+            self._activation_hooks.append(module.register_forward_hook(hook, with_kwargs=True))
+
+    def register_dgrad_hooks(self, model: Iterable[nn.Module] | nn.Module) -> None:
+        """Register backward hooks for dgrad raw moments."""
+        assert not self._dgrad_hooks
+        self._dgrads.sites.clear()
+        for module, input_site, output_site in _iter_hook_modules(model):
+            input_site = self._dgrads.sites.setdefault(input_site.name, input_site)
+            if output_site is not None:
+                output_site = self._dgrads.sites.setdefault(output_site.name, output_site)
+
+            def hook(_, grad_input, grad_output, input_site=input_site, output_site=output_site):
+                if output_site is not None:
+                    self._add_tensor(output_site, _first_item(grad_output))
+                self._add_tensor(input_site, _first_item(grad_input))
+
+            self._dgrad_hooks.append(module.register_full_backward_hook(hook))
+
+    def prepare_residual_logging(
+        self,
+        model: Iterable[nn.Module] | nn.Module,
+        capture_residuals: bool = True,
+        capture_dgrads: bool = False,
+    ) -> None:
+        """Create stable residual and residual-dgrad sites for the decoder layers in ``model``."""
+        assert not self._residual_dgrad_hooks
+        self._residuals.sites.clear()
+        self._residual_dgrads.sites.clear()
+        residual_sites_by_name: dict[str, _RawMomentSite] = {}
+        residual_dgrad_sites_by_name: dict[str, _RawMomentSite] = {}
+        for key, name, policy in _iter_residual_sites(model):
+            if capture_residuals:
+                self._residuals.sites[key] = residual_sites_by_name.setdefault(
+                    name, _RawMomentSite(name, policy)
+                )
+            if capture_dgrads:
+                self._residual_dgrads.sites[key] = residual_dgrad_sites_by_name.setdefault(
+                    name, _RawMomentSite(name, policy)
+                )
+
+    def record_residual_boundary(
+        self, stack: nn.Module, layer: nn.Module, boundary: LayerBoundary, tensor: torch.Tensor
+    ) -> None:
+        """Accumulate one observed residual-stream boundary."""
+        # Activation checkpointing reaches this path again under grad during recomputation.
+        if not torch.is_grad_enabled():
+            return
+        key = (id(stack), id(layer), boundary)
+        site = self._residuals.sites.get(key)
+        if site is not None:
+            self._add_tensor(site, tensor)
+        dgrad_site = self._residual_dgrads.sites.get(key)
+        if dgrad_site is not None and tensor.requires_grad:
+            # CUDA graph replay may reuse this Tensor object with a new grad_fn.
+            edge = torch.autograd.graph.get_gradient_edge(tensor)
+
+            def record_dgrad(
+                grad_outputs: tuple[torch.Tensor | None, ...],
+                output_nr: int = edge.output_nr,
+                site: _RawMomentSite = dgrad_site,
+            ) -> None:
+                self._add_tensor(site, grad_outputs[output_nr])
+
+            handle = edge.node.register_prehook(record_dgrad)
+            # Retain the edge's ownership token until backward and remove the hook afterward.
+            self._residual_dgrad_hooks.append((handle, edge))
+
+    def finalize_activation_raw_moments_by_layer(self) -> None:
+        """Reduce and cache activation raw moments for later logging."""
+        self._finalize_collection(self._activations)
+
+    def finalize_dgrad_raw_moments_by_layer(self) -> None:
+        """Reduce and cache dgrad raw moments for later logging."""
+        self._finalize_collection(self._dgrads)
+
+    def finalize_residual_raw_moments_by_layer(self) -> None:
+        """Reduce and cache residual-stream raw moments for later logging."""
+        self._finalize_collection(self._residuals)
+
+    def finalize_residual_dgrad_raw_moments_by_layer(self) -> None:
+        """Reduce and cache residual-stream dgrad raw moments for later logging."""
+        self._finalize_collection(self._residual_dgrads)
+        for hook, _ in self._residual_dgrad_hooks:
+            hook.remove()
+        self._residual_dgrad_hooks.clear()
+
+    def consume_activation_raw_moments_by_layer(self) -> list[tuple[str, dict[str, float]]] | None:
+        """Return and clear the latest activation raw moments."""
+        return self._consume_collection(self._activations)
+
+    def consume_dgrad_raw_moments_by_layer(self) -> list[tuple[str, dict[str, float]]] | None:
+        """Return and clear the latest dgrad raw moments."""
+        return self._consume_collection(self._dgrads)
+
+    def consume_residual_raw_moments_by_layer(self) -> list[tuple[str, dict[str, float]]] | None:
+        """Return and clear the latest residual-stream raw moments."""
+        return self._consume_collection(self._residuals)
+
+    def consume_residual_dgrad_raw_moments_by_layer(
+        self,
+    ) -> list[tuple[str, dict[str, float]]] | None:
+        """Return and clear the latest residual-stream dgrad raw moments."""
+        return self._consume_collection(self._residual_dgrads)
+
+    def remove_activation_hooks(self) -> None:
+        """Remove activation raw-moment hooks."""
+        for hook in self._activation_hooks:
+            hook.remove()
+        self._activation_hooks.clear()
+
+    def remove_dgrad_hooks(self) -> None:
+        """Remove dgrad raw-moment hooks."""
+        for hook in self._dgrad_hooks:
+            hook.remove()
+        self._dgrad_hooks.clear()
+
+    def _finalize_collection(self, collection: _RawMomentCollection) -> None:
+        unique_sites = OrderedDict((site.name, site) for site in collection.sites.values())
+        collection.latest = self._finalize_sites(unique_sites.values())
+        collection.sites.clear()
+
+    @staticmethod
+    def _consume_collection(
+        collection: _RawMomentCollection,
+    ) -> list[tuple[str, dict[str, float]]] | None:
+        values = collection.latest
+        collection.latest = None
+        return values
+
+    @torch.no_grad()
+    def _add_tensor(self, site: _RawMomentSite, tensor: torch.Tensor | None) -> None:
+        if tensor is None or not torch.is_tensor(tensor) or not torch.is_floating_point(tensor):
+            return
+        if tensor.numel() == 0:
+            return
+
+        row = raw_moment_row(tensor)
+        if site.moments is None:
+            site.moments = row
+        else:
+            site.moments.add_(row.to(device=site.moments.device))
+
+    def _finalize_sites(
+        self, sites: Iterable[_RawMomentSite]
+    ) -> list[tuple[str, dict[str, float]]] | None:
+        sites = list(sites)
+        if not sites:
+            return None
+
+        device = _select_device(sites)
+        reduced_rows: dict[str, torch.Tensor] = {}
+        sites_by_reduce_key = OrderedDict()
+        for site in sites:
+            key = tuple(id(group) for group in site.policy.reduce_groups)
+            if key not in sites_by_reduce_key:
+                sites_by_reduce_key[key] = (site.policy.reduce_groups, [])
+            sites_by_reduce_key[key][1].append(site)
+
+        for reduce_groups, group_sites in sites_by_reduce_key.values():
+            rows = [
+                site.moments.to(device=device)
+                if site.moments is not None
+                else torch.zeros(len(RAW_MOMENT_FIELDS), dtype=torch.float32, device=device)
+                for site in group_sites
+            ]
+            moments = torch.stack(rows)
+            if _distributed_is_initialized():
+                for group in reduce_groups:
+                    torch.distributed.all_reduce(
+                        moments, op=torch.distributed.ReduceOp.SUM, group=group
+                    )
+            for index, site in enumerate(group_sites):
+                reduced_rows[site.name] = moments[index]
+
+        writer_sites = [site for site in sites if _is_writer(site.policy)]
+        if not writer_sites:
+            return []
+
+        rows = torch.stack([reduced_rows[site.name] for site in writer_sites]).detach().cpu().tolist()
+        values = []
+        for site, row in zip(writer_sites, rows):
+            if row[0] == 0:
+                continue
+            values.append((site.name, raw_moment_row_to_dict(row)))
+        return values
+
+
+def _iter_hook_modules(
+    model: Iterable[nn.Module] | nn.Module,
+) -> Iterable[tuple[nn.Module, _RawMomentSite, _RawMomentSite | None]]:
+    model_chunks = model if isinstance(model, (list, tuple)) else [model]
+    for model_chunk in model_chunks:
+        unwrapped = unwrap_model(model_chunk)
+        for module_name, module in unwrapped.named_modules():
+            if not isinstance(module, LINEAR_TYPES):
+                continue
+            canonical_module_name = _canonical_module_name(unwrapped, module_name, module)
+            input_site_name = f"{canonical_module_name}/input0"
+            output_site_name = f"{canonical_module_name}/output0"
+            output_site = None
+            if not _is_output_layer_logits_site(canonical_module_name):
+                output_site = _RawMomentSite(
+                    output_site_name, _site_policy(module_name, module, "output0")
+                )
+            yield (
+                module,
+                _RawMomentSite(input_site_name, _site_policy(module_name, module, "input0")),
+                output_site,
+            )
+
+
+def _iter_residual_sites(
+    model: Iterable[nn.Module] | nn.Module,
+) -> Iterable[tuple[tuple[int, int, LayerBoundary], str, _SitePolicy]]:
+    for stack_name, stack in _iter_residual_stacks(model):
+        layer_prefix = f"{stack_name}." if stack_name else ""
+        for layer in getattr(stack, "layers"):
+            layer_number = getattr(layer, "layer_number", None)
+            if layer_number is None:
+                continue
+            policy = _residual_site_policy(layer)
+            if layer_number == 1:
+                input_name = f"{stack_name}/input0" if stack_name else "input0"
+                yield (id(stack), id(layer), "input"), input_name, policy
+            output_name = f"{layer_prefix}layers.{layer_number - 1}/output0"
+            yield (id(stack), id(layer), "output"), output_name, policy
+
+
+def _iter_residual_stacks(
+    model: Iterable[nn.Module] | nn.Module,
+) -> Iterable[tuple[str, nn.Module]]:
+    from megatron.core.models.hybrid.hybrid_block import HybridStack
+    from megatron.core.transformer.transformer_block import TransformerBlock
+
+    model_chunks = model if isinstance(model, (list, tuple)) else [model]
+    for model_chunk in model_chunks:
+        unwrapped = unwrap_model(model_chunk)
+        for module_name, module in unwrapped.named_modules():
+            if not isinstance(module, (TransformerBlock, HybridStack)):
+                continue
+            if getattr(module, "is_mtp_layer", False) or "mtp" in module_name.split("."):
+                continue
+            yield module_name, module
+
+
+def _is_output_layer_logits_site(module_name: str) -> bool:
+    return module_name.rsplit(".", maxsplit=1)[-1] == "output_layer"
+
+
+def _canonical_module_name(model_chunk: nn.Module, module_name: str, module: nn.Module) -> str:
+    if "mtp" in module_name or _LAYER_NAME_PATTERN.search(module_name) is None:
+        return module_name
+
+    from megatron.core.transformer.transformer_layer import TransformerLayer
+
+    for transformer_layer in model_chunk.modules():
+        if not isinstance(transformer_layer, TransformerLayer):
+            continue
+        for child in transformer_layer.modules():
+            if child is module:
+                return _LAYER_NAME_PATTERN.sub(
+                    f"layers.{transformer_layer.layer_number - 1}", module_name
+                )
+    return module_name
+
+
+def _site_policy(module_name: str, module: nn.Module, field: str) -> _SitePolicy:
+    reduce_groups = []
+    owner_groups = []
+    is_expert_site = _is_expert_site(module_name, module)
+
+    if is_expert_site:
+        expert_data_parallel_group = _expert_data_parallel_group()
+        if expert_data_parallel_group is not None:
+            reduce_groups.append(expert_data_parallel_group)
+            owner_groups.append(expert_data_parallel_group)
+
+        context_parallel_group = _context_parallel_group()
+        if context_parallel_group is not None:
+            reduce_groups.append(context_parallel_group)
+            owner_groups.append(context_parallel_group)
+    else:
+        dp_cp_group = _data_parallel_with_context_group()
+        if dp_cp_group is not None:
+            reduce_groups.append(dp_cp_group)
+            owner_groups.append(dp_cp_group)
+
+    tp_group = _module_tensor_parallel_group(module)
+    if tp_group is not None:
+        owner_groups.append(tp_group)
+        if _field_is_tensor_parallel_shard(module, field):
+            reduce_groups.append(tp_group)
+
+    if is_expert_site:
+        expert_group = _expert_model_parallel_group()
+        if expert_group is not None:
+            reduce_groups.append(expert_group)
+            owner_groups.append(expert_group)
+
+    return _SitePolicy(tuple(reduce_groups), tuple(owner_groups))
+
+
+def _residual_site_policy(layer: nn.Module) -> _SitePolicy:
+    reduce_groups = []
+    owner_groups = []
+
+    dp_cp_group = _data_parallel_with_context_group()
+    if dp_cp_group is not None:
+        reduce_groups.append(dp_cp_group)
+        owner_groups.append(dp_cp_group)
+
+    tp_group = _module_tensor_parallel_group(layer)
+    if tp_group is not None:
+        owner_groups.append(tp_group)
+        if _sequence_parallel_enabled(layer):
+            reduce_groups.append(tp_group)
+
+    return _SitePolicy(tuple(reduce_groups), tuple(owner_groups))
+
+
+def _field_is_tensor_parallel_shard(module: nn.Module, field: str) -> bool:
+    class_name = module.__class__.__name__
+    parallel_mode = getattr(module, "parallel_mode", None)
+    sequence_parallel = _sequence_parallel_enabled(module)
+
+    if isinstance(module, ColumnParallelLinear) or class_name in _COLUMN_PARALLEL_CLASS_NAMES:
+        if field == "output0":
+            return not bool(getattr(module, "gather_output", False))
+        return sequence_parallel
+
+    if isinstance(module, RowParallelLinear) or class_name in _ROW_PARALLEL_CLASS_NAMES:
+        if field == "input0":
+            return bool(getattr(module, "input_is_parallel", False)) or sequence_parallel
+        return sequence_parallel or bool(getattr(module, "explicit_expert_comm", False))
+
+    if parallel_mode == "column":
+        return field == "output0" or sequence_parallel
+    if parallel_mode == "row":
+        return field == "input0" or sequence_parallel
+
+    if isinstance(module, Router) or class_name in _SEQUENCE_PARALLEL_CLASS_NAMES:
+        return sequence_parallel
+
+    return False
+
+
+def _sequence_parallel_enabled(module: nn.Module) -> bool:
+    config = getattr(module, "config", None)
+    return bool(getattr(module, "sequence_parallel", False) or getattr(config, "sequence_parallel", False))
+
+
+def _is_expert_site(module_name: str, module: nn.Module) -> bool:
+    return bool(
+        getattr(module, "explicit_expert_comm", False)
+        or getattr(module, "is_expert", False)
+        or ".experts." in module_name
+    )
+
+
+def _select_device(sites: list[_RawMomentSite]) -> torch.device:
+    for site in sites:
+        if site.moments is not None:
+            return site.moments.device
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{torch.cuda.current_device()}")
+    return torch.device("cpu")
+
+
+def _first_item(value):
+    if isinstance(value, (list, tuple)):
+        return value[0] if value else None
+    return value
+
+
+def _distributed_is_initialized() -> bool:
+    return torch.distributed.is_available() and torch.distributed.is_initialized()
+
+
+def _model_parallel_is_initialized() -> bool:
+    return _distributed_is_initialized() and parallel_state.is_initialized()
+
+
+def _active_group(group: torch.distributed.ProcessGroup | None):
+    if group is None:
+        return None
+    try:
+        return group if group.size() > 1 else None
+    except RuntimeError:
+        return None
+
+
+def _data_parallel_with_context_group():
+    if not _model_parallel_is_initialized():
+        return None
+    return _active_group(parallel_state.get_data_parallel_group(with_context_parallel=True))
+
+
+def _expert_data_parallel_group():
+    if not _model_parallel_is_initialized():
+        return None
+    return _active_group(parallel_state.get_expert_data_parallel_group())
+
+
+def _context_parallel_group():
+    if not _model_parallel_is_initialized():
+        return None
+    return _active_group(parallel_state.get_context_parallel_group())
+
+
+def _module_tensor_parallel_group(module: nn.Module):
+    if not _model_parallel_is_initialized():
+        return None
+    group = getattr(module, "tp_group", None) or getattr(module, "_tp_group", None)
+    if group is None:
+        group = parallel_state.get_tensor_model_parallel_group()
+    return _active_group(group)
+
+
+def _expert_model_parallel_group():
+    if not _model_parallel_is_initialized():
+        return None
+    return _active_group(parallel_state.get_expert_model_parallel_group())
+
+
+def _is_writer(policy: _SitePolicy) -> bool:
+    if not _distributed_is_initialized():
+        return True
+    for group in policy.owner_groups:
+        if torch.distributed.get_rank(group=group) != 0:
+            return False
+    return True
+
+
+_LOGGER: RawMomentLogger | None = None
+
+
+def _get_logger() -> RawMomentLogger:
+    global _LOGGER
+    if _LOGGER is None:
+        _LOGGER = RawMomentLogger()
+    return _LOGGER
+
+
+def _require_logger() -> RawMomentLogger:
+    assert _LOGGER is not None, "No RawMomentLogger has been initialised"
+    return _LOGGER
+
+
+def enable_activation_raw_moment_logging(model: Iterable[nn.Module] | nn.Module) -> None:
+    """Enable activation raw-moment logging on ``model``."""
+    _get_logger().register_activation_hooks(model)
+
+
+def finalize_activation_raw_moments_by_layer() -> None:
+    """Reduce and cache activation raw moments."""
+    _require_logger().finalize_activation_raw_moments_by_layer()
+
+
+def consume_activation_raw_moments_by_layer() -> list[tuple[str, dict[str, float]]] | None:
+    """Return and clear the latest activation raw moments."""
+    return _require_logger().consume_activation_raw_moments_by_layer()
+
+
+def disable_activation_raw_moment_logging() -> None:
+    """Disable activation raw-moment logging."""
+    _require_logger().remove_activation_hooks()
+
+
+def enable_dgrad_raw_moment_logging(model: Iterable[nn.Module] | nn.Module) -> None:
+    """Enable dgrad raw-moment logging on ``model``."""
+    _get_logger().register_dgrad_hooks(model)
+
+
+def finalize_dgrad_raw_moments_by_layer() -> None:
+    """Reduce and cache dgrad raw moments."""
+    _require_logger().finalize_dgrad_raw_moments_by_layer()
+
+
+def consume_dgrad_raw_moments_by_layer() -> list[tuple[str, dict[str, float]]] | None:
+    """Return and clear the latest dgrad raw moments."""
+    return _require_logger().consume_dgrad_raw_moments_by_layer()
+
+
+def disable_dgrad_raw_moment_logging() -> None:
+    """Disable dgrad raw-moment logging."""
+    _require_logger().remove_dgrad_hooks()
+
+
+@contextmanager
+def capture_residual_raw_moments(
+    model: Iterable[nn.Module] | nn.Module,
+    capture_residuals: bool = True,
+    capture_dgrads: bool = False,
+) -> Iterator[None]:
+    """Capture decoder residual-stream and residual-dgrad raw moments within this context."""
+    logger = _get_logger()
+    logger.prepare_residual_logging(
+        model, capture_residuals=capture_residuals, capture_dgrads=capture_dgrads
+    )
+    with observe_transformer_layer_boundaries(logger.record_residual_boundary):
+        yield
+
+
+def finalize_residual_raw_moments_by_layer() -> None:
+    """Reduce and cache residual-stream raw moments."""
+    _require_logger().finalize_residual_raw_moments_by_layer()
+
+
+def consume_residual_raw_moments_by_layer() -> list[tuple[str, dict[str, float]]] | None:
+    """Return and clear the latest residual-stream raw moments."""
+    return _require_logger().consume_residual_raw_moments_by_layer()
+
+
+def finalize_residual_dgrad_raw_moments_by_layer() -> None:
+    """Reduce and cache residual-stream dgrad raw moments."""
+    _require_logger().finalize_residual_dgrad_raw_moments_by_layer()
+
+
+def consume_residual_dgrad_raw_moments_by_layer() -> list[tuple[str, dict[str, float]]] | None:
+    """Return and clear the latest residual-stream dgrad raw moments."""
+    return _require_logger().consume_residual_dgrad_raw_moments_by_layer()
diff --git a/megatron/training/statistics_logging.py b/megatron/training/statistics_logging.py
new file mode 100644
index 00000000000..20a09be32bb
--- /dev/null
+++ b/megatron/training/statistics_logging.py
@@ -0,0 +1,66 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+
+"""JSON Lines logging for high-cardinality training statistics."""
+
+import json
+import os
+from collections.abc import Iterable
+from contextlib import nullcontext
+
+import torch
+
+
+def _nvtx_range(message: str):
+    if torch.cuda.is_available():
+        return torch.cuda.nvtx.range(message)
+    return nullcontext()
+
+
+def _get_rank() -> int:
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    return 0
+
+
+def append_training_stat(log_dir: str, stat_name: str, record: dict, rank: int | None = None):
+    """Append one JSONL record for a training statistic.
+
+    Each rank writes to its own file under ``{log_dir}/training_stats/{stat_name}/``.
+    Callers decide which ranks should write; this function only handles the file layout.
+    """
+    rank = _get_rank() if rank is None else rank
+    stat_dir = os.path.join(log_dir, "training_stats", stat_name)
+    os.makedirs(stat_dir, exist_ok=True)
+    filepath = os.path.join(stat_dir, f"rank{rank}.jsonl")
+    with _nvtx_range("training_stats.json_dumps"):
+        payload = json.dumps(record) + "\n"
+    with _nvtx_range("training_stats.file_write"):
+        with open(filepath, "a") as f:
+            f.write(payload)
+
+
+def save_raw_moments_by_name(
+    log_dir: str,
+    stat_name: str,
+    iteration: int,
+    consumed_train_samples: int,
+    raw_moments_by_name: Iterable[tuple[str, dict[str, float]]],
+    rank: int | None = None,
+) -> None:
+    """Append one named raw moments record."""
+    with _nvtx_range("training_stats.dictionary_construction"):
+        values = {
+            name: {field: float(value) for field, value in raw_moments.items()}
+            for name, raw_moments in raw_moments_by_name
+        }
+        if not values:
+            return
+
+        record = {
+            "iter": iteration,
+            "consumed_train_samples": consumed_train_samples,
+            "stat": stat_name,
+            "values": values,
+        }
+
+    append_training_stat(log_dir, stat_name, record, rank=rank)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 7f91fb195a7..917a14dfd4a 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -114,7 +114,7 @@
     get_rerun_state_machine,
 )
 from megatron.core.resharding.refit import swap_model_weights
-from megatron.core.transformer.cuda_graphs import TECudaGraphHelper
+from megatron.core.transformer.cuda_graphs import TECudaGraphHelper, suspend_cuda_graph_replay
 from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper
 from megatron.core.transformer.module import Float16Module
 from megatron.core.transformer.moe import upcycling_utils
@@ -174,10 +174,27 @@
     get_tokenizer,
     get_wandb_writer,
 )
+from .raw_moment_logging import (
+    capture_residual_raw_moments,
+    consume_activation_raw_moments_by_layer,
+    consume_dgrad_raw_moments_by_layer,
+    consume_residual_dgrad_raw_moments_by_layer,
+    consume_residual_raw_moments_by_layer,
+    disable_activation_raw_moment_logging,
+    disable_dgrad_raw_moment_logging,
+    enable_activation_raw_moment_logging,
+    enable_dgrad_raw_moment_logging,
+    finalize_activation_raw_moments_by_layer,
+    finalize_dgrad_raw_moments_by_layer,
+    finalize_residual_dgrad_raw_moments_by_layer,
+    finalize_residual_raw_moments_by_layer,
+)
+from .statistics_logging import save_raw_moments_by_name
 from .theoretical_memory_usage import report_theoretical_memory
 from .utils import (
     append_to_progress_log,
     calc_params_l2_norm,
+    calc_params_raw_moments_by_param,
     check_adlr_autoresume_termination,
     is_last_rank,
     logical_and_across_model_parallel_group,
@@ -259,6 +276,8 @@
 num_checkpoints_memory_reported = 0
 MAX_NUM_CHECKPOINTS_MEMORY_REPORTED = 3
 
+_STATS_LOG_DIR_WARNING_SHOWN = False
+
 
 def set_startup_timestamps(program_start=None, main_entry=None):
     """Set startup timestamps from the entry script.
@@ -298,6 +317,47 @@ def print_datetime(string, override_timestamp=None):
     print_rank_0(f'[{string}] datetime: {time_str} ')
 
 
+def _get_statistics_log_dir(args):
+    return (
+        getattr(args, 'statistics_log_dir', None)
+        or getattr(args, 'tensorboard_dir', None)
+        or getattr(args, 'save', None)
+    )
+
+
+def _should_write_global_training_stats(args):
+    rank = getattr(args, 'rank', None)
+    if rank is None:
+        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+
+    world_size = getattr(args, 'world_size', None)
+    if world_size is None:
+        world_size = (
+            torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
+        )
+
+    return rank == world_size - 1
+
+
+def _get_expert_model_parallel_group(pg_collection):
+    if pg_collection is not None:
+        expert_model_parallel_group = getattr(pg_collection, 'ep', None)
+        if expert_model_parallel_group is not None:
+            return expert_model_parallel_group
+    return mpu.get_expert_model_parallel_group(check_initialized=False)
+
+
+def _warn_missing_statistics_log_dir():
+    global _STATS_LOG_DIR_WARNING_SHOWN
+    if not _STATS_LOG_DIR_WARNING_SHOWN:
+        print_rank_0(
+            "WARNING: raw-moment statistics logging was requested, but no statistics log "
+            "directory is available. Set --statistics-log-dir, --tensorboard-dir, or --save "
+            "to write high-cardinality JSONL statistics."
+        )
+        _STATS_LOG_DIR_WARNING_SHOWN = True
+
+
 def update_seqlen_stats_from_cu_seqlens(cu_seqlens):
     """Add ``sum(L_i)`` and ``sum(L_i ** 2)`` from one micro-batch's REAL ``cu_seqlens``.
 
@@ -2303,6 +2363,44 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
                                      (iteration + 1) % args.save_wgrads_interval == 0)
     save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and
                                      (iteration + 1) % args.save_dgrads_interval == 0)
+    activation_stats_interval = getattr(args, 'log_activation_stats_interval', 0)
+    dgrad_stats_interval = getattr(args, 'log_dgrad_stats_interval', 0)
+    residual_stats_interval = getattr(args, 'log_residual_stats_interval', 0)
+    residual_grad_stats_interval = getattr(args, 'log_residual_grad_stats_interval', 0)
+    log_activation_stats_in_this_iteration = (
+        activation_stats_interval > 0
+        and iteration is not None
+        and (iteration + 1) % activation_stats_interval == 0
+    )
+    log_dgrad_stats_in_this_iteration = (
+        dgrad_stats_interval > 0
+        and iteration is not None
+        and (iteration + 1) % dgrad_stats_interval == 0
+    )
+    log_residual_stats_in_this_iteration = (
+        residual_stats_interval > 0
+        and iteration is not None
+        and (iteration + 1) % residual_stats_interval == 0
+    )
+    log_residual_grad_stats_in_this_iteration = (
+        residual_grad_stats_interval > 0
+        and iteration is not None
+        and (iteration + 1) % residual_grad_stats_interval == 0
+    )
+    suspend_cuda_graphs_for_stats = (
+        log_activation_stats_in_this_iteration or log_dgrad_stats_in_this_iteration
+    )
+    if (
+        suspend_cuda_graphs_for_stats
+        or log_residual_stats_in_this_iteration
+        or log_residual_grad_stats_in_this_iteration
+    ) and getattr(
+        args, 'cuda_graph_impl', 'none'
+    ) == 'full_iteration':
+        raise RuntimeError(
+            "Activation, dgrad, residual, and residual-gradient statistics logging is not "
+            "supported with full-iteration CUDA graphs."
+        )
     while rerun_state_machine.should_run_forward_backward(data_iterator):
         # Set grad to zero.
         for model_chunk in model:
@@ -2352,20 +2450,42 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
             enable_tokens_per_expert_logging(model, args.save)
         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,
-            model=model,
-            num_microbatches=get_num_microbatches(),
-            seq_length=args.seq_length,
-            micro_batch_size=args.micro_batch_size,
-            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,
-            p2p_communicator=p2p_communicator,
-            pg_collection=pg_collection,
+        if log_activation_stats_in_this_iteration:
+            enable_activation_raw_moment_logging(model)
+        if log_dgrad_stats_in_this_iteration:
+            enable_dgrad_raw_moment_logging(model)
+        cuda_graph_context = (
+            suspend_cuda_graph_replay()
+            if suspend_cuda_graphs_for_stats
+            else nullcontext()
         )
+        residual_raw_moment_context = (
+            capture_residual_raw_moments(
+                model,
+                capture_residuals=log_residual_stats_in_this_iteration,
+                capture_dgrads=log_residual_grad_stats_in_this_iteration,
+            )
+            if (
+                log_residual_stats_in_this_iteration
+                or log_residual_grad_stats_in_this_iteration
+            )
+            else nullcontext()
+        )
+        with cuda_graph_context, residual_raw_moment_context:
+            losses_reduced = forward_backward_func(
+                forward_step_func=forward_step_func,
+                data_iterator=data_iterator,
+                model=model,
+                num_microbatches=get_num_microbatches(),
+                seq_length=args.seq_length,
+                micro_batch_size=args.micro_batch_size,
+                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,
+                p2p_communicator=p2p_communicator,
+                pg_collection=pg_collection,
+            )
         if save_activations_in_this_iteration:
             save_activations(iteration + 1)
             disable_activation_logging()
@@ -2375,6 +2495,16 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
         if save_dgrads_in_this_iteration:
             save_dgrads(iteration + 1)
             disable_dgrad_logging()
+        if log_activation_stats_in_this_iteration:
+            finalize_activation_raw_moments_by_layer()
+            disable_activation_raw_moment_logging()
+        if log_dgrad_stats_in_this_iteration:
+            finalize_dgrad_raw_moments_by_layer()
+            disable_dgrad_raw_moment_logging()
+        if log_residual_stats_in_this_iteration:
+            finalize_residual_raw_moments_by_layer()
+        if log_residual_grad_stats_in_this_iteration:
+            finalize_residual_dgrad_raw_moments_by_layer()
 
         # Advance the router tracer step if active.
         tracer = get_moe_router_tracer()
@@ -2419,6 +2549,16 @@ def _save_state_dict(attr_name, label):
     # Update parameters.
 
     timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time)
+    wgrad_stats_interval = getattr(args, 'log_wgrad_stats_interval', 0)
+    if (
+        optimizer is not None
+        and wgrad_stats_interval > 0
+        and iteration is not None
+        and (iteration + 1) % wgrad_stats_interval == 0
+    ):
+        optimizer.request_grad_raw_moments_by_param(
+            model, expert_model_parallel_group=_get_expert_model_parallel_group(pg_collection)
+        )
     update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
 
     # get max attention logit for logging and run clip_qk()
@@ -3882,9 +4022,118 @@ def trace_handler(p):
         else:
             loss_scale = 1.0
         params_norm = None
+        param_stats_interval = getattr(args, 'log_param_stats_interval', 0)
+        wgrad_stats_interval = getattr(args, 'log_wgrad_stats_interval', 0)
+        activation_stats_interval = getattr(args, 'log_activation_stats_interval', 0)
+        dgrad_stats_interval = getattr(args, 'log_dgrad_stats_interval', 0)
+        residual_stats_interval = getattr(args, 'log_residual_stats_interval', 0)
+        residual_grad_stats_interval = getattr(args, 'log_residual_grad_stats_interval', 0)
 
         if args.log_params_norm:
             params_norm = calc_params_l2_norm(model)
+        if (
+            param_stats_interval > 0
+            and iteration % param_stats_interval == 0
+        ):
+            param_raw_moments_by_param = calc_params_raw_moments_by_param(
+                model,
+                expert_model_parallel_group=_get_expert_model_parallel_group(model_pg_collection),
+            )
+            statistics_log_dir = _get_statistics_log_dir(args)
+            if statistics_log_dir is None:
+                _warn_missing_statistics_log_dir()
+            elif _should_write_global_training_stats(args):
+                save_raw_moments_by_name(
+                    statistics_log_dir,
+                    "param_raw_moments_by_param",
+                    iteration,
+                    args.consumed_train_samples,
+                    param_raw_moments_by_param,
+                )
+        if (
+            wgrad_stats_interval > 0
+            and iteration % wgrad_stats_interval == 0
+            and optimizer is not None
+        ):
+            grad_raw_moments_by_param = optimizer.consume_grad_raw_moments_by_param()
+            statistics_log_dir = _get_statistics_log_dir(args)
+            if statistics_log_dir is None:
+                _warn_missing_statistics_log_dir()
+            elif (
+                grad_raw_moments_by_param is not None
+                and _should_write_global_training_stats(args)
+            ):
+                save_raw_moments_by_name(
+                    statistics_log_dir,
+                    "grad_raw_moments_by_param",
+                    iteration,
+                    args.consumed_train_samples,
+                    grad_raw_moments_by_param,
+                )
+        if (
+            activation_stats_interval > 0
+            and iteration % activation_stats_interval == 0
+        ):
+            activation_raw_moments_by_layer = consume_activation_raw_moments_by_layer()
+            statistics_log_dir = _get_statistics_log_dir(args)
+            if statistics_log_dir is None:
+                _warn_missing_statistics_log_dir()
+            elif activation_raw_moments_by_layer:
+                save_raw_moments_by_name(
+                    statistics_log_dir,
+                    "activation_raw_moments_by_layer",
+                    iteration,
+                    args.consumed_train_samples,
+                    activation_raw_moments_by_layer,
+                )
+        if (
+            dgrad_stats_interval > 0
+            and iteration % dgrad_stats_interval == 0
+        ):
+            dgrad_raw_moments_by_layer = consume_dgrad_raw_moments_by_layer()
+            statistics_log_dir = _get_statistics_log_dir(args)
+            if statistics_log_dir is None:
+                _warn_missing_statistics_log_dir()
+            elif dgrad_raw_moments_by_layer:
+                save_raw_moments_by_name(
+                    statistics_log_dir,
+                    "dgrad_raw_moments_by_layer",
+                    iteration,
+                    args.consumed_train_samples,
+                    dgrad_raw_moments_by_layer,
+                )
+        if (
+            residual_stats_interval > 0
+            and iteration % residual_stats_interval == 0
+        ):
+            residual_raw_moments_by_layer = consume_residual_raw_moments_by_layer()
+            statistics_log_dir = _get_statistics_log_dir(args)
+            if statistics_log_dir is None:
+                _warn_missing_statistics_log_dir()
+            elif residual_raw_moments_by_layer:
+                save_raw_moments_by_name(
+                    statistics_log_dir,
+                    "residual_raw_moments_by_layer",
+                    iteration,
+                    args.consumed_train_samples,
+                    residual_raw_moments_by_layer,
+                )
+        if (
+            residual_grad_stats_interval > 0
+            and iteration % residual_grad_stats_interval == 0
+        ):
+            residual_dgrad_raw_moments_by_layer = consume_residual_dgrad_raw_moments_by_layer()
+            statistics_log_dir = _get_statistics_log_dir(args)
+            if statistics_log_dir is None:
+                _warn_missing_statistics_log_dir()
+            elif residual_dgrad_raw_moments_by_layer:
+                save_raw_moments_by_name(
+                    statistics_log_dir,
+                    "residual_dgrad_raw_moments_by_layer",
+                    iteration,
+                    args.consumed_train_samples,
+                    residual_dgrad_raw_moments_by_layer,
+                )
         if optimizer is not None:
             learning_rate = get_canonical_lr_for_logging(optimizer.param_groups)
         else:
diff --git a/megatron/training/utils/__init__.py b/megatron/training/utils/__init__.py
index 15cd4b26d4f..01a44419ef1 100644
--- a/megatron/training/utils/__init__.py
+++ b/megatron/training/utils/__init__.py
@@ -2,6 +2,7 @@
 
 from megatron.training.utils.common_utils import (
     calc_params_l2_norm,
+    calc_params_raw_moments_by_param,
     calc_dtensor_params_l2_norm,
     average_losses_across_data_parallel_group,
     reduce_max_stat_across_model_parallel_group,
diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py
index 316bf598fec..23ec01e60ec 100644
--- a/megatron/training/utils/common_utils.py
+++ b/megatron/training/utils/common_utils.py
@@ -16,6 +16,11 @@
 from megatron.core._slurm_utils import resolve_slurm_local_rank
 from megatron.core.dist_checkpointing.strategies.nvrx import has_nvrx_async_support
 from megatron.core.msc_utils import open_file
+from megatron.core.per_parameter_stats import (
+    NamedTensorBucket,
+    get_or_create_per_parameter_stat_registry,
+    reduce_raw_moments_by_param,
+)
 
 try:
     from transformer_engine.pytorch.optimizers import multi_tensor_applier, multi_tensor_l2norm
@@ -47,13 +52,67 @@
 )
 from megatron.training import get_adlr_autoresume, get_args, get_timers
 
+# Relative tolerance for the raw-moments self-check: sqrt(sum_2), recombined into an aggregate,
+# must match the independently-computed scalar norm to within this much.
+_RAW_MOMENTS_BY_PARAM_NORM_RTOL = 1e-2
+
 
 def calc_params_l2_norm(model, force_create_fp32_copy=False):
-    """Calculate l2 norm of parameters"""
+    """Calculate l2 norm of parameters."""
+    return _calc_params_l2_norm_or_raw_moments(
+        model, force_create_fp32_copy=force_create_fp32_copy, raw_moments_by_param=False
+    )
+
+
+def calc_params_raw_moments_by_param(
+    model, force_create_fp32_copy=False, expert_model_parallel_group=None
+):
+    """Calculate per-parameter raw moments of parameters."""
+    return _calc_params_l2_norm_or_raw_moments(
+        model,
+        force_create_fp32_copy=force_create_fp32_copy,
+        raw_moments_by_param=True,
+        expert_model_parallel_group=expert_model_parallel_group,
+    )
+
+
+def _calc_params_l2_norm_or_raw_moments(
+    model,
+    force_create_fp32_copy=False,
+    raw_moments_by_param=False,
+    expert_model_parallel_group=None,
+):
+    """Calculate scalar parameter norm or per-parameter raw moments.
+
+    If ``raw_moments_by_param`` is False, returns the aggregate l2 norm as a scalar float.
+
+    If ``raw_moments_by_param`` is True, returns a list of ``(parameter_name, moments)`` tuples.
+    The raw moments are reduced across the same process groups as the aggregate norm.
+
+    Expert parallelism: expert params are named by *local* expert index, which collides
+    across expert-parallel ranks. With ``--moe-grouped-gemm`` each rank's experts are stacked into
+    a single tensor, so the collision is benign. Sequential experts collide on distinct global
+    experts and are not supported here (see the asserts below).
+    """
     args = get_args()
     if not isinstance(model, list):
         model = [model]
 
+    if raw_moments_by_param and getattr(args, 'expert_model_parallel_size', 1) > 1:
+        assert getattr(args, 'moe_grouped_gemm', False), (
+            "calc_params_raw_moments_by_param() with expert parallelism is only supported with "
+            "--moe-grouped-gemm; sequential experts collide on local expert names across expert-"
+            "parallel ranks."
+        )
+
+    if raw_moments_by_param and (
+        getattr(args, 'use_megatron_fsdp', False) or getattr(args, 'use_torch_fsdp2', False)
+    ):
+        raise RuntimeError(
+            "calc_params_raw_moments_by_param() is not implemented for Megatron-FSDP "
+            "or Torch-FSDP2"
+        )
+
     if getattr(args, 'use_megatron_fsdp', False):
         # All Megatron FSDP parameters are expected to be PyTorch DTensor.
         # params_data is a dict of device_mesh -> list of local tensors.
@@ -70,51 +129,114 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
 
         return calc_dtensor_params_l2_norm(params)
 
+    if raw_moments_by_param and expert_model_parallel_group is None:
+        expert_model_parallel_group = mpu.get_expert_model_parallel_group(check_initialized=False)
+
+    raw_moments_registry = (
+        get_or_create_per_parameter_stat_registry(
+            model, expert_model_parallel_group=expert_model_parallel_group
+        )
+        if raw_moments_by_param
+        else None
+    )
+
     # Seperate moe and dense params
     params_data = []
     moe_params_data = []
     sharded_params_data = []
+    sharded_moe_params_data = []
+    # Parallel lists of parameter names, kept in lock-step with the *_params_data lists above.
+    # Only populated/used when raw_moments_by_param=True.
+    params_data_names = []
+    moe_params_data_names = []
+    sharded_params_data_names = []
+    sharded_moe_params_data_names = []
     data_parallel_group = None
 
-    for model_chunk in model:
-        for param in model_chunk.parameters():
-            data_parallel_group = get_data_parallel_group_if_dtensor(param, data_parallel_group)
-            is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param)
-            if not is_not_tp_duplicate:
-                continue
-            assert is_not_tp_duplicate
-            if not getattr(param, 'allreduce', True):
-                assert param_is_not_shared(param)
+    if raw_moments_by_param:
+        named_params = (
+            (param_name, param) for param, param_name in raw_moments_registry.param_to_name.items()
+        )
+    else:
+        named_params = (
+            (name, param)
+            for model_chunk in model
+            for name, param in unwrap_model(model_chunk).named_parameters()
+        )
+
+    for param_name, param in named_params:
+        data_parallel_group = get_data_parallel_group_if_dtensor(param, data_parallel_group)
+        is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param)
+        if not is_not_tp_duplicate:
+            continue
+        assert is_not_tp_duplicate
+        if not getattr(param, 'allreduce', True):
+            assert param_is_not_shared(param)
+            param = to_local_if_dtensor(param)
+            if args.bf16:
+                if not force_create_fp32_copy and hasattr(param, 'main_param'):
+                    if getattr(param, 'main_param_sharded', False):
+                        if param.main_param is not None:
+                            sharded_moe_params_data.append(param.main_param)
+                            sharded_moe_params_data_names.append(param_name)
+                    else:
+                        moe_params_data.append(param.main_param)
+                        moe_params_data_names.append(param_name)
+                else:
+                    # Fallback to original logic of making a fp32 copy of the
+                    # parameter if `.main_param` attribute is not available.
+                    moe_params_data.append(param.data.float())
+                    moe_params_data_names.append(param_name)
+            else:
+                moe_params_data.append(param.data)
+                moe_params_data_names.append(param_name)
+        else:
+            if param_is_not_shared(param):
                 param = to_local_if_dtensor(param)
                 if args.bf16:
                     if not force_create_fp32_copy and hasattr(param, 'main_param'):
                         if getattr(param, 'main_param_sharded', False):
                             if param.main_param is not None:
                                 sharded_params_data.append(param.main_param)
+                                sharded_params_data_names.append(param_name)
                         else:
-                            moe_params_data.append(param.main_param)
+                            params_data.append(param.main_param)
+                            params_data_names.append(param_name)
                     else:
                         # Fallback to original logic of making a fp32 copy of the
                         # parameter if `.main_param` attribute is not available.
-                        moe_params_data.append(param.data.float())
+                        params_data.append(param.data.float())
+                        params_data_names.append(param_name)
                 else:
-                    moe_params_data.append(param.data)
-            else:
-                if param_is_not_shared(param):
-                    param = to_local_if_dtensor(param)
-                    if args.bf16:
-                        if not force_create_fp32_copy and hasattr(param, 'main_param'):
-                            if getattr(param, 'main_param_sharded', False):
-                                if param.main_param is not None:
-                                    sharded_params_data.append(param.main_param)
-                            else:
-                                params_data.append(param.main_param)
-                        else:
-                            # Fallback to original logic of making a fp32 copy of the
-                            # parameter if `.main_param` attribute is not available.
-                            params_data.append(param.data.float())
-                    else:
-                        params_data.append(param.data)
+                    params_data.append(param.data)
+                    params_data_names.append(param_name)
+
+    # Dense params should sum across all model-parallel GPUs (tensor + pipeline).
+    dense_reduce_group = mpu.get_model_parallel_group()
+    # Expert params should sum across all model-parallel GPUs (expert + tensor + pipeline).
+    expert_reduce_group = mpu.get_expert_tensor_model_pipeline_parallel_group()
+
+    if raw_moments_by_param:
+        dense_reduce_groups = (
+            (data_parallel_group,) if data_parallel_group is not None else ()
+        ) + (dense_reduce_group,)
+        buckets = [
+            NamedTensorBucket(params_data_names, params_data, dense_reduce_groups),
+            NamedTensorBucket(
+                sharded_params_data_names,
+                sharded_params_data,
+                (mpu.get_data_parallel_group(with_context_parallel=True), dense_reduce_group),
+            ),
+            NamedTensorBucket(moe_params_data_names, moe_params_data, (expert_reduce_group,)),
+            NamedTensorBucket(
+                sharded_moe_params_data_names,
+                sharded_moe_params_data,
+                (mpu.get_expert_data_parallel_group(), expert_reduce_group),
+            ),
+        ]
+        raw_moments_by_param_result, aggregate_moments = reduce_raw_moments_by_param(
+            raw_moments_registry, buckets
+        )
 
     # Calculate norm.
     dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
@@ -170,12 +292,26 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
     else:
         moe_norm_2 = torch.zeros_like(norm_2)
 
+    if len(sharded_moe_params_data) > 0:
+        dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+        sharded_moe_norm, _ = multi_tensor_applier(
+            multi_tensor_l2norm,
+            dummy_overflow_buf,
+            [sharded_moe_params_data],
+            False,  # no per-parameter norm.
+        )
+        sharded_moe_norm_2 = sharded_moe_norm * sharded_moe_norm
+    else:
+        sharded_moe_norm_2 = torch.zeros_like(norm_2)
+    torch.distributed.all_reduce(
+        sharded_moe_norm_2,
+        op=torch.distributed.ReduceOp.SUM,
+        group=mpu.get_expert_data_parallel_group(),
+    )
+    moe_norm_2 += sharded_moe_norm_2
+
     # Reduce norm across model parallel groups (dense and expert).
-    # Dense params should sum across all model-parallel GPUs (tensor + pipeline).
-    dense_reduce_group = mpu.get_model_parallel_group()
     ranks_in_dense_reduce_group = torch.distributed.get_process_group_ranks(dense_reduce_group)
-    # Expert params should sum across all model-parallel GPUs (expert + tensor + pipeline).
-    expert_reduce_group = mpu.get_expert_tensor_model_pipeline_parallel_group()
     ranks_in_expert_reduce_group = torch.distributed.get_process_group_ranks(expert_reduce_group)
 
     # If dense and expert reduce groups are the same, sum then reduce.
@@ -194,7 +330,23 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
         )
         norm_2 += moe_norm_2
 
-    return norm_2.item() ** 0.5
+    scalar_norm = norm_2.item() ** 0.5
+
+    if raw_moments_by_param:
+        # Self-check: sqrt(sum_2) from the per-parameter raw moments should equal the scalar norm.
+        reconstructed_norm = aggregate_moments["sum_2"] ** 0.5
+        rel_diff = abs(reconstructed_norm - scalar_norm) / scalar_norm if scalar_norm > 0 else 0.0
+        if rel_diff > _RAW_MOMENTS_BY_PARAM_NORM_RTOL:
+            warn_rank_0(
+                "calc_params_raw_moments_by_param(): per-parameter sum_2 recombines to an "
+                f"aggregate of {reconstructed_norm:.6e}, but the directly-computed norm is "
+                f"{scalar_norm:.6e} (relative difference {rel_diff:.2e} > "
+                f"{_RAW_MOMENTS_BY_PARAM_NORM_RTOL:.0e}). The per-parameter reduction is likely "
+                "incorrect for this parallelism configuration; treat the raw moments with caution."
+            )
+        return raw_moments_by_param_result
+
+    return scalar_norm
 
 
 def calc_dtensor_params_l2_norm(params):
diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py
index 94613b7096c..3ad33e4bef5 100644
--- a/tests/unit_tests/test_optimizer.py
+++ b/tests/unit_tests/test_optimizer.py
@@ -505,6 +505,8 @@ def test_mtp_grad_separation():
     class MockOptimizer:
         """Minimal mock of MegatronOptimizer for testing grad filtering."""
 
+        _get_grad_for_grad_norm = MegatronOptimizer._get_grad_for_grad_norm
+        _include_param_in_grad_norm = MegatronOptimizer._include_param_in_grad_norm
         _filter_grads_for_norm = MegatronOptimizer._filter_grads_for_norm
         get_grads_for_grad_norm = MegatronOptimizer.get_grads_for_grad_norm
 
@@ -550,6 +552,8 @@ def test_mtp_grad_separation_no_mtp_params():
     from megatron.core.optimizer.optimizer import MegatronOptimizer
 
     class MockOptimizer:
+        _get_grad_for_grad_norm = MegatronOptimizer._get_grad_for_grad_norm
+        _include_param_in_grad_norm = MegatronOptimizer._include_param_in_grad_norm
         _filter_grads_for_norm = MegatronOptimizer._filter_grads_for_norm
         get_grads_for_grad_norm = MegatronOptimizer.get_grads_for_grad_norm
 
@@ -582,6 +586,8 @@ def test_unregistered_grad_norm_group_raises():
     from megatron.core.optimizer.optimizer import MegatronOptimizer
 
     class MockOptimizer:
+        _get_grad_for_grad_norm = MegatronOptimizer._get_grad_for_grad_norm
+        _include_param_in_grad_norm = MegatronOptimizer._include_param_in_grad_norm
         _filter_grads_for_norm = MegatronOptimizer._filter_grads_for_norm
         get_grads_for_grad_norm = MegatronOptimizer.get_grads_for_grad_norm
 
@@ -633,16 +639,22 @@ def test_mtp_grad_clipping_uses_separate_norms():
     from megatron.core.optimizer.optimizer import MegatronOptimizer
 
     class MockOptimizer:
+        _get_grad_for_grad_norm = MegatronOptimizer._get_grad_for_grad_norm
+        _include_param_in_grad_norm = MegatronOptimizer._include_param_in_grad_norm
         _filter_grads_for_norm = MegatronOptimizer._filter_grads_for_norm
         get_grads_for_grad_norm = MegatronOptimizer.get_grads_for_grad_norm
         get_grad_stats_parallel_group = MegatronOptimizer.get_grad_stats_parallel_group
         has_grad_norm_group = MegatronOptimizer.has_grad_norm_group
         _compute_grad_norms_by_group = MegatronOptimizer._compute_grad_norms_by_group
+        _maybe_record_grad_raw_moments_by_param = (
+            MegatronOptimizer._maybe_record_grad_raw_moments_by_param
+        )
         clip_grad_norm = MegatronOptimizer.clip_grad_norm
 
         def __init__(self, params):
             self.params = list(params)
             self.config = OptimizerConfig(optimizer='adam', lr=0.01)
+            self._per_param_grad_raw_moments_requested = False
 
         def get_parameters(self):
             return self.params
diff --git a/tests/unit_tests/test_parameter_names.py b/tests/unit_tests/test_parameter_names.py
new file mode 100644
index 00000000000..56c69e0130d
--- /dev/null
+++ b/tests/unit_tests/test_parameter_names.py
@@ -0,0 +1,149 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import pytest
+import torch
+
+from megatron.core.parameter_names import CanonicalParameterNameIndex, CanonicalParameterNameMap
+
+
+class _Layer(torch.nn.Module):
+    def __init__(self, layer_number):
+        super().__init__()
+        self.layer_number = layer_number
+        self.weight = torch.nn.Parameter(torch.zeros(1))
+
+
+class _PipelineModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = torch.nn.Module()
+        self.decoder.layers = torch.nn.ModuleList([_Layer(5), _Layer(8)])
+
+
+class _GroupedExpertModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = type("Config", (), {"num_moe_experts": 4})()
+        self.mlp = torch.nn.Module()
+        self.mlp.experts = torch.nn.Module()
+        self.mlp.experts.linear_fc1 = torch.nn.Module()
+        self.mlp.experts.linear_fc1.register_parameter(
+            "weight0", torch.nn.Parameter(torch.zeros(1))
+        )
+        self.mlp.experts.linear_fc1.register_parameter("bias1", torch.nn.Parameter(torch.zeros(1)))
+
+
+class _SequentialExpertModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = type("Config", (), {"num_moe_experts": 4})()
+        self.mlp = torch.nn.Module()
+        self.mlp.experts = torch.nn.Module()
+        self.mlp.experts.local_experts = torch.nn.ModuleList(
+            [torch.nn.Linear(1, 1, bias=False), torch.nn.Linear(1, 1, bias=False)]
+        )
+
+
+class _SingleParamModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.zeros(1))
+
+
+class _TwoParamModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.z = torch.nn.Parameter(torch.zeros(1))
+        self.a = torch.nn.Parameter(torch.zeros(1))
+
+
+def test_canonical_parameter_name_index_is_sorted_and_deduplicated():
+    index = CanonicalParameterNameIndex(["z", "a", "z"])
+
+    assert index.names == ("a", "z")
+    assert dict(index) == {"a": 0, "z": 1}
+
+
+def test_canonical_parameter_names_have_deterministic_local_index():
+    names = CanonicalParameterNameMap(_TwoParamModel())
+
+    assert names.local_index.names == ("a", "z")
+    assert names.local_index["a"] == 0
+    assert names.local_index["z"] == 1
+
+
+def test_canonical_parameter_names_use_global_layer_numbers():
+    model = _PipelineModel()
+    names = CanonicalParameterNameMap(model)
+
+    assert names.name_for_param(model.decoder.layers[0].weight) == "decoder.layers.4.weight"
+    assert names.name_for_param(model.decoder.layers[1].weight) == "decoder.layers.7.weight"
+
+
+def test_canonical_parameter_names_use_global_grouped_expert_numbers():
+    model = _GroupedExpertModel()
+    names = CanonicalParameterNameMap(model, expert_parallel_rank=1, expert_parallel_size=2)
+
+    assert set(names.values()) == {"mlp.experts.linear_fc1.weight2", "mlp.experts.linear_fc1.bias3"}
+
+
+def test_canonical_parameter_names_use_global_sequential_expert_numbers():
+    model = _SequentialExpertModel()
+    names = CanonicalParameterNameMap(model, expert_parallel_rank=1, expert_parallel_size=2)
+
+    assert set(names.values()) == {
+        "mlp.experts.local_experts.2.weight",
+        "mlp.experts.local_experts.3.weight",
+    }
+
+
+def test_canonical_parameter_names_reject_collisions_between_chunks():
+    with pytest.raises(ValueError, match="multiple distinct local parameters"):
+        CanonicalParameterNameMap([_SingleParamModel(), _SingleParamModel()])
+
+
+@pytest.mark.parametrize(
+    ("rank", "size", "match"), [(0, 0, "at least 1"), (-1, 2, "must be in"), (2, 2, "must be in")]
+)
+def test_canonical_parameter_names_validate_expert_topology(rank, size, match):
+    with pytest.raises(ValueError, match=match):
+        CanonicalParameterNameMap(
+            _SingleParamModel(), expert_parallel_rank=rank, expert_parallel_size=size
+        )
+
+
+def test_canonical_parameter_names_reject_uneven_expert_partition():
+    model = _GroupedExpertModel()
+    model.config.num_moe_experts = 5
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        CanonicalParameterNameMap(model, expert_parallel_rank=0, expert_parallel_size=2)
+
+
+def test_all_gather_index_returns_local_index_without_distributed(monkeypatch):
+    names = CanonicalParameterNameMap(_TwoParamModel())
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+
+    assert names.all_gather_index() is names.local_index
+
+
+def test_all_gather_index_collects_deduplicates_and_orders_names(monkeypatch):
+    names = CanonicalParameterNameMap(_SingleParamModel())
+    expected_group = object()
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 2)
+
+    def fake_all_gather_object(output, local_names, group):
+        assert local_names == ("weight",)
+        assert group is expected_group
+        output[:] = [("z", "shared"), ("a", "shared")]
+
+    monkeypatch.setattr(torch.distributed, "all_gather_object", fake_all_gather_object)
+
+    index = names.all_gather_index(expected_group)
+
+    assert index.names == ("a", "shared", "z")
+    assert dict(index) == {"a": 0, "shared": 1, "z": 2}
+    assert names.local_index.names == ("weight",)
diff --git a/tests/unit_tests/test_per_parameter_stats.py b/tests/unit_tests/test_per_parameter_stats.py
new file mode 100644
index 00000000000..45e990cc105
--- /dev/null
+++ b/tests/unit_tests/test_per_parameter_stats.py
@@ -0,0 +1,214 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import pytest
+import torch
+
+from megatron.core import per_parameter_stats as pps
+from megatron.core.parameter_names import CanonicalParameterNameIndex, CanonicalParameterNameMap
+from megatron.core.per_parameter_stats import (
+    NamedTensorBucket,
+    PerParameterStatRegistry,
+    get_or_create_per_parameter_stat_registry,
+    reduce_raw_moments_by_param,
+)
+
+
+class TwoParamModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.a = torch.nn.Parameter(torch.zeros(1))
+        self.b = torch.nn.Parameter(torch.zeros(1))
+
+
+class SequentialExpertModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = type("Config", (), {"num_moe_experts": 4})()
+        self.mlp = torch.nn.Module()
+        self.mlp.experts = torch.nn.Module()
+        self.mlp.experts.local_experts = torch.nn.ModuleList(
+            [torch.nn.Linear(1, 1, bias=False), torch.nn.Linear(1, 1, bias=False)]
+        )
+
+
+class _FakeCudaTensor:
+    device = torch.device("cuda:0")
+    dtype = torch.float32
+
+    def is_contiguous(self) -> bool:
+        return True
+
+
+class _FakeExpertGroup:
+    def __init__(self, rank=0, size=1):
+        self._rank = rank
+        self._size = size
+
+    def rank(self):
+        return self._rank
+
+    def size(self):
+        return self._size
+
+
+def test_reduce_raw_moments_by_param_on_cpu():
+    registry = PerParameterStatRegistry(TwoParamModel())
+
+    values, aggregate_moments = reduce_raw_moments_by_param(
+        registry,
+        [
+            NamedTensorBucket(
+                names=["a", "a", "b"],
+                tensors=[torch.tensor([1.0, 2.0]), torch.tensor([2.0, 4.0]), torch.tensor([3.0])],
+            )
+        ],
+    )
+
+    assert dict(values) == {
+        "a": {
+            "count": pytest.approx(4.0),
+            "sum_1": pytest.approx(9.0),
+            "sum_2": pytest.approx(25.0),
+            "sum_3": pytest.approx(81.0),
+            "sum_4": pytest.approx(289.0),
+        },
+        "b": {
+            "count": pytest.approx(1.0),
+            "sum_1": pytest.approx(3.0),
+            "sum_2": pytest.approx(9.0),
+            "sum_3": pytest.approx(27.0),
+            "sum_4": pytest.approx(81.0),
+        },
+    }
+    assert aggregate_moments == {
+        "count": pytest.approx(5.0),
+        "sum_1": pytest.approx(12.0),
+        "sum_2": pytest.approx(34.0),
+        "sum_3": pytest.approx(108.0),
+        "sum_4": pytest.approx(370.0),
+    }
+
+
+def test_reduce_raw_moments_by_param_rejects_mismatched_names_and_tensors():
+    registry = PerParameterStatRegistry(TwoParamModel())
+
+    with pytest.raises(ValueError, match="names but"):
+        reduce_raw_moments_by_param(
+            registry,
+            [NamedTensorBucket(names=["a"], tensors=[torch.tensor([1.0]), torch.tensor([2.0])])],
+        )
+
+
+def test_registry_cache_is_per_model_identity():
+    first_model = TwoParamModel()
+    second_model = TwoParamModel()
+
+    first_registry = get_or_create_per_parameter_stat_registry(first_model)
+    assert get_or_create_per_parameter_stat_registry(first_model) is first_registry
+    assert get_or_create_per_parameter_stat_registry(second_model) is not first_registry
+
+
+def test_registry_cache_includes_expert_group_identity():
+    model = TwoParamModel()
+    expert_group = _FakeExpertGroup(rank=0, size=2)
+
+    first_registry = get_or_create_per_parameter_stat_registry(model)
+    expert_registry = get_or_create_per_parameter_stat_registry(
+        model, expert_model_parallel_group=expert_group
+    )
+
+    assert expert_registry is not first_registry
+    assert (
+        get_or_create_per_parameter_stat_registry(model, expert_model_parallel_group=expert_group)
+        is expert_registry
+    )
+
+
+def test_registry_uses_explicit_expert_group_for_expert_names(monkeypatch):
+    expert_group = _FakeExpertGroup(rank=1, size=2)
+
+    monkeypatch.setattr(pps, "get_pg_rank", lambda group: 1 if group is expert_group else 0)
+    monkeypatch.setattr(pps, "get_pg_size", lambda group: 2 if group is expert_group else 1)
+
+    registry = PerParameterStatRegistry(
+        SequentialExpertModel(), expert_model_parallel_group=expert_group
+    )
+
+    assert isinstance(registry.parameter_names, CanonicalParameterNameMap)
+    assert set(registry.param_to_name.values()) == {
+        "mlp.experts.local_experts.2.weight",
+        "mlp.experts.local_experts.3.weight",
+    }
+
+
+def test_registry_uses_global_canonical_parameter_name_index(monkeypatch):
+    expected_index = CanonicalParameterNameIndex(["remote", "b", "a"])
+    monkeypatch.setattr(CanonicalParameterNameMap, "all_gather_index", lambda self: expected_index)
+
+    registry = PerParameterStatRegistry(TwoParamModel())
+
+    assert registry.parameter_name_index is expected_index
+    assert registry.name_to_index is expected_index
+    assert registry.index_to_name is expected_index.names
+
+
+def test_local_raw_moments_multi_tensor_path_preserves_order(monkeypatch):
+    calls = []
+
+    def fake_multi_tensor_applier(op, noop_flag_buffer, tensor_lists):
+        calls.append([tensor.dtype for tensor in tensor_lists[0]])
+        return op(0, noop_flag_buffer, tensor_lists)
+
+    def fake_multi_tensor_raw_moments(_, __, tensor_lists):
+        return torch.stack([pps._torch_raw_moment_row(tensor) for tensor in tensor_lists[0]])
+
+    monkeypatch.setattr(pps, "multi_tensor_applier", fake_multi_tensor_applier)
+    monkeypatch.setattr(pps, "multi_tensor_raw_moments", fake_multi_tensor_raw_moments)
+    monkeypatch.setattr(pps, "_can_use_multi_tensor_raw_moments", lambda tensors, device: True)
+
+    tensors = [
+        torch.tensor([1.0, 2.0], dtype=torch.float32),
+        torch.tensor([3.0], dtype=torch.bfloat16),
+        torch.tensor([4.0, 5.0], dtype=torch.float32),
+    ]
+    rows = pps._local_raw_moments(tensors, torch.device("cpu"))
+    expected = torch.stack([pps._torch_raw_moment_row(tensor) for tensor in tensors])
+
+    torch.testing.assert_close(rows, expected)
+    assert calls == [[torch.float32, torch.float32], [torch.bfloat16]]
+
+
+def test_local_raw_moments_multi_tensor_path_splits_oversized_tensors(monkeypatch):
+    calls = []
+
+    def fake_multi_tensor_applier(op, noop_flag_buffer, tensor_lists):
+        calls.append([tensor.numel() for tensor in tensor_lists[0]])
+        return op(0, noop_flag_buffer, tensor_lists)
+
+    def fake_multi_tensor_raw_moments(_, __, tensor_lists):
+        return torch.stack([pps._torch_raw_moment_row(tensor) for tensor in tensor_lists[0]])
+
+    monkeypatch.setattr(pps, "multi_tensor_applier", fake_multi_tensor_applier)
+    monkeypatch.setattr(pps, "multi_tensor_raw_moments", fake_multi_tensor_raw_moments)
+    monkeypatch.setattr(pps, "_can_use_multi_tensor_raw_moments", lambda tensors, device: True)
+    monkeypatch.setattr(pps, "_MAX_MULTI_TENSOR_RAW_MOMENTS_NUMEL", 4)
+
+    tensors = [torch.arange(1.0, 11.0), torch.tensor([11.0, 12.0, 13.0])]
+    rows = pps._local_raw_moments(tensors, torch.device("cpu"))
+    expected = torch.stack([pps._torch_raw_moment_row(tensor) for tensor in tensors])
+
+    torch.testing.assert_close(rows, expected)
+    assert calls == [[4, 4, 2, 3]]
+
+
+def test_multi_tensor_raw_moments_env_guard_disables_fast_path(monkeypatch):
+    tensor = _FakeCudaTensor()
+    device = torch.device("cuda:0")
+
+    monkeypatch.setattr(pps, "multi_tensor_applier", object())
+    monkeypatch.setattr(pps, "multi_tensor_raw_moments", object())
+    monkeypatch.delenv("MEGATRON_DISABLE_MULTI_TENSOR_RAW_MOMENTS", raising=False)
+    assert pps._can_use_multi_tensor_raw_moments([tensor], device)
+
+    monkeypatch.setenv("MEGATRON_DISABLE_MULTI_TENSOR_RAW_MOMENTS", "1")
+    assert not pps._can_use_multi_tensor_raw_moments([tensor], device)
diff --git a/tests/unit_tests/test_raw_moment_logging.py b/tests/unit_tests/test_raw_moment_logging.py
new file mode 100644
index 00000000000..72d050be78b
--- /dev/null
+++ b/tests/unit_tests/test_raw_moment_logging.py
@@ -0,0 +1,386 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from megatron.core.transformer.layer_boundary_observer import (
+    observe_transformer_layer_boundaries,
+    observe_transformer_layer_input,
+    observe_transformer_layer_output,
+)
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.training.arguments import _validate_raw_moment_logging_args
+from megatron.training.raw_moment_logging import RawMomentLogger
+
+_STATS_INTERVAL_FLAGS = (
+    'log_param_stats_interval',
+    'log_wgrad_stats_interval',
+    'log_activation_stats_interval',
+    'log_dgrad_stats_interval',
+    'log_residual_stats_interval',
+    'log_residual_grad_stats_interval',
+)
+
+
+class _LinearBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(2, 2, bias=False)
+        with torch.no_grad():
+            self.linear.weight.copy_(torch.eye(2))
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class _ToyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = nn.Module()
+        self.decoder.layers = nn.ModuleList([_LinearBlock()])
+
+    def forward(self, x):
+        return self.decoder.layers[0](x)
+
+
+class _EmbeddingModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = nn.Embedding(4, 2)
+
+    def forward(self, x):
+        return self.embedding(x)
+
+
+class _OutputLayerModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.output_layer = nn.Linear(2, 3, bias=False)
+
+    def forward(self, x):
+        return self.output_layer(x)
+
+
+class _ResidualLayer(nn.Module):
+    def __init__(self, layer_number):
+        super().__init__()
+        self.layer_number = layer_number
+        self.config = SimpleNamespace(sequence_parallel=False)
+
+
+class _ResidualModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        decoder = TransformerBlock.__new__(TransformerBlock)
+        nn.Module.__init__(decoder)
+        decoder.layers = nn.ModuleList([_ResidualLayer(1), _ResidualLayer(2)])
+        self.decoder = decoder
+
+
+def _values_dict(values):
+    return {name: moments for name, moments in values}
+
+
+def _raw_moment_logging_args():
+    return SimpleNamespace(
+        use_megatron_fsdp=False,
+        use_torch_fsdp2=False,
+        log_param_stats_interval=0,
+        log_wgrad_stats_interval=0,
+        log_activation_stats_interval=0,
+        log_dgrad_stats_interval=0,
+        log_residual_stats_interval=0,
+        log_residual_grad_stats_interval=0,
+        overlap_moe_expert_parallel_comm=False,
+    )
+
+
+@pytest.mark.parametrize('logging_flag', _STATS_INTERVAL_FLAGS)
+@pytest.mark.parametrize('fsdp_flag', ('use_megatron_fsdp', 'use_torch_fsdp2'))
+def test_raw_moment_logging_rejects_fsdp(logging_flag, fsdp_flag):
+    args = _raw_moment_logging_args()
+    setattr(args, logging_flag, 1)
+    setattr(args, fsdp_flag, True)
+
+    with pytest.raises(ValueError, match='Raw-moment statistics logging is not supported'):
+        _validate_raw_moment_logging_args(args)
+
+
+@pytest.mark.parametrize('logging_flag', _STATS_INTERVAL_FLAGS)
+@pytest.mark.parametrize('disabled_interval', (0, -1))
+@pytest.mark.parametrize('fsdp_flag', ('use_megatron_fsdp', 'use_torch_fsdp2'))
+def test_disabled_raw_moment_logging_allows_fsdp(logging_flag, disabled_interval, fsdp_flag):
+    args = _raw_moment_logging_args()
+    setattr(args, logging_flag, disabled_interval)
+    setattr(args, fsdp_flag, True)
+
+    _validate_raw_moment_logging_args(args)
+
+
+@pytest.mark.parametrize(
+    "logging_flag", ("log_residual_stats_interval", "log_residual_grad_stats_interval")
+)
+def test_residual_raw_moment_logging_rejects_fine_grained_ep_overlap(logging_flag):
+    args = _raw_moment_logging_args()
+    setattr(args, logging_flag, 1)
+    args.overlap_moe_expert_parallel_comm = True
+
+    with pytest.raises(ValueError, match='--overlap-moe-expert-parallel-comm'):
+        _validate_raw_moment_logging_args(args)
+
+
+def test_residual_raw_moments_capture_distinct_layer_boundaries():
+    model = _ResidualModel()
+    logger = RawMomentLogger()
+    logger.prepare_residual_logging(model)
+    first_layer, second_layer = model.decoder.layers
+
+    with observe_transformer_layer_boundaries(logger.record_residual_boundary):
+        observe_transformer_layer_input(model.decoder, first_layer, torch.tensor([1.0, 2.0]))
+        observe_transformer_layer_output(model.decoder, first_layer, torch.tensor([3.0, 4.0]))
+        observe_transformer_layer_input(model.decoder, second_layer, torch.tensor([30.0, 40.0]))
+        observe_transformer_layer_output(model.decoder, second_layer, torch.tensor([5.0, 6.0]))
+        observe_transformer_layer_output(model.decoder, first_layer, torch.tensor([7.0, 8.0]))
+
+    # The observer is scoped and does not retain events after the context exits.
+    observe_transformer_layer_output(model.decoder, second_layer, torch.tensor([50.0, 60.0]))
+    logger.finalize_residual_raw_moments_by_layer()
+
+    values = _values_dict(logger.consume_residual_raw_moments_by_layer())
+    assert set(values) == {"decoder/input0", "decoder.layers.0/output0", "decoder.layers.1/output0"}
+    assert values["decoder/input0"] == {
+        "count": 2.0,
+        "sum_1": 3.0,
+        "sum_2": 5.0,
+        "sum_3": 9.0,
+        "sum_4": 17.0,
+    }
+    assert values["decoder.layers.0/output0"] == {
+        "count": 4.0,
+        "sum_1": 22.0,
+        "sum_2": 138.0,
+        "sum_3": 946.0,
+        "sum_4": 6834.0,
+    }
+    assert values["decoder.layers.1/output0"]["sum_1"] == 11.0
+
+
+def test_residual_raw_moments_skip_no_grad_forward():
+    model = _ResidualModel()
+    logger = RawMomentLogger()
+    logger.prepare_residual_logging(model)
+    first_layer = model.decoder.layers[0]
+
+    with observe_transformer_layer_boundaries(logger.record_residual_boundary):
+        with torch.no_grad():
+            observe_transformer_layer_output(model.decoder, first_layer, torch.tensor([10.0, 20.0]))
+        observe_transformer_layer_output(model.decoder, first_layer, torch.tensor([1.0, 2.0]))
+
+    logger.finalize_residual_raw_moments_by_layer()
+    values = _values_dict(logger.consume_residual_raw_moments_by_layer())
+    assert values["decoder.layers.0/output0"]["count"] == 2.0
+    assert values["decoder.layers.0/output0"]["sum_1"] == 3.0
+
+
+def test_residual_dgrad_raw_moments_capture_boundary_gradients():
+    model = _ResidualModel()
+    logger = RawMomentLogger()
+    logger.prepare_residual_logging(model, capture_residuals=False, capture_dgrads=True)
+    first_layer, second_layer = model.decoder.layers
+
+    residual_input = torch.tensor([1.0, 2.0], requires_grad=True)
+    first_output = residual_input * torch.tensor([2.0, 3.0])
+    second_output = first_output * torch.tensor([5.0, 7.0])
+    with observe_transformer_layer_boundaries(logger.record_residual_boundary):
+        observe_transformer_layer_input(model.decoder, first_layer, residual_input)
+        observe_transformer_layer_output(model.decoder, first_layer, first_output)
+        observe_transformer_layer_output(model.decoder, second_layer, second_output)
+        second_output.sum().backward()
+
+    logger.finalize_residual_dgrad_raw_moments_by_layer()
+    values = _values_dict(logger.consume_residual_dgrad_raw_moments_by_layer())
+
+    assert values["decoder/input0"] == {
+        "count": 2.0,
+        "sum_1": 31.0,
+        "sum_2": 541.0,
+        "sum_3": 10261.0,
+        "sum_4": 204481.0,
+    }
+    assert values["decoder.layers.0/output0"]["sum_1"] == 12.0
+    assert values["decoder.layers.1/output0"]["sum_1"] == 2.0
+    assert residual_input.grad.tolist() == [10.0, 21.0]
+    assert not logger._residual_dgrad_hooks
+
+
+def test_residual_dgrad_raw_moments_capture_checkpoint_recomputation_once():
+    model = _ResidualModel()
+    logger = RawMomentLogger()
+    logger.prepare_residual_logging(model, capture_residuals=False, capture_dgrads=True)
+    first_layer = model.decoder.layers[0]
+
+    def checkpointed_layer(hidden_states):
+        observe_transformer_layer_input(model.decoder, first_layer, hidden_states)
+        output = hidden_states * 2
+        observe_transformer_layer_output(model.decoder, first_layer, output)
+        return output
+
+    hidden_states = torch.tensor([1.0, 2.0], requires_grad=True)
+    with observe_transformer_layer_boundaries(logger.record_residual_boundary):
+        checkpoint(checkpointed_layer, hidden_states, use_reentrant=True).sum().backward()
+
+    logger.finalize_residual_dgrad_raw_moments_by_layer()
+    values = _values_dict(logger.consume_residual_dgrad_raw_moments_by_layer())
+
+    assert values["decoder/input0"]["count"] == 2.0
+    assert values["decoder/input0"]["sum_1"] == 4.0
+    assert values["decoder.layers.0/output0"]["count"] == 2.0
+    assert values["decoder.layers.0/output0"]["sum_1"] == 2.0
+
+
+def test_residual_dgrad_raw_moments_support_reused_autograd_output_tensor():
+    class ReuseOutput(torch.autograd.Function):
+        output = torch.empty(2)
+
+        @staticmethod
+        def forward(ctx, tensor):
+            ReuseOutput.output.copy_(tensor)
+            return ReuseOutput.output
+
+        @staticmethod
+        def backward(ctx, grad):
+            return grad
+
+    model = _ResidualModel()
+    logger = RawMomentLogger()
+    first_layer = model.decoder.layers[0]
+    output_tensor = None
+
+    for input_values in ([1.0, 2.0], [3.0, 4.0]):
+        logger.prepare_residual_logging(model, capture_residuals=False, capture_dgrads=True)
+        tensor = ReuseOutput.apply(torch.tensor(input_values, requires_grad=True))
+        if output_tensor is None:
+            output_tensor = tensor
+        else:
+            assert tensor is output_tensor
+
+        with observe_transformer_layer_boundaries(logger.record_residual_boundary):
+            observe_transformer_layer_output(model.decoder, first_layer, tensor)
+            tensor.sum().backward()
+
+        logger.finalize_residual_dgrad_raw_moments_by_layer()
+        values = _values_dict(logger.consume_residual_dgrad_raw_moments_by_layer())
+        assert values["decoder.layers.0/output0"]["count"] == 2.0
+        assert values["decoder.layers.0/output0"]["sum_1"] == 2.0
+
+
+def test_activation_raw_moments_accumulate_by_module_site():
+    model = [_ToyModel()]
+    logger = RawMomentLogger()
+    logger.register_activation_hooks(model)
+
+    model[0](torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+    model[0](torch.tensor([[5.0, 6.0]]))
+
+    logger.finalize_activation_raw_moments_by_layer()
+    logger.remove_activation_hooks()
+
+    values = _values_dict(logger.consume_activation_raw_moments_by_layer())
+    moments = values["decoder.layers.0.linear/input0"]
+    assert moments == {"count": 6.0, "sum_1": 21.0, "sum_2": 91.0, "sum_3": 441.0, "sum_4": 2275.0}
+    assert values["decoder.layers.0.linear/output0"] == moments
+
+
+def test_activation_raw_moments_skip_no_grad_forward():
+    model = [_ToyModel()]
+    logger = RawMomentLogger()
+    logger.register_activation_hooks(model)
+
+    with torch.no_grad():
+        model[0](torch.tensor([[10.0, 20.0]]))
+    model[0](torch.tensor([[1.0, 2.0]]))
+
+    logger.finalize_activation_raw_moments_by_layer()
+    logger.remove_activation_hooks()
+
+    values = _values_dict(logger.consume_activation_raw_moments_by_layer())
+    assert values["decoder.layers.0.linear/input0"]["count"] == 2.0
+    assert values["decoder.layers.0.linear/input0"]["sum_1"] == 3.0
+
+
+def test_activation_raw_moments_skip_integer_inputs():
+    model = [_EmbeddingModel()]
+    logger = RawMomentLogger()
+    logger.register_activation_hooks(model)
+
+    model[0](torch.tensor([0, 1, 2], dtype=torch.long))
+
+    logger.finalize_activation_raw_moments_by_layer()
+    logger.remove_activation_hooks()
+
+    values = _values_dict(logger.consume_activation_raw_moments_by_layer())
+    assert "embedding/input0" not in values
+    assert "embedding/output0" in values
+
+
+def test_raw_moments_skip_output_layer_logits_site():
+    model = [_OutputLayerModel()]
+    logger = RawMomentLogger()
+    logger.register_activation_hooks(model)
+
+    model[0](torch.tensor([[1.0, 2.0]], requires_grad=True))
+
+    logger.finalize_activation_raw_moments_by_layer()
+    logger.remove_activation_hooks()
+
+    values = _values_dict(logger.consume_activation_raw_moments_by_layer())
+    assert "output_layer/input0" in values
+    assert "output_layer/output0" not in values
+
+
+def test_dgrad_raw_moments_skip_output_layer_logits_site():
+    model = [_OutputLayerModel()]
+    logger = RawMomentLogger()
+    logger.register_dgrad_hooks(model)
+
+    x = torch.tensor([[1.0, 2.0]], requires_grad=True)
+    model[0](x).sum().backward()
+
+    logger.finalize_dgrad_raw_moments_by_layer()
+    logger.remove_dgrad_hooks()
+
+    values = _values_dict(logger.consume_dgrad_raw_moments_by_layer())
+    assert "output_layer/input0" in values
+    assert "output_layer/output0" not in values
+
+
+def test_dgrad_raw_moments_accumulate_by_module_site():
+    model = [_ToyModel()]
+    logger = RawMomentLogger()
+    logger.register_dgrad_hooks(model)
+
+    x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
+    model[0](x).sum().backward()
+
+    logger.finalize_dgrad_raw_moments_by_layer()
+    logger.remove_dgrad_hooks()
+
+    values = _values_dict(logger.consume_dgrad_raw_moments_by_layer())
+    assert values["decoder.layers.0.linear/output0"] == {
+        "count": 4.0,
+        "sum_1": 4.0,
+        "sum_2": 4.0,
+        "sum_3": 4.0,
+        "sum_4": 4.0,
+    }
+    assert values["decoder.layers.0.linear/input0"] == {
+        "count": 4.0,
+        "sum_1": 4.0,
+        "sum_2": 4.0,
+        "sum_3": 4.0,
+        "sum_4": 4.0,
+    }
diff --git a/tests/unit_tests/test_statistics_logging.py b/tests/unit_tests/test_statistics_logging.py
new file mode 100644
index 00000000000..ff89687d144
--- /dev/null
+++ b/tests/unit_tests/test_statistics_logging.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import json
+
+from megatron.training.statistics_logging import save_raw_moments_by_name
+
+
+def _read_records(filepath):
+    return [json.loads(line) for line in filepath.read_text().strip().split("\n")]
+
+
+def test_save_raw_moments_by_name_creates_jsonl(tmp_path):
+    save_raw_moments_by_name(
+        str(tmp_path),
+        "custom_raw_moments",
+        iteration=100,
+        consumed_train_samples=8192,
+        raw_moments_by_name=[
+            (
+                "decoder.layers.0.self_attention.linear_qkv.weight",
+                {"count": 2, "sum_1": 1.5, "sum_2": 2.5, "sum_3": 3.5, "sum_4": 4.5},
+            ),
+            (
+                "decoder.layers.0.mlp.linear_fc1.weight",
+                {"count": 3, "sum_1": 2.25, "sum_2": 5.25, "sum_3": 8.25, "sum_4": 11.25},
+            ),
+        ],
+        rank=7,
+    )
+
+    filepath = tmp_path / "training_stats" / "custom_raw_moments" / "rank7.jsonl"
+    assert _read_records(filepath) == [
+        {
+            "iter": 100,
+            "consumed_train_samples": 8192,
+            "stat": "custom_raw_moments",
+            "values": {
+                "decoder.layers.0.self_attention.linear_qkv.weight": {
+                    "count": 2.0,
+                    "sum_1": 1.5,
+                    "sum_2": 2.5,
+                    "sum_3": 3.5,
+                    "sum_4": 4.5,
+                },
+                "decoder.layers.0.mlp.linear_fc1.weight": {
+                    "count": 3.0,
+                    "sum_1": 2.25,
+                    "sum_2": 5.25,
+                    "sum_3": 8.25,
+                    "sum_4": 11.25,
+                },
+            },
+        }
+    ]
+
+
+def test_save_raw_moments_by_name_appends_across_calls(tmp_path):
+    first_moments = {"count": 1, "sum_1": 1, "sum_2": 1, "sum_3": 1, "sum_4": 1}
+    second_moments = {"count": 1, "sum_1": 2, "sum_2": 4, "sum_3": 8, "sum_4": 16}
+    save_raw_moments_by_name(
+        str(tmp_path),
+        "param_raw_moments_by_param",
+        iteration=100,
+        consumed_train_samples=8192,
+        raw_moments_by_name=[("layer.weight", first_moments)],
+        rank=0,
+    )
+    save_raw_moments_by_name(
+        str(tmp_path),
+        "param_raw_moments_by_param",
+        iteration=200,
+        consumed_train_samples=16384,
+        raw_moments_by_name=[("layer.weight", second_moments)],
+        rank=0,
+    )
+
+    filepath = tmp_path / "training_stats" / "param_raw_moments_by_param" / "rank0.jsonl"
+    records = _read_records(filepath)
+    assert [record["iter"] for record in records] == [100, 200]
+    assert records[0]["values"] == {
+        "layer.weight": {"count": 1.0, "sum_1": 1.0, "sum_2": 1.0, "sum_3": 1.0, "sum_4": 1.0}
+    }
+    assert records[1]["values"] == {
+        "layer.weight": {"count": 1.0, "sum_1": 2.0, "sum_2": 4.0, "sum_3": 8.0, "sum_4": 16.0}
+    }
+
+
+def test_save_raw_moments_by_name_skips_empty_values(tmp_path):
+    save_raw_moments_by_name(
+        str(tmp_path),
+        "param_raw_moments_by_param",
+        iteration=100,
+        consumed_train_samples=8192,
+        raw_moments_by_name=[],
+        rank=0,
+    )
+
+    filepath = tmp_path / "training_stats" / "param_raw_moments_by_param" / "rank0.jsonl"
+    assert not filepath.exists()
diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py
index 52edbf9a264..4fb79248025 100644
--- a/tests/unit_tests/transformer/test_cuda_graphs.py
+++ b/tests/unit_tests/transformer/test_cuda_graphs.py
@@ -9,6 +9,7 @@
 from transformer_engine.pytorch.fp8 import check_fp8_support
 
 from megatron.core.enums import ModelType
+from megatron.core.models.common.utils import _BackwardDWWrapper
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_decoder_block_spec,
     get_gpt_layer_with_transformer_engine_spec,
@@ -34,15 +35,17 @@
     CudaGraphManager,
     TECudaGraphHelper,
     _CudagraphGlobalRecord,
+    is_cuda_graph_replay_suspended,
+    suspend_cuda_graph_replay,
 )
 from megatron.core.transformer.enums import CudaGraphModule, CudaGraphScope, InferenceCudaGraphScope
 from megatron.core.transformer.mlp import MLPSubmodules
-from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule
 from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer
 from megatron.core.transformer.spec_utils import ModuleSpec, get_submodules
 from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.transformer.transformer_layer import TransformerLayer
+from megatron.core.transformer.transformer_layer import MoETransformerLayer, TransformerLayer
 from megatron.core.utils import is_te_min_version
 from megatron.training import arguments as training_arguments
 from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args
@@ -1463,6 +1466,13 @@ def my_op(self, x):
         return x @ self.weight
 
 
+class _SimpleGraphableModule(GraphableMegatronModule):
+    """Minimal graphable module for testing eager fallback from TE replay."""
+
+    def forward(self, x):
+        return x + 1
+
+
 def _make_simple_module(config):
     return _SimpleModule(config).cuda().eval()
 
@@ -1471,6 +1481,91 @@ def _make_simple_non_module(config):
     return _SimpleNonModule(config)
 
 
+class TestCudaGraphReplaySuspension:
+    """Tests for temporarily dispatching graphable modules through eager execution."""
+
+    def _make_te_module(self):
+        config = TransformerConfig(
+            num_layers=1,
+            hidden_size=8,
+            num_attention_heads=1,
+            use_cpu_initialization=True,
+            cuda_graph_impl="transformer_engine",
+        )
+        module = _SimpleGraphableModule(config)
+        module.cuda_graphs = [object()]
+        return module
+
+    def test_te_replay_suspension_restores_eager_module_hooks(self):
+        module = self._make_te_module()
+        module._te_cuda_graph_replay = lambda x: x + 2
+        hook_calls = []
+        module.register_forward_hook(lambda *_: hook_calls.append(True))
+        x = torch.tensor(1.0)
+
+        assert module(x).item() == 3.0
+        assert hook_calls == []
+        with suspend_cuda_graph_replay():
+            assert is_cuda_graph_replay_suspended()
+            assert module(x).item() == 2.0
+        assert not is_cuda_graph_replay_suspended()
+        assert hook_calls == [True]
+
+    def test_suspension_precedes_overridden_graph_predicate(self):
+        module = self._make_te_module()
+        graph_calls = []
+        module._should_call_local_cudagraph = lambda *args, **kwargs: True
+        module.cudagraph_manager = lambda _module, args, kwargs: graph_calls.append(True) or args[0]
+        x = torch.tensor(1.0)
+
+        assert module(x) is x
+        assert graph_calls == [True]
+        with suspend_cuda_graph_replay():
+            assert module(x).item() == 2.0
+        assert graph_calls == [True]
+
+    def test_suspension_selects_eager_delayed_wgrad(self):
+        calls = []
+
+        def make_wrapper():
+            wrapper = object.__new__(_BackwardDWWrapper)
+            wrapper.layer = self._make_te_module()
+            wrapper.graphed_backward_dw_callable = lambda: calls.append("graph")
+            wrapper.attn_dw_callable = lambda: calls.append("attn")
+            wrapper.shared_expert_dw_callable = lambda: calls.append("shared_expert")
+            wrapper.cuda_graph_modules = [CudaGraphModule.attn, CudaGraphModule.moe_router]
+            return wrapper
+
+        make_wrapper().backward_dw()
+        assert calls == ["graph"]
+
+        calls.clear()
+        with suspend_cuda_graph_replay():
+            make_wrapper().backward_dw()
+        assert calls == ["shared_expert", "attn"]
+
+    def test_suspension_bypasses_partial_moe_graphs(self, monkeypatch):
+        layer = object.__new__(MoETransformerLayer)
+        object.__setattr__(layer, "use_partial_cudagraphs", True)
+        object.__setattr__(layer, "mlp", type("MLP", (), {"fwd_execution_map": "expert_compute"})())
+        eager_result = object()
+        eager_calls = []
+
+        def eager_forward(
+            _layer, hidden_states, inference_context=None, padding_mask=None, packed_seq_params=None
+        ):
+            eager_calls.append((hidden_states, padding_mask, packed_seq_params))
+            return eager_result
+
+        monkeypatch.setattr(TransformerLayer, "_forward_mlp", eager_forward)
+        with suspend_cuda_graph_replay():
+            result = layer._forward_mlp("hidden", padding_mask="mask", packed_seq_params="packed")
+
+        assert result is eager_result
+        assert eager_calls == [("hidden", "mask", "packed")]
+        assert layer.mlp.fwd_execution_map == ["route", "expert_compute", "postprocess"]
+
+
 class TestInlineCaptureManager:
     """Tests for CudaGraphManager with inline_capture, function_name, eager, and cache_key."""
 
@@ -1551,6 +1646,26 @@ def test_eager_bypass(self):
         _ = module.my_op(x, eager=True)
         assert len(mgr.cudagraph_runners) == 0, "eager=True should not create runners"
 
+    @torch.inference_mode()
+    def test_suspended_replay_bypasses_method_graph(self):
+        """The suspension context must bypass local method graph capture."""
+        config = self._make_config()
+        module = _SimpleModule(config).cuda().eval()
+
+        mgr = CudaGraphManager(
+            config,
+            base_module=module,
+            function_name="my_op",
+            inline_capture=True,
+            num_warmup_steps=0,
+            need_backward=False,
+        )
+
+        x = torch.randn(4, config.hidden_size, device="cuda")
+        with suspend_cuda_graph_replay():
+            _ = module.my_op(x)
+        assert len(mgr.cudagraph_runners) == 0
+
     @torch.inference_mode()
     def test_cache_key_routing(self):
         """Different cache_keys must create separate runners."""