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