Quantizing Transformer Architecture Below 8-bit (post training quantization)

username: pkadambi
cooked: <p>I’m trying to quantize BERT to 4 bits or mixed precision, and I don’t see available methods to to quantization aware training on BERT for any precision other than torch.uint8. This is given in the dynamic quantization tutorial.<br>
I want to use both post training quantization and dynamic quantization for lower than 8 bits.</p>
<p>Will I have to rewrite the modeling_bert.py (transformers/modeling_bert.py) layers with fake quantization added? How can lower than 8bit precision and mixed precision be implemented on BERT?</p>

username: tom
cooked: <p>The difficulty there is PyTorch inherently assumes that things are at least 1 byte when doing things with memory.<br>
I’d probably convert to <a href="https://tvm.ai/" rel="nofollow noopener">TVM</a> and see what can be done there.<br>
(QAT with fake quantization probably could work for 4 bits, too.)</p>

username: pkadambi
cooked: <p>It’s not an issue even if the weights are stored as FP32 values in memory.<br>
I’m trying to evaluate post training quantization or fine tune the model with quantization aware training, but do this all under under fake quantization to any bit width of my choosing.</p>

username: tom
cooked: <p>While I don’t think it works out of the box, you could try to adapt the observers and fake quant layers to be more flexible. For example, there are some obvious 8 bit hard coded values here:<br>
</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/a414bd69de8d01af44751bfe327703ec997dafd9/torch/quantization/observer.py#L146" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/a414bd69de8d01af44751bfe327703ec997dafd9/torch/quantization/observer.py#L146" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/a414bd69de8d01af44751bfe327703ec997dafd9/torch/quantization/observer.py#L146</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="136" style="counter-reset: li-counter 135 ;">
<li>    Learned Step Size Quantization: https://openreview.net/pdf?id=rkgO66VKDS</li>
<li>    Trained Quantization Thresholds: https://arxiv.org/pdf/1903.08066.pdf</li>
<li>    """</li>
<li>    # The variable names are prefixed with "initial" because their values (qmin and qmax) might be adjusted</li>
<li>    # based on whether quantization range is reduced and the datatype (signed/unsigned) used by the observer.</li>
<li>    initial_qmin, initial_qmax = initial_dynamic_qrange</li>
<li>    assert initial_qmin &lt;= 0 &lt;= initial_qmax, "Dynamic quantization range must include 0."</li>
<li>    assert initial_qmin &lt; initial_qmax, "qmin must be strictly less than qmax for dynamic quantization range."</li>
<li>
</li><li>@torch.jit.export</li>
<li class="selected">def _calculate_qmin_qmax(self):</li>
<li>    # type: () -&gt; Tuple[int, int]</li>
<li>    r"""Calculates actual qmin and qmax based on the quantization range,</li>
<li>    observer datatype and if range is reduced.</li>
<li>    """</li>
<li>    if self.is_dynamic_qrange:</li>
<li>        # This initialization here is to be resolve TorchScript compilation issues and allow</li>
<li>        # using of refinement to decouple initial_qmin and initial_qmax from quantization range.</li>
<li>        # The actual values of initial_qmin and initial_qmax will be reset below.</li>
<li>        initial_qmin, initial_qmax = 0, 255</li>
<li>        # The following assignment of initial_qrange to a local variable and the if check refine the</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: jerryzh168
cooked: <aside class="quote no-group" data-username="pkadambi" data-post="3" data-topic="91686">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/eb9ed0/48.png" class="avatar"> pkadambi:</div>
<blockquote>
<p>I’m trying to evaluate post training quantization or fine tune the model with quantization aware training, but do this all under under fake quantization to any bit width of my choosing.</p>
</blockquote>
</aside>
<p>we do have the support for lower bits in <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/observer.py#L185" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/torch/quantization/observer.py#L185</a> now, one of our interns just added this recently.</p>