Skip to content

Add experimental INT8/FP8 ONNX post-training quantization tools - #1231

Closed
zsqdx wants to merge 1 commit into
lightvector:masterfrom
zsqdx:agent/quantized-onnx-export
Closed

Add experimental INT8/FP8 ONNX post-training quantization tools#1231
zsqdx wants to merge 1 commit into
lightvector:masterfrom
zsqdx:agent/quantized-onnx-export

Conversation

@zsqdx

@zsqdx zsqdx commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

Summary

This draft adds a standalone, research-oriented post-training quantization path for ONNX graphs emitted by KataGo's TensorRT backend.

  • Calibrates explicit INT8 or FP8 Q/DQ graphs from official shuffled KataGo NPZ data.
  • Uses deterministic shard/position selection, history truncation, and board-symmetry augmentation.
  • Defaults to the weighted transformer projections while preserving attention QK/AV matmuls, Softmax, norms, the stem, and output heads at high precision.
  • Also exposes an aggressive all-weighted scope matching the executable MatMul/Conv coverage of zml24's reference script while still excluding activation-only attention matmuls.
  • Performs a semantic Q/DQ audit covering chain integrity, dtypes, scales, zero points, axes, channel counts, missing targets, orphans, and unexpected quantization.
  • Stages complete ONNX artifacts, verifies source/output hashes, and promotes sidecars before the primary ONNX file to avoid corrupting or mutating the source artifact.
  • Records calibration provenance, selected nodes, execution providers, validation metrics, and artifact hashes in JSON reports.

Motivation

KataGo currently consumes .bin.gz weights and constructs its TensorRT ONNX graph internally. This change provides an isolated way to study explicit INT8/FP8 PTQ against that exact five-output inference graph without changing KataGo's runtime, model format, or default precision behavior.

The default scope is intentionally accuracy-first. The broader scope is available for controlled coverage experiments, but promotion still requires held-out accuracy checks, a TensorRT build, production-batch throughput tests, and self-play.

Scope and safety

This PR contains only the core quantizer, shared helpers, pinned dependency, documentation, and focused tests. It does not include model weights, Q/DQ artifacts, TensorRT engines, benchmark logs, automatic scope-search experiments, or changes to KataGo runtime defaults.

Validation gates are descriptive unless the caller supplies explicit thresholds. The tool fails closed on source/validation overlap, requested-node count mismatches, provider fallback, malformed Q/DQ semantics, source mutation, and artifact publication errors.

Validation

  • 59 passed across the focused quantization, Q/DQ audit, and artifact-safety tests.
  • Ruff lint and formatting checks passed for all files in this draft.
  • Staged diff checks passed.
  • A separate b15 research benchmark used an RTX PRO 6000 with a 500 W board power-limit setting to approximate an RTX 5090-class community scenario. Absolute throughput remains hardware- and power-limit-specific, and this draft makes no portable speed or Elo claim.

Feedback requested

This is intentionally a draft. Feedback would be especially useful on whether this tooling belongs upstream, the desired CLI/API boundary, and whether the initial projection-only scope is narrow enough for a first experimental integration.

@zsqdx zsqdx closed this Aug 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant