models/micro_moe_llama.py

from typing import Optional, Tuple, Union, List, Callable
import logging 
import yaml

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist

from transformers import LlamaConfig, AutoModelForCausalLM, AutoConfig
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
from transformers.models.llama.modeling_llama import (
    LlamaRotaryEmbedding, 
    LlamaRMSNorm, 
    LlamaMLP,
    LlamaAttention,
    LlamaForCausalLM,
    LlamaPreTrainedModel, 
    GenerationMixin,
    apply_rotary_pos_emb,
    eager_attention_forward,

)
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.cache_utils import Cache, StaticCache, DynamicCache
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.processing_utils import Unpack
from transformers.utils import is_torchdynamo_compiling
from transformers.activations import ACT2FN
from models.modules import CausalLMOutputWithPast

logger = logging.getLogger(__name__)

def keep_alive_zero(model):
    z = 0.0
    for p in model.parameters():
        if p.requires_grad:
            # one scalar per param to avoid heavy sums
            z = z + (p.view(-1)[0] * 0.0)
    return z

class MiCRoLlamaMoEConfig(LlamaConfig):
    model_type = "micro_llama_moe"
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.num_experts = kwargs.get("num_experts", 4)
        self.use_router = kwargs.get("use_router", True)
        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 2)
        self.jitter_noise = kwargs.get("jitter_noise", 0.0)
        self.loss_method = kwargs.get("loss_method", "all")
        self.config_path = kwargs.get("config_path", None)

def _prepare_4d_causal_attention_mask_with_cache_position(
    attention_mask: torch.Tensor,
    sequence_length: int,
    target_length: int,
    dtype: torch.dtype,
    device: torch.device,
    min_dtype: float,
    cache_position: torch.Tensor,
    batch_size: int,
):
    """
    Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
    `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

    Args:
        attention_mask (`torch.Tensor`):
            A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
        sequence_length (`int`):
            The sequence length being processed.
        target_length (`int`):
            The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
        dtype (`torch.dtype`):
            The dtype to use for the 4D attention mask.
        device (`torch.device`):
            The device to plcae the 4D attention mask on.
        min_dtype (`float`):
            The minimum value representable with the dtype `dtype`.
        cache_position (`torch.Tensor`):
            Indices depicting the position of the input sequence tokens in the sequence.
        batch_size (`torch.Tensor`):
            Batch size.
    """
    if attention_mask is not None and attention_mask.dim() == 4:
        # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
        causal_mask = attention_mask
    else:
        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
        if sequence_length != 1:
            causal_mask = torch.triu(causal_mask, diagonal=1)
        causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
            mask_length = attention_mask.shape[-1]
            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
            padding_mask = padding_mask == 0
            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
                padding_mask, min_dtype
            )

    return causal_mask

