Quantization fails for custom backend

username: stuhlbein
cooked: <p>Hi all,<br>
I have issues trying to create a fully quantized model for my own backend (which will ultimately be a hardware AI accelerator). For simplicity, I wanted to purely use qint8 for now, the details will differ later as they depend a lot on memory bandwidth for different layers on hardware etc.<br>
So, what I want to do now is creating a simple model and quantize it completely (that means weights, inputs, outputs, biases… everything). I tried all kinds of different setups for the QConfigs but I never managed to create a model that is completely quantized.</p>
<pre><code class="lang-auto">import torch
from torch.ao.quantization import (
    default_weight_observer,
    get_default_qconfig_mapping,
    get_default_qconfig,
    MinMaxObserver,
    QConfig,
    QConfigMapping,
)
from torch.ao.quantization.backend_config import (
    BackendConfig,
    BackendPatternConfig,
    DTypeConfig,
    DTypeWithConstraints,
    ObservationType,
)
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.ao.quantization.fx.custom_config import PrepareCustomConfig
from torch.ao.quantization.observer import HistogramObserver, PerChannelMinMaxObserver

class Model(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.linear1 = torch.nn.Linear(input_size, 16)
        self.relu1 = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(16, 16)
        self.relu2 = torch.nn.ReLU()
        self.linear3 = torch.nn.Linear(16, output_size)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu1(x)
        x = self.linear2(x)
        x = self.relu2(x)
        x = self.linear3(x)
        return x

# Instantiate simple example model
input_size = 10
output_size = 3
model = Model(input_size, output_size)
model.eval()

# Define backend configuration, all QInt8 to keep things simple for now, no fusing
linear_int8_dtype = DTypeConfig(
    input_dtype=torch.qint8,
    output_dtype=torch.qint8,
    weight_dtype=torch.qint8,
    bias_dtype=torch.qint8)

linear_config = BackendPatternConfig(torch.nn.Linear) \
    .set_observation_type(ObservationType.OUTPUT_SHARE_OBSERVER_WITH_INPUT) \
    .add_dtype_config(linear_int8_dtype) \
    .set_root_module(torch.nn.Linear) \
    .set_reference_quantized_module(torch.ao.nn.quantized.reference.Linear)

backend_config = BackendConfig("my_backend") \
    .set_backend_pattern_config(linear_config)

# Create global QConfig mapping, using qint8 for everything
qconfig = QConfig(
    activation=MinMaxObserver.with_args(dtype=torch.qint8), 
    weight=MinMaxObserver.with_args(dtype=torch.qint8))

qconfig_mapping = QConfigMapping().set_global(qconfig)

# Setup s.t. input and outputs are expected to be already quantized
prepare_custom_config = PrepareCustomConfig() 
prepare_custom_config.set_input_quantized_indexes(list(range(input_size)))
prepare_custom_config.set_output_quantized_indexes(list(range(output_size)))

# generate some example data
example_input_float = torch.rand(20, input_size, dtype=torch.float)
example_input_int = torch.quantize_per_tensor(
                        example_input_float, 1.0/255, 0, torch.qint8)

# Preparing, calibrating, converting...
prepared = prepare_fx(model, qconfig_mapping, example_input_int, prepare_custom_config, backend_config=backend_config)
prepared(example_input_int) # fails here 
converted = convert_fx(prepared, backend_config=backend_config)
</code></pre>
<p>In the second to last line I get a runtime error <code>RuntimeError: Creation of quantized tensor requires quantized dtype like torch.quint8</code>.<br>
When adding some debugging output, it looks like the Observers are using 32b floats but not qint8 for some reason not quite clear to me.<br>
When calling the model with floats it works but after the conversion I do not get QuantizedLinear layers, so to me it looks like my configuration is wrong at some place, so preparing fails.</p>
<p>Does somebody know what the issue is?</p>

username: jerryzh168
cooked: <p>you don’t need to pass in a quantized tensor to prepared model</p>
<p>you can probably do <code>prepared(example_input_float)</code> instead</p>

username: xiangyun
cooked: <p>I have another question, I hope you can help me. Why did the code above only quantify activation without quantifying weights? as shown below<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/7/77d6a13fbf1519fed16e2c25b073a3e85ad2581a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/77d6a13fbf1519fed16e2c25b073a3e85ad2581a" title="1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/7/77d6a13fbf1519fed16e2c25b073a3e85ad2581a_2_689x272.jpeg" alt="1" data-base62-sha1="h68F0kY1krRQHyatLRkdbldHd3k" width="689" height="272" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/7/77d6a13fbf1519fed16e2c25b073a3e85ad2581a_2_689x272.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/7/77d6a13fbf1519fed16e2c25b073a3e85ad2581a_2_1033x408.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/7/77d6a13fbf1519fed16e2c25b073a3e85ad2581a.jpeg 2x" data-dominant-color="EFEDEE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">1</span><span class="informations">1290×510 53 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
thanks!</p>

username: jerryzh168
cooked: <p>weight is probably also quantized, how do you know it’s not quantized?</p>