Is fuse_fx supposed to preceed convert_fx in the quant pipeline?

pylNeuralNet: <p>I am trying to incorporate fuse_fx into my quantization pipeline:</p>
<pre data-code-wrap="python"><code class="lang-python">    qconfig_mapping = QConfigMapping().set_global(qconfig)
    model_fx = symbolic_trace(original_model)
    fused_model_fx = fuse_fx(model_fx)
    prepared_model = prepare_fx(fused_model_fx, qconfig_mapping, example_inputs)
    calibrate_model(prepared_model, calibration_loader, calibration_batches, device)
    quantized_model = convert_fx(prepared_model)
</code></pre>
<p>However, it seems that convert_fx expects the model to be modified as well, not just it’s graph representation,</p>
<pre><code class="lang-auto">in _lower_dynamic_weighted_ref_module
    type(named_modules[str(n.target)]) not in \
KeyError: 'layer1.0.conv1.1'
</code></pre>
<p>(such layer exists when I do prepared_fused_model.graph.print_tabular() but doesn’t when I print model nodes)</p>
<p>I can manually synchronize model nodes with its graph after fuse_fx, but I don’t think this is the point of using pytorch graph mode.</p>
<p>The tutorial <a href="https://pytorch.org/tutorials/prototype/fx_graph_mode_ptq_static.html" class="inline-onebox" rel="noopener nofollow ugc">(prototype) FX Graph Mode Post Training Static Quantization — PyTorch Tutorials 2.3.0+cu121 documentation</a> is very ambigious - they perform fuse_fx after convert_fx, which is not what is supposed to be done in a normal pipeline and it kind of hides potential problems. I can do fuse_fx on the original model as well.</p>
<p>Further it is unclear to me whether prepare_fx actually does what is indicated in the tutorial in the comment:<br>
prepared_model = prepare_fx(float_model, qconfig_mapping, example_inputs)  # fuse modules and insert observers</p>
<p>It is indicated here that prepare_fx does fusion of modules - so that would mean adding fuse_fx is pointless. Is that correct?</p>
<p>I would greatly appreciate your help and guidance.</p>
<p>Details: I work with torch resnet18, I expect the following output after aforementioned pipeline (fuse_fx → prepare_fx → convert_fx):<br>
Quantized model:</p>
<pre><code class="lang-auto">GraphModule(
  (conv1): QuantizedConvReLU2d(3, 64, kernel_size=(7, 7), stride=(2, 2), scale=0.011580255813896656, zero_point=0, padding=(3, 3))
  (bn1): Identity()
  (relu): Identity()
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Module(
    (0): Module(
      (conv1): QuantizedConvReLU2d(64, 64, kernel_size=(3, 3), stride=(1, 1), scale=0.00825551524758339, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), scale=0.021675927564501762, zero_point=151, padding=(1, 1))
    )
    (1): Module(
      (conv1): QuantizedConvReLU2d(64, 64, kernel_size=(3, 3), stride=(1, 1), scale=0.007387497462332249, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), scale=0.030672406777739525, zero_point=164, padding=(1, 1))
    )
  )
  (layer2): Module(
    (0): Module(
      (conv1): QuantizedConvReLU2d(64, 128, kernel_size=(3, 3), stride=(2, 2), scale=0.007220075465738773, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), scale=0.021219059824943542, zero_point=113, padding=(1, 1))
      (downsample): Module(
        (0): QuantizedConv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), scale=0.016369296237826347, zero_point=131)
      )
    )
    (1): Module(
      (conv1): QuantizedConvReLU2d(128, 128, kernel_size=(3, 3), stride=(1, 1), scale=0.008431993424892426, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), scale=0.02427751012146473, zero_point=131, padding=(1, 1))
    )
  )
  (layer3): Module(
    (0): Module(
      (conv1): QuantizedConvReLU2d(128, 256, kernel_size=(3, 3), stride=(2, 2), scale=0.00851518101990223, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), scale=0.025926809757947922, zero_point=93, padding=(1, 1))
      (downsample): Module(
        (0): QuantizedConv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), scale=0.00756411487236619, zero_point=166)
      )
    )
    (1): Module(
      (conv1): QuantizedConvReLU2d(256, 256, kernel_size=(3, 3), stride=(1, 1), scale=0.008150133304297924, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), scale=0.02820182591676712, zero_point=164, padding=(1, 1))
    )
  )
  (layer4): Module(
    (0): Module(
      (conv1): QuantizedConvReLU2d(256, 512, kernel_size=(3, 3), stride=(2, 2), scale=0.006357187870889902, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), scale=0.02578684873878956, zero_point=135, padding=(1, 1))
      (downsample): Module(
        (0): QuantizedConv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), scale=0.019855372607707977, zero_point=124)
      )
    )
    (1): Module(
      (conv1): QuantizedConvReLU2d(512, 512, kernel_size=(3, 3), stride=(1, 1), scale=0.014804747886955738, zero_point=0, padding=(1, 1))
      (conv2): QuantizedConv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), scale=0.12023842334747314, zero_point=94, padding=(1, 1))
    )
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): QuantizedLinear(in_features=512, out_features=1000, scale=0.14396773278713226, zero_point=68, qscheme=torch.per_channel_affine)
)
</code></pre>
<p>or something like that. This was achieved with <code>torch.ao.quantization.fuse_modules</code> API preceeding <code>prepare_fx</code> and <code>convert_fx</code>, although the benchmark shows it doesn’t work very good for</p>
<pre data-code-wrap="python"><code class="lang-python">    modules_to_fuse = [
        ['conv1', 'bn1', 'relu']
        ]
</code></pre>
<p>therefore I try to understand <code>fuse_fx</code> API. Am I supposed to be doing what I’m doing, or should I already switch to  PT2 Export workflow? I was under the impression Eager Mode is v1 and Graph Mode is v2, now it seems both are legacy and PT2 is the modern one?</p>

jerryzh168: <p>prepare_fx calls fuse_fx by default. yeah fx api is in maintainance, it’s better to switch to pt2 export flow now</p>