# Copyright 2024 MosaicML LLM Foundry authors
# SPDX-License-Identifier: Apache-2.0

"""muP-enabled MPT model.

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 __future__ import annotations

import copy
import logging
from dataclasses import asdict
from functools import cached_property
from logging import DEBUG, WARNING
from typing import TYPE_CHECKING, Any, cast

import torch
from composer.utils import dist
from flwr.common import log as flwr_log
from llmfoundry import ComposerMPTCausalLM, MPTForCausalLM
from llmfoundry.layers_registry import ffns_with_megablocks
from llmfoundry.models.layers.custom_embedding import SharedEmbedding
from llmfoundry.models.mpt import MPTConfig, MPTModel
from llmfoundry.models.utils.mpt_param_count import megablocks_n_active_params
from llmfoundry.registry import models, norms
from llmfoundry.utils.builders import (
    _extract_param_groups,  # noqa: PLC2701
)
from torch import nn
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from mt_dao.models.deptoe.init_config import ModelInitConfig
from mt_dao.models.deptoe.mup_config import MuPConfig
from mt_dao.models.embeddings.mup_shared_embeddings import MuPSharedEmbedding
from mt_dao.models.ffns.sigma_moe import SigmaMoE
from mt_dao.models.mpt.mup_completep_block import MPTCompletePBlock
from mt_dao.models.utils import dtensor_safe_check_numel

if TYPE_CHECKING:
    from collections.abc import Iterable, MutableMapping

    from transformers.configuration_utils import PretrainedConfig
    from transformers.modeling_outputs import (
        BaseModelOutputWithPast,
    )

log = logging.getLogger(__name__)


def sigma_moe_n_active_params(model: ComposerMPTMuPCausalLM) -> int:
    """Calculate the number of active parameters in a SigmaMoE model.

    This function iterates through the model's modules and counts the number of
    parameters in SigmaMoE modules, dividing by the number of active experts.

    Parameters
    ----------
    model : ComposerMPTMuPCausalLM
        The MPT model with SigmaMoE layers.

    Returns
    -------
    int
        The number of active parameters in the model.

    """
    n_active_params = 0
    n_active_experts = cast("MPTMuPConfig", model.config).ffn_config.get(
        "ff_n_active_experts",
    )
    n_total_experts = cast("MPTMuPConfig", model.config).ffn_config.get(
        "ff_n_experts",
    )
    assert n_active_experts is not None, "ff_n_active_experts must be set in the config"
    assert n_total_experts is not None, "ff_n_experts must be set in the config"
    n_active_experts = int(n_active_experts)
    assert n_active_experts > 0, (
        f"ff_n_active_experts must be greater than 0, not {n_active_experts}"
    )
    for module in model.modules():
        if isinstance(module, SigmaMoE):
            up_proj_n_params = dtensor_safe_check_numel(module.up_proj)
            down_proj_n_params = dtensor_safe_check_numel(module.down_proj)
            n_active_params += int(
                (up_proj_n_params + down_proj_n_params)
                / n_total_experts
                * n_active_experts,
            )
            # If the module exists, it is a torch.nn.Parameter
            if hasattr(module, "expert_sel"):
                n_active_params += dtensor_safe_check_numel(module.expert_sel)
        else:
            for p in module.parameters(recurse=False):
                n_active_params += p.numel()

    return n_active_params


def convert_sigma_moe_config_to_dense(
    sigma_moe_block_args: dict[str, Any],
) -> dict[str, Any]:
    """Convert a SigmaMoE block configuration to a dense block configuration.

    This function will use the same parameters as the SigmaMoE block, but will
    create a dense block instead. In particular, it will set the `ffn_config`
    keys of the block config to be:
        {
            "ffn_act_fn": {"name": "silu"},
            "ffn_type": "mptmlp",
            "fc_type": {"name": "torch"},
        }
    It will also correct the `ffn_hidden_size` to be equal to:
    `ff_n_experts * ff_expert_size` in case it differs.

    Parameters
    ----------
    sigma_moe_block_args : dict[str, Any]
        The SigmaMoE block configuration to convert.

    Returns
    -------
    dict[str, Any]
        The converted dense block configuration.

    """
    ff_n_experts = sigma_moe_block_args["ffn_config"].get("ff_n_experts", 1)
    ff_expert_size = sigma_moe_block_args["ffn_config"].get("ff_expert_size", 1)
    sigma_moe_block_args["ffn_config"] = {
        "ffn_act_fn": {"name": "silu"},
        "ffn_type": "mptmlp",
        "fc_type": {"name": "torch"},
    }
    ffn_hidden_size = sigma_moe_block_args["ffn_hidden_size"]
    if ff_n_experts * ff_expert_size != ffn_hidden_size:
        log.warning(
            "Converting SigmaMoE block to dense block, but ff_n_experts (%s) *"
            " ff_expert_size (%s) != ffn_hidden_size (%s). We will use ff_n_experts *"
            " ff_expert_size = %s as new hidden dimension to be consistent.",
            ff_n_experts,
            ff_expert_size,
            ffn_hidden_size,
            int(ff_n_experts * ff_expert_size),
        )
        sigma_moe_block_args["ffn_hidden_size"] = int(
            ff_n_experts * ff_expert_size,
        )
    return sigma_moe_block_args


class MPTMuPConfig(MPTConfig):
    """Configuration for MPT with muP support."""

    def __init__(
        self,
        *args: Any,  # noqa: ANN401
        n_non_expert_layers: int = 0,
        mup_config: dict | None = None,
        use_peri_norm: bool = True,
        use_embedding_norm: bool = True,
        **kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize the MPTMuPConfig.

        Parameters
        ----------
        *args: Any
            Positional arguments for the MPTConfig.
        n_non_expert_layers : int = 0
            The number of non-expert layers in the model. They will be positioned at the
            beginning of the stack. Defaults to 0.
        mup_config : dict | None, optional
            muP configuration dictionary. If None, defaults to MuPConfig.
            mup_enabled: bool = False,
                Description: Enable muP support.
            mup_disable_attention_scaling: bool = True,
                Description: Disable attention scaling for muP,
                True by default since it performs better.
            mup_disable_hidden_lr_scaling: bool = False,
                Description: Disable hidden layer learning rate scaling for muP.
            mup_width_multiplier: float = 1,
                Description: Width multiplier for muP.
            mup_input_alpha: float = 1,
                Description: Alpha scaling factor for inputs.
            mup_output_alpha: float = 1,
                Description: Alpha scaling factor for output layers.
            completep_depth_alpha_enabled: bool = False,
                Description: Enable depth alpha scaling for CompleteP.
            completep_depth_multiplier: float = 1,
                Description: Multiplier for depth alpha scaling.
            completep_depth_alpha_exp: float = 1,
                Description: Exponent for depth alpha scaling.
            completep_eps_scaling_enabled: bool = True
                Description: Enable epsilon scaling for CompleteP.
        use_peri_norm : bool, optional
            Whether to use peri normalization in the model. Defaults to False.
        use_embedding_norm : bool, optional
            Whether to use embedding normalization with peri norm. Defaults to False.
        **kwargs: Any
            Keyword arguments for the MPTConfig.

        Raises
        ------
        ValueError
            If logit_scale is set and muP is enabled, as it is not supported.

        """
        obj_mup_config = (
            MuPConfig(**mup_config) if mup_config is not None else MuPConfig()
        )

        self.mup_config = asdict(
            obj_mup_config,
        )
        self.n_non_expert_layers = n_non_expert_layers
        self.use_embedding_norm = use_embedding_norm
        self.use_peri_norm = use_peri_norm
        super().__init__(*args, **kwargs)

        # Consistency checks for the `n_non_expert_layers` config
        if self.n_non_expert_layers < 0:
            msg = f"n_non_expert_layers ({self.n_non_expert_layers}) must be >= 0."
            raise ValueError(msg)
        if self.n_non_expert_layers > self.n_layers:
            msg = (
                f"n_non_expert_layers ({self.n_non_expert_layers}) must be <= n_layers"
                f" ({self.n_layers})."
            )
            raise ValueError(msg)
        if self.n_non_expert_layers > 0:
            log.warning(
                "n_non_expert_layers (%s) is > 0, some layers may be converted.",
                self.n_non_expert_layers,
            )

        # NOTE: this is the only place where we must use
        # dict key access instead of attribute access
        if (
            self.mup_config["mup_enabled"]
            and not self.mup_config["mup_disable_attention_scaling"]
        ):
            head_dim = self.d_model // self.n_heads

            self.attn_config = cast("dict[str, Any]", self.attn_config)  # type: ignore[reportUnnecessaryCast, has-type]
            # NOTE: Assumes that d_keys is head_dim
            if self.attn_config.get("softmax_scale") is None:
                self.attn_config = dict(self.attn_config)
                self.attn_config["softmax_scale"] = 1.0 / float(head_dim)

        if self.logit_scale is not None and obj_mup_config.mup_enabled:
            msg = (
                "logit_scale is not supported with muP enabled. "
                "Please set logit_scale to None in the config and modify"
                "The MuP output scaling factor to match the desired logit scale."
            )
            raise ValueError(msg)


