Mirror facebook/MobileMoE-S-Base

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+mobilemoe_pareto.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,124 @@
+FAIR Noncommercial Research License
+v1 Last Updated: August 21, 2026
+
+“Acceptable Use Policy” means the FAIR Acceptable Use Policy, applicable to Research Materials, that is incorporated into this Agreement.
+
+“Agreement” means the terms and conditions for use, reproduction, distribution and modification of the Research Materials set forth herein.
+
+
+“Documentation” means the specifications, manuals and documentation accompanying 
+Research Materials distributed by Meta.
+
+
+“Licensee” or “you” means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
+
+
+“Meta” or “we” means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) and Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+“Noncommercial Research Uses” means noncommercial research use cases related to research, development, education, processing, or analysis and in each case, is not primarily intended for commercial advantage or monetary compensation to you or others.
+
+“Research Materials” means, collectively, Documentation and the models, software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, demonstration materials and other elements of the foregoing distributed by Meta and made available under this Agreement.
+
+By clicking “I Accept” below or by using or distributing any portion or element of the Research Materials, you agree to be bound by this Agreement.
+
+
+1. License Rights and Redistribution.
+
+
+a. Grant of Rights. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Meta’s intellectual property or other rights owned by Meta embodied in the Research Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the Research Materials.  
+ 
+b. Redistribution and Use.  
+ i. You will not use the Research Materials or any outputs or results of the Research Materials in connection with any commercial uses or for any uses other than Noncommercial Research Uses;
+
+
+ii. Distribution of Research Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the Research Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement. You shall also provide a copy of this Agreement to such third party.
+
+
+iii.  If you submit for publication the results of research you perform on, using, or otherwise in connection with Research Materials, you must acknowledge the use of Research Materials in your publication.
+
+
+iv. Your use of the Research Materials must comply with applicable laws and regulations (including Trade Control Laws) and adhere to the FAIR Acceptable Use Policy, which is hereby incorporated by reference into this Agreement.
+2. User Support. Your Noncommercial Research Use of the Research Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use.  Meta is under no obligation to provide any support services for the Research Materials. Any support provided is “as is”, “with all faults”, and without warranty of any kind.
+
+
+3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE RESEARCH MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE RESEARCH MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE RESEARCH MATERIALS AND ANY OUTPUT AND RESULTS.
+ 
+4. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+ 
+5. Intellectual Property.
+
+
+a. Subject to Meta’s ownership of Research Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the Research Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+ 
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Research Materials, outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the Research Materials.
+ 
+6. Term and Termination. The term of this Agreement will commence upon your acceptance of this Agreement or access to the Research Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the Research Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement. 
+ 
+7. Governing Law and Jurisdiction. This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement. 
+
+
+8. Modifications and Amendments. Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the Research Materials after any modification to this Agreement constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.
+
+
+FAIR Acceptable Use Policy 
+
+The Fundamental AI Research (FAIR) team at Meta seeks to further understanding of new and existing research domains with the mission of advancing the state-of-the-art in artificial intelligence through open research for the benefit of all. 
+
+As part of this mission, Meta makes certain research materials available for noncommercial research use. Meta is committed to promoting the safe and responsible use of such research materials.  
+
+Prohibited Uses
+
+You agree you will not use, or allow others to use, Research Materials to:
+
+ Violate the law or others’ rights, including to:
+Engage in, promote, generate, contribute to, encourage, plan, incite, or further illegal or unlawful activity or content, such as:
+Violence or terrorism
+Exploitation or harm to children, including the solicitation, creation, acquisition, or dissemination of child exploitative content or failure to report Child Sexual Abuse Material
+Human trafficking, exploitation, and sexual violence
+The illegal distribution of information or materials to minors, including obscene materials, or failure to employ legally required age-gating in connection with such information or materials.
+Sexual solicitation
+Any other criminal activity
+
+Engage in, promote, incite, or facilitate the harassment, abuse, threatening, or bullying of individuals or groups of individuals
+
+Engage in, promote, incite, or facilitate discrimination or other unlawful or harmful conduct in the provision of employment, employment benefits, credit, housing, other economic benefits, or other essential goods and services
+
+Engage in the unauthorized or unlicensed practice of any profession including, but not limited to, financial, legal, medical/health, or related professional practices
+
+Collect, process, disclose, generate, or infer health, demographic, or other sensitive personal or private information about individuals without rights and consents required by applicable laws
+
+Engage in or facilitate any action or generate any content that infringes, misappropriates, or otherwise violates any third-party rights, including the outputs or results of any technology using FAIR research materials
+
+Create, generate, or facilitate the creation of malicious code, malware, computer viruses or do anything else that could disable, overburden, interfere with or impair the proper working, integrity, operation or appearance of a website or computer system
+
+2. Engage in, promote, incite, facilitate, or assist in the planning or development of activities that present a risk of death or bodily harm to individuals, including use of research artifacts related to the following:
+
+Military, warfare, nuclear industries or applications, espionage, use for materials or activities that are subject to the International Traffic Arms Regulations (ITAR) maintained by the United States Department of State
+
+Guns and illegal weapons (including weapon development)
+
+Illegal drugs and regulated/controlled substances
+
+Operation of critical infrastructure, transportation technologies, or heavy machinery
+
+Self-harm or harm to others, including suicide, cutting, and eating disorders
+
+Any content intended to incite or promote violence, abuse, or any infliction of bodily harm to an individual
+
+3. Intentionally deceive or mislead others, including use of FAIR Research Materials related to the following:
+
+ Generating, promoting, or furthering fraud or the creation or promotion of disinformation
+
+ Generating, promoting, or furthering defamatory content, including the creation of defamatory statements, images, or other content
+
+Generating, promoting, or further distributing spam
+
+ Impersonating another individual without consent, authorization, or legal right
+
+Representing that outputs of FAIR research materials or outputs from technology using FAIR research materials are human-generated
+
+Generating or facilitating false online engagement, including fake reviews and other means of fake online engagement
+
+4. Fail to appropriately disclose to end users any known dangers of your Research Materials.
+
+Please report any violation of this Policy or other problems that could lead to a violation of this Policy by submitting a report here [https://docs.google.com/forms/d/e/1FAIpQLSeb11cryAopJ7LNrC4nxEUXrHY26hfkXQMf_uH-oFgA3WlYZQ/viewform]. 

README.md

@@ -0,0 +1,184 @@
+---
+license: fair-noncommercial-research-license
+language:
+- en
+library_name: transformers
+tags:
+- facebook
+- meta
+- pytorch
+- mixture-of-experts
+- MoE
+- on-device
+---
+> Mirror of [facebook/MobileMoE-S-Base](https://huggingface.co/facebook/MobileMoE-S-Base). All credit to the original authors. Not an official release from them. License unchanged (FAIR Noncommercial Research License, non-commercial use only); see `LICENSE`.
+
+# MobileMoE-S (Base) Model Card
+
+**MobileMoE** is a family of on-device Mixture-of-Experts (MoE) language models with sub-billion *active* parameters, designed to push the quality–efficiency Pareto frontier for on-device LLMs, including three model scales (S/M/L): 0.3B/0.5B/0.9B active parameters (1.3B/2.8B/5.3B total), with <3 GB INT4 weight footprints to fit in mobile DRAM. Each scale is released in three variants: a **Base** model (pre-training + mid-training), an **SFT** model (supervised fine-tuning), and a **QAT** model (quantization-aware training). You are currently in the **MobileMoE-S-Base** repository — the pre-trained 0.3B-active base model.
+
+| | **S** | **M** | **L** |
+|:---|:---:|:---:|:---:|
+| Active / total params | 272M / 1.3B | 528M / 2.8B | 922M / 5.3B |
+| Layers | 20 | 26 | 32 |
+| Model dimension | 768 | 1024 | 1280 |
+| Heads (Q / KV) | 12 / 4 | 16 / 4 | 20 / 4 |
+| Routed experts | 60 | 60 | 60 |
+| Top-k | 4 | 4 | 4 |
+| INT4 weight memory | 0.68 GB | 1.48 GB | 2.75 GB |
+| Base | **MobileMoE-S-Base** | [MobileMoE-M-Base](https://huggingface.co/facebook/MobileMoE-M-Base) | [MobileMoE-L-Base](https://huggingface.co/facebook/MobileMoE-L-Base) |
+| SFT | [MobileMoE-S-SFT](https://huggingface.co/facebook/MobileMoE-S-SFT) | [MobileMoE-M-SFT](https://huggingface.co/facebook/MobileMoE-M-SFT) | [MobileMoE-L-SFT](https://huggingface.co/facebook/MobileMoE-L-SFT) |
+| QAT (INT4) | [MobileMoE-S-QAT](https://huggingface.co/facebook/MobileMoE-S-QAT) | [MobileMoE-M-QAT](https://huggingface.co/facebook/MobileMoE-M-QAT) | [MobileMoE-L-QAT](https://huggingface.co/facebook/MobileMoE-L-QAT) |
+
+For the detailed technical report: 📝 [MobileMoE: Scaling On-Device Mixture of Experts](https://arxiv.org/abs/2605.27358)
+
+For more versions, check out the 🤗 [MobileMoE Collection](https://huggingface.co/collections/facebook/mobilemoe)
+
+![MobileMoE establishes a new Pareto frontier for on-device LLMs](mobilemoe_pareto.png)
+
+**MobileMoE establishes a new Pareto frontier for on-device LLMs.** Average benchmark accuracy, computed over 14 benchmarks spanning commonsense, knowledge, science, comprehension, and reasoning, is plotted against (a) per-token inference compute *F*<sub>inf</sub> = 2*N*<sub>act</sub> (GFLOPs) and (b) total parameters *N*<sub>total</sub> (B); in (b), x-axis tick labels show total params (B) | projected INT4 memory (GB).
+
+## Key Features
+
+- **A new Pareto frontier for on-device LLMs.** Across 14 foundational benchmarks, MobileMoE matches or exceeds leading on-device dense LLMs at 2–4× fewer inference FLOPs, and matches or surpasses the state-of-the-art MoE OLMoE-1B-7B with up to 60% fewer parameters.
+- **Scaling-law-derived architecture.** The architecture is derived from an on-device MoE scaling law that jointly optimizes under mobile *memory* and *compute* constraints, identifying an on-device sweet spot: moderate sparsity, with fine-grained experts and shared expert.
+- **Four-stage recipe.** Pre-training → mid-training → instruction fine-tuning → INT4 quantization-aware training, all on open-source datasets.
+
+## Model Information
+
+**Model:** MobileMoE-S-Base (pre-trained + mid-trained)<br>
+**Active Parameters:** 272M<br>
+**Total Parameters:** 1.3B<br>
+**Layers:** 20<br>
+**Model Dimension:** 768<br>
+**Attention Heads:** 12<br>
+**KV Heads:** 4 (GQA)<br>
+**Head Dimension:** 64<br>
+**Routed Experts:** 60 (fine-grained, FFN hidden dim 384 each)<br>
+**Active Experts per Token:** 4 (top-k sigmoid routing, with normalization)<br>
+**Shared Expert:** 1, always on (FFN hidden dim 1536)<br>
+**Vocabulary Size:** 128,256<br>
+**Other Features:** QK-Norm, tied input/output embeddings, RoPE (θ = 500,000)<br>
+**Input Modality:** Text<br>
+**Output Modality:** Text<br>
+**Languages:** English<br>
+**Training Stages:** Pre-training → mid-training<br>
+**Context Length:** 8,192 tokens<br>
+**Precision:** BF16<br>
+**Model Developer:** Meta<br>
+**Model Release Date:** Aug 2026<br>
+**License:** MobileMoE is FAIR NC licensed
+
+## Results
+
+All numbers below are for the **base (pre-trained)** models, re-evaluated under identical settings with greedy decoding using [`lm-eval`](https://github.com/EleutherAI/lm-evaluation-harness); few-shot counts are given in parentheses after the benchmark name, and benchmarks shown without one are evaluated 0-shot.
+
+### Foundational benchmarks
+
+| Capability | Benchmark | Gemma 3 270M | SmolLM2 360M | **MobileMoE-S** |
+|:---|:---|:---:|:---:|:---:|
+| *Active / total params* | | 270M | 362M | **272M / 1.3B** |
+| Commonsense Reasoning | HellaSwag | 41.4 | 56.5 | **58.9** |
+|  | PIQA | 68.3 | 71.7 | **75.4** |
+|  | SIQA | 40.2 | 40.7 | **46.8** |
+|  | WinoGrande | 53.7 | **59.0** | 58.6 |
+| Knowledge | MMLU (5-shot) | 26.7 | 25.2 | **43.7** |
+|  | NaturalQuestions (5-shot) | 4.1 | 7.4 | **12.6** |
+|  | TriviaQA (5-shot) | 14.3 | 26.8 | **33.2** |
+| Science | ARC-Challenge (25-shot) | 29.4 | 40.5 | **46.5** |
+|  | ARC-Easy | 56.8 | 68.1 | **73.9** |
+|  | OpenBookQA | 30.4 | **37.6** | 34.6 |
+| Reading | BoolQ | 58.3 | **61.8** | 60.2 |
+|  | DROP (3-shot) | 14.2 | 17.9 | **39.0** |
+| Reasoning | BIG-Bench Hard (3-shot) | 29.5 | 31.7 | **31.8** |
+|  | GSM8K (8-shot) | 1.8 | 5.3 | **36.2** |
+| **Average** | | 33.5 | 39.3 | **46.5** |
+
+## Training
+
+MobileMoE uses a four-stage recipe. This checkpoint is the output of stage 2 (mid-training).
+
+![MobileMoE four-stage training recipe](mobilemoe_recipe.png)
+
+**MobileMoE four-stage training recipe:** pre-training (PT) → mid-training (MT) → instruct supervised fine-tuning (SFT) → quantization-aware training (QAT) with INT4 precision.
+
+| | **Pre-training** | **Mid-training** | SFT | QAT |
+|---|---|---|---|---|
+| Context length | **2,048** | **8,192** | 8,192 | 8,192 |
+| Total tokens | **~6T** | **~500B** | ~126B | ~21B |
+| Peak learning rate | **4×10<sup>-4</sup>** | **4×10<sup>-5</sup>** | 4×10<sup>-6</sup> | 4×10<sup>-6</sup> |
+| LR schedule | **Cosine** | **Linear** | Cosine | Cosine |
+| Token dispatch | **drop-and-pad** | **drop-and-pad** | dropless | dropless |
+
+## How to use
+
+MobileMoE uses a custom architecture (`model_type: mobilemoe`) that is not yet part of upstream `transformers`, so **`trust_remote_code=True` is required**. The modeling code ships in this repo (`configuration_mobilemoe.py`, `modeling_mobilemoe.py`).
+
+### Requirements
+
+```bash
+pip install "torch>=2.1" "transformers>=4.57" "safetensors>=0.4" "accelerate>=1.0"
+```
+
+Verified with the following versions:
+
+| Package | Version |
+|---|---|
+| `torch` | 2.8.0 (cu128) |
+| `transformers` | 4.57.6 |
+| `tokenizers` | 0.22.2 |
+| `safetensors` | 0.7.0 |
+| `accelerate` | 1.13.0 |
+
+For batch evaluation we recommend vLLM (≥ 0.10.2) with `enforce_eager=True`.
+
+### Text generation
+
+This is a **base model** — it has no chat template and is not instruction-tuned. Prompt it with plain text continuation:
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+MODEL_ID = "facebook/MobileMoE-S-Base"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+)
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+model.eval()
+
+prompt = "Why are open-source on-device language models great?"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+outputs = model.generate(
+    input_ids=inputs["input_ids"],
+    attention_mask=inputs["attention_mask"],
+    max_new_tokens=64,
+    do_sample=False,
+    temperature=None,
+    top_p=None,
+    pad_token_id=tokenizer.eos_token_id,
+)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+**Known issues.** Loading the tokenizer on `transformers` 4.57.6 prints a `fix_mistral_regex=True` warning. Please ignore it and do not set the flag, as MobileMoE uses the Llama-3 tokenizer whose default tokenization is already correct.
+
+## Citation
+
+```bibtex
+@article{chen2026mobilemoe,
+  title={MobileMoE: Scaling On-Device Mixture of Experts},
+  author={Chen, Yanbei and Huang, Hanxian and Chang, Ernie and Szwejbka, Jacob and Desai, Digant and Liu, Zechun and Chandra, Vikas and Krishnamoorthi, Raghuraman},
+  journal={arXiv preprint arXiv:2605.27358},
+  year={2026}
+}
+```
+
+## License
+
+MobileMoE is distributed under the [FAIR Noncommercial Research License](https://huggingface.co/facebook/MobileMoE-S-Base/blob/main/LICENSE).

config.json

@@ -0,0 +1,87 @@
+{
+  "model_type": "mobilemoe",
+  "architectures": [
+    "MobileMoEForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mobilemoe.MobileMoEConfig",
+    "AutoModel": "modeling_mobilemoe.MobileMoETextModel",
+    "AutoModelForCausalLM": "modeling_mobilemoe.MobileMoEForCausalLM"
+  },
+  "vocab_size": 128256,
+  "bos_token_id": 128000,
+  "eos_token_id": 128001,
+  "hidden_size": 768,
+  "intermediate_size": 384,
+  "intermediate_size_mlp": 1536,
+  "num_hidden_layers": 20,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "max_position_embeddings": 8192,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 500000.0,
+  "rope_scaling": {
+    "rope_type": "llama3",
+    "factor": 16,
+    "original_max_position_embeddings": 8192,
+    "low_freq_factor": 1.0,
+    "high_freq_factor": 1.0
+  },
+  "num_local_experts": 60,
+  "num_experts_per_tok": 4,
+  "interleave_moe_layer_step": 1,
+  "no_rope_layers": [
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1
+  ],
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "use_qk_norm": true,
+  "attn_temperature_tuning": false,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "routed_scaling_factor": 2.5,
+  "norm_topk_prob": true
+}

configuration_mobilemoe.py

@@ -0,0 +1,110 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the FAIR Noncommercial Research License
+
+"""Configuration class for MobileMoE.
+
+MobileMoE is a Mixture-of-Experts language model with sigmoid-topk routing,
+postgate scoring, and a shared expert. This config extends PretrainedConfig
+with MoE-specific fields (routed_scaling_factor, norm_topk_prob,
+intermediate_size_mlp) as first-class attributes.
+"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class MobileMoEConfig(PretrainedConfig):
+    model_type = "mobilemoe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=128256,
+        hidden_size=768,
+        intermediate_size=384,
+        intermediate_size_mlp=1536,
+        num_hidden_layers=20,
+        num_attention_heads=12,
+        num_key_value_heads=4,
+        head_dim=64,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=True,
+        rope_theta=500000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        num_experts_per_tok=4,
+        num_local_experts=60,
+        interleave_moe_layer_step=1,
+        use_qk_norm=True,
+        attn_temperature_tuning=False,
+        no_rope_layers=None,
+        layer_types=None,
+        routed_scaling_factor=2.5,
+        norm_topk_prob=True,
+        output_router_logits=False,
+        floor_scale=8192,
+        attn_scale=0.1,
+        sliding_window=None,
+        attention_chunk_size=8192,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.intermediate_size_mlp = intermediate_size_mlp
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_local_experts = num_local_experts
+        self.interleave_moe_layer_step = interleave_moe_layer_step
+        self.use_qk_norm = use_qk_norm
+        self.attn_temperature_tuning = attn_temperature_tuning
+        self.routed_scaling_factor = routed_scaling_factor
+        self.norm_topk_prob = norm_topk_prob
+        self.output_router_logits = output_router_logits
+        self.floor_scale = floor_scale
+        self.attn_scale = attn_scale
+        self.sliding_window = sliding_window
+        self.attention_chunk_size = attention_chunk_size
+
+        if no_rope_layers is None:
+            no_rope_layers = [1] * num_hidden_layers
+        self.no_rope_layers = no_rope_layers
+
+        if layer_types is None:
+            layer_types = ["full_attention"] * num_hidden_layers
+        self.layer_types = layer_types
+
+        # Compute moe_layers from interleave_moe_layer_step
+        self.moe_layers = list(
+            range(
+                interleave_moe_layer_step - 1,
+                num_hidden_layers,
+                interleave_moe_layer_step,
+            )
+        )

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 128000,
+  "eos_token_id": 128001,
+  "do_sample": true,
+  "temperature": 0.6,
+  "top_p": 0.9,
+  "transformers_version": "4.57.6"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69d09b2b3e4845f7e44831799c74df45752189921021c039752ced2c53d319ac
+size 2526780232

modeling_mobilemoe.py

@@ -0,0 +1,721 @@
+# 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,
+        )

special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>"
+}

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "model_max_length": 8192,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "clean_up_tokenization_spaces": true,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ]
+}