Quantization and nn.Parameters

yannbane: <p>My model defines the following field:</p>
<pre><code class="lang-python">self.relative_position_bias_table = nn.Parameter(
    torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1),
                num_heads))  # 2*Wh-1 * 2*Ww-1, nH
</code></pre>
<p>And the following method:</p>
<pre><code class="lang-python">def init_weights(self):
    super(WindowMSA, self).init_weights()

    trunc_normal_(self.relative_position_bias_table, std=0.02)
</code></pre>
<p>I believe this method is only called at the start of training. I’m working with a pre-trained model, and yet <code>self.relative_position_bias_table</code> is not all zeros. I concluded that <code>nn.Parameter</code> objects are <em>also</em> saved in the model file, and loaded similar to weights.</p>
<p>For some reason, however, quantizing the model using <code>torch.quantization.convert</code> skips this field. It’s left as a normal tensor, instead of being converted to int8 like other parameters of the model. I cannot combine it in operations with other quantized tensors.</p>
<p><strong>What is the standard way of telling PyTorch to treat <code>nn.Parameter</code> fields as normal parameters and quantize them to int8?</strong></p>
<p>One option would be to explicitly quantize this tensor every time. I find that wasteful. Caching the quantized version on the other hand sounds tacky and seems like something that PyTorch should be doing for me anyway.</p>

Vasiliy_Kuznetsov: <p>Hi <a class="mention" href="/u/yannbane">@yannbane</a> , PyTorch quantization works on operations (conv, linear, etc) and weights used by those operations, it does not know how to quantize arbitrary parameters.</p>
<p>For the case you describe, the right way to quantize <code>self.relative_position_bias_table</code> would depend on how this parameter is used.  If you share which operations use this parameter, happy to try to help more based on that context.</p>

yannbane: <p>Thanks for the response. It’s just a tensor that gets summed with another tensor (<code>attn</code>). Here’s what I’m currently doing:</p>
<pre><code class="lang-python">        if self.use_qtable:
            relative_position_bias_table = self.qrelative_position_bias_table 
        else:
            relative_position_bias_table = self.relative_position_bias_table

        relative_position_bias = relative_position_bias_table[
            self.relative_position_index.view(-1)].view(
                self.window_size[0] * self.window_size[1],
                self.window_size[0] * self.window_size[1],
                -1)  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
        
        attn = self.f_add.add(attn, relative_position_bias.unsqueeze(0))
</code></pre>
<p>The model has a field called <code>use_qtable</code>, which starts in <code>False</code>. After quantization statistics are collected, I call a method once to quantize <code>self.relative_position_bias_table</code> and store it in <code>self.qrelative_position_bias_table</code>. This method also sets <code>use_qtable</code> to <code>True</code>, and thereon the model uses the quantized version.</p>
<p>This is a hacky solution. <code>self.relative_position_bias_table</code> is stored in the model file as an <code>nn.Parameter</code> and conceptually it’s almost the same as a model weight tensor-- except that it was auto-generated.</p>