MTP head graft for MLX Qwen3.8-27B quantizations that dropped theirs

README.md

@@ -0,0 +1,168 @@
+---
+license: apache-2.0
+library_name: mlx
+tags:
+  - mlx
+  - speculative-decoding
+  - multi-token-prediction
+  - apple-silicon
+  - qwen3.8
+base_model:
+  - Qwen/Qwen3.8-27B
+---
+
+# MTP graft — restoring Multi-Token-Prediction heads to MLX quantizations that dropped them
+
+Most MLX quantizations of Qwen3.5/3.6/3.8 keep `mtp_num_hidden_layers: 1` in
+`config.json` **while silently dropping the MTP weights**. The checkpoint loads,
+generates correctly, and gives you no reason to suspect anything — it just
+cannot use Lightning MTP in oMLX, and nothing tells you why.
+
+This repo contains the **MTP head only** (314 MB), plus a tool to graft it onto a
+recipient checkpoint. No base weights are redistributed.
+
+Measured on an Apple **M5 Max, 128 GB**, macOS 26.4.1, oMLX 0.6.3rc1.
+
+## Result
+
+Recipient: `Qwen3.8-27B` abliterated with [Heretic](https://github.com/p-e-w/heretic)
+(`trohrbaugh/Qwen3.8-27B-heretic-ara`, MLX affine 4-bit g64), which ships with
+**0 MTP tensors**. Served through oMLX with `mtp_enabled`, greedy, thinking off
+via `chat_template_kwargs`, A/B/A interleaved:
+
+| prompt | before (no MTP) | after graft | speedup |
+|---|---:|---:|---:|
+| Python LRU cache | 29.0 / 29.9 tok/s | **61.1 / 67.1** | **2.18×** |
+| Chinese prose (200 chars) | 29.4 tok/s | **41.6 / 40.3** | **1.39×** |
+
+For scale, the same machine and settings on a checkpoint that *shipped* with its
+MTP heads intact (`tozp/…-OBLITERATED-V2`) gives 63.3 / 39–41 tok/s — i.e. the
+graft lands the recipient in the same place, and the MTP multiplier itself is
+the same (2.15× vs 2.18×).
+
+**MTP is strongly content-dependent.** Code roughly doubles; Chinese prose gains
+~40%. That matches what a separate speculative-decoding stack (DFlash2 block
+diffusion) does on the same machine — two unrelated mechanisms hitting the same
+wall on CJK, which suggests the cause is CJK token entropy rather than any one
+draft head being undertrained.
+
+## Output fidelity — read this
+
+Speculative decoding is verify-then-accept: the head only *proposes*, the base
+model verifies. It cannot make the model say something the base model would
+reject. But batched K-token verification and single-token decoding take
+different floating-point paths, so **near-ties can flip at temperature 0**.
+
+Byte-for-byte comparison, fixed seed, `temperature=0, top_p=1, top_k=0`:
+
+| content | tokens (before / after) | identical | similarity |
+|---|---|---|---:|
+| Python code | 165 / 165 | **yes** | 1.000 |
+| Arithmetic | 3 / 3 | **yes** | 1.000 |
+| Chinese explanation | 70 / 70 | no — 1 token | 0.992 |
+| Chinese creative prose | 192 / 217 | no — diverges at ~50% | 0.663 |
+
+The single-token flip looks like this:
+
+```
+before: …因為它假設近期被使用過的資料[[，未來被再次使用的機率較高。]]
+after : …因為它假設近期被使用過的資料[[在未來被再次使用的機率較高。]]
+```
+
+Both are fluent; neither is better. Low-entropy text has no near-ties, so it
+reproduces exactly. Creative text is full of them, and once one token flips
+everything after it differs. **If you need bit-reproducible output, do not
+enable MTP.**
+
+## Known incompatibility: `mlx_lm` breaks on MTP tensors
+
+A checkpoint containing MTP tensors produces **garbage** under `mlx_lm.load`,
+silently — no error, no warning. Controlled test, same checkpoint, only the 46
+MTP tensors differing:
+
+```
+BigBang-v1-heretic-MLX-4bit      (MTP = 0)   -> normal Chinese output
+BigBang-v1-heretic-MLX-4bit-mtp  (MTP = 46)  -> ";av L WL two one lya"**开始中j位ode理ments谷雨谷…"
+```
+
+`mlx_vlm` and oMLX both handle the same checkpoints correctly. So:
+
+* **oMLX** — works, and is the only stack that actually *uses* the head.
+* **mlx_vlm** — works (ignores the head).
+* **mlx_lm** — **broken**. Keep an ungrafted copy if you need this path.
+
+This bit us: we first blamed vision tensors for the same garbage on a VL
+checkpoint that happened to have both. It is the MTP tensors.
+
+## Usage
+
+```bash
+python graft_mtp.py \
+  --recipient /path/to/model-without-mtp \
+  --donor     /path/to/donor-or-standalone-head \
+  --out       /path/to/model-with-mtp
+```
+
+The recipient's weights are **hard-linked**, not copied — the graft costs only
+the MTP shard on disk. The original directory is never modified.
+
+The tool refuses to graft across mismatched `hidden_size`, head counts,
+`vocab_size`, `num_experts`, `moe_intermediate_size`, or a different base
+quantization scheme (`--force` overrides). Two donor layouts are handled: heads
+already namespaced under `language_model.mtp.`, and repos that *are* just the
+head (`fc`/`layers.0`/`norm`/`pre_fc_norm_*`), such as
+[`mlx-community/Qwen3.6-35B-A3B-MTP-4bit`](https://huggingface.co/mlx-community/Qwen3.6-35B-A3B-MTP-4bit).
+
+Then in oMLX: set `mtp_enabled` for the model **and unload/reload it**. oMLX
+reads model settings at load time only — editing `model_settings.json` while the
+model is resident has no effect whatsoever (verified twice).
+
+### Finding recipients
+
+Check the weights, not the config — the config lies:
+
+```python
+import json
+wm = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"]
+print(sum(1 for k in wm if "mtp" in k.lower() or "nextn" in k.lower()))
+```
+
+Of eight Qwen3.5/3.6/3.8 MLX checkpoints on hand, **four** declared
+`mtp_num_hidden_layers: 1` with zero MTP tensors present.
+
+## Provenance
+
+The head in this repo is **Qwen's own**, not something trained here:
+
+1. `Qwen/Qwen3.8-27B` — original MTP heads
+2. `tozp/Qwen3.8-27B-OBLITERATED-V2-mlx-4bit` — MLX affine 4-bit g64; its
+   `abliteration_metadata.json` records *"Restore MTP + vision tensors from
+   stock"*, so the heads it carries are unmodified
+3. this repo — those 29 tensors, extracted unchanged
+
+29 tensors, 314.3 MB, `sha256[:16] = 505827ce17922d2d`.
+
+Applies to **`qwen3_5` dense, hidden 5120** (Qwen3.8-27B and derivatives).
+For `qwen3_5_moe` (hidden 2048, 256 experts), use
+`mlx-community/Qwen3.6-35B-A3B-MTP-4bit` as donor instead — same tool, 46
+tensors.
+
+## Limitations
+
+* One machine (M5 Max), one serving stack (oMLX), the prompt set above. No
+  long-context, no batch >1, no acceptance-rate instrumentation — oMLX does not
+  expose one.
+* Grafting the stock head onto an **abliterated** model is not obviously sound in
+  theory. The empirical argument: the recipient's KL to stock is 0.0535 (Heretic
+  reports it), while a checkpoint that ships this head natively sits at 0.9624 —
+  the recipient is ~18× closer to stock than a model where the head demonstrably
+  works. And the failure mode is bounded: a mismatched head lowers acceptance,
+  it cannot corrupt output.
+* Not tested with vision inputs on VL recipients.
+
+## Credits
+
+- [Qwen](https://huggingface.co/Qwen/Qwen3.8-27B) — the MTP heads
+- [tozp](https://huggingface.co/tozp/Qwen3.8-27B-OBLITERATED-V2-mlx-4bit) — the 4-bit quantization they were extracted from
+- [Heretic](https://github.com/p-e-w/heretic) / [trohrbaugh](https://huggingface.co/trohrbaugh/Qwen3.8-27B-heretic-ara) — the recipient
+- [oMLX](https://omlx.ai) — Lightning MTP

graft_mtp.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+"""Graft Multi-Token-Prediction heads onto an MLX checkpoint that lost them.
+
+Many MLX quantizations of Qwen3.5/3.6/3.8 keep `mtp_num_hidden_layers: 1` in
+config.json while silently dropping the MTP weights, because the standard
+conversion path only walks the main decoder. The result loads fine and runs
+fine -- it just cannot use Lightning MTP in oMLX, and nothing tells you why.
+
+This copies the MTP subtree from a donor that still has it, into a recipient
+that does not. The recipient's own weights are never modified: they are
+hard-linked into a new directory and the MTP tensors are added as one extra
+shard.
+
+Why this is safe: speculative decoding is verify-then-accept. The MTP head
+only *proposes* tokens; the base model verifies every one of them. A poorly
+matched head lowers the acceptance rate -- it costs throughput, it cannot
+corrupt output. The worst outcome is no speedup.
+
+Two caveats that are real and measured, see README.
+"""
+import argparse, json, os, shutil, sys
+import mlx.core as mx
+
+MTP_PREFIX = "language_model.mtp."
+
+
+def _load_all(d):
+    idx = os.path.join(d, "model.safetensors.index.json")
+    if os.path.exists(idx):
+        wm = json.load(open(idx))["weight_map"]
+        return wm, sorted(set(wm.values()))
+    files = sorted(f for f in os.listdir(d) if f.endswith(".safetensors"))
+    wm = {}
+    for f in files:
+        for k in mx.load(os.path.join(d, f)).keys():
+            wm[k] = f
+    return wm, files
+
+
+def collect_mtp(donor):
+    """Pull the MTP subtree out of `donor`, normalised to MTP_PREFIX.
+
+    Handles both layouts seen in the wild:
+      * in-model   -- keys already start with `language_model.mtp.`
+      * standalone -- the repo IS just the head (fc/layers.0/norm/pre_fc_*),
+                      e.g. mlx-community/Qwen3.6-35B-A3B-MTP-4bit
+    """
+    wm, files = _load_all(donor)
+    inline = {k: k for k in wm if k.startswith(MTP_PREFIX)}
+    if inline:
+        keys = inline
+    else:
+        roots = ("fc.", "layers.", "norm.", "pre_fc_norm_embedding.",
+                 "pre_fc_norm_hidden.")
+        keys = {k: MTP_PREFIX + k for k in wm if k.startswith(roots)}
+        if not keys:
+            sys.exit(f"no MTP tensors found in {donor}")
+    out = {}
+    for f in files:
+        w = mx.load(os.path.join(donor, f))
+        for src, dst in keys.items():
+            if wm[src] == f and src in w:
+                out[dst] = w[src]
+    if len(out) != len(keys):
+        sys.exit(f"expected {len(keys)} MTP tensors, recovered {len(out)}")
+    return out
+
+
+def check(recipient_cfg, donor_cfg):
+    """Refuse mismatched architectures. These are the fields the MTP head's
+    shapes actually depend on; a mismatch means the graft cannot work."""
+    r = recipient_cfg.get("text_config", recipient_cfg)
+    d = donor_cfg.get("text_config", donor_cfg)
+    bad = []
+    for k in ("hidden_size", "num_attention_heads", "num_key_value_heads",
+              "head_dim", "vocab_size", "num_experts", "moe_intermediate_size"):
+        if k in r and k in d and r[k] != d[k]:
+            bad.append(f"{k}: recipient={r[k]} donor={d[k]}")
+    rq = (recipient_cfg.get("quantization") or {})
+    dq = (donor_cfg.get("quantization") or {})
+    for k in ("bits", "group_size", "mode"):
+        # Per-tensor overrides are fine; only the base scheme has to agree.
+        if rq.get(k) != dq.get(k):
+            bad.append(f"quantization.{k}: recipient={rq.get(k)} donor={dq.get(k)}")
+    return bad
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--recipient", required=True, help="model missing its MTP head")
+    ap.add_argument("--donor", required=True, help="model (or standalone head) that has one")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--force", action="store_true", help="graft despite arch mismatch")
+    a = ap.parse_args()
+
+    rc = json.load(open(os.path.join(a.recipient, "config.json")))
+    dc = json.load(open(os.path.join(a.donor, "config.json")))
+    bad = check(rc, dc)
+    if bad:
+        print("architecture mismatch:")
+        for b in bad:
+            print("  " + b)
+        if not a.force:
+            sys.exit("refusing to graft (use --force to override)")
+
+    if os.path.exists(a.out):
+        shutil.rmtree(a.out)
+    os.makedirs(a.out)
+    # Hard-link the weights: the graft costs only the MTP shard on disk, and
+    # the recipient's files are shared, not copied. Small files are copied so
+    # a later in-place config edit cannot corrupt the original -- a real bug
+    # we hit with a tool that hard-linked config.json then rewrote it.
+    for f in os.listdir(a.recipient):
+        s, d = os.path.join(a.recipient, f), os.path.join(a.out, f)
+        if not os.path.isfile(s):
+            continue
+        (os.link if f.endswith(".safetensors") else shutil.copy2)(s, d)
+
+    mtp = collect_mtp(a.donor)
+    shard = "model-mtp-00001-of-00001.safetensors"
+    mx.save_safetensors(os.path.join(a.out, shard), mtp, metadata={"format": "mlx"})
+    size = os.path.getsize(os.path.join(a.out, shard))
+
+    ridx = os.path.join(a.recipient, "model.safetensors.index.json")
+    if os.path.exists(ridx):
+        base = json.load(open(ridx))
+    else:
+        wm, _ = _load_all(a.recipient)
+        base = {"metadata": {"total_size": 0}, "weight_map": wm}
+    wm = dict(base["weight_map"])
+    for k in mtp:
+        wm[k] = shard
+    meta = dict(base.get("metadata", {}))
+    meta["total_size"] = int(meta.get("total_size", 0)) + size
+    json.dump({"metadata": meta, "weight_map": wm},
+              open(os.path.join(a.out, "model.safetensors.index.json"), "w"), indent=2)
+
+    cfg = json.load(open(os.path.join(a.out, "config.json")))
+    tgt = cfg.get("text_config", cfg)
+    if not (cfg.get("mtp_num_hidden_layers") or tgt.get("mtp_num_hidden_layers")):
+        cfg["mtp_num_hidden_layers"] = 1
+        json.dump(cfg, open(os.path.join(a.out, "config.json"), "w"), indent=1)
+        print("  set mtp_num_hidden_layers = 1")
+
+    print(f"grafted {len(mtp)} MTP tensors ({size/1e6:.1f} MB) -> {a.out}")
+    print(f"  tensors: {len(base['weight_map'])} -> {len(wm)}")
+
+
+if __name__ == "__main__":
+    main()

qwen3.8-27b-mtp-4bit.safetensors

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