Problem in symbolically trace (torch.fx) nn.GRUCell/LSTMCell

Ahmed_Louati: <p>Hello,<br>
I am trying to symbolically trace a model containing an nn.GRUCell using torch.fx to prepare it for quantization.<br>
I tried 2 approaches:<br>
<strong>1. Define states of the mode using self.states</strong><br>
Here is the code:</p>
<pre><code class="lang-auto">import torch 
import torch.nn as nn
from pulse_ai.tiny.helper_functions import (prepare, profile)
from torch.ao.quantization import get_default_qconfig_mapping
import torch.ao.quantization.quantize_fx as quantize_fx


class GRU(nn.Module):
    def __init__(self, input_size, hidden_size=None):
        super().__init__()
        if hidden_size is None:
            hidden_size = input_size
        # Intialize empty states
        self.states = torch.zeros(1, hidden_size)
        #self.register_buffer('states', states)
        self.rnn = nn.GRUCell(input_size, hidden_size)

    def forward(self, x):
        self.states = self.rnn(x, self.states) 
        return self.states
    
TestGRU = GRU(128)

dummy_input = torch.rand(1,128)

outputs = TestGRU(dummy_input)

from torch.fx import symbolic_trace
# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced : torch.fx.GraphModule = symbolic_trace(TestGRU)

symbolic_traced.graph.print_tabular()

default_mapping = get_default_qconfig_mapping("qnnpack")

model_prepared = quantize_fx.prepare_fx(
                symbolic_traced,
                default_mapping,
                dummy_input,
)

