Upload folder using huggingface_hub

.gitattributes

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

LICENSE

@@ -0,0 +1,29 @@
+GLM-5.3 License
+
+Copyright (c) 2026 Z.AI
+
+Permission is hereby granted, free of charge, to any person or entity (the "Licensee") obtaining a copy of this software — including the model weights, parameters, configuration files, inference and training code, and associated documentation (collectively, the "Software") — to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software; to run, deploy, fine-tune, or otherwise modify the Software and create derivative works from it; and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+1. The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. The Licensee's use of the Software must comply with applicable laws and regulations.
+
+2. "Model as a Service" means giving a third party access to language model inference or fine-tuning (e.g., via API) in a manner that allows such third party to exercise meaningful control over the inputs, parameters, or training data. This does not include (a) end-user products with model capabilities solely embedded within specific features or harnesses, or (b) mere relaying of requests to models hosted by others.
+If the Licensee or any of its affiliates operates a Model as a Service business, and the aggregate revenue of the Licensee and its affiliates exceeds 10 billion US dollars (or the equivalent in other currencies) in total over any consecutive 12 months, the Licensee must pass Z.AI's security review before using the Software or its derivative works for any commercial purpose. The scope and method of the security review shall be reasonably determined by Z.AI.
+
+3. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL Z.AI OR ITS AFFILIATES OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+For any questions regarding this license, please contact glmlicense@z.ai.
+
+-----
+
+版权所有 (c) 2026 Z.AI
+
+特此免费授予任何获得本软件副本的个人或实体（"被许可方"）——包括模型权重、参数、配置文件、推理和训练代码及相关文档（统称"软件"）——不受限制地处理本软件的权利，包括但不限于：使用、复制、修改、合并、发布、分发、再许可和/或销售软件副本；运行、部署、微调或以其他方式修改软件并创建衍生作品；以及允许获得软件的其他人行使上述权利，但须遵守以下条件：
+
+1. 上述版权声明和本许可声明应包含在软件的所有副本或实质性部分中。被许可方对软件的使用必须符合适用法律法规。
+
+2. "模型即服务"指以允许第三方对输入、参数或训练数据行使实质性控制的方式，向第三方提供语言模型推理或微调服务（如通过API）。不包括：(a) 模型能力仅嵌入特定功能或框架中的终端用户产品，或(b) 单纯转发请求至他人托管的模型。
+若被许可方或其关联方运营"模型即服务"业务，且被许可方及关联方在任意连续12个月内累计总收入超过100亿美元（或等值其他货币），则被许可方在使用软件或其衍生作品进行任何商业用途之前，须通过Z.AI的安全审查。安全审查的范围和方式由Z.AI合理确定。
+
+3. 软件及其任何输出和结果均按"现状"提供，不附带任何形式的保证，无论是明示还是暗示，包括但不限于适销性、特定用途适用性和不侵权的保证。在任何情况下，Z.AI或其关联方或版权持有人均不对任何索赔、损害或其他责任承担责任，无论该责任是基于合同、侵权或其他方式，因软件或使用软件而产生或与之相关。
+
+如对本许可有任何疑问，请联系 glmlicense@z.ai。

README.md

@@ -0,0 +1,95 @@
+---
+license: other
+license_name: glm-5.3
+license_link: LICENSE
+base_model: zai-org/GLM-5.3
+base_model_relation: quantized
+tags:
+- mlx
+- apple-silicon
+- glm_moe_dsa
+- mixture-of-experts
+- 5-bit
+pipeline_tag: text-generation
+library_name: mlx
+---
+
+# GLM-5.3-MLX-5bit
+
+MLX (Apple Silicon) build of [**GLM-5.3**](https://huggingface.co/zai-org/GLM-5.3) — 744B-parameter
+`glm_moe_dsa` MoE (256 experts, top-8; MLA with DeepSeek-V3.2-style sparse attention) — quantized
+to **5-bit**.
+
+**These files are modified**: converted from the upstream **bfloat16** release
+([GLM-5.3-BF16](https://huggingface.co/zai-org/GLM-5.3-BF16)) to MLX and quantized; the architecture
+is unchanged. The multi-token-prediction layer (78) is not included.
+
+## Runtime — read this
+
+This checkpoint bundles `glm_moe_dsa.py` (declared via `model_file`) and needs it:
+
+```bash
+pip install -U mlx-lm
+mlx_lm.generate --model pipenetwork/GLM-5.3-MLX-5bit --trust-remote-code --prompt "..." --max-tokens 300
+```
+
+mlx-lm's own `glm_moe_dsa` builds a lightning indexer on all 78 layers, but GLM-5.2/5.3 ship
+indexer weights on 21 (`indexer_types`: the other 57 "shared" layers reuse the previous full layer's
+top-k selection). A strict load of the release fails with 285 missing parameters; `mlx_lm.load`
+loads leniently and leaves those 57 indexers at random initialisation. Prompts up to 2048 tokens are
+unaffected (the indexer is bypassed below `index_topk`); beyond that, 57 layers attend to keys
+chosen by random projections. The bundled runtime implements the schedule as the reference does,
+plus the reference's fp32 indexer scores and router logits and the indexer LayerNorm epsilon.
+Tiny-config parity against `transformers` 5.16 is **4e-7** with the sparse path live, cached decode
+exact; strict loading of this checkpoint reports zero missing and zero unexpected tensors. Details and
+tests: [https://github.com/PipeNetwork/glm53-mlx](https://github.com/PipeNetwork/glm53-mlx).
+
+## Size and what is quantized
+
+**511.5 GB** on disk. RAM: 768 GB (two machines).
+
+| group | share of parameters | this build |
+|---|---:|---|
+| routed experts (`switch_mlp`, 75 layers × 256) | 724.8B (97.5%) | 5-bit, group 64 |
+| attention (MLA), shared experts, dense layers 0–2, embeddings, `lm_head` | 18.4B (2.5%) | 5-bit, group 64 |
+| lightning indexer (21 layers), MoE router + correction bias, norms | 0.3B | as stored (bf16 / fp32) |
+
+Source precision: the FP8 release is a lossy derivative of the bf16 one (dequantized FP8 weights
+differ from bf16 by up to 1.6e-2 on values of 0.46 — half an e4m3 step). The ladder row `fp8` is the FP8 release itself measured against bf16: its error is the floor any FP8-sourced build inherits.
+
+## Quality
+
+Two measurements, because at 744B most of the ladder cannot be loaded on a 512 GB machine:
+
+**Per-layer divergence vs bf16** (`scripts/eval_ladder.py`): every decoder layer run in bf16 and in
+each recipe on identical inputs (16,384 tokens of wikitext-2), *teacher-forced* (each layer
+sees bf16 inputs — isolates its own damage) and *free-running* (each recipe feeds itself — what
+inference does). Relative L2 error of the layer output; lower is better.
+
+| recipe | teacher-forced (mean over layers) | free-running (final layer) | cosine (final) |
+|---|---:|---:|---:|
+| 8bit | 0.00685 | 0.13119 | 0.98945 |
+| 6bit | 0.01465 | 0.16736 | 0.98389 |
+| 5bit | 0.02651 | 0.22521 | 0.97272 |
+| 4bit | 0.05161 | 0.35740 | 0.93390 |
+| mixed-4_8bit | 0.02524 | 0.24951 | 0.96710 |
+| mixed-3_6bit | 0.05242 | 0.42380 | 0.90624 |
+| fp8 | 0.01741 | 0.17321 | 0.98320 |
+
+**Perplexity** on wikitext-2 (test), 288,627 tokens in 141 windows of 2048, for the builds
+that fit this machine, scored on identical windows:
+
+| build | size | perplexity [95% CI] |
+|---|---:|---|
+| [4bit](https://huggingface.co/pipenetwork/GLM-5.3-MLX-4bit) | 418.6 GB | 2.8636 [2.6681, 3.0714] |
+| [mixed-4_8bit](https://huggingface.co/pipenetwork/GLM-5.3-MLX-mixed-4_8bit) | 427.8 GB | 2.7420 [2.5533, 2.9477] |
+| [mixed-3_6bit](https://huggingface.co/pipenetwork/GLM-5.3-MLX-mixed-3_6bit) | 332.6 GB | 3.0338 [2.8366, 3.2386] |
+| [REAP50-4bit](https://huggingface.co/pipenetwork/GLM-5.3-REAP50-MLX-4bit) | 214.7 GB | 5.0295 [4.7571, 5.3137] |
+
+**Recommendation.** For a 512 GB Mac, **mixed 4/8-bit** (427.7 GB): perplexity 2.7420, a paired 4.3% better than uniform 4-bit (ratio 0.9575 [0.9537, 0.9612], better on 98.6% of windows) for 9 GB more — the 2.5% of non-expert weights are worth their 8 bits, as on every model we have measured. Uniform 4-bit (418.6 GB) is the fallback when those 9 GB matter. **Mixed 3/6-bit** (332.6 GB) is the 384 GB-class option, at a real cost: 3.0338, +5.9% over 4-bit and +10.6% over mixed 4/8 — it leads the ladder for the first ten layers and then 3-bit expert damage compounds. Among the builds that cannot be run here, the ladder puts 8-bit closest to bfloat16 (free-running error 0.131), then 6-bit (0.167); the **upstream FP8 release scores 0.173, between 6-bit and 5-bit**, which is why these are converted from the bf16 release. 5-bit (0.225) sits just above mixed 4/8 (0.250) at 100 GB more.
+
+Greedy generation (a collapse detector, not a ranking) is coherent on every published build.
+
+## License
+
+[GLM-5.3 license](LICENSE), as the upstream model. Port code: [https://github.com/PipeNetwork/glm53-mlx](https://github.com/PipeNetwork/glm53-mlx).

chat_template.jinja

@@ -0,0 +1,251 @@
+[gMASK]<sop>
+{%- set effective_reasoning_effort = reasoning_effort if reasoning_effort is defined and reasoning_effort in ['low', 'high'] else 'max' -%}
+{%- if effective_reasoning_effort is not none -%}<|system|>Reasoning Effort: {{ effective_reasoning_effort | capitalize }}{%- endif -%}
+{%- set clear_thinking = clear_thinking if clear_thinking is defined else false -%}
+{%- if tools -%}
+{%- macro tool_to_json(tool) -%}
+    {%- set ns_tool = namespace(first=true) -%}
+    {{ '{' -}}
+    {%- for k, v in tool.items() -%}
+        {%- if k != 'defer_loading' and k != 'strict' -%}
+            {%- if not ns_tool.first -%}{{- ', ' -}}{%- endif -%}
+            {%- set ns_tool.first = false -%}
+            "{{ k }}": {{ v | tojson(ensure_ascii=False) }}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- '}' -}}
+{%- endmacro -%}
+{%- macro tool_references_to_response(refs) -%}
+    {{- '<tool_response><tools>\n' -}}
+    {%- for tr in refs -%}
+        {%- for tool in tools -%}
+            {%- if 'function' in tool -%}
+                {%- set tool = tool['function'] -%}
+            {%- endif -%}
+            {%- if tool.name == tr.name -%}
+                {{- tool_to_json(tool) + '\n' -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+    {{- '</tools></tool_response>' -}}
+{%- endmacro -%}
+<|system|>
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{% for tool in tools %}
+{%- if 'function' in tool -%}
+    {%- set tool = tool['function'] -%}
+{%- endif -%}
+{% if tool.defer_loading is not defined or not tool.defer_loading %}
+{{ tool_to_json(tool) }}
+{% endif %}
+{% endfor %}
+</tools>
+
+For each function call, output the function name and arguments within the following XML format:
+<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>{%- endif -%}
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text }}
+            {%- elif item is string -%}
+                {{- item }}
+            {%- elif item is mapping and item.type in ['image', 'image_url', 'video', 'video_url', 'audio', 'audio_url', 'input_audio'] -%}
+                {%- set media_type = item.type | replace('_url', '') | replace('input_', '') -%}
+                {{- "<reminder>You are unable to process this " ~ media_type ~ " because you don't have multi-modal input ability. Try different methods.</reminder>" }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{- content }}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro tool_response(text) -%}
+{{- '<tool_response>' + text + '</tool_response>' -}}
+{%- endmacro -%}
+{%- macro render_tool_response(m) -%}
+{%- if m.content is string -%}
+    {{- tool_response(m.content) -}}
+{%- elif m.content and m.content is not mapping and m.content.0.type == "tool_reference" -%}
+    {{- tool_references_to_response(m.content) -}}
+{%- elif is_list_of_outputs(m) -%}
+    {%- for tr in m.content -%}
+        {%- if tr.output is iterable and tr.output is not string and tr.output is not mapping and tr.output and tr.output.0.type == "tool_reference" -%}
+            {{- tool_references_to_response(tr.output) -}}
+        {%- else -%}
+            {{- tool_response(visible_text(tr.output)) -}}
+        {%- endif -%}
+    {%- endfor -%}
+{%- else -%}
+    {{- tool_response(visible_text(m.content)) -}}
+{%- endif -%}
+{%- endmacro -%}
+{%- macro id_of(obj) -%}
+    {%- if obj.tool_call_id -%}
+        {{- obj.tool_call_id -}}
+    {%- elif obj.id -%}
+        {{- obj.id -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro is_list_of_outputs(m) -%}
+    {%- if m.content and m.content.0.output is defined -%}1{%- endif -%}
+{%- endmacro -%}
+{%- macro has_dup_tool_result_id(lo, hi, target) -%}
+    {%- set ns_cnt = namespace(n=0) -%}
+    {%- for k in range(lo, hi + 1) -%}
+        {%- set m = messages[k] -%}
+        {%- if is_list_of_outputs(m) -%}
+            {%- for entry in m.content -%}
+                {%- if id_of(entry) == target -%}
+                    {%- set ns_cnt.n = ns_cnt.n + 1 -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- elif id_of(m) == target -%}
+            {%- set ns_cnt.n = ns_cnt.n + 1 -%}
+        {%- endif -%}
+        {%- if ns_cnt.n > 1 -%}{%- break -%}{%- endif -%}
+    {%- endfor -%}
+    {%- if ns_cnt.n > 1 -%}1{%- endif -%}
+{%- endmacro -%}
+{%- macro tc_id_exists(tcs, target) -%}
+    {%- set ns_f = namespace(found=false) -%}
+    {%- for tc in tcs -%}
+        {%- if id_of(tc) == target -%}
+            {%- set ns_f.found = true -%}
+            {%- break -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if ns_f.found -%}1{%- endif -%}
+{%- endmacro -%}
+{%- set ns = namespace(last_user_index=-1) -%}
+{%- for m in messages %}
+    {%- if m.role == 'user' %}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{%- for m in messages -%}
+{%- if m.role == 'user' -%}<|user|>{{ visible_text(m.content) }}
+{%- elif m.role == 'assistant' -%}
+<|assistant|>
+{%- set content = visible_text(m.content) %}
+{%- if m.reasoning_content is string %}
+    {%- set reasoning_content = m.reasoning_content %}
+{%- elif '</think>' in content %}
+    {%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
+    {%- set content = content.split('</think>')[-1] %}
+{%- endif %}
+{%- if (not clear_thinking or loop.index0 > ns.last_user_index) and reasoning_content is defined -%}
+{{ '<think>' + reasoning_content +  '</think>'}}
+{%- else -%}
+{{ '<think></think>' }}
+{%- endif -%}
+{%- if content.strip() -%}
+{{ content.strip() }}
+{%- endif -%}
+{% if m.tool_calls %}
+{% for tc in m.tool_calls %}
+{%- if tc.function %}
+    {%- set tc = tc.function %}
+{%- endif %}
+{{- '<tool_call>' + tc.name -}}
+{% set _args = tc.arguments %}{% for k, v in _args.items() %}<arg_key>{{ k }}</arg_key><arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>{% endfor %}</tool_call>{% endfor %}
+{% endif %}
+{%- elif m.role == 'tool' -%}
+{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+    {{- '<|observation|>' -}}
+    {%- set block_start = loop.index0 -%}
+    {%- set ns_blk = namespace(end=block_start) -%}
+    {%- for j in range(block_start, messages|length) -%}
+        {%- if messages[j].role == 'tool' -%}
+            {%- set ns_blk.end = j -%}
+        {%- else -%}
+            {%- break -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- set ns_a = namespace(tool_calls=none) -%}
+    {%- if block_start > 0 and messages[block_start - 1].role == 'assistant' and messages[block_start - 1].tool_calls -%}
+        {%- set ns_a.tool_calls = messages[block_start - 1].tool_calls -%}
+    {%- endif -%}
+    {%- set ns_chk = namespace(can_sort=true) -%}
+    {%- if not ns_a.tool_calls -%}
+        {%- set ns_chk.can_sort = false -%}
+    {%- else -%}
+        {%- for k in range(block_start, ns_blk.end + 1) -%}
+            {%- set m = messages[k] -%}
+            {%- if is_list_of_outputs(m) -%}
+                {%- for entry in m.content -%}
+                    {%- set eid = id_of(entry) -%}
+                    {%- if not eid -%}
+                        {%- set ns_chk.can_sort = false -%}
+                    {%- elif has_dup_tool_result_id(block_start, ns_blk.end, eid) -%}
+                        {%- set ns_chk.can_sort = false -%}
+                    {%- elif not tc_id_exists(ns_a.tool_calls, eid) -%}
+                        {%- set ns_chk.can_sort = false -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- else -%}
+                {%- set tk_id = id_of(m) -%}
+                {%- if not tk_id -%}
+                    {%- set ns_chk.can_sort = false -%}
+                {%- elif has_dup_tool_result_id(block_start, ns_blk.end, tk_id) -%}
+                    {%- set ns_chk.can_sort = false -%}
+                {%- elif not tc_id_exists(ns_a.tool_calls, tk_id) -%}
+                    {%- set ns_chk.can_sort = false -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- for i in range(ns_a.tool_calls | length) -%}
+            {%- set tc_id = id_of(ns_a.tool_calls[i]) -%}
+            {%- if not tc_id -%}
+                {%- set ns_chk.can_sort = false -%}
+            {%- endif -%}
+            {%- for j in range(i + 1, ns_a.tool_calls | length) -%}
+                {%- if id_of(ns_a.tool_calls[j]) == tc_id -%}
+                    {%- set ns_chk.can_sort = false -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- if ns_chk.can_sort -%}
+        {%- for tc in ns_a.tool_calls -%}
+            {%- set tc_id = id_of(tc) -%}
+            {%- for k in range(block_start, ns_blk.end + 1) -%}
+                {%- set m = messages[k] -%}
+                {%- if is_list_of_outputs(m) -%}
+                    {%- for entry in m.content -%}
+                        {%- set eid = id_of(entry) -%}
+                        {%- if eid == tc_id -%}
+                            {%- if entry.output is iterable and entry.output is not string and entry.output is not mapping and entry.output and entry.output.0.type == "tool_reference" -%}
+                                {{- tool_references_to_response(entry.output) -}}
+                            {%- else -%}
+                                {{- tool_response(visible_text(entry.output)) -}}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                {%- else -%}
+                    {%- set tk_id = id_of(m) -%}
+                    {%- if tk_id == tc_id -%}
+                        {{- render_tool_response(m) -}}
+                    {%- endif -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endfor -%}
+    {%- else -%}
+        {%- for k in range(block_start, ns_blk.end + 1) -%}
+            {{- render_tool_response(messages[k]) -}}
+        {%- endfor -%}
+    {%- endif -%}
+{% endif -%}
+{%- elif m.role == 'system' -%}
+<|system|>{{ visible_text(m.content) }}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    <|assistant|>{{- '<think>' -}}
+{%- endif -%}

config.json

@@ -0,0 +1,233 @@
+{
+  "architectures": [
+    "GlmMoeDsaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "dtype": "bfloat16",
+  "eos_token_id": [
+    154820,
+    154827,
+    154829
+  ],
+  "ep_size": 1,
+  "first_k_dense_replace": 3,
+  "head_dim": 192,
+  "hidden_act": "silu",
+  "hidden_size": 6144,
+  "index_head_dim": 128,
+  "index_n_heads": 32,
+  "index_share_for_mtp_iteration": true,
+  "index_skip_topk_offset": 3,
+  "index_topk": 2048,
+  "index_topk_freq": 4,
+  "index_topk_pattern": null,
+  "indexer_rope_interleave": true,
+  "indexer_types": [
+    "full",
+    "full",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared",
+    "full",
+    "shared",
+    "shared",
+    "shared"
+  ],
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "kv_lora_rank": 512,
+  "max_position_embeddings": 1048576,
+  "mlp_layer_types": [
+    "dense",
+    "dense",
+    "dense",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse"
+  ],
+  "model_type": "glm_moe_dsa",
+  "moe_intermediate_size": 2048,
+  "moe_layer_freq": 1,
+  "moe_router_dtype": "float32",
+  "n_group": 1,
+  "n_routed_experts": 256,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 64,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 78,
+  "num_key_value_heads": 64,
+  "num_nextn_predict_layers": 1,
+  "pad_token_id": 154820,
+  "pretraining_tp": 1,
+  "q_lora_rank": 2048,
+  "qk_head_dim": 256,
+  "qk_nope_head_dim": 192,
+  "qk_rope_head_dim": 64,
+  "rms_norm_eps": 1e-05,
+  "rope_interleave": true,
+  "rope_parameters": {
+    "rope_theta": 8000000,
+    "rope_type": "default"
+  },
+  "routed_scaling_factor": 2.5,
+  "scoring_func": "sigmoid",
+  "tie_word_embeddings": false,
+  "topk_group": 1,
+  "topk_method": "noaux_tc",
+  "transformers_version": "5.15.0",
+  "use_cache": true,
+  "v_head_dim": 256,
+  "vocab_size": 154880,
+  "quantization": {
+    "group_size": 64,
+    "bits": 5
+  },
+  "quantization_config": {
+    "group_size": 64,
+    "bits": 5
+  },
+  "model_file": "glm_moe_dsa.py"
+}

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [
+    154820,
+    154827,
+    154829
+  ],
+  "pad_token_id": 154820,
+  "temperature": 1.0,
+  "top_p": 0.95,
+  "transformers_version": "5.12.0"
+}

glm_moe_dsa.py

@@ -0,0 +1,703 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+from mlx_lm.models.activations import swiglu
+from mlx_lm.models.base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from mlx_lm.models.cache import CacheList, KVCache
+from mlx_lm.models.mla import MultiLinear
+from mlx_lm.models.rope_utils import initialize_rope
+from mlx_lm.models.switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    index_head_dim: int
+    index_n_heads: int
+    index_topk: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: Optional[int]
+    n_routed_experts: Optional[int]
+    routed_scaling_factor: float
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    v_head_dim: int
+    qk_nope_head_dim: int
+    topk_method: str
+    scoring_func: str
+    norm_topk_prob: bool
+    n_group: int
+    topk_group: int
+    num_experts_per_tok: int
+    moe_layer_freq: int
+    first_k_dense_replace: int
+    max_position_embeddings: int
+    rms_norm_eps: float
+    rope_parameters: Dict
+    attention_bias: bool
+    rope_scaling: Dict = None
+    rope_theta: Optional[float] = None
+    indexer_rope_interleave: bool = True
+    rope_interleave: bool = True
+    indexer_types: Optional[list] = None
+    index_topk_freq: int = 1
+    index_skip_topk_offset: int = 0
+    index_topk_pattern: Optional[Any] = None
+    num_nextn_predict_layers: int = 0
+    mlp_layer_types: Optional[list] = None
+
+    def __post_init__(self):
+        self.rope_scaling = self.rope_parameters
+        self.rope_theta = self.rope_parameters["rope_theta"]
+        if self.indexer_types is None:
+            # Reference schedule: a layer runs its own indexer ("full") iff
+            # max(i - offset + 1, 0) % freq == 0; the others ("shared") reuse the most recent
+            # full layer's top-k selection and carry no indexer weights.
+            f, o = self.index_topk_freq, self.index_skip_topk_offset
+            self.indexer_types = [
+                "full" if max(i - o + 1, 0) % f == 0 else "shared"
+                for i in range(self.num_hidden_layers)
+            ]
+
+
+
+class Indexer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.dim = args.hidden_size
+        self.n_heads = args.index_n_heads
+        self.head_dim = args.index_head_dim
+        self.rope_head_dim = args.qk_rope_head_dim
+        self.index_topk = args.index_topk
+        self.q_lora_rank = args.q_lora_rank
+        self.wq_b = nn.Linear(
+            self.q_lora_rank, self.n_heads * self.head_dim, bias=False
+        )
+        self.wk = nn.Linear(self.dim, self.head_dim, bias=False)
+        self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6)
+        self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False)
+        self.softmax_scale = self.head_dim**-0.5
+        self.rope = initialize_rope(
+            dims=args.qk_rope_head_dim,
+            base=args.rope_theta,
+            traditional=args.indexer_rope_interleave,
+            max_position_embeddings=args.max_position_embeddings,
+            scaling_config=args.rope_scaling,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        qr: mx.array,
+        mask: Optional[mx.array],
+        cache: Optional[Any] = None,
+    ):
+        # Computes top_k indices for attention
+        b, s, _ = x.shape
+        q = self.wq_b(qr)
+        q = q.reshape(b, s, self.n_heads, self.head_dim).swapaxes(1, 2)
+        k = self.wk(x)
+        k = self.k_norm(k)
+        k = mx.reshape(k, (b, 1, s, self.head_dim))
+
+        offset = cache.offset if cache is not None else 0
+
+        q = self.rope(q, offset=offset)
+        k = self.rope(k, offset=offset)
+
+        if cache is not None:
+            k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
+        if k.shape[2] <= self.index_topk:
+            return None
+        # Scores in float32, as the reference (which also keeps weights_proj in fp32): in bf16
+        # keys at the top-k boundary flip.
+        scores = q.astype(mx.float32) @ k.astype(mx.float32).swapaxes(-1, -2)
+        scores = mx.maximum(scores, 0)
+        weights = self.weights_proj(x).astype(mx.float32) * (self.n_heads**-0.5 * self.softmax_scale)
+        weights = weights.swapaxes(-1, -2)[..., None]
+        scores = scores * weights
+        scores = scores.sum(axis=1, keepdims=True)
+        if mask is not None:
+            scores = mx.where(mask, scores, -float("inf"))
+        return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[
+            ..., -self.index_topk :
+        ]
+
+
+class DeepseekV32Attention(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int = 0):
+        super().__init__()
+        self.config = config
+        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 = config.q_lora_rank
+        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.qk_nope_head_dim + config.qk_rope_head_dim
+
+        self.scale = self.q_head_dim**-0.5
+
+        self.q_a_proj = nn.Linear(
+            self.hidden_size, self.q_lora_rank, bias=config.attention_bias
+        )
+        self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
+        self.q_b_proj = nn.Linear(
+            self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+        )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=config.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+        )
+
+        if self.config.rope_scaling is not None:
+            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+            if mscale_all_dim:
+                scaling_factor = self.config.rope_scaling["factor"]
+                if scaling_factor > 1:
+                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
+                    self.scale = self.scale * s * s
+
+        # "shared" layers carry no indexer weights: they reuse the previous full layer's top-k.
+        self.indexer = Indexer(config) if config.indexer_types[layer_idx] == "full" else None
+        self.rope = initialize_rope(
+            dims=self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=True,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_config=self.config.rope_scaling,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        prev_topk_indices: Optional[mx.array] = None,
+    ):
+        B, L, D = x.shape
+
+        qr = self.q_a_layernorm(self.q_a_proj(x))
+        q = self.q_b_proj(qr)
+
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
+
+        offset = cache[0].offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
+
+        if cache is not None:
+            kv_latent, k_pe = cache[0].update_and_fetch(kv_latent, k_pe)
+        else:
+            cache = [None] * 2
+
+        if self.indexer is not None:
+            topk_indices = self.indexer(x, qr, mask, cache=cache[1])
+        else:
+            topk_indices = prev_topk_indices
+        if topk_indices is not None:
+            if L == 1:
+                idx = topk_indices[:, :, 0, :, None]
+                kv_latent = mx.take_along_axis(
+                    kv_latent,
+                    mx.broadcast_to(idx, idx.shape[:-1] + (kv_latent.shape[-1],)),
+                    axis=2,
+                )
+                k_pe = mx.take_along_axis(
+                    k_pe,
+                    mx.broadcast_to(idx, idx.shape[:-1] + (k_pe.shape[-1],)),
+                    axis=2,
+                )
+                if mask is not None:
+                    mask = mx.take_along_axis(mask, topk_indices, axis=-1)
+            else:
+                shape = list(topk_indices.shape)
+                shape[-1] = kv_latent.shape[2]
+                sparse_mask = mx.zeros(shape, dtype=mx.bool_)
+                sparse_mask = mx.put_along_axis(
+                    sparse_mask, topk_indices, mx.array(True), axis=-1
+                )
+                if mask is not None:
+                    sparse_mask = sparse_mask & mask
+                mask = sparse_mask
+        # Ensure the indexer cache is evaluated even if the topk_indices are unused
+        # to keep the graph from getting too large
+        if self.indexer is not None and cache is not None and cache[0] is not None:
+            cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values))
+
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
+        output = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
+        )
+        if L == 1:
+            output = self.unembed_out(output)
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output), topk_indices
+
+
+class DeepseekV32MLP(nn.Module):
+    def __init__(
+        self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = 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)
+
+    def __call__(self, x):
+        down_proj = self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+        return down_proj
+
+
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + e_score_correction_bias
+    if n_group > 1:
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        denominator = scores.sum(axis=-1, keepdims=True)
+        scores = scores / denominator
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+        assert config.topk_method == "noaux_tc", "Unsupported topk method."
+
+    def __call__(self, x):
+        return group_expert_select(
+            x.astype(mx.float32) @ self.weight.astype(mx.float32).T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class DeepseekV32MoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+            self.shared_experts = DeepseekV32MLP(
+                config=config, intermediate_size=intermediate_size
+            )
+
+        self.sharding_group = None
+
+    def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
+
+
+class DeepseekV32DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = DeepseekV32Attention(config, layer_idx)
+        self.mlp = (
+            DeepseekV32MoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+                and layer_idx % config.moe_layer_freq == 0
+            )
+            else DeepseekV32MLP(config)
+        )
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        prev_topk_indices: Optional[mx.array] = None,
+    ):
+        r, topk_indices = self.self_attn(self.input_layernorm(x), mask, cache, prev_topk_indices)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r, topk_indices
+
+
+class DeepseekV32Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            DeepseekV32DecoderLayer(config, idx)
+            for idx in range(config.num_hidden_layers)
+        ]
+        self.start_idx = 0
+        self.end_idx = len(self.layers)
+        self.num_layers = self.end_idx
+
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pipeline_rank = 0
+        self.pipeline_size = 1
+
+    def pipeline(self, group):
+        # Split layers in reverse so rank=0 gets the last layers and
+        # rank=pipeline_size-1 gets the first
+        self.pipeline_rank = group.rank()
+        self.pipeline_size = group.size()
+        layers_per_rank = len(self.layers) // self.pipeline_size
+        extra = len(self.layers) - layers_per_rank * self.pipeline_size
+        if self.pipeline_rank < extra:
+            layers_per_rank += 1
+        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
+        self.end_idx = self.start_idx + layers_per_rank
+        self.layers = self.layers[: self.end_idx]
+        self.layers[: self.start_idx] = [None] * self.start_idx
+        self.num_layers = len(self.layers) - self.start_idx
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
+        if cache is None:
+            cache = [None] * self.num_layers
+        mask = create_attention_mask(
+            h, cache[0][0] if cache[0] else None, return_array=True
+        )
+
+        # Receive from the previous process in the pipeline
+
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        # A "shared" layer reuses the top-k selection of the most recent "full" layer.
+        topk = None
+        for i in range(self.num_layers):
+            h, topk = self.layers[self.start_idx + i](h, mask, cache[i], topk)
+
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1][0].keys = mx.depends(cache[-1][0].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = DeepseekV32Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        # Remove multi-token prediction layers
+        mpt_layer = self.args.num_hidden_layers
+        new_weights = {}
+        for k, v in weights.items():
+            parts = k.split(".")
+            if len(parts) >= 3 and parts[1] == "layers" and int(parts[2]) >= mpt_layer:
+                continue
+            new_weights[k] = v
+        weights = new_weights
+
+        def dequant(weight, scale_inv):
+            dtype = mx.bfloat16
+            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+            bs = 128  # block size
+            m, n = weight.shape
+            pad_bottom = (-m) % bs
+            pad_side = (-n) % bs
+            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+            weight = weight.reshape(
+                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+            )
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n].astype(dtype)
+
+        # Dequantize
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                weight = dequant(weight, scale_inv)
+                new_weights[wk] = weight
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.model.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
+
+        return weights
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+        for layer in self.model.layers:
+            layer.self_attn.q_b_proj = shard_linear(
+                layer.self_attn.q_b_proj, "all-to-sharded", group=group
+            )
+
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV32MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
+    @property
+    def layers(self):
+        return self.model.layers[self.model.start_idx : self.model.end_idx]
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, module):
+            # The lightning indexer (0.03% of parameters, reference keeps weights_proj in fp32) and the
+            # MoE router decide *what* is computed; keep them as stored.
+            if ".indexer." in path or path.endswith("mlp.gate"):
+                return False
+            return True
+
+        return predicate
+
+    def make_cache(self):
+        # Shared layers have no indexer, so no indexer key cache: an empty second KVCache would
+        # trip mlx-lm's `cache.state` evaluation (keys is None) during generation.
+        return [
+            CacheList(KVCache(), KVCache()) if layer.self_attn.indexer is not None else CacheList(KVCache())
+            for layer in self.layers
+        ]

model-00001.safetensors

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

model-00002.safetensors

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

model-00003.safetensors

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

model-00004.safetensors

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

model-00005.safetensors

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

model-00006.safetensors

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

model-00007.safetensors

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

model-00008.safetensors

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

model-00009.safetensors

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

model-00010.safetensors

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

model-00011.safetensors

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

model-00012.safetensors

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

model-00013.safetensors

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

model-00014.safetensors

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

model-00015.safetensors

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

model-00016.safetensors

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

model-00017.safetensors

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

model-00018.safetensors

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

model-00019.safetensors

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

model-00020.safetensors

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

model-00021.safetensors

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

model-00022.safetensors

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

model-00023.safetensors

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

model-00024.safetensors

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

model-00025.safetensors

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

model-00026.safetensors

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

model-00027.safetensors

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

model-00028.safetensors

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

model-00029.safetensors

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

model-00030.safetensors

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

model-00031.safetensors

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

model-00032.safetensors

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

model-00033.safetensors

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

model-00034.safetensors

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

model-00035.safetensors

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

model-00036.safetensors

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

model-00037.safetensors

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

model-00038.safetensors

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

model-00039.safetensors

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

model-00040.safetensors

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

model-00041.safetensors

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

model-00042.safetensors

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