Possible bug in fully quantized PyTorch implementation

username: mdoyle
cooked: <p>I would like to execute a PyTorch model trained with quantization-aware training (QAT) as a fully quantized model. However, the output of my fully quantized and fake quantized models do not match.</p>
<p>To highlight the problem, I defined a very simple experiment consisting of quantizing only a single fused Conv-ReLU operation with hard-coded weights and quantization parameters. What I noticed is that;</p>
<ul>
<li>Torch produces the expected result for a model prepared with <code>torch.quantization.prepare_qat</code> (fake quantized)</li>
<li>Torch produces unexpected results when the previously prepared model is converted to a fully quantized model with <code>torch.quantization.convert</code> (real quantized)</li>
</ul>
<p>I am wary that I might have an error in my implementation, so I provide a detailed example below and a repo of the example <a href="https://github.com/mylesDoyle/DebugTorchQuantization" rel="noopener nofollow ugc">here</a>;</p>
<h2>
<a name="experiment-1" class="anchor" href="#experiment-1"></a>Experiment</h2>
<p>Compare fake and real quantized model outputs;</p>
<ul>
<li>inferring with a normal QAT model (fake quantized) - produces expected results</li>
<li>inferring with a prepared and converted model to int8 (quantized) - produces unexpected results</li>
</ul>
<p>To highlight the issue, I have set up a simple toy example as follows;</p>
<h3>
<a name="model-2" class="anchor" href="#model-2"></a>Model</h3>
<p>A simple Conv-ReLU fused model is defined with</p>
<ul>
<li>bias set to zero</li>
<li>conv weights set to <code>k*I</code> where <code>k</code> is some floating point scalar multiplier and <code>I</code> represents an identity matrix of the correct shape for the conv layer weights</li>
<li>A quantization stub which quantizes the <code>fp32</code> inputs to <code>quint8</code>
</li>
<li>A dequantization stub which dequantizes the <code>quint8</code> outputs to <code>fp32</code> - note, this stub gets set to the identity for the fully int8 quantized model</li>
</ul>
<pre><code class="lang-auto">class FooConv1x1(nn.Module):
    
    def __init__(self, set_qconfig):
        super().__init__()

        self.conv = nn.Conv2d(3, 3, 1, 1)   # 1x1 Conv kernel
        self.act = nn.ReLU()

        self.quant = torch.quantization.QuantStub(CustomQConfigs.get_default_qconfig())
        self.dequant = torch.quantization.DeQuantStub()

        self.modules_to_fuse = [['conv', 'act']]

        if set_qconfig:
            self.set_qconfig()
    
    def forward(self, x):
        x = self.quant(x)
        output_quant = self.act(self.conv(x))
        return self.dequant(output_quant) 

    def fuse(self):
        torch.ao.quantization.fuse_modules(self, self.modules_to_fuse, inplace=True)
        return self

    def set_qconfig(self):
        self.qconfig = CustomQConfigs.get_default_qconfig()
        return self
    
    def set_weights(self, multiplier):
        # Set bias to zero and conv weights to k*Identity
        self.conv.bias = torch.nn.Parameter(torch.zeros_like(self.conv.bias))
        self.conv.weight = torch.nn.Parameter(multiplier * torch.eye(3).reshape(self.conv.weight.shape))
</code></pre>
<h3>
<a name="pytorch-qconfig-3" class="anchor" href="#pytorch-qconfig-3"></a>PyTorch QConfig</h3>
<p>The quantization config for the model was defined with;</p>
<ul>
<li>Per tensor affine quantization everywhere except for the Conv layer’s weights which are per tensor symmetric</li>
</ul>
<pre><code class="lang-auto">class CustomQConfigs:
    @staticmethod
    def get_default_qconfig():
        return torch.quantization.QConfig(activation=torch.quantization.FusedMovingAvgObsFakeQuantize.with_args(
                observer=torch.quantization.MovingAverageMinMaxObserver,
                quant_min=0,
                quant_max=255,
                reduce_range=False),
            weight=torch.quantization.FusedMovingAvgObsFakeQuantize.with_args(
                observer=torch.quantization.MovingAverageMinMaxObserver,
                quant_min=-128,
                quant_max=127,
                dtype=torch.qint8,
                qscheme=torch.per_tensor_symmetric))
