Duplicate from arcee-ai/Trinity-Large-Preview-NVFP4

Co-authored-by: Anneketh Vij <annekethvij@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text
+.quant_summary.txt filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.quant_summary.txt

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

LICENSE

@@ -0,0 +1,49 @@
+OpenMDW License Agreement, version 1.1 (OpenMDW-1.1)
+
+By exercising rights granted to you under this agreement, you accept and agree
+to its terms.
+
+As used in this agreement, "Model Materials" means the materials provided to
+you under this agreement, consisting of: (1) one or more machine learning
+models (including architecture and parameters); and (2) all related artifacts
+(including associated data, documentation and software) that are provided to
+you hereunder.
+
+Subject to your compliance with this agreement, permission is hereby granted,
+free of charge, to deal in the Model Materials without restriction, including
+under all copyright, patent, database, and trade secret rights included or
+embodied therein.
+
+If you distribute any portion of the Model Materials, you shall retain in your
+distribution (1) a copy of this agreement, and (2) all copyright notices and
+other notices of origin included in the Model Materials that are applicable to
+your distribution.
+
+If you file, maintain, or voluntarily participate in a lawsuit against any
+person or entity asserting that the Model Materials directly or indirectly
+infringe any patent or copyright, then all rights and grants made to you
+hereunder are terminated, unless that lawsuit was in response to a
+corresponding lawsuit first brought against you.
+
+This agreement does not impose any restrictions or obligations with respect to
+any use, modification, or sharing of any outputs generated by using the Model
+Materials.
+
+THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE
+ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL
+TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW.
+
+YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY
+APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION
+ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL
+MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS
+REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE
+DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR
+ANYTHING INCORPORATED OR EMBODIED THEREIN.
+
+IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE
+USE THEREOF OR OTHER DEALINGS THEREIN.

README.md

