Brevitas

Brevitas is a platform for modeling reduced precision hardware data-paths during training. It provides building blocks at various abstraction levels to support both research into new quantization-aware training (QAT) techniques and the practical application of existing techniques to models.

Key features include:

• Unified API: Supports a super-set of quantization schemes across different frameworks and compilers.
• Inference Acceleration: Supports exporting to FINN, onnxruntime, or PyTorch's native quantized operators for specific layer and quantization combinations.
• Quantization Style: Primarily implements affine quantization with a focus on uniform quantization. Note that non-uniform quantization is not supported out-of-the-box.

The ImageNet examples are categorized into two main workflows:

• PTQ (Post-Training Quantization): Demonstrates how to apply PTQ to models defined either 'by hand' (manually modified pretrained models) or 'programmatically' (starting from torchvision definitions).
• QAT (Quantization Aware Training): Provides examples of running inference on networks trained with QAT. These examples focus on low-precision (< 8-bit) models designed for export and acceleration in downstream compilers. Note that the full training code for these specific QAT models is not included; only the quantized model definitions and inference examples are provided.

Detailed information for each workflow can be found in the ptq and qat subdirectories. Quantized models defined by hand are located in the models folder.

The src/brevitas_examples directory contains several specialized example sets for different quantization and deployment scenarios:

• bnn_pynq: Focuses on simple Quantization Aware Training (QAT) at 1, 2, and 4 bits, specifically designed for deployment using the FINN framework.
• imagenet_classification: Provides Post-Training Quantization (PTQ) examples and pretrained 4-bit QAT models for ImageNet classification tasks.
• speech_to_text: Contains a pretrained 8-bit QAT MelGAN Vocoder intended for integration into text-to-speech pipelines.
• super_resolution: Features pretrained 8-bit QAT models for increasing RGB image resolution using Accumulator-Aware Quantization (A2Q).
• text_to_speech: Includes a pretrained 8-bit QAT QuartzNet model for Automatic Speech Recognition (ASR).

Brevitas provides Post Training Quantization (PTQ) features to quantize torchvision models or models manually defined with Brevitas quantized layers (e.g., MobileNet V1). This workflow can serve as a starting point for Quantization Aware Training (QAT) fine-tuning, particularly at lower precisions.

There are two primary workflows provided:

1. Evaluation script: Performs PTQ on a model given a name and quantization configuration, evaluates ImageNet top1 accuracy, and optionally exports the model to ONNX or TorchScript QCDQ.
2. Benchmark suite: Tests various quantization configurations across selected models.

Note: Programmatic quantization is experimental and may break on certain models or configurations. The implementation relies on torch.fx symbolic tracing, so the input model must support symbolic tracing.

Qronos (presented at ICLR 2026) is a Post-Training Quantization (PTQ) algorithm designed to address quantization errors in both weights and activations. Unlike standard PTQ techniques that solve a layerwise reconstruction problem in isolation, Qronos addresses the "mismatched" reconstruction problem by considering the potentially quantized inputs ($\tilde{X}$) from previously quantized layers.

The algorithm quantizes weights sequentially using two alternating steps:

1. Error Correction: Selecting a quantized weight that optimally corrects the current approximation error while holding other weights fixed.
2. Error Diffusion: Updating the remaining unquantized weights to compensate for the accumulated rounding error.

Qronos utilizes techniques like Cholesky decomposition and lazy batch updates to scale efficiently to large models.

The Learned Round implementation in Brevitas allows for learnable rounding during quantization. The experiments described in this directory compare Learned Round against several other techniques across two main phases:

Stage 1: Transform (Preprocessing)

• None: No preprocessing.
• HIP: Hadamard-based incoherence processing.
• MagR: Weight magnitude reduction.
• QuaRot: Rotation-based outlier reduction.
• SpinQuant: Cailey-optimized orthogonal rotations.

Stage 2: Rounding

• RTN: Round-To-Nearest.
• GPTQ: GPTQ quantization.
• Qronos: Qronos quantization.
• Learned Round: Brevitas' implementation of learnable rounding.

Brevitas uses custom torch.autograd.Function implementations. While Python implementations are used by default, you can enable C++ implementations to support end-to-end compilation of training-time quantizers when using JIT.

To enable this, set the environment variable BREVITAS_JIT=1.

Note on C++ Requirements: To use the C++ implementation, you must have a PyTorch compatible C++ compiler installed. Brevitas compiles the custom C++ autograd functions at runtime using torch.utils.cpp_extension.load(). This is primarily useful for distributed training scenarios.

This directory provides implementations for advanced quantization techniques designed to avoid accumulator overflow:

• A2Q (Accumulator-Aware Quantization): As proposed in A2Q: Accumulator-Aware Quantization with Guaranteed Overflow Avoidance.
• A2Q+: As proposed in A2Q+: Improving Accumulator-Aware Weight Quantization. This version introduces the AccumulatorAwareZeroCenterWeightQuant (a zero-centered weight quantizer) to improve performance.

Brevitas allows you to share instances of weight quantizers among multiple layers. This forces the layers to share the same scale, zero-point, and bit-width. When a quantizer is shared, it looks at all participating weight tensors to determine a single global scale factor (e.g., based on the overall maximum value).

This is designed to be eager-mode friendly; the quantizer is re-initialized appropriately whenever it is shared with a new layer.

Advanced users can define entirely new quantization algorithms using Brevitas's dependency-injection system. Instead of using enums, you inherit from ExtendedInjector and define a tensor_quant module (which must be a brevitas.jit.ScriptModule) and a proxy_class (which returns a QuantTensor).

The tensor_quant module is assembled by matching its keyword arguments to the attributes defined in your ExtendedInjector class. This allows for highly modular construction of quantization logic (rounding, scaling, clamping, etc.).

from brevitas.inject import ExtendedInjector
from brevitas.proxy import ActQuantProxyFromInjector

class MyCustomQuantizer(ExtendedInjector):
    proxy_class = ActQuantProxyFromInjector
    tensor_quant = RescalingIntQuant 
    int_quant = IntQuant
    float_to_int_impl = RoundSte 
    # ... other components ...

The diffusion quantization entrypoint supports several Post-Training Quantization (PTQ) techniques for both integer and floating point quantization:

• Activation Equalization: e.g., SmoothQuant, layerwise (requires addition of Mul ops).
• Activation Calibration: Used for static activation quantization.
• GPTQ
• SVDQuant
• Bias Correction