class MPTMuPModel(MPTModel):
    """MPT model with muP support."""

    config_class = MPTMuPConfig

    def __init__(self, config: MPTMuPConfig) -> None:
        """Initialize the MPT model with muP support."""
        config._validate_config()  # noqa: SLF001
        # NOTE: these must be above the super call so we have access to them
        # in the super call
        self.mup_config = MuPConfig(**config.mup_config)  # type: ignore[reportArgumentType]
        self.init_config = ModelInitConfig(**config.init_config)
        super().__init__(config)
        if self.mup_config.mup_enabled:
            # Begin muP code: scale input embeddings
            self.wte = MuPSharedEmbedding(
                config.vocab_size,
                config.d_model,
                padding_idx=config.pad_token_id,
                device=config.init_device,
                scale=self.mup_config.mup_input_alpha,
                use_embedding_norm=config.use_embedding_norm and config.use_peri_norm,
                norm_type=config.norm_type,
                norm_eps=config.norm_eps,
            )
            if config.init_device != "meta":
                self.param_init_fn(self.wte)
            # End muP code
        self.config: MPTMuPConfig = cast("MPTMuPConfig", self.config)  # type: ignore[reportUnnecessaryCast]

        # NOTE: Hack to remove any additional biases
        # stolen from MPT
        if config.no_bias:
            for module in self.modules():
                if hasattr(module, "bias") and isinstance(module.bias, nn.Parameter):
                    log.debug("Removing bias from %s", module)
                    module.register_parameter("bias", None)

                # For transformer engine
                if hasattr(module, "use_bias") and module.use_bias is True:
                    log.debug("Setting use_bias=False for %s", module)
                    module.use_bias = False  # type: ignore[reportArgumentType]

    @property
    def block_class(self) -> type[MPTCompletePBlock]:
        """Return the block class for this model.

        Returns:
            type[MPTCompletePBlock]: The block class used in this model.

        """
        return MPTCompletePBlock

    def construct_blocks(self, config: MPTConfig) -> nn.ModuleList:
        """Construct the nn.ModuleList with the Transformer blocks.

        Parameters
        ----------
        config : MPTConfig
            The configuration object.

        Returns
        -------
        nn.ModuleList
            The list of Transformer blocks.

        """
        block_args = self.extract_block_args(config.to_dict())
        self.kv_cache_layers: set[int] = set()
        self.blocks_fuse_norm_attn_norm = block_args.get(
            "fuse_norm_attn_norm",
            False,
        )

        if config.block_overrides is not None:
            block_args_list = self._get_override_block_args_list(
                config,
                block_args,
            )
        else:
            block_args_list = [(block_args) for _ in range(config.n_layers)]

        return nn.ModuleList(
            [
                self.block_class(
                    device=config.init_device,
                    **(
                        convert_sigma_moe_config_to_dense(copy.deepcopy(block_args_i))
                        if i < self.config.n_non_expert_layers
                        else block_args_i
                    ),
                )
                for i, block_args_i in enumerate(block_args_list)
            ],
        )

    def extract_block_args(self, block_args: dict[str, Any]) -> dict[str, Any]:
        """Set the block args to include completeP.

        Parameters
        ----------
        block_args : dict[str, Any]
            The block arguments to extract and modify.

        Returns
        -------
        dict[str, Any]
            The modified block arguments with completeP settings.

        """
        block_args = super().extract_block_args(block_args)
        block_args["depth_multiplier"] = self.mup_config.completep_depth_multiplier
        block_args["depth_alpha_enabled"] = (
            self.mup_config.completep_depth_alpha_enabled
        )
        block_args["depth_alpha_exp"] = self.mup_config.completep_depth_alpha_exp
        block_args["use_peri_norm"] = self.config.use_peri_norm
        if self.config.ffn_config["ffn_type"] == "sigma_moe":
            block_args["ffn_config"]["mup_config"] = self.config.mup_config
            block_args["ffn_config"]["d_model"] = self.config.d_model
            block_args["ffn_config"]["init_std"] = (
                self.init_config.init_std
                if self.init_config.init_std is not None
                else 0.02
            )
        return block_args

    # Replace forward with method which scales output
    def forward(
        self,
        *args: Any,  # noqa: ANN401
        **kwargs: Any,  # noqa: ANN401
    ) -> BaseModelOutputWithPast:
        """Forward pass of the model.

        Scaled by the width multiplier and output alpha from the muP config.

        Parameters
        ----------
        args: Any
            Positional arguments for the forward method.
        kwargs: Any
            Keyword arguments for the forward method.

        Returns
        -------
        BaseModelOutputWithPast: The output of the model
        which includes the last hidden state.
            and optionally past key values.

        """
        outputs = super().forward(*args, **kwargs)

        if self.mup_config.mup_enabled and outputs.last_hidden_state is not None:
            # Begin muP code: scale output logits
            outputs.last_hidden_state = outputs.last_hidden_state * (  # noqa: PLR6104
                self.mup_config.mup_output_alpha / self.mup_config.mup_width_multiplier
            )  # type: ignore[reportAttributeAccessIssue]

            assert outputs.last_hidden_state is not None, (
                "last_hidden_state should not be None when muP is enabled."
            )

            # Modify all_hidden_states if it exists
            if outputs.hidden_states:
                # Replace tuple with a version where the last hidden state is scaled
                outputs.hidden_states = (
                    *tuple(outputs.hidden_states[:-1]),
                    outputs.last_hidden_state,
                )

            # End muP code

        return outputs


