Accepted Model Formats - PyTorch

PyTorch

In 26.09 we shipped our initial support for PyTorch. The Chimera Graph Compiler (CGC) now accepts a quantized PyTorch model directly. Two libraries carry the model there: quadric_pyquant quantizes it, and ChiPy — CGC's Python front end — traces it with torch.export and lowers it to CGC's internal graph representation, without a conversion to ONNX.

Accepted Model Formats - Quantized ONNX remains the broader ingestion path: it accepts models from any training framework and covers the full operator surface CGC supports. Our PyTorch support begins with ResNet-18, which this path carries from a checkpoint all the way to a compiled program. Other architectures go through ONNX today, and future releases will have expanded PyTorch support.

Why a direct path

Quantization state lives on the PyTorch module — each quantized layer carries its own calibrated scales and zero points. Exporting to ONNX and re-reading it turns that state into graph structure, which the compiler then has to reconstruct by walking the graph. Reconstruction is lossy in both directions: a scale that cannot be recovered becomes a compile failure, and a scale recovered incorrectly can change the model's numerics.

ChiPy reads each value from the module that owns it:

• the dequantize belongs to the producer's output and uses the producer's own output scale
• the quantize belongs to the consumer's input and uses the consumer's own input scale

A layer supplying no quantization parameters is a float32 layer, and the compiler treats it as one.

Export also applies no decompositions, so operators keep the granularity they had in the source model and opaque submodules stay opaque. The compiler receives the module you quantized, with its structure intact.

Requirements

Python 3.10 or later. The Python packages come from the SDK's own extras, so install those rather than pinning versions by hand:

pip install -e ".[qq-core]"   # the quantization engine, for a model you wire yourself
pip install -e ".[qq]"        # adds torchvision and transformers, which the example below uses

setup.cfg is the authoritative source for what each extra requires.

quadric_pyquant is the quantization library and ChiPy is the front end. Both ship inside the Chimera SDK and carry no version of their own — their version is the SDK release you installed.

Both are importable from a working SDK installation:

from quadric_pyquant.model_conversion import quantize_model
from tvm.contrib.epu.chipy import export, Tensor

Preparing a PyTorch Model

Quantize the model, confirm it is in the state the exporter can consume, then export it.

Note that the code in these steps follows a worked example. It follows one model, ResNet-18 at 224x224 on a single Chimera GPNPU configuration, so that each step has something concrete to act on. Every model name, input shape, calibration setting and hardware value is a stand-in to illustrate how PyTorch compilation works.

Step 1: Quantize with quadric_pyquant

quadric_pyquant performs post-training quantization (PTQ). It replaces nn.Linear, nn.Conv2d, nn.LayerNorm, nn.Softmax, nn.RMSNorm, activations and matrix multiplications with quantized equivalents, calibrates their ranges on representative data, and freezes the weights to their integer form.

fold_batchnorm() runs first, on a model in eval mode. BatchNorm has no quantized equivalent, so a layer left unfolded stays in float and reaches the exporter as an operator it cannot convert. Folding it into the preceding convolution removes it from the graph and costs nothing numerically. It returns the number of pairs folded, and warns about any BatchNorm it could not fold.

import torch
from torchvision.models import resnet18
from torch.utils.data import DataLoader

from quadric_pyquant.model_conversion import (
    fold_batchnorm,
    generate_quant_config,
    quantize_model,
    freeze_model_weights,
)
from quadric_pyquant.calibration import run_calibration, tag_all_modules

model = resnet18(weights="IMAGENET1K_V1").eval()

## Fold BatchNorm into the convolution before it
fold_batchnorm(model)

## Assign a layer_name to every module
tag_all_modules(model)

## Auto-detect quantizable layers
config = generate_quant_config(model)

## Put every layer on the integer path ChiPy converts
config["defaults"].update(
    int8_matmul=True,
    quantize_output=True,
    w_per_channel=False,
)

## Replace nn.Linear / nn.Conv2d with QLinear / QConv2d
quantize_model(model, config)

## Calibrate ranges on representative data
run_calibration(model, DataLoader(calibration_dataset, batch_size=32), num_batches=50)

## Bake the integer weights
freeze_model_weights(model)

Note that export stops at the first layer whose quantization settings it cannot consume, and names the setting it needs. quadric_pyquant's README.md documents the full rule set as its compiler deployability gate. A configuration that does not pass it may still deploy through a hand-written custom CCL kernel — see Handling Unsupported or Custom Operators.

Note that calibration data should be representative of the deployment distribution. How much calibration data to use, which settings to reach for, and how to close an accuracy gap are covered by Tutorial: Quantizing and Validating a Neural Network.

Step 2: Confirm the export-ready state

A quantized model is export-ready when it is in production deployment state: post-calibration, post-freeze, and on the int8 deployment path. Make sure to check this explicitly, because the failure modes otherwise surface from inside the tracer, where the error names only PyTorch internals:

from quadric_pyquant.model_conversion import validate_model_export_ready

validate_model_export_ready(model)   # raises, naming the layer and the violated condition

Every quantized layer must satisfy:

A layer sizes its operand scale and zero-point buffers on its first forward pass, and decides whether to do so by inspecting the buffer's values. Tracing that first forward therefore makes the decision data-dependent, and torch.export aborts with a GuardOnDataDependentSymNode error that names only torch internals. A model loaded from a checkpoint must be run once before it is exported:

