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