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 replicate a highly stable workspace inside coding environments like **OpenCod
 
 ---
 
+## 🩹 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-oQ5-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 over deep contexts heavily loads the Apple Silicon SoC, raising internal core temperatures. Because native macOS fan curves prioritize absolute quiet over proactive temperature maintenance, they often delay full fan deployment until minor thermal throttling occurs.

model-00001-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a5d275076ed6374a4dc43724801bd7b0e55ba9c8a98b83602314fdadeaa1befe
-size 5002471642
+oid sha256:438aec814a8321aa261d66cd0b8c6ac83200bf9e01b3a0d7bf6f6a64a1b3e4fe
+size 5397355962

model-00002-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9fbea4ba957977df98a1b964bea477aa332074be188f4051bebff3836aaba6bb
-size 5002856294
+oid sha256:6c19022c91fc35088913db68535effaeed5eb5188230dbeae119093b3e34b148
+size 5419737013

model-00003-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b1580c307f04fbd39f2961464907421fecdc38bcb97d4c40a47f80198e730011
-size 5003346119
+oid sha256:c68d410b0fe4576e511781416792ec3d75c67052d79c8c4695d38431e7b33dc0
+size 5371602180

model-00004-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8daa20c45a9db36d0907cbf468dd5bfe406ab6ac6635ca8aef36a94cdabca83b
-size 5003277250
+oid sha256:054ff7af1198283a9b353b29649701875dd8b556d0e0a407bfb0bfb91e64193c
+size 5498433829

model-00005-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:db3b71aa42815be8fbe01352261e9a31298352a79f700bf5fe1da9442ec73f47
-size 4418205241
+oid sha256:81430f4bcefba95761d29e6c082a7a9356c91ee151099d6bafe6ae23f41d9eb9
+size 2730497548

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()