Quantization#

Embedl Deploy provides hardware-aware INT8 quantization through explicit QDQ (Quantize/DeQuantize) stub placement. Unlike uniform quantization approaches that insert QDQ nodes around every operator, Embedl Deploy places stubs only at positions controlled by each FusedModule’s inputs_to_quantize — ensuring that quantization does not break operator fusions in the target hardware compiler.

Quantization pipeline#

The quantization pipeline has two steps:

┌─────────────┐     ┌─────────────────────────────────────┐
│ transform() │ ──▶ │            quantize()               │
│(fuse model) │     │(configure + calibrate + freeze)     │
└─────────────┘     └─────────────────────────────────────┘

  1. Transform — apply conversions and fusions. Each FusedModule subclass declares inputs_to_quantize (a set of positional input indices), and FusedModule.__init__ creates the corresponding QuantStub instances.

  2. Quantize — quantize() configures each stub with the supplied QuantConfig, inserts Q/DQ nodes into the graph, runs SmoothQuant calibration (if enabled), then runs activation calibration to compute scale and zero-point for each quantizer.

Step 1: Transform and fuse#

import torch
from torchvision.models import resnet50

from embedl_deploy import transform

model = resnet50(weights="DEFAULT").eval()
example_input = torch.randn(1, 3, 224, 224)
result = transform(model, (example_input,))
fused_model = result.model

Step 2: Quantize#

from embedl_deploy.quantize import (
    Precision,
    QuantConfig,
    TensorQuantConfig,
    quantize,
)

def forward_loop(model):
    for batch_tensor, _ in calibration_loader:
        model(batch_tensor)

quantized_model = quantize(
    fused_model,
    (example_input,),
    config=QuantConfig(
        activation=TensorQuantConfig(Precision.INT8, symmetric=True),
        weight=TensorQuantConfig(Precision.INT8, symmetric=True, per_channel=True),
    ),
    forward_loop=forward_loop,
    freeze_weights=True,
)

quantize() chains these internal steps:

  1. Configure — enables QuantStub and WeightFakeQuantize modules on each FusedModule, applies skip rules from QuantConfig.skip, and inserts Q/DQ stub nodes into the graph.

  2. SmoothQuant calibration — if any SmoothQuantObserver is enabled, redistributes quantization difficulty from activations to weights at LayerNorm → Linear boundaries.

  3. Activation calibration — runs the forward loop to collect activation statistics and compute scale / zero_point for each QuantStub.

  4. Freeze weights (optional) — when freeze_weights=True, computes and stores weight scales as constant buffers, required for ONNX/TensorRT export.

QuantConfig#

QuantConfig controls how quantization stubs are configured:

Parameter

Type

Description

activation

TensorQuantConfig

Config for activation quantizers

weight

TensorQuantConfig

Config for weight quantizers

smooth_quant

SmoothQuantConfig

SmoothQuant migration settings

skip

ModulesToSkip

Modules or module types to leave unquantized

TensorQuantConfig#

Parameter

Type

Default

Description

precision

Precision

Precision.INT8

Quantized format (INT4/INT8)

symmetric

bool

True

Use symmetric quantization

per_channel

bool

True

Per-channel (weights) vs per-tensor

calibration_method

CalibrationMethod

MINMAX

Calibration algorithm (MINMAX, MOVING_AVERAGE_MINMAX, HISTOGRAM)

ModulesToSkip#

ModulesToSkip has three independent sets, each accepting module types, module instances, or compiled regex patterns:

Field

Controls

stub

Which activation stubs to leave disabled

weight

Which weight fake-quantizers to leave disabled

smooth

Which SmoothQuant observers to leave disabled

Recommended settings for TensorRT:

  • Activations: 8-bit, symmetric, per-tensor

  • Weights: 8-bit, symmetric, per-channel

  • Skip weight quantization for LayerNorm (it runs in FP16 on TensorRT)

from torch import nn
from embedl_deploy.quantize import (
    ModulesToSkip,
    Precision,
    QuantConfig,
    TensorQuantConfig,
)

config = QuantConfig(
    activation=TensorQuantConfig(Precision.INT8, symmetric=True),
    weight=TensorQuantConfig(Precision.INT8, symmetric=True, per_channel=True),
    skip=ModulesToSkip(weight={nn.LayerNorm}),
)

Mixed precision#

To keep specific layers in higher precision while quantizing the rest to INT8, add them to ModulesToSkip. Each set accepts module types, module instances, or compiled regex patterns matched against the original module’s qualified name.

By instance — useful when you know the exact module. Note that torch.fx.GraphModule deep-copies submodules during tracing, so you must take the reference from the fused graph, not from the original model:

from embedl_deploy import transform
from embedl_deploy.quantize import quantize, QuantConfig, ModulesToSkip

res = transform(model, (example_input,))

# Grab the conv instance from the fused graph (not from the original model)
first_conv = res.model.FusedConvBNActMaxPool_0.conv

config = QuantConfig(
    skip=ModulesToSkip(
        stub={first_conv},    # disables input activation quantization
        weight={first_conv},  # disables weight fake-quantization
    )
)
quantized_model = quantize(
    res.model, (example_input,), config=config, forward_loop=forward_loop
)

By regex — useful when you want to exclude layers by name without inspecting the fused graph. The pattern is matched with re.fullmatch against each child module’s original qualified name:

import re

config = QuantConfig(
    skip=ModulesToSkip(
        stub={re.compile(r"layer4\..*")},
        weight={re.compile(r"layer4\..*")},
    )
)

