Instructions to use ml-resources/MobileMoE-M-Base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use ml-resources/MobileMoE-M-Base with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="ml-resources/MobileMoE-M-Base", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("ml-resources/MobileMoE-M-Base", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use ml-resources/MobileMoE-M-Base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "ml-resources/MobileMoE-M-Base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "ml-resources/MobileMoE-M-Base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/ml-resources/MobileMoE-M-Base

SGLang

How to use ml-resources/MobileMoE-M-Base with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "ml-resources/MobileMoE-M-Base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "ml-resources/MobileMoE-M-Base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "ml-resources/MobileMoE-M-Base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "ml-resources/MobileMoE-M-Base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use ml-resources/MobileMoE-M-Base with Docker Model Runner:
```
docker model run hf.co/ml-resources/MobileMoE-M-Base
```

MobileMoE-M-Base / modeling_mobilemoe.py

mrfakename

Mirror facebook/MobileMoE-M-Base

de35c17 verified 13 days ago

Raw History Blame Contribute Delete

26.2 kB

# 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,
        )