Torch 2.0 compile not compatible with FX Graph Mode Quantization?

Archer_Z: <p>Hi!<br>
I’ve been trying the new compile feature with the FX quantization. However it seems compile the qat prepared model gives shape mismatch error.</p>
<pre><code class="lang-auto">from torchvision.models import resnet18
import torch
from torch.ao.quantization.quantize_fx import prepare_qat_fx

net = resnet18().to(torch.float32).cuda()
net.train()
qconfig_mapping = torch.ao.quantization.get_default_qat_qconfig_mapping('x86')
example_input = (torch.randn((10, 3, 128, 128)),)
net = prepare_qat_fx(net, qconfig_mapping, example_input)

net = torch.compile(net)
example_input = torch.randn((10, 3, 128, 128)).cuda()
output = net(example_input)
</code></pre>
<p>Running this code give me error</p>
<pre><code class="lang-auto">  File "/home/test_demo.py", line 14, in &lt;module&gt;
    output = net(example_input)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "&lt;eval_with_key&gt;.2", line 4, in forward
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2819, in forward
    return compiled_fn(full_args)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1222, in g
    return f(*args)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2386, in debug_compiled_function
    return compiled_function(*args)
  File "/home/.conda/envs/wzpy39/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1970, in runtime_wrapper
    original_inpt.copy_(updated_inpt)
RuntimeError: output with shape [1] doesn't match the broadcast shape [64]
</code></pre>
<p>if I comment out the following line then it runs fine.</p>
<pre><code class="lang-auto">qconfig_mapping = torch.ao.quantization.get_default_qat_qconfig_mapping('x86')
example_input = (torch.randn((10, 3, 128, 128)),)
net = prepare_qat_fx(net, qconfig_mapping, example_input)
</code></pre>

cokespace2: <p>I met similar problem.</p>
<p>I was tring to quantize a compiled model, for your example, code could be like:</p>
<pre><code class="lang-python">from torchvision.models import resnet18
import torch
from torch.ao.quantization.quantize_fx import prepare_qat_fx

net = resnet18().to(torch.float32)
net.train()

example_input = torch.randn((10, 3, 128, 128))
net = torch.compile(net)
output = net(example_input)

qconfig_mapping = torch.ao.quantization.get_default_qat_qconfig_mapping('x86')
net = prepare_qat_fx(net, qconfig_mapping, example_input)

output = net(example_input)
</code></pre>
<p>and the error message:</p>
<pre data-code-wrap="shell"><code class="lang-plaintext">Traceback (most recent call last):
  File "fx_demo.py", line 13, in &lt;module&gt;
    net = prepare_qat_fx(net, qconfig_mapping, example_input)
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/ao/quantization/quantize_fx.py", line 487, in prepare_qat_fx
    return _prepare_fx(
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/ao/quantization/quantize_fx.py", line 133, in _prepare_fx
    graph_module = GraphModule(model, tracer.trace(model))
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/fx/proxy.py", line 385, in __iter__
    return self.tracer.iter(self)
  File "/nfs1/fan.mo/00-pyenv/torch2.0/lib/python3.8/site-packages/torch/fx/proxy.py", line 285, in iter
    raise TraceError('Proxy object cannot be iterated. This can be '
torch.fx.proxy.TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors
</code></pre>

marksaroufim: <p>Quantization support in torch.compile is still early days but if you’re curious you can see how it works here <a href="https://github.com/facebookexperimental/protoquant" rel="noopener nofollow ugc">https://github.com/facebookexperimental/protoquant</a></p>

jcaip: <p>Hi <a class="mention" href="/u/archer_z">@Archer_Z</a></p>
<p>I’m not sure if torch.compile will work with the prepared QAT model. I do think however you should be able to follow the QAT workflow and call torch.compile on the returned converted model. There’s a QAT tutorial <a href="https://pytorch.org/docs/stable/quantization.html#quantization-aware-training-for-static-quantization" rel="noopener nofollow ugc">here</a> if it’s helpful.</p>
<p>cc <a class="mention" href="/u/andrewor">@andrewor</a> do you know if this is supported / in our roadmap?</p>

andrewor: <p>Hi <a class="mention" href="/u/archer_z">@Archer_Z</a>, there are no plans right now to support torch.compile with FX graph mode quantization, so it may not work well out of the box. We do have plans to support a new PT2.0 QAT flow that is compatible with torch.compile, but that is still in the works and may not be ready until later this year.</p>