| """Backbone weight-only quantization sweep (MatMulNBits) with the teacher-forced gate. |
| |
| Quantizes the fp32 per-layer graphs (backbone_build/fp32) with MatMulNBitsQuantizer, merges them with the same |
| export_backbone.merge() used for wo8, and runs the existing backbone gate. Writes only under |
| work_dir()/backbone_build/<tag>/. |
| |
| Tag: b{bits}_blk{block}_{algo}_{sym|asym}_acc{level} |
| |
| Usage (from the worktree root, under tools/memguard.py): |
| python -m tools.maya_export.quant_sweep_backbone build --bits 4 --block 128 --asym --accuracy-level 1 |
| python -m tools.maya_export.quant_sweep_backbone eval --bits 4 --block 128 --asym --accuracy-level 1 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
| import onnx |
|
|
| from tools.maya_export import eval_teacher_forced as E |
| from tools.maya_export import export_backbone as EB |
| from tools.maya_export.common import work_dir |
|
|
| N_LAYERS = EB.N_LAYERS |
|
|
|
|
| def tag(bits: int, block: int, algo: str, asym: bool, acc: int) -> str: |
| if algo == "hqq": |
| return f"b{bits}_blk{block}_hqq_default" |
| return f"b{bits}_blk{block}_{algo}_{'asym' if asym else 'sym'}_acc{acc}" |
|
|
|
|
| def out_dir(name: str) -> Path: |
| d = work_dir() / "backbone_build" / name |
| d.mkdir(parents=True, exist_ok=True) |
| return d |
|
|
|
|
| def make_quantizer(model, bits: int, block: int, algo: str, asym: bool, acc: int): |
| from onnxruntime.quantization.matmul_nbits_quantizer import ( |
| DefaultWeightOnlyQuantConfig, HQQWeightOnlyQuantConfig, MatMulNBitsQuantizer) |
| if algo == "rtn": |
| |
| cfg = DefaultWeightOnlyQuantConfig(block_size=block, is_symmetric=not asym, accuracy_level=acc, bits=bits) |
| else: |
| |
| cfg = HQQWeightOnlyQuantConfig(bits=bits, block_size=block) |
| return MatMulNBitsQuantizer(model, bits=bits, block_size=block, is_symmetric=not asym, accuracy_level=acc, |
| algo_config=cfg) |
|
|
|
|
| def quantize_layer(src: Path, dst: Path, **kw) -> None: |
| model = onnx.load(str(src)) |
| q = make_quantizer(model, **kw) |
| q.process() |
| EB._save_external(q.model.model, dst) |
|
|
|
|
| def build(bits: int, block: int, algo: str, asym: bool, acc: int) -> Path: |
| name = tag(bits, block, algo, asym, acc) |
| od = out_dir(name) |
| fp32 = work_dir() / "backbone_build" / "fp32" |
| kw = dict(bits=bits, block=block, algo=algo, asym=asym, acc=acc) |
| layer_paths = [] |
| for i in range(N_LAYERS): |
| dst = od / f"layer_{i:02d}.onnx" |
| t0 = time.time() |
| quantize_layer(fp32 / f"layer_{i:02d}.onnx", dst, **kw) |
| layer_paths.append(dst) |
| print(f"layer {i:02d} {name} ok {time.time() - t0:.1f}s", flush=True) |
| out = od / f"backbone_{name}.onnx" |
| EB.merge(layer_paths, fp32 / "head.onnx", out, name) |
| print(f"wrote {out} and {out.with_suffix('.data').name}", flush=True) |
| return out |
|
|
|
|
| def evaluate(bits: int, block: int, algo: str, asym: bool, acc: int, onnx_path: Path | None = None) -> dict: |
| name = tag(bits, block, algo, asym, acc) |
| out = onnx_path or out_dir(name) / f"backbone_{name}.onnx" |
| orig = E._run_candidate |
| E._run_candidate = lambda _unused, items, c, s: orig(out, items, c, s) |
| try: |
| res = E.backbone(name) |
| finally: |
| E._run_candidate = orig |
| res.update({"tag": name, "bits": bits, "block": block, "algo": algo, "asym": asym, "accuracy_level": acc, |
| "onnx": str(out)}) |
| return res |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("mode", choices=("build", "eval")) |
| ap.add_argument("--bits", type=int, choices=(4, 8), required=True) |
| ap.add_argument("--block", type=int, choices=(32, 64, 128), required=True) |
| ap.add_argument("--algo", choices=("rtn", "hqq"), default="rtn") |
| ap.add_argument("--asym", action="store_true") |
| ap.add_argument("--accuracy-level", type=int, choices=(0, 1, 2, 3, 4), default=None, |
| help="rtn only (default 4); not supported for hqq") |
| ap.add_argument("--onnx", type=Path, default=None, help="eval an existing merged graph instead of the tag's") |
| a = ap.parse_args() |
| if a.algo == "hqq" and (a.asym or a.accuracy_level is not None): |
| ap.error("--algo hqq ignores --asym and --accuracy-level (HQQWeightOnlyQuantConfig has no such options); omit them") |
| kw = dict(bits=a.bits, block=a.block, algo=a.algo, asym=a.asym, |
| acc=4 if a.accuracy_level is None else a.accuracy_level) |
| if a.mode == "build": |
| build(**kw) |
| else: |
| print(json.dumps(evaluate(**kw, onnx_path=a.onnx), indent=2), flush=True) |
|
|
