Custom LSTM PTSQ QconfigMapping

Ahmed_Louati: <p>Hello,<br>
I am working on quantizing LSTM using custom module quantization. This is the ObservedLSTM module:</p>
<pre><code class="lang-auto">class ObservedLSTM(torch.ao.nn.quantizable.LSTM):
    """
     the observed LSTM layer. This module needs to define a from_float function which defines
     how the observed module is created from the original fp32 module.
    """
    @classmethod
    def from_float(cls, float_lstm):
        assert isinstance(m.my_lstm, cls._FLOAT_MODULE)
        # uint16, [-16, 16)
        linear_output_obs_ctr = observer.FixedQParamsObserver.with_args(
            scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
        )
        # uint16, [0, 1)
        sigmoid_obs_ctr = observer.FixedQParamsObserver.with_args(
            scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
        )
        # uint16, [-1, 1)
        tanh_obs_ctr = observer.FixedQParamsObserver.with_args(
            scale=2.0 / 256.0, zero_point=128, dtype=torch.quint8
        )
        # int16, [-16, 16)
        cell_state_obs_ctr = observer.FixedQParamsObserver.with_args(
            scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
        )
        # uint8, [-1, 1)
        hidden_state_obs_ctr = observer.FixedQParamsObserver.with_args(
            scale=2**-7, zero_point=2**7, dtype=torch.quint8
        )
        example_inputs = (torch.rand(1, 50), (torch.rand(1, 50), torch.randn(1, 50)))
        return lstm_utils._get_lstm_with_individually_observed_parts(
            float_lstm=float_lstm,
            example_inputs=example_inputs,
            backend_config=my_backend_config,
            linear_output_obs_ctr=linear_output_obs_ctr,
            sigmoid_obs_ctr=sigmoid_obs_ctr,
            tanh_obs_ctr=tanh_obs_ctr,
            cell_state_obs_ctr=cell_state_obs_ctr,
            hidden_state_obs_ctr=hidden_state_obs_ctr,
        )
</code></pre>
<p>I made sure that qparams are aligned with _FIXED_QPARAMS_OP_TO_CONSTRAINTS as I am using <code>FixedQParamsObserver</code>.</p>
<p>I even add these configs:</p>
<pre><code class="lang-auto">conf_linear = qconfig.QConfig(
    activation=observer.FixedQParamsObserver.with_args(
        scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
    ),
    weight=observer.FixedQParamsObserver.with_args(
        scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
    ),
)

conf_sigmoid = qconfig.QConfig(
    activation=observer.FixedQParamsObserver.with_args(
        scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
    ),
    weight=observer.FixedQParamsObserver.with_args(
        scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8
    ),
)

