"""Module implements a DEPToE Causal Language Model for LLM Foundry."""

from __future__ import annotations

import logging
import math
import warnings
from logging import DEBUG
from typing import (
    TYPE_CHECKING,
    Any,
    cast,
)

import torch
import torch.nn.functional as F  # noqa: N812
from flwr.common import log as flwr_log
from llmfoundry.models.layers import LPLayerNorm
from llmfoundry.models.layers.custom_embedding import SharedEmbedding
from llmfoundry.models.utils.act_ckpt import (
    build_act_ckpt_mod_to_blocks,
    check_mapping_blocks_overlap,
)
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.deptoe import DEPToEModel, DEPToEPreTrainedModel
from mt_dao.models.deptoe.deptoe_block import DEPToEBlock
from mt_dao.models.deptoe.mup_config import MuPConfig
from mt_dao.models.ffns.sigma_moe import SigmaMoE
from mt_dao.models.switch_head.switch_head_core import SwitchHead
from mt_dao.models.utils import (
    CROSS_ENTROPY_IGNORE_INDEX,
)

if TYPE_CHECKING:
    from collections.abc import Iterable

    from transformers.cache_utils import Cache

    from mt_dao.models.deptoe.configuration_deptoe import DEPToEConfig

log = logging.getLogger(__name__)


