File size: 2,408 Bytes
ad024da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """int4 (weight-only, block-wise) quantization of the laya ONNX graph via ONNX Runtime's
MatMulNBits path — the same mechanism ORT/GenAI use for 4-bit LLM weights.
python int4_quantize.py --src laya-zh-fp32-clean.onnx --dst laya-zh-int4-b32.onnx \
--block-size 32 --symmetric
Only `MatMul` nodes whose weight input is a constant get replaced by `com.microsoft::MatMulNBits`
(attention QK^T / softmax·V keep their fp32 kernel). No calibration data needed (RTN).
"""
import argparse, collections, os, time
import onnx
from onnxruntime.quantization.matmul_nbits_quantizer import (
DefaultWeightOnlyQuantConfig, MatMulNBitsQuantizer)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--dst", required=True)
ap.add_argument("--bits", type=int, default=4)
ap.add_argument("--block-size", type=int, default=32)
ap.add_argument("--symmetric", action="store_true")
ap.add_argument("--accuracy-level", type=int, default=None)
a = ap.parse_args()
src, dst = os.path.expanduser(a.src), os.path.expanduser(a.dst)
cfg = DefaultWeightOnlyQuantConfig(block_size=a.block_size, is_symmetric=a.symmetric,
bits=a.bits, accuracy_level=a.accuracy_level)
print(f"int{a.bits} weight-only: block_size={a.block_size} symmetric={a.symmetric} "
f"accuracy_level={a.accuracy_level} | {src} -> {dst}", flush=True)
t0 = time.perf_counter()
q = MatMulNBitsQuantizer(model=src, algo_config=cfg) # first positional arg is `model`
print("config loaded, quantizing ...", flush=True)
q.process()
q.model.save_model_to_file(dst, use_external_data_format=True)
dt = time.perf_counter() - t0
m = onnx.load(dst, load_external_data=False)
ops = collections.Counter(n.op_type for n in m.graph.node)
size = os.path.getsize(dst) + (os.path.getsize(dst + ".data") if os.path.exists(dst + ".data") else 0)
src_size = os.path.getsize(src) + (os.path.getsize(src + ".data") if os.path.exists(src + ".data") else 0)
print(f"done in {dt:.1f}s | {src_size/1e6:.0f}MB -> {size/1e6:.0f}MB ({src_size/size:.1f}x smaller)")
print(f"node mix: MatMul={ops.get('MatMul',0)} MatMulNBits={ops.get('MatMulNBits',0)} "
f"Gemm={ops.get('Gemm',0)} total={sum(ops.values())}")
print("INT4_QUANTIZE_DONE")
if __name__ == "__main__":
main()
|