fix(weights): stack per-expert MoE tensors into switch_mlp layout for mlx-lm

Ornith-1.0-35B ships experts in the legacy per-expert HF layout, which
mlx-lm's qwen3_5_moe sanitize() does not stack, so the quantized build
failed to load with 'Received 92160 parameters not in model'. Stack the
92,160 per-expert tensors into 360 switch_mlp tensors (bitwise-identical
values, no requantization) and rewrite the per-path quantization override
keys in config.json to post-sanitize module paths. Ship
repair_moe_experts.py (in-place repair with bitwise self-check) so broken
earlier downloads can be fixed locally, and document it in the README.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

README.md

@@ -99,6 +99,27 @@ To seamlessly route this model through agentic development workspaces like **Ope
 
 ---
 
+## 🩹 Repair Script (`repair_moe_experts.py`)
+
+The weights in this repository are already fixed and load correctly — **you do not need to run this for normal use.**
+
+It's included only for anyone who cached an earlier broken download of this repo, or who runs into the same issue when quantizing another MoE model with a similar per-expert weight layout. Symptom: loading fails with an error like:
+
+```
+Received 92160 parameters not in model:
+language_model.model.layers.0.mlp.experts.0.down_proj.biases, ...
+```
+
+Cause: Ornith-1.0-35B ships its MoE experts as separate per-expert tensors (`mlp.experts.<0-255>.{gate,up,down}_proj`), but mlx-lm's `qwen3_5_moe` loader only understands the fused `switch_mlp` layout — so an affected build never quantizes them into a loadable shape, and its `config.json` also carries stale pre-sanitize key names for the per-path quantization overrides.
+
+Usage, if ever needed:
+
+```bash
+python3 repair_moe_experts.py /path/to/Ornith-1.0-35B-oQ8-fp16
+```
+
+Requires only `mlx` (any environment with `mlx-lm`/`omlx` installed has it). It streams the weights shard-by-shard, stacks the per-expert tensors into `switch_mlp` (bitwise-identical values — no requantization), and fixes the `config.json` key names. The repair happens **in place**: new shards are written alongside the originals, verified bitwise against the source tensors, and only then swapped in — a failure at any point leaves the original model untouched. It needs free disk roughly equal to the model size while running, and is safe to re-run (exits early on an already-repaired model).
+
 ## 🌡️ Thermal Optimization Notice
 
 Sustained execution across massive context windows heavily taxes the Apple Silicon GPU/CPU complexes, causing rapid heat buildup. Because Apple's default fan profiles emphasize near-silent operation, they delay ramping up system fans until thermal throttling has already begun to affect generation tokens-per-second (TPS).

model-00001-of-00008.safetensors → model-00001-of-00007.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8c765d88a30ee2f45e45f8f973b24249d270c6976e8e6d604152f5097c8aef72
-size 5002214088
+oid sha256:19546eef0eac6e498450aee132407b7270eb89857daafa8764544d7ec99fac43
+size 5425639655

model-00002-of-00008.safetensors → model-00002-of-00007.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3a3c8052e5361f34dd6e995c7efb5fc143c2ffb2d574fa218848cb523c7b14ad
-size 5001824658
+oid sha256:89fa11a0cbf88edff7ccd68203311ee2da26a9d7f4ac7e635f84600435be8751
+size 5613406117

model-00003-of-00008.safetensors → model-00003-of-00007.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ef490ac56cc38051a387b74c7c2d7de6069652a1a2df09425c39fd11bcd1c701
-size 5002272652
+oid sha256:0614d305661718a7815d0c4196992943763bc696c36b76c69ba7df9a712b4077
+size 5427797286

model-00004-of-00008.safetensors → model-00004-of-00007.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7ee28a9bed6fe63dd5d93b6303e23d88bdbc583d75b9aaec4b8e329effea30b6
-size 5002718081
+oid sha256:467c53dd16c1436c9aa84cb2f63cf90ad4796a66322aaab207f98cf7b9a023d4
+size 5577316492

model-00005-of-00007.safetensors

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

model-00005-of-00008.safetensors

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

model-00006-of-00007.safetensors

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

model-00006-of-00008.safetensors

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

model-00007-of-00007.safetensors

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

model-00007-of-00008.safetensors

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

model-00008-of-00008.safetensors

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

