Release Instella-MoE-16B-A3B-Base

.gitattributes

@@ -33,3 +33,7 @@ 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
+images/fig_cost_performance.png filter=lfs diff=lfs merge=lfs -text
+images/result_table_base_long.png filter=lfs diff=lfs merge=lfs -text
+images/result_table_base_standard.png filter=lfs diff=lfs merge=lfs -text
+images/result_table_instruct.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,306 @@
+RESEARCH-ONLY RAIL Model License
+
+Section I: PREAMBLE
+
+BY ACCESSING, DOWNLOADING, INSTALLING, OR USING THE ARTIFACT, YOU AGREE
+TO BE BOUND BY THIS LICENSE. IF YOU DO NOT AGREE TO ALL OF THE TERMS AND
+CONDITIONS OF THIS LICENSE, DO NOT ACCESS, DOWNLOAD, INSTALL, OR USE THE
+ARTIFACT.
+
+1. Definitions
+
+(a) "Application" refers to a sequence of instructions or statements
+    written in machine code language, including object code (that is the
+    product of a compiler), binary code (data using a two-symbol system)
+    or an intermediate language (such as register transfer language).
+
+(b) "Artifact" refers to a software application (in either binary or
+    source code format), Model, and/or Source Code that accompanies this
+    License.
+
+(c) "Contribution" means any work, including any modifications or
+    additions to an Artifact, that is intentionally submitted to
+    Licensor for inclusion or incorporation in the Artifact directly or
+    indirectly by the rights owner. For the purposes of this definition,
+    "submitted" means any form of electronic, verbal, or written
+    communication sent to the Licensor or its representatives, including
+    but not limited to communication on electronic mailing lists, source
+    code control systems, and issue tracking systems that are managed
+    by, or on behalf of, the Licensor for the purpose of discussing,
+    sharing and improving the Artifact, but excluding communication that
+    is conspicuously marked or otherwise designated in writing by the
+    contributor as "Not a Contribution."
+
+(d) "Contributor" means Licensor or any other individual or legal entity
+    that creates or owns a Contribution that is added to or incorporated
+    into an Artifact or its Derivative.
+
+(e) "Data" means a collection of information and/or content extracted
+    from the dataset used with a given Model, including to train,
+    pretrain, or otherwise evaluate the Model. The Data is not licensed
+    under this License.
+
+(f) "Derivative" means a work derived from or based upon an Artifact,
+    and includes all modified versions of such Artifact.
+
+(g) "Distribution" means any transmission, reproduction, publication or
+    other sharing of an Artifact or Derivative to a Third Party,
+    including providing a hosted service incorporating the Artifact,
+    which is made available by electronic or other remote means -
+    e.g. API-based or web access.
+
+(h) "Harm" includes but is not limited to physical, mental,
+    psychological, financial and reputational damage, pain, or loss.
+
+(i) "License" means the terms and conditions for use, reproduction, and
+    Distribution as defined in this document.
+
+(j) "Licensor" means the rights owner (by virtue of creation or
+    documented transfer of ownership) or entity authorized by the rights
+    owner (e.g., exclusive licensee) that is granting the rights in this
+    License.
+
+(k) "Model" means any machine-learning based assembly or assemblies
+    (including checkpoints), consisting of learnt weights, parameters
+    (including optimizer states), corresponding to the model
+    architecture as embodied in the Source Code.
+
+(l) "Output" means the results of operating a Model as embodied in
+    informational content resulting therefrom.
+
+(m) "Permitted Purpose" means for academic or research purposes only.
+
+(n) "Source Code" means any collection of text written using
+    human-readable programming language, including the code and scripts
+    used to define, run, load, benchmark or evaluate a Model or any
+    component thereof, and/or used to prepare data for training or
+    evaluation, if any. Source Code includes any accompanying
+    documentation, tutorials, examples, etc, if any. For clarity, the
+    term "Source Code" as used in this License includes any and all
+    Derivatives of such Source Code.
+
+(o) "Third Party(ies)" means an individual or legal entity that is not
+    under common control with Licensor or You.
+
+(p) "Use" includes accessing, using, copying, modifying, and/or
+    distributing an Artifact; in connection with a Model as Artifact,
+    Use also includes creating content, fine-tuning, updating, running,
+    training, evaluating and/or re-parametrizing such Model.
+
+(q) "You" (or "Your") means an individual or legal entity receiving and
+    exercising permissions granted by this License and/or making use of
+    the Artifact for permitted purposes and in any permitted field of
+    use, including usage of the Artifact in an end-use application -
+    e.g. chatbot, translator, image generator, etc.
+
+Section II: INTELLECTUAL PROPERTY RIGHTS
+
+Both copyright and patent grants may apply to the Artifact. The Artifact
+is subject to additional terms and conditions as described in Section III
+below. 
+
+2. Grant of Copyright License. Conditioned upon compliance with Section
+III below and subject to the terms and conditions of this License, each
+Contributor hereby grants to You, only in connection with the Permitted
+Purpose, a worldwide, non-exclusive, royalty-free copyright license to
+reproduce, use, publicly display, publicly perform, sublicense, and
+distribute the Artifact and Derivatives thereof.
+
+3. Grant of Patent License. Conditioned upon compliance with Section III
+below and subject to the terms and conditions of this License, and only
+where and as applicable, each Contributor hereby grants to You, only in
+connection with the Permitted Purpose, a worldwide, non-exclusive,
+royalty-free, irrevocable (except as stated in this paragraph) patent
+license to make, have made, use, sell, offer to sell, import, and
+otherwise transfer the Artifact where such license applies only to those
+patent claims licensable by such Contributor that are necessarily
+infringed by their Contribution(s) alone or by combination of their
+Contribution(s) with the Artifact to which such Contribution(s) was
+submitted. If You institute patent litigation against any entity
+(including a cross-claim or counterclaim in a lawsuit) alleging that the
+Artifact and/or a Contribution incorporated within the Artifact
+constitutes direct or contributory patent infringement, then any patent
+licenses granted to You under this License in connection with the
+Artifact shall terminate as of the date such litigation is asserted or
+filed.
+
+Section III: CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION
+
+4. Use-based restrictions. The restrictions set forth in Attachment A are
+mandatory Use-based restrictions. Therefore You may not Use the Artifact
+in violation of such restrictions. You may Use the Artifact only subject
+to this License. You shall require all of Your users who use the Artifact
+or its Derivative to comply with the terms this License, including those
+contained in this paragraph, and only for the Permitted Purpose.
+
+5. The Output You Generate with a Model (as Artifact). Except as set
+forth herein, Licensor claims no rights in the Output You generate. You
+are accountable for the Output You generate and its subsequent uses. No
+use of the Output may contravene any provision as stated in this
+License.
+
+6. Distribution and Redistribution. You may host for Third Party remote
+access purposes (e.g. software-as-a-service), reproduce and distribute
+copies of the Artifact or its Derivatives in any medium, with or without
+modifications, provided that You meet the following conditions:
+
+6.1.  Use-based restrictions in paragraph 4 MUST be included as a
+      condition precedent to effect any type of legal agreement (e.g. a
+      license) governing the use and/or distribution of the Artifact or
+      its Derivatives, and You shall give such notice to any subsequent
+      Third Party recipients;
+6.2.  You shall give any Third Party recipients of the Artifact or its
+      Derivatives a copy of this License;
+6.3.  You shall cause any modified files to carry prominent notices
+      stating that You changed the files;
+6.4.  You shall retain all copyright, patent, trademark, and attribution
+      notices excluding those notices that do not pertain to any part of
+      the Artifact or its Derivatives.
+6.5.  You and any Third Party recipients of the Artifact or its
+      Derivative shall adhere to the Permitted Purpose.
+
+You may add Your own copyright statement to Your modifications and may
+provide additional or different license terms and conditions with
+respect to paragraph 6.1., to govern the use, reproduction, or
+Distribution of Your modifications, or for any Derivative, provided that
+Your use, reproduction, and Distribution of the Artifact or its
+Derivative otherwise complies with the conditions stated in this
+License. In other words, the Use-based restrictions in Attachment A form
+the minimum set of terms for You to license to Third Parties any
+Artifact or its Derivative, but You may add more restrictive terms if
+You deem it necessary.
+
+Section IV: OTHER PROVISIONS
+
+7. Updates and Runtime Restrictions. To the maximum extent permitted by
+law, Licensor reserves the right to restrict (remotely or otherwise)
+usage of the Artifact in violation of this License or update the
+Artifact through electronic means.
+
+8. Trademarks and Related. Nothing in this License permits You to make
+use of trademarks, trade names, or logos of the Licensor or a
+Contributor, or to otherwise suggest an endorsement or misrepresent the
+relationship between You and the Licensor (or a Contributor); any rights
+not expressly granted herein are reserved.
+
+9. Disclaimer of Warranty. Unless required by applicable law or agreed
+to in writing, Licensor provides the Artifact (and each Contributor
+provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+CONDITIONS OF ANY KIND, either express or implied, including, without
+limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT,
+MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely
+responsible for determining the appropriateness of using the Artifact,
+and assume any risks associated with Your exercise of permissions under
+this License.
+
+10. Limitation of Liability. In no event and under no legal theory,
+whether in tort (including negligence), contract, or otherwise, unless
+required by applicable law (such as deliberate and grossly negligent
+acts) or agreed to in writing, shall any Contributor be liable to You
+for damages, including any direct, indirect, special, incidental, or
+consequential damages of any character arising as a result of this
+License or out of the use or inability to use the Artifact (including
+but not limited to damages for loss of goodwill, work stoppage, computer
+failure or malfunction, or any and all other commercial damages or
+losses), even if such Contributor has been advised of the possibility of
+such damages.
+
+11. If any provision of this License is held to be invalid, illegal or
+unenforceable, the remaining provisions shall be unaffected thereby and
+remain valid as if such provision had not been set forth herein.
+
+12. Term and Termination. The term of this License will commence upon
+the earlier of Your (a) acceptance of this License or (b) accessing the
+Artifact; and will continue in full force and effect until terminated in
+accordance with the terms and conditions herein. Licensor may terminate
+this License if You are in breach of any term or condition of this
+License. Upon termination of this License, all licenses granted to You
+will terminate and You must promptly delete and cease use of the
+Artifact. Sections 1, 7, 8, 9, 10, 11, and 12 survive termination of
+this License.
+
+END OF TERMS AND CONDITIONS
+
+Attachment A
+
+AMD Responsible AI Use Policy
+
+AMD is committed to the responsible use of its Artificial Intelligence
+(AI) products and technologies ("AMD AI").  AMD AI may include
+artificial intelligence or machine learning technologies that use
+algorithms to analyze data and generate output using predictions based
+on patterns in data.  This policy explains the uses that AMD
+specifically prohibits.
+
+If you use any AMD AI, you are agreeing to use
+the AMD AI in compliance with applicable laws and not for any of the
+following prohibited uses.
+
+Prohibited Uses:
+
+1) No Illegal Acts.  Do not use AMD AI in violation of any applicable
+national, state, local, or other jurisdictional law, rule, regulation,
+or sanction.
+
+2) No Explicit Content.  Do not use AMD AI to submit (as input),
+generate, or disseminate content depicting violent or sexually explicit
+content or to create sexual chatbots.
+
+3) No Harm.  Do not use AMD AI for any potentially harmful uses,
+   including fraud, deception, bias and discrimination, abuse, or harassment,
+   including the following:
+
+   a) Harm or abuse of a minor, including grooming and child sexual
+      exploitation.
+
+   b) Impersonation of human beings for purposes of deception.
+
+   c) Generation or dissemination of information you know to be false
+      for the purpose of harming others.
+
+   d) Intentionally defame, disparage, or otherwise harass others.
+
+   e) Intentionally attempting to materially distort the behavior of a
+      person in a manner that causes or is likely to cause that person
+      or another person physical or psychological harm.
+
+   f) Providing medical advice or interpretation of medical results that
+      is intended to be a substitute for professional medical advice,
+      diagnosis, or treatment.
+
+   g) Engaging in the unlawful or unauthorized practice of any
+      profession, including financial, legal, medical, health, or
+      related professional practices.
+
+   h) Judgment of, discrimination against, or harm to individuals or
+      groups based on legally protected characteristics or categories,
+      online or offline social behavior, or known or predicted personal
+      or personality characteristics, including any of the foregoing
+      uses in social credit systems.
+
+4) No High-Risk Activity.  Do not use AMD AI in any high-risk activities
+or applications that create risk of harm to individuals health, safety
+or fundamental rights or to safety of critical infrastructure or
+environmental damage, including in weapons or military applications.
+
+5) No Personal Information.  Do not use AMD AI to collect, process, or
+disclose personal data, including heath or sensitive personal
+information, without the necessary rights or consents.
+
+6) No Infringement.  Do not use AMD AI to generate or disseminate any
+information that infringes upon or misappropriates the intellectual
+property rights of others, including copyright, trademark, patent, and
+trade secret rights, rights to privacy, and publicity rights.
+
+7) No Malware.  Do not use AMD AI to generate or disseminate malware or
+any other content to be used for the purpose of facilitating unpermitted
+access to, or use of, computer systems or data.
+
+8) No Obfuscation.  Do not inappropriately obfuscate or fail to disclose
+to end users the presence of AI in any application in which AMD AI is
+deployed, along with any known risks or dangers of using AI without
+appropriate safeguards, oversight and human control.
+
+9) No Reliance.  Do not rely on any information generated using AMD AI
+without assessing it for accuracy, potential for harm, or other specific
+risks applicable to the use case.

