Add statistics logging for params and activations

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
         ):

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)(

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:

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
 

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-

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 = []