How to statically quantize a PyTorch model (Eager mode)

Overview (static quantization in PyTorch, eager mode / “ego mode” in subtitles)

PyTorch has multiple quantization frameworks:

• Eager mode (most familiar; runs like normal PyTorch eager execution)
  • Considered “dangerous” without robust unit testing.
• FX graph mode (more graph-based; transforms a computational graph / DAG to optimize).
• Torch export (work in progress in the Torch 2.0 ecosystem; also graph-based).

This talk focuses on the theory of static quantization in eager mode, and sets up future videos:

• A hands-on coding tutorial quantizing ResNet in eager mode and FX mode.
• A future comparison of strengths/weaknesses between eager and FX.
• A future video on dynamic quantization.

---

Required model architecture changes for eager-mode static quantization

To convert a floating-point model into a quantized (inference) model, several changes are highlighted:

1) Insert QuantStub and DeQuantStub

Add:

• a Quant stub at the start
• a DeQuant stub at the end

Flow:

1. Float tensor enters → QuantStub converts it to a quantized tensor.
2. Operations inside run in quantized space.
3. DeQuantStub converts output back to float.

Notes:

• Stubs correspond to quantization boundaries.
• The Quant wrapper (QuantWrapper) can wrap a whole module and automatically add stubs.

---

2) Replace unsupported tensor ops using torch.nn.quantized.FloatFunctional

Many standard ops (e.g., add/sub/mul/div/concat-like patterns) don’t directly work on quantized tensors because each op may require a requantization step with its own quantization parameters.

Solution: use torch.nn.quantized.FloatFunctional methods that model the extra quantization/requantization correctly.

Supported functional operations (as stated):

• add
• add_scalar
• mul
• mul_scalar
• add_relu (described as “add + ReLU”)
• concat

Important:

• Division is not supported in these float functionals (no quantized kernel mentioned).
  • Suggested workaround: perform the division in floating-point by temporarily using DeQuantStub, then re-quantize output (return it to quantized space).
  • Possible future implementation is hinted.

Practical tip:

• For easier model inspection / state_dict navigation, wrap float functionals inside an nn.Module with a descriptive name (e.g., "float_functional_add", "float_functional_concat").

---

Module fusion (highly recommended)

After architecture changes, fuse modules to:

• Minimize unnecessary requantization
• Reduce quantization error
• Speed up inference

Rationale example:

• Fusing patterns like Conv/ReLU avoids separate quantization/requantization around intermediate steps when quantization parameters could otherwise be redundant.

PyTorch provides fusion support for common patterns, via two approaches:

1. Name-based fusion (fragile; can silently fail if layer names change)
2. Graph-based fusion (more robust; traces/fuses patterns based on actual structure, e.g., Conv followed by ReLU)

Hardware consideration:

Fusion effort should be targeted to the hardware backend (e.g., if targeting NVIDIA, confirm fused ops are supported).

---

Attach QConfig (“quantization recipes”) per layer

QConfig defines:

• Symmetric vs affine quantization
• Initial quantization params
• Whether params are:
  • fixed
  • observed (PTQ)
  • learned (QAT)
• Choice of PTQ observer types (e.g., minmax, histogram-based, per-channel, etc.)
• Learned step size concepts, including grad scaling (related to the Learnable Step Size Quantization paper)

Example described:

• learnable fake quant (advanced QAT-like control)
• minmax observer for activation ranges
• per-tensor affine quantization in the example
• Weight quantization may use a standard observer by default, but the speaker recommends customizing weight QConfig when coding.

---

Prepare → Fake-quant model → Convert → True quantized model

Step 1: torch.ao.quantization.prepare_model_inplace(...)

One core line of code:

• Transforms the model into a “quantizable” model (still float computation, but with fake quant machinery).

Result:

• The state_dict structure largely remains similar, but layers gain attributes such as:
  • activation_post_process (activation quantization)
  • weight_fake_quant (weight quantization)

---

Step 2: Fake-quant control (PTQ vs QAT toggles)

Fake quant/observer components include flags like:

• observer_enabled / static_enabled (controls whether PTQ updates quantization parameters)

• fake_quant_enabled (controls whether fake quantization is applied)

Behavior described:

• If fake quant enabled:
  • Fake quantize tensors using scale and zero point
  • Apply rounding/clamping (quantization simulation)
• If fake quant disabled:
  • Layer outputs remain floating point.

Control granularity:

• Fake quant can be enabled for activations only, weights only, or both.

---

Step 3: torch.ao.quantization.convert(...)

Converts the fake-quant model into a true quantized inference model.

Common conversion/runtime issues mentioned:

• Device/backend constraints
  • Often need CPU tensors and CPU model
  • Example: QNNPACK often requires CPU
• Quant/DeQuant stubs must be correctly placed
  • Feeding float directly into quantized layers without QuantStub causes errors
  • Double-quantizing or dequantizing tensors already in the wrong state can cause errors
• Float functionals must be used correctly
  • Using + directly can fail; use the proper FloatFunctional API
• Errors about non-leaf tensors in the model state_dict
  • Conversion uses a deep copy and may fail with non-leaf tensors
  • Fix: remove/clear intermediary tensors from the model state_dict before converting

---

Future content preview

Next video: a hands-on “from scratch” tutorial quantizing official ResNet using:

• eager-mode quantization
• FX-mode quantization

Also planned:

• dynamic quantization video
• deeper comparison between eager and FX (strengths/weaknesses)

---

Main speakers/sources

• Speaker: the video presenter (no name provided in subtitles)
• Technical sources referenced:
  • PyTorch quantization frameworks: Eager mode, FX graph mode, Torch export
  • Learned Step Size Quantization (LSQ) paper
  • PyTorch quantization modules/APIs: QuantStub / DeQuantStub, FloatFunctional, prepare_model_inplace, convert, and fusion utilities