with torch.no_grad():
    model(example_input)

EXPORT_CONTRACT.md in quadric_pyquant/ is the authoritative contract: what the export-ready state is, what each converter reads, and which code paths are deliberately not exported.

Step 3: Export

export() traces the model with torch.export and maps each quantized layer onto a CGC operator through the module-converter registry. Input shapes and dtypes are declared per forward() parameter name.

The shape below is this example's, and yours will be whatever your model's forward() takes:

from tvm.contrib.epu.chipy import export, Tensor

relay_mod = export(
    model,
    x=Tensor(shape=(1, 3, 224, 224), dtype="float32"),
)

Parameters with default values may be omitted; torch.export uses the defaults.

The result is a graph of CGC operators with the quantization boundaries explicit at the edges. Here's an example of what this looks like for a model small enough to print in full, with one convolution and one linear layer.

def @main(%x: Tensor[(1, 3, 8, 8), float32], output_names=["out_0"]) {
  %0 = qnn.quantize(%x, 0.0244795f, -5, out_dtype="int8");
  %1 = contrib.epu.qlinear_conv2d(%0, meta[relay.Constant][0], bias=meta[relay.Constant][1],
         padding=[1, 1, 1, 1], channels=8, kernel_size=[3, 3], out_dtype="int8",
         bias_fusion=True, x_zero_point=-5, x_scale=0.0244795f, w_zero_point=0,
         w_scale=0.00149468f, y_zero_point=0, y_scale=0.0154259f);
  %2 = qnn.dequantize(%1, 0.0154259f, 0);
  %3 = nn.relu(%2);
  %4 = nn.global_avg_pool2d(%3);
  %5 = reshape(%4, newshape=[1, 8]);
  %6 = qnn.quantize(%5, 0.00104517f, -128, out_dtype="int8");
  %7 = contrib.epu.dense(%6, meta[relay.Constant][2], meta[relay.Constant][3], units=4,
         x_scale=0.00104517f, x_zero_point=-128, w_scale=0.00250289f,
         y_scale=0.00218769f, y_zero_point=13, out_dtype="int8");
  qnn.dequantize(%7, 0.00218769f, 13)
}

Step 4: Compile

The exported graph enters CGC through its Relay entry point. Build a target describing the Chimera GPNPU configuration you are compiling for, then build the graph under it:

import tvm
import tvm.relay.backend.contrib.epu.codegen as epu_codegen
import tvm.relay.backend.contrib.epu.util as epu_utils

## Every value below is a property of the Chimera GPNPU you are targeting.
ARRAY_SIZE = ...     # dimension of the PE array
MACS_PER_PE = ...    # multiply-accumulate units per processing element
OCM_SIZE = ...       # L2 memory, in bytes
LRM_SIZE = ...       # local register memory, in bytes
NUM_CORES = ...

target = tvm.target.epu(
    model="",
    array_size=ARRAY_SIZE,
    macs_per_pe=MACS_PER_PE,
    ocm_size=OCM_SIZE,
    lrm_size=LRM_SIZE,
    num_cores=NUM_CORES,
)

with target, epu_utils.add_epu_passcontext_attributes():
    epu_codegen.build_relay(relay_mod, module_name="resnet18_torch")

A target built from these arguments describes an array and a memory hierarchy, which is enough to compile. A licensed product configuration carries settings beyond the geometry, and the ONNX path expresses those through the HWConfig class shown in Tutorial: Using SDK as a Library.

build_relay returns a ChimeraModule and writes a complete program to disk:

The compiler writes everything under ccl_build in the current working directory, in a folder named for the module_name you passed to build_relay. These helpers return the two paths, so you do not have to construct them:

epu_codegen.get_module_source_path("resnet18_torch")   # ./ccl_build/resnet18_torch
epu_codegen.get_module_build_path("resnet18_torch")    # ./ccl_build/resnet18_torch/build

Set the build_path pass-context attribute to write somewhere other than ccl_build.

The build directory is nested inside the source directory, so a single tree holds the generated program and everything built from it.

The .cpp can be read, modified and rebuilt like any other CCL program, and compiled with the Chimera LLVM C++ Compiler. To measure the compiled program, see Chimera Instruction Set Simulator (ISS).

What ChiPy Supports

Quantized layers

These quantized layer types have module converters:

Framework operators

torch.export traces standard ATen operators as the glue between quantized layers. ChiPy converts these overloads:

Each overload is registered by its full name — view.default, add.Tensor, select.int — so an unregistered overload produces a clear "no converter" error naming the operator. Note that reshape has no converter of its own: PyTorch traces both reshape and flatten to view.default on a contiguous tensor.

Supported models

This is our initial PyTorch support, and it begins with ResNet-18: it goes from a PyTorch checkpoint, through quantization and export, to a compiled program, and it is the model the worked example above follows.

We'll keep this guide updated as the set of supported PyTorch models grows with future releases.

Related Documentation

• Accepted Model Formats - Quantized ONNX — the other ingestion path
• Tutorial: Quantizing and Validating a Neural Network — the quantization workflow end to end
• Handling Unsupported or Custom Operators — what to do with a layer no converter covers
• Overview of the CGC — what happens after ingestion