conf_tanh = qconfig.QConfig(
    activation=observer.FixedQParamsObserver.with_args(
        scale=2.0 / 256.0, zero_point=128, dtype=torch.quint8
    ),
    weight=observer.FixedQParamsObserver.with_args(
        scale=2.0 / 256.0, zero_point=128, dtype=torch.quint8
    ),
</code></pre>
<p>I still got this warning:</p>
<pre><code class="lang-auto">UserWarning: QConfig must specify a FixedQParamsObserver or a FixedQParamsFakeQuantize for fixed qparams ops, ignoring QConfig(activation=functools.partial(&lt;class 'torch.ao.quantization.observer.HistogramObserver'&gt;, reduce_range=False){'factory_kwargs': &lt;function _add_module_to_qconfig_obs_ctr.&lt;locals&gt;.get_factory_kwargs_based_on_module_device at 0x7f928dc8da20&gt;}, weight=functools.partial(&lt;class 'torch.ao.quantization.observer.MinMaxObserver'&gt;, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric){'factory_kwargs': &lt;function _add_module_to_qconfig_obs_ctr.&lt;locals&gt;.get_factory_kwargs_based_on_module_device at 0x7f928dc8da20&gt;}).
Please use torch.ao.quantization.get_default_qconfig_mapping or torch.ao.quantization.get_default_qat_qconfig_mapping. Example:
    qconfig_mapping = get_default_qconfig_mapping("fbgemm")
    model = prepare_fx(model, qconfig_mapping, example_inputs)
  warnings.warn(("QConfig must specify a FixedQParamsObserver or a FixedQParamsFakeQuantize "
</code></pre>
<p>From what I know the <code>.set_object_type()</code> overrides <code>.set_global()</code> configs.<br>
I am doing sthg wrong?</p>

HDCharles: <p>can you give a fuller repro? at what stage you do get this error?</p>

Ahmed_Louati: <p>Thank you <a class="mention" href="/u/hdcharles">@HDCharles</a>  for your reply!<br>
I am working on custom module quantization for LSTM layers.<br>
From what I understood, LSTM gates must have FixedQParamsObserver with certain ranges fo qparams as mentionned here <a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/backend_config/_common_operator_config_utils.py#L67" rel="noopener nofollow ugc">pytorch/torch/ao/quantization/backend_config/_common_operator_config_utils.py at main · pytorch/pytorch (github.com)</a>.</p>
<p>Besides, I added <code>Qconfig</code> which is compatible with those observers:</p>
<pre><code class="lang-auto">qconfig_mapping = (
    QConfigMapping()
    .set_global(get_default_qconfig("qnnpack"))
    .set_object_type(nn.Linear, conf_linear)
    .set_object_type(nn.Sigmoid, conf_sigmoid)
    .set_object_type(nn.Tanh, conf_tanh)
)
</code></pre>
<p>Normally <code>.set_global()</code> configs are overridden by those of <code>.set_object_type()</code>, So normally, observers of “qnnpack” (HistogramObserver for activations &amp; MinMaxObserver for weights) should be ignored which is normally the case!</p>
<p>I don’t know why I am still getting this warning!<br>
I hope that I made it a bit clearer!</p>

HDCharles: <p>without a full repro this is hard to diagnose.</p>
<p>Can you explain how what you are doing differs from: <a href="https://github.com/pytorch/pytorch/blob/main/test/quantization/fx/test_quantize_fx.py#L4886" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/test/quantization/fx/test_quantize_fx.py#L4886</a></p>

Ahmed_Louati: <p>I reproduced the same test function in a script and still get the same warning (I just copy-paste it!)</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy
from copy import deepcopy

import torch.ao.quantization.observer as observer
from torch.ao.quantization.backend_config._common_operator_config_utils import (
    _FIXED_QPARAMS_OP_TO_CONSTRAINTS,
)
import torch.ao.quantization.qconfig as qconfig

from torch.ao.quantization.backend_config import qnnpack, DTypeConfig, get_qnnpack_backend_config

from torch.ao.quantization.fx.custom_config import (
    PrepareCustomConfig,
    ConvertCustomConfig,
    QConfigMapping,
)
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.ao.quantization import get_default_qconfig_mapping, get_default_qconfig
from torch.ao.quantization.fx import lstm_utils
import torch.ao.ns._numeric_suite_fx as ns
import torch.ao.ns.fx

class MyModel(torch.nn.Module):
    def __init__(self):
                super().__init__()
                self.my_lstm = torch.nn.LSTM(50, 50, 1)

    def forward(self, inputs: torch.Tensor, h0: torch.Tensor, c0: torch.Tensor):
                x = self.my_lstm(inputs, (h0, c0))
                return x

        # Construct a BackendConfig that supports qint32 for certain ops
        # TODO: build a BackendConfig from scratch instead of modifying an existing one
qint32_dtype_config = DTypeConfig(input_dtype=torch.qint32, output_dtype=torch.qint32)
my_backend_config = get_qnnpack_backend_config()
for config in my_backend_config.configs:
    if config.pattern in [torch.nn.Sigmoid, torch.nn.Tanh, torch.add, torch.mul]:
                config.add_dtype_config(qint32_dtype_config)

class UserObservedLSTM(torch.ao.nn.quantizable.LSTM):
    """
    Example of user provided LSTM implementation that assigns fixed qparams
    to the inner ops.
    """
    @classmethod
    def from_float(cls, float_lstm):
        assert isinstance(float_lstm, cls._FLOAT_MODULE)
                # uint16, [-16, 16)
        linear_output_obs_ctr = observer.FixedQParamsObserver.with_args(scale=2 ** -11, zero_point=2 ** 15, dtype=torch.qint32)
                # uint16, [0, 1)
        sigmoid_obs_ctr = observer.FixedQParamsObserver.with_args(scale=2 ** -16, zero_point=0, dtype=torch.qint32)
                # uint16, [-1, 1)
        tanh_obs_ctr = observer.FixedQParamsObserver.with_args(scale=2 ** -15, zero_point=2 ** 15, dtype=torch.qint32)
                # int16, [-16, 16)
        cell_state_obs_ctr = observer.FixedQParamsObserver.with_args(scale=2 ** -11, zero_point=0, dtype=torch.qint32)
                # uint8, [-1, 1)
        hidden_state_obs_ctr = observer.FixedQParamsObserver.with_args(scale=2 ** -7, zero_point=2 ** 7, dtype=torch.quint8)
        example_inputs = (torch.rand(5, 3, 50), (torch.rand(1, 3, 50), torch.randn(1, 3, 50)))
        return torch.ao.quantization.fx.lstm_utils._get_lstm_with_individually_observed_parts(
                    float_lstm=float_lstm,
                    example_inputs=example_inputs,
                    backend_config=my_backend_config,
                    linear_output_obs_ctr=linear_output_obs_ctr,
                    sigmoid_obs_ctr=sigmoid_obs_ctr,
                    tanh_obs_ctr=tanh_obs_ctr,
                    cell_state_obs_ctr=cell_state_obs_ctr,
                    hidden_state_obs_ctr=hidden_state_obs_ctr,
                )
class UserQuantizedLSTM(torch.ao.nn.quantized.LSTM):
    """
    Example of user provided LSTM implementation that produces a reference
            quantized module from a `UserObservedLSTM`.
    """
    @classmethod
    def from_observed(cls, observed_lstm):
        assert isinstance(observed_lstm, cls._FLOAT_MODULE)
        return torch.ao.quantization.fx.lstm_utils._get_reference_quantized_lstm_module(
                    observed_lstm=observed_lstm,
                    backend_config=my_backend_config,
                )

        # FX graph mode quantization
m = MyModel()
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
example_inputs = (torch.rand(5, 3, 50), torch.rand(1, 3, 50), torch.randn(1, 3, 50))
prepare_custom_config = PrepareCustomConfig() \
                    .set_float_to_observed_mapping(torch.nn.LSTM, UserObservedLSTM)
convert_custom_config = ConvertCustomConfig() \
                    .set_observed_to_quantized_mapping(torch.ao.nn.quantizable.LSTM, UserQuantizedLSTM)
prepared = prepare_fx(
            m,
            qconfig_mapping,
            example_inputs,
            prepare_custom_config,
            backend_config=my_backend_config,
        )
prepared(*example_inputs)
converted = convert_fx(
            prepared,
            convert_custom_config,
            backend_config=my_backend_config,
        )
converted(*example_inputs)

</code></pre>
<p>This is the warning:</p>
<pre><code class="lang-auto">UserWarning: QConfig must specify a FixedQParamsObserver or a FixedQParamsFakeQuantize for fixed qparams ops, ignoring QConfig(activation=functools.partial(&lt;class 'torch.ao.quantization.observer.HistogramObserver'&gt;, reduce_range=False){'factory_kwargs': &lt;function _add_module_to_qconfig_obs_ctr.&lt;locals&gt;.get_factory_kwargs_based_on_module_device at 0x7f66b4195a20&gt;}, weight=functools.partial(&lt;class 'torch.ao.quantization.observer.MinMaxObserver'&gt;, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric){'factory_kwargs': &lt;function _add_module_to_qconfig_obs_ctr.&lt;locals&gt;.get_factory_kwargs_based_on_module_device at 0x7f66b4195a20&gt;}).
Please use torch.ao.quantization.get_default_qconfig_mapping or torch.ao.quantization.get_default_qat_qconfig_mapping. Example:
    qconfig_mapping = get_default_qconfig_mapping("fbgemm")
    model = prepare_fx(model, qconfig_mapping, example_inputs)
  warnings.warn(("QConfig must specify a FixedQParamsObserver or a FixedQParamsFakeQuantize "
</code></pre>
<p>torch version: 2.0.1</p>

HDCharles: <p>can you just run the test script, does that test pass for you?</p>
<p>python test/test_quantization.py TestQuantizeFx.test_static_lstm_with_custom_fixed_qparams</p>

HDCharles: <p>ahh i see.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/75a8c8a538287f1704f82fcfe64693e1ae955ba4/torch/ao/nn/quantizable/modules/rnn.py#L86">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/75a8c8a538287f1704f82fcfe64693e1ae955ba4/torch/ao/nn/quantizable/modules/rnn.py#L86" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/75a8c8a538287f1704f82fcfe64693e1ae955ba4/torch/ao/nn/quantizable/modules/rnn.py#L86" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/75a8c8a538287f1704f82fcfe64693e1ae955ba4/torch/ao/nn/quantizable/modules/rnn.py#L86</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="76" style="counter-reset: li-counter 75 ;">
          <li>    input_gate = self.input_gate(input_gate)</li>
          <li>    forget_gate = self.forget_gate(forget_gate)</li>
          <li>    cell_gate = self.cell_gate(cell_gate)</li>
          <li>    out_gate = self.output_gate(out_gate)</li>
          <li></li>
          <li>    fgate_cx = self.fgate_cx.mul(forget_gate, cx)</li>
          <li>    igate_cgate = self.igate_cgate.mul(input_gate, cell_gate)</li>
          <li>    fgate_cx_igate_cgate = self.fgate_cx_igate_cgate.add(fgate_cx, igate_cgate)</li>
          <li>    cy = fgate_cx_igate_cgate</li>
          <li></li>
          <li class="selected">    # TODO: make this tanh a member of the module so its qparams can be configured</li>
          <li>    tanh_cy = torch.tanh(cy)</li>
          <li>    hy = self.ogate_cy.mul(out_gate, tanh_cy)</li>
          <li>    return hy, cy</li>
          <li></li>
          <li>def initialize_hidden(self, batch_size: int, is_quantized: bool = False) -&gt; Tuple[Tensor, Tensor]:</li>
          <li>    h, c = torch.zeros((batch_size, self.hidden_size)), torch.zeros((batch_size, self.hidden_size))</li>
          <li>    if is_quantized:</li>
          <li>        (h_scale, h_zp) = self.initial_hidden_state_qparams</li>
          <li>        (c_scale, c_zp) = self.initial_cell_state_qparams</li>
          <li>        h = torch.quantize_per_tensor(h, scale=h_scale, zero_point=h_zp, dtype=self.hidden_state_dtype)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>there are 2 tanh ops, one cell_gate, one tanh_cy, only the cell_gate one gets a fixed_qparams qconfig. given the comment, it looks like its a known bug</p>

Ahmed_Louati: <p>Yeah, maybe!<br>
Could you please discard the warning message error until you fix it?<br>
Besides, can we have an explanation of the comments here?</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/29f856e3e02cea6b608167555719c907a80fb135/test/quantization/fx/test_quantize_fx.py#L4917C18-L4917C18">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/29f856e3e02cea6b608167555719c907a80fb135/test/quantization/fx/test_quantize_fx.py#L4917C18-L4917C18" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/29f856e3e02cea6b608167555719c907a80fb135/test/quantization/fx/test_quantize_fx.py#L4917C18-L4917C18" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/29f856e3e02cea6b608167555719c907a80fb135/test/quantization/fx/test_quantize_fx.py#L4917C18-L4917C18</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="4917" style="counter-reset: li-counter 4916 ;">
          <li># uint16, [-16, 16)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
The dtype is set to torch.qint32 (which is aligned with the backend config) but I don’t know why it is mentioned uint16 (are you just putting an uint16 in qint32? (something with casting)).</p>
<p>And as <a class="mention" href="/u/jerryzh168">@jerryzh168</a> mentioned, shouldn’t we use fixed qparams for at least sigmoid and tanh ops?<br>
<a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/backend_config/_common_operator_config_utils.py#L67" rel="noopener nofollow ugc">pytorch/torch/ao/quantization/backend_config/_common_operator_config_utils.py at main · pytorch/pytorch (github.com)</a></p>
<p>Thank you for the clarification!</p>

HDCharles: <aside class="quote no-group" data-username="HDCharles" data-post="7" data-topic="184489">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hdcharles/48/38408_2.png" class="avatar"> HDCharles:</div>
<blockquote>
<pre><code class="lang-auto">    # TODO: make this tanh a member of the module so its qparams can be configured
          
</code></pre>
</blockquote>
</aside>
<p>its just a test simulating some weird dtypes using a qint32.</p>
<p>The warning is working as intended, a fixedqparams op is getting matched with your global qconfig, you could either just make your qconfig_mapping handle it or ignore the warning.</p>