class MPTMuPForCausalLM(MPTForCausalLM):
    """MPT model for causal language modeling with muP support."""

    config_class = MPTMuPConfig

    def __init__(self, config: MPTMuPConfig) -> None:
        """Initialize the MPT model for causal language modeling with muP support.

        Parameters
        ----------
        config : MPTMuPConfig
            Configuration for the MPT model with muP support.

        """
        super().__init__(config)
        self.config: MPTMuPConfig = cast("MPTMuPConfig", self.config)  # type: ignore[reportUnnecessaryCast]
        self.mup_config = MuPConfig(**self.config.mup_config)
        self.init_config = ModelInitConfig(**self.config.init_config)

        self.transformer: MPTMuPModel = cast("MPTMuPModel", self.transformer)  # type: ignore[reportUnnecessaryCast]

        # Initialize the lm_head if it exists
        if config.init_device != "meta":
            self.reset_parameters()

    @property
    def backbone_model_class(self) -> type[MPTModel]:
        """Return the backbone model class for this composer model."""
        return MPTMuPModel

    def reset_parameters(self) -> None:
        """Reset the parameters of the model."""
        # Dispatch to param_init_fn
        for module in self.modules():
            self.param_init_fn(module)

    def param_init_fn(self, module: nn.Module) -> None:
        """Initialize the parameters of the module.

        This function complements the `reset_parameters` method above and
        it is strictly necessary
        when using FSDP with `init_device='meta'` parameters. Interestingly,
        using the standard wrapping strategy in `composer` as well as the standard
        wrapping of model classing that is adopted originally by the MPT models, this
        function is only called from this class and not the parents or the children
        modules. As such, this must completely initialize ALL the parameters whilst
        being careful on NOT double-initializing some leaf modules parameters.

        Parameters
        ----------
        ----------`
        module : nn.Module
            The module to initialize between the children of the DEPToEForCausalLM.

        """
        if not self.mup_config.mup_enabled:
            # Dispatch to the parent class initialization
            super().param_init_fn(module)
            return

        emb_init_std = self.init_config.emb_init_std

        base_std = (
            self.init_config.init_std if self.init_config.init_std is not None else 0.02
        )
        emb_init_std = emb_init_std if emb_init_std is not None else base_std

        width_mult = self.mup_config.mup_width_multiplier
        mup_std = base_std / (width_mult**0.5)

        with torch.no_grad():
            if (
                module is self.transformer.wte
                and isinstance(
                    module,
                    SharedEmbedding,
                )
            ) or (
                self.lm_head is not None
                and module is self.lm_head
                and isinstance(module, nn.Linear | nn.Embedding)
            ):
                # Initialize the shared embedding and/or lm head
                torch.nn.init.normal_(module.weight, mean=0.0, std=emb_init_std)
                if (
                    hasattr(module, "bias")
                    and module.bias is not None  # type: ignore[reportUnnecessaryComparison]
                    and isinstance(module.bias, nn.Parameter)
                ):  # type: ignore[reportUnnecessaryComparison]
                    torch.nn.init.zeros_(module.bias)
            elif isinstance(module, SigmaMoE):
                module.reset_parameters()
            elif isinstance(module, nn.Linear):
                # Initialize linear layers using muP std
                torch.nn.init.normal_(module.weight, mean=0.0, std=mup_std)
                if hasattr(module, "bias") and module.bias is not None:  # type: ignore[reportUnnecessaryComparison]
                    torch.nn.init.zeros_(module.bias)
            elif type(module) in norms.get_all().values():
                cast_module = cast("nn.LayerNorm", module)
                torch.nn.init.ones_(cast_module.weight)
                if hasattr(cast_module, "bias") and cast_module.bias is not None:  # type: ignore[reportUnnecessaryComparison]
                    torch.nn.init.zeros_(cast_module.bias)
            else:
                msg = f"Unsupported module type for muP initialization: {type(module)}"
                flwr_log(DEBUG, msg)

    def get_optimizer_param_groups(  # noqa: PLR0914, C901, PLR0915, PLR0912
        self,
        optimizer_config: dict[str, Any],
    ) -> tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]:
        """Create optimizer parameter groups for MPT models using muP/CompleteP scaling.

        This method categorizes model parameters into five groups:
        1. Embeddings
        2. Hidden LayerNorms (within transformer blocks)
        3. Hidden Weights (attn and ffn weights within blocks)
        4. Hidden Biases (attn and ffn biases within blocks)
        5. Final LayerNorm

        It then applies specific learning rate scaling and weight decay rules to each
        group based on the CompleteP and muP configuration. This approach avoids
        string-based name matching in favor of programmatic identification of layers.

        Parameters
        ----------
        optimizer_config : dict[str, Any]
            The base configuration for the optimizer, including 'lr' and 'weight_decay'.

        Returns
        -------
        tuple[list[dict[str, Any]], dict[str, Any]]
            A tuple containing:
            - A list of parameter group dictionaries for the optimizer.
            - The optimizer configuration, potentially modified with a scaled epsilon.

        Raises
        ------
        ValueError
            If unassigned parameters are found, indicating a bug in the parameter
            categorization logic.

        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

        """
        if not (
            self.mup_config.mup_enabled
            and not self.mup_config.mup_disable_hidden_lr_scaling
        ):
            return _extract_param_groups(self, optimizer_config), optimizer_config

        log.info("Creating optimizer parameter groups for muP/CompleteP.")

        # --- Initialize parameter lists based on CompleteP categories ---
        emb_params: list[nn.Parameter | torch.Tensor] = []
        hidden_ln_params: list[nn.Parameter] = []
        decay_lr_depth_width_scaling: list[nn.Parameter] = []
        hidden_bias_params: list[nn.Parameter] = []
        no_decay_params: list[nn.Parameter] = []

        # --- Programmatically categorize parameters ---

        emb_params.append(self.transformer.wte.weight)
        if hasattr(self.transformer, "wpe"):
            emb_params.append(self.transformer.wpe.weight)
        # Note: self.lm_head.weight is tied to wte.weight, so it's already included.

        # 1.5 Initial layer norm parameters
        if (
            hasattr(self.transformer.wte, "norm")
            and self.transformer.wte.norm is not None
        ):
            # NOTE: Not a hidden norm, should be treated
            # like the output norm
            no_decay_params.extend(p for p in self.transformer.wte.norm.parameters())

        # 2. Final LayerNorm parameters

        no_decay_params.extend(p for p in self.transformer.norm_f.parameters())

        # 3. Per-block parameters (hidden layers)

        for block in self.transformer.blocks:
            mup_block = cast("MPTCompletePBlock", block)

            # Categorize LayerNorms within the block
            norm_attr_names = [
                "norm_1",
                "norm_2",
                "post_attn_norm",
                "post_ffn_norm",
            ]
            for norm_attr_name in norm_attr_names:
                if hasattr(mup_block, norm_attr_name):
                    norm_module = getattr(mup_block, norm_attr_name)
                    if norm_module is not None:
                        hidden_ln_params.extend(p for p in norm_module.parameters())

            # Categorize Attention and FFN weights and biases
            linear_modules_to_process = []

            attn_module = mup_block.attn
            if hasattr(attn_module, "Wqkv"):
                linear_modules_to_process.append(attn_module.Wqkv)
            if hasattr(attn_module, "out_proj"):
                linear_modules_to_process.append(attn_module.out_proj)
            if hasattr(attn_module, "q_ln"):
                hidden_ln_params.extend(p for p in attn_module.q_ln.parameters())
            if hasattr(attn_module, "k_ln"):
                hidden_ln_params.extend(p for p in attn_module.k_ln.parameters())

            ffn_module = mup_block.ffn
            if hasattr(ffn_module, "up_proj"):
                linear_modules_to_process.append(ffn_module.up_proj)
            if hasattr(ffn_module, "down_proj"):
                linear_modules_to_process.append(ffn_module.down_proj)

            for linear_module in linear_modules_to_process:
                if linear_module is not None:
                    # Treat tensors with dim >= 2 as weights (e.g., for SigmaMoE)
                    if hasattr(linear_module, "weight"):
                        decay_lr_depth_width_scaling.append(linear_module.weight)
                    elif isinstance(linear_module, torch.Tensor):
                        decay_lr_depth_width_scaling.append(linear_module)  # type: ignore[reportArgumentType]

                    if (
                        hasattr(linear_module, "bias")
                        and linear_module.bias is not None  # type: ignore[reportAttributeAccessIssue]
                    ):
                        hidden_bias_params.append(linear_module.bias)  # type: ignore[reportAttributeAccessIssue]

            # Handle special SigmaMoE parameters
            if isinstance(
                mup_block.ffn,
                SigmaMoE,
            ) and hasattr(mup_block.ffn, "expert_sel"):
                if mup_block.ffn.config.mup_config.mup_scale_expert_sel:
                    decay_lr_depth_width_scaling.append(mup_block.ffn.expert_sel)
                else:
                    emb_params.append(mup_block.ffn.expert_sel)

        emb_params = list(dict.fromkeys(emb_params))  # Remove duplicates
        hidden_ln_params = list(dict.fromkeys(hidden_ln_params))  # Remove duplicates
        decay_lr_depth_width_scaling = list(dict.fromkeys(decay_lr_depth_width_scaling))
        hidden_bias_params = list(
            dict.fromkeys(hidden_bias_params),
        )  # Remove duplicates
        no_decay_params = list(dict.fromkeys(no_decay_params))  # Remove duplicates

        # --- Verify that all parameters have been assigned ---
        assigned_params = (
            set(emb_params)
            | set(hidden_ln_params)
            | set(decay_lr_depth_width_scaling)
            | set(hidden_bias_params)
            | set(no_decay_params)
        )
        all_params = set(self.parameters())
        unassigned_params = all_params - assigned_params  # type: ignore[reportOperatorIssue]

        # This part is crucial for debugging
        named_params = dict(self.named_parameters())
        # Create a reverse mapping from param id to name
        param_id_to_name = {id(p): name for name, p in named_params.items()}
        if unassigned_params:
            unassigned_names = [param_id_to_name[id(p)] for p in unassigned_params]
            flwr_log(
                WARNING,
                "Found %s unassigned parameters: %s",
                len(unassigned_params),
                unassigned_names,
            )
            msg = (
                f"Unassigned parameters found: {len(unassigned_params)}. "
                "This indicates a bug in the parameter categorization logic. "
                "Please check the model's parameter assignment logic."
            )
            raise ValueError(msg)

        # --- Calculate scaling factors ---
        base_lr = optimizer_config["lr"]
        weight_decay = optimizer_config.get("weight_decay", 0.0)

        width_lr_scaling = 1.0 / self.mup_config.mup_width_multiplier
        depth_lr_scaling = 1.0
        if self.mup_config.completep_depth_alpha_enabled:
            depth_lr_scaling = self.mup_config.completep_depth_multiplier ** (
                self.mup_config.completep_depth_alpha_exp - 1.0
            )

        # Optionally scale Adam's epsilon
        if self.mup_config.completep_eps_scaling_enabled:
            depth_eps_scaling = self.mup_config.completep_depth_multiplier ** (
                -1.0 * self.mup_config.completep_depth_alpha_exp
            )
            optimizer_config["eps"] = (
                optimizer_config.get("eps", 1e-8) * width_lr_scaling * depth_eps_scaling
            )

        # --- Create the final parameter groups ---
        param_groups = [
            {
                "params": emb_params,
                "weight_decay": weight_decay,
                "lr": base_lr,
            },
            {
                "params": hidden_ln_params,
                "weight_decay": 0.0,
                "lr": base_lr * depth_lr_scaling,
            },
            {
                "params": decay_lr_depth_width_scaling,
                "weight_decay": weight_decay / width_lr_scaling,
                "lr": base_lr * width_lr_scaling * depth_lr_scaling,
            },
            {
                "params": hidden_bias_params,
                "weight_decay": 0.0,
                "lr": base_lr * depth_lr_scaling,
            },
            {
                "params": no_decay_params,
                "weight_decay": 0.0,
                "lr": base_lr,
            },
        ]

        # Filter out any empty groups
        param_groups = [pg for pg in param_groups if pg["params"]]

        return param_groups, optimizer_config


