Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,195 @@
+Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf
+      of the copyright 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 and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      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 Work 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 Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing
+      the origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (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 or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. 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
+      Work (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.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

README.md

@@ -0,0 +1,233 @@
+---
+license: apache-2.0
+pipeline_tag: text-generation
+library_name: transformers
+---
+
+# Agnes 2.5 flash base: An Efficient Sparse Mixture-of-Experts Foundation Model
+
+**Agnes 2.5 flash base** is a **202B-parameter sparse Mixture-of-Experts (MoE) base model** with roughly **16B active parameters per token**. It is designed for long-context, high-throughput inference and is released here as an **FP8 checkpoint** that can be served out-of-the-box with `sglang`.
+
+## Introduction
+
+Agnes 2.5 flash base is a decoder-only Transformer that combines several efficiency-oriented components:
+
+1. **Sparse MoE feed-forward layers.** Each of the 48 layers routes every token to **6 of 160 experts** (plus one always-on shared expert). The first 3 layers use deterministic hash routing; the remaining 45 layers use a learned top-k router with auxiliary-loss-free load balancing (`noaux_tc`).
+2. **Parallel dense FFN branch.** Layers 3–47 additionally carry a lightweight dense FFN branch (intermediate size 2048) in parallel with the MoE block, increasing per-token capacity at negligible latency cost.
+3. **Multi-head Latent Attention (MLA) with KV compression.** Attention uses low-rank query/output projections and a per-layer compressor (compression ratios alternate between 4 and 128 across layers), together with a sparse top-512 token indexer, keeping the KV cache small at very long context.
+4. **Hyper-connections.** Residual streams use multi-stream hyper-connections (`hc_mult = 4`) with Sinkhorn-normalized mixing in place of a single residual path.
+5. **1M-token context.** YaRN rotary scaling (factor 16 over a 64K base window) extends the usable context to **1,048,576 tokens**.
+
+This repository contains the **base (pre-trained, non-instruction-tuned)** model. It is intended for continued pre-training, fine-tuning, and research; it has not undergone SFT or RLHF, so it should not be expected to follow chat-style instructions reliably.
+
+## Model Zoo
+
+
+| Model            | Precision | Layers | Experts (active / total) | Params (active / total) | Context   | Hugging Face Model Card |
+| ---------------- | --------- | ------ | ------------------------ | ----------------------- | --------- | ----------------------- |
+| Agnes 2.5 flash base | FP8       | 48     | 6 + 1 shared / 160       | ~16B / 202B             | 1,048,576 | ✅ this repository       |
+
+
+
+
+### Architecture at a glance
+
+
+| Hyper-parameter                  | Value                            |
+| -------------------------------- | -------------------------------- |
+| `hidden_size`                    | 4096                             |
+| `num_hidden_layers`              | 48                               |
+| `num_attention_heads`            | 64 (`head_dim` 512, RoPE dim 64) |
+| `q_lora_rank` / `o_lora_rank`    | 1024 / 1024                      |
+| `n_routed_experts`               | 160                              |
+| `num_experts_per_tok`            | 6                                |
+| `n_shared_experts`               | 1                                |
+| `moe_intermediate_size`          | 2048                             |
+| `parallel_ffn_intermediate_size` | 2048 (layers 3–47)               |
+| `num_hash_layers`                | 3                                |
+| `index_topk`                     | 512                              |
+| `hc_mult`                        | 4                                |
+| `vocab_size`                     | 129,292                          |
+| `max_position_embeddings`        | 1,048,576                        |
+
+
+
+
+## Quantization
+
+Weights are stored in **FP8 (e4m3)** with **128×128 block-wise UE8M0 scales** and **dynamic activation quantization**:
+
+```json
+"quantization_config": {
+  "quant_method": "fp8",
+  "fmt": "e4m3",
+  "scale_fmt": "ue8m0",
+  "weight_block_size": [128, 128],
+  "activation_scheme": "dynamic"
+}
+```
+
+Embeddings, the LM head, normalization layers, router weights and hyper-connection parameters are kept in BF16. Every FP8 linear weight `<name>.weight` is accompanied by a sibling `<name>.scale` tensor (fp32). The checkpoint is ~190 GB across 37 `safetensors` shards.
+
+## Getting Started: Serving with sglang
+
+The recommended way to run Agnes 2.5 flash base is with the **stock** `lmsysorg/sglang:v0.5.16` **Docker image**. Because Agnes support is not yet upstream in sglang, this repository ships the required support files under `[sglang_patch/](./sglang_patch)` together with a launcher script `[serve.sh](./serve.sh)` that overlays them onto the container's sglang package at start-up. **No custom image is needed, and the model directory itself is never modified.**
+
+**Hardware note:** the FP8 checkpoint needs ~190 GB of GPU memory for weights alone. The default configuration uses tensor parallelism over 8 GPUs (e.g. 8× H100/H200 80 GB+).
+
+### 1. Download the model
+
+```shell
+pip install -U "huggingface_hub[cli]"
+huggingface-cli download <org>/Agnes 2.5 flash base --local-dir ./Agnes 2.5 flash base
+```
+
+
+
+### 2. Launch the server
+
+```shell
+docker run --gpus all --shm-size 64g -p 30001:30002 \
+  -v $(pwd)/Agnes 2.5 flash base:/model \
+  lmsysorg/sglang:v0.5.16 bash /model/serve.sh
+```
+
+`serve.sh` copies `sglang_patch/srt` and `sglang_patch/kernels` into the container's `sglang` package and then execs:
+
+```shell
+sglang serve --model-path /model --trust-remote-code --tp 8 \
+  --context-length 1048576 --mem-fraction-static 0.90 \
+  --host 0.0.0.0 --port 30002
+```
+
+Any extra sglang flags can be appended after `serve.sh` and are passed straight through, e.g. a shorter context window to leave more room for the KV cache:
+
+```shell
+... lmsysorg/sglang:v0.5.16 bash /model/serve.sh --context-length 262144
+```
+
+Model loading takes roughly 10–15 minutes on 8 GPUs. The server is ready once `/health` returns `200`:
+
+```shell
+curl http://localhost:30001/health
+curl http://localhost:30001/get_model_info
+```
+
+
+
+### 3. Query the model
+
+Native `/generate` endpoint:
+
+```shell
+curl http://localhost:30001/generate \
+  -H "Content-Type: application/json" \
+  -d '{
+    "text": "The three laws of thermodynamics are",
+    "sampling_params": {"max_new_tokens": 128, "temperature": 0.7, "top_p": 0.95}
+  }'
+```
+
+OpenAI-compatible completions endpoint (this is a base model, so prefer `/v1/completions` over `/v1/chat/completions`):
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:30001/v1", api_key="EMPTY")
+resp = client.completions.create(
+    model="default",
+    prompt="The three laws of thermodynamics are",
+    max_tokens=128,
+    temperature=0.7,
+    top_p=0.95,
+)
+print(resp.choices[0].text)
+```
+
+
+
+### Manual variant (what `serve.sh` does)
+
+If you prefer not to use the launcher script:
+
+```shell
+docker run --gpus all --shm-size 64g -p 30001:30002 \
+  -v $(pwd)/Agnes 2.5 flash base:/model \
+  lmsysorg/sglang:v0.5.16 \
+  sh -c "cp -r /model/sglang_patch/srt /model/sglang_patch/kernels \
+           /sgl-workspace/sglang/python/sglang/ && \
+         exec sglang serve --model-path /model --trust-remote-code --tp 8 \
+           --context-length 1048576 --mem-fraction-static 0.90 \
+           --host 0.0.0.0 --port 30002"
+```
+
+**Important:** the image version must be **exactly** `lmsysorg/sglang:v0.5.16`. The overlay replaces a small set of version-specific files inside sglang; applying it to a different release is not supported.
+
+## Loading with transformers
+
+The repository ships `configuration_agnes.py` and `modeling_agnes.py`, so the model can also be loaded directly with 🤗 transformers using `trust_remote_code=True` (no sglang patch required). Note that the reference PyTorch implementation is intended for inspection, fine-tuning and research rather than high-throughput serving.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "<org>/Agnes 2.5 flash base"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+inputs = tokenizer("The three laws of thermodynamics are", return_tensors="pt").to(model.device)
+out = model.generate(**inputs, max_new_tokens=64, do_sample=True, temperature=0.7, top_p=0.95)
+print(tokenizer.decode(out[0], skip_special_tokens=True))
+```
+
+
+
+## Repository layout
+
+```
+Agnes 2.5 flash base/
+├── config.json                     # architecture + FP8 quantization_config
+├── generation_config.json
+├── configuration_agnes.py          # transformers remote code
+├── modeling_agnes.py
+├── tokenizer.json / tokenizer_config.json
+├── model-000xx-of-00037.safetensors
+├── model.safetensors.index.json
+├── serve.sh                        # one-command sglang launcher
+└── sglang_patch/                   # Agnes support overlay for sglang v0.5.16
+    ├── srt/...
+    └── kernels/...
+```
+
+
+
+## Limitations
+
+- **Base model.** No instruction tuning or safety alignment has been applied. Outputs may be incoherent, biased or unsafe; apply your own alignment and filtering before deployment.
+- **Memory.** The full FP8 checkpoint requires multi-GPU tensor parallelism; single-GPU inference is not supported.
+- **sglang version pin.** The bundled overlay targets sglang `v0.5.16` only.
+
+## License
+Both the code repository and the model weights are released under the [Apache License 2.0](LICENSE).
+
+
+## Citation
+
+If you use Agnes 2.5 flash base in your research, please cite:
+
+```bibtex
+@misc{agnes2026flash,
+      title={Agnes 2.5 flash base: An Efficient Sparse Mixture-of-Experts Foundation Model},
+      author={Agnes AI Team},
+      year={2026},
+      url={https://huggingface.co/<org>/Agnes 2.5 flash base},
+}
+```
+
+
+

config.json

@@ -0,0 +1,122 @@
+{
+  "model_type": "agnes",
+  "architectures": [
+    "AgnesForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_agnes.AgnesConfig",
+    "AutoModel": "modeling_agnes.AgnesModel",
+    "AutoModelForCausalLM": "modeling_agnes.AgnesForCausalLM"
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.1",
+  "vocab_size": 129292,
+  "hidden_size": 4096,
+  "num_hidden_layers": 48,
+  "num_attention_heads": 64,
+  "num_key_value_heads": 1,
+  "head_dim": 512,
+  "qk_rope_head_dim": 64,
+  "q_lora_rank": 1024,
+  "o_groups": 8,
+  "o_lora_rank": 1024,
+  "sliding_window": 128,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "index_n_heads": 64,
+  "index_head_dim": 128,
+  "index_topk": 512,
+  "compress_rope_theta": 160000,
+  "compress_ratios": [
+    0,
+    0,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    0
+  ],
+  "moe_intermediate_size": 2048,
+  "n_routed_experts": 160,
+  "n_shared_experts": 1,
+  "num_experts_per_tok": 6,
+  "num_hash_layers": 3,
+  "scoring_func": "sqrtsoftplus",
+  "topk_method": "noaux_tc",
+  "norm_topk_prob": true,
+  "routed_scaling_factor": 1.5,
+  "num_nextn_predict_layers": 0,
+  "parallel_ffn_intermediate_size": 2048,
+  "hidden_act": "silu",
+  "swiglu_limit": 10.0,
+  "hc_mult": 4,
+  "hc_sinkhorn_iters": 20,
+  "hc_eps": 1e-06,
+  "rms_norm_eps": 1e-06,
+  "initializer_range": 0.02,
+  "rope_theta": 10000,
+  "rope_scaling": {
+    "type": "yarn",
+    "factor": 16,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "original_max_position_embeddings": 65536
+  },
+  "max_position_embeddings": 1048576,
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "quantization_config": {
+    "activation_scheme": "dynamic",
+    "fmt": "e4m3",
+    "quant_method": "fp8",
+    "scale_fmt": "ue8m0",
+    "weight_block_size": [
+      128,
+      128
+    ]
+  }
+}

configuration_agnes.py

@@ -0,0 +1,330 @@
+# Copyright 2026 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 huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.modeling_rope_utils import RopeParameters
+from transformers.utils import auto_docstring
+
+
+# The two MoE block kinds and the three attention block kinds Agnes ships, plus
+# the legacy per-layer integer -> attention-kind map used when a checkpoint still
+# carries the old `compress_ratios` list.
+AGNES_MLP_LAYER_TYPES = ("agnes_hash_moe", "agnes_moe")
+
+AGNES_LAYER_TYPES = (
+    "agnes_local_attention",
+    "agnes_sparse_attention",
+    "agnes_pooled_attention",
+)
+
+_COMPRESS_RATIO_TO_LAYER_TYPE = {
+    0: "agnes_local_attention",
+    4: "agnes_sparse_attention",
+    128: "agnes_pooled_attention",
+}
+
+
+@auto_docstring
+@strict
+class AgnesConfig(PreTrainedConfig):
+    r"""
+    scoring_func (`str`):
+        Activation applied to the router logits — one of `sqrtsoftplus`,
+        `softmax`, `sigmoid`.
+    rope_theta (`float`):
+        Rotary base for the main (local) attention path.
+    layer_types (`list[str]`):
+        The per-layer attention schedule, drawn from `agnes_local_attention`,
+        `agnes_sparse_attention`, `agnes_pooled_attention`. Default schedule:
+        two pooled layers to bootstrap, then an alternating sparse/pooled tail.
+    compress_rates (`dict[str, int]`):
+        Compression stride per attention type; default
+        `{"agnes_sparse_attention": 4, "agnes_pooled_attention": 128}`. For
+        backward compatibility, top-level `compress_rate_csa` /
+        `compress_rate_hca` kwargs are absorbed in `__post_init__`.
+    compress_rope_theta (`float`):
+        Rotary base shared by the compressed branches (used together with
+        `rope_scaling` for YaRN).
+    hc_mult (`int`):
+        Number of parallel residual streams carried by the manifold-constrained
+        hyper-connections (always on).
+    hc_sinkhorn_iters (`int`):
+        Iteration budget for the Sinkhorn-Knopp projection of the stream-mixing
+        matrix onto the doubly-stochastic manifold.
+    hc_eps (`float`):
+        Small constant guarding the Sinkhorn normalisation against divide-by-zero.
+    mlp_layer_types (`list[str]`):
+        The per-layer MoE schedule over `agnes_hash_moe` / `agnes_moe`. Hash
+        layers pick experts from the frozen `tid2eid[input_ids]` table; plain
+        layers use learned top-k routing. Default: the first three layers hash,
+        the rest learned. Legacy `num_hash_layers` is folded in at
+        `__post_init__`.
+    swiglu_limit (`float`):
+        Clamp bound on the routed experts' gate/up pre-activations.
+    parallel_ffn_intermediate_size (`int`):
+        Width of the optional parallel dense-FFN branch on every non-hash MoE
+        layer (`0` turns it off). Its `down_proj` is exported as zeros, so the
+        branch is a no-op until trained.
+    sliding_window (`int`):
+        Local attention window width used by every block.
+    o_groups (`int`):
+        Head-group count for the grouped output projection.
+    o_lora_rank (`int`):
+        Per-group bottleneck width of the grouped output projection.
+    index_n_heads (`int`):
+        Query-head count of the Lightning Indexer.
+    index_head_dim (`int`):
+        Per-head width inside the Lightning Indexer.
+    index_topk (`int`):
+        How many compressed entries the indexer keeps per query.
+    num_nextn_predict_layers (`int`):
+        Number of MTP layers present in the source checkpoint (not built here).
+    partial_rotary_factor (`float`, *optional*):
+        Fraction of each head that is rotated. Defaults to
+        `qk_rope_head_dim / head_dim`, sizing cos/sin to `qk_rope_head_dim`.
+    """
+
+    model_type = "agnes"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # --- core dimensions ---
+    vocab_size: int = 129280
+    hidden_size: int = 4096
+    num_hidden_layers: int = 43
+
+    # --- self-attention (shared-KV MQA + grouped output projection) ---
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 1
+    head_dim: int = 512
+    default_partial_rotary_factor = 64 / 512  # `qk_rope_head_dim` (64) / `head_dim` (512)
+    partial_rotary_factor: float | None = None
+    q_lora_rank: int = 1024
+    o_groups: int = 8
+    o_lora_rank: int = 1024
+    sliding_window: int = 128
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+
+    # --- long-range compressor + lightning indexer ---
+    index_n_heads: int = 64
+    index_head_dim: int = 128
+    index_topk: int = 512
+    layer_types: list[str] | None = None
+    compress_rates: dict | None = None
+    default_compress_rates = {"agnes_sparse_attention": 4, "agnes_pooled_attention": 128}
+    compress_rope_theta: float | int = 160000.0
+
+    # --- mixture-of-experts routing ---
+    moe_intermediate_size: int = 2048
+    parallel_ffn_intermediate_size: int = 0
+    n_routed_experts: int = 256
+    n_shared_experts: int = 1
+    num_experts_per_tok: int = 6
+    mlp_layer_types: list[str] | None = None
+    default_num_hash_layers = 3
+    scoring_func: str = "sqrtsoftplus"
+    norm_topk_prob: bool = True
+    routed_scaling_factor: float = 1.5
+    num_nextn_predict_layers: int = 1
+    output_router_logits: bool = False
+    router_aux_loss_coef: float = 0.001
+    router_jitter_noise: float = 0.0
+
+    # --- feed-forward activation ---
+    hidden_act: str = "silu"
+    swiglu_limit: float = 10.0
+    mlp_bias: bool = False
+
+    # --- manifold-constrained hyper-connections ---
+    hc_mult: int = 4
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1.0e-6
+
+    # --- normalisation + initialisation ---
+    rms_norm_eps: float = 1.0e-6
+    initializer_range: float = 0.02
+
+    # --- rotary position embedding ---
+    rope_theta: float | int = 10000.0
+    rope_parameters: RopeParameters | dict | None = None
+    max_position_embeddings: int = 1048576
+
+    # --- special tokens + weight tying ---
+    pad_token_id: int | None = None
+    bos_token_id: int | None = 0
+    eos_token_id: int | list[int] | None = 1
+    tie_word_embeddings: bool = False
+
+    # --- runtime ---
+    use_cache: bool = True
+
+    # ------------------------------------------------------------------ #
+    # Non-field class attributes: alias map, parallelism plans, rope labels
+    # ------------------------------------------------------------------ #
+
+    # Expert-parallel plan. Agnes ships EP only — it is MoE, so there is no
+    # `base_model_tp_plan`. The gate routes, the routed experts run as a grouped
+    # GEMM sharded on the expert axis, and the experts module is wrapped so its
+    # output is all-reduced. Core attention stays replicated: it is shared-KV MQA
+    # broadcasting a single KV head to every query head, so colwise-sharding
+    # `q_b_proj` would desync the KV broadcast from the rank-local head count; the
+    # small shared MLP is not worth sharding either. The one exception is the
+    # Lightning Indexer, whose keys are replicated (its own compressor runs at
+    # index_head_dim on replicated hidden states): there `q_b_proj` and
+    # `scorer.weights_proj` go colwise and the `scorer` output is all-reduced so
+    # every rank picks the same top-k.
+    base_model_ep_plan = {
+        "layers.*.mlp.gate": "ep_router",
+        "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+        "layers.*.mlp.experts.down_proj": "grouped_gemm",
+        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.self_attn.compressor.indexer.q_b_proj": "colwise",
+        "layers.*.self_attn.compressor.indexer.scorer.weights_proj": "colwise",
+        "layers.*.self_attn.compressor.indexer.scorer": "all_reduce",
+    }
+
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    # `num_local_experts` / `intermediate_size` are the names the MoE FP8 / TP
+    # integrations and the shared-expert MLP read; Agnes only stores
+    # `n_routed_experts` / `moe_intermediate_size`, so alias them through.
+    attribute_map = {
+        "num_local_experts": "n_routed_experts",
+        "intermediate_size": "moe_intermediate_size",
+    }
+
+    # `rope_parameters` is keyed by rope label (`main` / `compress`), not by
+    # `layer_types`. The base `validate_rope` assumes the latter, so it is
+    # overridden below to walk the rope-label sub-dicts directly.
+    _rope_type_labels = ("main", "compress")
+
+    def __post_init__(self, **kwargs):
+        # Older checkpoints still ship a handful of legacy kwargs; pop them out of
+        # the dict before the strict parent init runs, then reconcile each into its
+        # current field afterwards.
+        leg_ratios = kwargs.pop("compress_ratios", None)
+        leg_csa = kwargs.pop("compress_rate_csa", None)
+        leg_hca = kwargs.pop("compress_rate_hca", None)
+        leg_hash = kwargs.pop("num_hash_layers", None)
+        leg_qk_rope = kwargs.pop("qk_rope_head_dim", None)
+        PreTrainedConfig.__post_init__(self, **kwargs)
+        n = self.num_hidden_layers
+
+        # compress_rates: seed the per-type defaults, then let a legacy scalar win.
+        if self.compress_rates is None:
+            self.compress_rates = dict(self.default_compress_rates)
+        if leg_csa is not None:
+            self.compress_rates["agnes_sparse_attention"] = leg_csa
+        if leg_hca is not None:
+            self.compress_rates["agnes_pooled_attention"] = leg_hca
+
+        # layer_types precedence: explicit field, then legacy 0/4/128 ratios, then
+        # the default (two pooled bootstrap layers + sparse/pooled interleave).
+        if self.layer_types is None and leg_ratios is not None:
+            self.layer_types = [_COMPRESS_RATIO_TO_LAYER_TYPE[r] for r in leg_ratios]
+        if self.layer_types is None:
+            tail = [
+                "agnes_sparse_attention" if i % 2 else "agnes_pooled_attention"
+                for i in range(max(n - 2, 0))
+            ]
+            self.layer_types = ["agnes_pooled_attention"] * min(n, 2) + tail
+        self.layer_types = list(self.layer_types[:n])
+
+        # mlp_layer_types: leading hash-routed layers, learned routing for the rest.
+        if self.mlp_layer_types is None:
+            n_hash = leg_hash if leg_hash is not None else self.default_num_hash_layers
+            self.mlp_layer_types = ["agnes_hash_moe"] * min(n, n_hash) + ["agnes_moe"] * max(0, n - n_hash)
+        self.mlp_layer_types = list(self.mlp_layer_types[:n])
+
+        # partial_rotary_factor from a legacy qk_rope_head_dim if provided, else the
+        # default; qk_rope_head_dim itself is only ever a runtime attr, not a field.
+        if self.partial_rotary_factor is None:
+            self.partial_rotary_factor = (
+                leg_qk_rope / self.head_dim if leg_qk_rope is not None else self.default_partial_rotary_factor
+            )
+        self.qk_rope_head_dim = int(self.head_dim * self.partial_rotary_factor)
+
+        # Normalise rope_parameters into the {main, compress} nesting. Local layers
+        # use plain rope at rope_theta; only the compress branch may be YaRN, and if
+        # it is we pin attention_factor=1.0 (Agnes never applies YaRN's mscale).
+        rp = self.rope_parameters or {}
+        if isinstance(rp.get("main"), dict) and isinstance(rp.get("compress"), dict):
+            self.rope_parameters = {"main": rp["main"], "compress": rp["compress"]}
+        else:
+            extra = {k: v for k, v in rp.items() if k not in ("main", "compress")}
+            main = {
+                "rope_type": "default",
+                "rope_theta": self.rope_theta,
+                "partial_rotary_factor": self.partial_rotary_factor,
+            }
+            compress = {
+                **extra,
+                "rope_theta": self.compress_rope_theta,
+                "partial_rotary_factor": self.partial_rotary_factor,
+            }
+            compress.setdefault("rope_type", "default")
+            if compress["rope_type"] == "yarn":
+                compress.setdefault("attention_factor", 1.0)
+            self.rope_parameters = {"main": main, "compress": compress}
+
+    def validate_layer_type(self):
+        """Keep `layer_types` / `mlp_layer_types` within the block kinds Agnes
+        actually builds, alongside the usual length check against
+        `num_hidden_layers`."""
+        if self.num_hidden_layers is None:
+            return
+        schedules = (
+            ("layer_types", self.layer_types, AGNES_LAYER_TYPES),
+            ("mlp_layer_types", self.mlp_layer_types, AGNES_MLP_LAYER_TYPES),
+        )
+        for name, schedule, allowed in schedules:
+            if schedule is None:
+                continue
+            if len(schedule) != self.num_hidden_layers:
+                raise ValueError(
+                    f"`num_hidden_layers` ({self.num_hidden_layers}) must equal `len({name})` ({len(schedule)})."
+                )
+            unknown = [t for t in schedule if t not in allowed]
+            if unknown:
+                raise ValueError(f"`{name}` entries must be one of {allowed} for Agnes; got {unknown}.")
+
+    def validate_rope(self):
+        # The stock validators index self.rope_parameters[<key>] directly, which
+        # breaks against the {main, compress} nesting, so run each rope-label
+        # sub-dict through its own validator with self.rope_parameters temporarily
+        # pointed at it, restoring the nested dict afterwards.
+        nested = getattr(self, "rope_parameters", None) or {}
+        ignore_keys = self.ignore_keys_at_rope_validation
+        for label in self._rope_type_labels:
+            sub = nested.get(label)
+            if not isinstance(sub, dict):
+                continue
+            kind = sub.get("rope_type", sub.get("type", "default"))
+            sub["rope_type"] = kind
+            validator = getattr(self, f"_validate_{kind}_rope_parameters", None)
+            if validator is None:
+                continue
+            self.rope_parameters = sub
+            try:
+                validator(sub, ignore_keys=ignore_keys)
+            finally:
+                self.rope_parameters = nested
+
+
+__all__ = ["AgnesConfig"]

generation_config.json

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

model-00001-of-00035.safetensors

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

model-00002-of-00035.safetensors

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

model-00003-of-00035.safetensors

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

model-00004-of-00035.safetensors

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

model-00005-of-00035.safetensors

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

model-00006-of-00035.safetensors

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

model-00007-of-00035.safetensors

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

model-00008-of-00035.safetensors

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

model-00009-of-00035.safetensors

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

model-00010-of-00035.safetensors

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

model-00011-of-00035.safetensors

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

model-00012-of-00035.safetensors

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

model-00013-of-00035.safetensors

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

model-00014-of-00035.safetensors

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

model-00015-of-00035.safetensors

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