</code></pre>
<p>This is the error I got</p>
<pre><code class="lang-auto">/home/ahmed/anaconda3/envs/tiny_pulse_env/bin/python /home/ahmed/PulseAudition_PFE/pulse_ai/scripts/tiny/Test.py
opcode       name    target    args         kwargs
-----------  ------  --------  -----------  --------
placeholder  x       x         ()           {}
get_attr     states  states    ()           {}
call_module  rnn     rnn       (x, states)  {}
output       output  output    (rnn,)       {}
Traceback (most recent call last):
  File "/home/ahmed/PulseAudition_PFE/pulse_ai/scripts/tiny/Test.py", line 36, in &lt;module&gt;
    model_prepared = quantize_fx.prepare_fx(
                     ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/ao/quantization/quantize_fx.py", line 380, in prepare_fx
    return _prepare_fx(
           ^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/ao/quantization/quantize_fx.py", line 137, in _prepare_fx
    graph_module = _fuse_fx(
                   ^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/ao/quantization/quantize_fx.py", line 89, in _fuse_fx
    return fuse(
           ^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/ao/quantization/fx/fuse.py", line 115, in fuse
    env[node.name] = fused_graph.node_copy(node, load_arg)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/graph.py", line 1163, in node_copy
    args = map_arg(node.args, arg_transform)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 621, in map_arg
    return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 629, in map_aggregate
    t = tuple(map_aggregate(elem, fn) for elem in a)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 629, in &lt;genexpr&gt;
    t = tuple(map_aggregate(elem, fn) for elem in a)
              ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 639, in map_aggregate
    return fn(a)
           ^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 621, in &lt;lambda&gt;
    return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
                                      ^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/ao/quantization/fx/fuse.py", line 86, in load_arg
    return map_arg(a, lambda node: env[node.name])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 621, in map_arg
    return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 639, in map_aggregate
    return fn(a)
           ^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/node.py", line 621, in &lt;lambda&gt;
    return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
                                      ^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/ao/quantization/fx/fuse.py", line 86, in &lt;lambda&gt;
    return map_arg(a, lambda node: env[node.name])
                                   ~~~^^^^^^^^^^^
KeyError: 'rnn'
</code></pre>
<p><strong>2. Define states using seld.register_buffer(‘states’,states)</strong><br>
Code</p>
<pre><code class="lang-auto">import torch 
import torch.nn as nn
from pulse_ai.tiny.helper_functions import (prepare, profile)
from torch.ao.quantization import get_default_qconfig_mapping
import torch.ao.quantization.quantize_fx as quantize_fx


class GRU(nn.Module):
    def __init__(self, input_size, hidden_size=None):
        super().__init__()
        if hidden_size is None:
            hidden_size = input_size
        # Intialize empty states
        states = torch.zeros(1, hidden_size)
        self.register_buffer('states', states)
        self.rnn = nn.GRUCell(input_size, hidden_size)

    def forward(self, x):
        self.states = self.rnn(x, self.states) 
        return self.states
    
TestGRU = GRU(128)

dummy_input = torch.rand(1,128)

outputs = TestGRU(dummy_input)

from torch.fx import symbolic_trace
# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced : torch.fx.GraphModule = symbolic_trace(TestGRU)

symbolic_traced.graph.print_tabular()

default_mapping = get_default_qconfig_mapping("qnnpack")

model_prepared = quantize_fx.prepare_fx(
                symbolic_traced,
                default_mapping,
                dummy_input,
)
</code></pre>
<p>Error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/ahmed/PulseAudition_PFE/pulse_ai/scripts/tiny/Test.py", line 30, in &lt;module&gt;
    symbolic_traced : torch.fx.GraphModule = symbolic_trace(TestGRU)
                                             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 1109, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/home/ahmed/PulseAudition_PFE/pulse_ai/scripts/tiny/Test.py", line 19, in forward
    self.states = self.rnn(x, self.states) 
    ^^^^^^^^^^^
  File "/home/ahmed/anaconda3/envs/tiny_pulse_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1665, in __setattr__
    raise TypeError("cannot assign '{}' as buffer '{}' "
TypeError: cannot assign 'Proxy(add)' as buffer 'states' (torch.Tensor or None expected)
</code></pre>
<p>In both approaches, I think the problem is when passing “SELF”.states in the forward function.<br>
Can you suggest any solution?<br>
( Maybe try to externally update the states of the cell! )</p>
<p>Hope to hearing from you soon,<br>
Ahmed</p>

jerryzh168: <p>yeah GRU is not symbolically traceable, but you can trace it with <code>torchdynamo</code> with the following code:</p>
<pre><code class="lang-auto">import torch
import torch._dynamo as torchdynamo
import copy
torchdynamo.config.allow_rnn = True


class RNNDynamicModel(torch.nn.Module):
    def __init__(self, mod_type):
        super().__init__()
        if mod_type == "GRU":
            self.mod = torch.nn.GRU(2, 2).to(dtype=torch.float)
        if mod_type == "LSTM":
            self.mod = torch.nn.LSTM(2, 2).to(dtype=torch.float)

    def forward(self, x):
        x = self.mod(x)
        return x


niter = 10
example_inputs = (
    torch.tensor([[100, -155], [-155, 100], [100, -155]], dtype=torch.float)
    .unsqueeze(0)
    .repeat(niter, 1, 1),
)
model = RNNDynamicModel("GRU")
model, guards = torchdynamo.export(
    model,
    *copy.deepcopy(example_inputs),
    aten_graph=True,
    tracing_mode="real",
)
</code></pre>
<p>this will work with our new quantization tool in pytorch 2.0 export (<a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/_quantize_pt2e.py#L50" class="inline-onebox" rel="noopener nofollow ugc">pytorch/_quantize_pt2e.py at main · pytorch/pytorch · GitHub</a>), I’m trying to figure out how we can quantize it right now actually.</p>

jerryzh168: <p>if you need to use fx graph mode quantization, the closet is LSTM custom module quantization:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/test/quantization/fx/test_quantize_fx.py#L4886">
  <header class="source">

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

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



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="4876" style="counter-reset: li-counter 4875 ;">
          <li>        # lstm[1][0].dequantize()</li>
          <li>        # lstm[1][1].dequantize()</li>
          <li>        ns.call_method("dequantize"): 2,</li>
          <li>        # lstm[1], lstm[1][0], lstm[1][1]</li>
          <li>        ns.call_function(operator.getitem): 3,</li>
          <li>        # tuple(hidden0_dq, hidden1_dq)</li>
          <li>        ns.call_function(tuple): 1,</li>
          <li>    }</li>
          <li>    self._test_static_lstm_helper(m2, prepare_node_occurrence, convert_node_occurrence2)</li>
          <li>
          </li>
<li class="selected">def test_static_lstm_with_custom_fixed_qparams(self):</li>
          <li>    """</li>
          <li>    Test statically quantized LSTM with custom fixed qparams assigned to each of the</li>
          <li>    inner submodules. This flow requires users to extend `torch.ao.nn.quantizable.LSTM`</li>
          <li>    and use the child class in the custom module mapping.</li>
          <li>    """</li>
          <li>    class MyModel(torch.nn.Module):</li>
          <li>        def __init__(self):</li>
          <li>            super().__init__()</li>
          <li>            self.my_lstm = torch.nn.LSTM(50, 50, 1)</li>
          <li>
      </li>
</ol>
    </code></pre>



  </article>

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

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

<p>where you need to write a custom observed and quantized module for GRU (break down the single GRU function call to calls to linear and non-linearities) in order to quantize it.</p>

Ahmed_Louati: <p>Thank you for your quick reply!<br>
Is it easier to quantize LSTM/GRU layers using Eager Mode?</p>

jerryzh168: <aside class="quote no-group" data-username="Ahmed_Louati" data-post="4" data-topic="181083">
<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/ahmed_louati/48/59732_2.png" class="avatar"> Ahmed_Louati:</div>
<blockquote>
<p>Is it easier to quantize LSTM/GRU layers using Eager Mode?</p>
</blockquote>
</aside>
<p>I don’t think it’srelated to eager mode or fx graph mode, it depends on whether you need to customize how to quantize each submodule of the quantizable LSTM or not, if you can use a global qconfig for all submodules then it will be simpler, both in fx graph mode and eager mode</p>

Ahmed_Louati: <p>Thank you <a class="mention" href="/u/jerryzh168">@jerryzh168</a> !<br>
I just wanted to make sure that after this custom quantization, we won’t be able to export the quantized layers to ONNX (at least for the moment!)</p>
<p>Thank you,</p>

jerryzh168: <p>yeah I haven’t tried, but it’s likely not exportable to ONNX</p>

cokespace2: <p>HI Jerry, just saw this, where can we have some docs on pt2e export? I just see the code now but don’t know the usage and it’s advantages.</p>

HDCharles: <p>its in development, its still experimental, as a result the docs are not complete yet as design changes may still occur.</p>

jerryzh168: <p>please stay tuned, we’ll have something after the end of this week and I’ll post here</p>

jerryzh168: <aside class="quote no-group" data-username="cokespace2" data-post="8" data-topic="181083" data-full="true">
<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/cokespace2/48/51244_2.png" class="avatar"> cokespace2:</div>
<blockquote>
<p>HI Jerry, just saw this, where can we have some docs on pt2e export? I just see the code now but don’t know the usage and it’s advantages.</p>
</blockquote>
</aside>
<p>OK two docs here:</p>
<ul>
<li><a href="https://pytorch.org/tutorials/prototype/pt2e_quant_ptq_static.html" class="inline-onebox" rel="noopener nofollow ugc">(prototype) PyTorch 2.0 Export Post Training Static Quantization — PyTorch Tutorials 2.0.1+cu117 documentation</a></li>
<li><a href="https://pytorch.org/tutorials/prototype/pt2e_quantizer.html" class="inline-onebox" rel="noopener nofollow ugc">How to Write a Quantizer for PyTorch 2 Export Quantization — PyTorch Tutorials 2.1.1+cu121 documentation</a></li>
</ul>
<p>There are still some formatting issues (e.g. some code didn’t display properly) that will be fixed next week</p>

cokespace2: <p>Cool, just read the tutorial of this pt2e, and very glad to see that we can use less configs to quantize a model.<br>
I struggled with fx’s too many configs for a while, I still cannot fully understand the differents of <code>qconfig</code>, <code>qconfig_mapping</code>, <code>equalization_config</code> and <code>backend_config</code>…</p>