@models.register_class("mpt_mup_causal_lm")
class ComposerMPTMuPCausalLM(ComposerMPTCausalLM):
    """Composer MPT model for causal language modeling with muP support."""

    @property
    def model_class(self) -> type[MPTMuPForCausalLM]:
        """Return the model class for this composer model."""
        return MPTMuPForCausalLM

    @property
    def config_class(self) -> type[MPTMuPConfig]:
        """Return the config class for this composer model."""
        return MPTMuPConfig

    def get_optimizer_param_groups(
        self,
        optimizer_config: dict[str, Any],
    ) -> tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]:
        """Extract the parameter groups for the optimizer when using MUP.

        Delegated to the underlying model's method.

        Parameters
        ----------
        optimizer_config : dict[str, Any]
            The configuration for the optimizer, which may include parameters like
            learning rate, weight decay, etc.

        Returns
        -------
        tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]
            A tuple containing the parameter groups and the optimizer configuration.

        """
        return cast(
            "MPTMuPForCausalLM",
            self.model,
        ).get_optimizer_param_groups(optimizer_config)

    @cached_property
    def n_active_params(self) -> int:
        """Count the number of active parameters in the model.

        This method calculates the number of active parameters in the model,
        which is useful for understanding the model's size and complexity.

        Returns
        -------
        int
            The number of active parameters in the model. This includes all parameters
            that are not part of the embedding layer, as embedding layers are typically
            lookup tables and do not contribute to the model's complexity in the same
            way as other layers.

        """
        if self.config.ffn_config["ffn_type"] in ffns_with_megablocks:
            params = megablocks_n_active_params(self)
        elif self.config.ffn_config["ffn_type"] == "sigma_moe":
            params = sigma_moe_n_active_params(self)
        else:
            params = sum(p.numel() for p in self.parameters())
        if not cast(
            "PretrainedConfig",
            cast("MPTMuPForCausalLM", self.model).transformer.config,
        ).tie_word_embeddings:
            # Embedding layers are lookup tables, therefore are not counted in the FLOP
            # computation
            params -= dtensor_safe_check_numel(
                cast("MPTMuPModel", self.model.transformer).wte.weight,
            )
        return params

    @cached_property
    def n_non_embedding_params(self) -> int:
        """Count the number of non-embedding parameters in the model.

        If the model uses tied word embeddings, this will exclude the embedding layer
        from the count. Otherwise, it will include all parameters except the embedding
        layer and the unembedding language modeling head.

        Returns
        -------
        int
            The number of non-embedding parameters in the model. This is useful for
            calculating the number of parameters that are not related to the embedding
            layer, which is often needed for FLOP calculations or model size
            estimations.

        """
        # Total number of parameters minus the embedding layer and the unembedding
        # language modeling head, if they exist.
        params = sum(p.numel() for p in self.parameters())
        if not cast(
            "PretrainedConfig",
            cast("MPTMuPForCausalLM", self.model).transformer.config,
        ).tie_word_embeddings:
            params -= dtensor_safe_check_numel(
                cast("MPTMuPModel", self.model.transformer).wte.weight,
            )
        if hasattr(self.model, "lm_head") and self.model.lm_head is not None:
            params -= dtensor_safe_check_numel(
                cast("torch.nn.Linear", self.model.lm_head).weight,
            )
        return params

    def forward(self, batch: MutableMapping) -> CausalLMOutputWithPast:
        """Forward the batch through the model.

        This method is used for both training and evaluation. It forwards the batch
        through the model and returns the output. The batch should contain keys like
        'input_ids', 'attention_mask', 'sequence_id', 'inputs_embeds', and
        'position_ids'. The method will log the keys of the batch being forwarded.

        Parameters
        ----------
        batch : MutableMapping
            The batch of data to forward through the model. It should contain keys like
            'input_ids', 'attention_mask', 'sequence_id', 'inputs_embeds', and
            'position_ids'.

        Returns
        -------
        CausalLMOutputWithPast
            The output of the model after forwarding the batch. This includes the last
            hidden state and optionally past key values.

        """
        # NOTE: Currently, the keys in the batch are ['input_ids', 'total_tokens',
        # 'loss_generating_tokens'] or ['input_ids', 'labels', 'total_tokens',
        # 'loss_generating_tokens'] for evaluation and training, respectively. Yes, I
        # didn't make a mistake in writing, the evaluation doesn't use labels.
        return self.model(
            input_ids=batch.get("input_ids", None),
            attention_mask=batch.get("attention_mask", None),
            sequence_id=batch.get("sequence_id", None),
            inputs_embeds=batch.get("inputs_embeds", None),
            position_ids=batch.get("position_ids", None),
        )

    def eval_forward(
        self,
        batch: MutableMapping,
        outputs: (
            dict[str, torch.Tensor]
            | tuple[torch.Tensor, ...]
            | CausalLMOutputWithPast
            | None
        ) = None,
    ) -> (
        list[str]
        | torch.Tensor
        | dict[str, torch.Tensor]
        | CausalLMOutputWithPast
        | tuple[torch.Tensor, ...]
    ):
        """Evaluate the model on a batch of data.

        This method handles both generation and evaluation modes. If the batch mode is
        'generate', it will generate text using the model's generate method. If the mode
        is 'icl_task' or if logits are used, it will compute the loss and return
        logits or the output of the forward pass.

        Parameters
        ----------
        batch : MutableMapping
            The batch of data to evaluate. It should contain keys like 'input_ids',
            'attention_mask', 'labels', and 'generation_kwargs' if applicable.
        outputs : (
            dict[str, torch.Tensor] |
            tuple[torch.Tensor, ...] |
            CausalLMOutputWithPast | None
        ), optional
            The outputs from the model's forward pass. If provided, it will be used
            instead of calling the forward method again. Defaults to None.

        Returns
        -------
        (
            list[str] |
            torch.Tensor |
            dict[str, torch.Tensor] |
            CausalLMOutputWithPast | tuple[torch.Tensor, ...]
        )
            The output of the evaluation. If the batch mode is 'generate', it will
            return generated text as a list of strings. If logits are used or the mode
            is 'icl_task', it will return logits or the output of the forward pass.

        Raises
        ------
        ValueError
            If the batch mode is 'generate' but no tokenizer is provided, or if the
            batch does not contain the required keys for evaluation.
        RuntimeError
            If the batch does not contain 'decoder_input_ids' for encoder-decoder
            models, or if the output from the forward pass is not in the expected
            format.

        """
        # NOTE: Currently, the keys in the batch are ['input_ids', 'labels',
        # 'total_tokens', 'loss_generating_tokens'] for evaluation.
        # If the batch mode is generate, we will generate a requested number of tokens
        # using the underlying model's generate function. Extra generation kwargs can be
        # passed in via the batch. Strings will be returned from eval_forward
        if batch.get("mode", None) == "generate":
            if self.tokenizer is None:
                msg = (
                    "Generation eval cannot be used without providing a tokenizer to"
                    " the model constructor."
                )
                raise ValueError(msg)

            self.labels = batch.pop("labels")
            generation = self.generate(
                batch["input_ids"],
                attention_mask=batch["attention_mask"],
                synced_gpus=dist.get_world_size() > 1,
                **batch.get("generation_kwargs", {}),
            )

            # don't remove prefix space to sentencepiece models
            if len(self.tokenizer.encode(" a", add_special_tokens=False)) == 1:
                return self.tokenizer.batch_decode(
                    generation[:, batch["input_ids"].shape[1] :],
                    skip_special_tokens=True,
                )
            return [
                " " + generation
                for generation in self.tokenizer.batch_decode(
                    generation[:, batch["input_ids"].shape[1] :],
                    skip_special_tokens=True,
                )
            ]

        if self.use_logits or batch.get("mode", None) == "icl_task":
            # Pop out labels first to avoid computing loss in the forward pass
            self.labels = batch.pop("labels")

            # HF encoder decoder models like T5 expect either decoder_input_ids or
            # labels, so we add decoder_input_ids to the batch if it is missing
            if self.config.is_encoder_decoder and "decoder_input_ids" not in batch:
                msg = (
                    "Encoder decoder models require that decoder_input_ids key"
                    " is present in the batch."
                )
                raise RuntimeError(msg)

            if self.shift_labels or batch.get("mode", None) == "icl_task":
                assert self.labels is not None
                # NOTE: HF CausalLM models internally shift labels before computing
                # loss, so we do the same here
                self.labels[:, :-1] = self.labels[:, 1:].clone()
                self.labels[:, -1] = -100

            output = outputs or self.forward(batch)

            if self.config.use_return_dict:
                assert isinstance(output, dict | CausalLMOutputWithPast), (
                    "Expected output to be a dict or CausalLMOutputWithPast when"
                    f" use_return_dict is True, but got {type(output)}"
                )
                output = output["logits"]
                assert isinstance(output, torch.Tensor), (
                    "Expected output['logits'] to be a torch.Tensor, "
                    f"but got {type(output)}"
                )
            else:
                assert isinstance(output, tuple), (
                    "Expected output to be a tuple when use_return_dict is False, "
                    f"but got {type(output)}"
                )
                # if loss was computed (cached outputs from forward), loss is at index 0
                # and logits are at index 1
                # if loss was not computed (no cached outputs during eval), loss is not
                # present and logits are at index 0
                output = output[1] if len(output[0].shape) == 0 else output[0]

            # NOTE: We remove any classes dimension in case of a single class
            if output.size(-1) == 1:
                output = output.squeeze(-1)
        else:
            output = outputs or self.forward(batch)

        return output
