megatron/core/optimizer/optimizer_config.py

# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

import fnmatch
from dataclasses import dataclass, field
from typing import Callable, Optional, Tuple, Union

import torch

from ..utils import is_te_min_version


@dataclass(frozen=True)
class ParamPredicate:
    """Wraps a matching function to make it hashable for ParamKey.
    Example:
        >>> shape_1_param = ParamPredicate(name="s1", fn=lambda param: len(param.shape) == 1)
        >>> shape_1_param(torch.empty(10))
        True
        >>> shape_1_param_copy = ParamPredicate(name="s1", fn=lambda param: len(param.shape) == 1)
        >>> shape_1_param == shape_1_param_copy  # name is used to match
        True
        >>> {shape_1_param, shape_1_param_copy} == {shape_1_param}  # set hashing works properly

    NOTE:
        __hash__ and __eq__ are automatically generated by @dataclass(frozen=True)
        based solely on 'name' because we set compare=False/hash=False on 'fn'.
    """

    name: str
    fn: Callable[[torch.nn.Parameter], bool] = field(compare=False, hash=False)

    def __call__(self, param: torch.nn.Parameter) -> bool:
        return self.fn(param)


@dataclass(frozen=True)
class ParamWithNamePredicate:
    """Wraps a matching function to make it hashable for ParamKey.
    Example:
        >>> shape_1_not_qkln_param = ParamWithNamePredicate(
                name="s1_not_qkln",
                fn=lambda param, name: (
                    len(param.shape) == 1 or name.endswith(".bias")
                    and not ("q_layernorm." in name or "k_layernorm." in name)
                )
            )
        >>> shape_1_not_qkln_param(torch.empty(10), "interesting.bias")
        True
        >>> shape_1_not_qkln_param(torch.empty(10), "interesting.q_layernorm.bias")
        False

    NOTE:
        __hash__ and __eq__ are automatically generated by @dataclass(frozen=True)
        based solely on 'name' because we set compare=False/hash=False on 'fn'.
    """

    name: str
    fn: Callable[[torch.nn.Parameter, str], bool] = field(compare=False, hash=False)

    def __call__(self, param: torch.nn.Parameter, name: str) -> bool:
        return self.fn(param, name)


@dataclass(frozen=True, slots=True)
class ParamKey:
    """Key to group parameters by. All such grouped parameters can share an
    optimizer config specification."""

    # TODO: Can add layer_id here later.

    name: Union[str, Tuple[str]] = field(default_factory=tuple)
    """Parameter name(s), will use unix filesystem path syntax for matching."""

    attr: Union[str, Tuple[str]] = field(default_factory=tuple)
    """Parameter attribute(s)."""

    predicate: Union[ParamPredicate, Tuple[ParamPredicate]] = field(default_factory=tuple)
    """Predicate(s) to match parameters by. If multiple predicates are provided, any must match."""

    with_name_predicate: Union[ParamWithNamePredicate, Tuple[ParamWithNamePredicate]] = field(
        default_factory=tuple
    )
    """
    Predicate(s) to match parameters with their name. If multiple predicates are provided, 
      any must match. This is useful if you need to filter out some parameters from an otherwise 
      positive match by their name.
    """

    def matches(self, param: torch.nn.Parameter, param_name: str) -> bool:
        """Returns true if passed-in parameter (with name) matches `param_key`.

        Args:
            param (torch.nn.Parameter): Handle to parameter object.
            param_name (str): Name of parameter in underlying PyTorch module.

        Returns:
            bool: True if parameter matches passed-in param_key.
        """

        # Check if name matches.
        if isinstance(self.name, str):
            target_names = [self.name]
        else:
            target_names = list(self.name)
        for target_name in target_names:
            if fnmatch.fnmatch(param_name, target_name):
                return True

        # Check if attribute matches.
        if isinstance(self.attr, str):
            target_attrs = [self.attr]
        else:
            target_attrs = list(self.attr)
        for target_attr in target_attrs:
            if getattr(param, target_attr, False):
                return True

        # Check if predicate matches.
        if isinstance(self.predicate, ParamPredicate):
            if self.predicate(param):
                return True
        else:
            for predicate in self.predicate:
                if predicate(param):
                    return True

        # Check if with_name_predicate matches.
        if isinstance(self.with_name_predicate, ParamWithNamePredicate):
            if self.with_name_predicate(param, param_name):
                return True
        else:
            for predicate in self.with_name_predicate:
                if predicate(param, param_name):
                    return True
        return False