class DEPToEForCausalLM(DEPToEPreTrainedModel):
    """DEPToEForCausalLM is a causal language model based on the DEPT architecture."""

    def __init__(self, config: DEPToEConfig) -> None:
        """Initialize the DEPToEForCausalLM model.

        Parameters
        ----------
        config : DEPToEConfig
            The configuration for the DEPToE model, which includes parameters such as
            `vocab_size`, `d_model`, `n_layers`, `n_heads`, and `tie_word_embeddings`.

        Raises
        ------
        ValueError
            If the configuration is invalid or if the `tie_word_embeddings` is set to
            True but the `lm_head` is not None.

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

        self.transformer: DEPToEModel = self.backbone_model_class(config)

        self.lm_head = None
        if not config.tie_word_embeddings:
            self.lm_head = nn.Linear(
                config.d_model,
                config.vocab_size,
                bias=False,
                device=config.init_device,
            )
            setattr(self.lm_head, "_fsdp_wrap", True)  # noqa: B010

        for child in self.transformer.children():
            if isinstance(child, torch.nn.ModuleList):
                setattr(child, "_fsdp_wrap", False)  # noqa: B010
            if isinstance(child, torch.nn.Module):
                setattr(child, "_fsdp_wrap", True)  # noqa: B010

        # NOTE: This enables scaling output logits; similar to a softmax "temperature".
        # PaLM paper uses scale 1/sqrt(config.d_model)
        self.logit_scale = None
        if config.logit_scale is not None:
            logit_scale = config.logit_scale
            if isinstance(logit_scale, str):
                if logit_scale == "inv_sqrt_d_model":
                    logit_scale = 1 / math.sqrt(config.d_model)
                else:
                    msg = f"{logit_scale=} is not recognized as an option;"
                    msg += " use numeric value or 'inv_sqrt_d_model'."
                    raise ValueError(msg)
            self.logit_scale = logit_scale
        self.final_logit_softcapping = config.final_logit_softcapping

        self.entropy_reg = config.entropy_reg
        self.att_entropy_reg = config.att_entropy_reg
        # Trigger initialization of all parameters if `init_device` is not "meta".
        if config.init_device != "meta":
            log.info(
                'We recommend using config.init_device="meta" with Composer + FSDP'
                " for faster initialization.",
            )
            self.reset_parameters()

    def reset_parameters(self) -> None:
        """Reset the parameters of the model.

        This method initializes the parameters of the model using a scale factor
        based on the number of repeats and the number of blocks. It also initializes
        the weights and biases of the LayerNorm layers to ones and zeros, respectively.
        The scale factor is calculated as the square root of 2 divided by the product
        of the number of repeats and the number of blocks.
        This method is called without gradients to ensure that the parameters are
        initialized correctly without affecting the gradient computation.

        """
        # 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 strincly
        necesssary 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.

        """
        with torch.no_grad():
            if isinstance(module, LPLayerNorm | torch.nn.LayerNorm):
                torch.nn.init.ones_(module.weight)
                if module.bias is not None:  # type: ignore[reportUnnecessaryComparison]
                    torch.nn.init.zeros_(module.bias)
            elif isinstance(module, (SwitchHead | SigmaMoE)):
                module.reset_parameters()
            # NOTE: double checking the lm_head to avoid potentially interfering
            # with any other nn.Linear layers
            elif isinstance(module, (SharedEmbedding | torch.nn.Embedding)) or (
                self.lm_head is not None
                and module is self.lm_head
                and isinstance(module, nn.Linear)
            ):
                DEPToEModel.init_trainable_embeddings(
                    module,
                    self.config.init_config,
                    mup_enabled=self.mup_config.mup_enabled,
                )

    @property
    def backbone_model_class(self) -> type[DEPToEModel]:
        """Return the class of the backbone model."""
        return DEPToEModel

    # TODO(anon): Add DEPT embeddings to the following methods
    def get_input_embeddings(self) -> SharedEmbedding | nn.Embedding:
        """Get the input embeddings of the model.

        Returns
        -------
        SharedEmbedding | nn.Embedding
            The input embeddings of the model, which can be either a SharedEmbedding or
            nn.Embedding instance.

        """
        return self.transformer.get_input_embeddings()

    def set_input_embeddings(  # type: ignore[reportIncompatibleMethodOverride]
        self,
        value: SharedEmbedding | nn.Embedding,
    ) -> None:
        """Set the input embeddings of the model.

        Parameters
        ----------
        value : SharedEmbedding | nn.Embedding
            The new input embeddings to set for the model. This can be either a
            SharedEmbedding or nn.Embedding instance.

        """
        self.transformer.set_input_embeddings(value)

    def get_output_embeddings(
        self,
    ) -> SharedEmbedding | nn.Embedding | nn.Linear:
        """Get the output embeddings of the model.

        If the model has a language model head (`lm_head`), it returns that.
        Otherwise, it returns the input embeddings of the transformer.

        Returns
        -------
        SharedEmbedding | nn.Embedding | nn.Linear
            The output embeddings of the model, which can be either a SharedEmbedding,
            nn.Embedding, or nn.Linear instance.

        """
        if self.lm_head is not None:
            return self.lm_head
        return self.transformer.get_input_embeddings()

    def set_output_embeddings(
        self,
        new_embeddings: SharedEmbedding | nn.Embedding | nn.Linear,
    ) -> None:
        """Set the output embeddings of the model.

        Parameters
        ----------
        new_embeddings : SharedEmbedding | nn.Embedding | nn.Linear
            The new output embeddings to set for the model. This can be either a
            SharedEmbedding, nn.Embedding, or nn.Linear instance.

        Raises
        ------
        ValueError
            If `new_embeddings` is not an instance of SharedEmbedding or nn.Embedding
            and the model does not have a language model head (`lm_head`).

        """
        if self.lm_head is not None:
            self.lm_head = new_embeddings
        else:
            if not isinstance(new_embeddings, (SharedEmbedding | nn.Embedding)):
                msg = "new_embeddings must be an instance of SharedEmbedding or "
                msg += f" nn.Embedding, but got {type(new_embeddings)}."
                raise ValueError(msg)
            warnings.warn(
                "Using `set_output_embeddings` to set the embedding layer of "
                "DEPToEForCausalLM with tied weights. Given weights are tied, "
                "using `set_input_embeddings` is recommended over using "
                "`set_output_embeddings`.",
                stacklevel=2,
            )
            self.transformer.set_input_embeddings(new_embeddings)

    def tie_weights(self) -> None:
        """Tie the weights of the input and output embeddings.

        This method sets the `lm_head` to None if `tie_word_embeddings` is True in the
        configuration. This is to ensure that the input and output embeddings share the
        same weights, which is a common practice in language models to reduce the number
        of parameters and improve training efficiency.

        """
        if getattr(self.config, "tie_word_embeddings", True):
            self.lm_head = None

    def set_decoder(self, decoder: DEPToEModel) -> None:
        """Set the decoder of the model.

        Parameters
        ----------
        decoder : DEPToEModel
            The new decoder to set for the model. This should be an instance of
            DEPToEModel, which is the backbone model class for DEPToEForCausalLM.

        """
        self.transformer = decoder

    def get_decoder(self) -> DEPToEModel:
        """Get the decoder of the model.

        Returns
        -------
        DEPToEModel
            The decoder of the model, which is an instance of DEPToEModel. This is the
            backbone model class for DEPToEForCausalLM.

        """
        return self.transformer

    def forward(  # noqa: PLR0913
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.ByteTensor | None = None,
        labels: torch.LongTensor | None = None,
        *,
        return_dict: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        use_cache: bool | None = None,
    ) -> CausalLMOutputWithPast:
        """Forward pass of the DEPToEForCausalLM model.

        Parameters
        ----------
        input_ids : torch.LongTensor, optional
            Input token IDs. If provided, `inputs_embeds` must not be set.
        attention_mask : torch.ByteTensor, optional
            Attention mask to avoid performing attention on padding token indices.
        labels : torch.LongTensor, optional
            Labels for computing the loss. If provided, the model will compute the loss.
        return_dict : bool, optional
            Whether to return a `CausalLMOutputWithPast` object or a plain tuple.
            Defaults to the value in the model's config.
        output_attentions : bool, optional
            Whether to return attention weights. Defaults to the value in the model's
            config.
        output_hidden_states : bool, optional
            Whether to return hidden states. Defaults to the value in the model's
            config.
        use_cache : bool, optional
            Whether to use caching for the past key and value states. Defaults to the
            value in the model's config.

        Returns
        -------
        CausalLMOutputWithPast
            The output of the model, which includes the loss (if `labels` are provided),
            logits, past key and value states, hidden states, and attention weights.

        """
        outputs = cast(
            "BaseModelOutputWithPast",
            self.transformer(
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_dict=return_dict,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                use_cache=use_cache,
            ),
        )

        if self.lm_head is not None:
            logits = cast("torch.Tensor", self.lm_head(outputs.last_hidden_state))
        else:
            # NOTE: We must move outputs to same device as weights for token embedding.
            # This is needed to support HF `device_map`
            out = outputs.last_hidden_state
            assert out is not None, "outputs.last_hidden_state must not be None"
            out = out.to(self.transformer.wte.weight.device)
            logits = cast("torch.Tensor", self.transformer.wte(out, unembed=True))

        if self.logit_scale is not None:
            if self.logit_scale == 0:
                warnings.warn(
                    f"Multiplying logits by {self.logit_scale=}. This will produce"
                    " uniform (uninformative) outputs.",
                    stacklevel=2,
                )
            logits *= self.logit_scale

        if self.final_logit_softcapping is not None:
            logits = self.final_logit_softcapping * torch.tanh(
                logits / self.final_logit_softcapping,
            )

        loss = torch.zeros(1, device=logits.device, dtype=logits.dtype)
        if labels is not None:
            shifted_labels = torch.roll(labels, shifts=-1)
            shifted_labels[:, -1] = CROSS_ENTROPY_IGNORE_INDEX
            loss += cast(
                "torch.FloatTensor",
                F.cross_entropy(
                    logits.view(-1, logits.size(-1)),
                    shifted_labels.to(logits.device).view(-1),
                ),
            )

        # Collect regularizaiton losses across the layers
        if self.transformer.n_repeats > 1:
            for layer in self.children():
                if isinstance(layer, SigmaMoE):
                    loss += self.entropy_reg * layer.get_reg_loss()
                elif isinstance(layer, SwitchHead):
                    loss += self.att_entropy_reg * layer.get_reg_loss()

        return CausalLMOutputWithPast(
            loss=cast("torch.FloatTensor", loss),
            logits=cast("torch.FloatTensor", logits),
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @staticmethod
    def fsdp_wrap_fn(module: nn.Module) -> bool:
        """FSDP wrap function for DEPToE models.

        Parameters
        ----------
        module : nn.Module
            The module to check for FSDP wrapping eligibility.

        Returns
        -------
        bool
            True if the module should be wrapped with FSDP, False otherwise.

        """
        if hasattr(module, "_fsdp_kwargs_dict"):
            return bool(module._fsdp_kwargs_dict)  # noqa: SLF001
        return isinstance(module, DEPToEBlock)

    # Activation Checkpointing
    def activation_checkpointing_fn(self, module: nn.Module) -> bool:
        """Implement DEPToE activation checkpointing (act ckpt) function.

        When `activation_checkpointing` in fsdp_config is set to true, this function
        will be called on all the modules in the FSDP wrapped model and determine
        whether a given module should be activation checkpointed. It checks the
        checkpointing target (`activation_checkpointing_target` in `model`) which can be
        specified as below:
            1. null (or no such field): The whole MPTBlock will be activation
                checkpointed on all layers
            2. a list of modules to act ckpt on all layers, e.g.,
                activation_checkpointing_target:
                    - grouped_query_attention
                    - mptmlp
            3. a dictionary of module name with target_blocks, e.g.,
                activation_checkpointing_target:
                    {
                            "mptblock": target_blocks_1,
                            "grouped_query_attention": target_blocks_2
                    }
                target_blocks (target_blocks_1, target_blocks_2 above) can be:
                - a single integer n: the first n transformer block will be activation
                    checkpointed
                - a string of first-n, middle-m, last-k, range-i-j: the first n, the
                    middle m,  the last k, or the range [i, j) layers will be activation
                    checkpointed. E.g, 'first-2, last-2' means the first 2 and last 2
                    transformer blocks will be activation checkpointed middle-m is range
                    [start, end) where ``start = max(max_block_idx // 2 - m // 2, 0),
                    end = min(start + m, max_block_idx + 1)``
                - a list of integers corresponds to the list of transformer block ids,
                    e.g., [2] means the second transformer block will be activation
                    checkpointed. [2, 3] means the second and third transformer blocks
                    will be activation checkpointed
                - a list of mixed integers and strings of first-n, middle-m, last-k,
                range-i-j

            An example in yaml config file:
                fsdp_config:
                    activation_checkpointing: true
                model:
                    activation_checkpointing_target:
                        {
                            "mptblock": 'first-5',
                            "grouped_query_attention": 'last-35'
                        }

        Parameters
        ----------
        module : nn.Module
            The module to check for activation checkpointing eligibility.

        Returns
        -------
        bool
            True if the module should be activation checkpointed, False otherwise.

        """
        if not hasattr(module, "block_idx"):
            log.debug(
                "%s cannot be activation checkpointed. Only transformer block or its"
                " submodules are eligible for activation checkpointing.",
                module.__class__.__name__,
            )
            return False

        act_ckpt_target = getattr(
            self.config,
            "activation_checkpointing_target",
            None,
        )
        act_ckpt_mod_to_blocks = build_act_ckpt_mod_to_blocks(
            act_ckpt_target,
            DEPToEBlock,
            int(cast("torch.Tensor", module.max_block_idx).item()),
        )

        check_mapping_blocks_overlap(
            act_ckpt_mod_to_blocks,
            int(cast("torch.Tensor", module.max_block_idx).item()),
        )

        for k in act_ckpt_mod_to_blocks:
            if isinstance(module, k):
                blocks = act_ckpt_mod_to_blocks[k]
                return True if blocks == -1 else module.block_idx in blocks

        return False

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        past_key_values: Cache | None = None,
        attention_mask: torch.LongTensor | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        cache_position: torch.LongTensor | None = None,  # noqa: ARG002
        **kwargs: Any,  # noqa: ANN401
    ) -> dict[str, Any]:
        """Prepare inputs for generation.

        This method prepares the inputs for the generation process, including handling
        the attention mask, sequence ID, and past key values. It ensures that the input
        IDs and attention mask are correctly formatted and that the model is ready for
        generation.

        Parameters
        ----------
        input_ids : torch.Tensor
            The input token IDs for the generation process.
        past_key_values : Cache, optional
            The past key and value states for the attention layers, used for caching.
            Defaults to None.
        attention_mask : torch.LongTensor, optional
            The attention mask to avoid performing attention on padding token indices.
            If not provided, it will be derived from `input_ids`. Defaults to None.
        inputs_embeds : torch.Tensor, optional
            The input embeddings to use instead of `input_ids`. If provided, it will be
            used in the first generation step only. Defaults to None.
        cache_position : torch.LongTensor, optional
            The position in the cache for the current generation step. This is used to
            determine which tokens to consider for the current generation step. Defaults
            to None.
        **kwargs : Any
            Additional keyword arguments that may include `attention_mask`, `use_cache`,
            and other parameters.

        Returns
        -------
        dict[str, Any]
            A dictionary containing the prepared inputs for the generation process,
            including `input_ids`, `attention_mask`, `sequence_id`, and
            `past_key_values`.

        Raises
        ------
        NotImplementedError
            If the attention mask indicates right padding, which is not supported by the
            model.

        """
        if (
            attention_mask is not None
            and attention_mask[:, -1].sum() != attention_mask.shape[0]
        ):
            msg = f"{type(self.backbone_model_class)} does not support generation with"
            msg += " right padding."
            raise NotImplementedError(msg)

        if self.transformer.attn_uses_sequence_id and self.training:
            sequence_id = torch.zeros_like(input_ids[:1])
        else:
            sequence_id = None

        # only last token for inputs_ids if past is defined in kwargs
        if past_key_values is not None:
            input_ids = input_ids[:, -1].unsqueeze(-1)

        model_inputs = {
            "attention_mask": attention_mask,
            "sequence_id": sequence_id,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", True),
        }

        # NOTE: If `inputs_embeds` are passed, we only want to use them in the
        # 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}
        return model_inputs

    def _reorder_cache(  # noqa: PLR6301
        self,
        past_key_values: list[tuple[torch.Tensor, torch.Tensor]],
        beam_idx: torch.LongTensor,
    ) -> list[tuple[torch.Tensor, ...]]:
        """Reorder the `past_key_values` cache for beam search.

        This function is used to re-order the `past_key_values` cache if
        [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called.
        This is required to match `past_key_values` with the beam_idx at every
        generation step. Used by HuggingFace generate when using beam search with
        kv-caching. See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
        for an example in transformers.

        Parameters
        ----------
        past_key_values : list[tuple[torch.Tensor, torch.Tensor]]
            The past key and value states for the attention layers, which is a list of
            tuples where each tuple contains the key and value tensors for a specific
            layer.
        beam_idx : torch.LongTensor
            The indices of the beams to reorder the past key values. This tensor should
            have the same length as the batch size and contains the indices of the beams
            to keep for each sample in the batch.

        Returns
        -------
        list[tuple[torch.Tensor, ...]]
            A reordered list of past key values, where each tuple contains the key and
            value tensors for a specific layer, reordered according to the `beam_idx`.
            Each tensor in the tuple is indexed by the `beam_idx`, allowing for
            efficient beam search with kv-caching.

        """
        reordered_past = []
        for layer_past in past_key_values:
            reordered_past += [
                tuple(
                    past_state.index_select(0, beam_idx) for past_state in layer_past
                ),
            ]
        return reordered_past

    def get_optimizer_param_groups(  # noqa: C901, PLR0915, PLR0914, PLR0912
        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.

        This method extracts the parameter groups for the optimizer when using
        MUP (Maximal Update Parametrisation). It separates the parameters into
        different groups based on their types and applies different learning rates
        and weight decay values to each group based on the model width and depth
        multipliers as well as the depth alpha exponent.

        If MUP is not enabled, it calls the `_extract_param_groups` function to
        get the default parameter groups.

        Parameters
        ----------
        optimizer_config : dict[str, Any]
            The configuration for the optimizer, which includes parameters like
            learning rate, weight decay, and epsilon.

        Returns
        -------
        tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]
            A tuple containing the parameter groups and the modified optimizer
            configuration. The parameter groups are separated into different
            categories such as embedding parameters, hidden weight parameters,
            hidden bias parameters, and layer normalization parameters.

        Raises
        ------
        ValueError
            If any model parameters are left unassigned after categorization.

        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
        ):
            flwr_log(
                DEBUG,
                "Using default optimizer param groups for DEPToEModel.",
            )
            return _extract_param_groups(self, optimizer_config), optimizer_config

        flwr_log(DEBUG, "Returning optimizer param groups for MuP.")

        weight_decay: float = optimizer_config.get("weight_decay", 0.0)
        base_lr: float = optimizer_config["lr"]

        # Initialize parameter lists with specified names
        unaltered_params: list[torch.nn.Parameter | torch.Tensor] = []
        hidden_ln_params: list[torch.nn.Parameter] = []
        decay_lr_depth_width_scaling: list[torch.nn.Parameter] = []
        hidden_bias_params: list[torch.nn.Parameter] = []
        no_decay_params: list[torch.nn.Parameter] = []

        # Handle embedding and language model head parameters
        unaltered_params.append(self.transformer.wte.weight)
        if self.lm_head is not None:
            unaltered_params.append(self.lm_head.weight)

        # Collect "outer" parameters (final norms) into no_decay_params
        no_decay_params.extend(
            p
            for p in [self.transformer.norm_f.weight, self.transformer.norm_f.bias]
            if p is not None
        )
        if hasattr(self, "embedding_peri_ln"):
            no_decay_params.extend(
                p
                for p in [
                    self.transformer.embedding_peri_ln.weight,
                    self.transformer.embedding_peri_ln.bias,
                ]
                if p is not None  # type: ignore[reportUnnecessaryComparison]
            )

        # Process each transformer block
        for block in self.transformer.blocks:
            cast_block = cast("DEPToEBlock", block)

            # Collect hidden LayerNorm parameters
            hidden_ln_params.extend(
                p
                for p in [
                    cast_block.switch_head.input_ln.weight,
                    cast_block.switch_head.input_ln.bias,
                    # NOTE: wrong now since we moved the ln
                    # for MPT, kept for reference
                    cast_block.sigma_moe.input_ln.weight,  # type: ignore[reportAttributeAccessIssue]
                    cast_block.sigma_moe.input_ln.bias,  # type: ignore[reportAttributeAccessIssue]
                ]
                if p is not None  # type: ignore[reportUnnecessaryComparison]
            )
            if cast_block.switch_head.use_peri_ln:
                hidden_ln_params.extend(
                    p
                    for p in [
                        cast_block.switch_head.peri_ln_out.weight,
                        cast_block.switch_head.peri_ln_out.bias,
                    ]
                    if p is not None  # type: ignore[reportUnnecessaryComparison]
                )
            if cast_block.sigma_moe.use_peri_ln:
                hidden_ln_params.extend(
                    p
                    for p in [
                        cast_block.sigma_moe.peri_ln_out.weight,  # type: ignore[reportAttributeAccessIssue]
                        cast_block.sigma_moe.peri_ln_out.bias,  # type: ignore[reportAttributeAccessIssue]
                    ]
                    if p is not None  # type: ignore[reportUnnecessaryComparison]
                )

            # Collect hidden weights
            decay_lr_depth_width_scaling.extend(
                p
                for p in [
                    cast_block.switch_head.k.weight,
                    cast_block.switch_head.q.weight,
                    cast_block.switch_head.v
                    if isinstance(cast_block.switch_head.v, nn.Parameter)
                    else cast_block.switch_head.v.weight,
                    cast_block.switch_head.o
                    if isinstance(cast_block.switch_head.o, nn.Parameter)
                    else cast_block.switch_head.o.weight,
                    cast_block.sigma_moe.up_proj
                    if isinstance(cast_block.sigma_moe.up_proj, nn.Parameter)
                    else cast_block.sigma_moe.up_proj.weight,
                    cast_block.sigma_moe.down_proj
                    if isinstance(cast_block.sigma_moe.down_proj, nn.Parameter)
                    else cast_block.sigma_moe.down_proj.weight,
                ]
                if p is not None
            )

            # Collect hidden attention biases
            hidden_bias_params.extend(
                p
                for p in [cast_block.switch_head.k.bias, cast_block.switch_head.q.bias]
                if p is not None  # type: ignore[reportUnnecessaryComparison]
            )

            # Handle expert selection parameters

            if hasattr(cast_block.switch_head, "sel_v"):
                if cast_block.switch_head.config.mup_config.mup_scale_expert_sel:
                    decay_lr_depth_width_scaling.append(cast_block.switch_head.sel_v)
                else:
                    unaltered_params.append(cast_block.switch_head.sel_v)

            if hasattr(cast_block.switch_head, "sel_o"):
                if cast_block.switch_head.config.mup_config.mup_scale_expert_sel:
                    decay_lr_depth_width_scaling.append(cast_block.switch_head.sel_o)
                else:
                    unaltered_params.append(cast_block.switch_head.sel_o)

            if hasattr(
                cast_block.sigma_moe,
                "expert_sel",
            ):
                if cast_block.sigma_moe.config.mup_config.mup_scale_expert_sel:
                    decay_lr_depth_width_scaling.append(cast_block.sigma_moe.expert_sel)
                else:
                    unaltered_params.append(cast_block.sigma_moe.expert_sel)

        # De-duplicate all parameter lists
        param_lists = [
            unaltered_params,
            hidden_ln_params,
            decay_lr_depth_width_scaling,
            hidden_bias_params,
            no_decay_params,
        ]
        (
            unaltered_params,
            hidden_ln_params,
            decay_lr_depth_width_scaling,
            hidden_bias_params,
            no_decay_params,
        ) = [list(dict.fromkeys(p_list)) for p_list in param_lists]

        # Calculate scaling factors
        width_lr_scaling = 1.0 / self.mup_config.mup_width_multiplier
        depth_lr_scaling = (
            self.mup_config.completep_depth_multiplier
            ** (self.mup_config.completep_depth_alpha_exp - 1.0)
            if self.mup_config.completep_depth_alpha_enabled
            else 1.0
        )
        if self.mup_config.completep_eps_scaling_enabled:
            optimizer_config["eps"] *= width_lr_scaling * (
                self.mup_config.completep_depth_multiplier
                ** (-1.0 * self.mup_config.completep_depth_alpha_exp)
            )

        # Create a reverse mapping from param ID to name for logging and debugging
        param_id_to_name = {id(p): name for name, p in self.named_parameters()}

        # Verify all parameters are assigned and raise error if not
        assigned_params = {p for p_list in param_lists for p in p_list}
        all_model_params = set(self.parameters())
        unassigned_params = all_model_params - assigned_params

        if unassigned_params:
            unassigned_names = [
                param_id_to_name.get(id(p), "Unknown") for p in unassigned_params
            ]
            msg = (
                f"Found {len(unassigned_params)} unassigned params: {unassigned_names}"
            )
            raise ValueError(msg)

        def get_names(p_list: list[nn.Parameter]) -> list[str]:
            return [param_id_to_name.get(id(p), "Unknown") for p in p_list]

        flwr_log(
            DEBUG,
            "Group 'unaltered_params' (%d params): lr=%.2e, wd=%.2e, names=%s",
            len(unaltered_params),
            base_lr,
            weight_decay,
            get_names(unaltered_params),  # type: ignore[reportArgumentType]
        )
        flwr_log(
            DEBUG,
            "Group 'hidden_ln_params' (%d params): lr=%.2e, wd=0.0, names=%s",
            len(hidden_ln_params),
            base_lr * depth_lr_scaling,
            get_names(hidden_ln_params),
        )
        flwr_log(
            DEBUG,
            "Group 'projections' (%d params): lr=%.2e, wd=%.2e, names=%s",
            len(decay_lr_depth_width_scaling),
            base_lr * width_lr_scaling * depth_lr_scaling,
            weight_decay / width_lr_scaling,
            get_names(decay_lr_depth_width_scaling),
        )
        flwr_log(
            DEBUG,
            "Group 'hidden_bias_params' (%d params): lr=%.2e, wd=0.0, names=%s",
            len(hidden_bias_params),
            base_lr * depth_lr_scaling,
            get_names(hidden_bias_params),
        )
        flwr_log(
            DEBUG,
            "Group 'no_decay_params' (%d params): lr=%.2e, wd=0.0, names=%s",
            len(no_decay_params),
            base_lr,
            get_names(no_decay_params),
        )

        # Create parameter groups in the specified order
        param_groups = [
            {"params": unaltered_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 before returning
        param_groups = [pg for pg in param_groups if pg["params"]]

        return param_groups, optimizer_config
