DOCSEARCHbyAlgolia
Submit Docs

QONNX (Quantized ONNX)

repository·main·Indexed 41 days ago

https://github.com/fastmachinelearning/qonnx

A framework for representing and manipulating arbitrary-precision quantized neural networks within the ONNX format. QONNX provides custom operators—including IntQuant, BipolarQuant, FloatQuant, and Trunc—to support diverse quantization schemes such as integer, minifloat, and OCP MX datatypes. It includes utilities for functional verification, shape inference, constant folding, and calculating inference costs (BOPS and memory volume), as well as tools for exporting Brevitas models to QONNX format.

Tokens
24.4K
Snippets
70
Records
91
Agent score
52%

What's inside QONNX

  1. What is QONNX?

    main

    QONNX (Quantized ONNX) is a framework for representing arbitrary-precision uniform quantization within ONNX. It uses four custom operators to enable representation of binary, ternary, 3-bit, 4-bit, 6-bit, or any other quantization, as well as quantized floating-point values.

    Key features:

    • Quantization as an operator: Can be applied to any parameter or layer input.
    • Flexible Granularity: Supports various choices for scaling factor and zero-point granularity.
    • ONNX Compatibility: Quantized values are carried using standard float datatypes to maintain ONNX protobuf compatibility.

    QONNX provides Python utilities for:

    • Functional verification (execution).
    • Shape inference, constant folding, and basic optimizations.
    • Summarizing inference costs (mixed-precision MACs, parameter/activation volume).
    • Writing transformations and defining custom ops.
    • (Experimental) Data layout conversion from NCHW to custom NHWC ops.

  2. QONNX Custom Quantization Operators

    main

    QONNX introduces four custom operators to handle arbitrary-precision quantization:

    • IntQuant: Used for 2-to-arbitrary-bit quantization, supporting scaling and zero-point.
    • BipolarQuant: Used for 1-bit (bipolar) quantization, supporting scaling and zero-point.
    • FloatQuant: Used for arbitrary-precision-float-quantized values.
    • Trunc: Used for truncating to a specified number of bits, supporting scaling and zero-point.

  3. Implement Analysis and Transformation passes

    main

    QONNX provides a framework for traversing and modifying graphs through two types of passes:

    Analysis Passes

    An analysis pass traverses the graph to extract properties.

    • Input: A ModelWrapper instance.
    • Output: A dictionary containing the extracted properties.
    • Module: qonnx.analysis contains existing passes.

    Transformation Passes

    A transformation pass modifies the graph structure or content.

    • Input: A ModelWrapper instance.
    • Output: A tuple containing the modified ModelWrapper and a boolean flag model_was_changed. If model_was_changed is True, the pass may need to be re-run.
    • Module: qonnx.transformation contains existing passes.

  4. Identify QONNX custom operations

    main
    QONNX utilizes custom operations that are not part of the standard ONNX operator schema. These nodes are identified in the ONNX protobuf by having a domain set to "qonnx.*". These operations may represent specific low-bit network requirements or hardware-specific optimizations.

  5. Understand QONNX custom quantization annotations

    main

    Since standard ONNX does not support datatypes smaller than 8-bit, QONNX uses the quantization_annotation field in the ONNX protobuf to store QONNX-specific DataType information (e.g., ternary, bipolar).

    Important Implementation Details:

    • Storage: All tensors are stored using single-precision floating point (float32) as a carrier type to represent integer values. This allows for the representation of sub-8-bit values.
    • Rounding Errors: Because floating-point arithmetic is used to represent integers, rounding errors can occur (e.g., (int_num * float_scale) / float_scale != int_num).
    • Sanitization: The custom QONNX execution flow automatically attempts to sanitize these rounding errors using qonnx.util.basic.sanitize_quant_values.
    • Disabling Sanitization: You can disable this behavior by setting the environment variable SANITIZE_QUANT_TENSORS=0 (not recommended).

  6. Reference implementation of BipolarQuant

    main

    If you are implementing or testing BipolarQuant logic manually, the reference behavior follows these two steps:

    1. Quantizing: Convert the input tensor to a tensor of 1.0 (where input $\ge 0$) or -1.0 (where input $< 0$).
    2. Scaling: Multiply the resulting integer-like tensor by the scale factor.

    Mathematical representation: $Y = \text{sign}(X) \times \text{scale}$ (where $\text{sign}(0) = 1$)

    def binary_quant(inp_tensor, scale):
    # Quantizing
    y_int = inp_tensor
    y_ones = np.ones(y_int.shape, dtype=y_int.dtype)
    y_int = np.where(y_int >= 0.0, y_ones, -y_ones)
    # Scaling
    out_tensor = y_int * scale

    return out_tensor

  7. Use the FloatQuant custom operator

    main

    The FloatQuant operator calculates arbitrary-precision floating-point quantized values for an input tensor. It maps one input tensor to one output tensor using five additional inputs that define the quantization parameters (scale, exponent bitwidth, mantissa bitwidth, maximum representable value, and exponent bias).

    These parameters can be scalars or tensors with shapes broadcastable to the input tensor, allowing for per-tensor or per-element (granularity) scaling.

    Note: This operator is specifically for floating-point quantization. For integer quantization, use the IntQuant custom operator instead.

    import numpy as np

# Example setup for FloatQuant
x = np.random.rand(100).astype(np.float32)
scale = 1
exponent_bitwidth = 4
mantissa_bitwidth = 3
exponent_bias = 0
# max_val must be calculated based on bitwidths and bias
max_val = 1.0 
rounding_mode = 'ROUND'

# xq = float_quantize(x, scale, exponent_bitwidth, mantissa_bitwidth, exponent_bias, max_val, rounding_mode)

  11. Export a Brevitas model to QONNX format

    main

    To export a quantized model from Brevitas to the QONNX format, you must first ensure the Brevitas model is in evaluation mode (.eval()). You then use the export_qonnx function, providing the model, a dummy input of the appropriate shape, and the desired export path.

    This process allows you to convert Brevitas-specific quantization parameters into a standardized QONNX ONNX graph suitable for FloatQuant execution.

    # important to put into eval mode before export
brevitas_model.eval()

# Export the Brevitas model to QONNX format
export_path = "qonnx_act_weight_fp8.onnx"
input_shape = (1, 3)  # Example input shape, adjust as needed
dummy_input = torch.randn(input_shape)
export_qonnx(brevitas_model, dummy_input, export_path)