| """Add CPU-oriented INT8 and INT4 variants to an exported artifact bundle.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import logging |
| import os |
| import shutil |
| import tempfile |
| from pathlib import Path |
|
|
| import onnxruntime as ort |
| from onnxruntime.quantization import QuantType, quantize_dynamic |
| from onnxruntime.quantization.matmul_nbits_quantizer import ( |
| DefaultWeightOnlyQuantConfig, |
| MatMulNBitsQuantizer, |
| ) |
|
|
| from .artifact import variant_metadata, write_manifest |
|
|
|
|
| DECISION_HEAD = "/openthai/head/SlotMatMul" |
| EXPECTED_BACKBONE_MATMULS = 186 |
|
|
|
|
| def _assert_runtime_contract(path: Path) -> None: |
| session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) |
| expected_inputs = { |
| "input_ids", |
| "attention_mask", |
| "answer_positions", |
| "option_counts", |
| "qtypes", |
| } |
| if {value.name for value in session.get_inputs()} != expected_inputs: |
| raise RuntimeError(f"quantized graph has an unexpected input contract: {path}") |
| if {value.name for value in session.get_outputs()} != {"logits", "probabilities"}: |
| raise RuntimeError(f"quantized graph has an unexpected output contract: {path}") |
|
|
|
|
| def _assert_operator_contract(metadata: dict, precision: str) -> None: |
| counts = metadata["operator_counts"] |
| expected = "ai.onnx::MatMulInteger" if precision == "int8" else "com.microsoft::MatMulNBits" |
| if counts.get(expected) != EXPECTED_BACKBONE_MATMULS: |
| raise RuntimeError( |
| f"{precision} graph must contain {EXPECTED_BACKBONE_MATMULS} {expected} nodes; " |
| f"found {counts.get(expected, 0)}" |
| ) |
| if counts.get("ai.onnx::MatMul") != 1: |
| raise RuntimeError(f"{precision} graph must retain exactly one FP32 decision-head MatMul") |
|
|
|
|
| def _build_int8(source: Path, destination: Path) -> None: |
| quantize_dynamic( |
| source, |
| destination, |
| op_types_to_quantize=["MatMul"], |
| per_channel=True, |
| weight_type=QuantType.QInt8, |
| nodes_to_exclude=[DECISION_HEAD], |
| use_external_data_format=True, |
| ) |
|
|
|
|
| def _build_int4(source: Path, destination: Path) -> None: |
| config = DefaultWeightOnlyQuantConfig( |
| bits=4, |
| block_size=32, |
| is_symmetric=False, |
| op_types_to_quantize=("MatMul",), |
| ) |
| quantizer = MatMulNBitsQuantizer( |
| str(source), |
| nodes_to_exclude=[DECISION_HEAD], |
| algo_config=config, |
| ) |
| quantizer.process() |
| quantizer.model.save_model_to_file(str(destination), True) |
|
|
|
|
| def _data_path(graph: Path) -> Path: |
| return Path(f"{graph}.data") |
|
|
|
|
| def _refresh_release_files(artifact: Path) -> None: |
| """Refresh packaged client/tooling and documentation from this checkout.""" |
| project_root = Path(__file__).resolve().parents[2] |
| for source_name, destination_name in ( |
| ("LICENSE", "LICENSE"), |
| ("MODEL_CARD.md", "README.md"), |
| ("MODEL_CARD.md", "MODEL_CARD.md"), |
| ("ATTRIBUTION.md", "ATTRIBUTION.md"), |
| ("INSTRUCTIONS.md", "INSTRUCTIONS.md"), |
| ("pyproject.toml", "pyproject.toml"), |
| ("uv.lock", "uv.lock"), |
| ): |
| shutil.copy2(project_root / source_name, artifact / destination_name) |
| for package in ("openthai_systemone_onnx", "openthai_systemone_onnx_tools"): |
| shutil.copytree( |
| project_root / "src" / package, |
| artifact / "src" / package, |
| dirs_exist_ok=True, |
| ignore=shutil.ignore_patterns("__pycache__", "*.pyc"), |
| ) |
| shutil.copytree( |
| project_root / "examples", |
| artifact / "examples", |
| dirs_exist_ok=True, |
| ignore=shutil.ignore_patterns("__pycache__", "*.pyc"), |
| ) |
|
|
|
|
| def quantize_artifact(artifact: Path) -> Path: |
| artifact = artifact.resolve() |
| source = artifact / "onnx" / "model.onnx" |
| manifest_path = artifact / "artifact-manifest.json" |
| if not source.is_file() or not manifest_path.is_file(): |
| raise FileNotFoundError("artifact must contain onnx/model.onnx and artifact-manifest.json") |
|
|
| onnx_dir = source.parent |
| with tempfile.TemporaryDirectory(prefix=".quantize-", dir=onnx_dir) as temporary: |
| temporary_dir = Path(temporary) |
| temporary_graphs = { |
| "int8": temporary_dir / "model.int8.onnx", |
| "int4": temporary_dir / "model.int4.onnx", |
| } |
| previous_logging_disable = logging.root.manager.disable |
| logging.disable(logging.INFO) |
| try: |
| _build_int8(source, temporary_graphs["int8"]) |
| _build_int4(source, temporary_graphs["int4"]) |
| finally: |
| logging.disable(previous_logging_disable) |
|
|
| temporary_variants = {} |
| quantization = { |
| "int8": { |
| "mode": "dynamic", |
| "weight_type": "qint8", |
| "activation_quantization": "dynamic", |
| "per_channel": True, |
| "operator": "MatMulInteger", |
| "excluded_nodes": [DECISION_HEAD], |
| }, |
| "int4": { |
| "mode": "weight_only", |
| "weight_type": "uint4", |
| "activation_type": "float32", |
| "block_size": 32, |
| "symmetric": False, |
| "operator": "MatMulNBits", |
| "excluded_nodes": [DECISION_HEAD], |
| }, |
| } |
| for precision, graph in temporary_graphs.items(): |
| if not _data_path(graph).is_file(): |
| raise RuntimeError(f"quantizer did not create external data for {precision}") |
| _assert_runtime_contract(graph) |
| metadata = variant_metadata( |
| temporary_dir, |
| precision=precision, |
| graph=graph.name, |
| external_data=_data_path(graph).name, |
| quantization=quantization[precision], |
| ) |
| _assert_operator_contract(metadata, precision) |
| temporary_variants[precision] = metadata |
|
|
| for precision, graph in temporary_graphs.items(): |
| final_graph = onnx_dir / graph.name |
| final_data = _data_path(final_graph) |
| os.replace(_data_path(graph), final_data) |
| os.replace(graph, final_graph) |
|
|
| old_manifest = json.loads(manifest_path.read_text(encoding="utf-8")) |
| variants = dict(old_manifest.get("variants", {})) |
| if "fp32" not in variants: |
| variants["fp32"] = variant_metadata( |
| artifact, |
| precision="fp32", |
| graph="onnx/model.onnx", |
| external_data="onnx/model.onnx.data", |
| quantization=None, |
| ) |
| for precision, metadata in temporary_variants.items(): |
| metadata["graph"] = f"onnx/{metadata['graph']}" |
| metadata["external_data"] = f"onnx/{metadata['external_data']}" |
| variants[precision] = metadata |
|
|
| _refresh_release_files(artifact) |
| report = artifact / "validation-report.json" |
| report.unlink(missing_ok=True) |
| return write_manifest(artifact, extra={"variants": variants}) |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--artifact", type=Path, required=True) |
| return parser |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| args = build_parser().parse_args(argv) |
| manifest = quantize_artifact(args.artifact) |
| print(f"Quantized bundle complete: {manifest}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
