"""Configuration for MuP (Maximal Update Parametrisation) and CompleteP.

This configuration is used to control various aspects of the MuP and CompleteP
mechanisms, including whether they are enabled, how they scale model parameters,
and how they affect the training process.

References
----------
- CompleteP: https://arxiv.org/abs/2505.01618
- MuP: https://arxiv.org/abs/2203.03466
- Practitioner's Guide to MuP: https://blog.eleuther.ai/mutransfer/
- Original Code: https://github.com/EleutherAI/nanoGPT-mup
- Our Code: https://github.com/anon-anon/nanoGPT-mt_dao

"""

from dataclasses import dataclass


@dataclass
class MuPConfig:
    """Configuration for MuP (Maximal Update Parametrisation) and CompleteP.

    Attributes
    ----------
    mup_enabled : bool
        Whether to enable MuP.
    mup_disable_attention_scaling : bool
        Whether to disable attention scaling in MuP.
        # NOTE: using 1/d instead of 1/sqrt(d) performs badly in practice.
        this is why the default is false
    mup_disable_hidden_lr_scaling: bool
        Whether to disable hidden layer learning rate scaling in MuP.
    mup_width_multiplier : float
        The width multiplier for MuP, controlling model hyperparameter scaling.
    mup_input_alpha : float
        The input alpha for MuP, controlling input layer output scaling.
    mup_output_alpha : float
        The output alpha for MuP, controlling output layer scaling.
    mup_scale_expert_sel: bool
        Whether to scale expert selection by 1/sqrt(mup_width_multiplier).
    completep_depth_alpha_enabled : bool
        Whether to enable depth alpha in CompleteP for residual stream scaling.
    completep_depth_multiplier : float
        The depth multiplier for CompleteP, scaling a block's contribution.
    completep_depth_alpha_exp : float
        The exponent for scaling a block's contribution in CompleteP.
    completep_eps_scaling_enabled : bool
        Whether to enable epsilon scaling for ADAM/ADOPT in CompleteP.

    """

    mup_enabled: bool = False
    mup_disable_attention_scaling: bool = True
    mup_disable_hidden_lr_scaling: bool = False
    mup_width_multiplier: float = 1.0
    mup_input_alpha: float = 1.0
    mup_output_alpha: float = 1.0
    mup_scale_expert_sel: bool = True
    completep_depth_alpha_enabled: bool = True
    completep_depth_multiplier: float = 1.0
    completep_depth_alpha_exp: float = 1.0
    completep_eps_scaling_enabled: bool = True
