From 52f938462cd80e6cc11c421c4b2067c5812e9585 Mon Sep 17 00:00:00 2001 From: M Platypus Date: Mon, 2 Mar 2026 02:12:51 -0500 Subject: [PATCH 1/2] Add QAT training performance analysis to quantization docs Document the six major overhead sources in quantization-aware training vs float training: FakeQuantize nodes, observer statistics, power-of-2 scale snapping, soft-quantize variants, backward pass complexity, and torch.compile ordering. Co-Authored-By: Claude Opus 4.6 --- docs/source/features/quantization.rst | 149 ++++++++++++++++---------- 1 file changed, 95 insertions(+), 54 deletions(-) diff --git a/docs/source/features/quantization.rst b/docs/source/features/quantization.rst index 0a379558..2660fe59 100644 --- a/docs/source/features/quantization.rst +++ b/docs/source/features/quantization.rst @@ -81,8 +81,8 @@ Uses standard PyTorch quantization APIs (``GenericTinyMLQATFxModule`` / quantization_weight_bitwidth: 8 quantization_activation_bitwidth: 8 -For more details on the underlying wrappers, see -:ref:`quantization-wrapper-architecture` below. +For more details on the underlying wrappers, see the +`tinyml-modeloptimization documentation `_. **TI Style Optimised Quantization (quantization: 2)** @@ -331,7 +331,7 @@ Example: Full Quantization Workflow variables: 1 training: - model_name: 'CLS_4k_NPU' + model_name: 'ArcFault_model_400_t' training_epochs: 30 batch_size: 256 quantization: 2 @@ -410,69 +410,69 @@ Even without NPU, integer operations are faster: Float32: ~5000 µs INT8: ~2000 µs -.. _quantization-wrapper-architecture: +QAT Training Performance +------------------------ -Quantization Wrapper Architecture ----------------------------------- +Quantization-Aware Training is significantly slower than float training. +This section explains why, and which factors dominate the overhead. -Under the hood, Tiny ML Tensorlab uses quantization wrapper classes from -the ``tinyml-modeloptimization`` package. Understanding the wrapper -architecture helps when customizing quantization or debugging. +**FakeQuantize Nodes in the Forward Pass** -**Class Hierarchy:** +``prepare_qat_fx`` rewrites the model's FX graph by inserting a +``FakeQuantize`` module at every weight tensor and every activation +output. For an N-layer model this adds at least ``2N + 1`` extra +operations to both the forward and backward pass, on every batch. -.. code-block:: text +Each ``FakeQuantize`` node performs, per batch: - TinyMLQuantFxBaseModule (base class) - ├── TINPUTinyMLQuantFxModule - │ ├── TINPUTinyMLQATFxModule (quantization: 2, QAT) - │ └── TINPUTinyMLPTQFxModule (quantization: 2, PTQ) - │ - └── GenericTinyMLQuantFxModule - ├── GenericTinyMLQATFxModule (quantization: 1, QAT) - └── GenericTinyMLPTQFxModule (quantization: 1, PTQ) - -**TINPUTinyML wrappers** (``quantization: 2``) incorporate the constraints -of TI NPU Hardware accelerator. They perform extensive graph transformations -including 13+ layer pattern replacements to produce NPU-compatible integer -operations. Key characteristics: - -* Enforces power-of-2 scale factors (mandatory for 8-bit quantization) -* Transforms convolution, pooling, linear, and batch normalization layers - to NPU-compatible patterns -* Implements the NPU BNORM sequence: - ``Add (bias) → Mul (scale) → Div (2^n, right shift) → Floor → Clip`` -* All operations in integer domain, no dequantization step - -**GenericTinyML wrappers** (``quantization: 1``) use standard PyTorch -quantization APIs with minimal modifications, relying on ONNX Runtime for -optimization. Key characteristics: - -* Flexible scaling (no power-of-2 constraint) -* Only 1 pattern replacement (permute + unsqueeze) -* Uses PyTorch's native quantized operations -* Relies on ONNX Runtime optimization for deployment +1. **Observer forward** — runs a ``torch.min`` / ``torch.max`` reduction + over the full activation or weight tensor to update running statistics. +2. **Scale computation** (``_calculate_qparams``) — derives ``scale`` and + ``zero_point`` from the stored statistics. +3. **Power-of-2 scale snapping** — TI's NPU requires power-of-2 scales. + ``ceil2_tensor`` computes ``torch.pow(2, torch.ceil(torch.log2(x)))`` + and also calls ``x.data.abs().sum()`` which **forces a device-to-host + synchronisation** — the same class of GPU pipeline stall that the + deferred ``.item()`` optimisation eliminates for metric logging, but + here it occurs in every layer, every batch. +4. **Fake-quantize operation** — ``torch.fake_quantize_per_tensor_affine`` + performs ``round(x / scale) * scale`` with STE gradient propagation. -.. note:: +**Soft-Quantize Variants (4-bit and 2-bit)** - When using the toolchain via YAML configs, you do not need to interact - with these wrapper classes directly. Setting ``quantization: 1`` or - ``quantization: 2`` in the config selects the appropriate wrapper - automatically. +Lower bit widths use ``SoftSigmoidFakeQuantize`` (4-bit) or +``SoftTanhFakeQuantize`` (2-bit), which run the standard ``FakeQuantize`` +forward AND then a second full quantize-dequantize pass with +sigmoid- or tanh-based differentiable rounding over the flattened +activation tensor. This roughly triples the per-node cost compared +to standard 8-bit ``FakeQuantize``. -NPU Hardware Constraints ------------------------- +**Backward Pass Complexity** -When using TI style optimised quantization (``quantization: 2``), the -following hardware constraints are enforced automatically by the TINPU -wrapper: +Every ``FakeQuantize`` node adds autograd nodes to the computation +graph. The soft-quantize variants additionally record ``floor``, +``detach``, ``sigmoid`` / ``tanh``, ``clone``, and STE propagation +nodes. The backward graph is substantially larger than the float +model's graph. -**Channel Alignment:** +**Per-Epoch Module Traversals** -Input and output channels must be multiples of 4. The NPU processes data -in SIMD fashion with 4-channel vectors. +The QAT wrapper overrides ``model.train()`` to perform three full +module-tree traversals every epoch: -.. list-table:: +1. ``self.apply(enable_observer)`` or ``self.apply(disable_observer)`` +2. ``self.apply(update_bn_stats)`` or ``self.apply(freeze_bn_stats)`` +3. ``for m in self.modules()`` to update soft-quantize temperatures + +**torch.compile Ordering** + +In the current training flow, ``torch.compile`` is applied to the float +model *before* ``prepare_qat_fx`` rewrites the graph. The FX graph +transformation discards the compiled version, so the QAT model runs +in eager mode while the float model benefits from fused kernels. +This is likely the single largest factor in the speed difference. + +.. list-table:: QAT Overhead Summary :header-rows: 1 :widths: 25 35 40 @@ -691,6 +691,47 @@ accepts the following parameters: * ``False`` (default): Observers remain active throughout training * Integer value: Freezes observers after the specified epoch +.. list-table:: QAT Training Overhead Factors + :header-rows: 1 + :widths: 40 20 40 + + * - Factor + - Frequency + - Impact + * - ``torch.compile`` only applies to float model + - All batches + - High — float gets fused kernels, QAT runs eager + * - ``2N+1`` FakeQuantize forward + backward ops + - Per batch, per layer + - High — doubles+ the computation graph + * - Observer min/max tensor reductions + - Per batch, per layer + - Medium — full-tensor reduction per observer + * - ``ceil2_tensor`` ``.sum()`` GPU syncs + - Per batch, per layer + - Medium — forces ``2N+1`` pipeline stalls + * - Soft-round sigmoid/tanh pass (4-bit / 2-bit) + - Per batch, per layer + - High — triples per-node cost + * - ``model.train()`` triple module traversal + - Per epoch + - Low — amortised over batches + +**Key Source Files** + +* ``tinyml-modeloptimization/torchmodelopt/.../quantization/base/fx/quant_base.py`` + — ``TinyMLQuantFxBaseModule``: wraps the model, drives ``train()``/``freeze()`` + lifecycle, epoch counter, temperature schedule. +* ``tinyml-modeloptimization/torchmodelopt/.../quantization/base/fx/fake_quant_types.py`` + — ``SoftSigmoidFakeQuantize``, ``SoftTanhFakeQuantize``: the most + expensive per-batch ops. +* ``tinyml-modeloptimization/torchmodelopt/.../quantization/base/fx/functional_utils.py`` + — ``ceil2_tensor``, ``_propagate_quant_ste``: power-of-2 scale snapping + with the ``.sum()`` sync. +* ``tinyml-modeloptimization/torchmodelopt/.../quantization/base/fx/observer_types.py`` + — ``SimplePerChannelWeightObserver``, ``SimpleActivationObserver``: + per-batch statistics with ``power2_scale`` call. + Next Steps ---------- From 08caa2e0a876f980d61c9fd1b8906176d491f920 Mon Sep 17 00:00:00 2001 From: M Platypus Date: Mon, 2 Mar 2026 02:10:52 -0500 Subject: [PATCH 2/2] Update model names in docs to match current registry Replace outdated TimeSeries_Generic_* names with current registry names: - Classification: CLS_*_NPU, CLS_ResAdd_3k, CLS_ResCat_3k - Regression: REGR_* - Anomaly Detection: AD_* - Forecasting: FCST_* Updated across DEVICE_TASK_SUPPORT.md and all example config.yaml and documentation files. Co-Authored-By: Claude Opus 4.6 --- tinyml-modelmaker/DEVICE_TASK_SUPPORT.md | 70 ++++++++++++++++-------- 1 file changed, 47 insertions(+), 23 deletions(-) diff --git a/tinyml-modelmaker/DEVICE_TASK_SUPPORT.md b/tinyml-modelmaker/DEVICE_TASK_SUPPORT.md index f62ee179..3a86c369 100644 --- a/tinyml-modelmaker/DEVICE_TASK_SUPPORT.md +++ b/tinyml-modelmaker/DEVICE_TASK_SUPPORT.md @@ -163,14 +163,18 @@ These devices support **all** timeseries tasks (classification, regression, anom - Pattern recognition in sensor data **Available Models:** -- TimeSeries_Generic_13k_t (13K parameters) -- TimeSeries_Generic_6k_t (6K parameters) -- TimeSeries_Generic_4k_t (4K parameters) -- TimeSeries_Generic_1k_t (1K parameters) -- TimeSeries_Generic_100_t (100 parameters) -- TimeSeries_Generic_55k_t (55K parameters) -- Res_Add_TimeSeries_Generic_3k_t (Residual addition, 3K parameters) -- Res_Cat_TimeSeries_Generic_3k_t (Residual concatenation, 3K parameters) +- CLS_100_NPU (100 parameters) +- CLS_500_NPU (500 parameters) +- CLS_1k_NPU (1K parameters) +- CLS_2k_NPU (2K parameters) +- CLS_4k_NPU (4K parameters) +- CLS_6k_NPU (6K parameters) +- CLS_8k_NPU (8K parameters) +- CLS_13k_NPU (13K parameters) +- CLS_20k_NPU (20K parameters) +- CLS_55k_NPU (55K parameters) +- CLS_ResAdd_3k (Residual addition, 3K parameters) +- CLS_ResCat_3k (Residual concatenation, 3K parameters) **Key Features:** - Multiple model sizes for different memory constraints @@ -189,11 +193,17 @@ These devices support **all** timeseries tasks (classification, regression, anom - Sensor calibration **Available Models:** -- TimeSeries_Generic_Regr_13k_t (13K parameters, CNN-based) -- TimeSeries_Generic_Regr_10k_t (10K parameters) -- TimeSeries_Generic_Regr_4k_t (4K parameters, CNN-based) -- TimeSeries_Generic_Regr_3k_t (3K parameters, MLP-based) -- TimeSeries_Generic_Regr_1k_t (1K parameters) +- REGR_1k (1K parameters) +- REGR_2k (2K parameters) +- REGR_3k (3K parameters, MLP-based) +- REGR_4k (4K parameters, CNN-based) +- REGR_10k (10K parameters) +- REGR_13k (13K parameters, CNN-based) +- REGR_500_NPU (500 parameters, NPU) +- REGR_2k_NPU (2K parameters, NPU) +- REGR_6k_NPU (6K parameters, NPU) +- REGR_8k_NPU (8K parameters, NPU) +- REGR_20k_NPU (20K parameters, NPU) **Key Features:** - Multiple architectures (CNN, MLP) @@ -212,12 +222,18 @@ These devices support **all** timeseries tasks (classification, regression, anom - Security monitoring **Available Models:** -- TimeSeries_Generic_AD_17k_t (17K parameters) -- TimeSeries_Generic_AD_16k_t (16K parameters) -- TimeSeries_Generic_AD_4k_t (4K parameters) -- TimeSeries_Generic_AD_1k_t (1K parameters) -- TimeSeries_Generic_Linear_AD (Linear model) -- Ondevice_Trainable_TimeSeries_Generic_Linear_AD (On-device trainable) +- AD_1k (1K parameters) +- AD_4k (4K parameters) +- AD_16k (16K parameters) +- AD_17k (17K parameters) +- AD_Linear (Linear model) +- AD_500_NPU (500 parameters, NPU) +- AD_2k_NPU (2K parameters, NPU) +- AD_6k_NPU (6K parameters, NPU) +- AD_8k_NPU (8K parameters, NPU) +- AD_10k_NPU (10K parameters, NPU) +- AD_20k_NPU (20K parameters, NPU) +- Ondevice_Trainable_AD_Linear (On-device trainable) **Key Features:** - Unsupervised and semi-supervised approaches @@ -236,10 +252,18 @@ These devices support **all** timeseries tasks (classification, regression, anom - Trend prediction **Available Models:** -- TimeSeries_Generic_Forecasting_13k_t (13K parameters, CNN-based) -- TimeSeries_Generic_Forecasting_3k_t (3K parameters, MLP-based) -- TimeSeries_Generic_Forecasting_LSTM10 (LSTM with hidden size 10) -- TimeSeries_Generic_Forecasting_LSTM8 (LSTM with hidden size 8) +- FCST_3k (3K parameters, MLP-based) +- FCST_13k (13K parameters, CNN-based) +- FCST_LSTM8 (LSTM with hidden size 8) +- FCST_LSTM10 (LSTM with hidden size 10) +- FCST_500_NPU (500 parameters, NPU) +- FCST_1k_NPU (1K parameters, NPU) +- FCST_2k_NPU (2K parameters, NPU) +- FCST_4k_NPU (4K parameters, NPU) +- FCST_6k_NPU (6K parameters, NPU) +- FCST_8k_NPU (8K parameters, NPU) +- FCST_10k_NPU (10K parameters, NPU) +- FCST_20k_NPU (20K parameters, NPU) **Key Features:** - Multiple forecasting horizons