| |
| |
| |
| """The .trq container format (spec section 10) — a real binary codec. |
| |
| Layout (all integers little-endian): |
| magic 4 bytes b'TRQ1' |
| header_len u32 |
| header JSON {version, arch, n_tensors, provenance, roll_meta, created} |
| per tensor: |
| name_len u16 |
| name bytes |
| ndim u8 |
| shape ndim x u64 |
| dtype_code u8 0=fp32 1=fp16 2=bf16 3=int8 4=int4-packed |
| quant_code u8 0=none 1=int8-asym 2=int4-group32 3=fp16-cast 4=bf16-cast |
| meta_len u32 |
| meta JSON {quantization metadata: scales/zeros as base64 f64, |
| roll metadata, codebooks} |
| data_len u64 |
| data bytes |
| checksum 32 bytes sha256(data) |
| footer 32 bytes sha256(all preceding bytes) |
| |
| read_trq() verifies the file checksum and every per-tensor checksum and |
| raises on any mismatch. Corrupt files never load silently. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import base64 |
| import hashlib |
| import json |
| import struct |
|
|
| import numpy as np |
|
|
| from .util import utc_now_iso |
|
|
| MAGIC = b"TRQ1" |
| VERSION = 1 |
|
|
| _DTYPE_CODES = {"fp32": 0, "fp16": 1, "bf16": 2, "int8": 3, "int4": 4} |
| _DTYPE_NAMES = {v: k for k, v in _DTYPE_CODES.items()} |
| _NP_DTYPES = {"fp32": np.float32, "fp16": np.float16, "bf16": np.uint16, |
| "int8": np.int8, "int4": np.uint8} |
| _QUANT_CODES = {"none": 0, "int8-asym": 1, "int4-group32": 2, |
| "fp16-cast": 3, "bf16-cast": 4} |
| _QUANT_NAMES = {v: k for k, v in _QUANT_CODES.items()} |
|
|
|
|
| def _b64(arr) -> str: |
| return base64.b64encode(np.asarray(arr, dtype=np.float64).tobytes()).decode() |
|
|
|
|
| def _unb64(s: str) -> np.ndarray: |
| return np.frombuffer(base64.b64decode(s.encode()), dtype=np.float64) |
|
|
|
|
| def _meta_to_json(meta: dict) -> dict: |
| m = dict(meta) |
| for k, v in list(m.items()): |
| if isinstance(v, np.ndarray): |
| |
| m[k] = {"__b64_f64__": _b64(np.asarray(v, dtype=np.float64)), |
| "__shape__": list(v.shape), "__key__": k} |
| return m |
|
|
|
|
| def _meta_from_json(m: dict) -> dict: |
| m = dict(m) |
| for k, v in list(m.items()): |
| if isinstance(v, dict) and "__b64_f64__" in v: |
| arr = _unb64(v["__b64_f64__"]) |
| if "__shape__" in v: |
| arr = arr.reshape(v["__shape__"]) |
| m[k] = arr |
| return m |
|
|
|
|
| class TrqError(ValueError): |
| pass |
|
|
|
|
| def write_trq(path: str, *, arch: dict, tensors: dict, provenance: dict, |
| roll_meta: dict = None) -> dict: |
| """tensors: {name: (payload_array, fmt_name, quant_meta_dict)}. |
| Returns {"path":..., "bytes":..., "sha256":...}.""" |
| buf = bytearray() |
| buf += MAGIC |
| header = { |
| "version": VERSION, |
| "arch": arch, |
| "n_tensors": len(tensors), |
| "provenance": provenance, |
| "roll_meta": roll_meta or {}, |
| "created": utc_now_iso(), |
| } |
| hbytes = json.dumps(header).encode() |
| buf += struct.pack("<I", len(hbytes)) + hbytes |
|
|
| for name, (payload, fmt, qmeta) in tensors.items(): |
| payload = np.ascontiguousarray(payload) |
| nbytes = payload.tobytes() |
| name_b = name.encode() |
| buf += struct.pack("<H", len(name_b)) + name_b |
| buf += struct.pack("<B", payload.ndim) |
| buf += struct.pack(f"<{payload.ndim}Q", *payload.shape) |
| buf += struct.pack("<B", _DTYPE_CODES[fmt]) |
| quant_name = {"fp32": "none", "fp16": "fp16-cast", "bf16": "bf16-cast", |
| "int8": "int8-asym", "int4": "int4-group32"}[fmt] |
| buf += struct.pack("<B", _QUANT_CODES[quant_name]) |
| meta = _meta_to_json({"fmt": fmt, **qmeta, |
| "roll": (roll_meta or {}).get(name, {})}) |
| mbytes = json.dumps(meta).encode() |
| buf += struct.pack("<I", len(mbytes)) + mbytes |
| buf += struct.pack("<Q", len(nbytes)) + nbytes |
| buf += hashlib.sha256(nbytes).digest() |
|
|
| digest = hashlib.sha256(bytes(buf)).digest() |
| buf += digest |
| with open(path, "wb") as f: |
| f.write(bytes(buf)) |
| return {"path": path, "bytes": len(buf), "sha256": digest.hex()} |
|
|
|
|
| def read_trq(path: str): |
| """Returns (header, tensors) with all checksums verified. |
| tensors: {name: (payload_array, meta_dict)}. Raises TrqError on corruption.""" |
| with open(path, "rb") as f: |
| blob = f.read() |
| if len(blob) < 4 + 32 or blob[:4] != MAGIC: |
| raise TrqError("not a .trq file (bad magic)") |
| file_sum, body = blob[-32:], blob[:-32] |
| if hashlib.sha256(body).digest() != file_sum: |
| raise TrqError("file checksum mismatch — refusing to load") |
| off = 4 |
| hlen = struct.unpack_from("<I", body, off)[0] |
| off += 4 |
| header = json.loads(body[off:off + hlen]) |
| off += hlen |
| tensors = {} |
| for _ in range(header["n_tensors"]): |
| nlen = struct.unpack_from("<H", body, off)[0] |
| off += 2 |
| name = body[off:off + nlen].decode() |
| off += nlen |
| ndim = struct.unpack_from("<B", body, off)[0] |
| off += 1 |
| shape = struct.unpack_from(f"<{ndim}Q", body, off) |
| off += 8 * ndim |
| dtype_code = struct.unpack_from("<B", body, off)[0] |
| off += 1 |
| quant_code = struct.unpack_from("<B", body, off)[0] |
| off += 1 |
| mlen = struct.unpack_from("<I", body, off)[0] |
| off += 4 |
| meta = _meta_from_json(json.loads(body[off:off + mlen])) |
| off += mlen |
| dlen = struct.unpack_from("<Q", body, off)[0] |
| off += 8 |
| data = body[off:off + dlen] |
| off += dlen |
| cksum = body[off:off + 32] |
| off += 32 |
| if hashlib.sha256(data).digest() != cksum: |
| raise TrqError(f"tensor {name!r} checksum mismatch — refusing to load") |
| fmt = _DTYPE_NAMES[dtype_code] |
| arr = np.frombuffer(data, dtype=_NP_DTYPES[fmt]) |
| if fmt == "int4": |
| arr = arr.reshape(-1) |
| else: |
| arr = arr.reshape(shape) |
| tensors[name] = (arr, meta) |
| return header, tensors |
|
|
|
|
| def total_params(tensors: dict) -> int: |
| """Total logical parameters across tensors (uses recorded shapes, so |
| packed INT4 counts its true parameters, not its packed bytes).""" |
| total = 0 |
| for _name, (arr, meta) in tensors.items(): |
| shape = meta.get("shape") or list(arr.shape) |
| total += int(np.prod(shape)) |
| return total |
|
|
|
|
| def param_budget_ok(tensors: dict, cap: float = 8_500_000_000) -> tuple: |
| """Verify the exported artifact obeys the hard parameter budget. |
| Returns (total_params, ok).""" |
| total = total_params(tensors) |
| return total, total <= cap |
|
|
