diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..a378a680a62eb0043fa519134a3613aeeb10fcbf 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +tokenizer.json filter=lfs diff=lfs merge=lfs -text +assets/sarvam-prism.png filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + 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. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..27616f7c8d33b64f47cf9ea306cdd3c9ceab8c78 --- /dev/null +++ b/README.md @@ -0,0 +1,145 @@ +--- +license: apache-2.0 +base_model: sarvamai/prism +pipeline_tag: text-classification +language: [en, hi, bn, ta, te, mr, gu, kn, ml, or, pa] +tags: [decision-model, classification, calibration, int4, gptq, quantized, moe] +--- + +# Sarvam-Prism INT4 (≤16 GB GPUs) + +An INT4 quantization of [sarvamai/prism](https://huggingface.co/sarvamai/prism), a decision model built on Sarvam-105B. +The original needs one FP8-capable GPU with ~120 GB of memory. This version runs on **any CUDA GPU with about 10 GB of +memory or more**, for example 16, 12 or 24 GB cards. It does not need FP8 tensor cores. + +You give it a context and one or more questions, each with a fixed set of answers. It returns a probability for every +answer in one forward pass. The API is the same as the original's: `decide(state, questions)`. + +## Requirements + +- A CUDA GPU with **≥ 10 GB** of memory. sm_80+ (Ampere or newer) is recommended. sm_70/75 also work, on a slower path. + It was tested with GPU memory hard-capped at 15 GiB and at 11 GiB. +- **≥ 64 GB of host RAM.** About 54 GB of expert weights live in pinned host memory. +- Python 3.10+, PyTorch ≥ 2.4, Triton ≥ 3.0, `transformers==4.52.1` (see `requirements.txt`), and a C compiler + (Triton builds its kernel on first use). +- About 56 GB of disk. + +## Quick start + +```bash +hf download jc-914/prism-quantized --local-dir prism-int4 +cd prism-int4 +pip install -r requirements.txt +python example.py . # optional 2nd arg: GPU memory cap in GiB, e.g. `python example.py . 12` +``` + +```python +from prism_int4 import Prism + +model = Prism.from_pretrained(".") # uses the free memory on cuda:0; or gpu_memory_gb=12 +answers = model.decide( + state={"ticket": "My card was charged twice for the same order and I still haven't got a refund."}, + questions={ + "team": {"type": "choice", "instructions": "Which team should handle this ticket?", + "criteria": {"billing": "payments, charges and refunds", "tech": "app or website problems", + "delivery": "shipping and order status"}}, + "urgent": {"type": "yes_no", "instructions": "Does this need a reply within the hour?"}, + "anger": {"type": "score", "instructions": "How upset is the customer?", + "criteria": ["calm", "annoyed", "angry", "furious"]}, + }, +) +# Many requests at once (amortises the weight streaming, see Speed): +results = model.decide_batch([(state1, questions1), (state2, questions2), ...]) +``` + +The request format (`choice` / `yes_no` / `score`) and the output format are the same as the +[original model card](https://huggingface.co/sarvamai/prism). + +## What was quantized + +| part | parameters | format here | vs. original | +|---|---|---|---| +| 128 routed experts × 31 MoE layers | ~100B | **INT4**, GPTQ, symmetric, group size 128, fp32 scales | was FP8 | +| attention, dense layer 0, shared experts | ~4B | FP8 (e4m3, one scale per output row), widened to bf16 at use | **unchanged** | +| embeddings, norms, MoE router, decision head | – | bf16 / fp32 | unchanged | + +- **GPTQ calibration.** 262,533 tokens of requests in Prism's own input format (state, question, options and + answer, with the decision attention mask). The text is multilingual Wikipedia in all 11 supported languages, and + the requests cover all three question types. Each layer's MoE input was recorded from the original FP8 model. + Every expert's down projection is calibrated on the output of its *quantized* gate/up projection. +- **Rare experts.** Experts that were routed fewer than 512 calibration tokens (180 of 3,968) use round-to-nearest + with an MSE-optimal clip search instead of GPTQ, so their error feedback is not fitted to a handful of samples. +- **Layer error.** The Hessian-weighted relative output error, averaged over layers, is 0.0014 for gate/up and + 0.0036 for down, against 0.0075 and 0.0109 for plain round-to-nearest. Per-layer numbers are in + `quantization/quant_stats.json`. + +## Accuracy against the original FP8 model + +The suite is 26 requests / 43 questions: English and all 10 Indic languages, all question types, and contexts from a +few tokens up to ~4k tokens. It is disjoint from the calibration data. The reference outputs come from the original +`prism.py` on an H200. + +| metric | INT4 (15 GiB cap) | INT4 (11 GiB cap) | +|---|---|---| +| top-1 answer agreement | **43 / 43** | 43 / 43 | +| mean of the largest per-option probability difference | 0.012 | 0.012 | +| worst largest per-option probability difference | 0.111 | 0.111 | +| mean KL(original ‖ int4) | 0.0016 | 0.0016 | +| peak GPU memory | 12.3 GiB | 9.1 GiB | + +The largest differences are on borderline questions. For example, a yes/no question at 0.53 / 0.47 in the original +comes out at 0.64 / 0.36 here. Clear-cut answers agree closely. This is a small suite: on any task where calibration +matters, validate against your own labelled data. Full outputs are in `quantization/eval_*.json`. + +## How it runs on a small GPU + +- Dense weights (~3.9 GiB) stay on the GPU. +- The INT4 experts (1.59 GiB per layer) stay in pinned host RAM. They are streamed to the GPU one layer at a time on + a side CUDA stream, double-buffered, so the next layer's copy overlaps the current layer's compute. +- As many layers as fit in `gpu_memory_gb − workspace_gb` stay resident. On a GPU with ~60 GB free, everything is + resident and nothing streams. +- The experts run in a Triton grouped W4A16 kernel: tensor-core dots with fp32 accumulation, using bf16 on sm_80+ and + fp16 on older GPUs. It compiles for sm_70 through sm_120. +- The 2 GB token-embedding table stays in host RAM. + +`from_pretrained` knobs: + +| knob | default | effect | +|---|---|---| +| `gpu_memory_gb` | all free memory | GPU memory the model may use | +| `workspace_gb` | 3.5 | memory kept for activations. Raise it for long contexts or big batches; lower it to keep more layers resident | +| `batch_size` | 16 | questions per forward pass | +| `max_batch_tokens` | 16384 | padded tokens per forward pass | +| `pin_memory` | `True` | set `False` if the host can't pin ~54 GB (copies become slower) | + +## Speed + +When layers stream, each forward pass moves the streamed expert weights over PCIe: ~50 GB with 2 resident layers, +53 GB with none. **Latency is therefore bounded by host-to-GPU bandwidth, not GPU compute.** + +- Measured on an H200 NVL host (PCIe Gen5) with the GPU capped at 15 GiB and at 11 GiB: median **2.6 s per request**. +- Expected at roughly 25 GB/s (PCIe Gen4 x16): about 2–2.5 s per forward pass. +- Expected at roughly 12 GB/s (PCIe Gen4 x8, or Gen3 x16): about 4.5–5 s per forward pass. + +A forward pass scores up to `batch_size` questions together, so put many questions, or many requests via +`decide_batch`, into one call. Throughput then scales almost linearly up to the batch limit. There are no CUDA graphs, +because weights move every pass. The first call also compiles the Triton kernel. + +## Files + +| file | contents | +|---|---| +| `model-layer-*.safetensors` | per layer: FP8 dense tensors (unchanged) plus `experts.{gate_up,down}_{qweight,scales}` (INT4) | +| `model-embed.safetensors`, `decision_head.safetensors`, `tokenizer*`, `config.json`, `*_sarvam_moe.py` | from the original release; `config.json` documents the quantization | +| `prism_int4.py`, `example.py`, `requirements.txt` | inference code | +| `quantization/` | the quantization script, calibration generator, per-layer stats and evaluation outputs | + +INT4 packing: `qweight` is uint8 `[E, N, K/2]`. In each 128-wide group along K, byte *j* holds element *j* in its low +nibble and element *j+64* in its high nibble; value = nibble − 8. `scales` is fp32 `[E, N, K/128]`, and +weight = value × scale. + +## License and credit + +Apache 2.0, the same as the original. All model weights derive from +[sarvamai/prism](https://huggingface.co/sarvamai/prism) by Sarvam AI. This is an unofficial quantization: for the +original, its evaluation and the hosted API, see Sarvam's [release blog](https://www.sarvam.ai/blogs/sarvam-prism). diff --git a/assets/sarvam-prism.png b/assets/sarvam-prism.png new file mode 100644 index 0000000000000000000000000000000000000000..49ca9ce9fabb4309f2e6a4d25ddc5dcc95d8435f --- /dev/null +++ b/assets/sarvam-prism.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aed1306570d04220db70c3edd53a5cf3296a0f9a02e5b42758e6a2870091fe34 +size 2406313 diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..94cc5f938bf30b35b5d814b0cf0169c66d61f602 --- /dev/null +++ b/config.json @@ -0,0 +1,75 @@ +{ + "architectures": [ + "SarvamMLAForCausalLM" + ], + "attention_dropout": 0.0, + "attn_implementation": null, + "auto_map": { + "AutoConfig": "configuration_sarvam_moe.SarvamMLAConfig", + "AutoModel": "modeling_sarvam_moe.SarvamMLAModel", + "AutoModelForCausalLM": "modeling_sarvam_moe.SarvamMLAForCausalLM" + }, + "default_theta": 10000.0, + "dtype": "bfloat16", + "embedding_dropout": 0.0, + "eos_token_id": 1, + "first_k_dense_replace": 1, + "head_dim": 576, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.006, + "intermediate_size": 16384, + "kv_lora_rank": 512, + "max_position_embeddings": 131072, + "model_type": "sarvam_mla", + "moe_intermediate_size": 2048, + "moe_router_enable_expert_bias": true, + "num_attention_heads": 64, + "num_experts": 128, + "num_experts_per_tok": 8, + "num_hidden_layers": 32, + "num_shared_experts": 1, + "output_dropout": 0.0, + "output_router_logits": false, + "pad_token_id": 0, + "q_head_dim": 192, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "deepseek_yarn" + }, + "rope_theta": 10000.0, + "routed_scaling_factor": 2.5, + "tie_word_embeddings": false, + "transformers_version": "4.57.2", + "use_cache": false, + "use_qk_norm": true, + "v_head_dim": 128, + "vocab_size": 262144, + "torch_dtype": "bfloat16", + "quantization": { + "routed_experts": { + "format": "int4", + "method": "GPTQ (damp 0.01, MSE clip search)", + "symmetric": true, + "group_size": 128, + "scale_dtype": "float32", + "packing": "uint8 [E, N, K/2]; in each 128-group byte j = elem j (low nibble) | elem j+64 (high nibble); value = nibble - 8", + "calibration": "262144 tokens of Prism-format requests over multilingual Wikipedia (11 languages)" + }, + "other_projections": { + "format": "float8_e4m3fn", + "weight_scale": "per output row (absmax / 448)", + "note": "unchanged from sarvamai/prism" + }, + "embeddings_norms": "bfloat16", + "router": "float32" + } +} \ No newline at end of file diff --git a/configuration_sarvam_moe.py b/configuration_sarvam_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..50f5389f62ac0c79132f4a5bdec45ba0c5d00e30 --- /dev/null +++ b/configuration_sarvam_moe.py @@ -0,0 +1,140 @@ +from transformers.configuration_utils import PretrainedConfig + + +class SarvamMLAConfig(PretrainedConfig): + model_type = "sarvam_mla" + + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + } + + def __init__( + self, + vocab_size: int = 262144, + hidden_size: int = 4096, + num_hidden_layers: int = 32, + intermediate_size: int = 16384, + moe_intermediate_size: int = 2048, + num_experts: int = 128, + num_experts_per_tok: int = 8, + num_shared_experts: int = 1, + first_k_dense_replace: int = 1, + num_attention_heads: int = 64, + qk_rope_head_dim: int = 64, + qk_nope_head_dim: int = 128, + kv_lora_rank: int = 512, + v_head_dim: int = 128, + max_position_embeddings: int = 4096, + rope_theta: float = 10000.0, + rope_scaling: dict = None, + attention_dropout: float = 0.0, + output_dropout: float = 0.0, + rms_norm_eps: float = 1e-6, + hidden_act: str = "silu", + use_cache: bool = True, + use_qk_norm: bool = True, + moe_router_enable_expert_bias: bool = True, + routed_scaling_factor: float = 2.5, + output_router_logits: bool = False, + tie_word_embeddings: bool = False, + pad_token_id: int = 0, + eos_token_id: int = 1, + embedding_dropout: float = 0.0, + initializer_range: float = 0.006, + attn_implementation: str = "eager", + **kwargs, + ): + # core geometry + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.intermediate_size = intermediate_size + self.num_attention_heads = num_attention_heads + self.max_position_embeddings = max_position_embeddings + + # MLA geometry + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.kv_lora_rank = kv_lora_rank + self.v_head_dim = v_head_dim + # convenient derived dim + self.q_head_dim = qk_rope_head_dim + qk_nope_head_dim + # vLLM MLA expects "head size" = Lkv + R, not hidden_size/num_heads. + self.head_dim = int(self.kv_lora_rank + self.qk_rope_head_dim) + + # MoE + self.moe_intermediate_size = moe_intermediate_size + self.num_experts = num_experts + self.num_experts_per_tok = num_experts_per_tok + self.num_shared_experts = num_shared_experts + self.first_k_dense_replace = first_k_dense_replace + + # Router + self.moe_router_enable_expert_bias = moe_router_enable_expert_bias + self.routed_scaling_factor = routed_scaling_factor + self.output_router_logits = output_router_logits + + # dropouts / norms / init + self.attention_dropout = attention_dropout + self.output_dropout = output_dropout + self.embedding_dropout = embedding_dropout + self.rms_norm_eps = rms_norm_eps + self.initializer_range = initializer_range + self.hidden_act = hidden_act + + # rope / cache + self.rope_theta = rope_theta + self.use_cache = use_cache + self.use_qk_norm = use_qk_norm + self.rope_scaling = rope_scaling + self.default_theta = 10000.0 + + if self.rope_scaling is None: + self.rope_scaling = { + 'beta_fast': 32, + 'beta_slow': 1, + 'factor': 40, + 'mscale': 1.0, + 'mscale_all_dim': 1.0, + 'original_max_position_embeddings': 4096, + 'rope_type': 'deepseek_yarn', + } + + self.attn_implementation = attn_implementation + self._attn_implementation = attn_implementation + + if "_attn_implementation" in kwargs: + self._attn_implementation = kwargs.pop("_attn_implementation") + if hasattr(self, "attn_implementation"): + self.attn_implementation = self._attn_implementation + + super().__init__( + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + self.validate_rope(ignore_keys=ignore_keys_at_rope_validation) + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs \ No newline at end of file diff --git a/decision_head.safetensors b/decision_head.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2dd44a891fef85126619a7067096e2d070287ff3 --- /dev/null +++ b/decision_head.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a84038851eafcf30ddd086f331fcda944fa8cb1b5a9aa780c078c3f6c5b5df10 +size 117841696 diff --git a/example.py b/example.py new file mode 100644 index 0000000000000000000000000000000000000000..b91aea97255585b73b3f7a1c933f01b25ade693a --- /dev/null +++ b/example.py @@ -0,0 +1,21 @@ +"""Run one request through Sarvam-Prism INT4 (any CUDA GPU with ~10-16 GB, >= 64 GB host RAM). + python example.py # from the Hugging Face Hub + python example.py /path/to/local/copy # from a local copy + python example.py . 12 # cap GPU memory use at 12 GiB +""" +import json, sys +from prism_int4 import Prism + +model = Prism.from_pretrained(sys.argv[1] if len(sys.argv) > 1 else "jc-914/prism-quantized", + gpu_memory_gb=float(sys.argv[2]) if len(sys.argv) > 2 else None) +answers = model.decide( + state={"ticket": "My card was charged twice for the same order and I still haven't got a refund."}, + questions={ + "team": {"type": "choice", "instructions": "Which team should handle this ticket?", + "criteria": {"billing": "payments, charges and refunds", "tech": "app or website problems", + "delivery": "shipping and order status"}}, + "urgent": {"type": "yes_no", "instructions": "Does this need a reply within the hour?"}, + "anger": {"type": "score", "instructions": "How upset is the customer?", + "criteria": ["calm", "annoyed", "angry", "furious"]}, + }) +print(json.dumps(answers, indent=2)) diff --git a/model-embed.safetensors b/model-embed.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f282e003ae28e6320f69318e61272d7b4a22a44c --- /dev/null +++ b/model-embed.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5c1e74d14913aa92d5e66728229976f3a0532393118ec4ace49cc16eeaf79ed0 +size 2147492072 diff --git a/model-layer-00.safetensors b/model-layer-00.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6fa312614da6196cdc7d8379e5f65b2d79407c42 --- /dev/null +++ b/model-layer-00.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b5b39439385eb5912b2ee14d1451ef269340cf1f2125386e941b3a9b6db896c +size 296260776 diff --git a/model-layer-01.safetensors b/model-layer-01.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f0095ced44c03d21bae7416d3cce8a169ec65f72 --- /dev/null +++ b/model-layer-01.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f3b2a3d87ac2b8dbd4ef1c0de9b92c2313f50c81bbfcc1736f275fd9e3ce08ad +size 1833359864 diff --git a/model-layer-03.safetensors b/model-layer-03.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..84f4023e66c2e4addc00e6fc653f56f603187736 --- /dev/null +++ b/model-layer-03.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d6aee4a1c0b7e9146967a71c0b780ef8d3af08c947be984cf69fb64dd4014b91 +size 1833359864 diff --git a/model-layer-04.safetensors b/model-layer-04.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a937c1720b4235dadd9669fc9436637a980acd8d --- /dev/null +++ b/model-layer-04.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e8832e679edbed37ede766dc88a5df515c59a481b964e66b10019a28141fa8e7 +size 1833359864 diff --git a/model-layer-05.safetensors b/model-layer-05.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ba079be5af91fffdf7589481a4d1ff9c7337f50a --- /dev/null +++ b/model-layer-05.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b100f7e77f9ff8aed5a276f593b9a5b35dc9d1d7278f9b3448cdfd8effe3f10e +size 1833359864 diff --git a/model-layer-06.safetensors b/model-layer-06.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e76f06821a0a49b08444b1d4fa6b3f8cdabf8569 --- /dev/null +++ b/model-layer-06.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4709620644b469259fa00b471d30e411bdd4fb80526a664e3e84b905e00200a +size 1833359864 diff --git a/model-layer-09.safetensors b/model-layer-09.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..aba9949dc9d0ac18dd5761384d44f5fbb4565cda --- /dev/null +++ b/model-layer-09.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8fe2d5a948200bf2d611432e2365704b5e5a21070c464500c4e730ebc98ffa1c +size 1833359864 diff --git a/model-layer-10.safetensors b/model-layer-10.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4af31617b9c0f9301126ed1ff123a7f277caf0eb --- /dev/null +++ b/model-layer-10.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e8a6a6cad527844b095686049a296bfb59e749edf5f9d6fb89b0439330a71106 +size 1833359880 diff --git a/model-layer-11.safetensors b/model-layer-11.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3b0229e9417a21a99440501c274a02745a7097fd --- /dev/null +++ b/model-layer-11.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ac31ddccc3e2ebe4b14d2dd3432f05c9cc5013129ac83299a5530852e5b1b3bf +size 1833359880 diff --git a/model-layer-12.safetensors b/model-layer-12.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1edc5d449651d342440c9d9bb4676d38cbd0da1e --- /dev/null +++ b/model-layer-12.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07bbc54d7c2a7dfd4e561b0b2c05317e9775e235d0763052a5ee4f606cf40ada +size 1833359880 diff --git a/model-layer-13.safetensors b/model-layer-13.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5d0c1f409a0af6977a97a49e6a901f445845315d --- /dev/null +++ b/model-layer-13.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6713191dc31b9b4a44a59d104c02d25e0df219991462c814fe3797ce5d1d6da5 +size 1833359880 diff --git a/model-layer-15.safetensors b/model-layer-15.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..90e7d5b91efafbdd8453867022c3f0f771db1428 --- /dev/null +++ b/model-layer-15.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6288a24d1cdf7a980edf1aaf09da7e5d7805424421adf25ee671b68a14263ef7 +size 1833359880 diff --git a/model-layer-16.safetensors b/model-layer-16.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8a226358b8caee3f340c37739ef503713cf698f8 --- /dev/null +++ b/model-layer-16.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab254c6675fa5e4c263126f807ff508d8856174bfa928b3a75906461321345db +size 1833359880 diff --git a/model-layer-17.safetensors b/model-layer-17.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..25eeff9414427896acfe31c6c43216663680e24a --- /dev/null +++ b/model-layer-17.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c6e0e8681e271b0a6faecb469e79c94680e2b70365c3e045383b93706f2e614 +size 1833359880 diff --git a/model-layer-18.safetensors b/model-layer-18.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7e1438057b1d090cda83afdae07eca19e7bcce79 --- /dev/null +++ b/model-layer-18.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:888d98c00b5786caccfd96fb7b3e057e45f2789b37bd3086a22634f127916c7a +size 1833359880 diff --git a/model-layer-19.safetensors b/model-layer-19.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8bd8665144f0946c1c0890c67cd863ae35924fff --- /dev/null +++ b/model-layer-19.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b8ad3f47ea73a30f4c448733d22bbac0fb304278ef893a33471bc53e7958de3a +size 1833359880 diff --git a/model-layer-21.safetensors b/model-layer-21.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..cfcdb64358fb8efedae9c143572e627102b94d0c --- /dev/null +++ b/model-layer-21.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f666deeb9a63e5c6925ca6a23fbf2c78e9f59ba5a5231ce0d95ad87ec8a7ac75 +size 1833359880 diff --git a/model-layer-22.safetensors b/model-layer-22.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4a57a5c98cf24ab6d8160a6c6a416e61129a34e9 --- /dev/null +++ b/model-layer-22.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ca505c768d507d39e5813af4e1af210d5c1fd3cd2cad7b9978ae598c8e1defcf +size 1833359880 diff --git a/model-layer-23.safetensors b/model-layer-23.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a3af9d4ca9866b4e13be09758986c67ed0f37dd7 --- /dev/null +++ b/model-layer-23.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f3866bc7af2e72e3e3287bee9a4802906aef4618d4b2dfdfd81955e1d6fc3cba +size 1833359880 diff --git a/model-layer-24.safetensors b/model-layer-24.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7b2eecf587505013867ce171870329fd81289d3e --- /dev/null +++ b/model-layer-24.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5424fbfcbf47d0b5b9ab6a610025358b2b9c751143a56adf1e16a3519cf4bc0b +size 1833359880 diff --git a/model-layer-25.safetensors b/model-layer-25.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7cddcff1012cc8d5cbd838d5c94aa12806344c6d --- /dev/null +++ b/model-layer-25.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a0eda1a562ca6c20aaa9146366b81a3bad80db4f9fe25e19b11673a10f0e9e08 +size 1833359880 diff --git a/model-layer-26.safetensors b/model-layer-26.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..dffe4a1389b35b2821254b102a6ae33ef35414c4 --- /dev/null +++ b/model-layer-26.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:17c941950bb822d84902008a249c77395b7ea585ebe54f470e033105bdef107b +size 1833359880 diff --git a/model-layer-27.safetensors b/model-layer-27.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c0f4880edb4ada0665f19211dd250a2cbfbc34a9 --- /dev/null +++ b/model-layer-27.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4bd934990c5c75c6ca8099b1e6f53e13c80bbfe2faa312e9df5a0a8d63cc96f +size 1833359880 diff --git a/model-layer-28.safetensors b/model-layer-28.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..bfdb7b98fdcf7611e6873353b8412e4c396252a3 --- /dev/null +++ b/model-layer-28.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:71362c040ec648cc61d980b5b3cebd51108f9eca4834e0d7cb8f388267ac7c58 +size 1833359880 diff --git a/model-layer-29.safetensors b/model-layer-29.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e2daa0868b03aa509c6ed6a46a78f0769ef974d6 --- /dev/null +++ b/model-layer-29.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:19789a9593c52702ef02f0158115b780dfa393982cc0d78165880d48554c676a +size 1833359880 diff --git a/model-layer-30.safetensors b/model-layer-30.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b5415eccf03e2e303ef8b1e4b80fa52722bc2dd1 --- /dev/null +++ b/model-layer-30.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1c08ae7ec56d446f523e3b0cdf339bd47c490fda2425a8a17964e5b647b9e907 +size 1833359880 diff --git a/model-layer-31.safetensors b/model-layer-31.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..91cd2e41ddd8cffca2ae600e881c9349c0c392fd --- /dev/null +++ b/model-layer-31.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3197f4d6bce9c3f3603228f4b7d0f5ba4b5a1803876a50f5580a0a6eb19cd51d +size 1833359880 diff --git a/model.safetensors.index.json b/model.safetensors.index.json new file mode 100644 index 0000000000000000000000000000000000000000..7068cbbde9511ed045a341cc5f09cddf8f1cc4ca --- /dev/null +++ b/model.safetensors.index.json @@ -0,0 +1,737 @@ +{ + "metadata": {}, + "weight_map": { + "model.embed_tokens.weight": "model-embed.safetensors", + "model.norm.weight": "model-embed.safetensors", + "model.layers.0.input_layernorm.weight": "model-layer-00.safetensors", + "model.layers.0.mlp.down_proj.weight": "model-layer-00.safetensors", + "model.layers.0.mlp.down_proj.weight_scale": "model-layer-00.safetensors", + "model.layers.0.mlp.gate_proj.weight": "model-layer-00.safetensors", + "model.layers.0.mlp.gate_proj.weight_scale": "model-layer-00.safetensors", + "model.layers.0.mlp.up_proj.weight": "model-layer-00.safetensors", + "model.layers.0.mlp.up_proj.weight_scale": "model-layer-00.safetensors", + "model.layers.0.post_attention_layernorm.weight": "model-layer-00.safetensors", + "model.layers.0.self_attn.kv_a_layernorm.weight": "model-layer-00.safetensors", + "model.layers.0.self_attn.kv_a_proj_with_mqa.weight": "model-layer-00.safetensors", + "model.layers.0.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-00.safetensors", + "model.layers.0.self_attn.kv_b_proj.weight": "model-layer-00.safetensors", + "model.layers.0.self_attn.kv_b_proj.weight_scale": "model-layer-00.safetensors", + "model.layers.0.self_attn.o_proj.weight": "model-layer-00.safetensors", + "model.layers.0.self_attn.o_proj.weight_scale": "model-layer-00.safetensors", + "model.layers.0.self_attn.q_proj.weight": "model-layer-00.safetensors", + "model.layers.0.self_attn.q_proj.weight_scale": "model-layer-00.safetensors", + "model.layers.1.input_layernorm.weight": "model-layer-01.safetensors", + "model.layers.1.mlp.experts.down_qweight": "model-layer-01.safetensors", + "model.layers.1.mlp.experts.down_scales": "model-layer-01.safetensors", + "model.layers.1.mlp.experts.gate_up_qweight": "model-layer-01.safetensors", + "model.layers.1.mlp.experts.gate_up_scales": "model-layer-01.safetensors", + "model.layers.1.mlp.gate.e_score_correction_bias": "model-layer-01.safetensors", + "model.layers.1.mlp.gate.weight": "model-layer-01.safetensors", + "model.layers.1.mlp.shared_experts.down_proj.weight": "model-layer-01.safetensors", + "model.layers.1.mlp.shared_experts.down_proj.weight_scale": "model-layer-01.safetensors", + "model.layers.1.mlp.shared_experts.gate_proj.weight": "model-layer-01.safetensors", + "model.layers.1.mlp.shared_experts.gate_proj.weight_scale": "model-layer-01.safetensors", + "model.layers.1.mlp.shared_experts.up_proj.weight": "model-layer-01.safetensors", + "model.layers.1.mlp.shared_experts.up_proj.weight_scale": "model-layer-01.safetensors", + "model.layers.1.post_attention_layernorm.weight": "model-layer-01.safetensors", + "model.layers.1.self_attn.kv_a_layernorm.weight": "model-layer-01.safetensors", + "model.layers.1.self_attn.kv_a_proj_with_mqa.weight": "model-layer-01.safetensors", + "model.layers.1.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-01.safetensors", + "model.layers.1.self_attn.kv_b_proj.weight": "model-layer-01.safetensors", + "model.layers.1.self_attn.kv_b_proj.weight_scale": "model-layer-01.safetensors", + "model.layers.1.self_attn.o_proj.weight": "model-layer-01.safetensors", + "model.layers.1.self_attn.o_proj.weight_scale": "model-layer-01.safetensors", + "model.layers.1.self_attn.q_proj.weight": "model-layer-01.safetensors", + "model.layers.1.self_attn.q_proj.weight_scale": "model-layer-01.safetensors", + "model.layers.2.input_layernorm.weight": "model-layer-02.safetensors", + "model.layers.2.mlp.experts.down_qweight": "model-layer-02.safetensors", + "model.layers.2.mlp.experts.down_scales": "model-layer-02.safetensors", + "model.layers.2.mlp.experts.gate_up_qweight": "model-layer-02.safetensors", + "model.layers.2.mlp.experts.gate_up_scales": "model-layer-02.safetensors", + "model.layers.2.mlp.gate.e_score_correction_bias": "model-layer-02.safetensors", + "model.layers.2.mlp.gate.weight": "model-layer-02.safetensors", + "model.layers.2.mlp.shared_experts.down_proj.weight": "model-layer-02.safetensors", + "model.layers.2.mlp.shared_experts.down_proj.weight_scale": "model-layer-02.safetensors", + "model.layers.2.mlp.shared_experts.gate_proj.weight": "model-layer-02.safetensors", + "model.layers.2.mlp.shared_experts.gate_proj.weight_scale": "model-layer-02.safetensors", + "model.layers.2.mlp.shared_experts.up_proj.weight": "model-layer-02.safetensors", + "model.layers.2.mlp.shared_experts.up_proj.weight_scale": "model-layer-02.safetensors", + "model.layers.2.post_attention_layernorm.weight": "model-layer-02.safetensors", + "model.layers.2.self_attn.kv_a_layernorm.weight": "model-layer-02.safetensors", + "model.layers.2.self_attn.kv_a_proj_with_mqa.weight": "model-layer-02.safetensors", + "model.layers.2.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-02.safetensors", + "model.layers.2.self_attn.kv_b_proj.weight": "model-layer-02.safetensors", + "model.layers.2.self_attn.kv_b_proj.weight_scale": "model-layer-02.safetensors", + "model.layers.2.self_attn.o_proj.weight": "model-layer-02.safetensors", + "model.layers.2.self_attn.o_proj.weight_scale": "model-layer-02.safetensors", + "model.layers.2.self_attn.q_proj.weight": "model-layer-02.safetensors", + "model.layers.2.self_attn.q_proj.weight_scale": "model-layer-02.safetensors", + "model.layers.3.input_layernorm.weight": "model-layer-03.safetensors", + "model.layers.3.mlp.experts.down_qweight": "model-layer-03.safetensors", + "model.layers.3.mlp.experts.down_scales": "model-layer-03.safetensors", + "model.layers.3.mlp.experts.gate_up_qweight": "model-layer-03.safetensors", + "model.layers.3.mlp.experts.gate_up_scales": "model-layer-03.safetensors", + "model.layers.3.mlp.gate.e_score_correction_bias": "model-layer-03.safetensors", + "model.layers.3.mlp.gate.weight": "model-layer-03.safetensors", + "model.layers.3.mlp.shared_experts.down_proj.weight": "model-layer-03.safetensors", + "model.layers.3.mlp.shared_experts.down_proj.weight_scale": "model-layer-03.safetensors", + "model.layers.3.mlp.shared_experts.gate_proj.weight": "model-layer-03.safetensors", + "model.layers.3.mlp.shared_experts.gate_proj.weight_scale": "model-layer-03.safetensors", + "model.layers.3.mlp.shared_experts.up_proj.weight": "model-layer-03.safetensors", + "model.layers.3.mlp.shared_experts.up_proj.weight_scale": "model-layer-03.safetensors", + "model.layers.3.post_attention_layernorm.weight": "model-layer-03.safetensors", + "model.layers.3.self_attn.kv_a_layernorm.weight": "model-layer-03.safetensors", + "model.layers.3.self_attn.kv_a_proj_with_mqa.weight": "model-layer-03.safetensors", + "model.layers.3.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-03.safetensors", + "model.layers.3.self_attn.kv_b_proj.weight": "model-layer-03.safetensors", + "model.layers.3.self_attn.kv_b_proj.weight_scale": "model-layer-03.safetensors", + "model.layers.3.self_attn.o_proj.weight": "model-layer-03.safetensors", + "model.layers.3.self_attn.o_proj.weight_scale": "model-layer-03.safetensors", + "model.layers.3.self_attn.q_proj.weight": "model-layer-03.safetensors", + "model.layers.3.self_attn.q_proj.weight_scale": "model-layer-03.safetensors", + "model.layers.4.input_layernorm.weight": "model-layer-04.safetensors", + "model.layers.4.mlp.experts.down_qweight": "model-layer-04.safetensors", + "model.layers.4.mlp.experts.down_scales": "model-layer-04.safetensors", + "model.layers.4.mlp.experts.gate_up_qweight": "model-layer-04.safetensors", + "model.layers.4.mlp.experts.gate_up_scales": "model-layer-04.safetensors", + "model.layers.4.mlp.gate.e_score_correction_bias": "model-layer-04.safetensors", + "model.layers.4.mlp.gate.weight": "model-layer-04.safetensors", + "model.layers.4.mlp.shared_experts.down_proj.weight": "model-layer-04.safetensors", + "model.layers.4.mlp.shared_experts.down_proj.weight_scale": "model-layer-04.safetensors", + "model.layers.4.mlp.shared_experts.gate_proj.weight": "model-layer-04.safetensors", + "model.layers.4.mlp.shared_experts.gate_proj.weight_scale": "model-layer-04.safetensors", + "model.layers.4.mlp.shared_experts.up_proj.weight": "model-layer-04.safetensors", + "model.layers.4.mlp.shared_experts.up_proj.weight_scale": "model-layer-04.safetensors", + "model.layers.4.post_attention_layernorm.weight": "model-layer-04.safetensors", + "model.layers.4.self_attn.kv_a_layernorm.weight": "model-layer-04.safetensors", + "model.layers.4.self_attn.kv_a_proj_with_mqa.weight": "model-layer-04.safetensors", + "model.layers.4.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-04.safetensors", + "model.layers.4.self_attn.kv_b_proj.weight": "model-layer-04.safetensors", + "model.layers.4.self_attn.kv_b_proj.weight_scale": "model-layer-04.safetensors", + "model.layers.4.self_attn.o_proj.weight": "model-layer-04.safetensors", + "model.layers.4.self_attn.o_proj.weight_scale": "model-layer-04.safetensors", + "model.layers.4.self_attn.q_proj.weight": "model-layer-04.safetensors", + "model.layers.4.self_attn.q_proj.weight_scale": "model-layer-04.safetensors", + "model.layers.5.input_layernorm.weight": "model-layer-05.safetensors", + "model.layers.5.mlp.experts.down_qweight": "model-layer-05.safetensors", + "model.layers.5.mlp.experts.down_scales": "model-layer-05.safetensors", + "model.layers.5.mlp.experts.gate_up_qweight": "model-layer-05.safetensors", + "model.layers.5.mlp.experts.gate_up_scales": "model-layer-05.safetensors", + "model.layers.5.mlp.gate.e_score_correction_bias": "model-layer-05.safetensors", + "model.layers.5.mlp.gate.weight": "model-layer-05.safetensors", + "model.layers.5.mlp.shared_experts.down_proj.weight": "model-layer-05.safetensors", + "model.layers.5.mlp.shared_experts.down_proj.weight_scale": "model-layer-05.safetensors", + "model.layers.5.mlp.shared_experts.gate_proj.weight": "model-layer-05.safetensors", + "model.layers.5.mlp.shared_experts.gate_proj.weight_scale": "model-layer-05.safetensors", + "model.layers.5.mlp.shared_experts.up_proj.weight": "model-layer-05.safetensors", + "model.layers.5.mlp.shared_experts.up_proj.weight_scale": "model-layer-05.safetensors", + "model.layers.5.post_attention_layernorm.weight": "model-layer-05.safetensors", + "model.layers.5.self_attn.kv_a_layernorm.weight": "model-layer-05.safetensors", + "model.layers.5.self_attn.kv_a_proj_with_mqa.weight": "model-layer-05.safetensors", + "model.layers.5.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-05.safetensors", + "model.layers.5.self_attn.kv_b_proj.weight": "model-layer-05.safetensors", + "model.layers.5.self_attn.kv_b_proj.weight_scale": "model-layer-05.safetensors", + "model.layers.5.self_attn.o_proj.weight": "model-layer-05.safetensors", + "model.layers.5.self_attn.o_proj.weight_scale": "model-layer-05.safetensors", + "model.layers.5.self_attn.q_proj.weight": "model-layer-05.safetensors", + "model.layers.5.self_attn.q_proj.weight_scale": "model-layer-05.safetensors", + "model.layers.6.input_layernorm.weight": "model-layer-06.safetensors", + "model.layers.6.mlp.experts.down_qweight": "model-layer-06.safetensors", + "model.layers.6.mlp.experts.down_scales": "model-layer-06.safetensors", + "model.layers.6.mlp.experts.gate_up_qweight": "model-layer-06.safetensors", + "model.layers.6.mlp.experts.gate_up_scales": "model-layer-06.safetensors", + "model.layers.6.mlp.gate.e_score_correction_bias": "model-layer-06.safetensors", + "model.layers.6.mlp.gate.weight": "model-layer-06.safetensors", + "model.layers.6.mlp.shared_experts.down_proj.weight": "model-layer-06.safetensors", + "model.layers.6.mlp.shared_experts.down_proj.weight_scale": "model-layer-06.safetensors", + "model.layers.6.mlp.shared_experts.gate_proj.weight": "model-layer-06.safetensors", + "model.layers.6.mlp.shared_experts.gate_proj.weight_scale": "model-layer-06.safetensors", + "model.layers.6.mlp.shared_experts.up_proj.weight": "model-layer-06.safetensors", + "model.layers.6.mlp.shared_experts.up_proj.weight_scale": "model-layer-06.safetensors", + "model.layers.6.post_attention_layernorm.weight": "model-layer-06.safetensors", + "model.layers.6.self_attn.kv_a_layernorm.weight": "model-layer-06.safetensors", + "model.layers.6.self_attn.kv_a_proj_with_mqa.weight": "model-layer-06.safetensors", + "model.layers.6.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-06.safetensors", + "model.layers.6.self_attn.kv_b_proj.weight": "model-layer-06.safetensors", + "model.layers.6.self_attn.kv_b_proj.weight_scale": "model-layer-06.safetensors", + "model.layers.6.self_attn.o_proj.weight": "model-layer-06.safetensors", + "model.layers.6.self_attn.o_proj.weight_scale": "model-layer-06.safetensors", + "model.layers.6.self_attn.q_proj.weight": "model-layer-06.safetensors", + "model.layers.6.self_attn.q_proj.weight_scale": "model-layer-06.safetensors", + "model.layers.7.input_layernorm.weight": "model-layer-07.safetensors", + "model.layers.7.mlp.experts.down_qweight": "model-layer-07.safetensors", + "model.layers.7.mlp.experts.down_scales": "model-layer-07.safetensors", + "model.layers.7.mlp.experts.gate_up_qweight": "model-layer-07.safetensors", + "model.layers.7.mlp.experts.gate_up_scales": "model-layer-07.safetensors", + "model.layers.7.mlp.gate.e_score_correction_bias": "model-layer-07.safetensors", + "model.layers.7.mlp.gate.weight": "model-layer-07.safetensors", + "model.layers.7.mlp.shared_experts.down_proj.weight": "model-layer-07.safetensors", + "model.layers.7.mlp.shared_experts.down_proj.weight_scale": "model-layer-07.safetensors", + "model.layers.7.mlp.shared_experts.gate_proj.weight": "model-layer-07.safetensors", + "model.layers.7.mlp.shared_experts.gate_proj.weight_scale": "model-layer-07.safetensors", + "model.layers.7.mlp.shared_experts.up_proj.weight": "model-layer-07.safetensors", + "model.layers.7.mlp.shared_experts.up_proj.weight_scale": "model-layer-07.safetensors", + "model.layers.7.post_attention_layernorm.weight": "model-layer-07.safetensors", + "model.layers.7.self_attn.kv_a_layernorm.weight": "model-layer-07.safetensors", + "model.layers.7.self_attn.kv_a_proj_with_mqa.weight": "model-layer-07.safetensors", + "model.layers.7.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-07.safetensors", + "model.layers.7.self_attn.kv_b_proj.weight": "model-layer-07.safetensors", + "model.layers.7.self_attn.kv_b_proj.weight_scale": "model-layer-07.safetensors", + "model.layers.7.self_attn.o_proj.weight": "model-layer-07.safetensors", + "model.layers.7.self_attn.o_proj.weight_scale": "model-layer-07.safetensors", + "model.layers.7.self_attn.q_proj.weight": "model-layer-07.safetensors", + "model.layers.7.self_attn.q_proj.weight_scale": "model-layer-07.safetensors", + "model.layers.8.input_layernorm.weight": "model-layer-08.safetensors", + "model.layers.8.mlp.experts.down_qweight": "model-layer-08.safetensors", + "model.layers.8.mlp.experts.down_scales": "model-layer-08.safetensors", + "model.layers.8.mlp.experts.gate_up_qweight": "model-layer-08.safetensors", + "model.layers.8.mlp.experts.gate_up_scales": "model-layer-08.safetensors", + "model.layers.8.mlp.gate.e_score_correction_bias": "model-layer-08.safetensors", + "model.layers.8.mlp.gate.weight": "model-layer-08.safetensors", + "model.layers.8.mlp.shared_experts.down_proj.weight": "model-layer-08.safetensors", + "model.layers.8.mlp.shared_experts.down_proj.weight_scale": "model-layer-08.safetensors", + "model.layers.8.mlp.shared_experts.gate_proj.weight": "model-layer-08.safetensors", + "model.layers.8.mlp.shared_experts.gate_proj.weight_scale": "model-layer-08.safetensors", + "model.layers.8.mlp.shared_experts.up_proj.weight": "model-layer-08.safetensors", + "model.layers.8.mlp.shared_experts.up_proj.weight_scale": "model-layer-08.safetensors", + "model.layers.8.post_attention_layernorm.weight": "model-layer-08.safetensors", + "model.layers.8.self_attn.kv_a_layernorm.weight": "model-layer-08.safetensors", + "model.layers.8.self_attn.kv_a_proj_with_mqa.weight": "model-layer-08.safetensors", + "model.layers.8.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-08.safetensors", + "model.layers.8.self_attn.kv_b_proj.weight": "model-layer-08.safetensors", + "model.layers.8.self_attn.kv_b_proj.weight_scale": "model-layer-08.safetensors", + "model.layers.8.self_attn.o_proj.weight": "model-layer-08.safetensors", + "model.layers.8.self_attn.o_proj.weight_scale": "model-layer-08.safetensors", + "model.layers.8.self_attn.q_proj.weight": "model-layer-08.safetensors", + "model.layers.8.self_attn.q_proj.weight_scale": "model-layer-08.safetensors", + "model.layers.9.input_layernorm.weight": "model-layer-09.safetensors", + "model.layers.9.mlp.experts.down_qweight": "model-layer-09.safetensors", + "model.layers.9.mlp.experts.down_scales": "model-layer-09.safetensors", + "model.layers.9.mlp.experts.gate_up_qweight": "model-layer-09.safetensors", + "model.layers.9.mlp.experts.gate_up_scales": "model-layer-09.safetensors", + "model.layers.9.mlp.gate.e_score_correction_bias": "model-layer-09.safetensors", + "model.layers.9.mlp.gate.weight": "model-layer-09.safetensors", + "model.layers.9.mlp.shared_experts.down_proj.weight": "model-layer-09.safetensors", + "model.layers.9.mlp.shared_experts.down_proj.weight_scale": "model-layer-09.safetensors", + "model.layers.9.mlp.shared_experts.gate_proj.weight": "model-layer-09.safetensors", + "model.layers.9.mlp.shared_experts.gate_proj.weight_scale": "model-layer-09.safetensors", + "model.layers.9.mlp.shared_experts.up_proj.weight": "model-layer-09.safetensors", + "model.layers.9.mlp.shared_experts.up_proj.weight_scale": "model-layer-09.safetensors", + "model.layers.9.post_attention_layernorm.weight": "model-layer-09.safetensors", + "model.layers.9.self_attn.kv_a_layernorm.weight": "model-layer-09.safetensors", + "model.layers.9.self_attn.kv_a_proj_with_mqa.weight": "model-layer-09.safetensors", + "model.layers.9.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-09.safetensors", + "model.layers.9.self_attn.kv_b_proj.weight": "model-layer-09.safetensors", + "model.layers.9.self_attn.kv_b_proj.weight_scale": "model-layer-09.safetensors", + "model.layers.9.self_attn.o_proj.weight": "model-layer-09.safetensors", + "model.layers.9.self_attn.o_proj.weight_scale": "model-layer-09.safetensors", + "model.layers.9.self_attn.q_proj.weight": "model-layer-09.safetensors", + "model.layers.9.self_attn.q_proj.weight_scale": "model-layer-09.safetensors", + "model.layers.10.input_layernorm.weight": "model-layer-10.safetensors", + "model.layers.10.mlp.experts.down_qweight": "model-layer-10.safetensors", + "model.layers.10.mlp.experts.down_scales": "model-layer-10.safetensors", + "model.layers.10.mlp.experts.gate_up_qweight": "model-layer-10.safetensors", + "model.layers.10.mlp.experts.gate_up_scales": "model-layer-10.safetensors", + "model.layers.10.mlp.gate.e_score_correction_bias": "model-layer-10.safetensors", + "model.layers.10.mlp.gate.weight": "model-layer-10.safetensors", + "model.layers.10.mlp.shared_experts.down_proj.weight": "model-layer-10.safetensors", + "model.layers.10.mlp.shared_experts.down_proj.weight_scale": "model-layer-10.safetensors", + "model.layers.10.mlp.shared_experts.gate_proj.weight": "model-layer-10.safetensors", + "model.layers.10.mlp.shared_experts.gate_proj.weight_scale": "model-layer-10.safetensors", + "model.layers.10.mlp.shared_experts.up_proj.weight": "model-layer-10.safetensors", + "model.layers.10.mlp.shared_experts.up_proj.weight_scale": "model-layer-10.safetensors", + "model.layers.10.post_attention_layernorm.weight": "model-layer-10.safetensors", + "model.layers.10.self_attn.kv_a_layernorm.weight": "model-layer-10.safetensors", + "model.layers.10.self_attn.kv_a_proj_with_mqa.weight": "model-layer-10.safetensors", + "model.layers.10.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-10.safetensors", + "model.layers.10.self_attn.kv_b_proj.weight": "model-layer-10.safetensors", + "model.layers.10.self_attn.kv_b_proj.weight_scale": "model-layer-10.safetensors", + "model.layers.10.self_attn.o_proj.weight": "model-layer-10.safetensors", + "model.layers.10.self_attn.o_proj.weight_scale": "model-layer-10.safetensors", + "model.layers.10.self_attn.q_proj.weight": "model-layer-10.safetensors", + "model.layers.10.self_attn.q_proj.weight_scale": "model-layer-10.safetensors", + "model.layers.11.input_layernorm.weight": "model-layer-11.safetensors", + "model.layers.11.mlp.experts.down_qweight": "model-layer-11.safetensors", + "model.layers.11.mlp.experts.down_scales": "model-layer-11.safetensors", + "model.layers.11.mlp.experts.gate_up_qweight": "model-layer-11.safetensors", + "model.layers.11.mlp.experts.gate_up_scales": "model-layer-11.safetensors", + "model.layers.11.mlp.gate.e_score_correction_bias": "model-layer-11.safetensors", + "model.layers.11.mlp.gate.weight": "model-layer-11.safetensors", + "model.layers.11.mlp.shared_experts.down_proj.weight": "model-layer-11.safetensors", + "model.layers.11.mlp.shared_experts.down_proj.weight_scale": "model-layer-11.safetensors", + "model.layers.11.mlp.shared_experts.gate_proj.weight": "model-layer-11.safetensors", + "model.layers.11.mlp.shared_experts.gate_proj.weight_scale": "model-layer-11.safetensors", + "model.layers.11.mlp.shared_experts.up_proj.weight": "model-layer-11.safetensors", + "model.layers.11.mlp.shared_experts.up_proj.weight_scale": "model-layer-11.safetensors", + "model.layers.11.post_attention_layernorm.weight": "model-layer-11.safetensors", + "model.layers.11.self_attn.kv_a_layernorm.weight": "model-layer-11.safetensors", + "model.layers.11.self_attn.kv_a_proj_with_mqa.weight": "model-layer-11.safetensors", + "model.layers.11.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-11.safetensors", + "model.layers.11.self_attn.kv_b_proj.weight": "model-layer-11.safetensors", + "model.layers.11.self_attn.kv_b_proj.weight_scale": "model-layer-11.safetensors", + "model.layers.11.self_attn.o_proj.weight": "model-layer-11.safetensors", + "model.layers.11.self_attn.o_proj.weight_scale": "model-layer-11.safetensors", + "model.layers.11.self_attn.q_proj.weight": "model-layer-11.safetensors", + "model.layers.11.self_attn.q_proj.weight_scale": "model-layer-11.safetensors", + "model.layers.12.input_layernorm.weight": "model-layer-12.safetensors", + "model.layers.12.mlp.experts.down_qweight": "model-layer-12.safetensors", + "model.layers.12.mlp.experts.down_scales": "model-layer-12.safetensors", + "model.layers.12.mlp.experts.gate_up_qweight": "model-layer-12.safetensors", + "model.layers.12.mlp.experts.gate_up_scales": "model-layer-12.safetensors", + "model.layers.12.mlp.gate.e_score_correction_bias": "model-layer-12.safetensors", + "model.layers.12.mlp.gate.weight": "model-layer-12.safetensors", + "model.layers.12.mlp.shared_experts.down_proj.weight": "model-layer-12.safetensors", + "model.layers.12.mlp.shared_experts.down_proj.weight_scale": "model-layer-12.safetensors", + "model.layers.12.mlp.shared_experts.gate_proj.weight": "model-layer-12.safetensors", + "model.layers.12.mlp.shared_experts.gate_proj.weight_scale": "model-layer-12.safetensors", + "model.layers.12.mlp.shared_experts.up_proj.weight": "model-layer-12.safetensors", + "model.layers.12.mlp.shared_experts.up_proj.weight_scale": "model-layer-12.safetensors", + "model.layers.12.post_attention_layernorm.weight": "model-layer-12.safetensors", + "model.layers.12.self_attn.kv_a_layernorm.weight": "model-layer-12.safetensors", + "model.layers.12.self_attn.kv_a_proj_with_mqa.weight": "model-layer-12.safetensors", + "model.layers.12.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-12.safetensors", + "model.layers.12.self_attn.kv_b_proj.weight": "model-layer-12.safetensors", + "model.layers.12.self_attn.kv_b_proj.weight_scale": "model-layer-12.safetensors", + "model.layers.12.self_attn.o_proj.weight": "model-layer-12.safetensors", + "model.layers.12.self_attn.o_proj.weight_scale": "model-layer-12.safetensors", + "model.layers.12.self_attn.q_proj.weight": "model-layer-12.safetensors", + "model.layers.12.self_attn.q_proj.weight_scale": "model-layer-12.safetensors", + "model.layers.13.input_layernorm.weight": "model-layer-13.safetensors", + "model.layers.13.mlp.experts.down_qweight": "model-layer-13.safetensors", + "model.layers.13.mlp.experts.down_scales": "model-layer-13.safetensors", + "model.layers.13.mlp.experts.gate_up_qweight": "model-layer-13.safetensors", + "model.layers.13.mlp.experts.gate_up_scales": "model-layer-13.safetensors", + "model.layers.13.mlp.gate.e_score_correction_bias": "model-layer-13.safetensors", + "model.layers.13.mlp.gate.weight": "model-layer-13.safetensors", + "model.layers.13.mlp.shared_experts.down_proj.weight": "model-layer-13.safetensors", + "model.layers.13.mlp.shared_experts.down_proj.weight_scale": "model-layer-13.safetensors", + "model.layers.13.mlp.shared_experts.gate_proj.weight": "model-layer-13.safetensors", + "model.layers.13.mlp.shared_experts.gate_proj.weight_scale": "model-layer-13.safetensors", + "model.layers.13.mlp.shared_experts.up_proj.weight": "model-layer-13.safetensors", + "model.layers.13.mlp.shared_experts.up_proj.weight_scale": "model-layer-13.safetensors", + "model.layers.13.post_attention_layernorm.weight": "model-layer-13.safetensors", + "model.layers.13.self_attn.kv_a_layernorm.weight": "model-layer-13.safetensors", + "model.layers.13.self_attn.kv_a_proj_with_mqa.weight": "model-layer-13.safetensors", + "model.layers.13.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-13.safetensors", + "model.layers.13.self_attn.kv_b_proj.weight": "model-layer-13.safetensors", + "model.layers.13.self_attn.kv_b_proj.weight_scale": "model-layer-13.safetensors", + "model.layers.13.self_attn.o_proj.weight": "model-layer-13.safetensors", + "model.layers.13.self_attn.o_proj.weight_scale": "model-layer-13.safetensors", + "model.layers.13.self_attn.q_proj.weight": "model-layer-13.safetensors", + "model.layers.13.self_attn.q_proj.weight_scale": "model-layer-13.safetensors", + "model.layers.14.input_layernorm.weight": "model-layer-14.safetensors", + "model.layers.14.mlp.experts.down_qweight": "model-layer-14.safetensors", + "model.layers.14.mlp.experts.down_scales": "model-layer-14.safetensors", + "model.layers.14.mlp.experts.gate_up_qweight": "model-layer-14.safetensors", + "model.layers.14.mlp.experts.gate_up_scales": "model-layer-14.safetensors", + "model.layers.14.mlp.gate.e_score_correction_bias": "model-layer-14.safetensors", + "model.layers.14.mlp.gate.weight": "model-layer-14.safetensors", + "model.layers.14.mlp.shared_experts.down_proj.weight": "model-layer-14.safetensors", + "model.layers.14.mlp.shared_experts.down_proj.weight_scale": "model-layer-14.safetensors", + "model.layers.14.mlp.shared_experts.gate_proj.weight": "model-layer-14.safetensors", + "model.layers.14.mlp.shared_experts.gate_proj.weight_scale": "model-layer-14.safetensors", + "model.layers.14.mlp.shared_experts.up_proj.weight": "model-layer-14.safetensors", + "model.layers.14.mlp.shared_experts.up_proj.weight_scale": "model-layer-14.safetensors", + "model.layers.14.post_attention_layernorm.weight": "model-layer-14.safetensors", + "model.layers.14.self_attn.kv_a_layernorm.weight": "model-layer-14.safetensors", + "model.layers.14.self_attn.kv_a_proj_with_mqa.weight": "model-layer-14.safetensors", + "model.layers.14.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-14.safetensors", + "model.layers.14.self_attn.kv_b_proj.weight": "model-layer-14.safetensors", + "model.layers.14.self_attn.kv_b_proj.weight_scale": "model-layer-14.safetensors", + "model.layers.14.self_attn.o_proj.weight": "model-layer-14.safetensors", + "model.layers.14.self_attn.o_proj.weight_scale": "model-layer-14.safetensors", + "model.layers.14.self_attn.q_proj.weight": "model-layer-14.safetensors", + "model.layers.14.self_attn.q_proj.weight_scale": "model-layer-14.safetensors", + "model.layers.15.input_layernorm.weight": "model-layer-15.safetensors", + "model.layers.15.mlp.experts.down_qweight": "model-layer-15.safetensors", + "model.layers.15.mlp.experts.down_scales": "model-layer-15.safetensors", + "model.layers.15.mlp.experts.gate_up_qweight": "model-layer-15.safetensors", + "model.layers.15.mlp.experts.gate_up_scales": "model-layer-15.safetensors", + "model.layers.15.mlp.gate.e_score_correction_bias": "model-layer-15.safetensors", + "model.layers.15.mlp.gate.weight": "model-layer-15.safetensors", + "model.layers.15.mlp.shared_experts.down_proj.weight": "model-layer-15.safetensors", + "model.layers.15.mlp.shared_experts.down_proj.weight_scale": "model-layer-15.safetensors", + "model.layers.15.mlp.shared_experts.gate_proj.weight": "model-layer-15.safetensors", + "model.layers.15.mlp.shared_experts.gate_proj.weight_scale": "model-layer-15.safetensors", + "model.layers.15.mlp.shared_experts.up_proj.weight": "model-layer-15.safetensors", + "model.layers.15.mlp.shared_experts.up_proj.weight_scale": "model-layer-15.safetensors", + "model.layers.15.post_attention_layernorm.weight": "model-layer-15.safetensors", + "model.layers.15.self_attn.kv_a_layernorm.weight": "model-layer-15.safetensors", + "model.layers.15.self_attn.kv_a_proj_with_mqa.weight": "model-layer-15.safetensors", + "model.layers.15.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-15.safetensors", + "model.layers.15.self_attn.kv_b_proj.weight": "model-layer-15.safetensors", + "model.layers.15.self_attn.kv_b_proj.weight_scale": "model-layer-15.safetensors", + "model.layers.15.self_attn.o_proj.weight": "model-layer-15.safetensors", + "model.layers.15.self_attn.o_proj.weight_scale": "model-layer-15.safetensors", + "model.layers.15.self_attn.q_proj.weight": "model-layer-15.safetensors", + "model.layers.15.self_attn.q_proj.weight_scale": "model-layer-15.safetensors", + "model.layers.16.input_layernorm.weight": "model-layer-16.safetensors", + "model.layers.16.mlp.experts.down_qweight": "model-layer-16.safetensors", + "model.layers.16.mlp.experts.down_scales": "model-layer-16.safetensors", + "model.layers.16.mlp.experts.gate_up_qweight": "model-layer-16.safetensors", + "model.layers.16.mlp.experts.gate_up_scales": "model-layer-16.safetensors", + "model.layers.16.mlp.gate.e_score_correction_bias": "model-layer-16.safetensors", + "model.layers.16.mlp.gate.weight": "model-layer-16.safetensors", + "model.layers.16.mlp.shared_experts.down_proj.weight": "model-layer-16.safetensors", + "model.layers.16.mlp.shared_experts.down_proj.weight_scale": "model-layer-16.safetensors", + "model.layers.16.mlp.shared_experts.gate_proj.weight": "model-layer-16.safetensors", + "model.layers.16.mlp.shared_experts.gate_proj.weight_scale": "model-layer-16.safetensors", + "model.layers.16.mlp.shared_experts.up_proj.weight": "model-layer-16.safetensors", + "model.layers.16.mlp.shared_experts.up_proj.weight_scale": "model-layer-16.safetensors", + "model.layers.16.post_attention_layernorm.weight": "model-layer-16.safetensors", + "model.layers.16.self_attn.kv_a_layernorm.weight": "model-layer-16.safetensors", + "model.layers.16.self_attn.kv_a_proj_with_mqa.weight": "model-layer-16.safetensors", + "model.layers.16.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-16.safetensors", + "model.layers.16.self_attn.kv_b_proj.weight": "model-layer-16.safetensors", + "model.layers.16.self_attn.kv_b_proj.weight_scale": "model-layer-16.safetensors", + "model.layers.16.self_attn.o_proj.weight": "model-layer-16.safetensors", + "model.layers.16.self_attn.o_proj.weight_scale": "model-layer-16.safetensors", + "model.layers.16.self_attn.q_proj.weight": "model-layer-16.safetensors", + "model.layers.16.self_attn.q_proj.weight_scale": "model-layer-16.safetensors", + "model.layers.17.input_layernorm.weight": "model-layer-17.safetensors", + "model.layers.17.mlp.experts.down_qweight": "model-layer-17.safetensors", + "model.layers.17.mlp.experts.down_scales": "model-layer-17.safetensors", + "model.layers.17.mlp.experts.gate_up_qweight": "model-layer-17.safetensors", + "model.layers.17.mlp.experts.gate_up_scales": "model-layer-17.safetensors", + "model.layers.17.mlp.gate.e_score_correction_bias": "model-layer-17.safetensors", + "model.layers.17.mlp.gate.weight": "model-layer-17.safetensors", + "model.layers.17.mlp.shared_experts.down_proj.weight": "model-layer-17.safetensors", + "model.layers.17.mlp.shared_experts.down_proj.weight_scale": "model-layer-17.safetensors", + "model.layers.17.mlp.shared_experts.gate_proj.weight": "model-layer-17.safetensors", + "model.layers.17.mlp.shared_experts.gate_proj.weight_scale": "model-layer-17.safetensors", + "model.layers.17.mlp.shared_experts.up_proj.weight": "model-layer-17.safetensors", + "model.layers.17.mlp.shared_experts.up_proj.weight_scale": "model-layer-17.safetensors", + "model.layers.17.post_attention_layernorm.weight": "model-layer-17.safetensors", + "model.layers.17.self_attn.kv_a_layernorm.weight": "model-layer-17.safetensors", + "model.layers.17.self_attn.kv_a_proj_with_mqa.weight": "model-layer-17.safetensors", + "model.layers.17.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-17.safetensors", + "model.layers.17.self_attn.kv_b_proj.weight": "model-layer-17.safetensors", + "model.layers.17.self_attn.kv_b_proj.weight_scale": "model-layer-17.safetensors", + "model.layers.17.self_attn.o_proj.weight": "model-layer-17.safetensors", + "model.layers.17.self_attn.o_proj.weight_scale": "model-layer-17.safetensors", + "model.layers.17.self_attn.q_proj.weight": "model-layer-17.safetensors", + "model.layers.17.self_attn.q_proj.weight_scale": "model-layer-17.safetensors", + "model.layers.18.input_layernorm.weight": "model-layer-18.safetensors", + "model.layers.18.mlp.experts.down_qweight": "model-layer-18.safetensors", + "model.layers.18.mlp.experts.down_scales": "model-layer-18.safetensors", + "model.layers.18.mlp.experts.gate_up_qweight": "model-layer-18.safetensors", + "model.layers.18.mlp.experts.gate_up_scales": "model-layer-18.safetensors", + "model.layers.18.mlp.gate.e_score_correction_bias": "model-layer-18.safetensors", + "model.layers.18.mlp.gate.weight": "model-layer-18.safetensors", + "model.layers.18.mlp.shared_experts.down_proj.weight": "model-layer-18.safetensors", + "model.layers.18.mlp.shared_experts.down_proj.weight_scale": "model-layer-18.safetensors", + "model.layers.18.mlp.shared_experts.gate_proj.weight": "model-layer-18.safetensors", + "model.layers.18.mlp.shared_experts.gate_proj.weight_scale": "model-layer-18.safetensors", + "model.layers.18.mlp.shared_experts.up_proj.weight": "model-layer-18.safetensors", + "model.layers.18.mlp.shared_experts.up_proj.weight_scale": "model-layer-18.safetensors", + "model.layers.18.post_attention_layernorm.weight": "model-layer-18.safetensors", + "model.layers.18.self_attn.kv_a_layernorm.weight": "model-layer-18.safetensors", + "model.layers.18.self_attn.kv_a_proj_with_mqa.weight": "model-layer-18.safetensors", + "model.layers.18.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-18.safetensors", + "model.layers.18.self_attn.kv_b_proj.weight": "model-layer-18.safetensors", + "model.layers.18.self_attn.kv_b_proj.weight_scale": "model-layer-18.safetensors", + "model.layers.18.self_attn.o_proj.weight": "model-layer-18.safetensors", + "model.layers.18.self_attn.o_proj.weight_scale": "model-layer-18.safetensors", + "model.layers.18.self_attn.q_proj.weight": "model-layer-18.safetensors", + "model.layers.18.self_attn.q_proj.weight_scale": "model-layer-18.safetensors", + "model.layers.19.input_layernorm.weight": "model-layer-19.safetensors", + "model.layers.19.mlp.experts.down_qweight": "model-layer-19.safetensors", + "model.layers.19.mlp.experts.down_scales": "model-layer-19.safetensors", + "model.layers.19.mlp.experts.gate_up_qweight": "model-layer-19.safetensors", + "model.layers.19.mlp.experts.gate_up_scales": "model-layer-19.safetensors", + "model.layers.19.mlp.gate.e_score_correction_bias": "model-layer-19.safetensors", + "model.layers.19.mlp.gate.weight": "model-layer-19.safetensors", + "model.layers.19.mlp.shared_experts.down_proj.weight": "model-layer-19.safetensors", + "model.layers.19.mlp.shared_experts.down_proj.weight_scale": "model-layer-19.safetensors", + "model.layers.19.mlp.shared_experts.gate_proj.weight": "model-layer-19.safetensors", + "model.layers.19.mlp.shared_experts.gate_proj.weight_scale": "model-layer-19.safetensors", + "model.layers.19.mlp.shared_experts.up_proj.weight": "model-layer-19.safetensors", + "model.layers.19.mlp.shared_experts.up_proj.weight_scale": "model-layer-19.safetensors", + "model.layers.19.post_attention_layernorm.weight": "model-layer-19.safetensors", + "model.layers.19.self_attn.kv_a_layernorm.weight": "model-layer-19.safetensors", + "model.layers.19.self_attn.kv_a_proj_with_mqa.weight": "model-layer-19.safetensors", + "model.layers.19.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-19.safetensors", + "model.layers.19.self_attn.kv_b_proj.weight": "model-layer-19.safetensors", + "model.layers.19.self_attn.kv_b_proj.weight_scale": "model-layer-19.safetensors", + "model.layers.19.self_attn.o_proj.weight": "model-layer-19.safetensors", + "model.layers.19.self_attn.o_proj.weight_scale": "model-layer-19.safetensors", + "model.layers.19.self_attn.q_proj.weight": "model-layer-19.safetensors", + "model.layers.19.self_attn.q_proj.weight_scale": "model-layer-19.safetensors", + "model.layers.20.input_layernorm.weight": "model-layer-20.safetensors", + "model.layers.20.mlp.experts.down_qweight": "model-layer-20.safetensors", + "model.layers.20.mlp.experts.down_scales": "model-layer-20.safetensors", + "model.layers.20.mlp.experts.gate_up_qweight": "model-layer-20.safetensors", + "model.layers.20.mlp.experts.gate_up_scales": "model-layer-20.safetensors", + "model.layers.20.mlp.gate.e_score_correction_bias": "model-layer-20.safetensors", + "model.layers.20.mlp.gate.weight": "model-layer-20.safetensors", + "model.layers.20.mlp.shared_experts.down_proj.weight": "model-layer-20.safetensors", + "model.layers.20.mlp.shared_experts.down_proj.weight_scale": "model-layer-20.safetensors", + "model.layers.20.mlp.shared_experts.gate_proj.weight": "model-layer-20.safetensors", + "model.layers.20.mlp.shared_experts.gate_proj.weight_scale": "model-layer-20.safetensors", + "model.layers.20.mlp.shared_experts.up_proj.weight": "model-layer-20.safetensors", + "model.layers.20.mlp.shared_experts.up_proj.weight_scale": "model-layer-20.safetensors", + "model.layers.20.post_attention_layernorm.weight": "model-layer-20.safetensors", + "model.layers.20.self_attn.kv_a_layernorm.weight": "model-layer-20.safetensors", + "model.layers.20.self_attn.kv_a_proj_with_mqa.weight": "model-layer-20.safetensors", + "model.layers.20.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-20.safetensors", + "model.layers.20.self_attn.kv_b_proj.weight": "model-layer-20.safetensors", + "model.layers.20.self_attn.kv_b_proj.weight_scale": "model-layer-20.safetensors", + "model.layers.20.self_attn.o_proj.weight": "model-layer-20.safetensors", + "model.layers.20.self_attn.o_proj.weight_scale": "model-layer-20.safetensors", + "model.layers.20.self_attn.q_proj.weight": "model-layer-20.safetensors", + "model.layers.20.self_attn.q_proj.weight_scale": "model-layer-20.safetensors", + "model.layers.21.input_layernorm.weight": "model-layer-21.safetensors", + "model.layers.21.mlp.experts.down_qweight": "model-layer-21.safetensors", + "model.layers.21.mlp.experts.down_scales": "model-layer-21.safetensors", + "model.layers.21.mlp.experts.gate_up_qweight": "model-layer-21.safetensors", + "model.layers.21.mlp.experts.gate_up_scales": "model-layer-21.safetensors", + "model.layers.21.mlp.gate.e_score_correction_bias": "model-layer-21.safetensors", + "model.layers.21.mlp.gate.weight": "model-layer-21.safetensors", + "model.layers.21.mlp.shared_experts.down_proj.weight": "model-layer-21.safetensors", + "model.layers.21.mlp.shared_experts.down_proj.weight_scale": "model-layer-21.safetensors", + "model.layers.21.mlp.shared_experts.gate_proj.weight": "model-layer-21.safetensors", + "model.layers.21.mlp.shared_experts.gate_proj.weight_scale": "model-layer-21.safetensors", + "model.layers.21.mlp.shared_experts.up_proj.weight": "model-layer-21.safetensors", + "model.layers.21.mlp.shared_experts.up_proj.weight_scale": "model-layer-21.safetensors", + "model.layers.21.post_attention_layernorm.weight": "model-layer-21.safetensors", + "model.layers.21.self_attn.kv_a_layernorm.weight": "model-layer-21.safetensors", + "model.layers.21.self_attn.kv_a_proj_with_mqa.weight": "model-layer-21.safetensors", + "model.layers.21.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-21.safetensors", + "model.layers.21.self_attn.kv_b_proj.weight": "model-layer-21.safetensors", + "model.layers.21.self_attn.kv_b_proj.weight_scale": "model-layer-21.safetensors", + "model.layers.21.self_attn.o_proj.weight": "model-layer-21.safetensors", + "model.layers.21.self_attn.o_proj.weight_scale": "model-layer-21.safetensors", + "model.layers.21.self_attn.q_proj.weight": "model-layer-21.safetensors", + "model.layers.21.self_attn.q_proj.weight_scale": "model-layer-21.safetensors", + "model.layers.22.input_layernorm.weight": "model-layer-22.safetensors", + "model.layers.22.mlp.experts.down_qweight": "model-layer-22.safetensors", + "model.layers.22.mlp.experts.down_scales": "model-layer-22.safetensors", + "model.layers.22.mlp.experts.gate_up_qweight": "model-layer-22.safetensors", + "model.layers.22.mlp.experts.gate_up_scales": "model-layer-22.safetensors", + "model.layers.22.mlp.gate.e_score_correction_bias": "model-layer-22.safetensors", + "model.layers.22.mlp.gate.weight": "model-layer-22.safetensors", + "model.layers.22.mlp.shared_experts.down_proj.weight": "model-layer-22.safetensors", + "model.layers.22.mlp.shared_experts.down_proj.weight_scale": "model-layer-22.safetensors", + "model.layers.22.mlp.shared_experts.gate_proj.weight": "model-layer-22.safetensors", + "model.layers.22.mlp.shared_experts.gate_proj.weight_scale": "model-layer-22.safetensors", + "model.layers.22.mlp.shared_experts.up_proj.weight": "model-layer-22.safetensors", + "model.layers.22.mlp.shared_experts.up_proj.weight_scale": "model-layer-22.safetensors", + "model.layers.22.post_attention_layernorm.weight": "model-layer-22.safetensors", + "model.layers.22.self_attn.kv_a_layernorm.weight": "model-layer-22.safetensors", + "model.layers.22.self_attn.kv_a_proj_with_mqa.weight": "model-layer-22.safetensors", + "model.layers.22.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-22.safetensors", + "model.layers.22.self_attn.kv_b_proj.weight": "model-layer-22.safetensors", + "model.layers.22.self_attn.kv_b_proj.weight_scale": "model-layer-22.safetensors", + "model.layers.22.self_attn.o_proj.weight": "model-layer-22.safetensors", + "model.layers.22.self_attn.o_proj.weight_scale": "model-layer-22.safetensors", + "model.layers.22.self_attn.q_proj.weight": "model-layer-22.safetensors", + "model.layers.22.self_attn.q_proj.weight_scale": "model-layer-22.safetensors", + "model.layers.23.input_layernorm.weight": "model-layer-23.safetensors", + "model.layers.23.mlp.experts.down_qweight": "model-layer-23.safetensors", + "model.layers.23.mlp.experts.down_scales": "model-layer-23.safetensors", + "model.layers.23.mlp.experts.gate_up_qweight": "model-layer-23.safetensors", + "model.layers.23.mlp.experts.gate_up_scales": "model-layer-23.safetensors", + "model.layers.23.mlp.gate.e_score_correction_bias": "model-layer-23.safetensors", + "model.layers.23.mlp.gate.weight": "model-layer-23.safetensors", + "model.layers.23.mlp.shared_experts.down_proj.weight": "model-layer-23.safetensors", + "model.layers.23.mlp.shared_experts.down_proj.weight_scale": "model-layer-23.safetensors", + "model.layers.23.mlp.shared_experts.gate_proj.weight": "model-layer-23.safetensors", + "model.layers.23.mlp.shared_experts.gate_proj.weight_scale": "model-layer-23.safetensors", + "model.layers.23.mlp.shared_experts.up_proj.weight": "model-layer-23.safetensors", + "model.layers.23.mlp.shared_experts.up_proj.weight_scale": "model-layer-23.safetensors", + "model.layers.23.post_attention_layernorm.weight": "model-layer-23.safetensors", + "model.layers.23.self_attn.kv_a_layernorm.weight": "model-layer-23.safetensors", + "model.layers.23.self_attn.kv_a_proj_with_mqa.weight": "model-layer-23.safetensors", + "model.layers.23.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-23.safetensors", + "model.layers.23.self_attn.kv_b_proj.weight": "model-layer-23.safetensors", + "model.layers.23.self_attn.kv_b_proj.weight_scale": "model-layer-23.safetensors", + "model.layers.23.self_attn.o_proj.weight": "model-layer-23.safetensors", + "model.layers.23.self_attn.o_proj.weight_scale": "model-layer-23.safetensors", + "model.layers.23.self_attn.q_proj.weight": "model-layer-23.safetensors", + "model.layers.23.self_attn.q_proj.weight_scale": "model-layer-23.safetensors", + "model.layers.24.input_layernorm.weight": "model-layer-24.safetensors", + "model.layers.24.mlp.experts.down_qweight": "model-layer-24.safetensors", + "model.layers.24.mlp.experts.down_scales": "model-layer-24.safetensors", + "model.layers.24.mlp.experts.gate_up_qweight": "model-layer-24.safetensors", + "model.layers.24.mlp.experts.gate_up_scales": "model-layer-24.safetensors", + "model.layers.24.mlp.gate.e_score_correction_bias": "model-layer-24.safetensors", + "model.layers.24.mlp.gate.weight": "model-layer-24.safetensors", + "model.layers.24.mlp.shared_experts.down_proj.weight": "model-layer-24.safetensors", + "model.layers.24.mlp.shared_experts.down_proj.weight_scale": "model-layer-24.safetensors", + "model.layers.24.mlp.shared_experts.gate_proj.weight": "model-layer-24.safetensors", + "model.layers.24.mlp.shared_experts.gate_proj.weight_scale": "model-layer-24.safetensors", + "model.layers.24.mlp.shared_experts.up_proj.weight": "model-layer-24.safetensors", + "model.layers.24.mlp.shared_experts.up_proj.weight_scale": "model-layer-24.safetensors", + "model.layers.24.post_attention_layernorm.weight": "model-layer-24.safetensors", + "model.layers.24.self_attn.kv_a_layernorm.weight": "model-layer-24.safetensors", + "model.layers.24.self_attn.kv_a_proj_with_mqa.weight": "model-layer-24.safetensors", + "model.layers.24.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-24.safetensors", + "model.layers.24.self_attn.kv_b_proj.weight": "model-layer-24.safetensors", + "model.layers.24.self_attn.kv_b_proj.weight_scale": "model-layer-24.safetensors", + "model.layers.24.self_attn.o_proj.weight": "model-layer-24.safetensors", + "model.layers.24.self_attn.o_proj.weight_scale": "model-layer-24.safetensors", + "model.layers.24.self_attn.q_proj.weight": "model-layer-24.safetensors", + "model.layers.24.self_attn.q_proj.weight_scale": "model-layer-24.safetensors", + "model.layers.25.input_layernorm.weight": "model-layer-25.safetensors", + "model.layers.25.mlp.experts.down_qweight": "model-layer-25.safetensors", + "model.layers.25.mlp.experts.down_scales": "model-layer-25.safetensors", + "model.layers.25.mlp.experts.gate_up_qweight": "model-layer-25.safetensors", + "model.layers.25.mlp.experts.gate_up_scales": "model-layer-25.safetensors", + "model.layers.25.mlp.gate.e_score_correction_bias": "model-layer-25.safetensors", + "model.layers.25.mlp.gate.weight": "model-layer-25.safetensors", + "model.layers.25.mlp.shared_experts.down_proj.weight": "model-layer-25.safetensors", + "model.layers.25.mlp.shared_experts.down_proj.weight_scale": "model-layer-25.safetensors", + "model.layers.25.mlp.shared_experts.gate_proj.weight": "model-layer-25.safetensors", + "model.layers.25.mlp.shared_experts.gate_proj.weight_scale": "model-layer-25.safetensors", + "model.layers.25.mlp.shared_experts.up_proj.weight": "model-layer-25.safetensors", + "model.layers.25.mlp.shared_experts.up_proj.weight_scale": "model-layer-25.safetensors", + "model.layers.25.post_attention_layernorm.weight": "model-layer-25.safetensors", + "model.layers.25.self_attn.kv_a_layernorm.weight": "model-layer-25.safetensors", + "model.layers.25.self_attn.kv_a_proj_with_mqa.weight": "model-layer-25.safetensors", + "model.layers.25.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-25.safetensors", + "model.layers.25.self_attn.kv_b_proj.weight": "model-layer-25.safetensors", + "model.layers.25.self_attn.kv_b_proj.weight_scale": "model-layer-25.safetensors", + "model.layers.25.self_attn.o_proj.weight": "model-layer-25.safetensors", + "model.layers.25.self_attn.o_proj.weight_scale": "model-layer-25.safetensors", + "model.layers.25.self_attn.q_proj.weight": "model-layer-25.safetensors", + "model.layers.25.self_attn.q_proj.weight_scale": "model-layer-25.safetensors", + "model.layers.26.input_layernorm.weight": "model-layer-26.safetensors", + "model.layers.26.mlp.experts.down_qweight": "model-layer-26.safetensors", + "model.layers.26.mlp.experts.down_scales": "model-layer-26.safetensors", + "model.layers.26.mlp.experts.gate_up_qweight": "model-layer-26.safetensors", + "model.layers.26.mlp.experts.gate_up_scales": "model-layer-26.safetensors", + "model.layers.26.mlp.gate.e_score_correction_bias": "model-layer-26.safetensors", + "model.layers.26.mlp.gate.weight": "model-layer-26.safetensors", + "model.layers.26.mlp.shared_experts.down_proj.weight": "model-layer-26.safetensors", + "model.layers.26.mlp.shared_experts.down_proj.weight_scale": "model-layer-26.safetensors", + "model.layers.26.mlp.shared_experts.gate_proj.weight": "model-layer-26.safetensors", + "model.layers.26.mlp.shared_experts.gate_proj.weight_scale": "model-layer-26.safetensors", + "model.layers.26.mlp.shared_experts.up_proj.weight": "model-layer-26.safetensors", + "model.layers.26.mlp.shared_experts.up_proj.weight_scale": "model-layer-26.safetensors", + "model.layers.26.post_attention_layernorm.weight": "model-layer-26.safetensors", + "model.layers.26.self_attn.kv_a_layernorm.weight": "model-layer-26.safetensors", + "model.layers.26.self_attn.kv_a_proj_with_mqa.weight": "model-layer-26.safetensors", + "model.layers.26.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-26.safetensors", + "model.layers.26.self_attn.kv_b_proj.weight": "model-layer-26.safetensors", + "model.layers.26.self_attn.kv_b_proj.weight_scale": "model-layer-26.safetensors", + "model.layers.26.self_attn.o_proj.weight": "model-layer-26.safetensors", + "model.layers.26.self_attn.o_proj.weight_scale": "model-layer-26.safetensors", + "model.layers.26.self_attn.q_proj.weight": "model-layer-26.safetensors", + "model.layers.26.self_attn.q_proj.weight_scale": "model-layer-26.safetensors", + "model.layers.27.input_layernorm.weight": "model-layer-27.safetensors", + "model.layers.27.mlp.experts.down_qweight": "model-layer-27.safetensors", + "model.layers.27.mlp.experts.down_scales": "model-layer-27.safetensors", + "model.layers.27.mlp.experts.gate_up_qweight": "model-layer-27.safetensors", + "model.layers.27.mlp.experts.gate_up_scales": "model-layer-27.safetensors", + "model.layers.27.mlp.gate.e_score_correction_bias": "model-layer-27.safetensors", + "model.layers.27.mlp.gate.weight": "model-layer-27.safetensors", + "model.layers.27.mlp.shared_experts.down_proj.weight": "model-layer-27.safetensors", + "model.layers.27.mlp.shared_experts.down_proj.weight_scale": "model-layer-27.safetensors", + "model.layers.27.mlp.shared_experts.gate_proj.weight": "model-layer-27.safetensors", + "model.layers.27.mlp.shared_experts.gate_proj.weight_scale": "model-layer-27.safetensors", + "model.layers.27.mlp.shared_experts.up_proj.weight": "model-layer-27.safetensors", + "model.layers.27.mlp.shared_experts.up_proj.weight_scale": "model-layer-27.safetensors", + "model.layers.27.post_attention_layernorm.weight": "model-layer-27.safetensors", + "model.layers.27.self_attn.kv_a_layernorm.weight": "model-layer-27.safetensors", + "model.layers.27.self_attn.kv_a_proj_with_mqa.weight": "model-layer-27.safetensors", + "model.layers.27.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-27.safetensors", + "model.layers.27.self_attn.kv_b_proj.weight": "model-layer-27.safetensors", + "model.layers.27.self_attn.kv_b_proj.weight_scale": "model-layer-27.safetensors", + "model.layers.27.self_attn.o_proj.weight": "model-layer-27.safetensors", + "model.layers.27.self_attn.o_proj.weight_scale": "model-layer-27.safetensors", + "model.layers.27.self_attn.q_proj.weight": "model-layer-27.safetensors", + "model.layers.27.self_attn.q_proj.weight_scale": "model-layer-27.safetensors", + "model.layers.28.input_layernorm.weight": "model-layer-28.safetensors", + "model.layers.28.mlp.experts.down_qweight": "model-layer-28.safetensors", + "model.layers.28.mlp.experts.down_scales": "model-layer-28.safetensors", + "model.layers.28.mlp.experts.gate_up_qweight": "model-layer-28.safetensors", + "model.layers.28.mlp.experts.gate_up_scales": "model-layer-28.safetensors", + "model.layers.28.mlp.gate.e_score_correction_bias": "model-layer-28.safetensors", + "model.layers.28.mlp.gate.weight": "model-layer-28.safetensors", + "model.layers.28.mlp.shared_experts.down_proj.weight": "model-layer-28.safetensors", + "model.layers.28.mlp.shared_experts.down_proj.weight_scale": "model-layer-28.safetensors", + "model.layers.28.mlp.shared_experts.gate_proj.weight": "model-layer-28.safetensors", + "model.layers.28.mlp.shared_experts.gate_proj.weight_scale": "model-layer-28.safetensors", + "model.layers.28.mlp.shared_experts.up_proj.weight": "model-layer-28.safetensors", + "model.layers.28.mlp.shared_experts.up_proj.weight_scale": "model-layer-28.safetensors", + "model.layers.28.post_attention_layernorm.weight": "model-layer-28.safetensors", + "model.layers.28.self_attn.kv_a_layernorm.weight": "model-layer-28.safetensors", + "model.layers.28.self_attn.kv_a_proj_with_mqa.weight": "model-layer-28.safetensors", + "model.layers.28.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-28.safetensors", + "model.layers.28.self_attn.kv_b_proj.weight": "model-layer-28.safetensors", + "model.layers.28.self_attn.kv_b_proj.weight_scale": "model-layer-28.safetensors", + "model.layers.28.self_attn.o_proj.weight": "model-layer-28.safetensors", + "model.layers.28.self_attn.o_proj.weight_scale": "model-layer-28.safetensors", + "model.layers.28.self_attn.q_proj.weight": "model-layer-28.safetensors", + "model.layers.28.self_attn.q_proj.weight_scale": "model-layer-28.safetensors", + "model.layers.29.input_layernorm.weight": "model-layer-29.safetensors", + "model.layers.29.mlp.experts.down_qweight": "model-layer-29.safetensors", + "model.layers.29.mlp.experts.down_scales": "model-layer-29.safetensors", + "model.layers.29.mlp.experts.gate_up_qweight": "model-layer-29.safetensors", + "model.layers.29.mlp.experts.gate_up_scales": "model-layer-29.safetensors", + "model.layers.29.mlp.gate.e_score_correction_bias": "model-layer-29.safetensors", + "model.layers.29.mlp.gate.weight": "model-layer-29.safetensors", + "model.layers.29.mlp.shared_experts.down_proj.weight": "model-layer-29.safetensors", + "model.layers.29.mlp.shared_experts.down_proj.weight_scale": "model-layer-29.safetensors", + "model.layers.29.mlp.shared_experts.gate_proj.weight": "model-layer-29.safetensors", + "model.layers.29.mlp.shared_experts.gate_proj.weight_scale": "model-layer-29.safetensors", + "model.layers.29.mlp.shared_experts.up_proj.weight": "model-layer-29.safetensors", + "model.layers.29.mlp.shared_experts.up_proj.weight_scale": "model-layer-29.safetensors", + "model.layers.29.post_attention_layernorm.weight": "model-layer-29.safetensors", + "model.layers.29.self_attn.kv_a_layernorm.weight": "model-layer-29.safetensors", + "model.layers.29.self_attn.kv_a_proj_with_mqa.weight": "model-layer-29.safetensors", + "model.layers.29.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-29.safetensors", + "model.layers.29.self_attn.kv_b_proj.weight": "model-layer-29.safetensors", + "model.layers.29.self_attn.kv_b_proj.weight_scale": "model-layer-29.safetensors", + "model.layers.29.self_attn.o_proj.weight": "model-layer-29.safetensors", + "model.layers.29.self_attn.o_proj.weight_scale": "model-layer-29.safetensors", + "model.layers.29.self_attn.q_proj.weight": "model-layer-29.safetensors", + "model.layers.29.self_attn.q_proj.weight_scale": "model-layer-29.safetensors", + "model.layers.30.input_layernorm.weight": "model-layer-30.safetensors", + "model.layers.30.mlp.experts.down_qweight": "model-layer-30.safetensors", + "model.layers.30.mlp.experts.down_scales": "model-layer-30.safetensors", + "model.layers.30.mlp.experts.gate_up_qweight": "model-layer-30.safetensors", + "model.layers.30.mlp.experts.gate_up_scales": "model-layer-30.safetensors", + "model.layers.30.mlp.gate.e_score_correction_bias": "model-layer-30.safetensors", + "model.layers.30.mlp.gate.weight": "model-layer-30.safetensors", + "model.layers.30.mlp.shared_experts.down_proj.weight": "model-layer-30.safetensors", + "model.layers.30.mlp.shared_experts.down_proj.weight_scale": "model-layer-30.safetensors", + "model.layers.30.mlp.shared_experts.gate_proj.weight": "model-layer-30.safetensors", + "model.layers.30.mlp.shared_experts.gate_proj.weight_scale": "model-layer-30.safetensors", + "model.layers.30.mlp.shared_experts.up_proj.weight": "model-layer-30.safetensors", + "model.layers.30.mlp.shared_experts.up_proj.weight_scale": "model-layer-30.safetensors", + "model.layers.30.post_attention_layernorm.weight": "model-layer-30.safetensors", + "model.layers.30.self_attn.kv_a_layernorm.weight": "model-layer-30.safetensors", + "model.layers.30.self_attn.kv_a_proj_with_mqa.weight": "model-layer-30.safetensors", + "model.layers.30.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-30.safetensors", + "model.layers.30.self_attn.kv_b_proj.weight": "model-layer-30.safetensors", + "model.layers.30.self_attn.kv_b_proj.weight_scale": "model-layer-30.safetensors", + "model.layers.30.self_attn.o_proj.weight": "model-layer-30.safetensors", + "model.layers.30.self_attn.o_proj.weight_scale": "model-layer-30.safetensors", + "model.layers.30.self_attn.q_proj.weight": "model-layer-30.safetensors", + "model.layers.30.self_attn.q_proj.weight_scale": "model-layer-30.safetensors", + "model.layers.31.input_layernorm.weight": "model-layer-31.safetensors", + "model.layers.31.mlp.experts.down_qweight": "model-layer-31.safetensors", + "model.layers.31.mlp.experts.down_scales": "model-layer-31.safetensors", + "model.layers.31.mlp.experts.gate_up_qweight": "model-layer-31.safetensors", + "model.layers.31.mlp.experts.gate_up_scales": "model-layer-31.safetensors", + "model.layers.31.mlp.gate.e_score_correction_bias": "model-layer-31.safetensors", + "model.layers.31.mlp.gate.weight": "model-layer-31.safetensors", + "model.layers.31.mlp.shared_experts.down_proj.weight": "model-layer-31.safetensors", + "model.layers.31.mlp.shared_experts.down_proj.weight_scale": "model-layer-31.safetensors", + "model.layers.31.mlp.shared_experts.gate_proj.weight": "model-layer-31.safetensors", + "model.layers.31.mlp.shared_experts.gate_proj.weight_scale": "model-layer-31.safetensors", + "model.layers.31.mlp.shared_experts.up_proj.weight": "model-layer-31.safetensors", + "model.layers.31.mlp.shared_experts.up_proj.weight_scale": "model-layer-31.safetensors", + "model.layers.31.post_attention_layernorm.weight": "model-layer-31.safetensors", + "model.layers.31.self_attn.kv_a_layernorm.weight": "model-layer-31.safetensors", + "model.layers.31.self_attn.kv_a_proj_with_mqa.weight": "model-layer-31.safetensors", + "model.layers.31.self_attn.kv_a_proj_with_mqa.weight_scale": "model-layer-31.safetensors", + "model.layers.31.self_attn.kv_b_proj.weight": "model-layer-31.safetensors", + "model.layers.31.self_attn.kv_b_proj.weight_scale": "model-layer-31.safetensors", + "model.layers.31.self_attn.o_proj.weight": "model-layer-31.safetensors", + "model.layers.31.self_attn.o_proj.weight_scale": "model-layer-31.safetensors", + "model.layers.31.self_attn.q_proj.weight": "model-layer-31.safetensors", + "model.layers.31.self_attn.q_proj.weight_scale": "model-layer-31.safetensors" + } +} \ No newline at end of file diff --git a/modeling_sarvam_moe.py b/modeling_sarvam_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..9b536883865545c120abccf52b89ec4bb6af4535 --- /dev/null +++ b/modeling_sarvam_moe.py @@ -0,0 +1,1002 @@ +# Copyright 2026 Sarvam AI team. All rights reserved. +# +# This code is based on Llama and Deepseek MoE implementations +# in this library. It has been modified from its original forms to +# accommodate Sarvam's MLA (multi-latent attention) MoE architecture. +# +# 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. + +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, + _prepare_4d_attention_mask, + _prepare_4d_causal_attention_mask, +) +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS +from transformers.pytorch_utils import ( + ALL_LAYERNORM_LAYERS, + is_torch_greater_or_equal_than_1_13, +) +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from transformers.utils.import_utils import is_torch_fx_available + +import torch.distributed as dist +import numpy as np + +from .configuration_sarvam_moe import SarvamMLAConfig + +if is_torch_fx_available(): + if not is_torch_greater_or_equal_than_1_13: + import torch.fx + + _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "SarvamMLAConfig" + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def _get_usable_past_kv_length(cache: Cache, new_seq_length: int, layer_idx: int = 0) -> int: + previous_length = cache.get_seq_length(layer_idx) + # Dynamic layers return -1, static layers return an int + max_length = cache.get_max_cache_shape(layer_idx) + if max_length is not None and max_length != -1 and previous_length + new_seq_length > max_length: + return max_length - new_seq_length + return previous_length + + +class SarvamMLARMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + SarvamMLARMSNorm 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) + + +ALL_LAYERNORM_LAYERS.append(SarvamMLARMSNorm) + + +class SarvamMLARotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + self._set_cos_sin_cache( + seq_len=max_position_embeddings, + device=self.inv_freq.device, + dtype=torch.get_default_dtype(), + ) + self.max_seq_len_cached = None + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + + freqs = torch.outer(t, self.inv_freq.to(t.device)) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + def forward(self, x, seq_len=None): + if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:seq_len].to(dtype=x.dtype), + self.sin_cached[:seq_len].to(dtype=x.dtype), + ) + + +def yarn_find_correction_dim(num_rotations, dim, base=10000, max_position_embeddings=2048): + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + +def yarn_find_correction_range(low_rot, high_rot, dim, base=10000, max_position_embeddings=2048): + low = math.floor(yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)) + high = math.ceil(yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)) + return max(low, 0), min(high, dim - 1) + + +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def yarn_linear_ramp_mask(min_val, max_val, dim): + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + +class SarvamMLAYarnRotaryEmbedding(SarvamMLARotaryEmbedding): + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=40.0, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1.0, + mscale_all_dim=1.0, + ): + self.scaling_factor = float(scaling_factor) + self.original_max_position_embeddings = int(original_max_position_embeddings) + self.beta_fast = float(beta_fast) + self.beta_slow = float(beta_slow) + self.mscale = float(mscale) + self.mscale_all_dim = float(mscale_all_dim) + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + dim = self.dim + + freq_extra = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + freq_inter = 1.0 / ( + self.scaling_factor * self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + + low, high = yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + + inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(device=device, dtype=torch.float32) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + + _mscale = float( + yarn_get_mscale(self.scaling_factor, self.mscale) + / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", (emb.cos() * _mscale).to(dtype), persistent=False) + self.register_buffer("sin_cached", (emb.sin() * _mscale).to(dtype), persistent=False) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): + cos = cos[position_ids].unsqueeze(unsqueeze_dim) + sin = sin[position_ids].unsqueeze(unsqueeze_dim) + + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class SarvamMLAMLP(nn.Module): + def __init__(self, config, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else 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): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class MoEGate(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.scoring_func = "sigmoid" + self.topk_method = "noaux_tc" + if ( + getattr(config, "n_group", None) is not None + or getattr(config, "topk_group", None) is not None + ): + raise NotImplementedError( + "Grouped routing is not supported; n_group and topk_group must both be None" + ) + self.n_group = None + self.topk_group = None + self.use_grouped_topk = False + if not 1 <= self.top_k <= self.n_routed_experts: + raise ValueError("num_experts_per_tok must be between 1 and num_experts") + + self.norm_topk_prob = True + self.gating_dim = config.hidden_size + self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) + if self.topk_method == "noaux_tc": + # This bias only changes discrete expert selection. It is persistent + # checkpoint state, but does not have a backpropagation gradient. + self.register_buffer( + "e_score_correction_bias", torch.empty((self.n_routed_experts)), persistent=True + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if hasattr(self, "e_score_correction_bias"): + init.zeros_(self.e_score_correction_bias) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + hidden_states = hidden_states.view(-1, h) + logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None) + if self.scoring_func == "sigmoid": + scores = logits.sigmoid() + else: + raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}") + + if self.topk_method == "noaux_tc": + scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0) + _, topk_idx = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) + topk_weight = scores.gather(1, topk_idx) + else: + raise NotImplementedError(f"insupportable TopK function for MoE gating: {self.topk_method}") + + ### norm gate to sum 1 + if self.top_k > 1 and self.norm_topk_prob: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor + + return topk_idx, topk_weight + + +class SarvamMLAMoE(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + + if hasattr(config, "ep_size") and config.ep_size > 1: + assert config.ep_size == dist.get_world_size() + self.ep_size = config.ep_size + self.experts_per_rank = config.num_experts // config.ep_size + self.ep_rank = dist.get_rank() + self.experts = nn.ModuleList( + [ + ( + SarvamMLAMLP(config, intermediate_size=config.moe_intermediate_size) + if i >= self.ep_rank * self.experts_per_rank and i < (self.ep_rank + 1) * self.experts_per_rank + else None + ) + for i in range(config.num_experts) + ] + ) + else: + self.ep_size = 1 + self.experts_per_rank = config.num_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + SarvamMLAMLP(config, intermediate_size=config.moe_intermediate_size) + for i in range(config.num_experts) + ] + ) + self.gate = MoEGate(config) + if ( + hasattr(config, "num_shared_experts") + and config.num_shared_experts is not None + and config.num_shared_experts > 0 + ): + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = SarvamMLAMLP(config=config, intermediate_size=intermediate_size) + else: + self.shared_experts = None + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + if not self.training: + y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape) + else: + y = self.moe_train(hidden_states, topk_idx, topk_weight).view(*orig_shape) + if self.shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + def moe_train(self, x, topk_ids, topk_weight): + if self.ep_size > 1: + raise NotImplementedError( + "Expert-parallel MoE training requires autograd-aware all-to-all collectives; " + "the current distributed inference dispatch is not differentiable" + ) + return self._moe_dispatch(x, topk_ids, topk_weight) + + @torch.no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + return self._moe_dispatch(x, topk_ids, topk_weight) + + def _moe_dispatch(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + sorted_tokens_shape = sorted_tokens.shape + if self.ep_size > 1: + tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1) + tokens_per_expert_group = tokens_per_expert.new_empty(tokens_per_expert.shape[0]) + dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert) + output_splits = tokens_per_expert_group.view(self.ep_size, -1).sum(1).cpu().numpy().tolist() + gathered_tokens = sorted_tokens.new_empty( + tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1] + ) + input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist() + dist.all_to_all( + list(gathered_tokens.split(output_splits)), + list(sorted_tokens.split(input_split_sizes)), + ) + tokens_per_expert_post_gather = tokens_per_expert_group.view(self.ep_size, self.experts_per_rank).sum(dim=0) + gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) + s = 0 + for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): + gatherd_idxs[s : s + k] = i % self.experts_per_rank + s += k + gatherd_idxs = gatherd_idxs.argsort() + sorted_tokens = gathered_tokens[gatherd_idxs] + tokens_per_expert = tokens_per_expert_post_gather + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + if expert is None: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + if self.ep_size > 1: + new_x = torch.empty_like(outs) + new_x[gatherd_idxs] = outs + gathered_tokens = new_x.new_empty(*sorted_tokens_shape) + dist.all_to_all( + list(gathered_tokens.split(input_split_sizes)), + list(new_x.split(output_splits)), + ) + outs = gathered_tokens + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + 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) + + +class SarvamMLAAttention(nn.Module): + is_causal = True + def __init__(self, config: SarvamMLAConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.q_lora_rank = getattr(config, "q_lora_rank", None) + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = config.q_head_dim + + if self.q_lora_rank is None: + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.q_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=getattr(config, "attention_bias", False) + ) + self.q_a_layernorm = SarvamMLARMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=getattr(config, "attention_bias", False), + ) + self.kv_a_layernorm = SarvamMLARMSNorm(config.kv_lora_rank) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=getattr(config, "attention_bias", False), + ) + self._init_rope() + + self.softmax_scale = self.q_head_dim ** (-0.5) + if self.config.rope_scaling is not None: + mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + + def _init_rope(self): + rope_scaling = getattr(self.config, "rope_scaling", None) + if rope_scaling is None or rope_scaling.get("type", None) in (None, "default"): + self.rotary_emb = SarvamMLARotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + return + + rope_type = rope_scaling.get("type") + if rope_type == "deepseek_yarn": + self.rotary_emb = SarvamMLAYarnRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + scaling_factor=rope_scaling.get("factor", 40.0), + original_max_position_embeddings=rope_scaling.get("original_max_position_embeddings", 4096), + beta_fast=rope_scaling.get("beta_fast", 32), + beta_slow=rope_scaling.get("beta_slow", 1), + mscale=rope_scaling.get("mscale", 1.0), + mscale_all_dim=rope_scaling.get("mscale_all_dim", 1.0), + ) + return + raise ValueError(f"Unknown rope_scaling type: {rope_type}") + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim).transpose(1, 2).contiguous() + + 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, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed in a previous version. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += _get_usable_past_kv_length(past_key_value, kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + assert attention_mask is not None + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is" + f" {attn_output.size()}" + ) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class SarvamMLADecoderLayer(nn.Module): + def __init__(self, config: SarvamMLAConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = SarvamMLAAttention(config=config, layer_idx=layer_idx) + + use_moe = ( + hasattr(config, "num_experts") + and config.num_experts is not None + and layer_idx >= getattr(config, "first_k_dense_replace", 0) + and layer_idx % getattr(config, "moe_layer_freq", 1) == 0 + ) + + self.mlp = SarvamMLAMoE(config) if use_moe else SarvamMLAMLP(config) + self.input_layernorm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + if use_cache: + outputs += (present_key_value,) + return outputs + + +class SarvamMLAPreTrainedModel(PreTrainedModel): + config_class = SarvamMLAConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["SarvamMLADecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = False # Not implemented yet + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class SarvamMLAModel(SarvamMLAPreTrainedModel): + def __init__(self, config: SarvamMLAConfig): + 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( + [SarvamMLADecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._use_flash_attention_2 = False # Not implemented yet + self.norm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + 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 = None, + 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, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = _get_usable_past_kv_length(past_key_values, seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + + hidden_states = inputs_embeds + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class SarvamMLAForCausalLM(SarvamMLAPreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = SarvamMLAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: torch.LongTensor = None, + 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, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + **kwargs, + ): + if past_key_values is not None: + if isinstance(past_key_values, Cache): + cache_length = past_key_values.get_seq_length() + past_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if hasattr(past_key_values, "get_max_length"): + max_cache_length = past_key_values.get_max_length() + else: + max_cache_length = None + else: + cache_length = past_length = past_key_values[0][0].shape[2] + max_cache_length = None + + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + + if ( + max_cache_length is not None + and attention_mask is not None + and cache_length + input_ids.shape[1] > max_cache_length + ): + attention_mask = attention_mask[:, -max_cache_length:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), + ) + return reordered_past diff --git a/prism_int4.py b/prism_int4.py new file mode 100644 index 0000000000000000000000000000000000000000..645fec219baec3bbea2d325f7c508b2abcb0ca92 --- /dev/null +++ b/prism_int4.py @@ -0,0 +1,627 @@ +"""Sarvam-Prism with INT4 routed experts, for any CUDA GPU with 16 GB or less (tested within a 15 GiB and an 11 GiB cap). + + from prism_int4 import Prism + model = Prism.from_pretrained("jc-914/prism-quantized") # or a local copy of the repo + model.decide(state={"ticket": "I was charged twice."}, + questions={"team": {"type": "choice", "instructions": "Route this ticket.", + "criteria": {"billing": "payments and refunds", "tech": "app problems"}}}) + +What changes from the original FP8 release (`prism.py` in sarvamai/prism, which needs ~120 GB of FP8-capable memory): + * the 128 routed experts of each of the 31 MoE layers (~100B of the ~105B parameters) are INT4, GPTQ-quantized + (symmetric, group 128, fp32 scales). They live in pinned host RAM (~54 GB) and are streamed to the GPU one layer + at a time on a side CUDA stream, double-buffered behind the previous layer's compute. As many layers as fit in + the GPU memory budget stay resident instead; with enough memory everything is resident and nothing streams. + * every other projection keeps the original FP8 weights (one scale per output row), widened to bf16 at use + (weight-only FP8, bf16 activations), so it runs on GPUs without FP8 tensor cores. These weights are unchanged. + * routed experts run as a Triton grouped W4A16 kernel with fp32 accumulation: bf16 tensor cores on sm_80+, + fp16 tensor cores on older GPUs (expert inputs are RMS-normalised, so fp16 range is not a concern). + * the token embedding table (2 GB) stays in host RAM; lookups are copied to the GPU. + +Needs: a CUDA GPU (sm_80+ recommended; sm_70/75 supported but slower), >= ~10 GB of GPU memory, >= 64 GB host RAM, +PyTorch >= 2.4, Triton >= 3.0, transformers 4.52.1. There are no CUDA graphs (weights move every forward), so when +layers stream, latency is bounded by host->GPU bandwidth: put many questions (or requests, `decide_batch`) per call. +""" +from __future__ import annotations + +import gc +import json +import math +import os +import sys +import time +import types +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from torch.nn.attention import SDPBackend, sdpa_kernel + +GROUP = 128 + +# --------------------------------------------------------------------------------------------------------------------- +# Input encoding (identical to prism.py) +# --------------------------------------------------------------------------------------------------------------------- +TAGS = {"state": "State:", "question": "\nQuestion:", "open": "\n", "answer": "\nAnswer:"} +CONTEXT, ANSWER_ROLE, PAD = -1, -2, -3 +SEG_STATE, SEG_QUESTION, SEG_OPTION, SEG_ANSWER = range(4) +TYPES = {"choice": 0, "yes_no": 1, "score": 2} + + +def _text(x): + return x if isinstance(x, str) else json.dumps(x, ensure_ascii=False, separators=(",", ":")) + + +@dataclass +class Encoded: + ids: list + pos: list + role: list + seg: list + close_idx: list + answer_idx: int + n_context: int + + +class Encoder: + def __init__(self, tok): + self.tok = tok + self.tags = {k: tok(v, add_special_tokens=False).input_ids for k, v in TAGS.items()} + + def _ids(self, s): + for t in ("", "