Add statistics logging for params and activations

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

megatron/core/optimizer/optimizer.py

@@ -46,14 +46,22 @@
     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 ..utils import get_data_parallel_group_if_dtensor, log_single_rank, to_local_if_dtensor
 from .clip_grads import clip_grad_by_total_norm_fp32, count_zeros_fp32, get_grad_norm_fp32
 from .grad_scaler import MegatronGradScaler
 from .optimizer_config import OptimizerConfig
 
 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 _filter_grads_for_norm(
         self,
         params: List[torch.nn.Parameter],
@@ -188,30 +231,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
 
@@ -257,6 +278,117 @@ 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 = []
+        data_parallel_group = None
+        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
+            data_parallel_group = get_data_parallel_group_if_dtensor(grad, data_parallel_group)
+            names.append(name)
+            grads.append(to_local_if_dtensor(grad).detach())
+
+        reduce_groups = ((data_parallel_group,) if data_parallel_group is not None else ()) + (
+            self.get_grad_stats_parallel_group(),
+        )
+        return [NamedTensorBucket(names, grads, reduce_groups)]
+
+    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).
 
@@ -324,6 +456,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.
@@ -744,6 +877,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.
@@ -754,6 +888,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()
 
@@ -1114,6 +1251,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.
@@ -1124,6 +1262,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()
 
@@ -1234,6 +1375,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:
@@ -1525,6 +1669,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:
@@ -1628,6 +1780,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()
@@ -1640,6 +1793,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: