From ff27077503440ed140a95258634c460075e739df Mon Sep 17 00:00:00 2001 From: zsqdx Date: Sun, 9 Aug 2026 00:24:00 -0700 Subject: [PATCH] Add experimental ONNX quantization tools --- .gitignore | 3 + python/README_quantization.md | 120 + python/katago/quantization.py | 2183 ++++++++++++++++++ python/quantize_onnx.py | 1937 ++++++++++++++++ python/requirements-quantization.txt | 3 + python/tests/test_quantization_qdq_audit.py | 543 +++++ python/tests/test_quantize_onnx.py | 698 ++++++ python/tests/test_quantize_onnx_artifacts.py | 539 +++++ 8 files changed, 6026 insertions(+) create mode 100644 python/README_quantization.md create mode 100644 python/katago/quantization.py create mode 100644 python/quantize_onnx.py create mode 100644 python/requirements-quantization.txt create mode 100644 python/tests/test_quantization_qdq_audit.py create mode 100644 python/tests/test_quantize_onnx.py create mode 100644 python/tests/test_quantize_onnx_artifacts.py diff --git a/.gitignore b/.gitignore index 7799d05291..bc14d99b06 100644 --- a/.gitignore +++ b/.gitignore @@ -68,6 +68,9 @@ cmake_install.cmake cpp/builds/ .DS_Store +/.venv-quant/ +/.quant-smoke/ +/artifacts/ GPATH GRTAGS GTAGS diff --git a/python/README_quantization.md b/python/README_quantization.md new file mode 100644 index 0000000000..a1e1030ea3 --- /dev/null +++ b/python/README_quantization.md @@ -0,0 +1,120 @@ +# Experimental INT8 and FP8 ONNX export + +`quantize_onnx.py` creates calibrated, explicit Q/DQ ONNX graphs for TensorRT research. It is deliberately accuracy-first: by default it quantizes only the weight projections inside KataGo transformer attention and SwiGLU blocks. Attention QK/AV matmuls, Softmax, RMSNorm, the input stem, outer bottleneck projections, trunk tip, and all policy/value heads remain FP32. + +This is not a replacement for a KataGo `.bin.gz`. The current runtime reads FP32 `.bin.gz` weights and emits its own ONNX graph, so the generated Q/DQ files must initially be built and benchmarked outside KataGo with a strongly typed TensorRT network. + +## 1. Dump the exact inference ONNX + +Add the following to a TensorRT benchmark config: + +```ini +trtDumpDebugPlanToDir = /path/to/onnx-dump +trtTransformerNHWC = true +``` + +Then run the normal KataGo benchmark with the exported `.bin.gz`. The backend writes files such as `plan_19x19_fp16_max.onnx`. Prefer the same exact/max-board variant that will be benchmarked; the `max` graph is useful when calibration includes masked smaller boards. + +Using this backend-emitted graph avoids maintaining a second PyTorch ONNX implementation and guarantees that quantization targets the same raw five-output graph used by KataGo. + +## 2. Install the pinned quantizer + +```bash +python -m venv .venv-quantization +source .venv-quantization/bin/activate +# Windows PowerShell: .venv-quantization\Scripts\Activate.ps1 +python -m pip install -r python/requirements-quantization.txt +``` + +Use a fresh Python 3.10-3.14 environment: Model Optimizer has a large PyTorch/ONNX dependency set that should not mutate the environment used for training. The tested version is NVIDIA Model Optimizer 0.45.0. The script refuses another version unless `--allow-unpinned-modelopt` is supplied, because graph rewrites and Q/DQ placement can change between versions. + +## 3. Calibrate and validate + +Use different shuffled-data shards for calibration and validation: + +```bash +python python/quantize_onnx.py \ + --onnx-input /path/to/onnx-dump/plan_19x19_fp16_max.onnx \ + --output-dir /path/to/quantized-b15 \ + --output-prefix b15c1024h16nbt3tflrs-fson-silu \ + --calibration-data /data/calibration \ + --validation-data /data/held-out \ + --calibration-samples 2048 \ + --validation-samples 512 \ + --batch-size 32 \ + --max-source-files 64 \ + --formats int8 fp8 \ + --expected-quantized-nodes 315 \ + --validation-ep cuda:0 +``` + +The reader samples positions from official shuffled training `.npz` files, expands the packed 22 spatial features, supplies `InputMask`, reshapes the 19 global features to NC11, and applies deterministic training-style history truncation. Since opening a compressed shard may decompress its complete arrays, the default deterministically limits each dataset to 64 weighted-random shards before sampling positions. This keeps a thousand-shard corpus fast; use `--max-source-files 0` when exact full-corpus uniform sampling is more important. The selected-shard count and paths are recorded. + +It also accepts `.npz` files that already contain the exact `InputMask`, `InputSpatial`, `InputGlobal`, and optional `InputMeta` arrays. The requested history transform is applied to those inputs too; use `--history-mode full` when they are already in the exact as-stored inference state. + +By default, `--symmetry-mode random` deterministically chooses one of KataGo's eight board symmetries for every selected source position. `InputMask` and `InputSpatial` receive the same transform; global and metadata inputs are unchanged. Use `--symmetry-mode all` for exhaustive coverage: `--calibration-samples 2048` still means 2048 unique source positions, but the quantizer receives 16384 effective rows (and validation expands the same way). Use `none` only when preserving stored orientation is intentional, such as with already-augmented input files. The manifest records the base/effective counts, symmetry histogram, and symmetry hash. + +Defaults are intentionally conservative: + +- signed symmetric INT8 or FP8 E4M3 Q/DQ through Model Optimizer; +- quantized activations per tensor and weights using Model Optimizer's TensorRT-oriented handling; +- FP32 model inputs, outputs, and fallback operations; +- entropy calibration for INT8 and max calibration for FP8 over real positions (Model Optimizer 0.45's FP8 conversion requires max-calibrated scales); +- seeded random D4 symmetry augmentation matching KataGo's training-style input distribution; +- no latency-only autotuning; +- held-out primary-policy KL/top-move agreement, value KL, raw Q-value/Q-score drift, per-channel score-value error, masked and unmasked ownership error, per-output p99/max/RMSE/relative-L2, and non-finite checks; +- a semantic Q/DQ audit covering Q-to-DQ chains, data types, scale/zero-point values, axes, channel counts, and unexpected quantization; +- a JSON manifest with immutable source/output hashes, sampled-position hashes, selected node names, package/GPU versions, actual ONNX Runtime providers, and validation results. + +Model Optimizer works only on a complete staging copy and writes each result as a staged artifact before promotion. This prevents its shape inference from changing the dumped source ONNX and prevents repeated `--overwrite` runs from appending another copy of a large external weight sidecar. + +Keep `--high-precision fp32` for the first experiments. Model Optimizer's FP16 option converts the entire non-quantized fallback graph, whereas KataGo's current TensorRT backend selectively keeps norms, trunk tip, and heads in FP32. The script therefore rejects global FP16 fallback unless `--allow-global-fp16-fallback` is also supplied, and any such variant must be evaluated as a separate mixed-precision experiment rather than attributed solely to INT8/FP8. + +Numerical release limits are model- and experiment-dependent, so the script does not invent them. Add explicit gates after an FP32/FP16 baseline is established, for example: + +```bash + --max-policy-kl-mean ... \ + --max-policy-kl-p99 ... \ + --max-value-kl-mean ... \ + --max-ownership-rmse ... \ + --max-score-mean-max-abs ... \ + --max-score-mean-sq-max-abs ... \ + --max-lead-max-abs ... \ + --max-q-value-rmse ... \ + --max-q-score-rmse ... \ + --min-policy-top1-agreement ... +``` + +Without explicit limits, the report is descriptive and must not be treated as release qualification. A separate held-out corpus, TensorRT build, throughput tests at production batch sizes, exhaustive symmetry validation (`--symmetry-mode all` or an equivalent independent check), `testgpuerror`, and self-play are still required. + +## TensorRT validation + +Pass a TensorRT 10.16 `trtexec` executable to add a parser/build check: + +```bash + --trtexec /path/to/trtexec --trt-opt-batch 32 --trt-max-batch 64 +``` + +The script uses `--stronglyTyped` and does not add `--fp16`, `--int8`, or `--fp8`; precision is encoded by the Q/DQ graph. It saves the engine plus detailed TensorRT layer information for precision inspection. Build success is only a parser/builder check: verify selected layer precision from that report and eventually run held-out inputs through TensorRT before making accuracy or speed claims. + +ONNX Runtime numerical comparison disables graph optimization for both reference and candidate, avoiding its FP8 rewrite bug. It also verifies that the requested CUDA/TRT provider actually survived session creation, rather than silently reporting a full CPU fallback. Individual unsupported nodes can still fall back and are called out in the manifest. + +INT8 can be researched on Ampere and newer. Hardware FP8 acceleration requires Ada or newer, so RTX 4090/5090 are suitable and RTX 3090 is not. + +## Background + +The initial KataGo-specific comparison was zml24's [`nano/quantize_int8.py`](https://github.com/zml24/KataGo_Transformer/blob/main/nano/quantize_int8.py). This implementation keeps its useful static-Q/DQ and real-position calibration ideas, but uses the current official KataGo input/output contract, deterministic disjoint data, an effective projection allowlist, NVIDIA Model Optimizer for both INT8 and FP8, and multi-output accuracy gates. + +For a direct coverage comparison, `--scope all-weighted` implements the +reference script's actual executable allowlist: on the KataGo-emitted graph, +every `MatMul` or `Conv` with a constant weight is eligible, while +activation-by-activation attention MatMuls (`QK^T` and attention-by-`V`) +remain unquantized. Weighted `Gemm` is supported too if a graph contains it. +This includes the stem, outer bottleneck projections, and policy/value heads. +The reference file declares additional stem/head skip patterns, but does not +pass those patterns to its quantization call. Use `--scope transformer` for +the narrower accuracy-first 315-projection b15 scope; use `all-weighted` for +the aggressive 358-node b15 comparison and validate all five raw heads before +promotion. + +TensorRT explicit-quantization semantics and supported Q/DQ patterns are documented in NVIDIA's [Working with Quantized Types](https://docs.nvidia.com/deeplearning/tensorrt/10.x.x/inference-library/work-quantized-types.html), and Model Optimizer provides the maintained [ONNX PTQ implementation](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/onnx_ptq). diff --git a/python/katago/quantization.py b/python/katago/quantization.py new file mode 100644 index 0000000000..bc78400aee --- /dev/null +++ b/python/katago/quantization.py @@ -0,0 +1,2183 @@ +"""Utilities for accuracy-first quantization of KataGo ONNX graphs. + +This module intentionally operates on the ONNX graph emitted by KataGo's +TensorRT backend. That graph is the inference graph KataGo actually builds, +including its five raw output tensors. It avoids maintaining a second, +slightly different PyTorch-to-ONNX exporter. + +The public helpers are kept independent of NVIDIA Model Optimizer so that data +loading, node selection, graph auditing, and accuracy metrics can be tested +without installing the optional quantization toolchain. +""" + +from __future__ import annotations + +import glob +import hashlib +import json +import math +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple + +import numpy as np +import onnx +from onnx import TensorProto + + +EXPECTED_INPUT_NAMES = ("InputMask", "InputSpatial", "InputGlobal") +OPTIONAL_INPUT_NAMES = ("InputMeta",) +EXPECTED_OUTPUT_NAMES = ( + "OutputPolicyPass", + "OutputPolicy", + "OutputValue", + "OutputScoreValue", + "OutputOwnership", +) + +_OUTPUT_CHANNEL_LABELS = { + "OutputPolicyPass": ("policy", "shortterm_optimistic", "q_value", "q_score"), + "OutputPolicy": ("policy", "shortterm_optimistic", "q_value", "q_score"), + "OutputValue": ("win", "loss", "no_result"), + "OutputScoreValue": ( + "score_mean", + "score_mean_sq", + "lead", + "variance_time", + "shortterm_value_error", + "shortterm_score_error", + ), + "OutputOwnership": ("ownership",), +} + +_TRANSFORMER_PROJECTION_RE = re.compile( + r"(?:^|[./])(?:q_proj|k_proj|v_proj|out_proj|" + r"ffn_linear1|ffn_linear_gate|ffn_linear2)(?=$|[./])" +) +_SUPPORTED_WEIGHTED_OPS = frozenset(("Conv", "MatMul", "Gemm")) +_CONSTANT_PASSTHROUGH_OPS = frozenset(("Cast", "Identity", "Reshape", "Transpose")) + + +@dataclass(frozen=True) +class InputSpec: + name: str + shape: Tuple[Optional[int], ...] + elem_type: int + + +@dataclass +class PositionDataset: + batches: List[Dict[str, np.ndarray]] + sample_count: int + base_sample_count: int + batch_size: int + seed: int + history_mode: str + symmetry_mode: str + symmetry_counts: Dict[str, int] + symmetry_sha256: str + position_sha256: str + selection_sha256: str + available_source_file_count: int + selected_source_file_count: int + max_source_files: int + source_files: List[Dict[str, Any]] + input_shapes: Dict[str, List[int]] + + def manifest(self) -> Dict[str, Any]: + return { + "sample_count": self.sample_count, + "base_sample_count": self.base_sample_count, + "batch_count": len(self.batches), + "batch_size": self.batch_size, + "seed": self.seed, + "history_mode": self.history_mode, + "symmetry_mode": self.symmetry_mode, + "symmetry_counts": self.symmetry_counts, + "symmetry_sha256": self.symmetry_sha256, + "position_sha256": self.position_sha256, + "selection_sha256": self.selection_sha256, + "available_source_file_count": self.available_source_file_count, + "selected_source_file_count": self.selected_source_file_count, + "max_source_files": self.max_source_files, + "source_files": self.source_files, + "input_shapes": self.input_shapes, + } + + +@dataclass +class NodeSelection: + scope: str + selected_names: List[str] + selected_by_op_type: Dict[str, int] + weighted_candidate_names: List[str] + weighted_candidates_by_op_type: Dict[str, int] + rejected: List[Dict[str, str]] + + def manifest(self) -> Dict[str, Any]: + return { + "scope": self.scope, + "selected_count": len(self.selected_names), + "selected_by_op_type": self.selected_by_op_type, + "selected_names": self.selected_names, + "weighted_candidate_count": len(self.weighted_candidate_names), + "weighted_candidates_by_op_type": self.weighted_candidates_by_op_type, + "rejected": self.rejected, + } + + +class ArrayCalibrationDataReader: + """Small ORT/ModelOpt-compatible reader backed by immutable numpy batches.""" + + def __init__(self, dataset: PositionDataset): + self.dataset = dataset + self._start = 0 + self._end = len(dataset.batches) + self._index = self._start + + def get_next(self) -> Optional[Dict[str, np.ndarray]]: + if self._index >= self._end: + return None + batch = self.dataset.batches[self._index] + self._index += 1 + return batch + + def rewind(self) -> None: + self._index = self._start + + def get_first(self) -> Dict[str, np.ndarray]: + if self._start >= self._end: + raise RuntimeError("Calibration dataset is empty") + return self.dataset.batches[self._start] + + def set_range(self, start_index: int, end_index: int) -> None: + if not (0 <= start_index <= end_index <= len(self.dataset.batches)): + raise ValueError( + f"Invalid calibration batch range [{start_index}, {end_index})" + ) + self._start = start_index + self._end = end_index + self.rewind() + + def __len__(self) -> int: + return self._end - self._start + + def __iter__(self): + return self + + def __next__(self) -> Dict[str, np.ndarray]: + value = self.get_next() + if value is None: + raise StopIteration + return value + + +def _dim_value(dim: onnx.TensorShapeProto.Dimension) -> Optional[int]: + if dim.HasField("dim_value") and dim.dim_value > 0: + return int(dim.dim_value) + return None + + +def read_input_specs(model: onnx.ModelProto) -> List[InputSpec]: + initializer_names = {initializer.name for initializer in model.graph.initializer} + specs: List[InputSpec] = [] + for value_info in model.graph.input: + if value_info.name in initializer_names: + continue + tensor_type = value_info.type.tensor_type + specs.append( + InputSpec( + name=value_info.name, + shape=tuple(_dim_value(dim) for dim in tensor_type.shape.dim), + elem_type=int(tensor_type.elem_type), + ) + ) + return specs + + +def validate_katago_io_contract( + model: onnx.ModelProto, *, require_producer_metadata: bool = True +) -> List[InputSpec]: + specs = read_input_specs(model) + names = [spec.name for spec in specs] + missing_inputs = [name for name in EXPECTED_INPUT_NAMES if name not in names] + unknown_inputs = [ + name + for name in names + if name not in EXPECTED_INPUT_NAMES + OPTIONAL_INPUT_NAMES + ] + if missing_inputs or unknown_inputs: + raise ValueError( + "Expected KataGo TensorRT-dump inputs " + f"{EXPECTED_INPUT_NAMES} plus optional {OPTIONAL_INPUT_NAMES}; " + f"got {names}. Missing={missing_inputs}, unknown={unknown_inputs}" + ) + for spec in specs: + if spec.elem_type != TensorProto.FLOAT: + raise ValueError( + f"Input {spec.name} must remain FLOAT, got ONNX elem_type={spec.elem_type}" + ) + if len(spec.shape) != 4: + raise ValueError(f"Input {spec.name} must be rank 4, got {spec.shape}") + + spec_by_name = {spec.name: spec for spec in specs} + expected_channels = {"InputMask": 1, "InputSpatial": 22, "InputGlobal": 19} + for name, channels in expected_channels.items(): + if spec_by_name[name].shape[1] != channels: + raise ValueError( + f"{name} must have {channels} channels for the current KataGo input format, " + f"got {spec_by_name[name].shape}" + ) + spatial_shape = spec_by_name["InputSpatial"].shape + if spatial_shape[2] is None or spatial_shape[3] is None: + raise ValueError("KataGo ONNX must have fixed board height and width") + if spec_by_name["InputMask"].shape[2:] != spatial_shape[2:]: + raise ValueError("InputMask and InputSpatial board dimensions differ") + for name in ("InputGlobal", "InputMeta"): + if name in spec_by_name and spec_by_name[name].shape[2:] != (1, 1): + raise ValueError( + f"{name} must use NC11 layout, got {spec_by_name[name].shape}" + ) + + output_names = [value_info.name for value_info in model.graph.output] + if output_names != list(EXPECTED_OUTPUT_NAMES): + raise ValueError( + "Expected the five raw outputs from KataGo's ONNX emitter in order " + f"{EXPECTED_OUTPUT_NAMES}, got {output_names}" + ) + + output_shapes = { + value_info.name: tuple( + _dim_value(dim) for dim in value_info.type.tensor_type.shape.dim + ) + for value_info in model.graph.output + } + for name, shape in output_shapes.items(): + if len(shape) != 4: + raise ValueError(f"Output {name} must be rank 4, got {shape}") + policy_channels = output_shapes["OutputPolicy"][1] + if policy_channels not in (1, 2, 4): + raise ValueError(f"Unexpected policy channel count: {policy_channels}") + if output_shapes["OutputPolicyPass"][1:] != (policy_channels, 1, 1): + raise ValueError("OutputPolicyPass does not match OutputPolicy channels") + if output_shapes["OutputPolicy"][2:] != spatial_shape[2:]: + raise ValueError("OutputPolicy board dimensions do not match InputSpatial") + if output_shapes["OutputValue"][1:] != (3, 1, 1): + raise ValueError( + f"OutputValue must be [N,3,1,1], got {output_shapes['OutputValue']}" + ) + if output_shapes["OutputScoreValue"][1:] != (6, 1, 1): + raise ValueError( + f"OutputScoreValue must be [N,6,1,1], got {output_shapes['OutputScoreValue']}" + ) + if output_shapes["OutputOwnership"][1:] != (1,) + spatial_shape[2:]: + raise ValueError("OutputOwnership board dimensions do not match InputSpatial") + + metadata = {entry.key: entry.value for entry in model.metadata_props} + if require_producer_metadata and ( + model.producer_name != "katago" or "modelVersion" not in metadata + ): + raise ValueError( + "The source must be the ONNX emitted by KataGo (producer_name=katago with modelVersion metadata)" + ) + return specs + + +def resolve_npz_files(paths: Sequence[str]) -> List[str]: + """Resolve files, directories, and glob expressions deterministically.""" + + found: List[Path] = [] + for raw in paths: + expanded = os.path.expandvars(os.path.expanduser(raw)) + candidate = Path(expanded) + if candidate.is_dir(): + found.extend(candidate.rglob("*.npz")) + elif candidate.is_file(): + if candidate.suffix.lower() != ".npz": + raise ValueError(f"Calibration input is not an .npz file: {candidate}") + found.append(candidate) + else: + matches = [Path(path) for path in glob.glob(expanded, recursive=True)] + found.extend( + path + for path in matches + if path.is_file() and path.suffix.lower() == ".npz" + ) + + unique: Dict[str, Path] = {} + for path in found: + resolved = path.resolve() + unique[os.path.normcase(str(resolved))] = resolved + result = [str(path) for _, path in sorted(unique.items(), key=lambda item: item[0])] + if not result: + raise ValueError(f"No .npz files found in: {list(paths)}") + return result + + +def _npz_sample_count(path: str) -> Tuple[int, str]: + with np.load(path, allow_pickle=False) as data: + if "InputSpatial" in data: + return int(data["InputSpatial"].shape[0]), "onnx-inputs" + if "binaryInputNCHWPacked" in data and "globalInputNC" in data: + return int(data["globalInputNC"].shape[0]), "katago-training" + raise ValueError( + f"{path} is neither an ONNX-input NPZ nor a KataGo training NPZ. " + "Expected InputSpatial or binaryInputNCHWPacked/globalInputNC." + ) + + +def _make_history_matrices() -> Tuple[np.ndarray, np.ndarray]: + """Numpy equivalent of data_processing_pytorch.build_history_matrices.""" + + diagonal = np.array( + [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1], + dtype=np.float32, + ) + h_base = np.diag(diagonal) + h_base[14, 15] = 1.0 + h_base[14, 16] = 1.0 + + builders = np.zeros((5, 22, 22), dtype=np.float32) + builders[0, 9, 9] = 1.0 + builders[0, 14, 15] = -1.0 + builders[0, 14, 16] = -1.0 + builders[0, 15, 15] = 1.0 + builders[0, 15, 16] = 1.0 + builders[1, 10, 10] = 1.0 + builders[1, 15, 16] = -1.0 + builders[1, 16, 16] = 1.0 + builders[2, 11, 11] = 1.0 + builders[3, 12, 12] = 1.0 + builders[4, 13, 13] = 1.0 + return h_base.reshape(1, 22, 22), builders + + +def _history_inclusion( + sample_count: int, history_mode: str, rng: np.random.Generator +) -> np.ndarray: + if history_mode == "full": + return np.ones((sample_count, 5), dtype=np.float32) + if history_mode == "none": + return np.zeros((sample_count, 5), dtype=np.float32) + if history_mode != "training": + raise ValueError(f"Unknown history mode: {history_mode}") + + # Match KataGo training's 2% chance of stopping at each successive history + # plane, but use an explicit generator so the exported scales are reproducible. + should_stop = rng.random((sample_count, 5)) >= 0.98 + return (np.cumsum(should_stop, axis=1, dtype=np.int32) == 0).astype(np.float32) + + +def _apply_history_selection( + spatial: np.ndarray, global_input: np.ndarray, include_history: np.ndarray +) -> Tuple[np.ndarray, np.ndarray]: + if spatial.shape[1] != 22 or global_input.shape[1] != 19: + raise ValueError( + f"History transform expects 22 spatial and 19 global features, got " + f"{spatial.shape[1]} and {global_input.shape[1]}" + ) + h_base, h_builder = _make_history_matrices() + matrices = h_base + np.einsum("bi,ijk->bjk", include_history, h_builder) + spatial = np.einsum("bijk,bil->bljk", spatial, matrices).astype( + np.float32, copy=False + ) + multiplier = np.pad(include_history, ((0, 0), (0, 14)), constant_values=1.0) + global_input = (global_input * multiplier).astype(np.float32, copy=False) + return spatial, global_input + + +def _apply_spatial_symmetry(value: np.ndarray, symmetry: int) -> np.ndarray: + """Apply KataGo's numbered D4 symmetry to the final two tensor axes.""" + + if symmetry == 0: + return value + if symmetry == 1: + return np.flip(np.swapaxes(value, -2, -1), axis=-2) + if symmetry == 2: + return np.flip(value, axis=(-2, -1)) + if symmetry == 3: + return np.flip(np.swapaxes(value, -2, -1), axis=-1) + if symmetry == 4: + return np.swapaxes(value, -2, -1) + if symmetry == 5: + return np.flip(value, axis=-1) + if symmetry == 6: + return np.flip(np.swapaxes(value, -2, -1), axis=(-2, -1)) + if symmetry == 7: + return np.flip(value, axis=-2) + raise ValueError(f"Symmetry must be in [0,7], got {symmetry}") + + +def _augment_symmetries( + inputs: Mapping[str, np.ndarray], + symmetry_mode: str, + seed: int, +) -> Tuple[Dict[str, np.ndarray], np.ndarray]: + """Apply deterministic inference symmetries and return effective symmetry ids. + + ``random`` applies one seeded symmetry to each source position. ``all`` + expands every source position into eight adjacent rows in symmetry order + 0..7. Non-spatial inputs are repeated but never transformed. + """ + + if symmetry_mode not in ("random", "all", "none"): + raise ValueError(f"Unknown symmetry mode: {symmetry_mode}") + spatial = np.asarray(inputs["InputSpatial"]) + if spatial.shape[-2] != spatial.shape[-1] and symmetry_mode != "none": + raise ValueError( + "Random/all D4 symmetry augmentation requires a square ONNX board; " + "use --symmetry-mode none for a rectangular graph" + ) + base_count = int(spatial.shape[0]) + spatial_names = frozenset(("InputMask", "InputSpatial")) + + if symmetry_mode == "none": + none_symmetry_ids = np.zeros(base_count, dtype=np.uint8) + return ( + {name: np.ascontiguousarray(value) for name, value in inputs.items()}, + none_symmetry_ids, + ) + + if symmetry_mode == "random": + symmetry_rng = np.random.default_rng( + np.random.SeedSequence([int(seed), 0x4B415441]) + ) + random_symmetry_ids = symmetry_rng.integers( + 0, 8, size=base_count, dtype=np.uint8 + ) + augmented: Dict[str, np.ndarray] = {} + for name, value in inputs.items(): + value = np.asarray(value) + if name not in spatial_names: + augmented[name] = np.ascontiguousarray(value) + continue + transformed = np.empty_like(value) + for symmetry in range(8): + indices = np.flatnonzero(random_symmetry_ids == symmetry) + if indices.size > 0: + transformed[indices] = _apply_spatial_symmetry( + value[indices], symmetry + ) + augmented[name] = np.ascontiguousarray(transformed) + return augmented, random_symmetry_ids + + all_symmetry_ids = np.tile(np.arange(8, dtype=np.uint8), base_count) + augmented = {} + for name, value in inputs.items(): + value = np.asarray(value) + if name not in spatial_names: + augmented[name] = np.ascontiguousarray(np.repeat(value, 8, axis=0)) + continue + transformed = np.empty((base_count, 8) + value.shape[1:], dtype=value.dtype) + for symmetry in range(8): + transformed[:, symmetry] = _apply_spatial_symmetry(value, symmetry) + augmented[name] = np.ascontiguousarray( + transformed.reshape((base_count * 8,) + value.shape[1:]) + ) + return augmented, all_symmetry_ids + + +def _normalize_nc11(value: np.ndarray, name: str) -> np.ndarray: + value = np.asarray(value, dtype=np.float32) + if value.ndim == 2: + value = value[:, :, None, None] + if value.ndim != 4 or value.shape[2:] != (1, 1): + raise ValueError( + f"{name} must have shape [N,C] or [N,C,1,1], got {value.shape}" + ) + return np.ascontiguousarray(value) + + +def _extract_training_rows( + data: Mapping[str, np.ndarray], + row_indices: np.ndarray, + specs: Sequence[InputSpec], + include_history: np.ndarray, +) -> Dict[str, np.ndarray]: + spatial_spec = next(spec for spec in specs if spec.name == "InputSpatial") + if spatial_spec.shape[2] is None or spatial_spec.shape[3] is None: + raise ValueError("KataGo ONNX must have fixed board dimensions") + height = int(spatial_spec.shape[2]) + width = int(spatial_spec.shape[3]) + packed = np.asarray(data["binaryInputNCHWPacked"][row_indices]) + spatial = np.unpackbits(packed, axis=2) + padded_area = ((height * width + 7) // 8) * 8 + if spatial.shape[2] != padded_area: + raise ValueError( + f"Packed spatial width is {spatial.shape[2]} bits, expected {padded_area} " + f"for {height}x{width}" + ) + spatial = spatial[:, :, : height * width].reshape( + len(row_indices), spatial.shape[1], height, width + ) + spatial = spatial.astype(np.float32, copy=False) + global_input = np.asarray(data["globalInputNC"][row_indices], dtype=np.float32) + spatial, global_input = _apply_history_selection( + spatial, global_input, include_history + ) + + result: Dict[str, np.ndarray] = { + "InputMask": np.ascontiguousarray(spatial[:, 0:1]), + "InputSpatial": np.ascontiguousarray(spatial), + "InputGlobal": _normalize_nc11(global_input, "InputGlobal"), + } + if any(spec.name == "InputMeta" for spec in specs): + if "metadataInputNC" not in data: + raise ValueError( + "The ONNX graph requires InputMeta but metadataInputNC is absent" + ) + result["InputMeta"] = _normalize_nc11( + np.asarray(data["metadataInputNC"][row_indices], dtype=np.float32), + "InputMeta", + ) + return result + + +def _extract_onnx_input_rows( + data: Mapping[str, np.ndarray], + row_indices: np.ndarray, + specs: Sequence[InputSpec], + include_history: np.ndarray, +) -> Dict[str, np.ndarray]: + extracted: Dict[str, np.ndarray] = {} + for spec in specs: + if spec.name not in data: + raise ValueError(f"Expanded input NPZ is missing {spec.name}") + extracted[spec.name] = np.asarray( + data[spec.name][row_indices], dtype=np.float32 + ) + + spatial = extracted["InputSpatial"] + global_input = extracted["InputGlobal"] + if global_input.ndim == 4 and global_input.shape[2:] == (1, 1): + global_input = global_input[:, :, 0, 0] + spatial, global_input = _apply_history_selection( + spatial, global_input, include_history + ) + extracted["InputSpatial"] = spatial + extracted["InputGlobal"] = global_input + + result: Dict[str, np.ndarray] = {} + for spec in specs: + value = extracted[spec.name] + if spec.name in ("InputGlobal", "InputMeta"): + result[spec.name] = _normalize_nc11(value, spec.name) + else: + result[spec.name] = np.ascontiguousarray(value) + return result + + +def _limit_source_files( + files: Sequence[str], + counts: Sequence[int], + formats: Sequence[str], + sample_count: int, + max_source_files: int, + seed: int, +) -> Tuple[List[str], List[int], List[str]]: + """Bound expensive NPZ decompressions while retaining seeded shard sampling.""" + + if max_source_files < 0: + raise ValueError("max_source_files must be nonnegative (0 means unlimited)") + if max_source_files == 0 or len(files) <= max_source_files: + return list(files), list(counts), list(formats) + + limit = min(max_source_files, len(files)) + largest = sorted(range(len(files)), key=lambda index: (-counts[index], index))[ + :limit + ] + largest_capacity = sum(counts[index] for index in largest) + if largest_capacity < sample_count: + raise ValueError( + f"No {limit} source files contain the requested {sample_count} positions; " + "increase --max-source-files or use 0 for exact full-corpus sampling" + ) + + rng = np.random.default_rng(np.random.SeedSequence([int(seed), 0x53484152])) + probabilities = np.asarray(counts, dtype=np.float64) + probabilities /= np.sum(probabilities) + selected = rng.choice( + len(files), size=limit, replace=False, p=probabilities + ).tolist() + # Highly uneven custom shards can produce a subset too small for the + # requested number of unique rows. Fall back to the largest eligible set; + # official shuffled shards are nearly equal-sized and do not hit this path. + if sum(counts[index] for index in selected) < sample_count: + selected = largest + selected.sort() + return ( + [files[index] for index in selected], + [int(counts[index]) for index in selected], + [formats[index] for index in selected], + ) + + +def _validate_batch_shapes( + batch: Mapping[str, np.ndarray], specs: Sequence[InputSpec] +) -> None: + expected_names = [spec.name for spec in specs] + if list(batch.keys()) != expected_names: + raise ValueError( + f"Input order mismatch: expected {expected_names}, got {list(batch.keys())}" + ) + batch_size: Optional[int] = None + for spec in specs: + value = batch[spec.name] + if value.dtype != np.float32: + raise ValueError(f"{spec.name} must be float32, got {value.dtype}") + if len(value.shape) != len(spec.shape): + raise ValueError( + f"{spec.name} rank mismatch: expected {spec.shape}, got {value.shape}" + ) + for axis, (actual, expected) in enumerate(zip(value.shape, spec.shape)): + if axis == 0: + continue + if expected is not None and actual != expected: + raise ValueError( + f"{spec.name} shape mismatch at axis {axis}: expected {spec.shape}, got {value.shape}" + ) + if batch_size is None: + batch_size = int(value.shape[0]) + elif value.shape[0] != batch_size: + raise ValueError("All ONNX inputs must have the same batch dimension") + + +def load_position_dataset( + paths: Sequence[str], + input_specs: Sequence[InputSpec], + sample_count: int, + batch_size: int, + seed: int, + history_mode: str = "training", + symmetry_mode: str = "random", + max_source_files: int = 64, +) -> PositionDataset: + """Uniformly sample real positions and materialize deterministic input batches. + + ``sample_count`` is the number of unique source positions. In ``all`` + symmetry mode, each source position produces eight effective samples. + """ + + if sample_count <= 0 or batch_size <= 0: + raise ValueError("sample_count and batch_size must be positive") + files = resolve_npz_files(paths) + available_source_file_count = len(files) + counts: List[int] = [] + formats: List[str] = [] + for path in files: + count, data_format = _npz_sample_count(path) + if count <= 0: + continue + counts.append(count) + formats.append(data_format) + if len(counts) != len(files): + raise ValueError("Empty .npz files are not supported") + if len(set(formats)) != 1: + raise ValueError( + "Do not mix expanded ONNX-input NPZs with KataGo training NPZs" + ) + files, counts, formats = _limit_source_files( + files, counts, formats, sample_count, max_source_files, seed + ) + total = int(sum(counts)) + if sample_count > total: + raise ValueError( + f"Requested {sample_count} positions but only {total} are available" + ) + + sampling_rng = np.random.default_rng(np.random.SeedSequence([int(seed), 0x504F53])) + history_rng = np.random.default_rng(np.random.SeedSequence([int(seed), 0x484953])) + selected_global = sampling_rng.choice(total, size=sample_count, replace=False) + include_history = _history_inclusion(sample_count, history_mode, history_rng) + cumulative = np.cumsum(np.array([0] + counts, dtype=np.int64)) + + positions_by_file: Dict[int, List[Tuple[int, int]]] = {} + for output_position, global_row in enumerate(selected_global.tolist()): + file_index = int(np.searchsorted(cumulative, global_row, side="right") - 1) + local_row = int(global_row - cumulative[file_index]) + positions_by_file.setdefault(file_index, []).append( + (output_position, local_row) + ) + + combined: Dict[str, Optional[np.ndarray]] = { + spec.name: None for spec in input_specs + } + source_manifest: List[Dict[str, Any]] = [] + selection_hasher = hashlib.sha256() + for file_index, path in enumerate(files): + pairs = positions_by_file.get(file_index, []) + stat = os.stat(path) + source_entry: Dict[str, Any] = { + "path": str(Path(path).resolve()), + "size": int(stat.st_size), + "mtime_ns": int(stat.st_mtime_ns), + "available_positions": counts[file_index], + "selected_positions": len(pairs), + } + source_manifest.append(source_entry) + if not pairs: + continue + + output_positions = np.array([pair[0] for pair in pairs], dtype=np.int64) + local_rows = np.array([pair[1] for pair in pairs], dtype=np.int64) + selection_hasher.update(str(Path(path).resolve()).encode("utf-8")) + selection_hasher.update(local_rows.tobytes(order="C")) + with np.load(path, allow_pickle=False) as data: + if formats[file_index] == "katago-training": + extracted = _extract_training_rows( + data, + local_rows, + input_specs, + include_history[output_positions], + ) + else: + extracted = _extract_onnx_input_rows( + data, + local_rows, + input_specs, + include_history[output_positions], + ) + + for spec in input_specs: + value = extracted[spec.name] + if combined[spec.name] is None: + combined[spec.name] = np.empty( + (sample_count,) + tuple(value.shape[1:]), dtype=np.float32 + ) + combined[spec.name][output_positions] = value # type: ignore[index] + + concrete = { + name: np.ascontiguousarray(value) + for name, value in combined.items() + if value is not None + } + if len(concrete) != len(input_specs): + raise RuntimeError("Failed to materialize every model input") + + concrete, symmetry_ids = _augment_symmetries(concrete, symmetry_mode, seed) + effective_sample_count = int(symmetry_ids.size) + symmetry_hasher = hashlib.sha256() + symmetry_hasher.update(symmetry_mode.encode("utf-8")) + symmetry_hasher.update(symmetry_ids.tobytes(order="C")) + symmetry_sha256 = symmetry_hasher.hexdigest() + selection_hasher.update(b"\x00symmetry\x00") + selection_hasher.update(symmetry_mode.encode("utf-8")) + selection_hasher.update(symmetry_ids.tobytes(order="C")) + symmetry_counts = { + str(symmetry): int(np.count_nonzero(symmetry_ids == symmetry)) + for symmetry in range(8) + } + + position_hasher = hashlib.sha256() + for spec in input_specs: + position_hasher.update(spec.name.encode("utf-8")) + position_hasher.update(concrete[spec.name].tobytes(order="C")) + + batches: List[Dict[str, np.ndarray]] = [] + for start in range(0, effective_sample_count, batch_size): + end = min(start + batch_size, effective_sample_count) + batch = { + spec.name: np.ascontiguousarray(concrete[spec.name][start:end]) + for spec in input_specs + } + _validate_batch_shapes(batch, input_specs) + batches.append(batch) + + return PositionDataset( + batches=batches, + sample_count=effective_sample_count, + base_sample_count=sample_count, + batch_size=batch_size, + seed=seed, + history_mode=history_mode, + symmetry_mode=symmetry_mode, + symmetry_counts=symmetry_counts, + symmetry_sha256=symmetry_sha256, + position_sha256=position_hasher.hexdigest(), + selection_sha256=selection_hasher.hexdigest(), + available_source_file_count=available_source_file_count, + selected_source_file_count=len(files), + max_source_files=max_source_files, + source_files=source_manifest, + input_shapes={name: list(value.shape) for name, value in concrete.items()}, + ) + + +def _constant_tensor_names(model: onnx.ModelProto) -> set[str]: + constants = {initializer.name for initializer in model.graph.initializer} + changed = True + while changed: + changed = False + for node in model.graph.node: + if not node.output or any(output in constants for output in node.output): + continue + if node.op_type == "Constant": + constants.update(node.output) + changed = True + elif node.op_type in _CONSTANT_PASSTHROUGH_OPS and node.input: + if all( + input_name in constants for input_name in node.input if input_name + ): + constants.update(node.output) + changed = True + return constants + + +def _node_has_weight(node: onnx.NodeProto, constant_names: set[str]) -> bool: + if node.op_type == "Conv": + return len(node.input) >= 2 and node.input[1] in constant_names + if node.op_type == "Gemm": + return len(node.input) >= 2 and node.input[1] in constant_names + if node.op_type == "MatMul": + return any(input_name in constant_names for input_name in node.input) + return False + + +def select_quantizable_nodes( + model: onnx.ModelProto, + scope: str = "transformer", + include_regexes: Sequence[str] = (), + exclude_regexes: Sequence[str] = (), + only_regexes: Sequence[str] = (), +) -> NodeSelection: + """Select weighted nodes while excluding attention activation matmuls by construction.""" + + if scope not in ("transformer", "all-weighted"): + raise ValueError(f"Unknown quantization scope: {scope}") + include_patterns = [re.compile(pattern) for pattern in include_regexes] + exclude_patterns = [re.compile(pattern) for pattern in exclude_regexes] + only_patterns = [re.compile(pattern) for pattern in only_regexes] + constants = _constant_tensor_names(model) + + selected: List[str] = [] + candidates: List[str] = [] + selected_counts: Dict[str, int] = {} + candidate_counts: Dict[str, int] = {} + rejected: List[Dict[str, str]] = [] + for index, node in enumerate(model.graph.node): + if node.op_type not in _SUPPORTED_WEIGHTED_OPS: + continue + node_name = node.name or f"__unnamed_{node.op_type}_{index}" + if not node.name: + raise ValueError( + f"Quantizable {node.op_type} node at index {index} has no name; " + "node-level reproducible selection is impossible" + ) + if not _node_has_weight(node, constants): + rejected.append( + {"name": node_name, "reason": "no constant weight (activation matmul)"} + ) + continue + candidates.append(node_name) + candidate_counts[node.op_type] = candidate_counts.get(node.op_type, 0) + 1 + + selected_by_scope = scope == "all-weighted" or bool( + _TRANSFORMER_PROJECTION_RE.search(node_name) + ) + if include_patterns and any( + pattern.search(node_name) for pattern in include_patterns + ): + selected_by_scope = True + if not selected_by_scope: + rejected.append({"name": node_name, "reason": "outside selected scope"}) + continue + if only_patterns and not any( + pattern.search(node_name) for pattern in only_patterns + ): + rejected.append({"name": node_name, "reason": "outside node restriction"}) + continue + if any(pattern.search(node_name) for pattern in exclude_patterns): + rejected.append({"name": node_name, "reason": "matched exclusion regex"}) + continue + selected.append(node_name) + selected_counts[node.op_type] = selected_counts.get(node.op_type, 0) + 1 + + if not selected: + raise ValueError( + "No weighted nodes were selected. Ensure this is the official KataGo-emitted ONNX " + "and inspect node names before overriding the selection regexes." + ) + return NodeSelection( + scope=scope, + selected_names=sorted(selected), + selected_by_op_type=dict(sorted(selected_counts.items())), + weighted_candidate_names=sorted(candidates), + weighted_candidates_by_op_type=dict(sorted(candidate_counts.items())), + rejected=rejected, + ) + + +def _external_locations(model_path: str, model: onnx.ModelProto) -> List[str]: + base = Path(model_path).resolve().parent + result: set[str] = set() + for initializer in model.graph.initializer: + if initializer.data_location != TensorProto.EXTERNAL: + continue + for entry in initializer.external_data: + if entry.key == "location": + result.add(str((base / entry.value).resolve())) + return sorted(result) + + +def sha256_files(paths: Sequence[str]) -> str: + hasher = hashlib.sha256() + for path in sorted(paths, key=os.path.normcase): + resolved = str(Path(path).resolve()) + hasher.update(Path(resolved).name.encode("utf-8")) + with open(resolved, "rb") as handle: + while True: + chunk = handle.read(8 * 1024 * 1024) + if not chunk: + break + hasher.update(chunk) + return hasher.hexdigest() + + +def artifact_manifest( + model_path: str, model: Optional[onnx.ModelProto] = None +) -> Dict[str, Any]: + if model is None: + model = onnx.load(model_path, load_external_data=False) + files = [str(Path(model_path).resolve())] + _external_locations(model_path, model) + missing = [path for path in files if not os.path.isfile(path)] + if missing: + raise FileNotFoundError(f"Missing ONNX external data files: {missing}") + return { + "path": str(Path(model_path).resolve()), + "files": [{"path": path, "size": int(os.path.getsize(path))} for path in files], + "total_size": int(sum(os.path.getsize(path) for path in files)), + "sha256": sha256_files(files), + "opsets": { + item.domain or "ai.onnx": int(item.version) for item in model.opset_import + }, + "metadata": {item.key: item.value for item in model.metadata_props}, + } + + +@dataclass(frozen=True) +class _ConstantValue: + array: np.ndarray + data_type: int + shape: Tuple[int, ...] + source: str + + +def _node_attribute_int(node: onnx.NodeProto, name: str, default: int) -> int: + for attribute in node.attribute: + if attribute.name == name: + return int(attribute.i) + return default + + +def _tensor_shape_from_value_info( + value: onnx.ValueInfoProto, +) -> Optional[Tuple[Optional[int], ...]]: + tensor_type = value.type.tensor_type + if not tensor_type.HasField("shape"): + return None + return tuple( + int(dimension.dim_value) if dimension.HasField("dim_value") else None + for dimension in tensor_type.shape.dim + ) + + +def audit_qdq_model(model_path: str, selected_names: Sequence[str]) -> Dict[str, Any]: + """Deeply audit the explicit Q/DQ contract expected by TensorRT. + + The audit deliberately checks graph semantics rather than only counting Q/DQ + nodes. In particular, every selected weighted input must follow + ``QuantizeLinear -> DequantizeLinear -> weighted op`` and use symmetric + INT8 or FP8 quantization. The returned ``errors`` list is the authoritative + pass/fail signal; the older summary fields are retained for report and CLI + compatibility. + """ + + model = onnx.load(model_path, load_external_data=False) + q_nodes = [node for node in model.graph.node if node.op_type == "QuantizeLinear"] + dq_nodes = [node for node in model.graph.node if node.op_type == "DequantizeLinear"] + initializer_by_name = { + initializer.name: initializer for initializer in model.graph.initializer + } + producer_by_output = { + output: node for node in model.graph.node for output in node.output if output + } + consumers_by_input: Dict[str, List[onnx.NodeProto]] = {} + for node in model.graph.node: + for input_name in node.input: + if input_name: + consumers_by_input.setdefault(input_name, []).append(node) + + constant_node_by_output = { + output: node + for node in model.graph.node + if node.op_type == "Constant" + for output in node.output + if output + } + constant_cache: Dict[str, Optional[_ConstantValue]] = {} + + def node_label(node: onnx.NodeProto) -> str: + if node.name: + return node.name + output = node.output[0] if node.output else "no-output" + return f"{node.op_type}[{output}]" + + def constant_value(name: str) -> Optional[_ConstantValue]: + if name in constant_cache: + return constant_cache[name] + tensor = initializer_by_name.get(name) + source = "initializer" + if tensor is None: + constant_node = constant_node_by_output.get(name) + if constant_node is None: + constant_cache[name] = None + return None + source = f"Constant node {node_label(constant_node)}" + for attribute in constant_node.attribute: + if attribute.name == "value" and attribute.HasField("t"): + tensor = attribute.t + break + if tensor is None: + constant_cache[name] = None + return None + try: + array = np.asarray( + onnx.numpy_helper.to_array( + tensor, base_dir=str(Path(model_path).resolve().parent) + ) + ) + except Exception: + constant_cache[name] = None + return None + value = _ConstantValue( + array=array, + data_type=int(tensor.data_type), + shape=tuple(int(dimension) for dimension in tensor.dims), + source=source, + ) + constant_cache[name] = value + return value + + tensor_shapes: Dict[str, Tuple[Optional[int], ...]] = { + initializer.name: tuple(int(dimension) for dimension in initializer.dims) + for initializer in model.graph.initializer + } + for value in ( + list(model.graph.input) + + list(model.graph.value_info) + + list(model.graph.output) + ): + shape = _tensor_shape_from_value_info(value) + if shape is not None: + tensor_shapes[value.name] = shape + for output_name in constant_node_by_output: + value = constant_value(output_name) + if value is not None: + tensor_shapes[output_name] = tuple(value.shape) + + error_groups: Dict[str, set[str]] = { + "qdq_chain": set(), + "scale": set(), + "zero_point": set(), + "axis": set(), + "granularity": set(), + "qtype": set(), + "unexpected_quantization": set(), + } + + def add_error(category: str, message: str) -> None: + error_groups[category].add(message) + + def dtype_name(data_type: Optional[int]) -> str: + if data_type is None: + return "UNSPECIFIED" + try: + return TensorProto.DataType.Name(int(data_type)) + except ValueError: + return f"UNKNOWN({data_type})" + + allowed_scale_types = { + TensorProto.FLOAT, + TensorProto.FLOAT16, + TensorProto.BFLOAT16, + } + + def validate_scale(node: onnx.NodeProto) -> Optional[_ConstantValue]: + label = node_label(node) + if len(node.input) < 2 or not node.input[1]: + add_error("scale", f"{label}: missing scale input") + return None + value = constant_value(node.input[1]) + if value is None: + add_error( + "scale", f"{label}: scale {node.input[1]!r} is not a readable constant" + ) + return None + if value.data_type not in allowed_scale_types: + add_error( + "scale", + f"{label}: scale has non-floating type {dtype_name(value.data_type)}", + ) + if value.array.size == 0: + add_error("scale", f"{label}: scale is empty") + return value + try: + numeric = np.asarray(value.array, dtype=np.float64) + except (TypeError, ValueError): + add_error( + "scale", f"{label}: scale cannot be interpreted as floating point" + ) + return value + if not np.all(np.isfinite(numeric)): + add_error("scale", f"{label}: scale contains non-finite values") + if np.any(numeric <= 0.0): + add_error("scale", f"{label}: scale must be strictly positive") + return value + + def validate_zero_point( + node: onnx.NodeProto, target_type: Optional[int] + ) -> Optional[_ConstantValue]: + label = node_label(node) + if len(node.input) < 3 or not node.input[2]: + add_error( + "zero_point", f"{label}: missing zero point for symmetric quantization" + ) + return None + value = constant_value(node.input[2]) + if value is None: + add_error( + "zero_point", + f"{label}: zero point {node.input[2]!r} is not a readable constant", + ) + return None + if target_type is not None and value.data_type != target_type: + add_error( + "zero_point", + f"{label}: zero point type {dtype_name(value.data_type)} does not match " + f"quantized type {dtype_name(target_type)}", + ) + try: + all_zero = value.array.size > 0 and bool(np.all(value.array == 0)) + except (TypeError, ValueError): + all_zero = False + if not all_zero: + add_error( + "zero_point", + f"{label}: zero point must be a non-empty all-zero constant", + ) + return value + + qtype_by_output: Dict[str, Optional[int]] = {} + + def quantized_type(node: onnx.NodeProto) -> Optional[int]: + output_dtype: Optional[int] = None + for attribute in node.attribute: + if attribute.name == "output_dtype": + output_dtype = int(attribute.i) + zero_point_type: Optional[int] = None + if len(node.input) >= 3 and node.input[2]: + zero_point = constant_value(node.input[2]) + if zero_point is not None: + zero_point_type = zero_point.data_type + if ( + output_dtype is not None + and zero_point_type is not None + and output_dtype != zero_point_type + ): + add_error( + "qtype", + f"{node_label(node)}: output_dtype {dtype_name(output_dtype)} does not match " + f"zero point type {dtype_name(zero_point_type)}", + ) + return output_dtype if output_dtype is not None else zero_point_type + + scale_by_output: Dict[str, Optional[_ConstantValue]] = {} + zero_point_by_output: Dict[str, Optional[_ConstantValue]] = {} + q_types: Dict[str, int] = {} + constant_names = _constant_tensor_names(model) + weight_q_count = 0 + activation_q_count = 0 + granularity_errors: List[str] = [] + + for node in q_nodes: + if not node.output: + add_error("qdq_chain", f"{node_label(node)}: QuantizeLinear has no output") + continue + output_name = node.output[0] + qtype = quantized_type(node) + qtype_by_output[output_name] = qtype + type_name = dtype_name(qtype) + q_types[type_name] = q_types.get(type_name, 0) + 1 + scale = validate_scale(node) + scale_by_output[output_name] = scale + zero_point_by_output[output_name] = validate_zero_point(node, qtype) + + is_weight = bool(node.input) and node.input[0] in constant_names + if is_weight: + weight_q_count += 1 + else: + activation_q_count += 1 + if scale is not None and scale.array.size != 1: + message = f"{node_label(node)}: activation scale is not per-tensor" + granularity_errors.append(message) + add_error("granularity", message) + + for node in dq_nodes: + if not node.output: + add_error( + "qdq_chain", f"{node_label(node)}: DequantizeLinear has no output" + ) + continue + output_name = node.output[0] + scale_by_output[output_name] = validate_scale(node) + upstream = producer_by_output.get(node.input[0]) if node.input else None + upstream_type = ( + qtype_by_output.get(upstream.output[0]) + if upstream is not None + and upstream.op_type == "QuantizeLinear" + and upstream.output + else None + ) + zero_point_by_output[output_name] = validate_zero_point(node, upstream_type) + + graph_outputs = {value.name for value in model.graph.output} + orphan_q_nodes: set[str] = set() + orphan_dq_nodes: set[str] = set() + for node in q_nodes: + if not node.output: + orphan_q_nodes.add(node_label(node)) + continue + dq_consumers = [ + consumer + for consumer in consumers_by_input.get(node.output[0], []) + if consumer.op_type == "DequantizeLinear" + ] + if not dq_consumers: + orphan_q_nodes.add(node_label(node)) + add_error( + "qdq_chain", + f"{node_label(node)}: QuantizeLinear output has no DQ consumer", + ) + + pair_by_dq_output: Dict[str, Dict[str, Any]] = {} + for dq_node in dq_nodes: + dq_label = node_label(dq_node) + if not dq_node.input: + orphan_dq_nodes.add(dq_label) + add_error( + "qdq_chain", f"{dq_label}: DequantizeLinear has no quantized input" + ) + continue + q_node = producer_by_output.get(dq_node.input[0]) + if q_node is None or q_node.op_type != "QuantizeLinear" or not q_node.output: + orphan_dq_nodes.add(dq_label) + add_error( + "qdq_chain", + f"{dq_label}: DQ input is not produced directly by QuantizeLinear", + ) + continue + if not dq_node.output: + orphan_dq_nodes.add(dq_label) + continue + + q_output = q_node.output[0] + dq_output = dq_node.output[0] + q_scale = scale_by_output.get(q_output) + dq_scale = scale_by_output.get(dq_output) + if q_scale is not None and dq_scale is not None: + if ( + q_scale.data_type != dq_scale.data_type + or q_scale.array.shape != dq_scale.array.shape + or not np.array_equal(q_scale.array, dq_scale.array) + ): + add_error( + "scale", f"{node_label(q_node)} -> {dq_label}: Q/DQ scales differ" + ) + + q_zero = zero_point_by_output.get(q_output) + dq_zero = zero_point_by_output.get(dq_output) + if q_zero is not None and dq_zero is not None: + if ( + q_zero.data_type != dq_zero.data_type + or q_zero.array.shape != dq_zero.array.shape + or not np.array_equal(q_zero.array, dq_zero.array) + ): + add_error( + "zero_point", + f"{node_label(q_node)} -> {dq_label}: Q/DQ zero points differ", + ) + + source_name = q_node.input[0] if q_node.input else "" + source_shape = tensor_shapes.get(source_name) + rank = len(source_shape) if source_shape is not None else None + q_axis = _node_attribute_int(q_node, "axis", 1) + dq_axis = _node_attribute_int(dq_node, "axis", 1) + + def normalize_axis(axis: int) -> int: + return axis + rank if rank is not None and axis < 0 else axis + + effective_q_axis = normalize_axis(q_axis) + effective_dq_axis = normalize_axis(dq_axis) + if rank is not None and not 0 <= effective_q_axis < rank: + add_error( + "axis", + f"{node_label(q_node)}: effective axis {effective_q_axis} is invalid for rank {rank}", + ) + if rank is not None and not 0 <= effective_dq_axis < rank: + add_error( + "axis", + f"{dq_label}: effective axis {effective_dq_axis} is invalid for rank {rank}", + ) + if effective_q_axis != effective_dq_axis: + add_error( + "axis", + f"{node_label(q_node)} -> {dq_label}: Q/DQ effective axes differ " + f"({effective_q_axis} vs {effective_dq_axis})", + ) + + pair_by_dq_output[dq_output] = { + "q_node": q_node, + "dq_node": dq_node, + "qtype": qtype_by_output.get(q_output), + "scale": q_scale, + "effective_axis": effective_q_axis, + "source_tensor": source_name, + "source_shape": source_shape, + } + + if not consumers_by_input.get(dq_output) and dq_output not in graph_outputs: + orphan_dq_nodes.add(dq_label) + add_error("qdq_chain", f"{dq_label}: DequantizeLinear output is unused") + + model_nodes_by_name: Dict[str, List[onnx.NodeProto]] = {} + for node in model.graph.node: + if node.name: + model_nodes_by_name.setdefault(node.name, []).append(node) + selected_missing = sorted( + name for name in selected_names if name not in model_nodes_by_name + ) + for name in selected_missing: + add_error("qdq_chain", f"selected node {name!r} is missing after quantization") + + selected_without_qdq: set[str] = set() + selected_qdq_input_counts: Dict[str, int] = {} + selected_chain_qtypes: Dict[str, Dict[str, str]] = {} + selected_input_chains: Dict[str, Dict[str, Any]] = {} + selected_set = set(selected_names) + supported_selected_qtypes = { + TensorProto.INT8, + getattr(TensorProto, "FLOAT8E4M3FN", 17), + } + + for selected_name in selected_names: + nodes = model_nodes_by_name.get(selected_name, []) + if not nodes: + continue + if len(nodes) != 1: + add_error( + "qdq_chain", f"selected node {selected_name!r} is not uniquely named" + ) + selected_without_qdq.add(selected_name) + continue + node = nodes[0] + weighted_inputs = list(node.input[:2]) + direct_dq_count = sum( + 1 + for input_name in weighted_inputs + if input_name in producer_by_output + and producer_by_output[input_name].op_type == "DequantizeLinear" + ) + selected_qdq_input_counts[selected_name] = direct_dq_count + if len(weighted_inputs) != 2 or direct_dq_count != 2: + selected_without_qdq.add(selected_name) + add_error( + "qdq_chain", + f"{selected_name}: selected weighted op does not have direct DQ on both inputs", + ) + + role_details: Dict[str, Any] = {} + role_qtypes: Dict[str, str] = {} + resolved_types: List[int] = [] + for input_index, role in ((0, "activation"), (1, "weight")): + if input_index >= len(node.input): + continue + input_name = node.input[input_index] + pair = pair_by_dq_output.get(input_name) + if pair is None: + selected_without_qdq.add(selected_name) + add_error( + "qdq_chain", + f"{selected_name}: {role} DQ is not fed directly by QuantizeLinear", + ) + continue + qtype = pair["qtype"] + role_qtypes[role] = dtype_name(qtype) + if qtype is None: + add_error( + "qtype", f"{selected_name}: {role} quantized type is unspecified" + ) + else: + resolved_types.append(qtype) + if qtype not in supported_selected_qtypes: + add_error( + "qtype", + f"{selected_name}: {role} uses unsupported quantized type {dtype_name(qtype)}", + ) + scale = pair["scale"] + scale_elements = int(scale.array.size) if scale is not None else None + role_details[role] = { + "q_node": node_label(pair["q_node"]), + "dq_node": node_label(pair["dq_node"]), + "qtype": dtype_name(qtype), + "scale_elements": scale_elements, + "effective_axis": pair["effective_axis"], + "source_tensor": pair["source_tensor"], + "source_shape": list(pair["source_shape"]) + if pair["source_shape"] is not None + else None, + } + + if role == "activation": + if scale_elements is not None and scale_elements != 1: + message = f"{selected_name}: activation scale is not per-tensor" + granularity_errors.append(message) + add_error("granularity", message) + continue + + source_shape = pair["source_shape"] + if pair["source_tensor"] not in constant_names: + add_error( + "granularity", + f"{selected_name}: weight Q input {pair['source_tensor']!r} is not constant", + ) + continue + if source_shape is None or any( + dimension is None for dimension in source_shape + ): + add_error( + "granularity", + f"{selected_name}: weight shape is not statically known", + ) + continue + + expected_axis: Optional[int] = None + output_channels: Optional[int] = None + if node.op_type == "MatMul": + if len(source_shape) != 2: + add_error( + "axis", + f"{selected_name}: MatMul weight must be rank 2, got {source_shape}", + ) + else: + expected_axis = 1 + output_channels = int(source_shape[1]) + elif node.op_type == "Conv": + if not source_shape: + add_error("axis", f"{selected_name}: Conv weight has no dimensions") + else: + expected_axis = 0 + output_channels = int(source_shape[0]) + elif node.op_type == "Gemm": + if len(source_shape) != 2: + add_error( + "axis", + f"{selected_name}: Gemm weight must be rank 2, got {source_shape}", + ) + else: + trans_b = _node_attribute_int(node, "transB", 0) + expected_axis = 0 if trans_b else 1 + output_channels = int(source_shape[expected_axis]) + + if expected_axis is not None and pair["effective_axis"] != expected_axis: + add_error( + "axis", + f"{selected_name}: {node.op_type} weight axis {pair['effective_axis']} " + f"does not match expected axis {expected_axis}", + ) + if ( + output_channels is not None + and scale_elements is not None + and scale_elements != output_channels + ): + message = ( + f"{selected_name}: weight scale length {scale_elements} does not match " + f"output channels {output_channels}" + ) + granularity_errors.append(message) + add_error("granularity", message) + + if len(set(resolved_types)) > 1: + add_error( + "qtype", + f"{selected_name}: selected input chains mix " + f"{sorted(dtype_name(item) for item in set(resolved_types))}", + ) + selected_chain_qtypes[selected_name] = role_qtypes + selected_input_chains[selected_name] = role_details + + unexpected_fully_quantized_weighted: List[str] = [] + for node in model.graph.node: + if node.op_type not in _SUPPORTED_WEIGHTED_OPS or node.name in selected_set: + continue + if len(node.input) < 2: + continue + pairs = [pair_by_dq_output.get(input_name) for input_name in node.input[:2]] + if any(pair is None for pair in pairs): + continue + if not any( + pair["source_tensor"] in constant_names + for pair in pairs + if pair is not None + ): + continue + label = node_label(node) + unexpected_fully_quantized_weighted.append(label) + add_error( + "unexpected_quantization", + f"{label}: non-selected weighted op unexpectedly has Q/DQ on both inputs", + ) + + io_types = { + "inputs": { + value.name: TensorProto.DataType.Name(value.type.tensor_type.elem_type) + for value in model.graph.input + }, + "outputs": { + value.name: TensorProto.DataType.Name(value.type.tensor_type.elem_type) + for value in model.graph.output + }, + } + bad_io = { + side: {name: dtype for name, dtype in values.items() if dtype != "FLOAT"} + for side, values in io_types.items() + } + bad_io = {side: values for side, values in bad_io.items() if values} + if bad_io: + add_error("qdq_chain", f"graph I/O types are not all FLOAT: {bad_io}") + + error_categories = { + category: sorted(messages) + for category, messages in error_groups.items() + if messages + } + errors = [ + f"{category}: {message}" + for category, messages in error_categories.items() + for message in messages + ] + + return { + "errors": errors, + "error_count": len(errors), + "error_categories": error_categories, + "quantize_linear_count": len(q_nodes), + "dequantize_linear_count": len(dq_nodes), + "qdq_pair_count": len(pair_by_dq_output), + "quantized_tensor_types": dict(sorted(q_types.items())), + "weight_quantize_linear_count": weight_q_count, + "activation_quantize_linear_count": activation_q_count, + "scale_granularity_errors": sorted(set(granularity_errors)), + "orphan_quantize_linear_nodes": sorted(orphan_q_nodes), + "orphan_dequantize_linear_nodes": sorted(orphan_dq_nodes), + "unexpected_fully_quantized_weighted_nodes": sorted( + unexpected_fully_quantized_weighted + ), + "io_types": io_types, + "non_float_io": bad_io, + "selected_nodes_missing_after_quantization": selected_missing, + "selected_nodes_without_qdq_inputs": sorted(selected_without_qdq), + "selected_nodes_with_qdq_inputs": len(selected_qdq_input_counts) + - len(selected_without_qdq), + "selected_qdq_input_counts": dict(sorted(selected_qdq_input_counts.items())), + "selected_chain_qtypes": dict(sorted(selected_chain_qtypes.items())), + "selected_input_chains": dict(sorted(selected_input_chains.items())), + } + + +def _protobuf_sha256(message: Any) -> str: + """Hash one protobuf without depending on map/dictionary iteration order.""" + + try: + payload = message.SerializeToString(deterministic=True) + except TypeError: # pragma: no cover - compatibility with old protobuf releases + payload = message.SerializeToString() + return hashlib.sha256(payload).hexdigest() + + +def _initializer_content_manifest( + initializer: onnx.TensorProto, model_path: str, roles: Sequence[str] +) -> Dict[str, Any]: + """Return a location-independent hash of an initializer's exact tensor bytes.""" + + array = np.asarray( + onnx.numpy_helper.to_array( + initializer, base_dir=str(Path(model_path).resolve().parent) + ) + ) + contiguous = np.ascontiguousarray(array) + byte_view = contiguous.view(np.uint8).reshape(-1) + return { + "data_type": int(initializer.data_type), + "data_type_name": TensorProto.DataType.Name(int(initializer.data_type)), + "dims": [int(value) for value in initializer.dims], + "byte_count": int(byte_view.size), + "byte_sha256": hashlib.sha256(byte_view).hexdigest(), + "roles": sorted(set(roles)), + } + + +def capture_existing_qdq_state( + model_path: str, selected_names: Optional[Sequence[str]] = None +) -> Dict[str, Any]: + """Fingerprint pre-existing fully-quantized weighted operators. + + This is intentionally stricter than :func:`audit_qdq_model`. It records + the weighted nodes, their direct Q/DQ chains, every scale/zero-point + constant node, and every referenced initializer (including quantized + weights). Initializer hashes are based on the resolved tensor bytes, so + moving an external-data artifact does not create a false mismatch. + """ + + model = onnx.load(model_path, load_external_data=False) + producer_by_output = { + output: node for node in model.graph.node for output in node.output if output + } + initializer_by_name = { + initializer.name: initializer for initializer in model.graph.initializer + } + constant_names = _constant_tensor_names(model) + node_by_name: Dict[str, onnx.NodeProto] = {} + for index, node in enumerate(model.graph.node): + if not node.name: + continue + if node.name in node_by_name: + raise ValueError(f"ONNX graph has duplicate node name {node.name!r}") + node_by_name[node.name] = node + + def direct_pair(input_name: str) -> Optional[Tuple[onnx.NodeProto, onnx.NodeProto]]: + dq_node = producer_by_output.get(input_name) + if ( + dq_node is None + or dq_node.op_type != "DequantizeLinear" + or not dq_node.input + ): + return None + q_node = producer_by_output.get(dq_node.input[0]) + if q_node is None or q_node.op_type != "QuantizeLinear" or not q_node.input: + return None + return q_node, dq_node + + discovered: List[str] = [] + for node in model.graph.node: + if ( + node.op_type not in _SUPPORTED_WEIGHTED_OPS + or not node.name + or len(node.input) < 2 + ): + continue + pairs = [direct_pair(input_name) for input_name in node.input[:2]] + if any(pair is None for pair in pairs): + continue + assert all(pair is not None for pair in pairs) + if not any( + pair[0].input[0] in constant_names for pair in pairs if pair is not None + ): + continue + discovered.append(node.name) + + if selected_names is None: + protected_names = sorted(discovered) + else: + protected_names = sorted(set(selected_names)) + missing = sorted(set(protected_names) - set(discovered)) + if missing: + raise ValueError( + "Previously quantized weighted nodes no longer have two direct Q/DQ inputs: " + f"{missing[:10]}" + ) + + if not protected_names: + raise ValueError("No pre-existing fully-quantized weighted nodes were found") + + weighted_nodes: Dict[str, Dict[str, Any]] = {} + qdq_nodes: Dict[str, Dict[str, Any]] = {} + constant_nodes: Dict[str, Dict[str, Any]] = {} + initializer_roles: Dict[str, set[str]] = {} + + def node_key(node: onnx.NodeProto) -> str: + if node.name: + return node.name + if node.output: + return f"{node.op_type}[{node.output[0]}]" + raise ValueError( + f"Protected {node.op_type} node has neither a name nor an output" + ) + + def record_constant(name: str, role: str) -> None: + if name in initializer_by_name: + initializer_roles.setdefault(name, set()).add(role) + return + node = producer_by_output.get(name) + if node is None or node.op_type != "Constant": + return + key = node_key(node) + entry = constant_nodes.setdefault( + key, + { + "op_type": node.op_type, + "protobuf_sha256": _protobuf_sha256(node), + "roles": [], + }, + ) + entry["roles"] = sorted(set(entry["roles"]) | {role}) + + for weighted_name in protected_names: + node = node_by_name.get(weighted_name) + if node is None: + raise ValueError(f"Protected weighted node {weighted_name!r} is missing") + weighted_nodes[weighted_name] = { + "op_type": node.op_type, + "protobuf_sha256": _protobuf_sha256(node), + "inputs": list(node.input), + "outputs": list(node.output), + } + for input_index, input_name in enumerate(node.input[:2]): + pair = direct_pair(input_name) + if pair is None: + raise ValueError( + f"Protected weighted node {weighted_name!r} input {input_index} lost Q/DQ" + ) + q_node, dq_node = pair + chain_role = f"{weighted_name}:input{input_index}" + for kind, chain_node in (("Q", q_node), ("DQ", dq_node)): + key = node_key(chain_node) + entry = qdq_nodes.setdefault( + key, + { + "op_type": chain_node.op_type, + "protobuf_sha256": _protobuf_sha256(chain_node), + "roles": [], + }, + ) + if entry["protobuf_sha256"] != _protobuf_sha256(chain_node): + raise ValueError(f"Protected Q/DQ identity {key!r} is ambiguous") + entry["roles"] = sorted(set(entry["roles"]) | {f"{chain_role}:{kind}"}) + + if q_node.input: + record_constant(q_node.input[0], f"{chain_role}:source") + for index, role in ((1, "scale"), (2, "zero_point")): + if index < len(q_node.input) and q_node.input[index]: + record_constant(q_node.input[index], f"{chain_role}:Q:{role}") + if index < len(dq_node.input) and dq_node.input[index]: + record_constant(dq_node.input[index], f"{chain_role}:DQ:{role}") + + initializers = { + name: _initializer_content_manifest( + initializer_by_name[name], model_path, sorted(roles) + ) + for name, roles in sorted(initializer_roles.items()) + } + sections = { + "weighted_nodes": weighted_nodes, + "qdq_nodes": qdq_nodes, + "constant_nodes": constant_nodes, + "initializers": initializers, + } + aggregate_payload = json.dumps( + sections, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + return { + "schema_version": 1, + "selected_count": len(protected_names), + "selected_names": protected_names, + "discovered_fully_quantized_weighted_count": len(discovered), + "discovered_fully_quantized_weighted_names": sorted(discovered), + "weighted_nodes": weighted_nodes, + "qdq_nodes": qdq_nodes, + "constant_nodes": constant_nodes, + "initializers": initializers, + "weighted_node_count": len(weighted_nodes), + "qdq_node_count": len(qdq_nodes), + "constant_node_count": len(constant_nodes), + "initializer_count": len(initializers), + "aggregate_sha256": hashlib.sha256(aggregate_payload).hexdigest(), + } + + +def compare_existing_qdq_state( + model_path: str, expected: Mapping[str, Any] +) -> Dict[str, Any]: + """Verify that an incremental PTQ pass preserved every protected byte.""" + + actual = capture_existing_qdq_state(model_path, expected["selected_names"]) + differences: List[str] = [] + for section in ("weighted_nodes", "qdq_nodes", "constant_nodes", "initializers"): + expected_items = expected.get(section, {}) + actual_items = actual.get(section, {}) + missing = sorted(set(expected_items) - set(actual_items)) + added = sorted(set(actual_items) - set(expected_items)) + changed = sorted( + name + for name in set(expected_items) & set(actual_items) + if expected_items[name] != actual_items[name] + ) + if missing: + differences.append(f"{section}: missing {missing[:10]}") + if added: + differences.append(f"{section}: added {added[:10]}") + if changed: + differences.append(f"{section}: changed {changed[:10]}") + return { + "status": "passed" if not differences else "failed", + "differences": differences, + "expected_aggregate_sha256": expected.get("aggregate_sha256"), + "actual_aggregate_sha256": actual["aggregate_sha256"], + "actual": actual, + } + + +def _softmax_log_probs( + logits: np.ndarray, axis: int = -1 +) -> Tuple[np.ndarray, np.ndarray]: + logits64 = np.asarray(logits, dtype=np.float64) + maximum = np.max(logits64, axis=axis, keepdims=True) + shifted = logits64 - maximum + log_sum = np.log(np.sum(np.exp(shifted), axis=axis, keepdims=True)) + log_probs = shifted - log_sum + return np.exp(log_probs), log_probs + + +def _summary(values: np.ndarray) -> Dict[str, float]: + values = np.asarray(values, dtype=np.float64).reshape(-1) + finite = values[np.isfinite(values)] + if finite.size == 0: + return {"mean": 0.0, "p99": 0.0, "max": 0.0, "nonfinite": int(values.size)} + return { + "mean": float(np.mean(finite)), + "p99": float(np.quantile(finite, 0.99)), + "max": float(np.max(finite)), + "nonfinite": int(values.size - finite.size), + } + + +def _output_error(reference: np.ndarray, candidate: np.ndarray) -> Dict[str, Any]: + reference64 = np.asarray(reference, dtype=np.float64) + candidate64 = np.asarray(candidate, dtype=np.float64) + if reference64.shape != candidate64.shape: + raise ValueError( + f"Output shape mismatch: {reference64.shape} vs {candidate64.shape}" + ) + finite_reference = np.isfinite(reference64) + finite_candidate = np.isfinite(candidate64) + finite = finite_reference & finite_candidate + diff = np.where(finite, candidate64 - reference64, 0.0) + abs_diff = np.abs(diff[finite]) + reference_finite = reference64[finite] + squared_error = float(np.sum(diff[finite] * diff[finite])) + reference_energy = float(np.sum(reference_finite * reference_finite)) + return { + "shape": list(reference64.shape), + "element_count": int(reference64.size), + "max_abs": float(np.max(abs_diff)) if abs_diff.size else 0.0, + "mean_abs": float(np.mean(abs_diff)) if abs_diff.size else 0.0, + "p99_abs": float(np.quantile(abs_diff, 0.99)) if abs_diff.size else 0.0, + "rmse": math.sqrt(squared_error / max(1, int(np.count_nonzero(finite)))), + "rel_l2": math.sqrt(squared_error / max(reference_energy, 1.0e-30)), + "reference_nonfinite": int( + reference64.size - np.count_nonzero(finite_reference) + ), + "candidate_nonfinite": int( + candidate64.size - np.count_nonzero(finite_candidate) + ), + } + + +def _masked_ownership_error( + reference: np.ndarray, candidate: np.ndarray, mask: np.ndarray +) -> Dict[str, Any]: + reference = np.asarray(reference) + candidate = np.asarray(candidate) + valid = np.asarray(mask) > 0.5 + valid = np.broadcast_to(valid, reference.shape) + inboard = _output_error(reference[valid], candidate[valid]) + unmasked = _output_error(reference, candidate) + offboard = _output_error(reference[~valid], candidate[~valid]) + # Quality gates use in-board ownership, but non-finite values anywhere in + # the tensor remain fatal and off-board drift stays visible in the report. + inboard["reference_nonfinite"] = unmasked["reference_nonfinite"] + inboard["candidate_nonfinite"] = unmasked["candidate_nonfinite"] + inboard["unmasked"] = unmasked + inboard["offboard"] = offboard + return inboard + + +def _per_channel_errors( + name: str, + reference: np.ndarray, + candidate: np.ndarray, + mask: Optional[np.ndarray] = None, +) -> Dict[str, Any]: + reference = np.asarray(reference) + candidate = np.asarray(candidate) + if reference.ndim < 2: + return {} + labels = _OUTPUT_CHANNEL_LABELS.get(name, ()) + result: Dict[str, Any] = {} + for channel in range(reference.shape[1]): + label = labels[channel] if channel < len(labels) else str(channel) + if mask is not None: + channel_mask = np.asarray(mask) > 0.5 + channel_mask = np.broadcast_to( + channel_mask, reference[:, channel : channel + 1].shape + ) + metric = _output_error( + reference[:, channel : channel + 1][channel_mask], + candidate[:, channel : channel + 1][channel_mask], + ) + else: + metric = _output_error(reference[:, channel], candidate[:, channel]) + result[label] = metric + return result + + +def _policy_metrics( + reference_outputs: Mapping[str, np.ndarray], + candidate_outputs: Mapping[str, np.ndarray], + mask: np.ndarray, +) -> Dict[str, Any]: + ref_board = np.asarray(reference_outputs["OutputPolicy"], dtype=np.float64) + cand_board = np.asarray(candidate_outputs["OutputPolicy"], dtype=np.float64) + ref_pass = np.asarray(reference_outputs["OutputPolicyPass"], dtype=np.float64) + cand_pass = np.asarray(candidate_outputs["OutputPolicyPass"], dtype=np.float64) + n, channels = ref_board.shape[:2] + ref_board = ref_board.reshape(n, channels, -1) + cand_board = cand_board.reshape(n, channels, -1) + valid = np.asarray(mask).reshape(n, -1) > 0.5 + ref_pass = ref_pass.reshape(n, channels, -1) + cand_pass = cand_pass.reshape(n, channels, -1) + if ref_pass.shape[2] != 1 or cand_pass.shape[2] != 1: + raise ValueError("OutputPolicyPass must contain exactly one logit per channel") + per_channel: Dict[str, Any] = {} + # Channels 0 and (when present) 1 are policy-logit distributions. Channels + # 2/3 in four-channel models are per-move Q values and must never be judged + # by a shift-invariant softmax KL. + policy_channel_count = min(channels, 2) + labels = _OUTPUT_CHANNEL_LABELS["OutputPolicy"] + for channel in range(policy_channel_count): + ref_board_channel = np.where(valid, ref_board[:, channel], -1.0e30) + cand_board_channel = np.where(valid, cand_board[:, channel], -1.0e30) + ref_logits = np.concatenate((ref_board_channel, ref_pass[:, channel]), axis=1) + cand_logits = np.concatenate( + (cand_board_channel, cand_pass[:, channel]), axis=1 + ) + ref_probs, ref_log_probs = _softmax_log_probs(ref_logits) + _, cand_log_probs = _softmax_log_probs(cand_logits) + kl = np.sum(ref_probs * (ref_log_probs - cand_log_probs), axis=1) + agreement = np.argmax(ref_logits, axis=1) == np.argmax(cand_logits, axis=1) + per_channel[labels[channel]] = { + "kl": _summary(kl), + "top1_agreement": float(np.mean(agreement)), + } + + primary = per_channel["policy"] + result: Dict[str, Any] = { + "kl": primary["kl"], + "top1_agreement": primary["top1_agreement"], + "per_channel": per_channel, + } + if channels == 4: + quantitative: Dict[str, Any] = {} + for channel, label in ((2, "q_value"), (3, "q_score")): + reference_values = np.concatenate( + (ref_board[:, channel][valid], ref_pass[:, channel].reshape(-1)) + ) + candidate_values = np.concatenate( + (cand_board[:, channel][valid], cand_pass[:, channel].reshape(-1)) + ) + quantitative[label] = _output_error(reference_values, candidate_values) + result["quantitative"] = quantitative + return result + + +def _value_metrics(reference: np.ndarray, candidate: np.ndarray) -> Dict[str, Any]: + reference = np.asarray(reference, dtype=np.float64).reshape(reference.shape[0], -1) + candidate = np.asarray(candidate, dtype=np.float64).reshape(candidate.shape[0], -1) + if reference.shape[1] != 3: + raise ValueError(f"OutputValue must have 3 channels, got {reference.shape}") + ref_probs, ref_log_probs = _softmax_log_probs(reference) + _, cand_log_probs = _softmax_log_probs(candidate) + kl = np.sum(ref_probs * (ref_log_probs - cand_log_probs), axis=1) + agreement = np.argmax(reference, axis=1) == np.argmax(candidate, axis=1) + return {"kl": _summary(kl), "top1_agreement": float(np.mean(agreement))} + + +def compute_validation_metrics( + reference_outputs: Mapping[str, np.ndarray], + candidate_outputs: Mapping[str, np.ndarray], + feed: Mapping[str, np.ndarray], +) -> Dict[str, Any]: + missing = [ + name + for name in EXPECTED_OUTPUT_NAMES + if name not in reference_outputs or name not in candidate_outputs + ] + if missing: + raise ValueError(f"Missing outputs for validation: {missing}") + if "InputMask" not in feed: + raise ValueError("InputMask is required for policy and ownership metrics") + + output_metrics: Dict[str, Any] = {} + for name in EXPECTED_OUTPUT_NAMES: + if name == "OutputOwnership": + output_metrics[name] = _masked_ownership_error( + reference_outputs[name], candidate_outputs[name], feed["InputMask"] + ) + output_metrics[name]["per_channel"] = _per_channel_errors( + name, + reference_outputs[name], + candidate_outputs[name], + feed["InputMask"], + ) + else: + output_metrics[name] = _output_error( + reference_outputs[name], candidate_outputs[name] + ) + output_metrics[name]["per_channel"] = _per_channel_errors( + name, reference_outputs[name], candidate_outputs[name] + ) + return { + "sample_count": int(np.asarray(feed["InputMask"]).shape[0]), + "outputs": output_metrics, + "policy": _policy_metrics( + reference_outputs, candidate_outputs, feed["InputMask"] + ), + "value": _value_metrics( + reference_outputs["OutputValue"], candidate_outputs["OutputValue"] + ), + } + + +def _get_nested(mapping: Mapping[str, Any], dotted_path: str) -> float: + current: Any = mapping + for key in dotted_path.split("."): + current = current[key] + return float(current) + + +def evaluate_accuracy_gates( + metrics: Mapping[str, Any], thresholds: Mapping[str, Optional[float]] +) -> Dict[str, Any]: + """Apply only user-specified gates; non-finite outputs always fail.""" + + checks: List[Dict[str, Any]] = [] + mapping = { + "max_policy_kl_mean": ("policy.kl.mean", "max"), + "max_policy_kl_p99": ("policy.kl.p99", "max"), + "max_value_kl_mean": ("value.kl.mean", "max"), + "max_ownership_rmse": ("outputs.OutputOwnership.rmse", "max"), + "max_scorevalue_max_abs": ("outputs.OutputScoreValue.max_abs", "max"), + "max_score_mean_max_abs": ( + "outputs.OutputScoreValue.per_channel.score_mean.max_abs", + "max", + ), + "max_score_mean_sq_max_abs": ( + "outputs.OutputScoreValue.per_channel.score_mean_sq.max_abs", + "max", + ), + "max_lead_max_abs": ( + "outputs.OutputScoreValue.per_channel.lead.max_abs", + "max", + ), + "max_q_value_rmse": ("policy.quantitative.q_value.rmse", "max"), + "max_q_value_max_abs": ("policy.quantitative.q_value.max_abs", "max"), + "max_q_score_rmse": ("policy.quantitative.q_score.rmse", "max"), + "max_q_score_max_abs": ("policy.quantitative.q_score.max_abs", "max"), + "min_policy_top1_agreement": ("policy.top1_agreement", "min"), + } + for threshold_name, (metric_path, direction) in mapping.items(): + threshold = thresholds.get(threshold_name) + if threshold is None: + continue + try: + value = _get_nested(metrics, metric_path) + except (KeyError, TypeError): + checks.append( + { + "threshold": threshold_name, + "metric": metric_path, + "value": None, + "limit": float(threshold), + "passed": False, + "reason": "metric is unavailable for this model/output contract", + } + ) + continue + passed = value <= threshold if direction == "max" else value >= threshold + checks.append( + { + "threshold": threshold_name, + "metric": metric_path, + "value": value, + "limit": float(threshold), + "passed": bool(passed), + } + ) + + reference_nonfinite = sum( + int(output["reference_nonfinite"]) for output in metrics["outputs"].values() + ) + candidate_nonfinite = sum( + int(output["candidate_nonfinite"]) for output in metrics["outputs"].values() + ) + checks.append( + { + "threshold": "reference_nonfinite", + "metric": "sum(outputs.*.reference_nonfinite)", + "value": reference_nonfinite, + "limit": 0, + "passed": reference_nonfinite == 0, + } + ) + checks.append( + { + "threshold": "candidate_nonfinite", + "metric": "sum(outputs.*.candidate_nonfinite)", + "value": candidate_nonfinite, + "limit": 0, + "passed": candidate_nonfinite == 0, + } + ) + configured = any(thresholds.get(name) is not None for name in mapping) + return { + "status": "passed" if all(check["passed"] for check in checks) else "failed", + "numeric_thresholds_configured": configured, + "checks": checks, + } + + +def concatenate_batches( + batches: Sequence[Mapping[str, np.ndarray]], names: Sequence[str] +) -> Dict[str, np.ndarray]: + return { + name: np.concatenate([np.asarray(batch[name]) for batch in batches], axis=0) + for name in names + } + + +def json_dump(path: str, value: Mapping[str, Any]) -> None: + temporary_path = path + ".tmp" + with open(temporary_path, "w", encoding="utf-8") as handle: + json.dump( + value, + handle, + indent=2, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + ) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary_path, path) diff --git a/python/quantize_onnx.py b/python/quantize_onnx.py new file mode 100644 index 0000000000..c9037f237b --- /dev/null +++ b/python/quantize_onnx.py @@ -0,0 +1,1937 @@ +#!/usr/bin/env python3 +"""Export calibrated INT8/FP8 Q/DQ variants of KataGo's emitted ONNX graph. + +The input must be the ONNX file produced by the TensorRT backend when +``trtDumpDebugPlanToDir`` is set. Current KataGo releases do not load the +resulting Q/DQ ONNX files directly; they are research artifacts for TensorRT +``--stronglyTyped`` builds and accuracy/throughput experiments. +""" + +from __future__ import annotations + +import argparse +import gc +import datetime +import importlib.metadata +import inspect +import logging +import os +import platform +import re +import shutil +import subprocess +import sys +import tempfile +import traceback +import uuid +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Sequence + +import numpy as np +import onnx + +from katago.quantization import ( + ArrayCalibrationDataReader, + EXPECTED_OUTPUT_NAMES, + PositionDataset, + artifact_manifest, + audit_qdq_model, + capture_existing_qdq_state, + compare_existing_qdq_state, + compute_validation_metrics, + concatenate_batches, + evaluate_accuracy_gates, + json_dump, + load_position_dataset, + resolve_npz_files, + select_quantizable_nodes, + sha256_files, + validate_katago_io_contract, +) + + +MODEL_OPT_VERSION = "0.45.0" + + +def _build_parser() -> argparse.ArgumentParser: + description = """ +Quantize the exact ONNX inference graph emitted by KataGo's TensorRT backend. +Only transformer q/k/v/out and SwiGLU FFN projections are selected by default; +attention score/value matmuls, norms, stem, outer bottlenecks, trunk tip, and +all output heads stay at the requested high precision. +""" + parser = argparse.ArgumentParser( + description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument("-onnx-input", "--onnx-input", required=True) + parser.add_argument("-output-dir", "--output-dir", required=True) + parser.add_argument( + "--output-prefix", + help="Artifact prefix; defaults to the source ONNX stem", + ) + parser.add_argument( + "--formats", + nargs="+", + choices=("int8", "fp8"), + default=("int8", "fp8"), + ) + parser.add_argument( + "--calibration-data", + nargs="+", + required=True, + help="KataGo training NPZ files/directories/globs, or expanded Input* NPZ files", + ) + parser.add_argument( + "--validation-data", + nargs="+", + help="Disjoint held-out NPZ files/directories/globs", + ) + parser.add_argument("--calibration-samples", type=int, default=2048) + parser.add_argument("--validation-samples", type=int, default=512) + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--seed", type=int, default=20260808) + parser.add_argument( + "--max-source-files", + type=int, + default=64, + help=( + "Bound NPZ shards decompressed per dataset; 0 samples uniformly from every shard. " + "Official shuffled shards make the bounded default much faster with little bias" + ), + ) + parser.add_argument( + "--history-mode", + choices=("training", "full", "none"), + default="training", + help="Deterministic training-style history truncation, or force all/no history planes", + ) + parser.add_argument( + "--symmetry-mode", + choices=("random", "all", "none"), + default="random", + help=( + "Apply one seeded D4 symmetry per source position, expand each source position " + "to all eight symmetries, or preserve stored orientation" + ), + ) + parser.add_argument( + "--allow-data-overlap", + action="store_true", + help="Allow calibration and validation to use any of the same NPZ shards", + ) + parser.add_argument( + "--skip-validation", + action="store_true", + help="Export without held-out numerical comparison (never sufficient for release)", + ) + + parser.add_argument( + "--scope", choices=("transformer", "all-weighted"), default="transformer" + ) + parser.add_argument( + "--include-node-regex", + action="append", + default=[], + help="Add weighted nodes matching this regex to the selected scope", + ) + parser.add_argument( + "--exclude-node-regex", + action="append", + default=[], + help="Remove selected weighted nodes matching this regex", + ) + parser.add_argument( + "--only-node-regex", + action="append", + default=[], + help=( + "Restrict the selected scope to nodes matching at least one regex. " + "Useful for reproducible block-level sensitivity searches" + ), + ) + parser.add_argument( + "--expected-quantized-nodes", + type=int, + help="Fail if node selection differs (315 for b15c1024h16nbt3tflrs)", + ) + parser.add_argument( + "--preserve-existing-qdq", + action="store_true", + help=( + "Incrementally calibrate only the newly selected INT8 nodes while strictly " + "preserving every pre-existing weighted Q/DQ chain and referenced initializer" + ), + ) + parser.add_argument( + "--expected-existing-qdq-nodes", + type=int, + help=( + "With --preserve-existing-qdq, fail unless exactly this many fully quantized " + "weighted nodes already exist in the input graph" + ), + ) + parser.add_argument( + "--calibration-method", + choices=("entropy", "max"), + default="entropy", + help="INT8 calibration method", + ) + parser.add_argument( + "--fp8-calibration-method", + choices=("max", "entropy"), + default="max", + help="FP8 calibration method; ModelOpt 0.45 supports conversion from max calibration", + ) + parser.add_argument( + "--fp8-scale-mode", + choices=("direct-amax", "modelopt-legacy"), + default="direct-amax", + help=( + "Use direct E4M3 amax/qmax scales, consistent with ModelOpt's Torch exporter, " + "or retain ModelOpt 0.45's lossy INT8-to-FP8 conversion for comparison" + ), + ) + parser.add_argument( + "--fp8-activation-qmax", + type=float, + default=448.0, + help=( + "Effective positive E4M3 activation range for direct-amax scaling. " + "Lower values reserve headroom for calibration-set outliers; weights always use 448" + ), + ) + parser.add_argument( + "--allow-fp8-nonmax-calibration", + action="store_true", + help="Permit an unsupported research experiment using non-max FP8 calibration", + ) + parser.add_argument( + "--calibration-eps", + default="cuda:0,cpu", + help="Ordered ModelOpt calibration execution providers", + ) + parser.add_argument( + "--high-precision", + choices=("fp32", "fp16"), + default="fp32", + help="Fallback precision; fp32 isolates quantization error", + ) + parser.add_argument( + "--allow-global-fp16-fallback", + action="store_true", + help=( + "Acknowledge that ModelOpt converts all non-quantized fallback ops to FP16; " + "this is not the same as KataGo's selective FP32 precision pins" + ), + ) + parser.add_argument( + "--calibrate-per-node", + action="store_true", + help="Lower calibration memory at substantial runtime cost", + ) + parser.add_argument("--keep-intermediate-files", action="store_true") + parser.add_argument( + "--allow-unpinned-modelopt", + action="store_true", + help=f"Permit a ModelOpt version other than the tested {MODEL_OPT_VERSION}", + ) + + parser.add_argument( + "--validation-ep", + default="cuda:0", + help="ONNX Runtime provider: cpu, cuda:N, or trt", + ) + parser.add_argument( + "--allow-validation-ep-fallback", + action="store_true", + help="Allow an unavailable requested CUDA/TRT EP to fall back entirely to CPU", + ) + parser.add_argument("--ort-intra-op-threads", type=int, default=0) + parser.add_argument("--max-policy-kl-mean", type=float) + parser.add_argument("--max-policy-kl-p99", type=float) + parser.add_argument("--max-value-kl-mean", type=float) + parser.add_argument("--max-ownership-rmse", type=float) + parser.add_argument("--max-scorevalue-max-abs", type=float) + parser.add_argument("--max-score-mean-max-abs", type=float) + parser.add_argument("--max-score-mean-sq-max-abs", type=float) + parser.add_argument("--max-lead-max-abs", type=float) + parser.add_argument("--max-q-value-rmse", type=float) + parser.add_argument("--max-q-value-max-abs", type=float) + parser.add_argument("--max-q-score-rmse", type=float) + parser.add_argument("--max-q-score-max-abs", type=float) + parser.add_argument("--min-policy-top1-agreement", type=float) + + parser.add_argument( + "--trtexec", + help="Optionally parse/build each Q/DQ graph using this trtexec executable", + ) + parser.add_argument("--trt-opt-batch", type=int, default=32) + parser.add_argument("--trt-max-batch", type=int, default=64) + parser.add_argument("--trt-workspace-mib", type=int, default=0) + parser.add_argument("--trt-timeout-seconds", type=int, default=3600) + + parser.add_argument( + "--skip-onnx-check", + action="store_true", + help="Skip ONNX checker on source/output", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Allow ModelOpt/report files at the selected paths to be replaced", + ) + parser.add_argument( + "--continue-on-error", + action="store_true", + help="Attempt the other format if one quantization path fails", + ) + return parser + + +def _setup_logging(output_dir: str, output_prefix: str) -> str: + log_path = os.path.join(output_dir, f"{output_prefix}.quantization.log") + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + handlers=[ + logging.StreamHandler(sys.stdout), + logging.FileHandler(log_path, mode="w", encoding="utf-8"), + ], + force=True, + ) + return log_path + + +def _git_revision() -> Optional[str]: + try: + return subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).resolve().parent.parent, + check=True, + capture_output=True, + text=True, + timeout=10, + ).stdout.strip() + except Exception: + return None + + +def _gpu_summary() -> List[Dict[str, str]]: + executable = shutil.which("nvidia-smi") + if executable is None: + return [] + try: + output = subprocess.run( + [ + executable, + "--query-gpu=index,name,compute_cap,memory.total,driver_version", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=15, + ).stdout + result = [] + for line in output.splitlines(): + values = [value.strip() for value in line.split(",")] + if len(values) == 5: + result.append( + dict( + index=values[0], + name=values[1], + compute_capability=values[2], + memory_mib=values[3], + driver_version=values[4], + ) + ) + return result + except Exception: + return [] + + +def _optional_package_versions() -> Dict[str, Optional[str]]: + distributions = ( + "onnxruntime-gpu", + "onnxruntime", + "nvidia-modelopt", + "onnx-graphsurgeon", + "polygraphy", + "tensorrt", + ) + versions: Dict[str, Optional[str]] = {} + for distribution in distributions: + try: + versions[distribution] = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + versions[distribution] = None + return versions + + +def _fp8_hardware_summary(gpus: Sequence[Mapping[str, str]]) -> Dict[str, Any]: + capable_indices: List[str] = [] + for gpu in gpus: + try: + major_text, minor_text = gpu["compute_capability"].split(".", 1) + capability = int(major_text) * 10 + int(minor_text) + except (KeyError, TypeError, ValueError): + continue + if capability >= 89: + capable_indices.append(str(gpu["index"])) + return { + "minimum_compute_capability": "8.9 (Ada or newer)", + "capable_gpu_indices": capable_indices, + "hardware_acceleration_detected": bool(capable_indices), + } + + +def _load_modelopt(allow_unpinned: bool): + try: + import modelopt # type: ignore[import-not-found] + from modelopt.onnx.quantization import quantize # type: ignore[import-not-found] + except ImportError as exc: + raise RuntimeError( + "NVIDIA Model Optimizer ONNX support is required. Install " + f"python/requirements-quantization.txt (nvidia-modelopt[onnx]=={MODEL_OPT_VERSION})." + ) from exc + version = getattr(modelopt, "__version__", "unknown") + if version != MODEL_OPT_VERSION and not allow_unpinned: + raise RuntimeError( + f"Expected NVIDIA Model Optimizer {MODEL_OPT_VERSION}, found {version}. " + "Use --allow-unpinned-modelopt only for an intentional compatibility experiment." + ) + return quantize, version + + +def _preload_ort_gpu_dependencies() -> None: + """Load CUDA/cuDNN wheels before ORT creates a GPU execution provider. + + Recent ``onnxruntime-gpu`` wheels intentionally do not depend on the large + NVIDIA runtime wheels unless the ``[cuda,cudnn]`` extra is installed. Even + when those wheels are present, their library directories are not normally + on Linux's loader path. ORT 1.21+ exposes this helper so a session does not + silently fall back to CPU merely because (for example) libcublasLt.so.12 + was not preloaded. + """ + + import onnxruntime as ort # type: ignore[import-untyped] + + preload = getattr(ort, "preload_dlls", None) + if callable(preload): + preload(directory="") + + +def _provider_spec(requested: str) -> tuple[List[Any], str]: + import onnxruntime as ort # type: ignore[import-untyped] + + if requested != "cpu": + _preload_ort_gpu_dependencies() + available = set(ort.get_available_providers()) + if requested == "cpu": + name = "CPUExecutionProvider" + providers: List[Any] = [name] + elif requested.startswith("cuda:"): + name = "CUDAExecutionProvider" + device_id = int(requested.split(":", 1)[1]) + providers = [(name, {"device_id": device_id}), "CPUExecutionProvider"] + elif requested == "trt": + name = "TensorrtExecutionProvider" + providers = [name, "CUDAExecutionProvider", "CPUExecutionProvider"] + else: + raise ValueError(f"Unknown validation EP: {requested}") + if name not in available: + raise RuntimeError( + f"Requested {name}, but ONNX Runtime only has {sorted(available)}" + ) + return providers, name + + +def _run_ort_model( + model_path: str, + dataset: PositionDataset, + requested_ep: str, + intra_op_threads: int, + allow_ep_fallback: bool, +) -> tuple[List[Dict[str, np.ndarray]], Dict[str, Any]]: + import onnxruntime as ort # type: ignore[import-untyped] + + options = ort.SessionOptions() + # Use the same unoptimized execution graph for the FP32 reference and every + # candidate. ORT 1.22's optimizer may incorrectly rewrite an FP8 Q/DQ + # MatMul into MatMulIntegerToFloat, whose input contract is INT8. Running + # the explicit Q/DQ nodes without graph rewrites correctly emulates FP8 for + # held-out error measurement. + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + if intra_op_threads > 0: + options.intra_op_num_threads = intra_op_threads + providers, required_provider = _provider_spec(requested_ep) + logging.info("Loading %s with ONNX Runtime providers %s", model_path, providers) + session = ort.InferenceSession( + model_path, sess_options=options, providers=providers + ) + active_providers = list(session.get_providers()) + fell_back_entirely = required_provider not in active_providers + if fell_back_entirely and not allow_ep_fallback: + raise RuntimeError( + f"Requested {required_provider}, but the created ONNX Runtime session only has " + f"{active_providers}. This commonly means a CUDA/cuDNN DLL failed to load. " + "Fix the environment or explicitly use --allow-validation-ep-fallback." + ) + actual_outputs = [output.name for output in session.get_outputs()] + if actual_outputs != list(EXPECTED_OUTPUT_NAMES): + raise RuntimeError( + f"Unexpected output contract in {model_path}: {actual_outputs}" + ) + results: List[Dict[str, np.ndarray]] = [] + for batch_index, batch in enumerate(dataset.batches): + values = session.run(list(EXPECTED_OUTPUT_NAMES), batch) + results.append(dict(zip(EXPECTED_OUTPUT_NAMES, values))) + if (batch_index + 1) % 10 == 0 or batch_index + 1 == len(dataset.batches): + logging.info( + "Validated %d/%d batches", batch_index + 1, len(dataset.batches) + ) + del session + return results, { + "requested": requested_ep, + "required_provider": required_provider, + "configured_providers": providers, + "active_providers": active_providers, + "entire_session_fallback": fell_back_entirely, + "fallback_allowed": allow_ep_fallback, + "graph_optimization": "ORT_DISABLE_ALL", + "note": ( + "Provider registration is verified. Individual unsupported nodes may still fall " + "back unless ONNX Runtime is configured separately to forbid per-node CPU fallback." + ), + } + + +def _validation_metrics( + reference_batches: Sequence[Mapping[str, np.ndarray]], + candidate_batches: Sequence[Mapping[str, np.ndarray]], + dataset: PositionDataset, +) -> Dict[str, Any]: + if len(reference_batches) != len(candidate_batches): + raise RuntimeError("Reference and candidate validation batch counts differ") + reference = concatenate_batches(reference_batches, EXPECTED_OUTPUT_NAMES) + candidate = concatenate_batches(candidate_batches, EXPECTED_OUTPUT_NAMES) + feed = concatenate_batches(dataset.batches, ["InputMask"]) + return compute_validation_metrics(reference, candidate, feed) + + +def _dynamic_shape_string(input_specs, batch: int) -> str: + items = [] + for spec in input_specs: + dims = [batch if axis == 0 else dim for axis, dim in enumerate(spec.shape)] + if any(dim is None for dim in dims): + raise ValueError( + f"trtexec validation needs fixed non-batch dimensions: {spec}" + ) + items.append(f"{spec.name}:" + "x".join(str(int(dim)) for dim in dims)) + return ",".join(items) + + +def _run_trtexec( + executable: str, + model_path: str, + input_specs, + opt_batch: int, + max_batch: int, + workspace_mib: int, + timeout_seconds: int, + overwrite: bool, +) -> Dict[str, Any]: + resolved = shutil.which(executable) or ( + executable if os.path.isfile(executable) else None + ) + if resolved is None: + raise FileNotFoundError(f"Could not find trtexec: {executable}") + if not (1 <= opt_batch <= max_batch): + raise ValueError("Require 1 <= --trt-opt-batch <= --trt-max-batch") + engine_path = str(Path(model_path).with_suffix(".engine")) + layer_info_path = str(Path(model_path).with_suffix(".layers.json")) + for generated_path in (engine_path, layer_info_path): + if os.path.exists(generated_path): + if not overwrite: + raise FileExistsError( + f"Refusing to replace TensorRT artifact {generated_path}; use --overwrite" + ) + os.remove(generated_path) + command = [ + resolved, + f"--onnx={model_path}", + "--stronglyTyped", + "--skipInference", + "--builderOptimizationLevel=5", + "--profilingVerbosity=detailed", + "--dumpLayerInfo", + f"--exportLayerInfo={layer_info_path}", + f"--saveEngine={engine_path}", + f"--minShapes={_dynamic_shape_string(input_specs, 1)}", + f"--optShapes={_dynamic_shape_string(input_specs, opt_batch)}", + f"--maxShapes={_dynamic_shape_string(input_specs, max_batch)}", + ] + if workspace_mib > 0: + command.append(f"--memPoolSize=workspace:{workspace_mib}MiB") + logging.info("Building strongly typed TensorRT network: %s", " ".join(command)) + version_process = subprocess.run( + [resolved, "--version"], + capture_output=True, + text=True, + timeout=min(timeout_seconds, 30), + ) + process = subprocess.run( + command, + capture_output=True, + text=True, + timeout=timeout_seconds, + ) + combined = (process.stdout or "") + "\n" + (process.stderr or "") + tail = combined.splitlines()[-250:] + passed = process.returncode == 0 + generated = {} + for label, path in (("engine", engine_path), ("layer_info", layer_info_path)): + if os.path.isfile(path): + generated[label] = { + "path": path, + "size": int(os.path.getsize(path)), + "sha256": sha256_files([path]), + } + return { + "command": command, + "returncode": int(process.returncode), + "passed": passed, + "output_tail": tail, + "version_output": ( + (version_process.stdout or "") + "\n" + (version_process.stderr or "") + ).strip(), + "artifacts": generated, + "precision_audit": ( + "Layer information is exported for inspection. A successful build alone does not " + "prove that every selected GEMM ran in INT8/FP8 or establish TensorRT accuracy." + ), + } + + +def _accuracy_thresholds(args: argparse.Namespace) -> Dict[str, Optional[float]]: + return { + "max_policy_kl_mean": args.max_policy_kl_mean, + "max_policy_kl_p99": args.max_policy_kl_p99, + "max_value_kl_mean": args.max_value_kl_mean, + "max_ownership_rmse": args.max_ownership_rmse, + "max_scorevalue_max_abs": args.max_scorevalue_max_abs, + "max_score_mean_max_abs": args.max_score_mean_max_abs, + "max_score_mean_sq_max_abs": args.max_score_mean_sq_max_abs, + "max_lead_max_abs": args.max_lead_max_abs, + "max_q_value_rmse": args.max_q_value_rmse, + "max_q_value_max_abs": args.max_q_value_max_abs, + "max_q_score_rmse": args.max_q_score_rmse, + "max_q_score_max_abs": args.max_q_score_max_abs, + "min_policy_top1_agreement": args.min_policy_top1_agreement, + } + + +def _external_artifact_members( + model_path: str, *, require_exists: bool +) -> List[tuple[Path, Path]]: + """Return safe ``(relative location, absolute path)`` external-data members. + + ONNX external-data locations are interpreted relative to the model file. A + quantization artifact must never be allowed to escape that directory: these + paths are later copied, replaced, and removed without globs. + """ + + primary = Path(model_path).resolve() + base = primary.parent + model = onnx.load(str(primary), load_external_data=False) + members: Dict[str, tuple[Path, Path]] = {} + for initializer in model.graph.initializer: + if initializer.data_location != onnx.TensorProto.EXTERNAL: + continue + locations = [ + entry.value + for entry in initializer.external_data + if entry.key == "location" + ] + if len(locations) != 1 or not locations[0]: + raise ValueError( + f"External initializer {initializer.name!r} must have exactly one location" + ) + relative = Path(locations[0]) + if relative.is_absolute() or relative.drive or ".." in relative.parts: + raise ValueError( + f"Unsafe external-data location {locations[0]!r} in {primary}" + ) + absolute = (base / relative).resolve() + try: + normalized_relative = absolute.relative_to(base) + except ValueError as exc: + raise ValueError( + f"External-data location {locations[0]!r} escapes {base}" + ) from exc + if absolute == primary: + raise ValueError(f"External data aliases its ONNX file: {absolute}") + key = os.path.normcase(str(absolute)) + members[key] = (normalized_relative, absolute) + + result = [members[key] for key in sorted(members)] + if require_exists: + missing = [str(absolute) for _, absolute in result if not absolute.is_file()] + if missing: + raise FileNotFoundError(f"Missing ONNX external data files: {missing}") + return result + + +def _copy_onnx_artifact(source_model_path: str, destination_dir: str) -> str: + """Copy an ONNX file and every referenced external-data file as one artifact.""" + + source = Path(source_model_path).resolve() + destination = Path(destination_dir).resolve() + destination.mkdir(parents=True, exist_ok=True) + target_model = destination / source.name + if target_model.exists(): + raise FileExistsError(f"Staging model already exists: {target_model}") + + shutil.copy2(source, target_model) + try: + for relative, external_source in _external_artifact_members( + str(source), require_exists=True + ): + external_target = (destination / relative).resolve() + try: + external_target.relative_to(destination) + except ValueError as exc: + raise ValueError(f"Unsafe staging target: {external_target}") from exc + if external_target == target_model: + raise ValueError( + f"External data aliases staged model: {external_target}" + ) + external_target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(external_source, external_target) + except Exception: + if target_model.exists(): + target_model.unlink() + raise + return str(target_model) + + +def _promote_staged_artifact( + staged_model_path: str, final_model_path: str, *, overwrite: bool +) -> Dict[str, Any]: + """Commit a complete staged artifact, swapping the primary ONNX file last. + + The staged output name is unique, therefore its external-data names are also + unique. Moving those files first leaves an existing artifact valid until + the atomic primary-file replacement. Only sidecars referenced by the old + primary file are removed afterwards. + """ + + staged_model = Path(staged_model_path).resolve() + final_model = Path(final_model_path).resolve() + final_dir = final_model.parent + if not staged_model.is_file(): + raise FileNotFoundError(f"Staged ONNX output is missing: {staged_model}") + if final_model.exists() and not overwrite: + raise FileExistsError(f"Refusing to replace {final_model}; use --overwrite") + + old_external: List[Path] = [] + if final_model.exists(): + old_external = [ + absolute + for _, absolute in _external_artifact_members( + str(final_model), require_exists=False + ) + ] + + new_members = _external_artifact_members(str(staged_model), require_exists=True) + staged_dir = staged_model.parent + new_targets: List[tuple[Path, Path]] = [] + old_keys = {os.path.normcase(str(path)) for path in old_external} + for relative, staged_external in new_members: + target = (final_dir / relative).resolve() + try: + target.relative_to(final_dir) + except ValueError as exc: + raise ValueError(f"Unsafe promoted external-data target: {target}") from exc + # Reusing the old sidecar name would invalidate the old primary before + # it is atomically swapped. Staged names are deliberately unique. + if target.exists() or os.path.normcase(str(target)) in old_keys: + raise FileExistsError( + f"Staged external-data name is not unique and cannot be safely promoted: {target}" + ) + try: + staged_external.relative_to(staged_dir) + except ValueError as exc: + raise ValueError( + f"Staged external data escapes its directory: {staged_external}" + ) from exc + new_targets.append((staged_external, target)) + + moved_targets: List[Path] = [] + primary_promoted = False + try: + for source, target in new_targets: + target.parent.mkdir(parents=True, exist_ok=True) + os.replace(source, target) + moved_targets.append(target) + os.replace(staged_model, final_model) + primary_promoted = True + finally: + if not primary_promoted: + for target in moved_targets: + try: + target.unlink() + except FileNotFoundError: + pass + + new_keys = {os.path.normcase(str(target)) for _, target in new_targets} + removed_old: List[str] = [] + for old_path in old_external: + if os.path.normcase(str(old_path)) in new_keys: + continue + try: + old_path.unlink() + removed_old.append(str(old_path)) + except FileNotFoundError: + pass + + return { + "promoted_external_data": [str(target) for _, target in new_targets], + "removed_previous_external_data": removed_old, + } + + +def _artifact_integrity(expected_manifest: Mapping[str, Any]) -> Dict[str, Any]: + """Re-hash an artifact using its originally resolved member list.""" + + expected_files = [str(item["path"]) for item in expected_manifest["files"]] + try: + missing = [path for path in expected_files if not os.path.isfile(path)] + if missing: + raise FileNotFoundError(f"Artifact members disappeared: {missing}") + current_hash = sha256_files(expected_files) + current_size = int(sum(os.path.getsize(path) for path in expected_files)) + passed = current_hash == expected_manifest["sha256"] and current_size == int( + expected_manifest["total_size"] + ) + return { + "status": "passed" if passed else "failed", + "expected_sha256": expected_manifest["sha256"], + "current_sha256": current_hash, + "expected_total_size": int(expected_manifest["total_size"]), + "current_total_size": current_size, + "files": expected_files, + } + except Exception as exc: + return { + "status": "failed", + "expected_sha256": expected_manifest.get("sha256"), + "error": str(exc), + "files": expected_files, + } + + +def _mode_report_failed(mode_report: Mapping[str, Any]) -> bool: + accuracy_gate = mode_report.get("accuracy_gate") + if accuracy_gate is not None and accuracy_gate.get("status") != "passed": + return True + trtexec = mode_report.get("trtexec") + return trtexec is not None and not bool(trtexec.get("passed")) + + +def _rewrite_fp8_direct_amax_scales( + model_path: str, + selected_names: Sequence[str], + *, + activation_qmax: float, +) -> Dict[str, Any]: + """Correct ModelOpt 0.45's INT8-to-FP8 scale conversion. + + ModelOpt's ONNX PTQ path first calibrates symmetric INT8 (scale roughly + ``amax / 127``), then multiplies that scale by ``448 / 127``. The result + uses only about 36 representable E4M3 levels. Direct E4M3 calibration is + ``amax / qmax``; ModelOpt's own Torch/ONNX exporters use ``qmax=448``. + + Since max calibration already encoded ``amax`` in every generated scale, + multiplying the legacy scales by ``127**2 / (448*qmax)`` recovers the + direct-amax result without rerunning calibration. Weight scales always use + the full exact range; activation qmax is exposed so held-out headroom can be + measured rather than guessed. + """ + + if not (1.0 <= activation_qmax <= 448.0) or not np.isfinite(activation_qmax): + raise ValueError("--fp8-activation-qmax must be finite and in [1, 448]") + model = onnx.load(model_path, load_external_data=False) + producers = { + output: node for node in model.graph.node for output in node.output if output + } + selected = set(selected_names) + roles: Dict[str, set[str]] = {} + for node in model.graph.node: + if node.name not in selected: + continue + if len(node.input) < 2: + raise RuntimeError(f"Selected node lacks data/weight inputs: {node.name}") + for input_index, role in ((0, "activation"), (1, "weight")): + dq = producers.get(node.input[input_index]) + if dq is None or dq.op_type != "DequantizeLinear" or len(dq.input) < 2: + raise RuntimeError( + f"Selected {role} input is not produced by DequantizeLinear: {node.name}" + ) + roles.setdefault(dq.input[1], set()).add(role) + + initializer_by_name = {item.name: item for item in model.graph.initializer} + legacy_effective_qmax = (127.0 * 127.0) / 448.0 + factors = { + "activation": legacy_effective_qmax / float(activation_qmax), + "weight": legacy_effective_qmax / 448.0, + } + counts = {"activation": 0, "weight": 0} + scale_elements = {"activation": 0, "weight": 0} + base_dir = str(Path(model_path).resolve().parent) + for scale_name, scale_roles in sorted(roles.items()): + if len(scale_roles) != 1: + raise RuntimeError( + f"FP8 scale {scale_name!r} is shared across incompatible roles: " + f"{sorted(scale_roles)}" + ) + role = next(iter(scale_roles)) + initializer = initializer_by_name.get(scale_name) + if initializer is None: + raise RuntimeError(f"FP8 scale is not a constant initializer: {scale_name}") + values = onnx.numpy_helper.to_array(initializer, base_dir=base_dir) + if values.size == 0 or not np.all(np.isfinite(values)) or np.any(values <= 0.0): + raise RuntimeError(f"FP8 scale must be finite and positive: {scale_name}") + corrected = np.asarray(values * factors[role], dtype=values.dtype) + if not np.all(np.isfinite(corrected)) or np.any(corrected <= 0.0): + raise RuntimeError(f"Corrected FP8 scale is invalid: {scale_name}") + initializer.CopyFrom(onnx.numpy_helper.from_array(corrected, name=scale_name)) + counts[role] += 1 + scale_elements[role] += int(values.size) + + if counts["activation"] == 0 or counts["weight"] == 0: + raise RuntimeError( + f"Did not find both activation and weight FP8 scales: {counts}" + ) + temporary = model_path + ".direct-amax.tmp.onnx" + onnx.save_model(model, temporary) + os.replace(temporary, model_path) + return { + "mode": "direct-amax", + "formula": "legacy_scale * 127^2 / (448 * qmax)", + "modelopt_legacy_effective_qmax": legacy_effective_qmax, + "activation_qmax": float(activation_qmax), + "weight_qmax": 448.0, + "scale_initializer_count": counts, + "scale_element_count": scale_elements, + "multipliers": factors, + "note": ( + "Accuracy-recovery experiment correcting ModelOpt 0.45 ONNX PTQ to direct " + "E4M3 amax scaling; modelopt-legacy remains available for A/B comparison." + ), + } + + +def _require_callable_parameters( + callable_object: Any, label: str, required: Sequence[str], runtime_version: str +) -> str: + """Fail clearly when a newer ORT changes an internal API we deliberately pin.""" + + try: + signature = inspect.signature(callable_object) + except (TypeError, ValueError) as exc: + raise RuntimeError( + f"Cannot inspect {label} with ONNX Runtime {runtime_version}; " + "incremental Q/DQ is disabled rather than guessing an internal API" + ) from exc + missing = sorted(set(required) - set(signature.parameters)) + if missing: + raise RuntimeError( + f"Unsupported ONNX Runtime {runtime_version} {label} signature {signature}; " + f"missing required parameters {missing}. Test and update the pinned incremental path." + ) + return str(signature) + + +def _incremental_calibration_providers( + specification: str, +) -> tuple[List[Any], List[str]]: + import onnxruntime as ort # type: ignore[import-untyped] + + available = set(ort.get_available_providers()) + providers: List[Any] = [] + skipped: List[str] = [] + for raw in specification.split(","): + token = raw.strip() + if not token: + continue + lowered = token.lower() + if lowered == "cpu": + provider: Any = "CPUExecutionProvider" + elif lowered.startswith("cuda"): + device_id = int(token.split(":", 1)[1]) if ":" in token else 0 + provider = ("CUDAExecutionProvider", {"device_id": device_id}) + elif lowered.startswith("dml"): + device_id = int(token.split(":", 1)[1]) if ":" in token else 0 + provider = ("DmlExecutionProvider", {"device_id": device_id}) + elif lowered == "trt": + provider = "TensorrtExecutionProvider" + else: + raise ValueError(f"Unsupported incremental calibration EP {token!r}") + provider_name = provider[0] if isinstance(provider, tuple) else provider + if provider_name not in available: + skipped.append(provider_name) + continue + if provider not in providers: + providers.append(provider) + if not providers: + raise RuntimeError( + "None of the requested incremental calibration execution providers are " + f"available; requested={specification!r}, available={sorted(available)}" + ) + return providers, sorted(set(skipped)) + + +def _node_filtered_calibration_tensors( + model: onnx.ModelProto, selected_names: Sequence[str] +) -> tuple[set[str], Dict[str, Any]]: + """ModelOpt 0.45's node-name calibration selection without global monkeypatching.""" + + selected = set(selected_names) + value_infos = {value.name: value for value in model.graph.value_info} + value_infos.update({value.name: value for value in model.graph.output}) + value_infos.update({value.name: value for value in model.graph.input}) + initializer_names = {item.name for item in model.graph.initializer} + tensors: set[str] = set() + found: set[str] = set() + allowed_types = {onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16} + for node in model.graph.node: + if node.name not in selected: + continue + found.add(node.name) + for tensor_name in list(node.input) + list(node.output): + value = value_infos.get(tensor_name) + if value is None or tensor_name in initializer_names: + continue + tensor_type = value.type.tensor_type + if tensor_type.elem_type in allowed_types: + tensors.add(tensor_name) + missing = sorted(selected - found) + if missing: + raise RuntimeError( + f"Incremental calibration nodes disappeared after shape inference: {missing[:10]}" + ) + if not tensors: + raise RuntimeError( + "Incremental calibration selected no floating activation tensors" + ) + return tensors, value_infos + + +def _selected_nodes_with_direct_qdq( + model_path: str, selected_names: Sequence[str] +) -> List[str]: + model = onnx.load(model_path, load_external_data=False) + producers = { + output: node for node in model.graph.node for output in node.output if output + } + result: List[str] = [] + selected = set(selected_names) + for node in model.graph.node: + if node.name not in selected or len(node.input) < 2: + continue + complete = True + for input_name in node.input[:2]: + dq_node = producers.get(input_name) + if ( + dq_node is None + or dq_node.op_type != "DequantizeLinear" + or not dq_node.input + ): + complete = False + break + q_node = producers.get(dq_node.input[0]) + if q_node is None or q_node.op_type != "QuantizeLinear": + complete = False + break + if complete: + result.append(node.name) + return sorted(result) + + +def _run_incremental_ort_quantization( + source_path: str, + output_path: str, + calibration_dataset: PositionDataset, + selected_names: Sequence[str], + selected_op_types: Sequence[str], + args: argparse.Namespace, + staging_dir: str, +) -> Dict[str, Any]: + """Quantize only new nodes in an existing Q/DQ graph. + + ModelOpt 0.45's public INT8 entry point skips every graph containing Q/DQ. + This follows its ORT configuration while using local calibrator subclasses, + avoiding ModelOpt's process-global ORT monkeypatches. + """ + + import onnxruntime as ort # type: ignore[import-untyped] + from onnxruntime.quantization.calibrate import ( # type: ignore[import-untyped] + EntropyCalibrater, + HistogramCollector, + MinMaxCalibrater, + TensorsData, + ) + from onnxruntime.quantization.qdq_quantizer import ( # type: ignore[import-untyped] + QDQQuantizer, + ) + from onnxruntime.quantization.quant_utils import ( # type: ignore[import-untyped] + QuantType, + add_infer_metadata, + ) + + runtime_version = str(ort.__version__) + signatures = { + "QDQQuantizer": _require_callable_parameters( + QDQQuantizer, + "QDQQuantizer", + ( + "model", + "per_channel", + "reduce_range", + "weight_qType", + "activation_qType", + "tensors_range", + "nodes_to_quantize", + "nodes_to_exclude", + "op_types_to_quantize", + "extra_options", + ), + runtime_version, + ), + "EntropyCalibrater": _require_callable_parameters( + EntropyCalibrater, + "EntropyCalibrater", + ( + "model_path", + "op_types_to_calibrate", + "augmented_model_path", + "use_external_data_format", + "symmetric", + "num_bins", + "num_quantized_bins", + ), + runtime_version, + ), + "MinMaxCalibrater": _require_callable_parameters( + MinMaxCalibrater, + "MinMaxCalibrater", + ( + "model_path", + "op_types_to_calibrate", + "augmented_model_path", + "use_external_data_format", + "symmetric", + ), + runtime_version, + ), + } + if args.calibrate_per_node: + raise RuntimeError( + "--calibrate-per-node is not supported with --preserve-existing-qdq; " + "the incremental calibrator is already restricted to the selected nodes" + ) + providers, skipped_providers = _incremental_calibration_providers( + args.calibration_eps + ) + + selected_set = set(selected_names) + + class _NodeFilterMixin: + def select_tensors_to_calibrate(self, model): + return _node_filtered_calibration_tensors(model, selected_names) + + class _StreamingEntropy(_NodeFilterMixin, EntropyCalibrater): + def collect_data(self, data_reader): + collected = False + output_names = [item.name for item in self.infer_session.get_outputs()] + input_names = {item.name for item in self.infer_session.get_inputs()} + while True: + inputs = data_reader.get_next() + if not inputs: + break + outputs = self.infer_session.run(None, inputs) + values = { + name: [np.copy(value) if name in input_names else value] + for name, value in zip(output_names, outputs) + if name in self.tensors_to_calibrate + } + if not self.collector: + self.collector = HistogramCollector( + method=self.method, + symmetric=self.symmetric, + num_bins=self.num_bins, + num_quantized_bins=self.num_quantized_bins, + percentile=self.percentile, + scenario=self.scenario, + ) + self.collector.collect(values) + collected = True + if not collected: + raise ValueError( + "No incremental entropy calibration data was collected" + ) + + class _StreamingMinMax(_NodeFilterMixin, MinMaxCalibrater): + def collect_data(self, data_reader): + collected = False + while True: + inputs = data_reader.get_next() + if not inputs: + break + self.intermediate_outputs.append(self.infer_session.run(None, inputs)) + result = self.compute_data() + if not isinstance(result, TensorsData): + raise TypeError( + f"Expected TensorsData from incremental min/max calibration, got {type(result)}" + ) + self.clear_collected_data() + collected = True + if not collected: + raise ValueError( + "No incremental min/max calibration data was collected" + ) + + calibration_dir = os.path.join(staging_dir, "incremental-calibration") + os.makedirs(calibration_dir, exist_ok=False) + augmented_path = os.path.join(calibration_dir, "augmented.onnx") + calibrator_cls = ( + _StreamingEntropy if args.calibration_method == "entropy" else _StreamingMinMax + ) + calibrator_kwargs: Dict[str, Any] = dict( + model_path=source_path, + # Selection is by node name in the local mixin, matching ModelOpt 0.45. + op_types_to_calibrate=list(selected_names), + augmented_model_path=augmented_path, + use_external_data_format=True, + symmetric=False, + ) + if args.calibration_method == "entropy": + calibrator_kwargs.update(num_bins=128, num_quantized_bins=128) + calibrator = calibrator_cls(**calibrator_kwargs) + calibrator.augment_graph() + calibrator.set_execution_providers(providers) + reader = ArrayCalibrationDataReader(calibration_dataset) + reader.rewind() + calibrator.collect_data(reader) + tensors_range = calibrator.compute_data() + if not isinstance(tensors_range, TensorsData): + raise TypeError( + f"Expected TensorsData from incremental calibration, got {type(tensors_range)}" + ) + calibration_tensor_names = sorted(tensors_range.data) + expected_tensors, _ = _node_filtered_calibration_tensors( + calibrator.model, selected_names + ) + unexpected_calibrated = sorted(set(calibration_tensor_names) - expected_tensors) + if unexpected_calibrated: + raise RuntimeError( + "Incremental calibrator collected tensors outside the selected nodes: " + f"{unexpected_calibrated[:10]}" + ) + + # The calibrator's ModelProto may retain external-data locations relative + # to its temporary augmented path. Keep only its proven shape information, + # then reload and fully materialize the immutable parent from its real path. + # A b15 parent is ~813 MiB, comfortably below the 64 GiB research host, and + # the resulting output is self-contained instead of sharing a fragile link. + inferred_value_info = [ + value.SerializeToString() for value in calibrator.model.graph.value_info + ] + calibrator.infer_session = None + del calibrator + gc.collect() + shutil.rmtree(calibration_dir) + + quantization_model = onnx.load(source_path, load_external_data=True) + existing_value_info = {value.name for value in quantization_model.graph.value_info} + for payload in inferred_value_info: + value = onnx.ValueInfoProto() + value.ParseFromString(payload) + if value.name not in existing_value_info: + quantization_model.graph.value_info.append(value) + existing_value_info.add(value.name) + add_infer_metadata(quantization_model) + onnx.external_data_helper.convert_model_from_external_data(quantization_model) + remaining_external = [ + initializer.name + for initializer in quantization_model.graph.initializer + if initializer.data_location == onnx.TensorProto.EXTERNAL + or initializer.external_data + ] + if remaining_external: + raise RuntimeError( + "Failed to materialize parent external data before incremental Q/DQ: " + f"{remaining_external[:10]}" + ) + materialized_raw_bytes = int( + sum( + len(initializer.raw_data) + for initializer in quantization_model.graph.initializer + ) + ) + + graph_op_types = sorted({node.op_type for node in quantization_model.graph.node}) + extra_options = { + "QuantizeBias": False, + "ActivationSymmetric": True, + "WeightSymmetric": True, + "OpTypesToExcludeOutputQuantization": graph_op_types, + "AddQDQPairToWeight": True, + "QDQOpTypePerChannelSupportToAxis": {"Conv": 0, "ConvTranspose": 1}, + "DedicatedQDQPair": False, + "ForceQuantizeNoInputCheck": True, + "QDQDisableWeightAdjustForInt32Bias": True, + } + quantizer = QDQQuantizer( + model=quantization_model, + per_channel=True, + reduce_range=False, + weight_qType=QuantType.QInt8, + activation_qType=QuantType.QInt8, + tensors_range=tensors_range, + nodes_to_quantize=list(selected_names), + nodes_to_exclude=[], + op_types_to_quantize=list(selected_op_types), + extra_options=extra_options, + ) + quantizer.quantize_model() + save_method = quantizer.model.save_model_to_file + signatures["save_model_to_file"] = _require_callable_parameters( + save_method, + "ONNXModel.save_model_to_file", + ("output_path", "use_external_data_format"), + runtime_version, + ) + save_method(output_path=output_path, use_external_data_format=True) + del quantizer + del quantization_model + gc.collect() + + newly_quantized = _selected_nodes_with_direct_qdq(output_path, selected_names) + if newly_quantized != sorted(selected_set): + raise RuntimeError( + "Incremental ORT quantization produced no usable change; refusing to promote " + f"a parent clone. expected={sorted(selected_set)}, actual={newly_quantized}" + ) + return { + "backend": "onnxruntime-node-filtered-qdq", + "onnxruntime_version": runtime_version, + "api_signatures": signatures, + "providers": providers, + "skipped_unavailable_providers": skipped_providers, + "calibration_tensor_count": len(calibration_tensor_names), + "calibration_tensor_names": calibration_tensor_names, + "newly_quantized_nodes": newly_quantized, + "materialized_parent_raw_bytes": materialized_raw_bytes, + "self_contained_output": True, + "configuration": { + "calibration_method": args.calibration_method, + "activation_type": "QInt8", + "weight_type": "QInt8", + "activation_symmetric": True, + "weight_symmetric": True, + "per_channel_weights": True, + "reduce_range": False, + "quantize_bias": False, + "output_quantization": False, + }, + } + + +def _quantize_one( + mode: str, + source_path: str, + output_path: str, + calibration_dataset: PositionDataset, + selected_names: Sequence[str], + selected_op_types: Sequence[str], + args: argparse.Namespace, + quantize, + preservation_snapshot: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + if any( + item.strip().lower() != "cpu" + for item in args.calibration_eps.split(",") + if item.strip() + ): + _preload_ort_gpu_dependencies() + if os.path.exists(output_path) and not args.overwrite: + raise FileExistsError(f"Refusing to replace {output_path}; use --overwrite") + staging_dir = tempfile.mkdtemp( + prefix=f".katago-quant-{args.output_prefix}.{mode}-", + dir=args.output_dir, + ) + details: Dict[str, Any] = {} + try: + # ModelOpt shape-infers its input in place, so normal PTQ requires a + # complete clone. The incremental ORT path is read-only and uses the + # original parent directly, avoiding another ~800 MiB full315 copy. + staged_source = ( + source_path + if preservation_snapshot is not None + else _copy_onnx_artifact(source_path, staging_dir) + ) + token = uuid.uuid4().hex + staged_output = os.path.join( + staging_dir, + f"{Path(output_path).stem}.{token}.onnx", + ) + reader = ArrayCalibrationDataReader(calibration_dataset) + exact_node_patterns = [f"^{re.escape(name)}$" for name in selected_names] + calibration_method = ( + args.fp8_calibration_method if mode == "fp8" else args.calibration_method + ) + kwargs = dict( + # ModelOpt 0.45 shape-infers this path in place, so it must always + # receive the complete staged copy rather than the user's dump. + onnx_path=staged_source, + quantize_mode=mode, + calibration_data_reader=reader, + calibration_method=calibration_method, + calibration_eps=[ + item.strip() for item in args.calibration_eps.split(",") if item.strip() + ], + op_types_to_quantize=list(selected_op_types), + nodes_to_quantize=exact_node_patterns, + use_external_data_format=True, + keep_intermediate_files=args.keep_intermediate_files, + output_path=staged_output, + log_level="INFO", + log_file=os.path.join( + args.output_dir, f"{args.output_prefix}.{mode}.modelopt.log" + ), + high_precision_dtype=args.high_precision, + mha_accumulation_dtype="fp32", + disable_mha_qdq=True, + use_zero_point=False, + passes=[], + simplify=False, + calibrate_per_node=args.calibrate_per_node, + direct_io_types=False, + # The selected KataGo projections operate on H*W tokens and are not + # GEMV even at batch 1. Keep ModelOpt's generic shape heuristic from + # silently overriding the explicit, audited node allowlist. + enable_gemv_detection_for_trt=False, + ) + if mode == "fp8": + kwargs["opset"] = 21 + logging.info( + "Quantizing %d %s nodes to %s with %s calibration and %s fallback in %s", + len(selected_names), + "/".join(selected_op_types), + mode.upper(), + calibration_method, + args.high_precision.upper(), + staging_dir, + ) + if preservation_snapshot is not None: + details["incremental_quantizer"] = _run_incremental_ort_quantization( + staged_source, + staged_output, + calibration_dataset, + selected_names, + selected_op_types, + args, + staging_dir, + ) + else: + quantize(**kwargs) + if not os.path.isfile(staged_output): + raise RuntimeError(f"ModelOpt returned without creating {staged_output}") + if mode == "fp8" and args.fp8_scale_mode == "direct-amax": + details["fp8_scale_rewrite"] = _rewrite_fp8_direct_amax_scales( + staged_output, + selected_names, + activation_qmax=args.fp8_activation_qmax, + ) + if preservation_snapshot is not None: + union_names = sorted( + set(preservation_snapshot["selected_names"]) | set(selected_names) + ) + staged_union_audit = audit_qdq_model(staged_output, union_names) + if staged_union_audit["errors"]: + raise RuntimeError( + "Incremental Q/DQ union audit failed before artifact promotion: " + f"{staged_union_audit['errors'][:10]}" + ) + preservation = compare_existing_qdq_state( + staged_output, preservation_snapshot + ) + if preservation["status"] != "passed": + raise RuntimeError( + "ModelOpt changed pre-existing Q/DQ state; refusing to promote artifact: " + f"{preservation['differences'][:10]}" + ) + details["incremental_union_selected_count"] = len(union_names) + details["incremental_union_selected_names"] = union_names + details["incremental_staged_qdq_audit"] = staged_union_audit + details["existing_qdq_preservation"] = preservation + details.update( + _promote_staged_artifact( + staged_output, + output_path, + overwrite=args.overwrite, + ) + ) + if args.keep_intermediate_files: + details["modelopt_intermediate_directory"] = staging_dir + return details + finally: + if not args.keep_intermediate_files: + staging_path = Path(staging_dir).resolve() + output_root = Path(args.output_dir).resolve() + if staging_path.parent != output_root or not staging_path.name.startswith( + ".katago-quant-" + ): + raise RuntimeError( + f"Refusing to clean unsafe staging path: {staging_path}" + ) + try: + shutil.rmtree(staging_path) + except FileNotFoundError: + pass + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + args.onnx_input = str(Path(args.onnx_input).resolve()) + args.output_dir = str(Path(args.output_dir).resolve()) + if args.max_source_files < 0: + parser.error("--max-source-files must be nonnegative (0 means unlimited)") + if args.expected_existing_qdq_nodes is not None and not args.preserve_existing_qdq: + parser.error("--expected-existing-qdq-nodes requires --preserve-existing-qdq") + if args.preserve_existing_qdq: + if set(args.formats) != {"int8"}: + parser.error( + "--preserve-existing-qdq currently supports exactly --formats int8" + ) + if not args.only_node_regex: + parser.error( + "--preserve-existing-qdq requires --only-node-regex so ModelOpt receives " + "an explicit incremental node allowlist" + ) + if args.expected_existing_qdq_nodes is None: + parser.error( + "--preserve-existing-qdq requires --expected-existing-qdq-nodes" + ) + if args.expected_existing_qdq_nodes <= 0: + parser.error("--expected-existing-qdq-nodes must be positive") + if ( + "fp8" in args.formats + and args.fp8_calibration_method != "max" + and not args.allow_fp8_nonmax_calibration + ): + parser.error( + "ModelOpt 0.45 FP8 conversion is documented for max-calibrated INT8 scales. " + "Use --fp8-calibration-method max, or explicitly acknowledge an unsupported " + "experiment with --allow-fp8-nonmax-calibration." + ) + if not np.isfinite(args.fp8_activation_qmax) or not ( + 1.0 <= args.fp8_activation_qmax <= 448.0 + ): + parser.error("--fp8-activation-qmax must be finite and in [1, 448]") + if ( + "fp8" in args.formats + and args.fp8_scale_mode == "direct-amax" + and args.fp8_calibration_method != "max" + ): + parser.error("--fp8-scale-mode direct-amax requires max calibration") + if args.high_precision == "fp16" and not args.allow_global_fp16_fallback: + parser.error( + "--high-precision fp16 globally converts fallback ops, including sensitive heads/norms. " + "Use FP32 for quantization-only error, or explicitly add " + "--allow-global-fp16-fallback for a separately gated experiment." + ) + if not os.path.isfile(args.onnx_input): + parser.error(f"ONNX input does not exist: {args.onnx_input}") + os.makedirs(args.output_dir, exist_ok=True) + if args.output_prefix is None: + args.output_prefix = Path(args.onnx_input).stem + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._+-]*", args.output_prefix): + parser.error( + "--output-prefix must be a filename-safe model name containing only " + "letters, digits, dot, underscore, plus, or minus" + ) + planned_output_paths = { + mode: str( + (Path(args.output_dir) / f"{args.output_prefix}.{mode}.qdq.onnx").resolve() + ) + for mode in dict.fromkeys(args.formats) + } + source_key = os.path.normcase(args.onnx_input) + colliding_formats = [ + mode + for mode, path in planned_output_paths.items() + if os.path.normcase(path) == source_key + ] + if colliding_formats: + parser.error( + "Quantized output would overwrite the source ONNX artifact for format(s) " + f"{colliding_formats}. Choose a different --output-dir or --output-prefix." + ) + report_path = os.path.join( + args.output_dir, f"{args.output_prefix}.quantization-report.json" + ) + if os.path.exists(report_path) and not args.overwrite: + raise FileExistsError(f"Refusing to replace {report_path}; use --overwrite") + log_path = _setup_logging(args.output_dir, args.output_prefix) + + gpu_summary = _gpu_summary() + report: Dict[str, Any] = { + "schema_version": 1, + "created_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "status": "running", + "research_only": True, + "runtime_integration": ( + "Current KataGo loads FP32 .bin.gz and builds its own ONNX. These Q/DQ ONNX artifacts " + "must be benchmarked with a standalone strongly typed TensorRT build until a separate " + "runtime integration is implemented." + ), + "command": [sys.executable, str(Path(__file__).resolve())] + + list(argv or sys.argv[1:]), + "environment": { + "python": sys.version, + "platform": platform.platform(), + "numpy": np.__version__, + "onnx": onnx.__version__, + "packages": _optional_package_versions(), + "git_revision": _git_revision(), + "gpus": gpu_summary, + "fp8_hardware": _fp8_hardware_summary(gpu_summary), + }, + "options": { + key: value + for key, value in vars(args).items() + if key not in ("calibration_data", "validation_data") + }, + "log_path": log_path, + "formats": {}, + } + + source_manifest: Optional[Dict[str, Any]] = None + existing_qdq_state: Optional[Dict[str, Any]] = None + try: + if not args.skip_onnx_check: + logging.info("Running ONNX checker on source graph") + onnx.checker.check_model(args.onnx_input, full_check=False) + logging.info("Loading source ONNX: %s", args.onnx_input) + source_model = onnx.load(args.onnx_input, load_external_data=False) + input_specs = validate_katago_io_contract( + source_model, + require_producer_metadata=not args.preserve_existing_qdq, + ) + source_manifest = artifact_manifest(args.onnx_input, source_model) + report["source"] = source_manifest + + if args.preserve_existing_qdq: + logging.info("Strictly auditing pre-existing weighted Q/DQ state") + existing_qdq_state = capture_existing_qdq_state(args.onnx_input) + if existing_qdq_state["selected_count"] != args.expected_existing_qdq_nodes: + raise RuntimeError( + "Found " + f"{existing_qdq_state['selected_count']} pre-existing fully quantized " + f"weighted nodes, expected {args.expected_existing_qdq_nodes}. " + "Refusing incremental quantization of a changed parent graph." + ) + existing_audit = audit_qdq_model( + args.onnx_input, existing_qdq_state["selected_names"] + ) + if existing_audit["errors"]: + raise RuntimeError( + "Pre-existing Q/DQ semantic audit failed: " + f"{existing_audit['errors'][:10]}" + ) + non_int8_types = { + qtype: count + for qtype, count in existing_audit["quantized_tensor_types"].items() + if qtype != "INT8" + } + if non_int8_types or not existing_audit["quantized_tensor_types"].get( + "INT8" + ): + raise RuntimeError( + "--preserve-existing-qdq requires a purely INT8 parent graph; " + f"quantized tensor types={existing_audit['quantized_tensor_types']}" + ) + report["existing_qdq"] = { + "state": existing_qdq_state, + "audit": existing_audit, + } + + selection = select_quantizable_nodes( + source_model, + scope=args.scope, + include_regexes=args.include_node_regex, + exclude_regexes=args.exclude_node_regex, + only_regexes=args.only_node_regex, + ) + report["node_selection"] = selection.manifest() + logging.info( + "Selected %d/%d weighted nodes: %s", + len(selection.selected_names), + len(selection.weighted_candidate_names), + selection.selected_by_op_type, + ) + if ( + args.expected_quantized_nodes is not None + and len(selection.selected_names) != args.expected_quantized_nodes + ): + raise RuntimeError( + f"Selected {len(selection.selected_names)} nodes, expected " + f"{args.expected_quantized_nodes}. Refusing to quantize a changed graph." + ) + if existing_qdq_state is not None: + overlap = sorted( + set(existing_qdq_state["selected_names"]) + & set(selection.selected_names) + ) + if overlap: + raise RuntimeError( + "Incremental selection overlaps pre-existing Q/DQ nodes: " + f"{overlap[:10]}" + ) + union_names = sorted( + set(existing_qdq_state["selected_names"]) + | set(selection.selected_names) + ) + report["union_node_selection"] = { + "existing_count": len(existing_qdq_state["selected_names"]), + "incremental_count": len(selection.selected_names), + "selected_count": len(union_names), + "selected_names": union_names, + } + # The b15 graph contains roughly 680 MB of FP32 initializers. ModelOpt + # loads its own graph copies, so release this inspection copy before + # calibration/quantization to keep peak host memory under control. + del source_model + gc.collect() + + calibration_files = resolve_npz_files(args.calibration_data) + validation_files = ( + resolve_npz_files(args.validation_data) if args.validation_data else [] + ) + overlap = sorted( + set(map(os.path.normcase, calibration_files)) + & set(map(os.path.normcase, validation_files)) + ) + if overlap and not args.allow_data_overlap: + raise RuntimeError( + "Calibration and validation share NPZ shards. Use disjoint held-out files; " + f"overlap={overlap[:10]}" + ) + if not args.skip_validation and not validation_files: + raise RuntimeError( + "--validation-data is required unless --skip-validation is explicitly set" + ) + + logging.info("Sampling %d calibration positions", args.calibration_samples) + calibration_dataset = load_position_dataset( + calibration_files, + input_specs, + sample_count=args.calibration_samples, + batch_size=args.batch_size, + seed=args.seed, + history_mode=args.history_mode, + symmetry_mode=args.symmetry_mode, + max_source_files=args.max_source_files, + ) + report["calibration_dataset"] = calibration_dataset.manifest() + + validation_dataset: Optional[PositionDataset] = None + reference_batches: Optional[List[Dict[str, np.ndarray]]] = None + if not args.skip_validation: + logging.info( + "Sampling %d held-out validation positions", args.validation_samples + ) + validation_dataset = load_position_dataset( + validation_files, + input_specs, + sample_count=args.validation_samples, + batch_size=args.batch_size, + seed=args.seed + 1, + history_mode=args.history_mode, + symmetry_mode=args.symmetry_mode, + max_source_files=args.max_source_files, + ) + report["validation_dataset"] = validation_dataset.manifest() + logging.info("Computing FP32 reference outputs once") + reference_batches, reference_execution = _run_ort_model( + args.onnx_input, + validation_dataset, + args.validation_ep, + args.ort_intra_op_threads, + args.allow_validation_ep_fallback, + ) + report["reference_execution"] = reference_execution + + quantize, modelopt_version = _load_modelopt(args.allow_unpinned_modelopt) + report["environment"]["nvidia_modelopt"] = modelopt_version + thresholds = _accuracy_thresholds(args) + if not any(value is not None for value in thresholds.values()): + logging.warning( + "No numerical accuracy thresholds were supplied. Metrics will be reported, " + "but the artifact must not be called release-qualified." + ) + + any_failure = False + for mode in dict.fromkeys(args.formats): + output_path = planned_output_paths[mode] + mode_report: Dict[str, Any] = {"status": "running", "path": output_path} + mode_report["calibration_method"] = ( + args.fp8_calibration_method + if mode == "fp8" + else args.calibration_method + ) + report["formats"][mode] = mode_report + json_dump(report_path, report) + try: + quant_details = _quantize_one( + mode, + args.onnx_input, + output_path, + calibration_dataset, + selection.selected_names, + sorted(selection.selected_by_op_type), + args, + quantize, + existing_qdq_state, + ) + mode_report.update(quant_details) + if not args.skip_onnx_check: + logging.info("Running ONNX checker on %s", output_path) + onnx.checker.check_model(output_path, full_check=False) + mode_report["artifact"] = artifact_manifest(output_path) + audited_names = ( + report["union_node_selection"]["selected_names"] + if existing_qdq_state is not None + else selection.selected_names + ) + mode_report["qdq_audit"] = audit_qdq_model(output_path, audited_names) + if mode_report["qdq_audit"]["non_float_io"]: + raise RuntimeError( + f"Quantization changed graph I/O types: " + f"{mode_report['qdq_audit']['non_float_io']}" + ) + if mode_report["qdq_audit"]["quantize_linear_count"] == 0: + raise RuntimeError( + "Quantized graph contains no QuantizeLinear nodes" + ) + audit = mode_report["qdq_audit"] + if audit["errors"]: + raise RuntimeError( + f"Q/DQ semantic audit failed: {audit['errors'][:10]}" + ) + if audit["selected_nodes_missing_after_quantization"]: + raise RuntimeError( + "ModelOpt removed or renamed explicitly selected nodes: " + f"{audit['selected_nodes_missing_after_quantization'][:10]}" + ) + if audit["selected_nodes_without_qdq_inputs"]: + raise RuntimeError( + "Some selected nodes do not have Q/DQ on both data and weight inputs: " + f"{audit['selected_nodes_without_qdq_inputs'][:10]}" + ) + if audit["scale_granularity_errors"]: + raise RuntimeError( + "Unexpected Q/DQ scale granularity: " + f"{audit['scale_granularity_errors'][:10]}" + ) + expected_qtype = "INT8" if mode == "int8" else "FLOAT8E4M3FN" + unexpected_qtypes = { + qtype: count + for qtype, count in audit["quantized_tensor_types"].items() + if qtype != expected_qtype + } + if unexpected_qtypes: + raise RuntimeError( + f"{mode.upper()} graph contains unexpected quantized tensor types: " + f"{unexpected_qtypes}" + ) + + if validation_dataset is not None and reference_batches is not None: + candidate_batches, candidate_execution = _run_ort_model( + output_path, + validation_dataset, + args.validation_ep, + args.ort_intra_op_threads, + args.allow_validation_ep_fallback, + ) + mode_report["validation_execution"] = candidate_execution + metrics = _validation_metrics( + reference_batches, candidate_batches, validation_dataset + ) + mode_report["validation"] = metrics + mode_report["accuracy_gate"] = evaluate_accuracy_gates( + metrics, thresholds + ) + + if args.trtexec: + mode_report["trtexec"] = _run_trtexec( + args.trtexec, + output_path, + input_specs, + args.trt_opt_batch, + args.trt_max_batch, + args.trt_workspace_mib, + args.trt_timeout_seconds, + args.overwrite, + ) + mode_failed = _mode_report_failed(mode_report) + any_failure = any_failure or mode_failed + mode_report["status"] = "failed" if mode_failed else "complete" + except Exception as exc: + any_failure = True + mode_report["status"] = "error" + mode_report["error"] = str(exc) + mode_report["traceback"] = traceback.format_exc() + logging.exception("%s quantization failed", mode.upper()) + json_dump(report_path, report) + if not args.continue_on_error: + raise + finally: + json_dump(report_path, report) + + assert source_manifest is not None + report["source_integrity"] = _artifact_integrity(source_manifest) + if report["source_integrity"]["status"] != "passed": + raise RuntimeError( + "Source ONNX artifact changed during quantization; see source_integrity" + ) + report["status"] = "failed" if any_failure else "complete" + json_dump(report_path, report) + logging.info("Wrote reproducibility/accuracy report: %s", report_path) + return 2 if any_failure else 0 + except Exception as exc: + if source_manifest is not None and "source_integrity" not in report: + report["source_integrity"] = _artifact_integrity(source_manifest) + report["status"] = "error" + report["error"] = str(exc) + if ( + source_manifest is not None + and report["source_integrity"]["status"] != "passed" + and "Source ONNX artifact changed" not in report["error"] + ): + report["error"] += ( + "; source ONNX artifact also changed during the failed run; " + "see source_integrity" + ) + report["traceback"] = traceback.format_exc() + json_dump(report_path, report) + logging.exception("Quantization pipeline failed") + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/requirements-quantization.txt b/python/requirements-quantization.txt new file mode 100644 index 0000000000..d8fa89d094 --- /dev/null +++ b/python/requirements-quantization.txt @@ -0,0 +1,3 @@ +# Reproducible ONNX post-training quantization toolchain. +# Model Optimizer's ONNX extra supplies its compatible ONNX Runtime and graph tools. +nvidia-modelopt[onnx]==0.45.0 diff --git a/python/tests/test_quantization_qdq_audit.py b/python/tests/test_quantization_qdq_audit.py new file mode 100644 index 0000000000..c933e4d1f7 --- /dev/null +++ b/python/tests/test_quantization_qdq_audit.py @@ -0,0 +1,543 @@ +"""Adversarial tests for the TensorRT explicit Q/DQ graph audit.""" + +from __future__ import annotations + +import numpy as np +import pytest + + +onnx = pytest.importorskip("onnx") +quantization = pytest.importorskip("katago.quantization") + + +def _make_qdq_model( + *, + op_type: str = "MatMul", + output_channels: int = 3, + trans_b: int = 0, + scalar_weight_scale: bool = False, +): + helper = onnx.helper + tensor_proto = onnx.TensorProto + + if op_type == "Conv": + input_shape = [None, 4, 5, 5] + weight_shape = [output_channels, 4, 1, 1] + weight_axis = 0 + elif op_type == "Gemm" and trans_b: + input_shape = [None, 4] + weight_shape = [output_channels, 4] + weight_axis = 0 + else: + input_shape = [None, 4] + weight_shape = [4, output_channels] + weight_axis = 1 + + weight_scale_shape = [] if scalar_weight_scale else [output_channels] + weight_scale_values = [0.05] if scalar_weight_scale else [0.05] * output_channels + weight_zero_shape = [] if scalar_weight_scale else [output_channels] + weight_zero_values = [0] if scalar_weight_scale else [0] * output_channels + + initializers = [ + helper.make_tensor( + "weight", + tensor_proto.FLOAT, + weight_shape, + np.linspace(-1.0, 1.0, int(np.prod(weight_shape)), dtype=np.float32), + ), + helper.make_tensor("act_scale", tensor_proto.FLOAT, [], [0.1]), + helper.make_tensor("act_zp", tensor_proto.INT8, [], [0]), + helper.make_tensor( + "weight_scale", tensor_proto.FLOAT, weight_scale_shape, weight_scale_values + ), + helper.make_tensor( + "weight_zp", tensor_proto.INT8, weight_zero_shape, weight_zero_values + ), + ] + nodes = [ + helper.make_node( + "QuantizeLinear", + ["input", "act_scale", "act_zp"], + ["input_q"], + name="act_q", + ), + helper.make_node( + "DequantizeLinear", + ["input_q", "act_scale", "act_zp"], + ["input_dq"], + name="act_dq", + ), + helper.make_node( + "QuantizeLinear", + ["weight", "weight_scale", "weight_zp"], + ["weight_q"], + name="weight_q", + axis=weight_axis, + ), + helper.make_node( + "DequantizeLinear", + ["weight_q", "weight_scale", "weight_zp"], + ["weight_dq"], + name="weight_dq", + axis=weight_axis, + ), + ] + attributes = {"transB": trans_b} if op_type == "Gemm" else {} + nodes.append( + helper.make_node( + op_type, + ["input_dq", "weight_dq"], + ["output"], + name="selected", + **attributes, + ) + ) + graph = helper.make_graph( + nodes, + "qdq_audit", + [helper.make_tensor_value_info("input", tensor_proto.FLOAT, input_shape)], + [helper.make_tensor_value_info("output", tensor_proto.FLOAT, None)], + initializer=initializers, + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + + +def _node(model, name: str): + return next(node for node in model.graph.node if node.name == name) + + +def _replace_initializer(model, name: str, values, dtype) -> None: + replacement = onnx.numpy_helper.from_array( + np.asarray(values, dtype=dtype), name=name + ) + initializer = next(item for item in model.graph.initializer if item.name == name) + initializer.CopyFrom(replacement) + + +def _set_axis(node, axis: int) -> None: + for attribute in node.attribute: + if attribute.name == "axis": + attribute.i = axis + return + node.attribute.append(onnx.helper.make_attribute("axis", axis)) + + +def _add_initializer(model, name: str, values, dtype) -> None: + model.graph.initializer.append( + onnx.numpy_helper.from_array(np.asarray(values, dtype=dtype), name=name) + ) + + +def _audit(tmp_path, model): + path = tmp_path / "model.onnx" + onnx.save(model, path) + return quantization.audit_qdq_model(str(path), ["selected"]) + + +@pytest.mark.parametrize( + ("op_type", "trans_b", "expected_axis"), + [("MatMul", 0, 1), ("Conv", 0, 0), ("Gemm", 0, 1), ("Gemm", 1, 0)], +) +def test_accepts_valid_qdq_and_operator_weight_axis( + tmp_path, op_type, trans_b, expected_axis +): + report = _audit(tmp_path, _make_qdq_model(op_type=op_type, trans_b=trans_b)) + + assert report["errors"] == [] + assert report["qdq_pair_count"] == 2 + assert report["selected_nodes_with_qdq_inputs"] == 1 + assert report["selected_chain_qtypes"] == { + "selected": {"activation": "INT8", "weight": "INT8"} + } + assert ( + report["selected_input_chains"]["selected"]["weight"]["effective_axis"] + == expected_axis + ) + + +def test_allows_scalar_scale_for_single_output_channel(tmp_path): + model = _make_qdq_model(op_type="Conv", output_channels=1, scalar_weight_scale=True) + + report = _audit(tmp_path, model) + + assert report["errors"] == [] + assert report["selected_input_chains"]["selected"]["weight"]["scale_elements"] == 1 + + +def test_reads_scale_and_zero_point_from_external_data_relative_to_model(tmp_path): + model = _make_qdq_model(output_channels=3) + _replace_initializer(model, "act_scale", 0.1, np.float32) + _replace_initializer(model, "act_zp", 0, np.int8) + _replace_initializer(model, "weight_scale", [0.05] * 3, np.float32) + _replace_initializer(model, "weight_zp", [0] * 3, np.int8) + model_dir = tmp_path / "nested" + model_dir.mkdir() + path = model_dir / "model.onnx" + onnx.save_model( + model, + path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="payload.bin", + size_threshold=0, + ) + unloaded = onnx.load(path, load_external_data=False) + unloaded_by_name = { + initializer.name: initializer for initializer in unloaded.graph.initializer + } + assert all( + unloaded_by_name[name].data_location == onnx.TensorProto.EXTERNAL + for name in ("act_scale", "act_zp", "weight_scale", "weight_zp") + ) + + report = quantization.audit_qdq_model(str(path), ["selected"]) + + assert report["errors"] == [] + + +def test_accepts_equal_qdq_constants_with_different_names(tmp_path): + model = _make_qdq_model() + _add_initializer(model, "act_scale_dq", 0.1, np.float32) + _add_initializer(model, "act_zp_dq", 0, np.int8) + _node(model, "act_dq").input[1] = "act_scale_dq" + _node(model, "act_dq").input[2] = "act_zp_dq" + + assert _audit(tmp_path, model)["errors"] == [] + + +def test_accepts_scale_from_constant_node(tmp_path): + model = _make_qdq_model() + scale_tensor = onnx.numpy_helper.from_array(np.asarray(0.1, dtype=np.float32)) + model.graph.node.insert( + 0, + onnx.helper.make_node( + "Constant", + [], + ["act_scale_constant"], + name="act_scale_constant", + value=scale_tensor, + ), + ) + _node(model, "act_q").input[1] = "act_scale_constant" + _node(model, "act_dq").input[1] = "act_scale_constant" + + assert _audit(tmp_path, model)["errors"] == [] + + +@pytest.mark.parametrize( + ("values", "expected_fragment"), + [ + ([np.nan], "non-finite"), + ([-0.1], "strictly positive"), + ([0.0], "strictly positive"), + ], +) +def test_rejects_nonfinite_or_nonpositive_scale(tmp_path, values, expected_fragment): + model = _make_qdq_model() + _replace_initializer(model, "act_scale", values, np.float32) + + report = _audit(tmp_path, model) + + assert any(expected_fragment in error for error in report["errors"]) + + +def test_rejects_nonconstant_scale(tmp_path): + model = _make_qdq_model() + _node(model, "act_q").input[1] = "input" + _node(model, "act_dq").input[1] = "input" + + report = _audit(tmp_path, model) + + assert any("not a readable constant" in error for error in report["errors"]) + + +def test_rejects_different_qdq_scales(tmp_path): + model = _make_qdq_model() + _add_initializer(model, "different_scale", 0.2, np.float32) + _node(model, "act_dq").input[1] = "different_scale" + + report = _audit(tmp_path, model) + + assert any("Q/DQ scales differ" in error for error in report["errors"]) + + +def test_rejects_nonzero_and_different_qdq_zero_points(tmp_path): + model = _make_qdq_model() + _add_initializer(model, "different_zp", 1, np.int8) + _node(model, "act_dq").input[2] = "different_zp" + + report = _audit(tmp_path, model) + + assert any("all-zero" in error for error in report["errors"]) + assert any("Q/DQ zero points differ" in error for error in report["errors"]) + + +def test_rejects_zero_point_with_wrong_target_dtype(tmp_path): + model = _make_qdq_model() + _replace_initializer(model, "weight_zp", [0, 0, 0], np.uint8) + _node(model, "weight_q").attribute.append( + onnx.helper.make_attribute("output_dtype", onnx.TensorProto.INT8) + ) + + report = _audit(tmp_path, model) + + assert any( + "zero point type UINT8 does not match quantized type INT8" in error + for error in report["errors"] + ) + + +def test_rejects_dq_not_directly_fed_by_q_and_reports_orphans(tmp_path): + model = _make_qdq_model() + _node(model, "act_dq").input[0] = "input" + + report = _audit(tmp_path, model) + + assert "act_q" in report["orphan_quantize_linear_nodes"] + assert "act_dq" in report["orphan_dequantize_linear_nodes"] + assert report["selected_nodes_without_qdq_inputs"] == ["selected"] + assert any( + "not produced directly by QuantizeLinear" in error for error in report["errors"] + ) + + +def test_rejects_qdq_axis_mismatch_using_default_axis_one(tmp_path): + model = _make_qdq_model(op_type="MatMul") + _set_axis(_node(model, "weight_dq"), 0) + + report = _audit(tmp_path, model) + + assert any( + "Q/DQ effective axes differ (1 vs 0)" in error for error in report["errors"] + ) + + +@pytest.mark.parametrize( + ("op_type", "trans_b", "wrong_axis", "expected_axis"), + [("MatMul", 0, 0, 1), ("Conv", 0, 1, 0), ("Gemm", 0, 0, 1), ("Gemm", 1, 1, 0)], +) +def test_rejects_wrong_weight_axis_for_operator( + tmp_path, op_type, trans_b, wrong_axis, expected_axis +): + model = _make_qdq_model(op_type=op_type, trans_b=trans_b) + _set_axis(_node(model, "weight_q"), wrong_axis) + _set_axis(_node(model, "weight_dq"), wrong_axis) + + report = _audit(tmp_path, model) + + assert any( + f"does not match expected axis {expected_axis}" in error + for error in report["errors"] + ) + + +def test_normalizes_negative_axis_before_comparison(tmp_path): + model = _make_qdq_model(op_type="MatMul") + _set_axis(_node(model, "weight_q"), -1) + _set_axis(_node(model, "weight_dq"), -1) + + report = _audit(tmp_path, model) + + assert report["errors"] == [] + assert report["selected_input_chains"]["selected"]["weight"]["effective_axis"] == 1 + + +def test_rejects_weight_scale_length_not_equal_to_output_channels(tmp_path): + model = _make_qdq_model(output_channels=3) + _replace_initializer(model, "weight_scale", [0.05, 0.05], np.float32) + _replace_initializer(model, "weight_zp", [0, 0], np.int8) + + report = _audit(tmp_path, model) + + assert any( + "scale length 2 does not match output channels 3" in error + for error in report["errors"] + ) + + +def test_rejects_per_channel_activation_scale(tmp_path): + model = _make_qdq_model() + _replace_initializer(model, "act_scale", [0.1, 0.1], np.float32) + _replace_initializer(model, "act_zp", [0, 0], np.int8) + + report = _audit(tmp_path, model) + + assert any( + "activation scale is not per-tensor" in error for error in report["errors"] + ) + + +def test_reports_selected_chain_qtype_and_rejects_mixed_types(tmp_path): + model = _make_qdq_model() + _replace_initializer(model, "weight_zp", [0, 0, 0], np.uint8) + + report = _audit(tmp_path, model) + + assert report["selected_chain_qtypes"]["selected"] == { + "activation": "INT8", + "weight": "UINT8", + } + assert any( + "unsupported quantized type UINT8" in error for error in report["errors"] + ) + assert any("selected input chains mix" in error for error in report["errors"]) + + +def test_rejects_nonselected_weighted_node_with_two_qdq_inputs(tmp_path): + model = _make_qdq_model() + model.graph.node.append( + onnx.helper.make_node( + "MatMul", + ["input_dq", "weight_dq"], + ["unexpected_output"], + name="unselected", + ) + ) + + report = _audit(tmp_path, model) + + assert report["unexpected_fully_quantized_weighted_nodes"] == ["unselected"] + assert any("non-selected weighted op" in error for error in report["errors"]) + + +def _append_unquantized_outer_p(model): + weight = onnx.numpy_helper.from_array( + np.arange(6, dtype=np.float32).reshape(3, 2), name="outer_p_weight" + ) + model.graph.initializer.append(weight) + model.graph.node.append( + onnx.helper.make_node( + "MatMul", + ["output", "outer_p_weight"], + ["outer_p_output"], + name="model.blocks.0.normactconvp.conv/nhwc", + ) + ) + model.graph.output[0].name = "outer_p_output" + + +def _quantize_outer_p_in_place(model): + outer = _node(model, "model.blocks.0.normactconvp.conv/nhwc") + activation_scale = onnx.numpy_helper.from_array( + np.asarray(0.2, dtype=np.float32), name="outer_p_act_scale" + ) + activation_zp = onnx.numpy_helper.from_array( + np.asarray(0, dtype=np.int8), name="outer_p_act_zp" + ) + weight_scale = onnx.numpy_helper.from_array( + np.asarray([0.03, 0.04], dtype=np.float32), name="outer_p_weight_scale" + ) + weight_zp = onnx.numpy_helper.from_array( + np.asarray([0, 0], dtype=np.int8), name="outer_p_weight_zp" + ) + model.graph.initializer.extend( + [activation_scale, activation_zp, weight_scale, weight_zp] + ) + nodes = [ + onnx.helper.make_node( + "QuantizeLinear", + ["output", "outer_p_act_scale", "outer_p_act_zp"], + ["outer_p_act_q_value"], + name="outer_p_act_q", + ), + onnx.helper.make_node( + "DequantizeLinear", + ["outer_p_act_q_value", "outer_p_act_scale", "outer_p_act_zp"], + ["outer_p_act_dq_value"], + name="outer_p_act_dq", + ), + onnx.helper.make_node( + "QuantizeLinear", + ["outer_p_weight", "outer_p_weight_scale", "outer_p_weight_zp"], + ["outer_p_weight_q_value"], + name="outer_p_weight_q", + axis=1, + ), + onnx.helper.make_node( + "DequantizeLinear", + ["outer_p_weight_q_value", "outer_p_weight_scale", "outer_p_weight_zp"], + ["outer_p_weight_dq_value"], + name="outer_p_weight_dq", + axis=1, + ), + ] + outer.input[0] = "outer_p_act_dq_value" + outer.input[1] = "outer_p_weight_dq_value" + outer_index = list(model.graph.node).index(outer) + for offset, node in enumerate(nodes): + model.graph.node.insert(outer_index + offset, node) + + +def test_incremental_qdq_snapshot_finds_only_existing_parent_and_freezes_bytes( + tmp_path, +): + parent = _make_qdq_model() + _append_unquantized_outer_p(parent) + parent_path = tmp_path / "parent.onnx" + onnx.save_model( + parent, + parent_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="parent.data", + size_threshold=0, + ) + + selection = quantization.select_quantizable_nodes( + parent, + scope="all-weighted", + only_regexes=(r"normactconvp\.conv",), + ) + assert selection.selected_names == ["model.blocks.0.normactconvp.conv/nhwc"] + state = quantization.capture_existing_qdq_state(str(parent_path)) + assert state["selected_names"] == ["selected"] + assert state["weighted_node_count"] == 1 + assert state["qdq_node_count"] == 4 + assert set(state["initializers"]) == { + "act_scale", + "act_zp", + "weight", + "weight_scale", + "weight_zp", + } + + candidate = onnx.load_model(parent_path, load_external_data=True) + _quantize_outer_p_in_place(candidate) + candidate_path = tmp_path / "candidate.onnx" + onnx.save_model( + candidate, + candidate_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="candidate.data", + size_threshold=0, + ) + + preservation = quantization.compare_existing_qdq_state(str(candidate_path), state) + assert preservation["status"] == "passed" + union_audit = quantization.audit_qdq_model( + str(candidate_path), + ["selected", "model.blocks.0.normactconvp.conv/nhwc"], + ) + assert union_audit["errors"] == [] + assert union_audit["selected_nodes_with_qdq_inputs"] == 2 + + +def test_incremental_qdq_snapshot_detects_old_scale_byte_change(tmp_path): + parent = _make_qdq_model() + _append_unquantized_outer_p(parent) + parent_path = tmp_path / "parent.onnx" + onnx.save_model(parent, parent_path) + state = quantization.capture_existing_qdq_state(str(parent_path)) + + changed = onnx.load_model(parent_path) + _replace_initializer(changed, "act_scale", 0.125, np.float32) + changed_path = tmp_path / "changed.onnx" + onnx.save_model(changed, changed_path) + + comparison = quantization.compare_existing_qdq_state(str(changed_path), state) + assert comparison["status"] == "failed" + assert any( + "initializers: changed ['act_scale']" in difference + for difference in comparison["differences"] + ) diff --git a/python/tests/test_quantize_onnx.py b/python/tests/test_quantize_onnx.py new file mode 100644 index 0000000000..0b0c64adfb --- /dev/null +++ b/python/tests/test_quantize_onnx.py @@ -0,0 +1,698 @@ +"""Focused tests for the research ONNX quantization exporter. + +The synthetic graphs below mirror the names emitted by onnxmodelbuilder.cpp. In +particular, an nbt3 block has three attention/FFN pairs, so the target b15 model +has 15 * 3 * (4 attention projections + 3 FFN projections) = 315 eligible +weight projections. Attention score/value MatMuls deliberately have no +initializer and must remain unquantized. +""" + +from __future__ import annotations + +import numpy as np +import pytest + + +onnx = pytest.importorskip("onnx") +quantize_onnx = pytest.importorskip( + "katago.quantization", reason="KataGo quantization helpers have not been added yet" +) + + +def _make_target_projection_graph(*, nhwc: bool): + helper = onnx.helper + tensor_proto = onnx.TensorProto + nodes = [] + initializers = [] + expected = [] + current = "input" + + def add_weight_node(name: str, op_type: str) -> None: + nonlocal current + weight_name = name + (".Wnhwc" if nhwc else ".W") + if op_type == "MatMul": + weight = helper.make_tensor(weight_name, tensor_proto.FLOAT, [1, 1], [1.0]) + else: + weight = helper.make_tensor( + weight_name, tensor_proto.FLOAT, [1, 1, 1, 1], [1.0] + ) + output = name + "/output" + nodes.append( + helper.make_node(op_type, [current, weight_name], [output], name=name) + ) + initializers.append(weight) + current = output + + projection_op = "MatMul" if nhwc else "Conv" + suffix = "/nhwc" if nhwc else "" + for outer_idx in range(15): + for pair_idx in range(3): + attn_idx = pair_idx * 2 + ffn_idx = attn_idx + 1 + attn_base = f"model.blocks.{outer_idx}.blockstack.{attn_idx}" + ffn_base = f"model.blocks.{outer_idx}.blockstack.{ffn_idx}" + for projection in ("q_proj", "k_proj", "v_proj", "out_proj"): + name = f"{attn_base}.{projection}{suffix}" + add_weight_node(name, projection_op) + expected.append(name) + for projection in ("ffn_linear1", "ffn_linear_gate", "ffn_linear2"): + name = f"{ffn_base}.{projection}{suffix}" + add_weight_node(name, projection_op) + expected.append(name) + + # Activation x activation attention MatMuls are precision-sensitive and are + # never weight projections, despite living in the transformer scope. + nodes.append( + helper.make_node( + "MatMul", + [current, current], + ["scores"], + name="model.blocks.0.blockstack.0/scores", + ) + ) + nodes.append( + helper.make_node( + "MatMul", + ["scores", current], + ["sv"], + name="model.blocks.0.blockstack.0/sv", + ) + ) + + # These all have constant weights, but are outside the deliberately narrow + # transformer-projection scope. + add_weight_node("model.conv_spatial", "Conv") + add_weight_node("model.blocks.0.normactconvp.conv" + suffix, projection_op) + add_weight_node("model.policy_head.conv1p", "Conv") + add_weight_node("model.value_head.linear2" + suffix, projection_op) + + # A projection-looking activation MatMul and a non-projection op are both + # negative controls for name-only selection. + nodes.append( + helper.make_node( + "MatMul", + [current, current], + ["fake_projection"], + name="model.blocks.0.blockstack.0.q_proj/activation_only", + ) + ) + fake_weight = helper.make_tensor("fake_add.W", tensor_proto.FLOAT, [1], [1.0]) + initializers.append(fake_weight) + nodes.append( + helper.make_node( + "Add", + [current, "fake_add.W"], + ["final"], + name="model.blocks.0.blockstack.0.q_proj" + suffix, + ) + ) + + graph = helper.make_graph( + nodes, + "target_projection_selection", + [helper.make_tensor_value_info("input", tensor_proto.FLOAT, [None, 1, 1, 1])], + [helper.make_tensor_value_info("final", tensor_proto.FLOAT, None)], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)]) + return model, expected + + +@pytest.mark.parametrize("nhwc", [False, True]) +def test_selects_exactly_315_b15_nbt3_weight_projections(nhwc): + model, expected = _make_target_projection_graph(nhwc=nhwc) + + selected = quantize_onnx.select_quantizable_nodes(model, scope="transformer") + + assert selected.selected_names == sorted(expected) + assert len(selected.selected_names) == 315 + assert selected.selected_by_op_type == {"MatMul" if nhwc else "Conv": 315} + + +@pytest.mark.parametrize("nhwc", [False, True]) +def test_all_weighted_matches_reference_matmul_conv_scope(nhwc): + """Match the executable scope in zml24's reference INT8 script. + + That script passes every weighted MatMul and Conv to ORT and excludes only + activation-only attention MatMuls. Its declared stem/head skip patterns are + not wired into the quantization call. + """ + + model, transformer = _make_target_projection_graph(nhwc=nhwc) + suffix = "/nhwc" if nhwc else "" + expected = transformer + [ + "model.conv_spatial", + "model.blocks.0.normactconvp.conv" + suffix, + "model.policy_head.conv1p", + "model.value_head.linear2" + suffix, + ] + + selected = quantize_onnx.select_quantizable_nodes(model, scope="all-weighted") + + assert selected.selected_names == sorted(expected) + assert len(selected.selected_names) == 319 + assert "model.blocks.0.blockstack.0/scores" not in selected.selected_names + assert "model.blocks.0.blockstack.0/sv" not in selected.selected_names + assert not any("activation_only" in name for name in selected.selected_names) + + +def test_only_node_regex_restricts_selection_to_one_outer_block_ffn(): + model, _ = _make_target_projection_graph(nhwc=True) + + selected = quantize_onnx.select_quantizable_nodes( + model, + scope="transformer", + only_regexes=(r"^model\.blocks\.7\..*ffn_",), + ) + assert len(selected.selected_names) == 9 + assert all( + name.startswith("model.blocks.7.") and "ffn_" in name + for name in selected.selected_names + ) + + +def _make_five_output_contract_without_provenance(): + helper = onnx.helper + tp = onnx.TensorProto + inputs = [ + helper.make_tensor_value_info("InputMask", tp.FLOAT, [None, 1, 19, 19]), + helper.make_tensor_value_info("InputSpatial", tp.FLOAT, [None, 22, 19, 19]), + helper.make_tensor_value_info("InputGlobal", tp.FLOAT, [None, 19, 1, 1]), + ] + outputs = [ + helper.make_tensor_value_info("OutputPolicyPass", tp.FLOAT, [None, 4, 1, 1]), + helper.make_tensor_value_info("OutputPolicy", tp.FLOAT, [None, 4, 19, 19]), + helper.make_tensor_value_info("OutputValue", tp.FLOAT, [None, 3, 1, 1]), + helper.make_tensor_value_info("OutputScoreValue", tp.FLOAT, [None, 6, 1, 1]), + helper.make_tensor_value_info("OutputOwnership", tp.FLOAT, [None, 1, 19, 19]), + ] + return helper.make_model(helper.make_graph([], "derived", inputs, outputs)) + + +def test_incremental_parent_may_lack_provenance_but_not_five_output_contract(): + model = _make_five_output_contract_without_provenance() + with pytest.raises(ValueError, match="producer_name=katago"): + quantize_onnx.validate_katago_io_contract(model) + + specs = quantize_onnx.validate_katago_io_contract( + model, require_producer_metadata=False + ) + assert [spec.name for spec in specs] == [ + "InputMask", + "InputSpatial", + "InputGlobal", + ] + + model.graph.output[4].name = "WrongOwnership" + with pytest.raises(ValueError, match="Expected the five raw outputs"): + quantize_onnx.validate_katago_io_contract( + model, require_producer_metadata=False + ) + + +def _input_specs(height: int, width: int): + return [ + quantize_onnx.InputSpec( + "InputMask", (None, 1, height, width), onnx.TensorProto.FLOAT + ), + quantize_onnx.InputSpec( + "InputSpatial", (None, 22, height, width), onnx.TensorProto.FLOAT + ), + quantize_onnx.InputSpec( + "InputGlobal", (None, 19, 1, 1), onnx.TensorProto.FLOAT + ), + ] + + +def _write_training_npz(path, spatial: np.ndarray, global_input: np.ndarray) -> None: + packed = np.packbits( + spatial.reshape(spatial.shape[0], spatial.shape[1], -1), axis=2 + ) + np.savez_compressed( + path, + binaryInputNCHWPacked=packed, + globalInputNC=global_input, + ) + + +def _concatenate_inputs(dataset): + return { + name: np.concatenate([batch[name] for batch in dataset.batches], axis=0) + for name in ("InputMask", "InputSpatial", "InputGlobal") + } + + +def _make_identifiable_positions(sample_count: int, height: int, width: int): + spatial = np.zeros((sample_count, 22, height, width), dtype=np.uint8) + for row in range(sample_count): + # Different channels and rows exercise both pack-byte boundaries and + # the final padding bits when H*W is not divisible by eight. + flat = spatial[row].reshape(22, -1) + for channel in range(22): + flat[channel, (row * 3 + channel * 5) % (height * width)] = 1 + global_input = np.ones((sample_count, 19), dtype=np.float32) + global_input[:, 18] = np.arange(sample_count, dtype=np.float32) + return spatial, global_input + + +def test_training_npz_unpack_and_full_history_round_trip(tmp_path): + height, width, sample_count = 3, 5, 7 + spatial, global_input = _make_identifiable_positions(sample_count, height, width) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + + dataset = quantize_onnx.load_position_dataset( + [str(path)], + _input_specs(height, width), + sample_count=sample_count, + batch_size=3, + seed=12345, + history_mode="full", + symmetry_mode="none", + ) + actual = _concatenate_inputs(dataset) + + assert [batch["InputSpatial"].shape[0] for batch in dataset.batches] == [3, 3, 1] + assert dataset.sample_count == sample_count + assert dataset.history_mode == "full" + assert all(value.dtype == np.float32 for value in actual.values()) + assert all(value.flags.c_contiguous for value in actual.values()) + + # The sampler intentionally shuffles rows. The last global feature is a + # stable row id, allowing an exact comparison without relying on order. + row_ids = actual["InputGlobal"][:, 18, 0, 0].astype(np.int64) + np.testing.assert_array_equal(actual["InputSpatial"], spatial[row_ids]) + np.testing.assert_array_equal(actual["InputMask"], spatial[row_ids, 0:1]) + np.testing.assert_array_equal( + actual["InputGlobal"][:, :, 0, 0], global_input[row_ids] + ) + + +def test_none_history_matches_katago_history_matrix_semantics(tmp_path): + height, width, sample_count = 3, 5, 8 + spatial, global_input = _make_identifiable_positions(sample_count, height, width) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + + dataset = quantize_onnx.load_position_dataset( + [str(path)], + _input_specs(height, width), + sample_count=sample_count, + batch_size=sample_count, + seed=9, + history_mode="none", + symmetry_mode="none", + ) + actual = _concatenate_inputs(dataset) + row_ids = actual["InputGlobal"][:, 18, 0, 0].astype(np.int64) + transformed = actual["InputSpatial"] + + np.testing.assert_array_equal(transformed[:, 9:14], 0.0) + np.testing.assert_array_equal(transformed[:, 14], spatial[row_ids, 14]) + np.testing.assert_array_equal(transformed[:, 15], spatial[row_ids, 14]) + np.testing.assert_array_equal(transformed[:, 16], spatial[row_ids, 14]) + np.testing.assert_array_equal(actual["InputGlobal"][:, :5], 0.0) + np.testing.assert_array_equal( + actual["InputGlobal"][:, 5:, 0, 0], global_input[row_ids, 5:] + ) + + +def test_training_history_is_seeded_and_prefix_shaped(tmp_path): + height, width, sample_count = 3, 5, 256 + spatial, global_input = _make_identifiable_positions(sample_count, height, width) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + + kwargs = dict( + paths=[str(path)], + input_specs=_input_specs(height, width), + sample_count=sample_count, + batch_size=31, + seed=20260808, + history_mode="training", + symmetry_mode="none", + ) + first = quantize_onnx.load_position_dataset(**kwargs) + second = quantize_onnx.load_position_dataset(**kwargs) + first_inputs = _concatenate_inputs(first) + second_inputs = _concatenate_inputs(second) + + assert first.position_sha256 == second.position_sha256 + assert first.selection_sha256 == second.selection_sha256 + for name in first_inputs: + np.testing.assert_array_equal(first_inputs[name], second_inputs[name]) + + history_flags = first_inputs["InputGlobal"][:, :5, 0, 0] + assert np.any(history_flags == 0.0) + assert np.any(np.all(history_flags == 1.0, axis=1)) + # Included history is always a prefix: 1,1,...,1,0,...,0. + assert np.all(np.diff(history_flags, axis=1) <= 0.0) + + +def test_calibration_reader_supports_modelopt_reader_protocol(tmp_path): + height, width, sample_count = 3, 5, 7 + spatial, global_input = _make_identifiable_positions(sample_count, height, width) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + dataset = quantize_onnx.load_position_dataset( + [str(path)], + _input_specs(height, width), + sample_count=sample_count, + batch_size=3, + seed=12, + history_mode="full", + symmetry_mode="none", + ) + reader = quantize_onnx.ArrayCalibrationDataReader(dataset) + + assert len(reader) == 3 + assert reader.get_first() is dataset.batches[0] + assert reader.get_next() is dataset.batches[0] + reader.rewind() + assert reader.get_next() is dataset.batches[0] + reader.set_range(1, 3) + assert len(reader) == 2 + assert reader.get_first() is dataset.batches[1] + ranged = list(reader) + assert len(ranged) == 2 + assert all( + actual is expected for actual, expected in zip(ranged, dataset.batches[1:3]) + ) + + +def _expected_spatial_symmetry(value: np.ndarray, symmetry: int) -> np.ndarray: + if symmetry == 0: + return value + if symmetry == 1: + return np.rot90(value, k=1, axes=(-2, -1)) + if symmetry == 2: + return np.rot90(value, k=2, axes=(-2, -1)) + if symmetry == 3: + return np.rot90(value, k=3, axes=(-2, -1)) + if symmetry == 4: + return np.swapaxes(value, -2, -1) + if symmetry == 5: + return np.flip(value, axis=-1) + if symmetry == 6: + return np.flip(np.swapaxes(value, -2, -1), axis=(-2, -1)) + if symmetry == 7: + return np.flip(value, axis=-2) + raise AssertionError(symmetry) + + +def test_all_symmetries_expand_each_source_position_in_numbered_order(tmp_path): + height = width = 5 + spatial, global_input = _make_identifiable_positions(4, height, width) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + + kwargs = dict( + paths=[str(path)], + input_specs=_input_specs(height, width), + sample_count=3, + batch_size=10, + seed=42, + history_mode="full", + symmetry_mode="all", + ) + first = quantize_onnx.load_position_dataset(**kwargs) + second = quantize_onnx.load_position_dataset(**kwargs) + actual = _concatenate_inputs(first) + + assert first.base_sample_count == 3 + assert first.sample_count == 24 + assert [batch["InputSpatial"].shape[0] for batch in first.batches] == [10, 10, 4] + assert first.symmetry_mode == "all" + assert first.symmetry_counts == {str(symmetry): 3 for symmetry in range(8)} + assert first.symmetry_sha256 == second.symmetry_sha256 + assert first.selection_sha256 == second.selection_sha256 + assert first.position_sha256 == second.position_sha256 + + row_ids = actual["InputGlobal"][:, 18, 0, 0].astype(np.int64).reshape(3, 8) + for base_index in range(3): + assert np.all(row_ids[base_index] == row_ids[base_index, 0]) + source_row = row_ids[base_index, 0] + for symmetry in range(8): + output_index = base_index * 8 + symmetry + np.testing.assert_array_equal( + actual["InputSpatial"][output_index], + _expected_spatial_symmetry(spatial[source_row], symmetry), + ) + np.testing.assert_array_equal( + actual["InputMask"][output_index], + actual["InputSpatial"][output_index, 0:1], + ) + np.testing.assert_array_equal( + actual["InputGlobal"][output_index, :, 0, 0], global_input[source_row] + ) + + +def test_random_symmetries_are_seeded_and_transform_mask_with_spatial(tmp_path): + height = width = 5 + sample_count = 64 + spatial, global_input = _make_identifiable_positions(sample_count, height, width) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + kwargs = dict( + paths=[str(path)], + input_specs=_input_specs(height, width), + sample_count=sample_count, + batch_size=17, + seed=20260808, + history_mode="full", + symmetry_mode="random", + ) + + first = quantize_onnx.load_position_dataset(**kwargs) + second = quantize_onnx.load_position_dataset(**kwargs) + first_inputs = _concatenate_inputs(first) + second_inputs = _concatenate_inputs(second) + + assert first.sample_count == first.base_sample_count == sample_count + assert sum(first.symmetry_counts.values()) == sample_count + assert sum(count > 0 for count in first.symmetry_counts.values()) > 1 + assert first.symmetry_sha256 == second.symmetry_sha256 + assert first.position_sha256 == second.position_sha256 + for name in first_inputs: + np.testing.assert_array_equal(first_inputs[name], second_inputs[name]) + + row_ids = first_inputs["InputGlobal"][:, 18, 0, 0].astype(np.int64) + for output_index, source_row in enumerate(row_ids): + possible = [ + _expected_spatial_symmetry(spatial[source_row], symmetry) + for symmetry in range(8) + ] + assert any( + np.array_equal(first_inputs["InputSpatial"][output_index], candidate) + for candidate in possible + ) + np.testing.assert_array_equal( + first_inputs["InputMask"][output_index], + first_inputs["InputSpatial"][output_index, 0:1], + ) + np.testing.assert_array_equal( + first_inputs["InputGlobal"][output_index, :, 0, 0], global_input[source_row] + ) + + +def test_symmetry_augmentation_rejects_rectangular_graph(tmp_path): + spatial, global_input = _make_identifiable_positions(2, 3, 5) + path = tmp_path / "positions.npz" + _write_training_npz(path, spatial, global_input) + + with pytest.raises(ValueError, match="square ONNX board"): + quantize_onnx.load_position_dataset( + [str(path)], + _input_specs(3, 5), + sample_count=2, + batch_size=2, + seed=1, + history_mode="full", + symmetry_mode="random", + ) + + +def test_expanded_inputs_apply_requested_history_transform(tmp_path): + height = width = 3 + spatial, global_input = _make_identifiable_positions(4, height, width) + path = tmp_path / "expanded.npz" + np.savez_compressed( + path, + InputMask=spatial[:, 0:1].astype(np.float32), + InputSpatial=spatial.astype(np.float32), + InputGlobal=global_input[:, :, None, None], + ) + + dataset = quantize_onnx.load_position_dataset( + [str(path)], + _input_specs(height, width), + sample_count=4, + batch_size=4, + seed=33, + history_mode="none", + symmetry_mode="none", + ) + actual = _concatenate_inputs(dataset) + + np.testing.assert_array_equal(actual["InputSpatial"][:, 9:14], 0.0) + np.testing.assert_array_equal(actual["InputGlobal"][:, :5], 0.0) + + +def test_source_file_limit_is_seeded_and_bounds_npz_decompression(tmp_path): + height = width = 3 + for file_index in range(5): + spatial, global_input = _make_identifiable_positions(10, height, width) + global_input[:, 18] = file_index * 100 + np.arange(10, dtype=np.float32) + _write_training_npz(tmp_path / f"shard-{file_index}.npz", spatial, global_input) + + kwargs = dict( + paths=[str(tmp_path)], + input_specs=_input_specs(height, width), + sample_count=8, + batch_size=4, + seed=99, + history_mode="full", + symmetry_mode="none", + max_source_files=2, + ) + first = quantize_onnx.load_position_dataset(**kwargs) + second = quantize_onnx.load_position_dataset(**kwargs) + + assert first.available_source_file_count == 5 + assert first.selected_source_file_count == 2 + assert first.max_source_files == 2 + assert len(first.source_files) == 2 + assert first.selection_sha256 == second.selection_sha256 + assert first.position_sha256 == second.position_sha256 + + +def _raw_outputs(sample_count=2, height=2, width=2, policy_channels=2): + return { + "OutputPolicyPass": np.zeros( + (sample_count, policy_channels, 1, 1), dtype=np.float32 + ), + "OutputPolicy": np.zeros( + (sample_count, policy_channels, height, width), dtype=np.float32 + ), + "OutputValue": np.zeros((sample_count, 3, 1, 1), dtype=np.float32), + "OutputScoreValue": np.zeros((sample_count, 6, 1, 1), dtype=np.float32), + "OutputOwnership": np.zeros((sample_count, 1, height, width), dtype=np.float32), + } + + +def test_validation_metrics_mask_offboard_policy_and_ownership(): + reference = _raw_outputs() + candidate = {name: value.copy() for name, value in reference.items()} + mask = np.ones((2, 1, 2, 2), dtype=np.float32) + mask[0, 0, 1, 1] = 0.0 + candidate["OutputPolicy"][0, :, 1, 1] = 1000.0 + candidate["OutputOwnership"][0, 0, 1, 1] = 1000.0 + + metrics = quantize_onnx.compute_validation_metrics( + reference, candidate, {"InputMask": mask} + ) + + assert metrics["policy"]["kl"]["max"] == pytest.approx(0.0, abs=1.0e-12) + assert metrics["outputs"]["OutputOwnership"]["max_abs"] == 0.0 + assert metrics["outputs"]["OutputOwnership"]["unmasked"]["max_abs"] == 1000.0 + assert metrics["outputs"]["OutputOwnership"]["offboard"]["max_abs"] == 1000.0 + # The generic raw-output audit intentionally remains unmasked, making it + # possible to spot unexpectedly large off-board activations too. + assert metrics["outputs"]["OutputPolicy"]["max_abs"] == 1000.0 + + +def test_four_channel_q_outputs_use_raw_error_not_policy_kl(): + reference = _raw_outputs(policy_channels=4) + candidate = {name: value.copy() for name, value in reference.items()} + mask = np.ones((2, 1, 2, 2), dtype=np.float32) + candidate["OutputPolicy"][:, 2] += 3.0 + candidate["OutputPolicyPass"][:, 2] += 3.0 + candidate["OutputPolicy"][:, 3] -= 2.0 + candidate["OutputPolicyPass"][:, 3] -= 2.0 + + metrics = quantize_onnx.compute_validation_metrics( + reference, candidate, {"InputMask": mask} + ) + + assert metrics["policy"]["kl"]["max"] == pytest.approx(0.0, abs=1.0e-12) + assert set(metrics["policy"]["per_channel"]) == { + "policy", + "shortterm_optimistic", + } + assert metrics["policy"]["quantitative"]["q_value"]["rmse"] == pytest.approx(3.0) + assert metrics["policy"]["quantitative"]["q_score"]["max_abs"] == pytest.approx(2.0) + gated = quantize_onnx.evaluate_accuracy_gates( + metrics, + { + "max_q_value_rmse": 2.9, + "max_q_score_max_abs": 2.1, + }, + ) + assert gated["status"] == "failed" + assert ( + next( + check + for check in gated["checks"] + if check["threshold"] == "max_q_value_rmse" + )["passed"] + is False + ) + + +def test_validation_metrics_and_accuracy_gates_detect_real_drift(): + reference = _raw_outputs() + candidate = {name: value.copy() for name, value in reference.items()} + mask = np.ones((2, 1, 2, 2), dtype=np.float32) + candidate["OutputPolicyPass"][0, 0, 0, 0] = 1.0 + candidate["OutputValue"][0, 1, 0, 0] = 1.0 + candidate["OutputScoreValue"][0, 2, 0, 0] = 0.25 + candidate["OutputOwnership"][0, 0, 0, 0] = 0.1 + + metrics = quantize_onnx.compute_validation_metrics( + reference, candidate, {"InputMask": mask} + ) + assert metrics["policy"]["kl"]["mean"] > 0.0 + assert metrics["value"]["kl"]["mean"] > 0.0 + assert metrics["outputs"]["OutputScoreValue"]["max_abs"] == 0.25 + assert metrics["outputs"]["OutputOwnership"]["rmse"] > 0.0 + + failed = quantize_onnx.evaluate_accuracy_gates( + metrics, + { + "max_policy_kl_mean": 0.0, + "max_value_kl_mean": 0.0, + "max_scorevalue_max_abs": 0.1, + "max_ownership_rmse": 0.001, + }, + ) + assert failed["status"] == "failed" + assert any(not check["passed"] for check in failed["checks"]) + + permissive = quantize_onnx.evaluate_accuracy_gates( + metrics, + { + "max_policy_kl_mean": 1.0, + "max_value_kl_mean": 1.0, + "max_scorevalue_max_abs": 1.0, + "max_ownership_rmse": 1.0, + }, + ) + assert permissive["status"] == "passed" + + +def test_nonfinite_candidate_always_fails_without_optional_thresholds(): + reference = _raw_outputs() + candidate = {name: value.copy() for name, value in reference.items()} + candidate["OutputScoreValue"][0, 0, 0, 0] = np.nan + metrics = quantize_onnx.compute_validation_metrics( + reference, + candidate, + {"InputMask": np.ones((2, 1, 2, 2), dtype=np.float32)}, + ) + + result = quantize_onnx.evaluate_accuracy_gates(metrics, {}) + + assert result["status"] == "failed" + assert result["numeric_thresholds_configured"] is False + assert result["checks"][-1]["threshold"] == "candidate_nonfinite" + assert result["checks"][-1]["passed"] is False diff --git a/python/tests/test_quantize_onnx_artifacts.py b/python/tests/test_quantize_onnx_artifacts.py new file mode 100644 index 0000000000..c30be7a781 --- /dev/null +++ b/python/tests/test_quantize_onnx_artifacts.py @@ -0,0 +1,539 @@ +"""Offline tests for safe ONNX artifact staging and replacement.""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + + +onnx = pytest.importorskip("onnx") +exporter = pytest.importorskip("quantize_onnx") + + +def _write_model( + path: Path, + value: float, + *, + external_location: str | None = None, +) -> None: + helper = onnx.helper + tensor_proto = onnx.TensorProto + weight = onnx.numpy_helper.from_array( + np.full((1, 16), value, dtype=np.float32), name="weight" + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Add", ["input", "weight"], ["output"], name="add")], + "artifact-test", + [helper.make_tensor_value_info("input", tensor_proto.FLOAT, [1, 16])], + [helper.make_tensor_value_info("output", tensor_proto.FLOAT, [1, 16])], + [weight], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + if external_location is None: + onnx.save_model(model, str(path)) + else: + onnx.save_model( + model, + str(path), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=external_location, + size_threshold=0, + ) + + +def _weight_value(path: Path) -> float: + model = onnx.load(str(path), load_external_data=True) + return float(onnx.numpy_helper.to_array(model.graph.initializer[0]).flat[0]) + + +def _write_fp8_qdq_matmul(path: Path) -> None: + helper = onnx.helper + tp = onnx.TensorProto + activation_scale = onnx.numpy_helper.from_array( + np.asarray(0.027776, dtype=np.float32), name="activation_scale" + ) + weight_scale = onnx.numpy_helper.from_array( + np.asarray([0.01, 0.02, 0.03], dtype=np.float32), name="weight_scale" + ) + weight = onnx.numpy_helper.from_array( + np.arange(6, dtype=np.float32).reshape(2, 3), name="weight" + ) + nodes = [ + helper.make_node( + "QuantizeLinear", + ["input", "activation_scale"], + ["input_q"], + name="input_q", + output_dtype=tp.FLOAT8E4M3FN, + ), + helper.make_node( + "DequantizeLinear", + ["input_q", "activation_scale"], + ["input_dq"], + name="input_dq", + ), + helper.make_node( + "QuantizeLinear", + ["weight", "weight_scale"], + ["weight_q"], + name="weight_q", + axis=1, + output_dtype=tp.FLOAT8E4M3FN, + ), + helper.make_node( + "DequantizeLinear", + ["weight_q", "weight_scale"], + ["weight_dq"], + name="weight_dq", + axis=1, + ), + helper.make_node( + "MatMul", + ["input_dq", "weight_dq"], + ["output"], + name="model.blocks.0.blockstack.1.ffn_linear1/nhwc", + ), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "fp8-scale-test", + [helper.make_tensor_value_info("input", tp.FLOAT, [None, 2])], + [helper.make_tensor_value_info("output", tp.FLOAT, [None, 3])], + [activation_scale, weight_scale, weight], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.save_model(model, str(path)) + + +def _write_incremental_int8_parent(path: Path) -> None: + helper = onnx.helper + tp = onnx.TensorProto + initializers = [ + onnx.numpy_helper.from_array( + np.arange(6, dtype=np.float32).reshape(2, 3), name="old_weight" + ), + onnx.numpy_helper.from_array(np.asarray(0.1, dtype=np.float32), name="old_as"), + onnx.numpy_helper.from_array(np.asarray(0, dtype=np.int8), name="old_az"), + onnx.numpy_helper.from_array( + np.asarray([0.02, 0.03, 0.04], dtype=np.float32), name="old_ws" + ), + onnx.numpy_helper.from_array( + np.asarray([0, 0, 0], dtype=np.int8), name="old_wz" + ), + onnx.numpy_helper.from_array( + np.arange(3 * 512, dtype=np.float32).reshape(3, 512), name="new_weight" + ), + ] + nodes = [ + helper.make_node( + "QuantizeLinear", ["input", "old_as", "old_az"], ["old_aq"], name="old_aq" + ), + helper.make_node( + "DequantizeLinear", + ["old_aq", "old_as", "old_az"], + ["old_adq"], + name="old_adq", + ), + helper.make_node( + "QuantizeLinear", + ["old_weight", "old_ws", "old_wz"], + ["old_wq"], + name="old_wq", + axis=1, + ), + helper.make_node( + "DequantizeLinear", + ["old_wq", "old_ws", "old_wz"], + ["old_wdq"], + name="old_wdq", + axis=1, + ), + helper.make_node("MatMul", ["old_adq", "old_wdq"], ["old_output"], name="old"), + helper.make_node( + "MatMul", + ["old_output", "new_weight"], + ["output"], + name="model.blocks.0.normactconvp.conv/nhwc", + ), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "incremental-parent", + [helper.make_tensor_value_info("input", tp.FLOAT, [None, 2])], + [helper.make_tensor_value_info("output", tp.FLOAT, [None, 512])], + initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + model.ir_version = 10 + onnx.save_model( + model, + path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=path.name + ".parent.data", + size_threshold=0, + ) + + +def _add_incremental_qdq(model) -> None: + outer = next( + node + for node in model.graph.node + if node.name == "model.blocks.0.normactconvp.conv/nhwc" + ) + model.graph.initializer.extend( + [ + onnx.numpy_helper.from_array( + np.asarray(0.2, dtype=np.float32), name="new_as" + ), + onnx.numpy_helper.from_array(np.asarray(0, dtype=np.int8), name="new_az"), + onnx.numpy_helper.from_array( + np.asarray([0.05, 0.06], dtype=np.float32), name="new_ws" + ), + onnx.numpy_helper.from_array( + np.asarray([0, 0], dtype=np.int8), name="new_wz" + ), + ] + ) + new_nodes = [ + onnx.helper.make_node( + "QuantizeLinear", + ["old_output", "new_as", "new_az"], + ["new_aq"], + name="new_aq", + ), + onnx.helper.make_node( + "DequantizeLinear", + ["new_aq", "new_as", "new_az"], + ["new_adq"], + name="new_adq", + ), + onnx.helper.make_node( + "QuantizeLinear", + ["new_weight", "new_ws", "new_wz"], + ["new_wq"], + name="new_wq", + axis=1, + ), + onnx.helper.make_node( + "DequantizeLinear", + ["new_wq", "new_ws", "new_wz"], + ["new_wdq"], + name="new_wdq", + axis=1, + ), + ] + outer.input[0] = "new_adq" + outer.input[1] = "new_wdq" + index = list(model.graph.node).index(outer) + for offset, node in enumerate(new_nodes): + model.graph.node.insert(index + offset, node) + + +def test_direct_amax_fp8_scale_rewrite_uses_separate_activation_headroom( + tmp_path: Path, +) -> None: + path = tmp_path / "fp8.onnx" + _write_fp8_qdq_matmul(path) + original = onnx.load(str(path)) + original_values = { + item.name: onnx.numpy_helper.to_array(item).copy() + for item in original.graph.initializer + } + + details = exporter._rewrite_fp8_direct_amax_scales( + str(path), + ["model.blocks.0.blockstack.1.ffn_linear1/nhwc"], + activation_qmax=224.0, + ) + + rewritten = onnx.load(str(path)) + values = { + item.name: onnx.numpy_helper.to_array(item) + for item in rewritten.graph.initializer + } + legacy_qmax = 127.0**2 / 448.0 + np.testing.assert_allclose( + values["activation_scale"], + original_values["activation_scale"] * legacy_qmax / 224.0, + ) + np.testing.assert_allclose( + values["weight_scale"], + original_values["weight_scale"] * legacy_qmax / 448.0, + ) + np.testing.assert_array_equal(values["weight"], original_values["weight"]) + assert details["scale_initializer_count"] == {"activation": 1, "weight": 1} + assert details["activation_qmax"] == 224.0 + + +def test_copy_artifact_stages_external_data_without_touching_source( + tmp_path: Path, +) -> None: + source_dir = tmp_path / "source" + source_dir.mkdir() + source = source_dir / "model.onnx" + _write_model(source, 3.0, external_location="weights.bin") + expected = exporter.artifact_manifest(str(source)) + + staging_dir = tmp_path / "staging" + staged = Path(exporter._copy_onnx_artifact(str(source), str(staging_dir))) + + assert staged == staging_dir / source.name + assert (staging_dir / "weights.bin").is_file() + assert _weight_value(staged) == pytest.approx(3.0) + + # Simulate ModelOpt's in-place shape-inference rewrite of its input path. + staged_model = onnx.load(str(staged), load_external_data=False) + staged_model.producer_name = "mutated-staging-copy" + onnx.save_model(staged_model, str(staged)) + + assert exporter._artifact_integrity(expected)["status"] == "passed" + + +def test_promote_replaces_only_exact_previous_artifact(tmp_path: Path) -> None: + output_dir = tmp_path / "output" + output_dir.mkdir() + final_model = output_dir / "network.int8.qdq.onnx" + _write_model(final_model, 1.0, external_location="old.weights.bin") + old_sidecar = output_dir / "old.weights.bin" + unrelated = output_dir / "do-not-delete.weights.bin" + unrelated.write_bytes(b"unrelated") + + staging_dir = output_dir / ".katago-quant-test" + staging_dir.mkdir() + staged_model = staging_dir / "candidate.0123456789abcdef.onnx" + new_location = "candidate.0123456789abcdef.onnx_data" + _write_model(staged_model, 9.0, external_location=new_location) + + details = exporter._promote_staged_artifact( + str(staged_model), str(final_model), overwrite=True + ) + + new_sidecar = output_dir / new_location + assert _weight_value(final_model) == pytest.approx(9.0) + assert new_sidecar.is_file() + assert not old_sidecar.exists() + assert unrelated.read_bytes() == b"unrelated" + assert str(old_sidecar.resolve()) in details["removed_previous_external_data"] + assert details["promoted_external_data"] == [str(new_sidecar.resolve())] + + +def test_promote_without_overwrite_has_no_side_effects(tmp_path: Path) -> None: + final_model = tmp_path / "existing.onnx" + _write_model(final_model, 2.0, external_location="existing.weights.bin") + expected = exporter.artifact_manifest(str(final_model)) + + staging_dir = tmp_path / "staging" + staging_dir.mkdir() + staged_model = staging_dir / "candidate.onnx" + _write_model(staged_model, 7.0, external_location="candidate.weights.bin") + + with pytest.raises(FileExistsError, match="--overwrite"): + exporter._promote_staged_artifact( + str(staged_model), str(final_model), overwrite=False + ) + + assert exporter._artifact_integrity(expected)["status"] == "passed" + assert staged_model.is_file() + assert (staging_dir / "candidate.weights.bin").is_file() + + +def test_quantize_one_never_passes_source_artifact_to_modelopt(tmp_path: Path) -> None: + source = tmp_path / "source.onnx" + _write_model(source, 4.0) + expected = exporter.artifact_manifest(str(source)) + output_dir = tmp_path / "output" + output_dir.mkdir() + output = output_dir / "safe.int8.qdq.onnx" + seen: dict[str, Path] = {} + + def mutating_quantizer(**kwargs) -> None: + staged_source = Path(kwargs["onnx_path"]).resolve() + seen["source"] = staged_source + assert staged_source != source.resolve() + staged_model = onnx.load(str(staged_source), load_external_data=False) + staged_model.producer_name = "modelopt-mutated-this-file" + onnx.save_model(staged_model, str(staged_source)) + shutil.copy2(staged_source, kwargs["output_path"]) + + args = SimpleNamespace( + overwrite=True, + output_dir=str(output_dir), + output_prefix="safe", + calibration_method="entropy", + calibration_eps="cpu", + keep_intermediate_files=False, + high_precision="fp32", + calibrate_per_node=False, + ) + dataset = SimpleNamespace(batches=[{}]) + + exporter._quantize_one( + "int8", + str(source), + str(output), + dataset, + [], + [], + args, + mutating_quantizer, + ) + + assert seen["source"].parent != source.parent + assert output.is_file() + assert exporter._artifact_integrity(expected)["status"] == "passed" + assert not list(output_dir.glob(".katago-quant-*")) + + +def test_incremental_quantize_one_passes_only_new_allowlist_and_preserves_parent( + tmp_path: Path, +) -> None: + source = tmp_path / "parent.onnx" + _write_incremental_int8_parent(source) + snapshot = exporter.capture_existing_qdq_state(str(source)) + output_dir = tmp_path / "output" + output_dir.mkdir() + output = output_dir / "incremental.int8.qdq.onnx" + selected = ["model.blocks.0.normactconvp.conv/nhwc"] + + def incremental_quantizer(**kwargs) -> None: + raise AssertionError("preserve mode must bypass ModelOpt's Q/DQ short circuit") + + args = SimpleNamespace( + overwrite=True, + output_dir=str(output_dir), + output_prefix="incremental", + calibration_method="entropy", + calibration_eps="cpu", + keep_intermediate_files=False, + high_precision="fp32", + calibrate_per_node=False, + ) + details = exporter._quantize_one( + "int8", + str(source), + str(output), + SimpleNamespace( + batches=[ + {"input": np.asarray([[1.0, -2.0], [0.5, 3.0]], dtype=np.float32)}, + {"input": np.asarray([[-1.5, 0.25], [2.0, -0.75]], dtype=np.float32)}, + ] + ), + selected, + ["MatMul"], + args, + incremental_quantizer, + snapshot, + ) + + incremental = details["incremental_quantizer"] + assert incremental["backend"] == "onnxruntime-node-filtered-qdq" + assert incremental["newly_quantized_nodes"] == selected + assert incremental["calibration_tensor_names"] == ["old_output", "output"] + assert incremental["materialized_parent_raw_bytes"] > 0 + assert incremental["self_contained_output"] is True + assert incremental["configuration"] == { + "calibration_method": "entropy", + "activation_type": "QInt8", + "weight_type": "QInt8", + "activation_symmetric": True, + "weight_symmetric": True, + "per_channel_weights": True, + "reduce_range": False, + "quantize_bias": False, + "output_quantization": False, + } + assert details["incremental_union_selected_count"] == 2 + assert details["existing_qdq_preservation"]["status"] == "passed" + assert details["incremental_staged_qdq_audit"]["errors"] == [] + assert output.is_file() + output_members = exporter._external_artifact_members( + str(output), require_exists=True + ) + assert output_members + parent_members = exporter._external_artifact_members( + str(source), require_exists=True + ) + assert parent_members + assert not any( + output_member.samefile(parent_member) + for _, output_member in output_members + for _, parent_member in parent_members + ) + + # Prove the promoted result is not accidentally dependent on the parent payload. + for _, parent_member in parent_members: + parent_member.unlink() + loaded = onnx.load_model(output, load_external_data=True) + onnx.checker.check_model(loaded, full_check=False) + + +def test_incremental_quantize_one_refuses_noop_clone_and_cleans_staging( + tmp_path: Path, + monkeypatch, +) -> None: + source = tmp_path / "parent.onnx" + _write_incremental_int8_parent(source) + snapshot = exporter.capture_existing_qdq_state(str(source)) + output_dir = tmp_path / "output" + output_dir.mkdir() + output = output_dir / "bad.int8.qdq.onnx" + + def no_change_runner(source_path, output_path, *args, **kwargs): + shutil.copy2(source_path, output_path) + raise RuntimeError( + "Incremental ORT quantization produced no usable change; refusing to promote a parent clone" + ) + + monkeypatch.setattr(exporter, "_run_incremental_ort_quantization", no_change_runner) + + args = SimpleNamespace( + overwrite=True, + output_dir=str(output_dir), + output_prefix="bad", + calibration_method="entropy", + calibration_eps="cpu", + keep_intermediate_files=False, + high_precision="fp32", + calibrate_per_node=False, + ) + with pytest.raises(RuntimeError, match="produced no usable change"): + exporter._quantize_one( + "int8", + str(source), + str(output), + SimpleNamespace(batches=[{}]), + ["model.blocks.0.normactconvp.conv/nhwc"], + ["MatMul"], + args, + lambda **kwargs: None, + snapshot, + ) + + assert not output.exists() + assert not list(output_dir.glob(".katago-quant-*")) + + +@pytest.mark.parametrize( + ("mode_report", "expected"), + [ + ({}, False), + ({"accuracy_gate": {"status": "passed"}}, False), + ({"accuracy_gate": {"status": "failed"}}, True), + ({"trtexec": {"passed": True}}, False), + ({"trtexec": {"passed": False}}, True), + ], +) +def test_mode_report_failed(mode_report, expected: bool) -> None: + assert exporter._mode_report_failed(mode_report) is expected