How to do qat after ptq in PyTorch2 quantization?

xiguadong: <p>As mentioned in <a href="https://pytorch.org/docs/stable/quantization.html" rel="noopener nofollow ugc">Quantization — PyTorch 2.6 documentation</a> , supports 3 mode: Eager / Fx / Pytorch2.<br>
And in Eager and Fx mode, the demo shows pipeline for ptq → qat, it’s useful to get bertter precision.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.ao.quantization import (
  get_default_qconfig_mapping,
  get_default_qat_qconfig_mapping,
  QConfigMapping,
)
import torch.ao.quantization.quantize_fx as quantize_fx
import copy

model_fp = UserModel()

#
# post training dynamic/weight_only quantization
#

# we need to deepcopy if we still want to keep model_fp unchanged after quantization since quantization apis change the input model
model_to_quantize = copy.deepcopy(model_fp)
model_to_quantize.eval()
qconfig_mapping = QConfigMapping().set_global(torch.ao.quantization.default_dynamic_qconfig)
# a tuple of one or more example inputs are needed to trace the model
example_inputs = (input_fp32)
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# no calibration needed when we only have dynamic/weight_only quantization
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# post training static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model_to_quantize.eval()
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# calibrate (not shown)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# quantization aware training for static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qat_qconfig_mapping("qnnpack")
model_to_quantize.train()
# prepare
model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_mapping, example_inputs)
# training loop (not shown)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# fusion
#
model_to_quantize = copy.deepcopy(model_fp)
model_fused = quantize_fx.fuse_fx(model_to_quantize)
</code></pre>
<p><a href="https://pytorch.org/docs/stable/quantization.html#prototype-maintenance-mode-fx-graph-mode-quantization" rel="noopener nofollow ugc">Quantization — PyTorch 2.6 documentation</a></p>
<p>But in PyTorch2 ，the api changes to <code>prepare_pt2e</code> and <code>prepare_qat_pt2e</code> , i think they are different , so how to get a pipeline to do this?</p>
<p>In my imagine, it’s like:</p>
<pre><code class="lang-auto">float_model(Python)                          Example Input
    \                                              /
     \                                            /
—-------------------------------------------------------
|                        export                        |
—-------------------------------------------------------
                            |
                    FX Graph in ATen     Backend Specific Quantizer
                            |                       /
—--------------------------------------------------------
|                     prepare_pt2e                      |
—--------------------------------------------------------
                            |
                     Calibrate
                            |
—--------------------------------------------------------
|                     convert back to Fx ？          |
—--------------------------------------------------------
                            |
                     FX Graph？
                            |
—--------------------------------------------------------
|                     prepare_qat_pt2e                      |
—--------------------------------------------------------
                            |
                     Calibrate
                            |
—--------------------------------------------------------
|                    convert_pt2e                       |
—--------------------------------------------------------
                            |
                    Quantized Model
                            |
—--------------------------------------------------------
|                       Lowering                        |
—--------------------------------------------------------
                            |
        Executorch, Inductor or &lt;Other Backends&gt;
</code></pre>
<p>thanks~</p>

jerryzh168: <p>for QAT, we will start from the original model, not the quantized model I think, here is the doc for QAT: <a href="https://pytorch.org/tutorials/prototype/pt2e_quant_qat.html" class="inline-onebox" rel="noopener nofollow ugc">(prototype) PyTorch 2 Export Quantization-Aware Training (QAT) — PyTorch Tutorials 2.7.0+cu126 documentation</a></p>