From 34f71bc9cdb3f1ca741134091af6c84913d4589d Mon Sep 17 00:00:00 2001 From: Philip Monk Date: Wed, 24 Jun 2026 18:03:31 -0700 Subject: [PATCH 1/2] Add raw moment statistics logging Signed-off-by: Philip Monk --- megatron/core/optimizer/__init__.py | 2 + megatron/core/optimizer/distrib_optimizer.py | 20 + .../core/optimizer/layer_wise_optimizer.py | 16 + megatron/core/optimizer/optimizer.py | 194 +++++++- megatron/core/per_parameter_stats.py | 358 ++++++++++++++ megatron/training/arguments.py | 2 + megatron/training/config/training_config.py | 24 + megatron/training/raw_moment_logging.py | 461 ++++++++++++++++++ megatron/training/statistics_logging.py | 162 ++++++ megatron/training/training.py | 168 +++++++ megatron/training/utils/__init__.py | 1 + megatron/training/utils/common_utils.py | 201 ++++++-- tests/unit_tests/test_per_parameter_stats.py | 146 ++++++ tests/unit_tests/test_raw_moment_logging.py | 161 ++++++ tests/unit_tests/test_statistics_logging.py | 224 +++++++++ 15 files changed, 2081 insertions(+), 59 deletions(-) create mode 100644 megatron/core/per_parameter_stats.py create mode 100644 megatron/training/raw_moment_logging.py create mode 100644 megatron/training/statistics_logging.py create mode 100644 tests/unit_tests/test_per_parameter_stats.py create mode 100644 tests/unit_tests/test_raw_moment_logging.py create mode 100644 tests/unit_tests/test_statistics_logging.py 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 374b1aab096..e8bcf2dbb0e 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 e03992e0657..36a7cc813ce 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/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,107 @@ 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 + ) -> 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) + 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) -> 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) + 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 +446,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 +867,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 +878,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 +1241,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 +1252,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 +1365,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 +1659,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 +1770,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 +1783,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/per_parameter_stats.py b/megatron/core/per_parameter_stats.py new file mode 100644 index 00000000000..3e084e87468 --- /dev/null +++ b/megatron/core/per_parameter_stats.py @@ -0,0 +1,358 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared helpers for high-cardinality per-parameter statistics.""" + +from __future__ import annotations + +import os +import re +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 import parallel_state +from megatron.core.utils import unwrap_model + +_LAYER_NAME_PATTERN = re.compile(r"layers\.(\d+)") +_GROUPED_EXPERT_PATTERN = re.compile(r"^(.*\.mlp\.experts\.linear_fc\d\.weight)(\d+)(.*)$") +_SEQUENTIAL_EXPERT_PATTERN = re.compile(r"^(.*\.mlp\.experts\.local_experts\.)(\d+)(\..*)$") +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): + self.model_chunks = unwrap_model(_normalize_model_chunks(model_chunks)) + self.cache_key = tuple(id(model_chunk) for model_chunk in self.model_chunks) + self.param_to_name = self._build_local_param_to_name() + self.name_to_index = self._build_name_to_index() + self.index_to_name = sorted(self.name_to_index, key=self.name_to_index.get) + + 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 _build_local_param_to_name(self) -> dict[torch.nn.Parameter, str]: + param_to_name = {} + num_experts = _get_num_moe_experts(self.model_chunks) + for model_chunk in self.model_chunks: + for local_name, param in model_chunk.named_parameters(): + param_to_name[param] = _canonical_param_name( + model_chunk, local_name, param, num_experts + ) + return param_to_name + + def _build_name_to_index(self) -> dict[str, int]: + local_names = list(self.param_to_name.values()) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + gathered_names = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(gathered_names, local_names) + all_names = set() + for names in gathered_names: + all_names.update(names) + else: + all_names = set(local_names) + return {name: idx for idx, name in enumerate(sorted(all_names))} + + +def get_or_create_per_parameter_stat_registry( + model_chunks: Iterable[torch.nn.Module] | torch.nn.Module, +) -> 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.") + + cache_key = tuple(id(model_chunk) for model_chunk in unwrapped_model_chunks) + 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) + 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): + local_tensor = getattr(tensor, "_local_tensor", None) + if local_tensor is None: + local_tensor = tensor + flat_tensor = local_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 _canonical_param_name( + model_chunk: torch.nn.Module, + local_name: str, + param: torch.nn.Parameter, + num_experts: int | None, +) -> str: + name = _global_layer_param_name(model_chunk, local_name, param) + return _global_expert_param_name(name, num_experts) + + +def _global_layer_param_name( + model_chunk: torch.nn.Module, local_name: str, param: torch.nn.Parameter +) -> str: + if "mtp" in local_name or _LAYER_NAME_PATTERN.search(local_name) is None: + return local_name + + from megatron.core.transformer.transformer_layer import TransformerLayer + + for module in model_chunk.modules(): + if not isinstance(module, TransformerLayer): + continue + for module_param in module.parameters(): + if module_param is param: + return _LAYER_NAME_PATTERN.sub(f"layers.{module.layer_number - 1}", local_name) + return local_name + + +def _global_expert_param_name(local_name: str, num_experts: int | None) -> str: + if not num_experts: + return local_name + + expert_offset = _get_local_expert_offset(num_experts) + if expert_offset == 0: + return local_name + + grouped_match = _GROUPED_EXPERT_PATTERN.match(local_name) + if grouped_match is not None: + prefix, local_expert_index, suffix = grouped_match.groups() + return f"{prefix}{int(local_expert_index) + expert_offset}{suffix}" + + sequential_match = _SEQUENTIAL_EXPERT_PATTERN.match(local_name) + if sequential_match is not None: + prefix, local_expert_index, suffix = sequential_match.groups() + return f"{prefix}{int(local_expert_index) + expert_offset}{suffix}" + + return local_name + + +def _get_local_expert_offset(num_experts: int) -> int: + expert_group = parallel_state.get_expert_model_parallel_group(check_initialized=False) + if expert_group is None: + return 0 + expert_parallel_size = parallel_state.get_expert_model_parallel_world_size() + if expert_parallel_size <= 1: + return 0 + local_experts = num_experts // expert_parallel_size + return parallel_state.get_expert_model_parallel_rank() * local_experts + + +def _get_num_moe_experts(model_chunks: Sequence[torch.nn.Module]) -> int | None: + if not model_chunks: + return None + config = getattr(model_chunks[0], "config", None) + return getattr(config, "num_moe_experts", None) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f305a5a7668..00bf3e8b400 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1277,6 +1277,8 @@ def validate_args(args, defaults={}): assert not args.overlap_param_gather if args.log_memory_interval is not None: assert args.log_memory_interval % args.log_interval == 0 + if args.activation_log_interval is not None: + assert args.activation_log_interval > 0, '--activation-log-interval must be greater than 0' # Mixed precision checks. if args.fp16_lm_cross_entropy: assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' diff --git a/megatron/training/config/training_config.py b/megatron/training/config/training_config.py index cd273275353..e070917ad51 100644 --- a/megatron/training/config/training_config.py +++ b/megatron/training/config/training_config.py @@ -251,6 +251,30 @@ class LoggerConfig: log_params_norm: bool = False """If set, calculate and log parameters norm.""" + log_param_raw_moments_by_param: bool = False + """If set, calculate and log count and raw sums of powers 1-4 for each parameter separately + (keyed by parameter name) to JSONL statistics files.""" + + log_grad_raw_moments_by_param: bool = False + """If set, calculate and log count and raw sums of powers 1-4 for each parameter's + pre-clipping gradient separately (keyed by parameter name) to JSONL statistics files.""" + + log_activation_raw_moments_by_layer: bool = False + """If set, calculate and log count and raw sums of powers 1-4 for activations keyed by + module site to JSONL statistics files.""" + + log_dgrad_raw_moments_by_layer: bool = False + """If set, calculate and log count and raw sums of powers 1-4 for backward data gradients + keyed by module site to JSONL statistics files.""" + + activation_log_interval: int | None = None + """Interval for activation and dgrad raw-moment statistics. If unset, uses + tensorboard_log_interval.""" + + statistics_log_dir: str | None = None + """Directory for high-cardinality JSONL statistics. If unset, statistics use + tensorboard_dir when available, then save 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..464a54546b5 --- /dev/null +++ b/megatron/training/raw_moment_logging.py @@ -0,0 +1,461 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Activation and dgrad raw-moment logging using module hooks.""" + +from __future__ import annotations + +import re +from collections import OrderedDict +from dataclasses import dataclass +from typing import Iterable + +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.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 + + +class RawMomentLogger: + """Collect activation and dgrad raw moments from module hooks.""" + + def __init__(self): + self._activation_hooks: list[torch.utils.hooks.RemovableHandle] = [] + self._dgrad_hooks: list[torch.utils.hooks.RemovableHandle] = [] + self._activation_sites: OrderedDict[str, _RawMomentSite] = OrderedDict() + self._dgrad_sites: OrderedDict[str, _RawMomentSite] = OrderedDict() + self._latest_activation_raw_moments_by_layer: list[tuple[str, dict[str, float]]] | None = ( + None + ) + self._latest_dgrad_raw_moments_by_layer: list[tuple[str, dict[str, float]]] | None = None + self._latest_dgrad_loss_scale: float | None = None + self._pending_dgrad_loss_scale: float | None = None + + 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._activation_sites.clear() + for module, input_site, output_site in _iter_hook_modules(model): + input_key = input_site.name + self._activation_sites[input_key] = input_site + output_key = None + if output_site is not None: + output_key = output_site.name + self._activation_sites[output_key] = output_site + + def hook(_, args, __kwargs, output, input_key=input_key, output_key=output_key): + if not torch.is_grad_enabled(): + return + self._add_tensor(self._activation_sites[input_key], _first_item(args)) + if output_key is not None: + self._add_tensor(self._activation_sites[output_key], _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, loss_scale: float | None + ) -> None: + """Register backward hooks for dgrad raw moments.""" + assert not self._dgrad_hooks + self._dgrad_sites.clear() + self._pending_dgrad_loss_scale = loss_scale + for module, input_site, output_site in _iter_hook_modules(model): + input_key = input_site.name + self._dgrad_sites[input_key] = input_site + output_key = None + if output_site is not None: + output_key = output_site.name + self._dgrad_sites[output_key] = output_site + + def hook(_, grad_input, grad_output, input_key=input_key, output_key=output_key): + if output_key is not None: + self._add_tensor(self._dgrad_sites[output_key], _first_item(grad_output)) + self._add_tensor(self._dgrad_sites[input_key], _first_item(grad_input)) + + self._dgrad_hooks.append(module.register_full_backward_hook(hook)) + + def finalize_activation_raw_moments_by_layer(self) -> None: + """Reduce and cache activation raw moments for later logging.""" + self._latest_activation_raw_moments_by_layer = self._finalize_sites( + self._activation_sites.values() + ) + self._activation_sites.clear() + + def finalize_dgrad_raw_moments_by_layer(self) -> None: + """Reduce and cache dgrad raw moments for later logging.""" + self._latest_dgrad_raw_moments_by_layer = self._finalize_sites(self._dgrad_sites.values()) + self._latest_dgrad_loss_scale = self._pending_dgrad_loss_scale + self._pending_dgrad_loss_scale = None + self._dgrad_sites.clear() + + def consume_activation_raw_moments_by_layer( + self, + ) -> list[tuple[str, dict[str, float]]] | None: + """Return and clear the latest activation raw moments.""" + values = self._latest_activation_raw_moments_by_layer + self._latest_activation_raw_moments_by_layer = None + return values + + def consume_dgrad_raw_moments_by_layer( + self, + ) -> tuple[list[tuple[str, dict[str, float]]], float | None] | None: + """Return and clear the latest dgrad raw moments and loss scale.""" + values = self._latest_dgrad_raw_moments_by_layer + loss_scale = self._latest_dgrad_loss_scale + self._latest_dgrad_raw_moments_by_layer = None + self._latest_dgrad_loss_scale = None + if values is None: + return None + return values, loss_scale + + 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() + + @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 _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 _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, loss_scale: float | None = None +) -> None: + """Enable dgrad raw-moment logging on ``model``.""" + _get_logger().register_dgrad_hooks(model, loss_scale) + + +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( +) -> tuple[list[tuple[str, dict[str, float]]], float | None] | None: + """Return and clear the latest dgrad raw moments and loss scale.""" + 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() diff --git a/megatron/training/statistics_logging.py b/megatron/training/statistics_logging.py new file mode 100644 index 00000000000..32909b405c4 --- /dev/null +++ b/megatron/training/statistics_logging.py @@ -0,0 +1,162 @@ +# 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 + +import torch + + +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 open(filepath, "a") as f: + f.write(json.dumps(record) + "\n") + + +def save_raw_moments_by_param( + log_dir: str, + stat_name: str, + iteration: int, + consumed_train_samples: int, + raw_moments_by_param: Iterable[tuple[str, dict[str, float]]], + rank: int | None = None, + extra_record_fields: dict | None = None, +) -> None: + """Append one per-parameter raw moments record.""" + save_raw_moments_by_name( + log_dir, + stat_name, + iteration, + consumed_train_samples, + raw_moments_by_param, + rank=rank, + extra_record_fields=extra_record_fields, + ) + + +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, + extra_record_fields: dict | None = None, +) -> None: + """Append one named raw moments record.""" + 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, + } + if extra_record_fields is not None: + record.update(extra_record_fields) + + append_training_stat( + log_dir, + stat_name, + record, + rank=rank, + ) + + +def save_param_raw_moments_by_param( + log_dir: str, + iteration: int, + consumed_train_samples: int, + param_raw_moments_by_param: Iterable[tuple[str, dict[str, float]]], + rank: int | None = None, +) -> None: + """Append one per-parameter parameter raw moments record.""" + save_raw_moments_by_param( + log_dir, + "param_raw_moments_by_param", + iteration, + consumed_train_samples, + param_raw_moments_by_param, + rank=rank, + ) + + +def save_grad_raw_moments_by_param( + log_dir: str, + iteration: int, + consumed_train_samples: int, + grad_raw_moments_by_param: Iterable[tuple[str, dict[str, float]]], + rank: int | None = None, +) -> None: + """Append one pre-clipping per-parameter gradient raw moments record.""" + save_raw_moments_by_param( + log_dir, + "grad_raw_moments_by_param", + iteration, + consumed_train_samples, + grad_raw_moments_by_param, + rank=rank, + extra_record_fields={"gradient_stage": "pre_clip"}, + ) + + +def save_activation_raw_moments_by_layer( + log_dir: str, + iteration: int, + consumed_train_samples: int, + activation_raw_moments_by_layer: Iterable[tuple[str, dict[str, float]]], + rank: int | None = None, +) -> None: + """Append one activation raw moments record keyed by module site.""" + save_raw_moments_by_name( + log_dir, + "activation_raw_moments_by_layer", + iteration, + consumed_train_samples, + activation_raw_moments_by_layer, + rank=rank, + ) + + +def save_dgrad_raw_moments_by_layer( + log_dir: str, + iteration: int, + consumed_train_samples: int, + dgrad_raw_moments_by_layer: Iterable[tuple[str, dict[str, float]]], + rank: int | None = None, + loss_scale: float | None = None, +) -> None: + """Append one backward dgrad raw moments record keyed by module site.""" + extra_record_fields = {"gradient_stage": "backward_scaled"} + if loss_scale is not None: + extra_record_fields["loss_scale"] = float(loss_scale) + save_raw_moments_by_name( + log_dir, + "dgrad_raw_moments_by_layer", + iteration, + consumed_train_samples, + dgrad_raw_moments_by_layer, + rank=rank, + extra_record_fields=extra_record_fields, + ) diff --git a/megatron/training/training.py b/megatron/training/training.py index e825853fdb1..46f300fe61c 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -174,9 +174,26 @@ get_wandb_writer, ) from .theoretical_memory_usage import report_theoretical_memory +from .raw_moment_logging import ( + consume_activation_raw_moments_by_layer, + consume_dgrad_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, +) +from .statistics_logging import ( + save_activation_raw_moments_by_layer, + save_dgrad_raw_moments_by_layer, + save_grad_raw_moments_by_param, + save_param_raw_moments_by_param, +) 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, @@ -256,6 +273,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. @@ -295,6 +314,46 @@ 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 _warn_missing_statistics_log_dir(): + global _STATS_LOG_DIR_WARNING_SHOWN + if not _STATS_LOG_DIR_WARNING_SHOWN: + print_rank_0( + "WARNING: per-parameter 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 _get_activation_log_interval(args): + activation_log_interval = getattr(args, 'activation_log_interval', None) + if activation_log_interval is not None: + return activation_log_interval + return getattr(args, 'tensorboard_log_interval', None) + + 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``. @@ -2250,6 +2309,25 @@ 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_log_interval = _get_activation_log_interval(args) + log_activation_raw_moments_in_this_iteration = ( + getattr(args, 'log_activation_raw_moments_by_layer', False) + and iteration is not None + and activation_log_interval is not None + and (iteration + 1) % activation_log_interval == 0 + ) + log_dgrad_raw_moments_in_this_iteration = ( + getattr(args, 'log_dgrad_raw_moments_by_layer', False) + and iteration is not None + and activation_log_interval is not None + and (iteration + 1) % activation_log_interval == 0 + ) + if ( + log_activation_raw_moments_in_this_iteration or log_dgrad_raw_moments_in_this_iteration + ) and getattr(args, 'cuda_graph_impl', 'none') != 'none': + raise RuntimeError( + "Activation/dgrad raw moment logging is not supported with CUDA graph modes." + ) while rerun_state_machine.should_run_forward_backward(data_iterator): # Set grad to zero. for model_chunk in model: @@ -2299,6 +2377,13 @@ 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) + if log_activation_raw_moments_in_this_iteration: + enable_activation_raw_moment_logging(model) + if log_dgrad_raw_moments_in_this_iteration: + loss_scale_for_dgrad_raw_moments = None + if optimizer is not None and not optimizer.is_stub_optimizer: + loss_scale_for_dgrad_raw_moments = optimizer.get_loss_scale().item() + enable_dgrad_raw_moment_logging(model, loss_scale=loss_scale_for_dgrad_raw_moments) losses_reduced = forward_backward_func( forward_step_func=forward_step_func, data_iterator=data_iterator, @@ -2322,6 +2407,12 @@ 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_raw_moments_in_this_iteration: + finalize_activation_raw_moments_by_layer() + disable_activation_raw_moment_logging() + if log_dgrad_raw_moments_in_this_iteration: + finalize_dgrad_raw_moments_by_layer() + disable_dgrad_raw_moment_logging() # Reset force_all_reduce field. for model_chunk in model: @@ -2361,6 +2452,13 @@ def _save_state_dict(attr_name, label): # Update parameters. timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time) + if ( + optimizer is not None + and getattr(args, 'log_grad_raw_moments_by_param', False) + and iteration is not None + and (iteration + 1) % args.tensorboard_log_interval == 0 + ): + optimizer.request_grad_raw_moments_by_param(model) update_successful, grad_norm, num_zeros_in_grad = optimizer.step() # get max attention logit for logging and run clip_qk() @@ -3799,6 +3897,76 @@ def trace_handler(p): if args.log_params_norm: params_norm = calc_params_l2_norm(model) + if ( + getattr(args, 'log_param_raw_moments_by_param', False) + and iteration % args.tensorboard_log_interval == 0 + ): + param_raw_moments_by_param = calc_params_raw_moments_by_param(model) + 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_param_raw_moments_by_param( + statistics_log_dir, + iteration, + args.consumed_train_samples, + param_raw_moments_by_param, + ) + if ( + getattr(args, 'log_grad_raw_moments_by_param', False) + and iteration % args.tensorboard_log_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_grad_raw_moments_by_param( + statistics_log_dir, + iteration, + args.consumed_train_samples, + grad_raw_moments_by_param, + ) + activation_log_interval = _get_activation_log_interval(args) + if ( + getattr(args, 'log_activation_raw_moments_by_layer', False) + and activation_log_interval is not None + and iteration % activation_log_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_activation_raw_moments_by_layer( + statistics_log_dir, + iteration, + args.consumed_train_samples, + activation_raw_moments_by_layer, + ) + if ( + getattr(args, 'log_dgrad_raw_moments_by_layer', False) + and activation_log_interval is not None + and iteration % activation_log_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 is not None: + dgrad_raw_moments, dgrad_loss_scale = dgrad_raw_moments_by_layer + if dgrad_raw_moments: + save_dgrad_raw_moments_by_layer( + statistics_log_dir, + iteration, + args.consumed_train_samples, + dgrad_raw_moments, + loss_scale=dgrad_loss_scale, + ) 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 d6e2fe7c246..0e06e06ef95 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..5a201ee4a3a 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,16 +52,58 @@ ) 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): + """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 + ) + + +def _calc_params_l2_norm_or_raw_moments( + model, force_create_fp32_copy=False, raw_moments_by_param=False +): + """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 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. + if raw_moments_by_param: + raise RuntimeError( + "calc_params_raw_moments_by_param() is not implemented for --use-megatron-fsdp" + ) params = [] for model_chunk in model: model_chunk.stop_communication() @@ -70,51 +117,107 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): return calc_dtensor_params_l2_norm(params) + raw_moments_registry = ( + get_or_create_per_parameter_stat_registry(model) 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 +273,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 +311,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_per_parameter_stats.py b/tests/unit_tests/test_per_parameter_stats.py new file mode 100644 index 00000000000..35ec3340033 --- /dev/null +++ b/tests/unit_tests/test_per_parameter_stats.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core import per_parameter_stats as pps +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 _FakeCudaTensor: + device = torch.device("cuda:0") + dtype = torch.float32 + + def is_contiguous(self) -> bool: + return True + + +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(21.0), + "sum_3": pytest.approx(81.0), + "sum_4": pytest.approx(321.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(30.0), + "sum_3": pytest.approx(108.0), + "sum_4": pytest.approx(402.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_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..78a47f1db33 --- /dev/null +++ b/tests/unit_tests/test_raw_moment_logging.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch +import torch.nn as nn + +from megatron.training.raw_moment_logging import RawMomentLogger + + +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) + + +def _values_dict(values): + return {name: moments for name, moments in values} + + +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, loss_scale=None) + + 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, loss_scale = logger.consume_dgrad_raw_moments_by_layer() + values = _values_dict(values) + assert loss_scale is None + assert "output_layer/input0" in values + assert "output_layer/output0" not in values + + +def test_dgrad_raw_moments_accumulate_by_module_site_with_loss_scale(): + model = [_ToyModel()] + logger = RawMomentLogger() + logger.register_dgrad_hooks(model, loss_scale=128.0) + + 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, loss_scale = logger.consume_dgrad_raw_moments_by_layer() + values = _values_dict(values) + assert loss_scale == 128.0 + 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..1fb3e3d6f9e --- /dev/null +++ b/tests/unit_tests/test_statistics_logging.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import json + +from megatron.training.statistics_logging import ( + save_activation_raw_moments_by_layer, + save_dgrad_raw_moments_by_layer, + save_grad_raw_moments_by_param, + save_param_raw_moments_by_param, +) + + +def _read_records(filepath): + return [json.loads(line) for line in filepath.read_text().strip().split("\n")] + + +class TestSaveParamRawMomentsByParam: + def test_creates_jsonl(self, tmp_path): + save_param_raw_moments_by_param( + str(tmp_path), + iteration=100, + consumed_train_samples=8192, + param_raw_moments_by_param=[ + ( + "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" / "param_raw_moments_by_param" / "rank7.jsonl" + assert filepath.exists() + + records = _read_records(filepath) + assert records == [ + { + "iter": 100, + "consumed_train_samples": 8192, + "stat": "param_raw_moments_by_param", + "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_appends_across_calls(self, 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_param_raw_moments_by_param( + str(tmp_path), + iteration=100, + consumed_train_samples=8192, + param_raw_moments_by_param=[("layer.weight", first_moments)], + rank=0, + ) + save_param_raw_moments_by_param( + str(tmp_path), + iteration=200, + consumed_train_samples=16384, + param_raw_moments_by_param=[("layer.weight", second_moments)], + rank=0, + ) + + filepath = tmp_path / "training_stats" / "param_raw_moments_by_param" / "rank0.jsonl" + records = _read_records(filepath) + assert len(records) == 2 + assert records[0]["iter"] == 100 + 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]["iter"] == 200 + 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_empty_values_do_not_create_file(self, tmp_path): + save_param_raw_moments_by_param( + str(tmp_path), + iteration=100, + consumed_train_samples=8192, + param_raw_moments_by_param=[], + rank=0, + ) + + filepath = tmp_path / "training_stats" / "param_raw_moments_by_param" / "rank0.jsonl" + assert not filepath.exists() + + +class TestSaveGradRawMomentsByParam: + def test_creates_jsonl(self, tmp_path): + save_grad_raw_moments_by_param( + str(tmp_path), + iteration=100, + consumed_train_samples=8192, + grad_raw_moments_by_param=[ + ( + "decoder.layers.0.self_attention.linear_qkv.weight", + {"count": 4, "sum_1": 3.5, "sum_2": 6.5, "sum_3": 9.5, "sum_4": 12.5}, + ), + ( + "decoder.layers.0.mlp.linear_fc1.weight", + {"count": 5, "sum_1": 4.25, "sum_2": 8.25, "sum_3": 12.25, "sum_4": 16.25}, + ), + ], + rank=3, + ) + + filepath = tmp_path / "training_stats" / "grad_raw_moments_by_param" / "rank3.jsonl" + records = _read_records(filepath) + assert records == [ + { + "iter": 100, + "consumed_train_samples": 8192, + "stat": "grad_raw_moments_by_param", + "gradient_stage": "pre_clip", + "values": { + "decoder.layers.0.self_attention.linear_qkv.weight": { + "count": 4.0, + "sum_1": 3.5, + "sum_2": 6.5, + "sum_3": 9.5, + "sum_4": 12.5, + }, + "decoder.layers.0.mlp.linear_fc1.weight": { + "count": 5.0, + "sum_1": 4.25, + "sum_2": 8.25, + "sum_3": 12.25, + "sum_4": 16.25, + }, + }, + } + ] + + +class TestSaveActivationRawMomentsByLayer: + def test_creates_jsonl(self, tmp_path): + save_activation_raw_moments_by_layer( + str(tmp_path), + iteration=100, + consumed_train_samples=8192, + activation_raw_moments_by_layer=[ + ( + "decoder.layers.0.self_attention.linear_qkv/output0", + {"count": 2, "sum_1": 1.5, "sum_2": 2.5, "sum_3": 3.5, "sum_4": 4.5}, + ) + ], + rank=2, + ) + + filepath = tmp_path / "training_stats" / "activation_raw_moments_by_layer" / "rank2.jsonl" + records = _read_records(filepath) + assert records == [ + { + "iter": 100, + "consumed_train_samples": 8192, + "stat": "activation_raw_moments_by_layer", + "values": { + "decoder.layers.0.self_attention.linear_qkv/output0": { + "count": 2.0, + "sum_1": 1.5, + "sum_2": 2.5, + "sum_3": 3.5, + "sum_4": 4.5, + } + }, + } + ] + + +class TestSaveDgradRawMomentsByLayer: + def test_creates_jsonl_with_loss_scale(self, tmp_path): + save_dgrad_raw_moments_by_layer( + str(tmp_path), + iteration=100, + consumed_train_samples=8192, + dgrad_raw_moments_by_layer=[ + ( + "decoder.layers.0.self_attention.linear_qkv/input0", + {"count": 3, "sum_1": 2.5, "sum_2": 4.5, "sum_3": 6.5, "sum_4": 8.5}, + ) + ], + rank=5, + loss_scale=128.0, + ) + + filepath = tmp_path / "training_stats" / "dgrad_raw_moments_by_layer" / "rank5.jsonl" + records = _read_records(filepath) + assert records == [ + { + "iter": 100, + "consumed_train_samples": 8192, + "stat": "dgrad_raw_moments_by_layer", + "gradient_stage": "backward_scaled", + "loss_scale": 128.0, + "values": { + "decoder.layers.0.self_attention.linear_qkv/input0": { + "count": 3.0, + "sum_1": 2.5, + "sum_2": 4.5, + "sum_3": 6.5, + "sum_4": 8.5, + } + }, + } + ] From 27c3516360462752d4edd2b8f1dcb07a1ee69dd0 Mon Sep 17 00:00:00 2001 From: Philip Monk Date: Wed, 24 Jun 2026 19:49:11 -0700 Subject: [PATCH 2/2] thread EP groups into per_parameter_stats Signed-off-by: Philip Monk --- megatron/core/optimizer/optimizer.py | 18 ++++- megatron/core/per_parameter_stats.py | 81 ++++++++++++++++---- megatron/training/training.py | 17 +++- megatron/training/utils/common_utils.py | 23 +++++- tests/unit_tests/test_per_parameter_stats.py | 57 ++++++++++++++ 5 files changed, 169 insertions(+), 27 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 36a7cc813ce..f6295d0229e 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -321,18 +321,28 @@ def get_raw_moment_buckets_for_grad_norm( return [NamedTensorBucket(names, grads, reduce_groups)] def get_grad_raw_moments_by_param( - self, registry: PerParameterStatRegistry | None = None + 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) + 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) -> None: + 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) + 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 diff --git a/megatron/core/per_parameter_stats.py b/megatron/core/per_parameter_stats.py index 3e084e87468..dd3346387de 100644 --- a/megatron/core/per_parameter_stats.py +++ b/megatron/core/per_parameter_stats.py @@ -17,8 +17,7 @@ multi_tensor_applier = None multi_tensor_raw_moments = None -from megatron.core import parallel_state -from megatron.core.utils import unwrap_model +from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model _LAYER_NAME_PATTERN = re.compile(r"layers\.(\d+)") _GROUPED_EXPERT_PATTERN = re.compile(r"^(.*\.mlp\.experts\.linear_fc\d\.weight)(\d+)(.*)$") @@ -46,9 +45,22 @@ class NamedTensorBucket: class PerParameterStatRegistry: """Canonical parameter-name registry for per-parameter statistics.""" - def __init__(self, model_chunks: Iterable[torch.nn.Module] | torch.nn.Module): + 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.cache_key = tuple(id(model_chunk) for model_chunk in self.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.param_to_name = self._build_local_param_to_name() self.name_to_index = self._build_name_to_index() self.index_to_name = sorted(self.name_to_index, key=self.name_to_index.get) @@ -65,10 +77,13 @@ def num_params(self) -> int: def _build_local_param_to_name(self) -> dict[torch.nn.Parameter, str]: param_to_name = {} num_experts = _get_num_moe_experts(self.model_chunks) + expert_offset = _get_local_expert_offset( + num_experts, self.expert_model_parallel_rank, self.expert_model_parallel_size + ) for model_chunk in self.model_chunks: for local_name, param in model_chunk.named_parameters(): param_to_name[param] = _canonical_param_name( - model_chunk, local_name, param, num_experts + model_chunk, local_name, param, num_experts, expert_offset ) return param_to_name @@ -87,17 +102,28 @@ def _build_name_to_index(self) -> dict[str, int]: 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.") - cache_key = tuple(id(model_chunk) for model_chunk in unwrapped_model_chunks) + 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) + registry = PerParameterStatRegistry( + unwrapped_model_chunks, expert_model_parallel_group=expert_model_parallel_group + ) cache_owner._per_parameter_stat_registry = registry return registry @@ -297,9 +323,10 @@ def _canonical_param_name( local_name: str, param: torch.nn.Parameter, num_experts: int | None, + expert_offset: int, ) -> str: name = _global_layer_param_name(model_chunk, local_name, param) - return _global_expert_param_name(name, num_experts) + return _global_expert_param_name(name, num_experts, expert_offset) def _global_layer_param_name( @@ -319,11 +346,12 @@ def _global_layer_param_name( return local_name -def _global_expert_param_name(local_name: str, num_experts: int | None) -> str: +def _global_expert_param_name( + local_name: str, num_experts: int | None, expert_offset: int +) -> str: if not num_experts: return local_name - expert_offset = _get_local_expert_offset(num_experts) if expert_offset == 0: return local_name @@ -340,15 +368,34 @@ def _global_expert_param_name(local_name: str, num_experts: int | None) -> str: return local_name -def _get_local_expert_offset(num_experts: int) -> int: - expert_group = parallel_state.get_expert_model_parallel_group(check_initialized=False) - if expert_group is None: - return 0 - expert_parallel_size = parallel_state.get_expert_model_parallel_world_size() - if expert_parallel_size <= 1: +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, + ) + + +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 local_experts = num_experts // expert_parallel_size - return parallel_state.get_expert_model_parallel_rank() * local_experts + return expert_parallel_rank * local_experts def _get_num_moe_experts(model_chunks: Sequence[torch.nn.Module]) -> int | None: diff --git a/megatron/training/training.py b/megatron/training/training.py index 46f300fe61c..9dcc6d3bb1a 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -336,6 +336,14 @@ def _should_write_global_training_stats(args): 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: @@ -2458,7 +2466,9 @@ def _save_state_dict(attr_name, label): and iteration is not None and (iteration + 1) % args.tensorboard_log_interval == 0 ): - optimizer.request_grad_raw_moments_by_param(model) + 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() @@ -3901,7 +3911,10 @@ def trace_handler(p): getattr(args, 'log_param_raw_moments_by_param', False) and iteration % args.tensorboard_log_interval == 0 ): - param_raw_moments_by_param = calc_params_raw_moments_by_param(model) + 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() diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 5a201ee4a3a..4f3b26165c0 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -64,15 +64,23 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): ) -def calc_params_raw_moments_by_param(model, force_create_fp32_copy=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 + 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 + 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. @@ -117,8 +125,15 @@ def _calc_params_l2_norm_or_raw_moments( 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) if raw_moments_by_param else None + 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 diff --git a/tests/unit_tests/test_per_parameter_stats.py b/tests/unit_tests/test_per_parameter_stats.py index 35ec3340033..f5b4555bf75 100644 --- a/tests/unit_tests/test_per_parameter_stats.py +++ b/tests/unit_tests/test_per_parameter_stats.py @@ -19,6 +19,20 @@ def __init__(self): 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 @@ -27,6 +41,18 @@ 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()) @@ -84,6 +110,37 @@ def test_registry_cache_is_per_model_identity(): 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 set(registry.param_to_name.values()) == { + "mlp.experts.local_experts.2.weight", + "mlp.experts.local_experts.3.weight", + } + + def test_local_raw_moments_multi_tensor_path_preserves_order(monkeypatch): calls = []