BatchNorm and ConvTranspose Fusion for QAT with FX Graph Mode

username: Enceladus
cooked: <h2><a name="description-1" class="anchor" href="#description-1"></a>Description</h2>
<p>I’m trying to quantize a model for training using FX Graph Mode. (ultimately, I want to run it with int8 precision using TensorRT, but that’s not the issue for now). Tracing works fine, the problem is during the fusion stage. My model uses BatchNorm and ConvTranspose modules, for which fusion is not yet supported for QAT. I tried just skipping the fusion of these specific modules, but without success. I’m not familiar enough with this part of PyTorch. My questions are:</p>
<ul>
<li>Can I skip the BatchNorm-ConvTranspose fusion and continue with QAT? How would I do that?</li>
<li>Do I need to fuse in general before finetuning/QAT? Or can I get away with fusing after QAT?</li>
</ul>
<p>I’m running:</p>
<pre><code class="lang-auto">qconfig_mapping = get_default_qat_qconfig_mapping("qnnpack")
custom_config = PrepareCustomConfig()
model_prepared = prepare_qat_fx(model, qconfig_mapping, example_input, custom_config)
</code></pre>
<p>The error that I get:</p>
<pre><code class="lang-auto">model_prepared = prepare_qat_fx(model, qconfig_mapping, example_input, custom_config)
  File "\venv\lib\site-packages\torch\ao\quantization\quantize_fx.py", line 489, in prepare_qat_fx
    return _prepare_fx(
  File "\venv\lib\site-packages\torch\ao\quantization\quantize_fx.py", line 139, in _prepare_fx
    graph_module = _fuse_fx(
  File "\venv\lib\site-packages\torch\ao\quantization\quantize_fx.py", line 91, in _fuse_fx
    return fuse(
  File "\venv\lib\site-packages\torch\ao\quantization\fx\fuse.py", line 112, in fuse
    env[node.name] = obj.fuse(
  File "\venv\lib\site-packages\torch\ao\quantization\fx\fuse_handler.py", line 102, in fuse
    fused_module = fuser_method(is_qat, *matched_modules)
  File "\venv\lib\site-packages\torch\ao\quantization\fuser_method_mappings.py", line 195, in reversed
    return f(is_qat, y, x)
  File "\venv\lib\site-packages\torch\ao\quantization\fuser_method_mappings.py", line 149, in fuse_convtranspose_bn
    raise Exception("Fusing ConvTranspose+BatchNorm not yet supported in QAT.")
Exception: Fusing ConvTranspose+BatchNorm not yet supported in QAT.
</code></pre>
<p>Thanks!</p>
<h2><a name="environment-2" class="anchor" href="#environment-2"></a>Environment</h2>
<p>Operating System: Windows<br>
Python Version: 3.9.12<br>
PyTorch Version: 2.1.1+cu121</p>
<h2><a name="steps-to-reproduce-3" class="anchor" href="#steps-to-reproduce-3"></a>Steps To Reproduce</h2>
<p>I’m happy to add an example model if it is needed/helpful.</p>

username: jerryzh168
cooked: <p>fx graph mode quantization is in maintainence mode currently, I think you can start with our new flow: <a href="https://pytorch.org/docs/main/quantization.html#prototype-pytorch-2-export-quantization" class="inline-onebox" rel="noopener nofollow ugc">Quantization — PyTorch main documentation</a>, this flow does not come with conv transpose - bn fusion in QAT currently</p>

username: Enceladus
cooked: <p>It seems that I need to switch to a Linux system for this, but that should be doable. Thanks!<br>
Do you suspect any difficulties in exporting to ONNX and using TensorRT with this approach?</p>
<pre><code class="lang-auto">  File "\venv\lib\site-packages\torch\_export\__init__.py", line 145, in capture_pre_autograd_graph
    m = torch._dynamo.export(
  File "\venv\lib\site-packages\torch\_dynamo\eval_frame.py", line 1031, in inner
    check_if_dynamo_supported()
  File "\venv\lib\site-packages\torch\_dynamo\eval_frame.py", line 535, in check_if_dynamo_supported
    raise RuntimeError("Windows not yet supported for torch.compile")
RuntimeError: Windows not yet supported for torch.compile
</code></pre>

username: jerryzh168
cooked: <aside class="quote no-group" data-username="Enceladus" data-post="3" data-topic="193784">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/e/bc79bd/48.png" class="avatar"> Enceladus:</div>
<blockquote>
<pre><code class="lang-auto">    raise RuntimeError("Windows not yet supported for torch.compile")
</code></pre>
</blockquote>
</aside>
<p>yeah windows is not supported…</p>
<blockquote>
<p>Do you suspect any difficulties in exporting to ONNX and using TensorRT with this approach?</p>
</blockquote>
<p>I don’t see there will be unfixable issues, but I haven’t tried this before</p>