</code></pre>
<h3>
<a name="inputs-4" class="anchor" href="#inputs-4"></a>Inputs</h3>
<p>Inputs are provided to the model in single precision floating point units in all cases. To highlight the issue, I consider passing a range of input values between 0 and 255 across an input image of size <code>[1,3,256,256]</code> and scaling the values to between 0 and 1 by dividing by 255.</p>
<h3>
<a name="setup-and-execution-5" class="anchor" href="#setup-and-execution-5"></a>Setup and execution</h3>
<p>The example can be executed using the following snippets or by cloning the repo and following the instructions <a href="https://github.com/mylesDoyle/DebugTorchQuantization" rel="noopener nofollow ugc">here</a>.</p>
<p>Note, I use <code>torch.quantization.prepare_qat</code> instead of <code>torch.quantization.prepare</code> so that observers get added the Conv layer so that I can hard code the quantization parameters to their analytically calculated values.</p>
<pre><code class="lang-auto">import torch

backend = "fbgemm"
# backend = "qnnpack"
torch.backends.quantized.engine = backend

torch.manual_seed(0)

# Hard code relevant quantization parameters
def set_qconfig_params(model_prepared, k):
    # Conv weight
    model_prepared.conv.weight_fake_quant.scale = torch.Tensor([2.0*k/255.0])   # Symmetric, hence multiply by 2
    model_prepared.conv.weight_fake_quant.activation_post_process.min_val = torch.tensor(0.0)
    model_prepared.conv.weight_fake_quant.activation_post_process.max_val = torch.tensor(k)

    # Requantization
    model_prepared.conv.activation_post_process.scale = torch.Tensor([k/255.0])
    model_prepared.conv.activation_post_process.min_val = torch.tensor(0.0)
    model_prepared.conv.activation_post_process.max_val = torch.tensor(k)
    model_prepared.conv.activation_post_process.activation_post_process.min_val = torch.tensor(0.0)
    model_prepared.conv.activation_post_process.activation_post_process.max_val = torch.tensor(k)

    # Input quant stub
    model_prepared.quant.activation_post_process.scale = torch.Tensor([1.0/255.0])
    model_prepared.quant.activation_post_process.activation_post_process.min_val = torch.tensor(0.0)
    model_prepared.quant.activation_post_process.activation_post_process.max_val = torch.tensor(1.0)

if __name__ == "__main__":
   
    input_fp32 = torch.arange(0,256).repeat(1,3,256,1)/255.0   # 0 to 255 repeated across rows, then normalized to [0,1]

    model = FooConv1x1(set_qconfig=True)    # Prepare model with QConfig defined
    k = 1.0 # Set Conv layer multiplier
    model.set_weights(k)    # Set bias to zero and conv weights to k*Identity
    model.fuse()    # fuse conv and ReLU 

    model_prepared = torch.quantization.prepare_qat(model).train()  # prepare_qat required to set weight qparams
    model_prepared.eval()

    model_prepared.apply(torch.quantization.disable_observer).eval()    # Disable QConfig Observers

    set_qconfig_params(model_prepared, k)   # Set quantization parameters to theoretical values

    expected_output_fp32 = model_prepared(input_fp32)
    expected_output_quint8 = (expected_output_fp32*(k*255)).to(torch.uint8)
    
    model_prepared.dequant = torch.nn.Identity()    # Disable the output dequant stub

    # Convert model so that it runs as fully quantized model
    model_quant = torch.quantization.convert(model_prepared, inplace=False)

    output_quint8_fp32 = model_quant(input_fp32) # fp32 outputs with scale and shift parameterising it to quint8

    error = torch.abs(expected_output_fp32 - output_quint8_fp32.dequantize())
    error_mean = torch.mean(error)
    error_max = torch.max(error)
    first_nonzero_index = error.nonzero()[0].tolist()

    print(f"{error_mean=}")
    print(f"{error_max=}")
    print(f"First nonzero: index: ({first_nonzero_index}")
    print(f"\tvalue fp32: {error[*first_nonzero_index]}")
    print(f"\tvalue expected quint8: {expected_output_quint8[*first_nonzero_index]}")
    print(f"\tvalue outputed quint8: {output_quint8_fp32.int_repr()[*first_nonzero_index]}")

    # import ipdb; ipdb.set_trace()