class DummyModule(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x
    
class LlamaSparseMiCRoMoEBlock(nn.Module):
    """
    This implementation is
    strictly equivalent to standard MoE with full capacity (no
    dropped tokens). It's faster since it formulates MoE operations
    in terms of block-sparse operations to accommodate imbalanced
    assignments of tokens to experts, whereas standard MoE either
    (1) drop tokens at the cost of reduced performance or (2) set
    capacity factor to number of experts and thus waste computation
    and memory on padding.
    """

    def __init__(self, config):
        super().__init__()
        self.hidden_dim = config.hidden_size
        self.ffn_dim = config.intermediate_size
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_tok
        self.use_router = config.use_router
        self.ablate = config.ablate

        # gating
        self.gate = nn.Sequential(
            nn.Linear(self.hidden_dim, self.hidden_dim, bias=False),
            nn.Linear(self.hidden_dim, self.num_experts, bias=False)
        )

        self.experts = nn.ModuleList([LlamaMLP(config) for _ in range(self.num_experts)])

        self.dummy = DummyModule()

        # Jitter parameters
        self.jitter_noise = config.jitter_noise

    def forward(self, hidden_states: torch.Tensor, routing_weights: Optional[torch.Tensor] = None) -> torch.Tensor:
        """ """
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        if self.training and self.jitter_noise > 0:
            hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
        hidden_states = hidden_states.view(-1, hidden_dim)
        
        if self.use_router:
            router_logits = self.gate(hidden_states)
            if "logic" in self.ablate:
                router_logits[..., 0] = -torch.inf
            if "social" in self.ablate:
                router_logits[..., 1] = -torch.inf
            if "world" in self.ablate:
                router_logits[..., 2] = -torch.inf
            if "language" in self.ablate:
                router_logits[..., 3] = -torch.inf
            routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
        else:
            routing_weights = routing_weights.reshape(-1, 4).float()
            router_logits = routing_weights
        # router_logits: (batch * sequence_length, n_experts)
        
        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
        # we cast back to the input dtype
        routing_weights = routing_weights.to(hidden_states.dtype)

        final_hidden_states = torch.zeros(
            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
        )

        H_up = self.experts[0].up_proj.out_features
        Y_up = hidden_states.new_zeros((batch_size, sequence_length, self.num_experts, H_up))


        # One hot encode the selected experts to create an expert mask
        # this will be used to easily index which expert is going to be sollicitated
        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)

        expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
        for expert_idx in expert_hitted:
            expert_layer = self.experts[expert_idx]
            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
            # Index the correct hidden states and compute the expert hidden state for
            # the current expert. We need to make sure to multiply the output hidden
            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)

            # --- Hook to capture up-proj output BEFORE nonlinearity ---
            captured_up = []
            def _up_hook(m, inp, out):
                # out shape: [N_e, H_up]
                captured_up.append(out.detach())

            h = expert_layer.up_proj.register_forward_hook(_up_hook)

            current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
            h.remove()

            # Scatter captured up-proj per-token into Y_up[b, t, expert, :]
            if captured_up:
                up = captured_up[-1]  # [N_e, H_up]
                b_idx = top_x // sequence_length
                t_idx = top_x %  sequence_length
                # Y_up[b,t,e,:] = up[n,:]
                Y_up[b_idx, t_idx, expert_idx, :] = up
            
            # However `index_add_` only support torch tensors for indexing so we'll use
            # the `top_x` tensor here.
            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)

        self.dummy(Y_up)
        return final_hidden_states, router_logits

class LlamaMiCRoMoEDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: MiCRoLlamaMoEConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size

        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)

        self.block_sparse_moe = LlamaSparseMiCRoMoEBlock(config)
        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        routing_weights: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[tuple[torch.Tensor]] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> torch.FloatTensor:
        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            cache_position=cache_position,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states, router_logits = self.block_sparse_moe(hidden_states, routing_weights)
        hidden_states = residual + hidden_states

        return hidden_states, router_logits

    