@dataclass
class OptimizerConfig:
    """Base optimizer configuration object."""

    ##############
    # General
    ##############

    lr: Optional[float] = None
    """Initial learning rate. Depending on decay style and initial warmup, the learning rate at each
       iteration would be different.
    """

    min_lr: Optional[float] = None
    """Minumum value for learning rate. The scheduler clip values below this threshold."""

    decoupled_lr: Optional[float] = None
    """Separate learning rate for the input and output layer."""

    decoupled_min_lr: Optional[float] = None
    """Minimum value for learning rate for the input and output layer. The scheduler clip values
       below this threshold.
    """

    weight_decay: float = 0.01
    """Weight decay coefficient for L2 regularization."""

    apply_wd_to_qk_layernorm: bool = False
    """If true, apply weight decay to qk layernorm as a special case."""

    ##############
    # Precision
    ##############
    fp8_recipe: Optional[str] = None
    """The type of fp8 recipe will affect the processing logic inside distributed optimizer."""

    fp16: bool = False
    """If true, train with fp16 mixed precision training. Defaults to False."""

    bf16: bool = False
    """If true, train with bf16 mixed precision training. Defaults to False."""

    pure_bf16_optimizer: bool = False
    """If true (with bf16), keep main params and optimizer states in bf16."""

    reuse_grad_buf_for_mxfp8_param_ag: bool = False
    """If true, reuse the grad buffer for param AG when using mxfp8 recipe. Should be 
       set to True only when fp8_recipe is mxfp8 and fp8_param_gather is True."""

    params_dtype: torch.dtype = torch.float32
    """dtype used when intializing the weights. Defaults to torch.float32."""

    use_precision_aware_optimizer: bool = False
    """If true, allows optimizer-related tensors (master_param, gradients and optimizer states)
    to be set to lower precision. Defaults to False.
    """

    store_param_remainders: bool = True
    """If true, store the 16-bit FP32 parameter remainders in the optimizer state, excluding the
        16 bits shared with the BF16 parameters. This lowers GPU memory usage. Defaults to True.
    """

    main_grads_dtype: torch.dtype = torch.float32
    """dtype of main grads when enabling precision-aware-optimizer"""

    main_params_dtype: torch.dtype = torch.float32
    """dtype of main params when enabling precision-aware-optimizer"""

    exp_avg_dtype: torch.dtype = torch.float32
    """dtype of exp_avg when enabling precision-aware-optimizer"""

    exp_avg_sq_dtype: torch.dtype = torch.float32
    """dtype of exp_avg_sq when enabling precision-aware-optimizer"""

    optimizer: str = 'adam'
    """Optimizer name. NOTE: Deprecated, use individual optimizer classes instead."""

    ###############
    # Loss scaling
    ###############
    loss_scale: Optional[float] = None
    """Static loss scaling, positive power of 2 values can improve fp16 convergence. If None,
       dynamic loss scaling is used.
    """

    initial_loss_scale: float = 2**32
    """Initial loss-scale for dynamic loss scaling."""

    min_loss_scale: float = 1.0
    """Minimum loss scale for dynamic loss scaling."""

    loss_scale_window: float = 1000
    """Window over which to raise/lower dynamic scale."""

    hysteresis: int = 2
    """Hysteresis for dynamic loss scaling."""

    ###################################################################################
    # Optimizer (NOTE: Deprecated, use individual optimizer classes instead.).
    ###################################################################################
    # Adam.
    adam_beta1: float = 0.9
    """First coefficient for computing running averages of gradient and its square in Adam
    optimizer.
    """

    adam_beta2: float = 0.999
    """Second coefficient for computing running averages of gradient and its square in Adam
    optimizer.
    """

    adam_eps: float = 1e-08
    """Term added to the denominator to improve numerical stability in Adam optimizer."""

    decoupled_weight_decay: bool = True
    """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false,
    original Adam update rule will be used. Defaults to True.
    """

    min_singular_reg_interval: int = 0
    """If > 0, every this many steps add sigma_min*u_min*v_min^T to 2D params before weight decay
    to smooth spectrum and reduce condition number. When 0, this regularization is disabled."""

    min_singular_reg_scale: float = 1.0
    """Scale for the min-singular rank-1 update: param += scale * sigma_min * u_min * v_min^T."""

    # SGD.
    sgd_momentum: float = 0.9
    """Momentum factor for SGD optimizer."""

    # Muon.
    # TODO: move muon configs to it's own `MuonConfig`.
    muon_momentum: float = 0.95
    """The momentum used by the internal SGD."""

    muon_split_qkv: bool = True
    """Whether to split QKV parameters for Muon optimizer."""

    muon_use_nesterov: bool = False
    """Whether to use Nesterov-style momentum in the internal SGD."""

    muon_scale_mode: str = "spectral"
    """The mode to use for the scale factor. Defaults to "spectral"."""

    muon_fp32_matmul_prec: str = "medium"
    """The precision to use for the fp32 matmul. Defaults to "medium"."""

    muon_num_ns_steps: int = 5
    """The number of iteration steps to use in the Newton-Schulz iteration."""

    muon_tp_mode: str = "blockwise"
    """How to perform NS calculation for tensor parallel weights. Defaults to "blockwise"."""

    muon_extra_scale_factor: float = 1.0
    """Additional scale factor for the muon update."""

    # Pion.
    pion_degree: int = 2
    """The degree of the Pion optimizer."""

    pion_beta1: float = 0.9
    """First momentum coefficient for Pion."""

    pion_beta2: float = 0.999
    """Second momentum coefficient for Pion."""

    pion_split_qkv: bool = True
    """Whether to split QKV parameters for Pion optimizer."""

    pion_split_gate: bool = True
    """When True, split up_project and gate_project for SwiGLU linear_fc1 in Pion and in spectral norm init (same style as pion_split_qkv)."""

    pion_split_qkv_per_head: bool = True
    """Whether to split Q/K/V per head (per group in GQA) in Pion and in spectral norm init. Same style as pion_split_gate."""

    pion_qkv_split_granularity: Optional[str] = None
    """Attention QKV update granularity for Pion: head | qkv | group. None follows pion_split_qkv_per_head."""

    pion_first_momentum: str = "none"
    """Compatibility selector for Pion first-momentum geometry: lie | ambient | ambient_transport."""

    pion_second_momentum: str = "none"
    """Compatibility selector for enabling second moment: none disables it, any other value enables it."""

    pion_12_momentum: str = "none"
    """Joint first+second order modes: none | lie_lie | transported_ambient_ambient."""

    pion_momentum: str = "none"
    """Unified Pion momentum geometry selector: lie | lie_lie | transported_ambient | transported_ambient_ambient."""

    pion_use_second_momentum: Optional[bool] = None
    """Whether Pion applies second-moment normalization inside the selected geometry. None preserves compatibility inference."""

    pion_exp_map: str = "exp_truncated"
    """Matrix map for Pion updates. Unified pion.py uses exp_truncated."""

    pion_spectrum_reset_interval: int = 0
    """If > 0, every this many steps reset matrix singular values to the initialization spectrum in Pion."""

    pion_scaling: str = "rms"
    """Scaling mode for Pion biside updates."""

    pion_update_side: str = "both"
    """Pion update side: both | alternate. "alternate" switches in/out by parameter step."""

    pion_update_csv: Optional[str] = None
    """If set, write Pion update stats to this CSV path."""

    pion_update_csv_interval: int = 1
    """Write Pion update stats every this many steps."""
    
    #######################
    # Distributed optimizer
    #######################
    use_distributed_optimizer: bool = False
    """Distribute optimizer state over data-parallel replicas."""

    overlap_param_gather: bool = False
    """If true, overlap param all-gather with forward compute. 
        This argument is intended to have the same value as the "overlap_param_gather" argument 
        in the "distributed_data_parallel_config.py" file. In the optimizer, this argument is 
        only used when "reuse_grad_buf_for_mxfp8_param_ag=True & fp8_param_gather=True".
    """

    overlap_param_gather_with_optimizer_step: bool = False
    """If true, overlap param all-gather of first bucket with optimizer step."""

    #######################
    # Optimizer Offload
    #######################

    optimizer_cpu_offload: bool = False
    """If True, offload optimizer states tensor and compute to CPU."""

    optimizer_offload_fraction: float = 0.0
    """Specifies the fraction of optimizer states to offload from GPU memory to CPU."""

    use_torch_optimizer_for_cpu_offload: bool = False
    """If True, use torch.optim.Optimizer for CPU offload."""

    overlap_cpu_optimizer_d2h_h2d: bool = False
    """
    When set to `True`, this flag enables overlapping of the CPU optimizer
    update process with the data transfer operations. This can help improve
    overall training efficiency by reducing idle time during data movement,
    allowing the optimizer to perform updates while gradients and parameters
    are being transferred between devices.
    """

    pin_cpu_grads: bool = True
    """If True, pin the optimizer gradients to CPU memory."""

    pin_cpu_params: bool = True
    """If True, pin the optimizer parameters to CPU memory."""

    ################
    # Miscellaneous
    ################
    clip_grad: float = 1.0
    """Gradient clipping based on global L2 norm."""

    log_num_zeros_in_grad: bool = False
    """If true, calculate and log the number of zeros in gradient."""

    barrier_with_L1_time: bool = False
    """If true, use barrier with level 1 time measurements."""

    timers: Optional[Callable] = None
    """Function to get timers."""

    config_logger_dir: str = ""
    """When non-empty, dumps entry-point configs to config_logger_dir"""

    def __post_init__(self):
        """Check the validity of the config."""
        
        if self.pure_bf16_optimizer:
            assert self.bf16, "--pure-bf16-optimizer requires --bf16."
            assert not self.fp16, "--pure-bf16-optimizer is incompatible with --fp16."
            assert (
                not self.use_distributed_optimizer
            ), "--pure-bf16-optimizer currently supports non-distributed optimizer only."
            assert (
                not self.use_precision_aware_optimizer
            ), "--pure-bf16-optimizer should not be combined with --use-precision-aware-optimizer."

        # The following condition is used to avoid repetition in distrib_optimizer.py.
        # This is because in distrib_optimizer.py, the process to handle parameters are
        # different for different training precision settings. FP8 cases require different
        # handling while FP8 delayed scaling is an exception because the Adam optimizer in
        # TransformerEngine supports it in the kernel computation.
        # This is also the flag to determine the usage of param.grad or param.decoupled_grad
        self.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = (
            self.use_precision_aware_optimizer
            and (
                self.main_params_dtype != torch.float32
                or (self.fp8_recipe is None or self.fp8_recipe == "delayed")
                or self.optimizer_cpu_offload
            )
        )

        if self.fp8_recipe == "mxfp8":
            if not self.reuse_grad_buf_for_mxfp8_param_ag:
                import warnings

                warnings.warn(
                    "mxfp8 without using reuse_grad_buf_for_mxfp8_param_ag and fp8_param_gather"
                    "will use significant amount additional GPU memory."
                    "Setting --reuse-grad-buf-for-mxfp8-param-ag and --fp8-param-gather is "
                    "recommended for mxfp8 training."
                )

        if self.use_precision_aware_optimizer:
            assert (
                self.optimizer == 'adam'
            ), '--use-precision-aware-optimizer only supported with adam'
            assert (
                self.use_distributed_optimizer
            ), '--use-precision-aware-optimizer only supported with distributed optimizer'

            if not is_te_min_version("2.1.0"):
                self.store_param_remainders = False

            # Only the FusedAdam in TE and HybridDeviceOptimizer supports
            # --use-precision-aware-optimizer.
            # TODO: Remove this check when apex's FusedAdam is no longer used.
            if self.optimizer_cpu_offload:
                return
            try:
                import inspect

                # TODO: Move this below?
                from transformer_engine.pytorch.optimizers import FusedAdam as Adam

                adam_args = inspect.signature(Adam).parameters
                arg_names = [
                    'master_weight_dtype',
                    'exp_avg_dtype',
                    'exp_avg_sq_dtype',
                    'use_decoupled_grad',
                ]
                for name in arg_names:
                    assert name in adam_args, (
                        "Current FusedAdam of TE doesn't support --use-precision-aware-optimizer, "
                        "please update TE version."
                    )
            except ImportError:
                raise RuntimeError(
                    '--use-precision-aware-optimizer requires FusedAdam from TransformerEngine, '
                    'but not found.'
                )
        else:
            assert (
                self.main_grads_dtype == torch.float32
            ), "main_grads_dtype can only be fp32 when not using precision-aware optimizer"
            assert (
                self.main_params_dtype == torch.float32
            ), "main_params_dtype can only be fp32 when not using precision-aware optimizer"
            assert (
                self.exp_avg_dtype == torch.float32
            ), "exp_avg_dtype can only be fp32 when not using precision-aware optimizer"
            assert (
                self.exp_avg_sq_dtype == torch.float32
            ), "exp_avg_sq_dtype can only be fp32 when not using precision-aware optimizer"


@dataclass
class AdamOptimizerConfig(OptimizerConfig):
    """Adam optimizer configuration object."""

    optimizer: str = 'adam'
    """Optimizer name."""

    adam_beta1: float = 0.9
    """First coefficient for computing running averages of gradient and its square in Adam
    optimizer.
    """

    adam_beta2: float = 0.999
    """Second coefficient for computing running averages of gradient and its square in Adam
    optimizer.
    """

    adam_eps: float = 1e-08
    """Term added to the denominator to improve numerical stability in Adam optimizer."""


@dataclass
class SGDOptimizerConfig(OptimizerConfig):
    """SGD optimizer configuration object."""

    optimizer: str = 'sgd'
    """Optimizer name."""

    sgd_momentum: float = 0.9
    """Momentum factor for SGD optimizer."""