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)