class MiCRoLlamaMoE(LlamaPreTrainedModel, GenerationMixin):
    config_class = MiCRoLlamaMoEConfig
    def __init__(self, config):
        with open(config.config_path, 'r', encoding="utf-8") as file:
            run_config = yaml.load(file.read(), Loader=yaml.FullLoader)

        self.config: MiCRoLlamaMoEConfig = config
        self.config.torch_dtype = torch.bfloat16
        self.config.use_bfloat16 = True
        self.config._attn_implementation = "eager" # {sdpa, flash_attention_2, eager}
        self.config.use_cache = True
        self.config.backbone_num_layers = self.config.num_hidden_layers
        self.config.num_hidden_layers = self.config.num_hidden_layers
        self.config.loss_type = "ForCausalLMLoss"

        super(MiCRoLlamaMoE, self).__init__(self.config)
        self.build_model(run_config)

    def build_model(self, run_config):
    
        self.config.num_experts = run_config["num-experts"]
        self.config.use_router = run_config["use-router"]
        self.config.num_experts_per_tok = run_config["top-k-experts"]
        print(f">> Top-K Experts Per Token: {self.config.num_experts_per_tok}")
        self.config.jitter_noise = run_config["jitter-noise"]
        self.config.loss_method = run_config.get("loss", "all")
        self.router_aux_loss_coef = run_config["router-aux-loss-coef"]
        self.use_load_balancing = run_config.get("use-load-balancing", False)

        self.config.gradient_checkpointing = run_config.get("gradient-checkpointing", False)
        self.gradient_checkpointing = self.config.gradient_checkpointing

        print(f">> Gradient Checkpointing: {self.config.gradient_checkpointing}")

        self.run_config = run_config
        self.padding_idx = 2 if "smollm2" in run_config["model"] else 128004
        
        # LlamaMoE model
        self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([LlamaMiCRoMoEDecoderLayer(self.config, layer_idx) for layer_idx in range(self.config.backbone_num_layers)])
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        self.rotary_emb = LlamaRotaryEmbedding(config=self.config)
        self.final_norm = LlamaRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
        
        if "model" not in run_config["trainable"]:
            print(">> Freezing Model Except Experts + Routing Gates")
            for param in self.parameters():
                param.requires_grad = False

            for layer in self.layers:
                layer: LlamaMiCRoMoEDecoderLayer
                for param in layer.block_sparse_moe.parameters():
                    param.requires_grad = True

        if "experts" not in run_config["trainable"]:
            print(">> Freezing Experts")
            for layer in self.layers:
                layer: LlamaMiCRoMoEDecoderLayer
                for param in layer.block_sparse_moe.experts.parameters():
                    param.requires_grad = False

        if "experts-router" not in run_config["trainable"]:
            print(">> Freezing Routing Gates")
            for layer in self.layers:
                layer: LlamaMiCRoMoEDecoderLayer
                for param in layer.block_sparse_moe.gate.parameters():
                    param.requires_grad = False


    def forward(self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        experts_ablate: Optional[List[str]] = None,
        routing_weights: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs: Unpack[FlashAttentionKwargs],
    ):

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if self.gradient_checkpointing and self.training and use_cache:
            logger.warning_once(
                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
            )
            use_cache = False

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        causal_mask = self._update_causal_mask(
            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
        )

        hidden_states = inputs_embeds

        # create position embeddings to be shared across the decoder layers
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None

        all_routing_weights = ()

        for decoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs, router_logits = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    position_embeddings,
                    routing_weights,
                    causal_mask,
                    position_ids,
                    past_key_values,
                    cache_position,
                )
            else:
                layer_outputs, router_logits = decoder_layer(
                    hidden_states,
                    position_embeddings=position_embeddings,
                    routing_weights=routing_weights,
                    attention_mask=causal_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    cache_position=cache_position,
                    **kwargs,
                )

            hidden_states = layer_outputs

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

            all_routing_weights += (router_logits,)

        hidden_states = self.final_norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
    
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

            loss += keep_alive_zero(self)
            
            aux_loss = None
            if self.use_load_balancing:
                aux_loss = load_balancing_loss_func(
                    all_routing_weights,
                    self.config.num_experts,
                    self.config.num_experts_per_tok,
                    attention_mask,
                )
                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device

        if not return_dict:
            output = (logits,) + (past_key_values, all_hidden_states, all_self_attns, all_routing_weights) if use_cache else (logits, all_hidden_states, all_self_attns, all_routing_weights) 
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            routing_weights=all_routing_weights,
        )

    def _update_causal_mask(
        self,
        attention_mask: torch.Tensor,
        input_tensor: torch.Tensor,
        cache_position: torch.Tensor,
        past_key_values: Cache,
        output_attentions: bool,
    ):
        if self.config._attn_implementation == "flash_attention_2":
            if attention_mask is not None and 0.0 in attention_mask:
                return attention_mask
            return None

        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
        # to infer the attention mask.
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
        using_static_cache = isinstance(past_key_values, StaticCache)

        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
        if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
            if AttentionMaskConverter._ignore_causal_mask_sdpa(
                attention_mask,
                inputs_embeds=input_tensor,
                past_key_values_length=past_seen_tokens,
                is_training=self.training,
            ):
                return None

        dtype, device = input_tensor.dtype, input_tensor.device
        min_dtype = torch.finfo(dtype).min
        sequence_length = input_tensor.shape[1]
        if using_static_cache:
            target_length = past_key_values.get_max_length()
        else:
            target_length = (
                attention_mask.shape[-1]
                if isinstance(attention_mask, torch.Tensor)
                else past_seen_tokens + sequence_length + 1
            )

        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
        causal_mask = _prepare_4d_causal_attention_mask_with_cache_position(
            attention_mask,
            sequence_length=sequence_length,
            target_length=target_length,
            dtype=dtype,
            device=device,
            min_dtype=min_dtype,
            cache_position=cache_position,
            batch_size=input_tensor.shape[0],
        )

        if (
            self.config._attn_implementation == "sdpa"
            and attention_mask is not None
            and attention_mask.device.type == "cuda"
            and not output_attentions
        ):
            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
            # Details: https://github.com/pytorch/pytorch/issues/110213
            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)

        return causal_mask

    def load_pretrained(self, model_name):
        base_model: LlamaForCausalLM = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
        self.lm_head.load_state_dict(base_model.lm_head.state_dict())
        self.embed_tokens.load_state_dict(base_model.get_input_embeddings().state_dict())
        self.rotary_emb.load_state_dict(base_model.model.rotary_emb.state_dict())
        self.final_norm.load_state_dict(base_model.model.norm.state_dict())

        for layer_idx, layer in enumerate(self.layers):
            
            attn_layer = base_model.model.layers[layer_idx].self_attn.state_dict()
            layer.self_attn.load_state_dict(attn_layer)
            
            input_layernorm = base_model.model.layers[layer_idx].input_layernorm.state_dict()
            layer.input_layernorm.load_state_dict(input_layernorm)

            post_attention_layernorm = base_model.model.layers[layer_idx].post_attention_layernorm.state_dict()
            layer.post_attention_layernorm.load_state_dict(post_attention_layernorm)
            
            mlp_model_layer = base_model.model.layers[layer_idx].mlp.state_dict()
            for expert in layer.block_sparse_moe.experts:
                expert.load_state_dict(mlp_model_layer)

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        experts_ablate=None,
        use_cache=True,
        num_logits_to_keep=None,
        **kwargs,
    ):
        
        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        # Exception 1: when passing input_embeds, input_ids may be missing entries
        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        if past_key_values is not None:
            if inputs_embeds is not None:  # Exception 1
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -input_ids.shape[1] :]

                # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s  `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture.
                position_ids = position_ids.clone(memory_format=torch.contiguous_format)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and cache_position[0] == 0:
            model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
        else:
            # The clone here is for the same reason as for `position_ids`.
            model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None}

        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
            if model_inputs["inputs_embeds"] is not None:
                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
                device = model_inputs["inputs_embeds"].device
            else:
                batch_size, sequence_length = model_inputs["input_ids"].shape
                device = model_inputs["input_ids"].device

            dtype = self.lm_head.weight.dtype
            min_dtype = torch.finfo(dtype).min

            attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
                attention_mask,
                sequence_length=sequence_length,
                target_length=past_key_values.get_max_length(),
                dtype=dtype,
                device=device,
                min_dtype=min_dtype,
                cache_position=cache_position,
                batch_size=batch_size,
            )

        if num_logits_to_keep is not None:
            model_inputs["num_logits_to_keep"] = num_logits_to_keep

        model_inputs.update(
            {
                "experts_ablate": experts_ablate,
                "position_ids": position_ids,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
            }
        )
        return model_inputs
    
