Qat_swap_modules should consider the non_traceable_module_class

BobChen: <p>If there is <code>non_traceable_module_class</code> in <code>prepare_custom_config_dict</code>, <code>qat_swap_modules</code> would still swap the module under the <code>non_traceable_module_class</code>. I think this should be an unexpected behavior, since if a module is non-traceable, it should not be quantized.</p>
<p>example codes:</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.quantization.quantize_fx import prepare_qat_fx


class UnTraceableModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 2)

    def forward(self, x):
        return self.linear(x)


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 2)
        self.untraceable_module = UnTraceableModule()

    def forward(self, x):
        x = self.linear(x)
        x = self.untraceable_module(x)
        return x


model = MyModel()
qconfig_dict = {
    "": torch.quantization.get_default_qat_qconfig()
}
prepare_custom_config_dict = {
    "non_traceable_module_class": [UnTraceableModule]
}
prepared_model = prepare_qat_fx(model.train(), qconfig_dict, prepare_custom_config_dict)
print(type(prepared_model.untraceable_module.linear))    # &lt;---- get &lt;class 'torch.nn.qat.modules.linear.Linear'&gt;
</code></pre>

HDCharles: <p>made an issue:</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/74240">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/74240" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/74240" target="_blank" rel="noopener nofollow ugc">FX QAT Untraceable Modules are still converted to qat versions if possible</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2022-03-15" data-time="15:58:41" data-timezone="UTC">03:58PM - 15 Mar 22 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/HDCharles" target="_blank" rel="noopener nofollow ugc">
          <img alt="HDCharles" src="https://avatars.githubusercontent.com/u/39544797?v=4" class="onebox-avatar-inline" width="20" height="20">
          HDCharles
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: quantization
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          fx
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🐛 Describe the bug

untraced modules should be excluded for the QAT conver<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">t in fx quantization.

repro:
```
from torch.quantization.quantize_fx import prepare_qat_fx


class UnTraceableModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 2)

    def forward(self, x):
        return self.linear(x)


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 2)
        self.untraceable_module = UnTraceableModule()

    def forward(self, x):
        x = self.linear(x)
        x = self.untraceable_module(x)
        return x


model = MyModel()
qconfig_dict = {
    "": torch.quantization.get_default_qat_qconfig()
}
prepare_custom_config_dict = {
    "non_traceable_module_class": [UnTraceableModule]
}
prepared_model = prepare_qat_fx(model.train(), qconfig_dict, prepare_custom_config_dict)
print(type(prepared_model.untraceable_module.linear))    # &lt;---- get &lt;class 'torch.nn.qat.modules.linear.Linear'&gt;
```

the inner linear module should not be converted to the qat version

This happens because for prepare() in torch/ao/quantization/fx/prepare.py 
1) propagate_qconfig_() doesn't use the graph and propagates qconfig by named modules, so it adds a qconfig to the hidden Linear module (https://github.com/pytorch/pytorch/blob/3b421af85a876f727599c647d0111739aac028f0/torch/ao/quantization/fx/prepare.py#L1392)
2) qat_swap_modules() eventually hits swap_module() (https://github.com/pytorch/pytorch/blob/7ab482f2df7a34224b24d5f8fe16a3ee70f1ec4a/torch/ao/quantization/quantize.py#L553) which only checks for a qconfig



### Versions

Collecting environment information...
PyTorch version: 1.12.0a0+git730fef2
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.2.1 (x86_64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2)
CMake version: version 3.19.6
Libc version: N/A

Python version: 3.8.8 (default, Apr 13 2021, 12:59:45)  [Clang 10.0.0 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.1
[pip3] numpydoc==1.1.0
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] torch==1.11.0a0+git20f2e55
[pip3] torchvision==0.12.0a0+bf073e7
[conda] blas                      1.0                         mkl  
[conda] libblas                   3.9.0                     9_mkl    conda-forge
[conda] libcblas                  3.9.0                     9_mkl    conda-forge
[conda] liblapack                 3.9.0                     9_mkl    conda-forge
[conda] mkl                       2021.2.0           hecd8cb5_269  
[conda] mkl-include               2021.4.0           hecd8cb5_637  
[conda] mkl-service               2.3.0            py38h9ed2024_1  
[conda] mkl_fft                   1.3.0            py38h4a7008c_2  
[conda] mkl_random                1.2.1            py38hb2f4e1b_2  
[conda] mypy                      0.910              pyhd3eb1b0_0  
[conda] mypy_extensions           0.4.3                    py38_0  
[conda] numpy                     1.20.1           py38hd6e1bb9_0  
[conda] numpy-base                1.20.1           py38h585ceec_0  
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] pytorch-sphinx-theme      0.0.24                   pypi_0    pypi
[conda] torch                     1.11.0a0+git20f2e55          pypi_0    pypi
[conda] torchvision               0.12.0a0+bf073e7           dev_0    &lt;develop&gt;

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo</span></p>
  </div>

  </article>

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

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

<p>I’ll see if I can put together a fix</p>