Quantization#

Embedl Deploy provides hardware-aware integer and FP8 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

output

TensorQuantConfig

Config for output 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 integer affine or FP8 E4M3 explicit QDQ)

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)

FP8 uses E4M3 with symmetric, amax-based MINMAX calibration. It is available for PTQ/export only: pass a forward_loop to quantize() and freeze weights before export. FP8 requires PyTorch 2.1 or newer, and classic ONNX export requires opset 19 or newer. The QDQ representation is backend-independent, but the selected target compiler must support FP8 execution.

ModulesToSkip#

ModulesToSkip has four 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

output

Which output quantizers to leave disabled

smooth

Which SmoothQuant observers to leave disabled

On a module whose weight (or output) quantizer is synced to its input stub, only stub disables quantization for that module – weight/output have no effect there, since a synced quantizer’s enabled state always follows the stub. weight/output only matter for modules with an unsynced weight or output quantizer.

On TensorRT, no module attaches an output quantizer at all, so output and skip.output currently have no effect regardless of sync state – configure() logs a warning if either is set. They exist for backends that attach their own output quantizer.

Recommended settings for TensorRT:

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

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

  • FusedLayerNorm never attaches a quantizer, so no skip entry is needed for it. Its floating-point type comes from the exported ONNX graph. To also disable its SmoothQuant migration, skip it via config.skip.smooth

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(smooth={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(
        # disables the whole lock-stepped set, input stub and weight alike
        stub={first_conv},
    )
)
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. This qualified name is attached to a module by tracing with the default trace=Trace.EXPORT; Trace.SYMBOLIC never attaches one. Trace.NONE skips tracing entirely and deep-copies the GraphModule as-is, so it only preserves a qualified name if the module already had one going in (e.g. it was produced by a prior Trace.EXPORT pass) – otherwise, as with Trace.SYMBOLIC, regex entries silently match nothing:

import re

config = QuantConfig(
    skip=ModulesToSkip(
        # disables the whole lock-stepped set, input stub and weight alike
        stub={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 ↔ floating point), 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.

TensorRT compilation#

Embedl Deploy exports explicit ONNX QuantizeLinear and DequantizeLinear nodes. These nodes specify where quantization occurs and which quantized type is used. Compile the model as a strongly typed network so TensorRT preserves the tensor types encoded in ONNX. See NVIDIA’s documentation for explicit quantization and strongly typed networks.

For TensorRT 10.x:

trtexec \
  --onnx=model_quantized.onnx \
  --stronglyTyped \
  --saveEngine=model_quantized.plan \
  --profilingVerbosity=detailed \
  --dumpLayerInfo \
  --exportLayerInfo=layer_info.json

Do not combine --stronglyTyped with --fp16, --int8, --fp8, or --best. Those flags permit opportunistic precision selection and can cause unquantized regions to run in an unintended reduced precision.

TensorRT 11 always uses strong typing and removes the precision-selection flags, so omit --stronglyTyped as well:

trtexec \
  --onnx=model_quantized.onnx \
  --saveEngine=model_quantized.plan \
  --profilingVerbosity=detailed \
  --dumpLayerInfo \
  --exportLayerInfo=layer_info.json

Floating-point fallback#

Precision.UNQUANTIZED means that Deploy does not insert QDQ at that location; it does not select FP16 or FP32. The exported ONNX graph determines the floating-point type. With an FP32 source model, Deploy’s FP32 quantization scales make DequantizeLinear outputs FP32, and strong typing keeps unquantized regions in FP32. Use explicit casts in the source graph if selected regions should use another floating-point type.

Start with the intended floating-point model and inputs before transformation and quantization. Do not cast the complete model after QDQ insertion because that can also cast quantization buffers.

Validation#

Engine creation alone does not prove that TensorRT used the intended precision. Validate with representative held-out data:

  1. Check that every output is finite.

  2. Compare against the FP32 model and the framework QDQ model.

  3. Measure the task-level accuracy metric.

  4. Inspect the exported layer information for the intended quantized types and floating-point fallbacks.

  5. Compare latency with a separately compiled floating-point baseline.

If strong typing does not restore accuracy, improve calibration coverage or use ModulesToSkip to exclude the first layer whose output diverges.

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 the explicit QDQ model with strong typing on TensorRT 10.x:

trtexec --onnx=resnet50_int8.onnx --stronglyTyped

Exporting in float16#

convert_float_to_float16 returns a half-precision copy of a quantized model, so the export writes a float16 graph directly. Weight quantization must be frozen first, which also means the export has to use the torch.export-based exporter:

from embedl_deploy.quantize import convert_float_to_float16, quantize

quantized = quantize(
    result.model, (example_input,), forward_loop=loop, freeze_weights=True
)
half = convert_float_to_float16(quantized)

torch.onnx.export(
    half,
    (example_input.half(),),
    "resnet50_int8_fp16.onnx",
    dynamo=True,
    opset_version=20,
)

The Q/DQ nodes keep their quantized precision, integer or FP8; what changes is the floating-point tensors around them. Operations that half precision degrades – interpolation, cumulative sums, top-k, and the torchvision detection ops – keep running in float32, with casts inserted at their boundaries. Everything else follows the model to float16, including tensors tracing lifted out of forward and dtype literals the trace recorded from the float32 model (a softmax(dtype=torch.float32), say), so the exported graph carries no stray float32 islands and no Cast pairs around the Q/DQ nodes. Build the resulting ONNX strongly typed (trtexec --stronglyTyped) to keep every tensor at the precision the file says.

To hold a module at float32 as well, name it with skip, which accepts the same entries as ModulesToSkip:

import re

half = convert_float_to_float16(
    quantized,
    skip={re.compile(r"layer4\..*"), torch.nn.Linear},
)

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)