# Copyright (c) Meta Platforms, Inc. and affiliates. # This source code is licensed under the FAIR Noncommercial Research License """MobileMoE model implementation. MobileMoE is a decoder-only Mixture-of-Experts (MoE) language model designed for on-device inference. Only the MoE layer is custom; all other components (attention, layernorm, embeddings) follow standard transformer design. Usage: from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained(ckpt_dir, trust_remote_code=True) """ import math from typing import Callable, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache from transformers.generation import GenerationMixin from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS try: from .configuration_mobilemoe import MobileMoEConfig except ImportError: from configuration_mobilemoe import MobileMoEConfig # --------------------------------------------------------------------------- # Building blocks # --------------------------------------------------------------------------- class MobileMoEL2Norm(nn.Module): """L2 normalization for QK norm (no learnable parameters).""" def __init__(self, eps: float = 1e-6): super().__init__() self.eps = eps def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): return self._norm(x.float()).type_as(x) class MobileMoERMSNorm(nn.Module): """RMS normalization with learnable scale.""" def __init__(self, hidden_size, eps=1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(hidden_size)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): output = self._norm(x.float()).type_as(x) return output * self.weight class MobileMoERotaryEmbedding(nn.Module): """Rotary Position Embedding (RoPE) with optional Llama3-style scaling.""" def __init__(self, config, device=None): super().__init__() self.rope_type = "llama3" if config.rope_scaling is not None else "default" self.max_seq_len_cached = config.max_position_embeddings self.config = config dim = config.head_dim base = config.rope_theta inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.float, device=device) / dim) ) # Llama3 rope scaling self.attention_scaling = 1.0 if config.rope_scaling is not None: factor = config.rope_scaling.get("factor", 1.0) low_freq_factor = config.rope_scaling.get("low_freq_factor", 1.0) high_freq_factor = config.rope_scaling.get("high_freq_factor", 4.0) old_context_len = config.rope_scaling.get( "original_max_position_embeddings", config.max_position_embeddings ) low_freq_wavelen = old_context_len / low_freq_factor high_freq_wavelen = old_context_len / high_freq_factor new_freqs = [] for freq, wavelen in zip(inv_freq, 2 * math.pi / inv_freq): if wavelen < high_freq_wavelen: new_freqs.append(freq) elif wavelen > low_freq_wavelen: new_freqs.append(freq / factor) else: if high_freq_factor != low_freq_factor: smooth = (old_context_len / wavelen - low_freq_factor) / ( high_freq_factor - low_freq_factor ) new_freqs.append((1 - smooth) * freq / factor + smooth * freq) else: new_freqs.append(freq) inv_freq = torch.tensor(new_freqs, dtype=inv_freq.dtype, device=device) # attention_scaling stays 1.0 for llama3 rope type # (factor ** 0.5 is only used for dynamic/yarn types) self.register_buffer("inv_freq", inv_freq, persistent=False) @torch.no_grad() def forward(self, x, position_ids): inv_freq_expanded = ( self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) ) position_ids_expanded = position_ids[:, None, :].float() device_type = ( x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" ) with torch.autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded.to(x.device) @ position_ids_expanded).transpose( 1, 2 ) freqs_cis = torch.polar(torch.ones_like(freqs), freqs) freqs_cis = freqs_cis * self.attention_scaling return freqs_cis def _apply_rotary_emb(xq, xk, freqs_cis): """Apply rotary embeddings using complex multiplication.""" xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) xq_out = torch.view_as_real(xq_ * freqs_cis[:, :, None, :]).flatten(3) xk_out = torch.view_as_real(xk_ * freqs_cis[:, :, None, :]).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) def _repeat_kv(hidden_states, n_rep): """Expand KV heads for grouped-query attention.""" batch, num_kv_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_kv_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim) def _eager_attention_forward( module, query, key, value, attention_mask, scaling, dropout=0.0, **kwargs ): """Eager attention forward (no FlashAttention dependency).""" key_states = _repeat_kv(key, module.num_key_value_groups) value_states = _repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = F.softmax(attn_weights, dim=-1) attn_weights = F.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights # --------------------------------------------------------------------------- # MLP / Expert layers # --------------------------------------------------------------------------- class MobileMoEMLP(nn.Module): """SwiGLU feed-forward network.""" def __init__(self, config, intermediate_size=None): super().__init__() if intermediate_size is None: intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False) self.act_fn = nn.SiLU() def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class MobileMoEExperts(nn.Module): """Parameter container for routed experts (fused gate+up projection).""" def __init__(self, config): super().__init__() self.num_experts = config.num_local_experts self.hidden_size = config.hidden_size self.expert_dim = config.intermediate_size self.gate_up_proj = nn.Parameter( torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim) ) self.down_proj = nn.Parameter( torch.empty(self.num_experts, self.expert_dim, self.hidden_size) ) self.act_fn = nn.SiLU() class MobileMoERouter(nn.Linear): """Router for expert selection.""" def __init__(self, config): super().__init__(config.hidden_size, config.num_local_experts, bias=False) self.num_experts = config.num_local_experts self.top_k = config.num_experts_per_tok # --------------------------------------------------------------------------- # MoE layer (custom MobileMoE routing) # --------------------------------------------------------------------------- class MobileMoETextMoe(nn.Module): """MobileMoE Mixture-of-Experts layer. Routing: sigmoid -> topk -> normalize -> scale Gating: postgate — routing scores applied to expert OUTPUT Shared expert: always-active dense FFN with independent intermediate dimension """ def __init__(self, config): super().__init__() self.top_k = config.num_experts_per_tok self.hidden_dim = config.hidden_size self.num_experts = config.num_local_experts self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 2.5) self.norm_topk_prob = getattr(config, "norm_topk_prob", True) self.experts = MobileMoEExperts(config) self.router = MobileMoERouter(config) # Expert bias for load-balancing-free routing # Added post-sigmoid, pre-topk: affects expert SELECTION only, not weights self.expert_bias = nn.Parameter( torch.zeros(config.num_local_experts), requires_grad=False, ) shared_size = getattr(config, "intermediate_size_mlp", config.intermediate_size) self.shared_expert = MobileMoEMLP(config, intermediate_size=shared_size) def forward(self, hidden_states): hidden_states = hidden_states.reshape(-1, self.hidden_dim) # Custom routing order (bypass default router forward) # Compute router matmul in fp32 for numerical stability router_logits = F.linear(hidden_states.float(), self.router.weight.float()) # MobileMoE routing: sigmoid FIRST, then topk scores = torch.sigmoid(router_logits.float()) # Add expert bias for topk selection (load-balancing-free routing) # Bias only affects which experts are SELECTED, not the routing weights scores_for_topk = scores + self.expert_bias _, topk_indices = torch.topk(scores_for_topk, self.top_k, dim=-1) # Gather UNBIASED scores at selected indices topk_scores = scores.gather(1, topk_indices) # Normalize per-token and scale if self.norm_topk_prob: topk_scores = topk_scores / (topk_scores.sum(dim=-1, keepdim=True) + 1e-20) topk_scores = topk_scores * self.routed_scaling_factor # Sparse per-expert loop with postgate scoring (reference implementation). # Processes only the top-k selected tokens per expert. out = torch.zeros_like(hidden_states) for e in range(self.num_experts): mask = (topk_indices == e).any(dim=-1) if not mask.any(): continue idx = mask.nonzero(as_tuple=True)[0] x_sel = hidden_states[idx] # Expert FFN: gate_up_proj[e] is [D, 2F], down_proj[e] is [F, D] gate_up = x_sel @ self.experts.gate_up_proj[e] gate, up = gate_up.chunk(2, dim=-1) expert_out = (self.experts.act_fn(gate) * up) @ self.experts.down_proj[e] # Postgate: score applied to expert OUTPUT pos = (topk_indices[idx] == e).nonzero(as_tuple=True) weights = topk_scores[idx][pos[0], pos[1]] out[idx] += expert_out * weights.unsqueeze(-1) # Shared expert (additive, applied to all tokens) out = out + self.shared_expert(hidden_states) return out, router_logits # --------------------------------------------------------------------------- # Attention # --------------------------------------------------------------------------- class MobileMoEAttention(nn.Module): """Multi-head grouped-query attention with RoPE and optional QK L2 norm.""" def __init__(self, config, layer_idx): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr( config, "head_dim", config.hidden_size // config.num_attention_heads ) self.num_attention_heads = config.num_attention_heads self.num_key_value_groups = ( config.num_attention_heads // config.num_key_value_heads ) self.num_key_value_heads = config.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attn_scale = config.attn_scale self.floor_scale = config.floor_scale self.attn_temperature_tuning = config.attn_temperature_tuning self.attention_dropout = config.attention_dropout self.is_causal = True self.use_rope = config.no_rope_layers[layer_idx] self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias, ) self.k_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias, ) self.v_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias, ) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias, ) if config.use_qk_norm and self.use_rope: self.qk_norm = MobileMoEL2Norm(config.rms_norm_eps) def forward( self, hidden_states, position_embeddings, attention_mask=None, past_key_values=None, cache_position=None, **kwargs, ): input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape) key_states = self.k_proj(hidden_states).view(*input_shape, -1, self.head_dim) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) if self.use_rope: query_states, key_states = _apply_rotary_emb( query_states, key_states, position_embeddings.to(query_states.device) ) if hasattr(self, "qk_norm"): query_states = self.qk_norm(query_states) key_states = self.qk_norm(key_states) if self.attn_temperature_tuning and not self.use_rope: attn_scales = ( torch.log1p( torch.floor((cache_position.float() + 1.0) / self.floor_scale) ) * self.attn_scale + 1.0 ) attn_scales = attn_scales.view((1, input_shape[-1], 1, 1)).expand( (*input_shape, 1, 1) ) query_states = (query_states * attn_scales).to(query_states.dtype) query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) if past_key_values is not None: cache_kwargs = {"cache_position": cache_position} key_states, value_states = past_key_values.update( key_states, value_states, self.layer_idx, cache_kwargs ) attention_interface: Callable = _eager_attention_forward if self.config._attn_implementation != "eager": if self.config._attn_implementation in ALL_ATTENTION_FUNCTIONS: attention_interface = ALL_ATTENTION_FUNCTIONS[ self.config._attn_implementation ] attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights # --------------------------------------------------------------------------- # Decoder layer # --------------------------------------------------------------------------- class MobileMoEDecoderLayer(nn.Module): """MobileMoE decoder layer with custom MoE routing.""" def __init__(self, config, layer_idx): super().__init__() self.hidden_size = config.hidden_size self.layer_idx = layer_idx self.attention_type = config.layer_types[layer_idx] self.self_attn = MobileMoEAttention(config, layer_idx) self.is_moe_layer = layer_idx in config.moe_layers if self.is_moe_layer: self.feed_forward = MobileMoETextMoe(config) else: self.feed_forward = MobileMoEMLP( config, intermediate_size=config.intermediate_size_mlp ) self.input_layernorm = MobileMoERMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.post_attention_layernorm = MobileMoERMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, hidden_states, attention_mask=None, position_ids=None, past_key_values=None, use_cache=False, cache_position=None, position_embeddings=None, **kwargs, ): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention attn_output, _ = self.self_attn( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, cache_position=cache_position, **kwargs, ) hidden_states = residual + attn_output # Feed-Forward residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.feed_forward(hidden_states) if self.is_moe_layer: hidden_states, _ = hidden_states hidden_states = residual + hidden_states.view(residual.shape) return hidden_states # --------------------------------------------------------------------------- # Full model # --------------------------------------------------------------------------- class MobileMoEPreTrainedModel(PreTrainedModel): """Base class for MobileMoE models.""" config_class = MobileMoEConfig base_model_prefix = "model" supports_gradient_checkpointing = True _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = False _supports_sdpa = True _supports_attention_backend = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, MobileMoERMSNorm): module.weight.data.fill_(1.0) elif isinstance(module, MobileMoEExperts): module.gate_up_proj.data.normal_(mean=0.0, std=std) module.down_proj.data.normal_(mean=0.0, std=std) class MobileMoETextModel(MobileMoEPreTrainedModel): """MobileMoE text model.""" _no_split_modules = ["MobileMoEDecoderLayer"] def __init__(self, config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [MobileMoEDecoderLayer(config, i) for i in range(config.num_hidden_layers)] ) self.norm = MobileMoERMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = MobileMoERotaryEmbedding(config) self.gradient_checkpointing = False self.post_init() def forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, cache_position=None, **kwargs, ): 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(config=self.config) 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) # Build causal mask causal_mask = create_causal_mask( config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids, ) hidden_states = inputs_embeds freq_cis = self.rotary_emb(hidden_states, position_ids) for decoder_layer in self.layers: hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=freq_cis, **kwargs, ) hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, ) def _build_causal_mask( self, attention_mask, inputs_embeds, cache_position, past_key_values ): """Build 4D causal attention mask.""" dtype, device = inputs_embeds.dtype, inputs_embeds.device min_dtype = torch.finfo(dtype).min seq_length = inputs_embeds.shape[1] target_length = ( cache_position[-1].item() + 1 if cache_position is not None else seq_length ) # Create causal mask [1, 1, Q, K] causal_mask = torch.full( (seq_length, target_length), fill_value=min_dtype, dtype=dtype, device=device, ) if seq_length != 1: causal_mask = torch.triu( causal_mask, diagonal=target_length - seq_length + 1 ) causal_mask = ( causal_mask[None, None, :, :] .expand(inputs_embeds.shape[0], 1, -1, -1) .clone() ) # Apply padding mask if attention_mask is not None and attention_mask.dim() == 2: mask_length = attention_mask.shape[-1] padding_mask = attention_mask[:, None, None, :mask_length].eq(0) causal_mask[:, :, :, :mask_length].masked_fill_(padding_mask, min_dtype) return causal_mask class MobileMoEForCausalLM(MobileMoEPreTrainedModel, GenerationMixin): """MobileMoE causal language model for text generation.""" _no_split_modules = ["MobileMoEDecoderLayer"] _tied_weights_keys = ["lm_head.weight"] _tp_plan = {"lm_head": "colwise_rep"} def __init__(self, config): super().__init__(config) self.model = MobileMoETextModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, cache_position=None, logits_to_keep=0, **kwargs, ): use_cache = use_cache if use_cache is not None else self.config.use_cache outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, cache_position=cache_position, ) hidden_states = outputs[0] 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, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, )