How to approach quantizing torch.matmul/aten::bmm?

yannbane: <p>I have two quantized tensors:</p>
<pre><code class="lang-auto">In [14]: q.size()
Out[14]: torch.Size([64, 3, 49, 32])

In [15]: k.size()
Out[15]: torch.Size([64, 3, 49, 32])
</code></pre>
<p>I’m trying to run the following operation:</p>
<pre><code class="lang-python">torch.matmul(q, k.transpose(-2, -1))
</code></pre>
<p>Which yields the usual error:</p>
<pre><code class="lang-auto">RuntimeError: Could not run 'aten::bmm' with arguments from the 'QuantizedCPU' backend.
</code></pre>
<p>I’m aware that matmul apparently isn’t supported in quantization. <strong>But is there a way to transform or use another operation that would allow these activations/tensors to remain quantized, but achieve the same outcome as matmul?</strong> Of course while retaining the potential performance benefit of quantization overall.</p>
<p><strong>Also is there a timetable on when matmul ops will be supported?</strong></p>

Vasiliy_Kuznetsov: <p>If you use FX graph mode quantization, quantizing <code>torch.matmul</code> is supported.  It’s not supported in Eager mode quantization.</p>

yannbane: <p>Thanks for the info! FX graph mode quantization seemed really tough to set up for this model. Is there a reason why <code>matmul</code> isn’t supported in eager mode? I’m currently researching how to add the support myself, but if it’s a big technical issue I’d consider FX again.</p>

Vasiliy_Kuznetsov: <p>It’s possible to do this in Eager mode, it’s just not something we plan to land in the codebase as the Eager mode workflow is not designed for handling functions, it’s better at handling modules.  We have handling for some very common functions, but matmul is not one of them.  If someone is interested in adding matmul to these workarounds, we’d be happy to accept a PR.</p>
<p>If you want to get it working for your model, you can follow what is done for functions such as torch.add.  Specifically:</p>
<ol>
<li>Start with the <code>FloatFunctional</code> and <code>QFunctional</code> classes (<a href="https://github.com/pytorch/pytorch/blob/2f04ba2c7c8920418ad77ebb1ab09d93374e6578/torch/ao/nn/quantized/modules/functional_modules.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch/functional_modules.py at 2f04ba2c7c8920418ad77ebb1ab09d93374e6578 · pytorch/pytorch · GitHub</a>)</li>
<li>extend the classes for your operation (in your case, add a method for <code>matmul</code>)</li>
<li>rewrite your model, instead of calling <code>torch.matmul</code> you need to create an instance of your new class, and create <code>float_functional.matmul(a, b)</code> on it.  Each call of the matmul should have its own instance of the class, so that they can collect statistics separately.</li>
<li>add your new classes to <a href="https://github.com/pytorch/pytorch/blob/2f04ba2c7c8920418ad77ebb1ab09d93374e6578/torch/ao/quantization/quantization_mappings.py#L51" class="inline-onebox" rel="noopener nofollow ugc">pytorch/quantization_mappings.py at 2f04ba2c7c8920418ad77ebb1ab09d93374e6578 · pytorch/pytorch · GitHub</a> . This will tell the <code>prepare</code> and <code>convert</code> functions that your new classes exist.</li>
<li>run the quantization flow</li>
</ol>
<p>Here is a test case which uses this logic, in case that is helpful: <a href="https://github.com/pytorch/pytorch/blob/7088a98fba3a5031a2afc293cbf25cec09f248a5/test/quantization/eager/test_quantize_eager_qat.py#L555" class="inline-onebox" rel="noopener nofollow ugc">pytorch/test_quantize_eager_qat.py at 7088a98fba3a5031a2afc293cbf25cec09f248a5 · pytorch/pytorch · GitHub</a></p>

Vasiliy_Kuznetsov: <p>To clarify, you don’t need to change PyTorch source code for this, you just need to create some custom classes, tell prepare/convert about these classes, and rewrite your model to call these classes.</p>

yannbane: <p>Wow thanks! I kind of went through that workflow to add support for a quantized softmax. But I am confused: the bindings for quantized softmax were already accessible: <code>torch.ops.quantized.softmax(x, self.dim, self.scale, self.zero_point)</code>, so I just had to instruct Pytorch to convert <code>nn.Softmax</code> into my extension of <code>FloatFunctional</code>.</p>
<p>But as I understand it, there are no such bindings for <code>matmul</code> exposed by ATen available. I would actually have to add these to ATen myself, and then rebuild Pytorch.</p>
<p>Am I mistaken?</p>

Vasiliy_Kuznetsov: <p>The quantized version of <code>torch.matmul</code> was added in <a href="https://github.com/pytorch/pytorch/pull/71783" class="inline-onebox" rel="noopener nofollow ugc">[PyTorch Edge] Add Quantized Matmul Op (Naive Implementation) by salilsdesai · Pull Request #71783 · pytorch/pytorch · GitHub</a>, if you are on the latest version of PyTorch (1.12) you should be able to use it.</p>

yannbane: <p>Ah I missed that. But doesn’t it say:</p>
<blockquote>
<p>Summary: Adds Quantized Matmul op which just naively performs dequantize → matmul → quantize. (To be optimized in the future)</p>
</blockquote>
<p>So basically it’s kind of doing what I’m already doing? And if I want a true quantized op I’d still have to make it myself?</p>

Vasiliy_Kuznetsov: <p>I just linked the PR which added the original function, there have been some improvements on it landed.  You could check the current state here: <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/quantized/cpu/qmatmul.cpp" class="inline-onebox" rel="noopener nofollow ugc">pytorch/qmatmul.cpp at master · pytorch/pytorch · GitHub</a> .  Whether the improvements would be useful for you would depend on whether your backend supports the ruy library.</p>

yannbane: <p>In the meantime I’ve discovered that yes, but now I’m trying to compile Pytorch with that flag enabled and can’t to get the development docker image built. Made a new thread about that: <a href="https://discuss.pytorch.org/t/cant-build-pytorch-using-the-dockerfile-from-the-repo/162062/3" class="inline-onebox">Can't build Pytorch using the Dockerfile from the repo - #3 by yannbane</a></p>