def load_balancing_loss_func(
    gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],
    num_experts: Optional[int] = None,
    top_k=2,
    attention_mask: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, int]:
    r"""
    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.

    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
    experts is too unbalanced.

    Args:
        gate_logits:
            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
            shape [batch_size X sequence_length, num_experts].
        num_experts:
            Number of experts
        top_k:
            The number of experts to route per-token, can be also interpreted as the `top-k` routing
            parameter.
        attention_mask (`torch.Tensor`, *optional*):
            The attention_mask used in forward function
            shape [batch_size X sequence_length] if not None.

    Returns:
        The auxiliary loss.
    """
    if gate_logits is None or not isinstance(gate_logits, tuple):
        return 0

    if isinstance(gate_logits, tuple):
        compute_device = gate_logits[0].device
        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)

    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)

    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)

    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)

    if attention_mask is None:
        # Compute the percentage of tokens routed to each experts
        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)

        # Compute the average probability of routing to these experts
        router_prob_per_expert = torch.mean(routing_weights, dim=0)
    else:
        batch_size, sequence_length = attention_mask.shape
        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)

        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
        expert_attention_mask = (
            attention_mask[None, :, :, None, None]
            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
            .reshape(-1, top_k, num_experts)
            .to(compute_device)
        )

        # Compute the percentage of tokens routed to each experts
        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
            expert_attention_mask, dim=0
        )

        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
        router_per_expert_attention_mask = (
            attention_mask[None, :, :, None]
            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
            .reshape(-1, num_experts)
            .to(compute_device)
        )

        # Compute the average probability of routing to these experts
        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
            router_per_expert_attention_mask, dim=0
        )

    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
    return overall_loss * num_experts


AutoConfig.register("micro_llama_moe", MiCRoLlamaMoEConfig)
AutoModelForCausalLM.register(MiCRoLlamaMoEConfig, MiCRoLlamaMoE)