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