| """Weight-only block quantization (MatMulNBits) of the fp16-storage depth graphs, streamed one tensor at a time. |
| |
| Input: an fp16-storage depth graph from export_depth_all (depth_frame_fp16.onnx, the sampled runtime graph, or |
| depth_all.onnx, the teacher-forced graph used by the quality gate). In both, each Linear weight W (N, K) is an fp16 |
| initializer used as Cast(W -> fp32) followed by either Transpose(perm 1,0) -> MatMul(s) or Gemm(transB=1). |
| Each such use becomes MatMulNBits(A, Wq, scales) with Wq quantized from the fp32 upcast of W (lossless upcast of |
| the fp16 source). Shared weights (group 8 for levels 9..15, the cross blocks for all levels) are quantized once |
| and referenced by every MatMulNBits node that used them. |
| |
| Weights matching --keep (regex on the initializer name) stay fp16-stored behind their Cast. Embeddings (Gather) |
| are left as they are. Graph I/O is unchanged. |
| |
| Never holds more than one weight (plus its quantized copy) in memory: initializers are read from the source |
| external-data file by offset and written sequentially to the destination data file. |
| |
| PYTHONPATH=. .venv-tts/bin/python -m tools.maya_export.quant_depth_frame \ |
| --src ~/.cache/maya_fast/depth_build/fp16/depth_frame_fp16.onnx \ |
| --dst ~/.cache/maya_fast/depth_build/q8a1/depth_frame_q8a1.onnx --bits 8 --block 32 --accuracy-level 1 |
| |
| Note (ORT MLAS, checked on main): 8-bit MatMulNBits has native kernels only for accuracy_level 4 (int8 compute); |
| accuracy_level 1 runs the generic path on every CPU: dequantize B to an fp32 scratch buffer on each call, then |
| SGEMM. That path is the same C++ code on x86 and ARM. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import collections |
| import os |
| import re |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| from onnx import TensorProto, helper |
|
|
| NP = {TensorProto.FLOAT: np.float32, TensorProto.FLOAT16: np.float16, TensorProto.INT64: np.int64, |
| TensorProto.INT32: np.int32, TensorProto.INT8: np.int8, TensorProto.UINT8: np.uint8, |
| TensorProto.BOOL: np.bool_, TensorProto.DOUBLE: np.float64} |
|
|
|
|
| def _attr(node, name, default=None): |
| for a in node.attribute: |
| if a.name == name: |
| return helper.get_attribute_value(a) |
| return default |
|
|
|
|
| def _read(t: TensorProto, src_dir: str) -> np.ndarray: |
| if t.data_location == TensorProto.EXTERNAL: |
| info = {e.key: e.value for e in t.external_data} |
| off = int(info.get("offset", 0)) |
| n = int(np.prod(list(t.dims), dtype=np.int64)) * np.dtype(NP[t.data_type]).itemsize |
| length = int(info.get("length", n)) |
| with open(os.path.join(src_dir, info["location"]), "rb") as fh: |
| fh.seek(off) |
| raw = fh.read(length) |
| return np.frombuffer(raw, dtype=NP[t.data_type]).reshape(tuple(t.dims)) |
| return onnx.numpy_helper.to_array(t) |
|
|
|
|
| def quant_block(w_kn: np.ndarray, bits: int, block: int, symmetric: bool): |
| """w_kn (K, N) fp32 (MatMul B layout). Returns packed (N, k_blocks, blob) uint8, scales (N, k_blocks) fp32, |
| zero points (N, ceil(k_blocks / kpack)) uint8. Same C++ kernels as onnxruntime's MatMulNBitsQuantizer.""" |
| from onnxruntime.capi._pybind_state import quantize_matmul_4bits, quantize_matmul_8bits |
|
|
| K, N = w_kn.shape |
| kpack = 8 // bits |
| k_blocks = (K + block - 1) // block |
| blob = (block + kpack - 1) // kpack |
| pad = k_blocks * block - K |
| if pad: |
| w_kn = np.pad(w_kn, ((0, pad), (0, 0))) |
| w_kn = np.ascontiguousarray(w_kn, dtype=np.float32) |
| packed = np.zeros((N, k_blocks, blob), dtype=np.uint8) |
| zp = np.zeros((N, (k_blocks + kpack - 1) // kpack), dtype=np.uint8) |
| scales = np.zeros((N, k_blocks), dtype=np.float32) |
| fn = quantize_matmul_8bits if bits == 8 else quantize_matmul_4bits |
| fn(packed, w_kn, scales, zp, block, N, K, symmetric) |
| return packed, scales, zp |
|
|
|
|
| def plan(model: onnx.ModelProto, keep: re.Pattern | None): |
| """Find the quantizable weights. Returns {init_name: [(kind, node, extra)]} and the set of node names to drop.""" |
| g = model.graph |
| inits = {t.name: t for t in g.initializer} |
| cons = collections.defaultdict(list) |
| for n in g.node: |
| for i in n.input: |
| cons[i].append(n) |
| out = {} |
| unsupported = set() |
| drop = set() |
| for n in g.node: |
| if n.op_type != "Cast" or n.input[0] not in inits: |
| continue |
| t = inits[n.input[0]] |
| if t.data_type != TensorProto.FLOAT16 or len(t.dims) != 2 or _attr(n, "to") != TensorProto.FLOAT: |
| continue |
| if keep is not None and keep.search(t.name): |
| continue |
| uses, ok, local_drop = [], True, {n.name} |
| for c in cons[n.output[0]]: |
| if c.op_type == "Transpose" and list(_attr(c, "perm", [])) == [1, 0]: |
| local_drop.add(c.name) |
| for mm in cons[c.output[0]]: |
| if mm.op_type == "MatMul" and mm.input[1] == c.output[0]: |
| uses.append(("matmul", mm)) |
| else: |
| ok = False |
| elif (c.op_type == "Gemm" and c.input[1] == n.output[0] and _attr(c, "transB", 0) == 1 |
| and _attr(c, "transA", 0) == 0 and _attr(c, "alpha", 1.0) == 1.0 and _attr(c, "beta", 1.0) == 1.0): |
| uses.append(("gemm", c)) |
| else: |
| ok = False |
| if ok and uses: |
| out.setdefault(t.name, []).extend(uses) |
| drop |= local_drop |
| else: |
| unsupported.add(t.name) |
| mixed = sorted(set(out) & unsupported) |
| if mixed: |
| raise ValueError(f"weights with both convertible and unconvertible Cast uses: {mixed}") |
| return out, drop |
|
|
|
|
| def convert(src: Path, dst: Path, bits: int, block: int, accuracy_level: int, symmetric: bool, |
| keep: str | None) -> dict: |
| model = onnx.load(str(src), load_external_data=False) |
| g = model.graph |
| keep_re = re.compile(keep) if keep else None |
| targets, drop_nodes = plan(model, keep_re) |
| replaced = {} |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| loc = dst.name + ".data" |
| data_path = dst.with_name(loc) |
| data_path.unlink(missing_ok=True) |
| src_dir = str(src.parent) |
| new_inits, pos = [], 0 |
| stats = collections.Counter() |
| used = set() |
| for n in g.node: |
| if n.name not in drop_nodes: |
| used.update(n.input) |
|
|
| def add_tensor(fh, name, arr, dtype): |
| nonlocal pos |
| raw = np.ascontiguousarray(arr).tobytes() |
| t = TensorProto() |
| t.name, t.data_type = name, dtype |
| t.dims.extend(arr.shape) |
| if len(raw) < 1024: |
| t.raw_data = raw |
| else: |
| fh.write(raw) |
| for k, v in (("location", loc), ("offset", str(pos)), ("length", str(len(raw)))): |
| e = t.external_data.add() |
| e.key, e.value = k, v |
| t.data_location = TensorProto.EXTERNAL |
| pos += len(raw) |
| new_inits.append(t) |
|
|
| with open(data_path, "wb") as fh: |
| for t in g.initializer: |
| if t.name in targets: |
| w = _read(t, src_dir).astype(np.float32) |
| N, K = w.shape |
| packed, scales, zp = quant_block(w.T, bits, block, symmetric) |
| del w |
| qn, sn, zn = t.name + f"_Q{bits}", t.name + "_scales", t.name + "_zp" |
| add_tensor(fh, qn, packed, TensorProto.UINT8) |
| add_tensor(fh, sn, scales, TensorProto.FLOAT) |
| ins = [qn, sn] |
| if not symmetric: |
| add_tensor(fh, zn, zp, TensorProto.UINT8) |
| ins.append(zn) |
| stats["quantized_weights"] += 1 |
| stats["quantized_params"] += N * K |
| for kind, node in targets[t.name]: |
| attrs = dict(K=K, N=N, bits=bits, block_size=block) |
| if accuracy_level: |
| attrs["accuracy_level"] = accuracy_level |
| nodes = [] |
| if kind == "gemm" and len(node.input) > 2 and node.input[2]: |
| mid = node.output[0] + "_nbits" |
| nodes.append(helper.make_node("MatMulNBits", [node.input[0]] + ins, [mid], |
| name=node.name + "_nbits", domain="com.microsoft", **attrs)) |
| nodes.append(helper.make_node("Add", [mid, node.input[2]], [node.output[0]], |
| name=node.name + "_bias")) |
| else: |
| nodes.append(helper.make_node("MatMulNBits", [node.input[0]] + ins, [node.output[0]], |
| name=node.name + "_nbits", domain="com.microsoft", **attrs)) |
| replaced[node.name] = nodes |
| stats["matmulnbits_nodes"] += 1 |
| continue |
| if t.name not in used: |
| stats["dropped_inits"] += 1 |
| continue |
| arr = _read(t, src_dir) |
| if t.data_type == TensorProto.FLOAT16 and arr.ndim == 2: |
| stats["kept_fp16_params"] += arr.size |
| add_tensor(fh, t.name, arr, t.data_type) |
| nodes = [] |
| for n in g.node: |
| if n.name in drop_nodes: |
| continue |
| nodes.extend(replaced.get(n.name, [n])) |
| del g.node[:] |
| g.node.extend(nodes) |
| del g.initializer[:] |
| g.initializer.extend(new_inits) |
| names = {t.name for t in new_inits} |
| keep_inputs = [i for i in g.input if i.name not in names or i.name in {"h", "cb0", "mem", "noise", "codes"}] |
| del g.input[:] |
| g.input.extend(keep_inputs) |
| if not any(o.domain == "com.microsoft" for o in model.opset_import): |
| model.opset_import.append(helper.make_opsetid("com.microsoft", 1)) |
| onnx.save_model(model, str(dst)) |
| stats["data_bytes"] = pos |
| return dict(stats) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--src", type=Path, required=True) |
| ap.add_argument("--dst", type=Path, required=True) |
| ap.add_argument("--bits", type=int, choices=(4, 8), default=8) |
| ap.add_argument("--block", type=int, default=32) |
| ap.add_argument("--accuracy-level", type=int, default=1) |
| ap.add_argument("--asym", action="store_true") |
| ap.add_argument("--keep", default=None, help="regex: initializers kept fp16-stored (e.g. 'heads\\.|cross\\.')") |
| a = ap.parse_args() |
| res = convert(a.src.expanduser(), a.dst.expanduser(), a.bits, a.block, a.accuracy_level, not a.asym, a.keep) |
| print(res, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
