Quantization (torch.ao)

Active contributors: jerryzh168, hdcharles

Purpose

torch.ao ("Architecture Optimization") is the home of PyTorch's quantization, sparsity, and pruning. The quantization stack supports three distinct workflows:

1. Eager Mode Quantization — module-replacement based; the user manually inserts QuantStub/DeQuantStub and runs the quantization workflow on an nn.Module.
2. FX Graph Mode Quantization — graph-rewrite based on top of torch.fx. The user provides a config; FX rewrites the graph automatically.
3. PT2 Export Quantization (PT2E) — newer, built on torch.export and torch.compile. The recommended path going forward; backends register Quantizer plugins that pick patterns to quantize.

The package lives at torch/ao/ (for the modern code) with a deprecated torch/quantization/ shim still pointing at the old paths.

Directory layout

Path	Contents
torch/ao/quantization/	Quantization
torch/ao/quantization/observer.py	Min/max, histogram, per-channel observers
torch/ao/quantization/fake_quantize.py	Fake-quant modules used in QAT
torch/ao/quantization/quantize.py	Eager-mode quantization API
torch/ao/quantization/quantize_fx.py	FX-mode quantization API
torch/ao/quantization/pt2e/	PT2E quantization
torch/ao/quantization/quantizer/	Quantizer ABCs and built-in quantizers (X86, XNNPACK, …)
torch/ao/quantization/qconfig.py	QConfig (per-module quantization config)
torch/ao/quantization/backend_config/	Backend-specific op patterns and dtypes
torch/ao/nn/	Quantized nn.Module variants
torch/ao/pruning/	Structured / unstructured pruning
torch/ao/sparsity/	Sparsity utilities
aten/src/ATen/native/quantized/	C++ quantized op kernels (CPU, CUDA, cudnn)
aten/src/ATen/native/ao_sparse/	Sparse quantized kernels
torch/quantization/	Deprecated shim that imports from torch.ao.quantization

Key abstractions

Type	File	Purpose
Observer	torch/ao/quantization/observer.py	Records statistics (min/max, histogram) of activations
FakeQuantize	torch/ao/quantization/fake_quantize.py	Simulates quantization in float during training (QAT)
QConfig	torch/ao/quantization/qconfig.py	(activation_obs, weight_obs) pair for a module
BackendConfig	torch/ao/quantization/backend_config/backend_config.py	What ops/dtypes a backend supports
Quantizer	torch/ao/quantization/quantizer/quantizer.py	PT2E plugin: selects patterns and emits annotations
QuantStub / DeQuantStub	torch/ao/quantization/stubs.py	Manual eager-mode markers
prepare_fx / convert_fx	torch/ao/quantization/quantize_fx.py	FX flow entry points
prepare_pt2e / convert_pt2e	torch/ao/quantization/pt2e/prepare.py	PT2E flow entry points

How it works

Eager-mode flow

1. User wraps a model with QuantStub/DeQuantStub and assigns a qconfig to relevant modules.
2. prepare(model) swaps each leaf module with an observed version that records activation statistics during calibration runs.
3. The user calibrates by running representative inputs.
4. convert(model) swaps observed modules with their quantized counterparts (nn.quantized.Linear, nn.quantized.Conv2d, …) using the recorded scale/zero-point.

This is the oldest API and the one with the most mileage in production.

FX flow

graph LR
    Model[nn.Module] -->|symbolic_trace| FX[FX graph]
    FX -->|prepare_fx<br/>insert observers| FXObs[Observed graph]
    FXObs -->|calibration runs| Stats[Observer stats]
    Stats -->|convert_fx<br/>rewrite to quantized ops| FXQ[Quantized graph]
    FXQ -->|jit.script or jit.trace| Deploy[Deployable model]

prepare_fx/convert_fx use the BackendConfig to know which op patterns to fuse and what dtypes are valid for the target backend (FBGEMM/QNNPACK/X86/XNNPACK). Pattern matching lives in torch/ao/quantization/fx/.

PT2E flow

graph LR
    Model[nn.Module] -->|torch.export| EP[ExportedProgram]
    EP -->|prepare_pt2e<br/>quantizer.annotate| Annotated[Annotated graph]
    Annotated -->|calibration| Stats[Observer stats]
    Stats -->|convert_pt2e<br/>fold into quantized ops| Q[Quantized ExportedProgram]
    Q -->|torch.compile| Compiled[Compiled inference]

A Quantizer (e.g., XNNPACKQuantizer, X86InductorQuantizer) tells the framework which subgraphs to quantize and how. The framework inserts the right observer/fake-quant nodes; calibration fills them in; convert swaps in quantized ops. PT2E composes with torch.compile so quantized models can be Inductor-compiled.

Quantized ops

The quantized op kernels live in aten/src/ATen/native/quantized/. Backends:

• cpu/ — FBGEMM, QNNPACK, XNNPACK
• cuda/ — limited CUDA kernels
• cudnn/ — cuDNN-backed quantized convs

Quantized tensors carry quint8 / qint8 / qint32 / quint4x2 dtypes (with associated quantization params) and use a per-tensor or per-channel scale + zero-point representation.

Integration points

• torch.export is the upstream graph capture for PT2E.
• torch.compile / Inductor consumes PT2E-quantized graphs; Inductor has dedicated lowerings for quantized ops (in torch/_inductor/fx_passes/).
• ONNX export can export both FX-mode and (better) PT2E-quantized graphs.
• Mobile / ExecuTorch is the primary deployment target for XNNPACK-quantized models.

Entry points for modification

• New observer/fake-quant → torch/ao/quantization/observer.py or fake_quantize.py.
• New quantized op → C++ kernel under aten/src/ATen/native/quantized/<backend>/, schema in native_functions.yaml (with the Quantized* dispatch keys), and a Python nn.quantized wrapper.
• New backend → write a Quantizer subclass plus a BackendConfig.
• For debugging quantization correctness, torch.ao.ns (numeric suite) compares per-layer outputs across float and quantized versions.

Key source files

File	Purpose
torch/ao/quantization/observer.py	Observers
torch/ao/quantization/fake_quantize.py	Fake-quant modules
torch/ao/quantization/quantize_fx.py	FX flow
torch/ao/quantization/pt2e/prepare.py	PT2E prepare
torch/ao/quantization/pt2e/convert.py	PT2E convert
torch/ao/quantization/quantizer/xnnpack_quantizer.py	Reference Quantizer
torch/ao/quantization/backend_config/backend_config.py	BackendConfig
aten/src/ATen/native/quantized/cpu/	CPU quantized kernels
torch/ao/nn/quantized/	Quantized nn.Modules