QDQ stub placement#

Each FusedModule subclass declares inputs_to_quantize — the set of positional input indices for which QuantStub instances are created. This is the key difference from uniform quantization:

Fused module

QDQ points

Rationale

FusedConvBN

INPUT

Quantize activations entering the conv

FusedConvBNAct

INPUT

Quantize input; activation is fused, no separate stub needed

FusedConvBNAddAct

INPUT, RESIDUAL

Both paths into the add must be quantized

FusedConvBNActMaxPool

INPUT

First layer — quantize input image

FusedLinear

INPUT

Quantize activations entering the linear

FusedLinearAct

INPUT

Quantize input; activation is fused

FusedLayerNorm

(none)

Not quantized — hurts accuracy, no latency benefit

FusedMHAInProjection

INPUT

Quantize input to Q/K/V projection

FusedScaledDotProductAttention

(none)

Not quantized

FusedAdaptiveAvgPool2d

(none)

Not quantized — memory-bound, no INT8 benefit

FusedActAdd

INPUT, RESIDUAL

Both paths into the add must be quantized

Why pattern-aware QDQ matters#

The problem with uniform quantization#

Tools like NVIDIA ModelOpt apply QDQ stubs around every operator indiscriminately. This causes several issues:

  1. Broken fusions — QDQ nodes between operators that should be fused (e.g., between Conv and BN) prevent the hardware compiler from merging them into a single kernel.

  2. Reformatting overhead — quantizing memory-bound operators like depthwise convolutions or global average pooling forces TensorRT to insert data reformatting layers (INT8 ↔ FP16), which can cost more than the quantization saves.

  3. Accuracy loss without latency gain — quantizing operators that don’t benefit from INT8 (LayerNorm, element-wise ops) reduces accuracy with no performance improvement.

Embedl Deploy’s approach#

Pattern-declared QDQ points ensure:

  • Stubs are placed outside fused operator groups, not between them.

  • Memory-bound operators can be left unquantized when beneficial.

  • The hardware compiler sees exactly the QDQ topology it expects.

Full example: ResNet50 INT8 PTQ#

import torch
from torchvision.models import resnet50

from embedl_deploy import transform
from embedl_deploy.quantize import (
    Precision,
    QuantConfig,
    QuantStub,
    TensorQuantConfig,
    WeightFakeQuantize,
    quantize,
)

# 1. Load and fuse
model = resnet50(weights="DEFAULT").eval()
example_input = torch.randn(1, 3, 224, 224)
result = transform(model, (example_input,))

# 2. Quantize (configure + calibrate + freeze)
def forward_loop(model):
    for batch in calibration_batches[:32]:
        model(batch)

quantized = quantize(
    result.model,
    (example_input,),
    config=QuantConfig(
        activation=TensorQuantConfig(Precision.INT8, symmetric=True),
        weight=TensorQuantConfig(Precision.INT8, symmetric=True, per_channel=True),
    ),
    forward_loop=forward_loop,
    freeze_weights=True,
)

# Count inserted stubs
n_quant = sum(1 for m in quantized.modules() if isinstance(m, QuantStub))
n_wfq = sum(1 for m in quantized.modules() if isinstance(m, WeightFakeQuantize))
print(f"QuantStubs: {n_quant}, WeightFakeQuantize: {n_wfq}")

# 3. Export
torch.onnx.export(
    quantized.cpu().eval(),
    torch.randn(1, 3, 224, 224),
    "resnet50_int8.onnx",
    opset_version=20,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
)

Compile with TensorRT using --best (FP16 + INT8):

trtexec --onnx=resnet50_int8.onnx --best

SmoothQuant#

SmoothQuant redistributes quantization difficulty from activations to weights at LayerNorm → Linear boundaries. This improves INT8 accuracy for transformer models where activation distributions have large outliers.

from embedl_deploy.quantize import (
    Precision,
    QuantConfig,
    SmoothQuantConfig,
    TensorQuantConfig,
    quantize,
)

quantized = quantize(
    result.model,
    (example_input,),
    config=QuantConfig(
        activation=TensorQuantConfig(Precision.INT8, symmetric=True),
        weight=TensorQuantConfig(Precision.INT8, symmetric=True, per_channel=True),
        smooth_quant=SmoothQuantConfig(alpha=0.5),
    ),
    forward_loop=forward_loop,
    freeze_weights=True,
)

SmoothQuantConfig.alpha controls the migration strength: 0 keeps all difficulty on activations, 1 pushes it entirely to weights. 0.5 is a good starting point for most models.

Quantization-Aware Training (QAT)#

For higher accuracy, you can fine-tune the quantized model with QAT. Call quantize() without a forward_loop to insert and configure the Q/DQ stubs while skipping calibration — the scales are learned during training instead:

from embedl_deploy.quantize import (
    enable_fake_quant,
    freeze_bn_stats,
    freeze_weight_quantization,
    prepare_qat,
    quantize,
)

# Insert and configure Q/DQ stubs (no calibration)
quantized = quantize(result.model, (example_input,))

# Prepare for QAT (enables fake quantization in training mode)
prepare_qat(quantized)

# Fine-tune with your training loop
quantized.train()
enable_fake_quant(quantized)
freeze_bn_stats(quantized)

for epoch in range(num_epochs):
    for images, targets in train_loader:
        output = quantized(images)
        loss = criterion(output, targets)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

# Freeze weights and export
quantized.eval()
freeze_weight_quantization(quantized)