</code></pre>
<h2>
<a name="dependencies-6" class="anchor" href="#dependencies-6"></a>Dependencies</h2>
<p>The example in this repo was tested using</p>
<ul>
<li>Python 3.11.0</li>
<li>Python packages installed with pip which are listed in the <code>requirements.txt</code> in the <a href="https://github.com/mylesDoyle/DebugTorchQuantization/blob/main/requirements.txt" rel="noopener nofollow ugc">repo provided</a>
</li>
</ul>
<p>Operating system:</p>
<ul>
<li>Ubuntu 22.04.1 LTS</li>
</ul>
<p>Hardware architecture:</p>
<ul>
<li>x86_64 - Intel</li>
</ul>
<h2>
<a name="observations-7" class="anchor" href="#observations-7"></a>Observations</h2>
<p>For simplicity, I compare the quantized outputs, but the same can be observed for the dequantized outputs.</p>
<ul>
<li>For the output of a model (fake or real quantized), I expect each row to be identical across all rows and channels. This was observed in all cases indicating determinism within a model’s execution</li>
<li>The quantized outputs of the (fake quantized) model prepared with <code>torch.quantization.prepare_qat</code> were as expected;
<ul>
<li>values ranging from 0 to 255, indicating a unique bin for each of the outputs which get dequantized into the expected output value</li>
<li>a summary of the first row of the first channel depicts the beginning, middle and end of that row</li>
</ul>
</li>
</ul>
<pre><code class="lang-auto">tensor([  0,   1, ..., 126, 127, 128, 129, 130, 131, ..., 253, 254, 255], dtype=torch.uint8)
</code></pre>
<ul>
<li>The quantized outputs of the (quantized) model converted with <code>torch.quantization.convert</code> were not quite as expected;
<ul>
<li>values ranging from 0 to 254, implying we are losing information somewhere within the quantized execution</li>
<li>a summary of the first row of the first channel depicts the beginning, middle and end of that row</li>
</ul>
</li>
</ul>
<pre><code class="lang-auto">tensor([  0,   1, ..., 126, 127, 127, 128, 129, 130, ..., 252, 253, 254], dtype=torch.uint8)
</code></pre>
<ul>
<li>Comparing the quantized outputs of the two models, we observe in the real quantized model;
<ul>
<li>a discrepancy can be seen as the value 127 is repeated twice, and all other values are shifted after that</li>
<li>this results in the repeated 127 value being the incorrect bin for its expected dequantized value, and all other values following this duplication have been shifted to incorrect bins as well</li>
<li>this behaviour is unexpected and results in non-determinism across both model’s executions</li>
<li>note, it is interesting that this discrepancy appears for <code>128=255/2+1</code> and all following values due to 128 being halway bin of the possible range of bins</li>
</ul>
</li>
</ul>
<h2>
<a name="conclusion-8" class="anchor" href="#conclusion-8"></a>Conclusion</h2>
<p>One of the main reasons for using quantization is to ensure determinism across different compute platforms, so the non-deterministic behaviour between a fake and real quantized model is extremely problematic, especially when it comes to deploying quantized models.</p>
<p>It is clear from this example that the real quantized model is not working as expected. This must either be due to an error I have in my implementation or a bug within PyTorch. Any help determining the issue would be really appreciated, and if you think it is a bug in PyTorch, I will happily create a bug report, but I wanted to check here first.</p>

username: andrewor
cooked: <p>Hi Myles, there is actually no guarantee that the numerics of fake quantized will 100% match those of the real quantized values. How far apart are the values? It’s possible that there are rounding errors in fake quantize, but I’m not sure if this is a bug per se. The purpose of QAT in general is to improve the accuracy of the final model. Is that something you observed in your experiments as well?</p>

username: mdoyle
cooked: <p>Hi <a class="mention" href="/u/andrewor">@andrewor</a>, thank you for your reply.</p>
<blockquote>
<p>How far apart are the values?</p>
</blockquote>
<p>The values can differ by 1, with all values after also being shifted by 1. But also this shift can occur several times, resulting in errors greater than 1. Overflow can also occur where the error will become significant.</p>
<blockquote>
<p>It’s possible that there are rounding errors in fake quantize, but I’m not sure if this is a bug per se. The purpose of QAT in general is to improve the accuracy of the final model.</p>
</blockquote>
<p>I think you are misunderstanding my post;</p>
<ul>
<li>I have a QAT model and I am porting it to a fully quantized model</li>
<li>The issue is that the fully quantized model gives the incorrect results and therefore also results in non-determinism when the same model (ie weights and quantization parameters) is executed on other compute platforms or using other frameworks for quantization
<ul>
<li>Specifically, I tested this by verifying my model ran and produced the correct expected outputs when run in <code>int8</code>/<code>uint8</code> when on Qualcomm’s Snapdrgaon processor’s DSP (tested for the 888 or 8 Gen1 chips).</li>
</ul>
</li>
<li>I am aware fake quantize does not guarantee determinism, and can include errors. In my Identity example above, it does not encounter apparent errors and it just happens that Fake Quantize produces the expected analytical results</li>
<li>The issue is that real quantization in PyTorch (with the compute occurring in <code>int8</code>/<code>uint8</code>, ie. using the <code>torch.quantization.convert</code> model) produces incorrect and non-deterministic results.
<ul>
<li>The results do not match the analytical solution.</li>
<li>Nor do they match the expected results which I produced running in <code>int8</code> when running in other frameworks</li>
</ul>
</li>
</ul>
<p>This implies that PyTorch’s implementation of real quantization is incorrect/bugged.</p>