Variable-bit (sub 8-bits) quantization for custom hardware deployment with power-of-two (pot) scales

d0uwe: <p>Hey everyone! I am looking for a way to perform Quantization-Aware Training (QAT) using PyTorch.</p>
<p>My usecase concerns deploying trained PyTorch models on custom hardware (silicon) and so I have a few requirements:</p>
<ul>
<li>Needs to support <code>nn.Conv1d</code> (as this is part of the network that I want to deploy)</li>
<li>Needs to support some form of batch-norm folding</li>
<li>Needs to have power-of-two scales (as this avoids integer divisions in hardware)</li>
<li>Preferably does not require me to redefine my model with quantized modules</li>
<li>Supports up-to-date Python and PyTorch versions</li>
</ul>
<p>In my search, I checked out all possible quantization-aware-training frameworks I could find and made a list of them (see below). In brackets () indicates the last commit date and in square brackets [] it shows whether I tried it or not.</p>
<ul>
<li>QONNX: <a href="https://github.com/fastmachinelearning/qonnx" class="inline-onebox" rel="noopener nofollow ugc">GitHub - fastmachinelearning/qonnx</a> (3 weeks ago)
<ul>
<li>[-] ONNX-only</li>
</ul>
</li>
<li>QKeras: <a href="https://github.com/google/qkeras" class="inline-onebox" rel="noopener nofollow ugc">GitHub - google/qkeras: QKeras: a quantization deep learning library for Tensorflow Keras</a> (yesterday)
<ul>
<li>[+] Power of two scaling</li>
<li>[-] Need to rewrite network</li>
<li>[-] Keras-only</li>
</ul>
</li>
<li>TorchQuant: <a href="https://github.com/camlsys/torchquant" class="inline-onebox" rel="noopener nofollow ugc">GitHub - camlsys/torchquant: A Hackable Quantization Library for PyTorch</a> (29 March 2021)
<ul>
<li>[-] Doesnt support Conv1d</li>
</ul>
</li>
<li>MQBench: <a href="https://github.com/ModelTC/MQBench" class="inline-onebox" rel="noopener nofollow ugc">GitHub - ModelTC/MQBench: Model Quantization Benchmark</a> (14 Feb 2023) [tried]
<ul>
<li>[-] Doesnt work with PyTorch 1.13.1</li>
</ul>
</li>
<li>NNCF: <a href="https://github.com/openvinotoolkit/nncf" class="inline-onebox" rel="noopener nofollow ugc">GitHub - openvinotoolkit/nncf: Neural Network Compression Framework for enhanced OpenVINO™ inference</a> (2 days ago) [tried]
<ul>
<li>[-] Doesnt support power of-two scaling</li>
<li>[+] Amazing Github support</li>
<li>[-] Hard to export quantized model parameters</li>
</ul>
</li>
<li>QPyTorch: <a href="https://github.com/Tiiiger/QPyTorch" class="inline-onebox" rel="noopener nofollow ugc">GitHub - Tiiiger/QPyTorch: Low Precision Arithmetic Simulation in PyTorch</a> (14 Jan 2022)
<ul>
<li>[-] Not maintained</li>
</ul>
</li>
<li>AI Model Efficiency Toolkit (AIMET): <a href="https://github.com/quic/aimet" class="inline-onebox" rel="noopener nofollow ugc">GitHub - quic/aimet: AIMET is a library that provides advanced quantization and compression techniques for trained neural network models.</a> (19 hours ago) [tried]
<ul>
<li>[-] Only supports Python 3.8</li>
<li>[-] Really bad support via Github</li>
<li>[+] Very large feature set</li>
<li>[+] Tutorial videos</li>
<li>[+/-] Supports PyTorch 1.13 since latest commit, but no pre-built package available yet</li>
<li>[-] Building is on Python 3.10 is not trivial!</li>
</ul>
</li>
<li>Brevitas: <a href="https://github.com/Xilinx/brevitas" class="inline-onebox" rel="noopener nofollow ugc">GitHub - Xilinx/brevitas: Brevitas: quantization-aware training in PyTorch</a> (10 Jan 2023)
<ul>
<li>[-] Doesnt support batch norm</li>
<li>[-] Requires you to rewrite your entire network</li>
</ul>
</li>
<li>pytorch-quantization: <a href="https://docs.nvidia.com/deeplearning/tensorrt/pytorch-quantization-toolkit/docs/index.html" class="inline-onebox" rel="noopener nofollow ugc">pytorch-quantization’s documentation — pytorch-quantization master documentation</a> (very recent)
<ul>
<li>[-] Only int8</li>
</ul>
</li>
<li>NEMO: <a href="https://github.com/pulp-platform/nemo" class="inline-onebox" rel="noopener nofollow ugc">GitHub - pulp-platform/nemo: NEural Minimizer for pytOrch</a> (23 Feb 2022)
<ul>
<li>[-] Not maintained anymore</li>
</ul>
</li>
<li>model_optimization: <a href="https://github.com/sony/model_optimization#experimental-features" class="inline-onebox" rel="noopener nofollow ugc">GitHub - sony/model_optimization: Model Compression Toolkit (MCT) is an open source project for neural network model optimization under efficient, constrained hardware. This project provides researchers, developers, and engineers advanced quantization and compression tools for deploying state-of-the-art neural networks.</a> (1 hour ago) [paper: <a href="https://arxiv.org/pdf/2109.09113.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/2109.09113.pdf</a>, tried]
<ul>
<li>[+] Power-of-two quantization</li>
<li>[+] Simple to use</li>
<li>[+] Very active repository and members</li>
<li>[-] Not possible to get access to quantized weights?</li>
</ul>
</li>
<li>PyTorch quantization:
<ul>
<li>[-] Need to implement custom observer + fake quantizer</li>
</ul>
</li>
</ul>
<p>However, none of these options really work or have all the features that I need. Does anyone else have suggestions on what I can use/do?</p>

jcaip: <p>Hi <a class="mention" href="/u/d0uwe">@d0uwe</a>,</p>
<p>We have a tutorial for QAT here that you <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#quantization-aware-training" rel="noopener nofollow ugc">might</a> find helpful.</p>

d0uwe: <p>Thanks for the quick response! I have read this tutorial but it is unclear to me how I can quantize the weights lower than 8 bits (for example 5).</p>

111357: <p>As far as I know, PyTorch 2.0 does not support quantized weight lower than 8 bits natively. But you can emulate it numerically with a customized <a href="https://pytorch.org/docs/stable/generated/torch.ao.quantization.observer.MinMaxObserver.html#torch.ao.quantization.observer.MinMaxObserver" rel="noopener nofollow ugc">observer</a>.</p>
<p>For example, if you want to quantize weight to <code>int4</code>, you can try the following setting:</p>
<pre><code class="lang-python">from torch.ao.quantization.observer import MinMaxObserver
custom_observer = MinMaxObserver(quant_min=-8, quant_max=7)

</code></pre>

d0uwe: <p>Thanks, that looks good! How would I then for example implement that that the scales of the quantizer are always a power of two? I.e a scale factor of 1.9 is rounded to 2 (2^1), a scale factor of 59 is rounded to 64 (2^6) etc.</p>

Vasiliy_Kuznetsov: <p>You could create your own observer class and override this function (<a href="https://github.com/pytorch/pytorch/blob/5cc2e4d7c939852f6de6f8497dc89d311e333dce/torch/ao/quantization/observer.py#L296" class="inline-onebox" rel="noopener nofollow ugc">pytorch/observer.py at 5cc2e4d7c939852f6de6f8497dc89d311e333dce · pytorch/pytorch · GitHub</a>) to calculate the scales with a power of two restriction.</p>

d0uwe: <p>So if I understand correctly, the scale and zero-point are currently not learned but calculated right?</p>
<p>For creating my own observer class, from which class would I inherit? From <code>UniformQuantizationObserverBase</code> (or any other observer) or from <code>ObserverBase</code>?</p>
<p>Then for a learned scale and zero-point, is the below the correct approach?</p>
<pre><code class="lang-python">def quantize_to_closest_power_of_two(x: torch.Tensor):
   ...

   return quantized_value

class MyObserverClass(...):
   def __init__(...):
      ...
      
      self.scale = nn.Parameter(torch.ones(required_shape))
      self.zero_point = nn.Parameter(torch.ones(required_shape))

   @torch.jit.export
   def calculate_qparams(self):
        return quantize_to_closest_power_of_two(self.scale), self.zero_point
</code></pre>
<p>Then a related question: what if I also want to only allow the weights to be powers of two, next to the scales being powers of two? Do I create a new class with <code>FakeQuantize</code> as base and then in the forward pass I clamp the weights to the closest powers of two?</p>

Tobias-Fischer: <p>Hi <a class="mention" href="/u/d0uwe">@d0uwe</a> <a class="mention" href="/u/vasiliy_kuznetsov">@Vasiliy_Kuznetsov</a> - did you ever figure this out? Just running into the same issue at the moment.</p>

jerryzh168: <p>you can create a custom Observer/FakeQuantize class that makes sure weights are close to power of two I think</p>

d0uwe: <p>Hi Tobias, no, I didn’t figure out how to properly use the PyTorch APIs for this. I finally decided to use <a href="https://github.com/Xilinx/brevitas" rel="noopener nofollow ugc">Brevitas </a> because it offers a lot of flexibility.</p>
<p>However, since PoT weight quantization was not supported, I built that myself <a href="https://github.com/V0XNIHILI/brevitas-utils/blob/main/brevitas_utils/custom_quantizers/Int8WeightPerTensorPowerOfTwo.py" rel="noopener nofollow ugc">here</a> as part of a mini-library to to use Brevitas to quantize neural networks for ASICs.</p>