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-oQ4-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-00005.safetensors → model-00001-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d772f04e18e4705dcaa317c4e054eaeb8af9bdb184e4b6a8fcfc98b55b98bc8a
-size 5005868403
+oid sha256:c9b6651178c7eafbf36afca441f6015fe237b20af0b8c7b66ff16ed79e253ae3
+size 5445093729

model-00002-of-00005.safetensors → model-00002-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:132022bd8c275d6901bebc7156cffbe7b10fda558590f21f66a95a26565f281d
-size 5003882154
+oid sha256:b0051c2108c949a365edec5ecd6ab64feb17286ed74c09b5ca8144b9861ce27e
+size 5372082282

model-00003-of-00005.safetensors → model-00003-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bf933742439697a5a635b954692431e8fbe9efdf7256f3819876edbac30ecda0
-size 5003483080
+oid sha256:1225ab315b23a821dd68f724d92c1fecc77adc28f5b377ee1c6e225e5c5c3c89
+size 5389445564

model-00004-of-00005.safetensors → model-00004-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:142bb0b29d53912fb71efdefb5bb24f6dc0d62634085bc2522c4453fa6ba80d2
-size 5003348234
+oid sha256:8cc12d547e07abb4e59b06768b66ee82b875a9f6cbc4c82c85618bb436d85e07
+size 4162092468

model-00005-of-00005.safetensors

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

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