File size: 7,370 Bytes
7bf0e3b 28fd117 7bf0e3b 28fd117 7bf0e3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Export the no-rotation experimental trits to a standalone MLX checkpoint.
Reads only the compressed files, head.pt, config and tokenizer. No base model
weights, PEFT weights, dequantized backbone or training/evaluation data.
"""
import argparse
import hashlib
import json
import shutil
from pathlib import Path
import mlx.core as mx
import numpy as np
import torch
from kev.packed_model import FORMAT
def pack_native(codes, bits):
"""Unsigned integer codes -> MLX little-endian packed uint32 (continuous bit stream; 2/3/4/8 bits), without quantization."""
if codes.ndim != 2 or (codes.shape[1] * bits) % 32 or np.any(codes >= 1 << bits):
raise ValueError("invalid quantized codes")
c = codes.astype(np.uint64)
out = np.zeros((c.shape[0], c.shape[1] * bits // 32), dtype=np.uint64)
for j in range(c.shape[1]):
word, off = divmod(j * bits, 32)
out[:, word] |= (c[:, j] << np.uint64(off)) & np.uint64(0xFFFFFFFF)
if off + bits > 32:
out[:, word + 1] |= c[:, j] >> np.uint64(32 - off)
return out.astype(np.uint32)
def sha256(path):
with Path(path).open("rb") as source:
return hashlib.file_digest(source, "sha256").hexdigest()
def build(packed, checkpoint, config, output):
packed, checkpoint, config, output = map(Path, (packed, checkpoint, config, output))
if output.exists():
raise FileExistsError(f"output directory already exists: {output}")
if not (checkpoint / "tokenizer.json").is_file():
raise ValueError("source checkpoint must carry its tokenizer.json")
tensors, meta = mx.load(str(packed), return_metadata=True)
if meta.get("format") != "kev-ternary-v1" or meta.get("rot") != "none":
raise ValueError("standalone native export requires kev-ternary-v1 with rot=none")
if any(k.endswith(".perm") for k in tensors):
raise ValueError("activation-order packed tensors are not supported")
kinds, groups = json.loads(meta["kinds"]), json.loads(meta["groups"])
emb_bits = int(meta["emb_bits"])
if emb_bits not in (2, 4, 8):
raise ValueError("the standalone export requires a quantized embedding")
head = torch.load(checkpoint / "head.pt", map_location="cpu", weights_only=True)
if head.get("option_isolation") or head.get("special_embeddings"):
raise ValueError("unsupported checkpoint option isolation or token embeddings")
base_config = json.loads(config.read_text(encoding="utf-8"))
source_run = Path(meta["run"])
if not (source_run / "head.pt").exists() or sha256(source_run / "head.pt") != sha256(checkpoint / "head.pt"):
raise ValueError("head.pt does not match the packed file's source checkpoint")
text = base_config.get("text_config", base_config)
if text.get("model_type") != "qwen3_5_text" or text.get("hidden_size") != head["head"]["q.weight"].shape[1]:
raise ValueError("config does not match the Qwen3.5 pointer head")
cloq_path = packed.with_name(packed.stem + ".cloq.safetensors")
cloq = mx.load(str(cloq_path)) if cloq_path.exists() else {}
weights, linears = {}, {}
for pname, kind in kinds.items():
name = pname.removesuffix(".weight")
group = groups[pname]
if kind.startswith("mlx"):
bits = int(kind[3:])
weights[name + ".weight"] = tensors.pop(pname + ".q")
weights[name + ".scales"] = tensors.pop(pname + ".scales")
weights[name + ".biases"] = tensors.pop(pname + ".biases")
else:
n, k = map(int, np.asarray(tensors.pop(pname + ".shape")))
scales = tensors.pop(pname + ".scales")
if kind == "tern":
bits = 2
raw = np.asarray(tensors.pop(pname + ".trits"))
if np.any(raw > 242):
raise ValueError("invalid base-3 packed byte")
codes = ((raw[:, None].astype(np.uint16) // np.array([1, 3, 9, 27, 81], dtype=np.uint16)) % 3).reshape(-1)[:n * k].astype(np.uint8).reshape(n, k)
weights[name + ".scales"] = scales
weights[name + ".biases"] = -scales
elif kind in ("affine4", "affine8"):
bits = int(kind[6:])
raw = np.asarray(tensors.pop(pname + ".codes"))
per = 8 // bits
codes = ((raw[:, None] >> (np.arange(per, dtype=np.uint8) * bits)) & ((1 << bits) - 1)).reshape(-1)[:n * k].reshape(n, k)
weights[name + ".scales"] = scales[:, :, 0]
weights[name + ".biases"] = scales[:, :, 1]
else:
raise ValueError(f"unsupported packed kind {kind}")
weights[name + ".weight"] = mx.array(pack_native(codes, bits))
mx.eval(weights[name + ".weight"], weights[name + ".scales"], weights[name + ".biases"])
del codes, raw, scales
rank = 0
if pname + ".lora_a" in cloq:
a, b = cloq.pop(pname + ".lora_a"), cloq.pop(pname + ".lora_b")
rank = a.shape[0]
weights[name + ".comp_a"], weights[name + ".comp_b"] = a, b
linears[name] = {"bits": bits, "group_size": group, "rank": rank}
if cloq:
raise ValueError("unused compensation tensors")
emb = "language_model.model.embed_tokens"
for source, target in (("q", "weight"), ("scales", "scales"), ("biases", "biases")):
weights[emb + "." + target] = tensors.pop(emb + ".weight." + source)
weights.update(tensors)
del tensors, cloq
output.mkdir(parents=True, exist_ok=False)
(output / "config.json").write_text(json.dumps(base_config, indent=2) + "\n", encoding="utf-8")
for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "added_tokens.json", "vocab.json", "merges.txt", "chat_template.jinja"):
if (checkpoint / name).exists():
shutil.copyfile(checkpoint / name, output / name)
head.update(weights="packed", weights_dtype="bf16", lora=0)
head["compression"] = {"format": FORMAT, "source_packed_sha256": sha256(packed), "source_cloq_sha256": sha256(cloq_path) if cloq_path.exists() else None,
"arithmetic": "native affine quantized matmul plus separate fp32 low-rank compensation; not merged-bf16 parity"}
torch.save(head, output / "head.pt")
mx.save_safetensors(str(output / "model.safetensors"), weights, metadata={"format": FORMAT})
spec = {"format": FORMAT, "linears": linears, "embedding": {"bits": emb_bits, "group_size": 64},
"source": head["compression"], "model_sha256": sha256(output / "model.safetensors")}
(output / "packed.json").write_text(json.dumps(spec, indent=2) + "\n", encoding="utf-8")
print(json.dumps({"directory": str(output), "total_bytes": sum(f.stat().st_size for f in output.iterdir()), "mlx_peak_bytes": mx.get_peak_memory()}, indent=2))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--packed", required=True)
parser.add_argument("--checkpoint", required=True, help="head and tokenizer only; backbone weights are never read")
parser.add_argument("--config", required=True, help="base config.json only; no base weights")
parser.add_argument("--out", required=True)
args = parser.parse_args()
build(args.packed, args.checkpoint, args.config, args.out)
|