README.md

@@ -0,0 +1,179 @@
+---
+license: other
+license_name: researchrail
+license_link: LICENSE
+pipeline_tag: text-generation
+library_name: transformers
+language:
+  - en
+tags:
+  - instella
+  - moe
+  - amd
+  - rocm
+---
+
+<div align="center">
+  <br>
+  <br>
+
+  <h1>Instella-MoE✨: Fully Open State-of-the-Art Mixture-of-Experts Language Model</h1>
+<a href='https://huggingface.co/collections/amd/instella-moe'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-blue'></a>
+<a href='https://rocm.blogs.amd.com/artificial-intelligence/instella-moe/README.html'><img src='https://img.shields.io/badge/Technical-Blog-red'></a>
+<a href='https://github.com/AMD-AGI/Instella-MoE'><img src='https://img.shields.io/badge/GitHub-Code-black?logo=github'></a> 
+
+</div>
+
+
+Instella-MoE is a state-of-the-art fully open Mixture-of-Experts (MoE) language model with 16 billion total parameters and 2.8 billion active parameters trained end-to-end from pre-training to RL. 
+Trained from scratch on AMD Instinct™ MI300X and MI325X GPUs using AMD's Primus framework, Instella-MoE combines a sparsely activated MoE design with architectural innovations such as Gated Multi-head Latent Attention (Gated MLA) and [FarSkip-Collective](https://github.com/AMD-AGI/FarSkip-Collective). 
+
+<div align="center">
+  <img src="images/fig_cost_performance.png" alt="Instella-MoE cost vs. performance" width="720"/>
+</div>
+
+*Figure 1: Pre-trained and Post-trained Instella-MoE model performance compared with other similar size state-of-the-art models.*
+
+## Takeaways
+
+- Instella-MoE is a new state-of-the-art fully open Mixture-of-Experts language model developed by AMD, with 16 billion total parameters and 2.8 billion active parameters per token, trained from scratch on AMD Instinct™ MI300X and MI325X GPUs.
+- The Instella-MoE model checkpoint release spans every stage of the model training pipeline, including pre-training, mid-training, long-context extension, SFT, DPO and RL.
+- Built entirely on the AMD ROCm™ software stack on top of the Primus training and Miles RL frameworks, Instella-MoE incorporates cutting-edge architecture and systems innovations—including Gated Multi-head Latent Attention (Gated MLA) and extreme communication-computation overlap through FarSkip-Collective—for efficient large-scale training and inference on AMD hardware.
+- Fully open and accessible: we provide our complete training recipe across all training stages, including training frameworks, data mixtures, intermediate checkpoints and inference code.
+
+The release includes checkpoints from the following Instella-MoE training pipeline stages as shown in Table 1 below:
+
+| Model | Stage | Description |
+| ----- | ----- | ----------- |
+| Instella-MoE-16B-A3B-Pretrain ([Link](https://huggingface.co/amd/Instella-MoE-16B-A3B-Pretrain)) | Pre-training | MoE base model trained from scratch on a large and diverse training corpus. |
+| Instella-MoE-16B-A3B-Midtrain ([Link](https://huggingface.co/amd/Instella-MoE-16B-A3B-Midtrain)) | Mid-training | Pretrained model further trained on high-quality data mixtures to refine key capabilities. |
+| Instella-MoE-16B-A3B-Base ([Link](https://huggingface.co/amd/Instella-MoE-16B-A3B-Base)) | Long-context | Long-context training to extend the model’s ability to process and reason over longer sequences. **We use this as our final base checkpoint**. |
+| Instella-MoE-16B-A3B-SFT ([Link](https://huggingface.co/amd/Instella-MoE-16B-A3B-SFT)) | SFT | Base checkpoint extended via supervised fine-tuning (SFT) to enable instruction following and chain-of-thought reasoning capabilities. |
+| Instella-MoE-16B-A3B-DPO ([Link](https://huggingface.co/amd/Instella-MoE-16B-A3B-DPO)) | DPO | Direct preference optimization (DPO) on contrastive preference data to improve model performance. |
+| Instella-MoE-16B-A3B-Think ([Link](https://huggingface.co/amd/Instella-MoE-16B-A3B-Think)) | RL | Final thinking checkpoint refined with reinforcement learning (RL) to further strengthen instruction following and overall response quality. |
+
+<em><strong>Table 1:</strong> Instella-MoE-16B-A3B models and training stages.</em>
+
+### Model Summary
+
+| Parameter | Value |
+| :---- | :---- |
+| Total Parameters | 16B |
+| Active Parameters per Token | 2.8B |
+| Decoder Layers | 27 |
+| Hidden Size | 2048 |
+| Attention Heads | 16 |
+| Number of Experts | 64 |
+| Shared Experts | 2 |
+| Activated Experts per Token | 6 |
+| Vocabulary Size | 128,896 |
+| Attention | Gated Multi-head Latent Attention (Gated MLA) |
+| MoE Connectivity | FarSkip-Collective |
+
+
+## Results
+
+### Pretraining Results
+
+
+<div align="center">
+  <img src="images/result_table_base_standard.png" alt="Base results on standard benchmarks" width="900"/>
+</div>
+
+*Table 2: Instella-MoE-16B-A3B-Base Results on Standard Benchmarks.*
+
+<div align="center">
+  <img src="images/result_table_base_long.png" alt="Base results on long-context benchmarks" width="900"/>
+</div>
+
+*Table 3: Instella-MoE-16B-A3B-Base Results on Long Context HELMET and RULER Benchmarks.*
+
+### Post-training Results
+
+
+<div align="center">
+  <img src="images/result_table_instruct.png" alt="Instella-MoE-Think post-training results" width="900"/>
+</div>
+
+*Table 4: Instella-MoE-Think results. We evaluate all models using the OLMES framework, generating up to a maximum of 32768 tokens.*
+
+## Getting Started
+
+### Example Usage
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+checkpoint = "amd/Instella-MoE-16B-A3B-Think"
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", trust_remote_code=True)
+
+prompt = [{"role": "user", "content": "What are the computational benefits of Mixture-of-Experts models?"}]
+inputs = tokenizer.apply_chat_template(
+    prompt,
+    add_generation_prompt=True,
+    return_tensors='pt'
+)
+
+tokens = model.generate(
+    inputs.to(model.device),
+    max_new_tokens=1024,
+    temperature=0.8,
+    do_sample=True
+)
+
+print(tokenizer.decode(tokens[0], skip_special_tokens=False))
+```
+For high-throughput inference with [SGLang](https://github.com/sgl-project/sglang), please refer to the setup and usage instructions in our [GitHub repository](https://github.com/AMD-AGI/Instella-MoE).
+
+## Training Details
+
+Instella-MoE was trained end-to-end on AMD Instinct™ MI300X and MI325X GPUs using the AMD ROCm™ software stack, built on top of the [Primus](https://github.com/AMD-AGI/Primus) training framework and the [Miles](https://github.com/radixark/miles) RL framework. Training proceeds through a multi-stage pipeline—pre-training, mid-training, long-context extension, SFT, DPO, and RL—with each stage progressively strengthening the model's capabilities.
+
+For the complete training recipe, including per-stage data mixtures, hyperparameters, training frameworks, and inference code, please refer to our [GitHub repository](https://github.com/AMD-AGI/Instella-MoE) and [technical blog](https://rocm.blogs.amd.com/artificial-intelligence/instella-moe/README.html).
+
+## Acknowledgements
+
+We are deeply grateful to the LLM360 team and the Miles team for their invaluable support throughout the development of our model.
+
+## License
+
+- The Instella-MoE models are licensed for academic and research purposes under a ResearchRAIL license.
+- Refer to the [LICENSE](./LICENSE) for more information.
+
+## Bias, Risks, and Limitations
+
+- The models are being released for research purposes only. They are not intended for use cases requiring high levels of factual accuracy, safety-critical applications, or health and medical applications. They must not be used to generate false information or facilitate toxic conversations.
+- Model checkpoints are made accessible without any safety promises. It is crucial for users to conduct comprehensive evaluations and implement safety filtering mechanisms as per their respective use cases.
+- It may be possible to prompt the model to generate content that may be factually inaccurate, harmful, violent, toxic, biased, or otherwise objectionable. Such content may also be generated in response to prompts that were not intended to elicit it. Users are thus requested to be aware of this and exercise caution and responsible thinking when using the model.
+- The multilingual abilities of the models have not been tested and thus may misunderstand and generate erroneous responses across different languages.
+
+
+## Contributors
+
+> **Core contributors**:
+> Jiang Liu, Sudhanshu Ranjan, Prakamya Mishra, Yonatan Dukler, Gowtham Ramesh, Zicheng Liu
+>
+> **Contributors:**
+> Jialian Wu, Ximeng Sun, Wen Xie, Chaojun Hou, Vikram Appia, Zhenyu Gu, Emad Barsoum
+
+## Citations
+
+Instella-MoE technical report is coming soon. In the meantime feel free to cite:
+
+```text
+@article{instella,
+  title={Instella: Fully Open Language Models with Stellar Performance},
+  author={Liu, Jiang and Wu, Jialian and Yu, Xiaodong and Su, Yusheng and Mishra, Prakamya and Ramesh, Gowtham and Ranjan, Sudhanshu and Manem, Chaitanya and Sun, Ximeng and Wang, Ze and Brahma, Pratik Prabhanjan and Liu, Zicheng and Barsoum, Emad},
+  journal={arXiv preprint arXiv:2511.10628},
+  year={2025}
+}
+
+@inproceedings{
+dukler2026farskipcollective,
+title={FarSkip-Collective: Unhobbling Blocking Communication in Mixture of Experts Models},
+author={Yonatan Dukler and Guihong Li and Deval Shah and Jiang Liu and Vikram Appia and Emad Barsoum},
+booktitle={Ninth Conference on Machine Learning and Systems},
+year={2026},
+url={https://openreview.net/forum?id=ruOpvLzsGV}
+}
+```

__init__.py

@@ -0,0 +1,24 @@
+from typing import TYPE_CHECKING
+
+from transformers.utils import _LazyModule
+
+_import_structure = {
+    "configuration_instella_moe": ["InstellaMoEConfig"],
+    "modeling_instella_moe": [
+        "InstellaMoEForCausalLM",
+        "InstellaMoEModel",
+        "InstellaMoEPreTrainedModel",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .configuration_instella_moe import InstellaMoEConfig
+    from .modeling_instella_moe import (
+        InstellaMoEForCausalLM,
+        InstellaMoEModel,
+        InstellaMoEPreTrainedModel,
+    )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)

config.json

@@ -0,0 +1,66 @@
+{
+  "architectures": [
+    "InstellaMoEForCausalLM"
+  ],
+  "attention_bias": false,
+  "farskip": true,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_instella_moe.InstellaMoEConfig",
+    "AutoModel": "modeling_instella_moe.InstellaMoEModel",
+    "AutoModelForCausalLM": "modeling_instella_moe.InstellaMoEForCausalLM"
+  },
+  "aux_loss_alpha": 0.001,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "ep_size": 1,
+  "first_k_dense_replace": 1,
+  "gated_attention": true,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 10944,
+  "kv_lora_rank": 512,
+  "qk_layernorm": true,
+  "max_position_embeddings": 65536,
+  "model_type": "deepseek_v3",
+  "moe_intermediate_size": 1408,
+  "moe_layer_freq": 1,
+  "n_group": 1,
+  "n_routed_experts": 64,
+  "n_shared_experts": 2,
+  "norm_topk_prob": true,
+  "num_attention_heads": 16,
+  "num_experts_per_tok": 6,
+  "num_hidden_layers": 27,
+  "num_key_value_heads": 16,
+  "num_nextn_predict_layers": 1,
+  "q_lora_rank": null,
+  "pretraining_tp": 1,
+  "qk_head_dim": 128,
+  "qk_nope_head_dim": 96,
+  "qk_rope_head_dim": 32,
+  "rope_interleave": true,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn"
+  },
+  "rope_theta": 8000000,
+  "routed_scaling_factor": 2.5,
+  "scoring_func": "sigmoid",
+  "seq_aux": true,
+  "tie_word_embeddings": false,
+  "topk_group": 1,
+  "topk_method": "noaux_tc",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "v_head_dim": 128,
+  "vocab_size": 128896
+}

configuration_instella_moe.py

@@ -0,0 +1,237 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+class InstellaMoEConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`InstellaMoEModel`]. It is used to instantiate an DeepSeek
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the DeepSeek-V3.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 129280):
+            Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`InstellaMoEModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        moe_intermediate_size (`int`, *optional*, defaults to 1407):
+            Dimension of the MoE representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_nextn_predict_layers (`int`, *optional*, defaults to 1):
+            Number of nextn predict layers in the InstellaMoE Model.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        n_shared_experts (`int`, *optional*, defaults to None):
+            Number of shared experts, None means dense model.
+        n_routed_experts (`int`, *optional*, defaults to None):
+            Number of routed experts, None means dense model.
+        routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor or routed experts.
+        topk_method (`str`, *optional*, defaults to `gready`):
+            Topk method used in routed gate.
+        n_group (`int`, *optional*, defaults to None):
+            Number of groups for routed experts.
+        topk_group (`int`, *optional*, defaults to None):
+            Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups).
+        num_experts_per_tok (`int`, *optional*, defaults to None):
+            Number of selected experts, None means dense model.
+        moe_layer_freq (`int`, *optional*, defaults to 1):
+            The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers.
+        first_k_dense_replace (`int`, *optional*, defaults to 0):
+            Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head).
+                                                            \--k dense layers--/
+        norm_topk_prob (`bool`, *optional*, defaults to False):
+            Whether to normalize the weights of the routed experts.
+        scoring_func (`str`, *optional*, defaults to 'softmax'):
+            Method of computing expert weights.
+        aux_loss_alpha (`float`, *optional*, defaults to 0.001):
+            Auxiliary loss weight coefficient.
+        seq_aux = (`bool`, *optional*, defaults to True):
+            Whether to compute the auxiliary loss for each individual sample.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            End of stream token id.
+        pretraining_tp (`int`, *optional*, defaults to 1):
+            Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
+            document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+            necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+            issue](https://github.com/pytorch/pytorch/issues/76232).
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import InstellaMoEModel, InstellaMoEConfig
+
+    >>> # Initializing a Deepseek-V3 style configuration
+    >>> configuration = InstellaMoEConfig()
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "deepseek_v3"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=129280,
+        hidden_size=7168,
+        intermediate_size=18432,
+        moe_intermediate_size = 2048,
+        num_hidden_layers=61,
+        num_nextn_predict_layers=1,
+        num_attention_heads=128,
+        num_key_value_heads=128,
+        n_shared_experts = 1,
+        n_routed_experts = 256,
+        ep_size = 1,
+        routed_scaling_factor = 2.5,
+        kv_lora_rank = 512,
+        q_lora_rank = None,
+        qk_rope_head_dim = 64,
+        qk_layernorm = True,
+        v_head_dim = 128,
+        qk_nope_head_dim = 128,
+        qk_head_dim = None,
+        rope_interleave = True,
+        topk_method = 'noaux_tc',
+        n_group = 8,
+        topk_group = 4,
+        num_experts_per_tok = 8,
+        moe_layer_freq = 1,
+        first_k_dense_replace = 3,
+        gated_attention = True,
+        norm_topk_prob = True,
+        scoring_func = 'sigmoid',
+        aux_loss_alpha = 0.001,
+        seq_aux = True,
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=0,
+        eos_token_id=1,
+        pretraining_tp=1,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        farskip=False,
+        tp_split=None,
+        farskip_start_idx=0,
+        farskip_end_idx=1e4,
+        apply_all_reduce=False,
+        attn_only_farskip=False,
+        mlp_only_farskip=False,
+        verbose=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+        self.num_attention_heads = num_attention_heads
+        self.n_shared_experts = n_shared_experts
+        self.n_routed_experts = n_routed_experts
+        self.ep_size = ep_size
+        self.routed_scaling_factor = routed_scaling_factor
+        self.kv_lora_rank = kv_lora_rank
+        self.q_lora_rank = q_lora_rank
+        self.qk_layernorm = qk_layernorm
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_head_dim = qk_head_dim if qk_head_dim is not None else qk_nope_head_dim + qk_rope_head_dim
+        self.rope_interleave = rope_interleave
+        # yarn rope init reads config.head_dim (falls back to hidden_size//num_heads
+        # otherwise). MLA rope operates on qk_rope_head_dim, so pin it here to match.
+        self.head_dim = qk_rope_head_dim
+        self.topk_method = topk_method
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.num_experts_per_tok = num_experts_per_tok
+        self.moe_layer_freq = moe_layer_freq
+        self.first_k_dense_replace = first_k_dense_replace
+        self.norm_topk_prob = norm_topk_prob
+        self.scoring_func = scoring_func
+        self.aux_loss_alpha = aux_loss_alpha
+        self.seq_aux = seq_aux
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.pretraining_tp = pretraining_tp
+        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.farskip = farskip
+        self.tp_split = tp_split
+        self.farskip_start_idx = farskip_start_idx
+        self.farskip_end_idx = farskip_end_idx
+        self.apply_all_reduce = apply_all_reduce
+        self.attn_only_farskip = attn_only_farskip
+        self.mlp_only_farskip = mlp_only_farskip
+        self.verbose = verbose
+        self.gated_attention = gated_attention
+
+        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,
+        )

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "do_sample": true,
+  "temperature": 0.3,
+  "top_p": 0.95,
+  "transformers_version": "4.39.3"
+}

merge_metadata.json

@@ -0,0 +1,19 @@
+{
+  "checkpoints": [
+    "/wekafs/prakmish/pretain_checkpoints_long_hf/exp_instella_16b_gated_farskip_overlap_midtraining_dolmino_100B_2e-4_iter_0021000-hf-with-mtp",
+    "/wekafs/prakmish/pretain_checkpoints_long_hf/exp_instella_16b_gated_farskip_overlap_midtraining_dolmino_100B_v1_2e-4_iter_0021000-hf-with-mtp",
+    "/wekafs/prakmish/pretain_checkpoints_long_hf/exp_instella_16b_gated_farskip_overlap_midtraining_dolmino_100B_v2_2e-4_iter_0022125-hf-with-mtp"
+  ],
+  "weights_before_normalization": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "normalized_weights": [
+    0.3333333333333333,
+    0.3333333333333333,
+    0.3333333333333333
+  ],
+  "merge_dtype": "float32",
+  "total_tensors": 5568
+}

model-00000-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59a5edd7fa807e7a283d012b8993a7fa579da73eafe6921fe5355625d2c983c6
+size 5363712448

model-00001-of-00006.safetensors

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

model-00002-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82baf52fa5264d854b5a00e36212b259eac260ee3deb36b232185a4f7096b31d
+size 5365546904

model-00003-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fe8c67cb1c8777f7b2c30b4e6a317379faf0ac20ece411f82a8d3f77adaf942
+size 5364822728

model-00004-of-00006.safetensors

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

model-00005-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:024dad22bb70e39a5b780490d0a8378be75519b14ee257b25e3750fd24f727a1
+size 4907096096

modeling_instella_moe.py

@@ -0,0 +1,397 @@
+"""
+Modular Instella-MoE for transformers 4.57.1.
+
+This file imports the numerically-unchanged building blocks from the installed
+`transformers.models.deepseek_v3` package and defines ONLY the classes that carry
+a FarSkip or gated-attention delta:
+
+  * InstellaMoEForCausalLM
+  * InstellaMoEPreTrainedModel
+  * InstellaMoEModel      - unwraps the residual tuple before the final norm
+  * FarSkipDecoderLayer - tuple-residual (residual, residual_no_routed) dataflow
+  * FarSkipMoE        - returns (routed, shared) separately (FarSkip needs both)
+  * MLAGatedAttention  - adds sigmoid `gate_proj` before `o_proj`
+"""
+
+from typing import Optional, Union
+
+import torch
+from torch import nn
+
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import check_model_inputs
+
+from transformers.models.deepseek_v3.modeling_deepseek_v3 import (
+    DeepseekV3RMSNorm,
+    DeepseekV3RotaryEmbedding,
+    DeepseekV3MLP,
+    DeepseekV3MoE,
+    DeepseekV3TopkRouter,
+    DeepseekV3Attention,
+    apply_rotary_pos_emb,
+    apply_rotary_pos_emb_interleave,
+    eager_attention_forward,
+)
+
+from .configuration_instella_moe import InstellaMoEConfig
+
+
+class MLAGatedAttention(DeepseekV3Attention):
+    """DeepSeek-V3 MLA with optional gated attention (attn_output * sigmoid(gate_proj(x)))."""
+
+    def __init__(self, config: InstellaMoEConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.gated_attention = getattr(config, "gated_attention", False)
+        if self.gated_attention:
+            self.gate_proj = nn.Linear(
+                config.hidden_size, self.num_heads * self.v_head_dim, bias=False
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, seq_length = hidden_states.shape[:-1]
+        query_shape = (batch_size, seq_length, -1, self.qk_head_dim)
+        key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
+
+        if self.q_lora_rank is None:
+            q_states = self.q_proj(hidden_states)
+        else:
+            q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+        q_states = q_states.view(query_shape).transpose(1, 2)
+        q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+        k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+
+        k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)
+        k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
+
+        cos, sin = position_embeddings
+        if self.config.rope_interleave:
+            q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
+        else:
+            q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
+        k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
+
+        query_states = torch.cat((q_pass, q_rot), dim=-1)
+        key_states = torch.cat((k_pass, k_rot), dim=-1)
+
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
+            value_states = nn.functional.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            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,
+        )
+
+        if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:
+            attn_output = attn_output[:, :, :, : self.v_head_dim]
+
+        attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
+
+        if self.gated_attention:
+            attn_output = attn_output * torch.sigmoid(self.gate_proj(hidden_states))
+
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class FarSkipMoE(DeepseekV3MoE):
+    """FarSkip-Collective MoE that returns routed and shared outputs separately so FarSkip can route them
+    into the two residual streams independently."""
+
+    def forward(self, hidden_states):
+        residuals = hidden_states
+        orig_shape = hidden_states.shape
+        topk_indices, topk_weights = self.gate(hidden_states)
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        routed = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape)
+        shared = self.shared_experts(residuals)
+        # First element is the full MoE output (routed + shared) so the main residual
+        # stream stays numerically identical to stock DeepSeek-V3; `shared` is returned
+        # separately so FarSkip can build the routed-free residual stream.
+        return routed + shared, shared
+
+
+class FarSkipDecoderLayer(GradientCheckpointingLayer):
+    """
+    FarSkip-Collective connectivity decoder layer
+    """
+    def __init__(self, config: InstellaMoEConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = MLAGatedAttention(config=config, layer_idx=layer_idx)
+
+        if layer_idx >= config.first_k_dense_replace:
+            self.mlp = FarSkipMoE(config)
+        else:
+            self.mlp = DeepseekV3MLP(config)
+
+        self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.config = config
+        self.layer_idx = layer_idx
+        self.farskip = (
+            config.farskip
+            and layer_idx >= config.farskip_start_idx
+            and layer_idx <= min(config.farskip_end_idx, config.num_hidden_layers - 1)
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        if self.farskip:
+            if not isinstance(hidden_states, tuple):  # first farskip layer
+                residual = hidden_states
+                input_to_attn = hidden_states
+                input_to_mlp = hidden_states
+            else:
+                residual = hidden_states[0]
+                input_to_attn = hidden_states[1]
+                input_to_mlp = residual
+            if self.config.attn_only_farskip:
+                input_to_mlp = None
+            if self.config.mlp_only_farskip:
+                input_to_attn = residual
+        else:
+            if isinstance(hidden_states, tuple):
+                hidden_states = hidden_states[0]
+            residual = hidden_states
+            input_to_attn = hidden_states
+            input_to_mlp = None
+
+        input_to_attn = self.input_layernorm(input_to_attn)
+
+        attn_output, _ = self.self_attn(
+            hidden_states=input_to_attn,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        residual = residual + attn_output
+
+        if input_to_mlp is None:
+            input_to_mlp = residual
+        input_to_mlp = self.post_attention_layernorm(input_to_mlp)
+
+        if isinstance(self.mlp, FarSkipMoE):
+            mlp_output, mlp_shared_output = self.mlp(input_to_mlp)
+            residual_no_routed = residual + mlp_shared_output
+            residual = residual + mlp_output
+            # residual_no_routed is combine-free and feeds the next block's attention
+            hidden_states = (residual, residual_no_routed)
+        else:
+            hidden_states = residual + self.mlp(input_to_mlp)
+
+        return hidden_states
+
+
+@auto_docstring
+class InstellaMoEPreTrainedModel(PreTrainedModel):
+    config: InstellaMoEConfig
+    config_class = InstellaMoEConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["FarSkipDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = False
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": FarSkipDecoderLayer,
+        "attentions": MLAGatedAttention,
+    }
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, DeepseekV3TopkRouter):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+
+
+@auto_docstring
+class InstellaMoEModel(InstellaMoEPreTrainedModel):
+    def __init__(self, config: InstellaMoEConfig):
+        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(
+            [FarSkipDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = DeepseekV3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        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)
+
+        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
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        if isinstance(hidden_states, tuple):
+            hidden_states = hidden_states[0]  # routed-inclusive residual stream
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+@auto_docstring
+class InstellaMoEForCausalLM(InstellaMoEPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = InstellaMoEModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        outputs: BaseModelOutputWithPast = 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,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        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,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "InstellaMoEPreTrainedModel",
+    "InstellaMoEModel",
+    "InstellaMoEForCausalLM",
+]

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,35 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 131072,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "unk_token": null,
+  "tokenizer_class": "LlamaTokenizerFast",
+  "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true) %}{%- for message in messages %}{%- if message['role'] == 'system' %}{%- if ns.is_first_sp %}{% set ns.system_prompt = ns.system_prompt + message['content'] %}{% set ns.is_first_sp = false %}{%- else %}{% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}{%- endif %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{{'<｜Assistant｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}"
+}