Failed to prepare model for QAT in Pt2e

[jimmy133719]

Description of the bug:

I met the same issue as #355 , but still can't find the way to solve it, but I didn't use SiLU or GELU.

Environment:

• torch==2.5.1+cu124
• ai-edge-torch==0.4.0

import torch, numpy as np
from torch import nn
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import prepare_qat_pt2e, convert_pt2e

import ai_edge_torch
from ai_edge_torch.quantize import pt2e_quantizer, quant_config

class SmallCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, 1, 1)
        self.bn   = nn.BatchNorm2d(16)
        self.act  = nn.ReLU(inplace=False)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc   = nn.Linear(16, num_classes)

    def forward(self, x):
        x = self.act(self.bn(self.conv(x)))
        x = self.pool(x)
        x = torch.flatten(x, 1)
        return self.fc(x)

model = SmallCNN().eval()
example_inputs = (torch.randn(1, 3, 224, 224),)

exported = capture_pre_autograd_graph(model, example_inputs)

quantizer = pt2e_quantizer.PT2EQuantizer().set_global(
    pt2e_quantizer.get_symmetric_quantization_config(
        is_per_channel=True,  
        is_dynamic=False,     
        is_qat=True           
    )
)

qat_model = prepare_qat_pt2e(exported, quantizer)  

Actual vs expected behavior:

Traceback (most recent call last):
  File "/media/ipcam/RTK-IPCAM/jimmy/MediaPipePyTorch/qat_example.py", line 38, in <module>
    qat_model = prepare_qat_pt2e(exported, quantizer)
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/torch/ao/quantization/quantize_pt2e.py", line 169, in prepare_qat_pt2e
    quantizer.annotate(model)
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/ai_edge_torch/quantize/pt2e_quantizer.py", line 384, in annotate
    model = self._annotate_for_static_quantization_config(model)
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/ai_edge_torch/quantize/pt2e_quantizer.py", line 432, in _annotate_for_static_quantization_config
    self._annotate_all_static_patterns(
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/ai_edge_torch/quantize/pt2e_quantizer.py", line 399, in _annotate_all_static_patterns
    OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn)
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/ai_edge_torch/quantize/pt2e_quantizer_utils.py", line 461, in _annotate_conv_bn_relu
    return _do_annotate_conv_bn(gm, quantization_config, filter_fn, has_relu=True)
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/ai_edge_torch/quantize/pt2e_quantizer_utils.py", line 517, in _do_annotate_conv_bn
    pattern = _get_aten_graph_module_for_pattern(
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/torch/ao/quantization/pt2e/utils.py", line 375, in _get_aten_graph_module_for_pattern
    aten_pattern = capture_pre_autograd_graph(
  File "/home/ipcam/anaconda3/envs/ai_edge_torch/lib/python3.10/site-packages/torch/_export/__init__.py", line 123, in capture_pre_autograd_graph
    assert isinstance(f, torch.nn.Module), "Expected an nn.Module instance."
AssertionError: Expected an nn.Module instance.

Any other information you'd like to share?

No response

[gaikwadrahul8]

Hi @jimmy133719, Thank you for your patience and Apologies for the delay. I have successfully reproduced the issue on my end and can confirm the behavior you reported. Our team is now investigating this further to identify the root cause. I will post an update on this thread as soon as more information is available. For reference, here is the Colab Gist I used for replication: Gist File Link.

We appreciate you bringing this to our attention.

Here is error log for reference :

/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py:335: UserWarning: Device capability of jax unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.
  warnings.warn(
W0924 18:32:32.077000 9706 torch/_export/__init__.py:64] +============================+
W0924 18:32:32.079000 9706 torch/_export/__init__.py:65] |     !!!   WARNING   !!!    |
W0924 18:32:32.081000 9706 torch/_export/__init__.py:66] +============================+
W0924 18:32:32.083000 9706 torch/_export/__init__.py:67] capture_pre_autograd_graph() is deprecated and doesn't provide any function guarantee moving forward.
W0924 18:32:32.085000 9706 torch/_export/__init__.py:68] Please switch to use torch.export.export_for_training instead.
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
[/tmp/ipython-input-2628353422.py](https://localhost:8080/#) in <cell line: 0>()
     35 )
     36 
---> 37 qat_model = prepare_qat_pt2e(exported, quantizer)

7 frames
[/usr/local/lib/python3.12/dist-packages/torch/_export/__init__.py](https://localhost:8080/#) in capture_pre_autograd_graph(f, args, kwargs, dynamic_shapes)
    121         raise RuntimeError("capture_pre_autograd_graph not yet supported on Windows")
    122 
--> 123     assert isinstance(f, torch.nn.Module), "Expected an nn.Module instance."
    124 
    125     if kwargs is None:

AssertionError: Expected an nn.Module instance.