@@ -0,0 +1,181 @@
+---
+license: other
+language:
+- en
+- es
+- fr
+- de
+- it
+- pt
+- ru
+- ar
+- hi
+- ko
+- zh
+library_name: transformers
+base_model:
+- arcee-ai/Trinity-Large-Preview
+base_model_relation: quantized
+tags:
+- moe
+- nvfp4
+- modelopt
+- blackwell
+- vllm
+license_link: LICENSE
+license_name: openmdw-1.1
+---
+
+<div align="center">
+  <picture>
+    <img
+      src="https://cdn-uploads.huggingface.co/production/uploads/6435718aaaef013d1aec3b8b/i-v1KyAMOW_mgVGeic9WJ.png"
+      alt="Arcee Trinity Large"
+      style="max-width: 100%; height: auto;"
+    >
+  </picture>
+</div>
+<hr>
+
+# Trinity-Large-Preview-NVFP4
+
+**This repository contains the NVFP4 quantized weights of Trinity-Large-Preview for deployment on NVIDIA Blackwell GPUs.**
+
+Trinity-Large-Preview is a 398B-parameter sparse Mixture-of-Experts (MoE) model with approximately 13B active parameters per token. It is the largest model in Arcee AI's Trinity family, trained on more than 17 trillion tokens and delivering frontier-level performance with strong long-context comprehension.
+Trinity-Large-Preview is a lightly post-trained model based on Trinity-Large-Base.
+
+Try it at [chat.arcee.ai](http://chat.arcee.ai/)
+
+More details on the training of Trinity Large are available in the [technical report](https://arxiv.org/abs/2602.17004).
+
+***
+
+## Quantization Details
+
+- **Scheme:** NVFP4 (`nvfp4_mlp_only` — MLP/expert weights only, attention remains BF16)
+- **Tool:** [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer)
+- **Calibration:** 512 samples, seq_length=2048, all-expert calibration enabled
+- **KV cache:** Not quantized
+
+***
+
+## Model Variants
+
+The Trinity Large family consists of three checkpoints from the same training run:
+
+- **[Trinity-Large-Preview](https://huggingface.co/arcee-ai/Trinity-Large-Preview)**: Lightly post-trained, chat-ready model undergoing active RL
+- **[Trinity-Large-TrueBase](https://huggingface.co/arcee-ai/Trinity-Large-TrueBase)**: 10T-token pre-anneal pretraining checkpoint
+- **[Trinity-Large-Base](https://huggingface.co/arcee-ai/Trinity-Large-Base)**: Full 17T-token pretrained foundation model with mid-training anneals
+
+## Architecture
+
+| Hyperparameter | Value |
+|:---|:---:|
+| Total parameters | ~398B |
+| Active parameters per token | ~13B |
+| Experts | 256 (1 shared) |
+| Active experts | 4 |
+| Routing strategy | 4-of-256 (1.56% sparsity) |
+| Dense layers | 6 |
+| Pretraining context length | 8,192 |
+| Context length after extension | 512k |
+| Architecture | Sparse MoE (AfmoeForCausalLM) |
+
+## Benchmarks
+
+| Benchmark | Llama 4 Maverick | Trinity-Large Preview |
+|-----------|------------------|----------------------|
+| MMLU | 85.5 | 87.2 |
+| MMLU-Pro | 80.5 | 75.2 |
+| GPQA-Diamond | 69.8 | 63.3 |
+| AIME 2025 | 19.3 | 24.0 |
+
+***
+
+## Running with vLLM
+
+Requires [vLLM](https://github.com/vllm-project/vllm) >= 0.18.0. Native FP4 compute requires Blackwell GPUs; older GPUs fall back to Marlin weight decompression automatically.
+
+### Blackwell GPUs (B200/B300/GB300) — Docker (recommended)
+
+```bash
+docker run --runtime nvidia --gpus all -p 8000:8000 \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  vllm/vllm-openai:v0.18.0-cu130 \
+  arcee-ai/Trinity-Large-Preview-NVFP4 \
+  --trust-remote-code \
+  --tensor-parallel-size 8 \
+  --gpu-memory-utilization 0.90 \
+  --max-model-len 8192
+```
+
+### Hopper GPUs (H100/H200) and others
+
+```bash
+vllm serve arcee-ai/Trinity-Large-Preview-NVFP4 \
+  --trust-remote-code \
+  --tensor-parallel-size 8 \
+  --gpu-memory-utilization 0.90 \
+  --max-model-len 8192 \
+  --host 0.0.0.0 \
+  --port 8000
+```
+
+ **Note (Blackwell pip installs):** If installing vLLM via pip on Blackwell rather than using Docker, native FP4 kernels may produce incorrect output due to package version mismatches. As a workaround, force the Marlin backend:
+
+ ```bash
+ export VLLM_NVFP4_GEMM_BACKEND=marlin
+
+ vllm serve arcee-ai/Trinity-Large-Preview-NVFP4 \
+   --trust-remote-code \
+   --tensor-parallel-size 8 \
+   --moe-backend marlin \
+   --gpu-memory-utilization 0.90 \
+   --max-model-len 8192 \
+   --host 0.0.0.0 \
+   --port 8000
+```
+
+Marlin decompresses FP4 weights to BF16 for compute, providing the full memory compression benefit but not native FP4 compute speedup. On Hopper GPUs (H100/H200), Marlin is selected automatically and no extra flags are needed.
+
+### API
+
+Available on OpenRouter:
+
+```bash
+curl -X POST "https://openrouter.ai/v1/chat/completions" \
+  -H "Authorization: Bearer $OPENROUTER_API_KEY" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "arcee-ai/trinity-large-preview",
+    "messages": [
+      {
+        "role": "user",
+        "content": "What are some fun things to do in New York?"
+      }
+    ]
+  }'
+```
+
+***
+
+## License
+
+Trinity-Large-Preview-NVFP4 is released under the OpenMDW License, version 1.1 (OpenMDW-1.1).
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{singh2026arceetrinity,
+  title        = {Arcee Trinity Large Technical Report},
+  author       = {Varun Singh and Lucas Krauss and Sami Jaghouar and Matej Sirovatka and Charles Goddard and Fares Obied and Jack Min Ong and Jannik Straube and Fern and Aria Harley and Conner Stewart and Colin Kealty and Maziyar Panahi and Simon Kirsten and Anushka Deshpande and Anneketh Vij and Arthur Bresnu and Pranav Veldurthi and Raghav Ravishankar and Hardik Bishnoi and DatologyAI Team and Arcee AI Team and Prime Intellect Team and Mark McQuade and Johannes Hagemann and Lucas Atkins},
+  year         = {2026},
+  eprint       = {2602.17004},
+  archivePrefix= {arXiv},
+  primaryClass = {cs.LG},
+  doi          = {10.48550/arXiv.2602.17004},
+  url          = {https://arxiv.org/abs/2602.17004}
+}
+```

chat_template.jinja

@@ -0,0 +1,65 @@
+{{ bos_token }}{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {% generation %}
+        {{- content}}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>' }}
+        {% endgeneration%}
+        {{- '\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,252 @@
+{
+    "architectures": [
+        "AfmoeForCausalLM"
+    ],
+    "attention_dropout": 0.0,
+    "auto_map": {
+        "AutoConfig": "configuration_afmoe.AfmoeConfig",
+        "AutoModel": "modeling_afmoe.AfmoeModel",
+        "AutoModelForCausalLM": "modeling_afmoe.AfmoeForCausalLM"
+    },
+    "dtype": "bfloat16",
+    "global_attn_every_n_layers": 4,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 3072,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "layer_types": [
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention"
+    ],
+    "load_balance_coeff": 5e-05,
+    "max_position_embeddings": 262144,
+    "model_type": "afmoe",
+    "moe_intermediate_size": 3072,
+    "mup_enabled": true,
+    "n_group": 1,
+    "num_attention_heads": 48,
+    "num_dense_layers": 6,
+    "num_expert_groups": 1,
+    "num_experts": 256,
+    "num_experts_per_tok": 4,
+    "num_hidden_layers": 60,
+    "num_key_value_heads": 8,
+    "num_limited_groups": 1,
+    "num_shared_experts": 1,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "rope_theta": 10000,
+    "route_norm": true,
+    "route_scale": 2.448,
+    "score_func": "sigmoid",
+    "sliding_window": 4096,
+    "tie_word_embeddings": false,
+    "topk_group": 1,
+    "transformers_version": "4.57.6",
+    "use_cache": true,
+    "use_grouped_mm": true,
+    "vocab_size": 200192,
+    "quantization_config": {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {
+                    "dynamic": false,
+                    "num_bits": 4,
+                    "type": "float",
+                    "group_size": 16
+                },
+                "weights": {
+                    "dynamic": false,
+                    "num_bits": 4,
+                    "type": "float",
+                    "group_size": 16
+                },
+                "targets": [
+                    "Linear"
+                ]
+            }
+        },
+        "ignore": [
+            "lm_head",
+            "model.layers.0.self_attn*",
+            "model.layers.1.self_attn*",
+            "model.layers.10.mlp.router*",
+            "model.layers.10.self_attn*",
+            "model.layers.11.mlp.router*",
+            "model.layers.11.self_attn*",
+            "model.layers.12.mlp.router*",
+            "model.layers.12.self_attn*",
+            "model.layers.13.mlp.router*",
+            "model.layers.13.self_attn*",
+            "model.layers.14.mlp.router*",
+            "model.layers.14.self_attn*",
+            "model.layers.15.mlp.router*",
+            "model.layers.15.self_attn*",
+            "model.layers.16.mlp.router*",
+            "model.layers.16.self_attn*",
+            "model.layers.17.mlp.router*",
+            "model.layers.17.self_attn*",
+            "model.layers.18.mlp.router*",
+            "model.layers.18.self_attn*",
+            "model.layers.19.mlp.router*",
+            "model.layers.19.self_attn*",
+            "model.layers.2.self_attn*",
+            "model.layers.20.mlp.router*",
+            "model.layers.20.self_attn*",
+            "model.layers.21.mlp.router*",
+            "model.layers.21.self_attn*",
+            "model.layers.22.mlp.router*",
+            "model.layers.22.self_attn*",
+            "model.layers.23.mlp.router*",
+            "model.layers.23.self_attn*",
+            "model.layers.24.mlp.router*",
+            "model.layers.24.self_attn*",
+            "model.layers.25.mlp.router*",
+            "model.layers.25.self_attn*",
+            "model.layers.26.mlp.router*",
+            "model.layers.26.self_attn*",
+            "model.layers.27.mlp.router*",
+            "model.layers.27.self_attn*",
+            "model.layers.28.mlp.router*",
+            "model.layers.28.self_attn*",
+            "model.layers.29.mlp.router*",
+            "model.layers.29.self_attn*",
+            "model.layers.3.self_attn*",
+            "model.layers.30.mlp.router*",
+            "model.layers.30.self_attn*",
+            "model.layers.31.mlp.router*",
+            "model.layers.31.self_attn*",
+            "model.layers.32.mlp.router*",
+            "model.layers.32.self_attn*",
+            "model.layers.33.mlp.router*",
+            "model.layers.33.self_attn*",
+            "model.layers.34.mlp.router*",
+            "model.layers.34.self_attn*",
+            "model.layers.35.mlp.router*",
+            "model.layers.35.self_attn*",
+            "model.layers.36.mlp.router*",
+            "model.layers.36.self_attn*",
+            "model.layers.37.mlp.router*",
+            "model.layers.37.self_attn*",
+            "model.layers.38.mlp.router*",
+            "model.layers.38.self_attn*",
+            "model.layers.39.mlp.router*",
+            "model.layers.39.self_attn*",
+            "model.layers.4.self_attn*",
+            "model.layers.40.mlp.router*",
+            "model.layers.40.self_attn*",
+            "model.layers.41.mlp.router*",
+            "model.layers.41.self_attn*",
+            "model.layers.42.mlp.router*",
+            "model.layers.42.self_attn*",
+            "model.layers.43.mlp.router*",
+            "model.layers.43.self_attn*",
+            "model.layers.44.mlp.router*",
+            "model.layers.44.self_attn*",
+            "model.layers.45.mlp.router*",
+            "model.layers.45.self_attn*",
+            "model.layers.46.mlp.router*",
+            "model.layers.46.self_attn*",
+            "model.layers.47.mlp.router*",
+            "model.layers.47.self_attn*",
+            "model.layers.48.mlp.router*",
+            "model.layers.48.self_attn*",
+            "model.layers.49.mlp.router*",
+            "model.layers.49.self_attn*",
+            "model.layers.5.self_attn*",
+            "model.layers.50.mlp.router*",
+            "model.layers.50.self_attn*",
+            "model.layers.51.mlp.router*",
+            "model.layers.51.self_attn*",
+            "model.layers.52.mlp.router*",
+            "model.layers.52.self_attn*",
+            "model.layers.53.mlp.router*",
+            "model.layers.53.self_attn*",
+            "model.layers.54.mlp.router*",
+            "model.layers.54.self_attn*",
+            "model.layers.55.mlp.router*",
+            "model.layers.55.self_attn*",
+            "model.layers.56.mlp.router*",
+            "model.layers.56.self_attn*",
+            "model.layers.57.mlp.router*",
+            "model.layers.57.self_attn*",
+            "model.layers.58.mlp.router*",
+            "model.layers.58.self_attn*",
+            "model.layers.59.mlp.router*",
+            "model.layers.59.self_attn*",
+            "model.layers.6.mlp.router*",
+            "model.layers.6.self_attn*",
+            "model.layers.7.mlp.router*",
+            "model.layers.7.self_attn*",
+            "model.layers.8.mlp.router*",
+            "model.layers.8.self_attn*",
+            "model.layers.9.mlp.router*",
+            "model.layers.9.self_attn*"
+        ],
+        "quant_algo": "NVFP4",
+        "producer": {
+            "name": "modelopt",
+            "version": "0.43.0rc2.dev24+gd69886413.d20260326"
+        },
+        "quant_method": "modelopt"
+    }
+}

configuration_afmoe.py

@@ -0,0 +1,133 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.configuration_utils import layer_type_validation
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+class AfmoeConfig(PretrainedConfig):
+    """
+    n_group (`int`, *optional*, defaults to 1):
+            Number of groups for routed experts.
+    topk_group (`int`, *optional*, defaults to 1):
+        Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups).
+    """
+    model_type = "afmoe"
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        num_hidden_layers: int = 32,
+        vocab_size: int = 200192,
+        hidden_size: int = 2048,
+        intermediate_size: int = 6144,
+        moe_intermediate_size=1408,
+        num_dense_layers=1,
+        num_attention_heads=16,
+        num_key_value_heads=None,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=16384,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        num_experts=64,
+        num_experts_per_tok=6,
+        num_shared_experts=2,
+        num_expert_groups=1,
+        num_limited_groups=1,
+        score_func="sigmoid",
+        route_norm=True,
+        route_scale=1.0,
+        global_attn_every_n_layers=4,
+        sliding_window=1024,
+        mup_enabled=False,
+        layer_types=None,
+        attention_dropout: float = 0.0,
+        n_group: int = 1,
+        topk_group: int = 1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_dense_layers = num_dense_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        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
+
+
+        # MoE specific
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_expert_groups = num_expert_groups
+        self.num_limited_groups = num_limited_groups
+        self.score_func = score_func
+        self.route_norm = route_norm
+        self.route_scale = route_scale
+
+
+        # Attention specific
+        self.attention_dropout = attention_dropout
+        self.global_attn_every_n_layers = global_attn_every_n_layers
+        self.sliding_window = sliding_window
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            self.layer_types = [
+                "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types)
+
+        # muP specific
+        self.mup_enabled = mup_enabled
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+
+
+        # Validate rope configs
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["AfmoeConfig"]

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 3,
+  "pad_token_id": 12,
+  "transformers_version": "4.57.3",
+  "temperature": 0.8,
+  "top_p": 0.8
+}

hf_quant_config.json

@@ -0,0 +1,128 @@
+{
+    "producer": {
+        "name": "modelopt",
+        "version": "0.43.0rc2.dev24+gd69886413.d20260326"
+    },
+    "quantization": {
+        "quant_algo": "NVFP4",
+        "kv_cache_quant_algo": null,
+        "group_size": 16,
+        "exclude_modules": [
+            "lm_head",
+            "model.layers.0.self_attn*",
+            "model.layers.1.self_attn*",
+            "model.layers.10.mlp.router*",
+            "model.layers.10.self_attn*",
+            "model.layers.11.mlp.router*",
+            "model.layers.11.self_attn*",
+            "model.layers.12.mlp.router*",
+            "model.layers.12.self_attn*",
+            "model.layers.13.mlp.router*",
+            "model.layers.13.self_attn*",
+            "model.layers.14.mlp.router*",
+            "model.layers.14.self_attn*",
+            "model.layers.15.mlp.router*",
+            "model.layers.15.self_attn*",
+            "model.layers.16.mlp.router*",
+            "model.layers.16.self_attn*",
+            "model.layers.17.mlp.router*",
+            "model.layers.17.self_attn*",
+            "model.layers.18.mlp.router*",
+            "model.layers.18.self_attn*",
+            "model.layers.19.mlp.router*",
+            "model.layers.19.self_attn*",
+            "model.layers.2.self_attn*",
+            "model.layers.20.mlp.router*",
+            "model.layers.20.self_attn*",
+            "model.layers.21.mlp.router*",
+            "model.layers.21.self_attn*",
+            "model.layers.22.mlp.router*",
+            "model.layers.22.self_attn*",
+            "model.layers.23.mlp.router*",
+            "model.layers.23.self_attn*",
+            "model.layers.24.mlp.router*",
+            "model.layers.24.self_attn*",
+            "model.layers.25.mlp.router*",
+            "model.layers.25.self_attn*",
+            "model.layers.26.mlp.router*",
+            "model.layers.26.self_attn*",
+            "model.layers.27.mlp.router*",
+            "model.layers.27.self_attn*",
+            "model.layers.28.mlp.router*",
+            "model.layers.28.self_attn*",
+            "model.layers.29.mlp.router*",
+            "model.layers.29.self_attn*",
+            "model.layers.3.self_attn*",
+            "model.layers.30.mlp.router*",
+            "model.layers.30.self_attn*",
+            "model.layers.31.mlp.router*",
+            "model.layers.31.self_attn*",
+            "model.layers.32.mlp.router*",
+            "model.layers.32.self_attn*",
+            "model.layers.33.mlp.router*",
+            "model.layers.33.self_attn*",
+            "model.layers.34.mlp.router*",
+            "model.layers.34.self_attn*",
+            "model.layers.35.mlp.router*",
+            "model.layers.35.self_attn*",
+            "model.layers.36.mlp.router*",
+            "model.layers.36.self_attn*",
+            "model.layers.37.mlp.router*",
+            "model.layers.37.self_attn*",
+            "model.layers.38.mlp.router*",
+            "model.layers.38.self_attn*",
+            "model.layers.39.mlp.router*",
+            "model.layers.39.self_attn*",
+            "model.layers.4.self_attn*",
+            "model.layers.40.mlp.router*",
+            "model.layers.40.self_attn*",
+            "model.layers.41.mlp.router*",
+            "model.layers.41.self_attn*",
+            "model.layers.42.mlp.router*",
+            "model.layers.42.self_attn*",
+            "model.layers.43.mlp.router*",
+            "model.layers.43.self_attn*",
+            "model.layers.44.mlp.router*",
+            "model.layers.44.self_attn*",
+            "model.layers.45.mlp.router*",
+            "model.layers.45.self_attn*",
+            "model.layers.46.mlp.router*",
+            "model.layers.46.self_attn*",
+            "model.layers.47.mlp.router*",
+            "model.layers.47.self_attn*",
+            "model.layers.48.mlp.router*",
+            "model.layers.48.self_attn*",
+            "model.layers.49.mlp.router*",
+            "model.layers.49.self_attn*",
+            "model.layers.5.self_attn*",
+            "model.layers.50.mlp.router*",
+            "model.layers.50.self_attn*",
+            "model.layers.51.mlp.router*",
+            "model.layers.51.self_attn*",
+            "model.layers.52.mlp.router*",
+            "model.layers.52.self_attn*",
+            "model.layers.53.mlp.router*",
+            "model.layers.53.self_attn*",
+            "model.layers.54.mlp.router*",
+            "model.layers.54.self_attn*",
+            "model.layers.55.mlp.router*",
+            "model.layers.55.self_attn*",
+            "model.layers.56.mlp.router*",
+            "model.layers.56.self_attn*",
+            "model.layers.57.mlp.router*",
+            "model.layers.57.self_attn*",
+            "model.layers.58.mlp.router*",
+            "model.layers.58.self_attn*",
+            "model.layers.59.mlp.router*",
+            "model.layers.59.self_attn*",
+            "model.layers.6.mlp.router*",
+            "model.layers.6.self_attn*",
+            "model.layers.7.mlp.router*",
+            "model.layers.7.self_attn*",
+            "model.layers.8.mlp.router*",
+            "model.layers.8.self_attn*",
+            "model.layers.9.mlp.router*",
+            "model.layers.9.self_attn*"
+        ]
+    }
+}

model-00001-of-00024.safetensors

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

model-00002-of-00024.safetensors

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

model-00003-of-00024.safetensors

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

model-00004-of-00024.safetensors

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

model-00005-of-00024.safetensors

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

model-00006-of-00024.safetensors

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

model-00007-of-00024.safetensors

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

model-00008-of-00024.safetensors

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

model-00009-of-00024.safetensors

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

model-00010-of-00024.safetensors

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

model-00011-of-00024.safetensors

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

model-00012-of-00024.safetensors

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

model-00013-of-00024.safetensors

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

model-00014-of-00024.safetensors

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

model-00015-of-00024.safetensors

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

model-00016-of-00024.safetensors

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

model-00017-of-00024.safetensors

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

model-00018-of-00024.safetensors

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

model-00019-of-00024.safetensors

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

model-00020-of-00024.safetensors

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

model-00021-of-00024.safetensors

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

model-00022-of-00024.safetensors

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

model-00023-of-00024.safetensors

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

model-00024-of-00024.safetensors

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

model.safetensors.index.json

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

modeling_afmoe.py

@@ -0,0 +1,680 @@
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    MoeCausalLMOutputWithPast,
+    MoeModelOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.integrations import use_kernel_forward_from_hub
+
+
+try:
+    from .configuration_afmoe import AfmoeConfig
+except:
+    from configuration_afmoe import AfmoeConfig
+
+class AfmoeRotaryEmbedding(nn.Module):
+
+    def __init__(self, config: AfmoeConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    def _dynamic_frequency_update(self, position_ids, device):
+        """
+        dynamic RoPE layers should recompute `inv_freq` in the following situations:
+        1 - growing beyond the cached sequence length (allow scaling)
+        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
+        """
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:  # growth
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
+            self.max_seq_len_cached = seq_len
+
+        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
+            # This .to() is needed if the model has been moved to a device after being initialized (because
+            # the buffer is automatically moved, but not the original copy)
+            self.original_inv_freq = self.original_inv_freq.to(device)
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            self._dynamic_frequency_update(position_ids, device=x.device)
+
+        # Core RoPE block
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
+        device_type = x.device.type
+        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
+        cos = cos * self.attention_scaling
+        sin = sin * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+@use_kernel_forward_from_hub("RMSNorm")
+class AfmoeRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float):
+        """
+        AfmoeRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    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 = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.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
+
+
+class AfmoeMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class AfmoeTokenChoiceRouter(nn.Module):
+    """Token-choice top-K router for MoE routing."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.score_func = config.score_func
+        self.route_norm = config.route_norm
+        self.route_scale = config.route_scale
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)     
+
+    def forward(self, hidden_states, expert_bias: torch.Tensor | None):
+        _, _, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        scores = self.gate(hidden_states)
+
+        # Apply scoring function in float32 for stability
+        if self.score_func == "sigmoid":
+            scores = torch.sigmoid(scores.to(torch.float32))
+        else:
+            scores = F.softmax(scores.to(torch.float32), dim=-1)
+
+        if expert_bias is not None:
+            _, selected_experts = torch.topk(scores + expert_bias, k=self.top_k, dim=1)
+            top_scores = scores.gather(dim=1, index=selected_experts)
+        else:
+            top_scores, selected_experts = torch.topk(scores, k=self.top_k, dim=1)
+
+        # Normalize weights if using sigmoid
+        if self.score_func == "sigmoid" and self.route_norm:
+            denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20
+            top_scores = top_scores / denominator
+
+        top_scores = top_scores * self.route_scale
+        return top_scores, selected_experts
+
+class AfmoeMoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.router = AfmoeTokenChoiceRouter(config)
+
+        self.shared_experts = None
+        if config.num_shared_experts > 0:
+            self.shared_experts = AfmoeMLP(
+                config, config.moe_intermediate_size * config.num_shared_experts
+            )
+        self.experts = nn.ModuleList(
+            [AfmoeMLP(
+                config, intermediate_size=config.moe_intermediate_size
+            ) for _ in range(config.num_experts)]
+        )
+        self.expert_bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False)
+        
+
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states_flat = hidden_states.view(-1, hidden_dim)
+
+        # Get routing decisions
+        top_scores, selected_experts = self.router(hidden_states, self.expert_bias)
+
+        # Process through shared experts
+        if self.shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states_flat)
+        else:
+            shared_output = torch.zeros_like(hidden_states_flat)
+
+        # Reorder tokens by expert for efficient processing
+        token_indices_sorted = torch.argsort(selected_experts.view(-1), stable=True)
+        top_scores_sorted = top_scores.view(-1)[token_indices_sorted]
+        token_to_expert = selected_experts.view(-1)[token_indices_sorted]
+        token_indices_sorted = token_indices_sorted // self.config.num_experts_per_tok
+
+        # Gather input tokens
+        token_indices_expanded = token_indices_sorted.unsqueeze(-1).expand(
+            -1, hidden_dim
+        )
+        routed_input = torch.gather(
+            hidden_states_flat, dim=0, index=token_indices_expanded
+        )
+
+        routed_output = torch.zeros_like(routed_input)
+        for expert_id in range(self.config.num_experts):
+            mask = token_to_expert == expert_id
+            if mask.any():
+                expert_input = routed_input[mask]
+                expert_out = self.experts[expert_id](expert_input)
+                routed_output[mask] = expert_out
+          
+        routed_output = (
+            routed_output.to(torch.float32) * top_scores_sorted.unsqueeze(-1)
+        ).to(hidden_states.dtype)
+
+        # Scatter back to original positions
+        output = shared_output.scatter_add(
+            dim=0, index=token_indices_expanded, src=routed_output
+        )
+
+        return output.view(batch_size, seq_len, hidden_dim)
+
+
+class AfmoeAttention(nn.Module):
+    """Multi-headed attention with local/global pattern and gating."""
+
+    def __init__(self, config: AfmoeConfig, layer_idx: int):
+        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_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention"
+        self.sliding_window = config.sliding_window if self.is_local_attention else None
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, config.hidden_size, bias=False
+        )
+
+        self.q_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+        self.gate_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.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_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:   
+
+        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(hidden_shape)
+        value_states = self.v_proj(hidden_states).view(hidden_shape)
+        gate_states = self.gate_proj(hidden_states)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+        
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if self.is_local_attention:
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        output, _ = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask=attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        output = output.view(*input_shape, -1).contiguous()
+        output = output * F.sigmoid(gate_states)
+        return self.o_proj(output)
+
+
+class AfmoeDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: AfmoeConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+
+        self.self_attn = AfmoeAttention(config=config, layer_idx=layer_idx)
+        self.attention_type = config.layer_types[layer_idx]
+
+        # Dual normalization for attention
+        self.input_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Dual normalization for FFN
+        self.pre_mlp_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_mlp_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # MoE or dense FFN
+        self.moe_enabled = layer_idx >= config.num_dense_layers
+        if self.moe_enabled:
+            self.mlp = AfmoeMoE(config)
+        else:
+            self.mlp = AfmoeMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+
+        # Self Attention with dual normalization
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        # FFN with dual normalization
+        residual = hidden_states
+        hidden_states = self.pre_mlp_layernorm(hidden_states)
+
+        if self.moe_enabled:
+            hidden_states = self.mlp(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
+
+        hidden_states = self.post_mlp_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class AfmoePreTrainedModel(PreTrainedModel):
+    config_class = AfmoeConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["AfmoeDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _keep_in_fp32_modules = [
+        "input_layernorm",
+        "post_attention_layernorm",
+        "pre_mlp_layernorm",
+        "post_mlp_layernorm",
+        "q_norm",
+        "k_norm",
+        "norm",
+    ]
+    _supports_sdpa = True
+    _supports_attention_backend = True
+    supports_gradient_checkpointing = True
+
+
+class AfmoeModel(AfmoePreTrainedModel):
+    _no_split_modules = ["AfmoeDecoderLayer"]
+
+    def __init__(self, config: AfmoeConfig):
+        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(
+            [
+                AfmoeDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = AfmoeRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        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)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+            }
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
+            }
+
+        hidden_states = inputs_embeds
+
+        # Apply muP input scaling if enabled
+        if self.config.mup_enabled:
+            hidden_states = hidden_states * (self.config.hidden_size**0.5)
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class AfmoeForCausalLM(AfmoePreTrainedModel, 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 = AfmoeModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        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 set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        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,
+        token_type_ids: Optional[torch.Tensor] = None,  # will be ignored
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+        outputs: MoeModelOutputWithPast = 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
+        # Only compute necessary logits
+        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, labels, self.vocab_size, **kwargs)
+
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = [
+    "AfmoeForCausalLM",
+    "AfmoeModel",
+    "AfmoePreTrainedModel",
+]

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|begin_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,272 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|begin_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<|eot_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<|start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<|channel|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<|message|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<|end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<|fitm_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<|fitm_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<|fitm_hole|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "14": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "15": {
+      "content": "<tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "16": {
+      "content": "</tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "17": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "18": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "19": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "20": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "21": {
+      "content": "<|reserved_special_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "22": {
+      "content": "<|reserved_special_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "23": {
+      "content": "<|reserved_special_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "24": {
+      "content": "<|reserved_special_11|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "25": {
+      "content": "<|reserved_special_12|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "26": {
+      "content": "<|reserved_special_13|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "27": {
+      "content": "<|reserved_special_14|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "28": {
+      "content": "<|reserved_special_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "29": {
+      "content": "<|reserved_special_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "30": {
+      "content": "<|reserved_special_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "31": {
+      "content": "<|reserved_special_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "model_max_length": 65536,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "use_default_system_prompt": false,
+  "chat_template": "{{ bos_token }}{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {{- content}}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>' }}\n        {{- '\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}"
+}