Quark Onnx Model Intake

quark-onnx-model-intake

Purpose

Validate the target .onnx model and extract the structural facts that an ONNX quant plan needs to make correct quantization decisions. This step exists because different ONNX models have very different quantization requirements — a CNN destined for AMD NPU CNN needs NHWC layout and PoF2 scales, an LLM exported to ONNX needs MatMulNBits or SmoothQuant + GPTQ, a >2 GB model needs external-data handling, and a model that already contains QDQ nodes (or Quark's custom-op domains) must not be silently re-quantized.

Inputs

• Model path — a single .onnx file, or a directory containing model.onnx plus its model.onnx_data / external-weight blobs
• env_context.json for Python and accelerator constraints (CUDA major / ROCm major / NPU presence)
• workspace_context.json for the validated model path
• onnx_install_result.json (optional) — the installed onnxruntime* variant determines which execution providers are realistic deployment targets

Outputs: model_analysis.json

Captures the model's graph metadata, op-type coverage, quantizable-op count, deployment-target compatibility, and risks. Schema is shared with the Torch intake; ONNX-specific fields live under model.onnx_specific and quantization_targets.onnx_specific.

Schema: model_analysis.schema.json

{
  "analysis_status": "complete",
  "model": {
    "model_path": "/data/models/resnet50/model.onnx",
    "model_type": "onnx",
    "loading_class": "onnx.load",
    "estimated_size_gb": 0.10,
    "onnx_specific": {
      "ir_version": 8,
      "producer_name": "pytorch",
      "producer_version": "2.5.1",
      "opsets": [{"domain": "ai.onnx", "version": 17}],
      "uses_external_data": false,
      "external_data_files": [],
      "input_layout_guess": "NCHW",
      "inputs": [
        {"name": "input", "dtype": "float32", "shape": [1, 3, 224, 224]}
      ],
      "outputs": [
        {"name": "output", "dtype": "float32", "shape": [1, 1000]}
      ],
      "has_dynamic_dims": false,
      "has_qdq_already": false,
      "has_quark_custom_ops": false,
      "architecture_guess": "cnn"
    }
  },
  "quantization_targets": {
    "linear_layer_count": 53,
    "onnx_specific": {
      "op_counts": {"Conv": 53, "BatchNormalization": 53, "Relu": 49, "MaxPool": 1, "GlobalAveragePool": 1, "Gemm": 1},
      "quantizable_op_counts": {"Conv": 53, "Gemm": 1, "MatMul": 0, "Add": 16, "Mul": 0},
      "preset_candidates": ["XINT8", "A8W8", "INT8_CNN_DEFAULT"],
      "exclude_defaults": []
    }
  },
  "risks": [
    {
      "severity": "low",
      "message": "Inputs are NCHW. AMD NPU CNN target expects NHWC.",
      "recovery_hint": "Run quark.onnx.tools.convert_nchw_to_nhwc before quantization."
    }
  ]
}

Supported Domains and Preset Families

Quark ONNX ships a large DefaultConfigMapping of named presets in quark/onnx/quantization/config/custom_config.py. The intake should narrow the candidate set based on what the model actually looks like:

preset_candidates in model_analysis.json should list the 1–3 presets that are realistic given architecture + deployment target; never list all of them.

What to Extract

From the ONNX model proto (without loading weights)

• IR version: model.ir_version
• Producer: model.producer_name, model.producer_version
• Opsets: [(opset.domain or 'ai.onnx', opset.version) for opset in model.opset_import] — for QDQ insertion, ai.onnx opset must be ≥ 13 (≥ 19 recommended for full INT4/FP8 support)
• Custom domains: presence of com.amd.quark or com.microsoft indicates the model has already been processed by Quark or ORT and should not be naively re-quantized
• External data: os.path.exists(model_path + '_data') or any tensor with tensor.data_location == TensorProto.EXTERNAL; total file size > 2 GB without external data is invalid
• Inputs / outputs: name, dtype (from type.tensor_type.elem_type), shape (dim_value or symbolic dim_param)
• Dynamic dims: any dim_param or dim_value == 0 in inputs → flag as risk for NPU targets
• Input layout guess: 4D float input with channel dim == 3 or small power-of-2 → likely NCHW (channel-first) or NHWC (channel-last); use the convention of the producer
• Architecture guess: histogram of op types — Conv-dominated → CNN; MatMul/Gemm + softmax → Transformer; LSTM / GRU / RNN → recurrent; mix of Conv + MatMul + Attention → hybrid

Quantization Targets

• Quantizable op counts: count Conv, Gemm, MatMul, ConvTranspose, Add, Mul, BatchNormalization (will be folded into Conv during pre-processing)
• linear_layer_count: sum of MatMul + Gemm + Conv (the canonical "linear ops" in ONNX)
• exclude_defaults: usually empty for ONNX (Quark doesn't have a lm_head-equivalent default), but flag the final classifier Gemm/MatMul if accuracy is at risk
• Already-quantized check: presence of QuantizeLinear / DequantizeLinear / BFPQuantizeDequantize / MXQuantizeDequantize / ExtendedQuantizeLinear → set model.onnx_specific.has_qdq_already = true or has_quark_custom_ops = true

Deployment-Target Compatibility

Cross-reference op types against the chosen target:

• NPU CNN (enable_npu_cnn=True, XINT8): expects NHWC, PoF2 scales, per-tensor symmetric INT8, no MatMul-based attention. Flag NCHW inputs, dynamic dims, MatMul-heavy graphs.
• NPU Transformer (enable_npu_transformer=True): QDQ on Gemm/MatMul only; flag heavy use of unsupported ops (e.g. custom attention kernels).
• CPU/CUDA/ROCm general: largely permissive; flag custom-domain ops that require matching custom-ops library at inference.
• Block formats (BFP16 / MX / MXFP)*: require Quark's custom-ops library compiled for the target EP. If onnx_install_result.json shows the custom-ops compile failed, suppress these presets from preset_candidates and emit a risk.

Calibration Data Considerations

Surface what the user must prepare separately:

• Input names + dtypes + shapes — the CalibrationDataReader must yield batches keyed by these names with these dtypes
• Suggested sample count: 64–512 for CV; 128–1024 for LLM; reduce if optimize_mem=True is needed
• Flag dynamic batch / seq dims that the data reader must concretize

Risks

Flag anything that could cause failures downstream:

• Model file > 2 GB without external data → ProtoBuf serialization will fail
• Opset < 13 → QDQ insertion may fail; suggest onnx.version_converter.convert_version(model, 13+)
• Dynamic dims in input shape → NPU targets and many post-processing steps don't tolerate unconcretized dims; recommend quark.onnx.tools.fix_shapes
• NCHW inputs for an NPU CNN target → recommend quark.onnx.tools.convert_nchw_to_nhwc
• FP16-only model → some calibrators expect FP32; recommend quark.onnx.tools.convert_fp16_to_fp32 first
• Model already contains QDQ or Quark custom ops → do not re-quantize; suggest quark.onnx.tools.remove_qdq or treat as a no-op
• Custom-domain ops present (com.microsoft, third-party) → flag, may require op_types_to_quantize curation
• Unusual loading requirements (encrypted / crypto mode) → require the user to pass an onnx.ModelProto rather than a path (see input_check.py:137)