repair_moe_experts.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python3
+"""Repair oMLX-quantized Qwen3.5-MoE checkpoints that store routed experts
+in the legacy per-expert layout (mlp.experts.<E>.{gate,up,down}_proj.*).
+
+mlx-lm's qwen3_5_moe sanitize() only stacks the fused `experts.gate_up_proj`
+layout, so per-expert checkpoints fail to load with
+"Received NNNNN parameters not in model". This script repairs the model
+in place:
+
+  1. Stacks every `<prefix>.mlp.experts.<E>.<proj>.<tensor>` group along a new
+     leading axis into `<prefix>.mlp.switch_mlp.<proj>.<tensor>`.
+  2. Rewrites per-path quantization overrides in config.json from raw HF key
+     names (model.language_model.*) to post-sanitize module paths
+     (language_model.model.*), which is how mlx-lm looks them up at load time.
+
+New shards are written alongside the originals, verified bitwise against the
+source tensors, and only then swapped in (originals deleted). A failure at any
+point leaves the original model untouched. Needs free disk roughly equal to
+the model size while running.
+
+Usage:
+    python3 repair_moe_experts.py <model_dir>
+
+Requires mlx; no other dependencies.
+"""
+
+import json
+import re
+import struct
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+
+EXPERT_RE = re.compile(
+    r"^(?P<prefix>.+\.mlp)\.experts\.(?P<e>\d+)\.(?P<proj>\w+_proj)\.(?P<t>weight|scales|biases)$"
+)
+SHARD_BYTES = 5 * 1024**3
+TMP_PREFIX = "tmp-shard-"
+
+
+def quant_key_to_module_path(key):
+    if key.startswith("model.language_model"):
+        return key.replace("model.language_model", "language_model.model", 1)
+    if key.startswith("language_model."):
+        return key
+    return "language_model." + key
+
+
+def read_header(path):
+    with open(path, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        return json.loads(f.read(n))
+
+
+def nbytes(a):
+    return a.size * a.dtype.size
+
+
+class ShardWriter:
+    def __init__(self, out_dir, metadata):
+        self.out_dir = out_dir
+        self.metadata = metadata
+        self.buffer = {}
+        self.buffer_bytes = 0
+        self.files = []  # [(tmp_path, [keys])]
+
+    def add(self, key, array):
+        self.buffer[key] = array
+        self.buffer_bytes += nbytes(array)
+        if self.buffer_bytes >= SHARD_BYTES:
+            self.flush()
+
+    def flush(self):
+        if not self.buffer:
+            return
+        tmp = self.out_dir / f"{TMP_PREFIX}{len(self.files):05d}.safetensors"
+        mx.save_safetensors(str(tmp), self.buffer, metadata=self.metadata)
+        mx.clear_cache()
+        self.files.append((tmp, list(self.buffer)))
+        self.buffer = {}
+        self.buffer_bytes = 0
+
+    def tmp_weight_map(self):
+        return {k: tmp for tmp, keys in self.files for k in keys}
+
+    def commit(self, old_shards):
+        """Delete the original shards and move tmp shards to final names."""
+        for shard in old_shards:
+            (self.out_dir / shard).unlink()
+        n = len(self.files)
+        weight_map, total = {}, 0
+        for i, (tmp, keys) in enumerate(self.files, 1):
+            name = f"model-{i:05d}-of-{n:05d}.safetensors"
+            tmp.rename(self.out_dir / name)
+            hdr = read_header(self.out_dir / name)
+            for k, v in hdr.items():
+                if k != "__metadata__":
+                    total += v["data_offsets"][1] - v["data_offsets"][0]
+            for k in keys:
+                weight_map[k] = name
+        index = {"metadata": {"total_size": total}, "weight_map": weight_map}
+        with open(self.out_dir / "model.safetensors.index.json", "w") as f:
+            json.dump(index, f, indent=2)
+        return weight_map
+
+
+def cleanup_tmp(model_dir):
+    for p in model_dir.glob(f"{TMP_PREFIX}*.safetensors"):
+        p.unlink()
+
+
+def main():
+    if len(sys.argv) != 2:
+        sys.exit(__doc__)
+    model_dir = Path(sys.argv[1]).resolve()
+    index_file = model_dir / "model.safetensors.index.json"
+    if not index_file.exists():
+        sys.exit(f"error: {index_file} not found")
+    cleanup_tmp(model_dir)  # leftovers from an interrupted run
+
+    weight_map = json.load(open(index_file))["weight_map"]
+    expert_keys = [k for k in weight_map if EXPERT_RE.match(k)]
+    if not expert_keys:
+        print("no per-expert tensors found — model is already repaired")
+        return
+
+    # group per-expert keys by their stacked target
+    groups = {}  # stacked_key -> {expert_idx: source_key}
+    for k in expert_keys:
+        m = EXPERT_RE.match(k)
+        stacked = f"{m['prefix']}.switch_mlp.{m['proj']}.{m['t']}"
+        groups.setdefault(stacked, {})[int(m["e"])] = k
+    n_experts = {len(v) for v in groups.values()}
+    if len(n_experts) != 1:
+        sys.exit(f"error: inconsistent expert counts per group: {sorted(n_experts)}")
+    n_experts = n_experts.pop()
+    key_to_stacked = {sk: stacked for stacked, exps in groups.items() for sk in exps.values()}
+    print(f"{len(expert_keys)} per-expert tensors -> {len(groups)} stacked tensors ({n_experts} experts)")
+
+    mx.set_default_device(mx.cpu)
+    old_shards = sorted({v for v in weight_map.values()})
+    src_metadata = read_header(model_dir / old_shards[0]).get("__metadata__") or {"format": "mlx"}
+    writer = ShardWriter(model_dir, src_metadata)
+    pending = {}  # stacked_key -> {expert_idx: array}
+
+    try:
+        for i, shard in enumerate(old_shards, 1):
+            print(f"[{i}/{len(old_shards)}] {shard}")
+            tensors = mx.load(str(model_dir / shard))
+            for key, array in tensors.items():
+                stacked = key_to_stacked.get(key)
+                if stacked is None:
+                    writer.add(key, array)
+                    continue
+                e = int(EXPERT_RE.match(key)["e"])
+                pending.setdefault(stacked, {})[e] = array
+                if len(pending[stacked]) == n_experts:
+                    parts = pending.pop(stacked)
+                    assert sorted(parts) == list(range(n_experts)), f"non-contiguous experts for {stacked}"
+                    writer.add(stacked, mx.stack([parts[j] for j in range(n_experts)]))
+            del tensors
+        if pending:
+            raise RuntimeError(f"incomplete expert groups: {list(pending)[:3]}")
+        writer.flush()
+
+        # verify before touching the originals
+        tmp_map = writer.tmp_weight_map()
+        expected = len(weight_map) - len(expert_keys) + len(groups)
+        if len(tmp_map) != expected:
+            raise RuntimeError(f"tensor count mismatch: {len(tmp_map)} != {expected}")
+        check = [(sk, e) for sk in list(groups)[::max(1, len(groups) // 4)] for e in (0, n_experts - 1)]
+        by_tmp = {}
+        for sk, e in check:
+            by_tmp.setdefault(tmp_map[sk], []).append((sk, e))
+        for tmp, items in by_tmp.items():
+            out_tensors = mx.load(str(tmp))
+            for sk, e in items:
+                src_key = groups[sk][e]
+                orig = mx.load(str(model_dir / weight_map[src_key]))[src_key]
+                if not mx.array_equal(out_tensors[sk][e], orig).item():
+                    raise RuntimeError(f"bitwise mismatch at {sk}[{e}]")
+            del out_tensors
+        print(f"spot-check passed ({len(check)} slices)")
+    except Exception as e:
+        cleanup_tmp(model_dir)
+        sys.exit(f"error: {e} — original model left untouched")
+
+    # point of no return: swap repaired shards in, rewrite index and config
+    writer.commit(old_shards)
+    config_file = model_dir / "config.json"
+    config = json.load(open(config_file))
+    for section in ("quantization", "quantization_config"):
+        if isinstance(config.get(section), dict):
+            config[section] = {
+                (quant_key_to_module_path(k) if isinstance(v, dict) else k): v
+                for k, v in config[section].items()
+            }
+    with open(config_file, "w") as f:
+        json.dump(config, f, indent=2)
+    print(f"done: {model_dir} repaired in place")
+
+
+if __name__ == "__main__":
+    main()