Repository: pytorch
Issue: [ONNX] Trying to export QAT model with torch.onnx.export, running into Problems that workarounds cant fix. Is ONNX export planned for QAT?
Labels: module: onnx,triaged

### 🐛 Describe the bug

The bug occurs when trying to export a model fully converted model to onnx. I have tried the workarounds suggested in (https://pytorch.org/docs/main/quantization.html#frequently-asked-questions) but even after trying all of them and debugging further I end up at various torch.onnx.errors.SymbolicValueError, which I cannot really explain due to the nature of these workarounds feeling like patchwork.

Imports: 
```
import torch.onnx
import torch
import torch.nn as nn 
from torchinfo import summary
import numpy as np
import onnx
from torch.ao.quantization import fuse_modules, QuantStub, DeQuantStub
from torch.ao.quantization import (
    get_default_qat_qconfig,
    prepare_qat,
    convert
)
```


My Code:
```
class Conv1d(nn.Conv1d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.causal_padding = self.dilation[0] * (self.kernel_size[0] - 1)
        
        self.conv1d = nn.Conv1d(
            in_channels=self.in_channels,
            out_channels=self.out_channels,
            kernel_size=self.kernel_size[0],
            stride=self.stride[0],
            dilation=self.dilation[0],
            padding= 0,
            padding_mode='zeros' 
        )
    def forward(self, x):
        x = nn.functional.pad(x, (self.causal_padding//2, self.causal_padding-self.causal_padding//2), 'constant')
        return self.conv1d(x)

class ResidualUnit(nn.Module):
    def __init__(self, in_channels, out_channels, dilation):
        super().__init__()
        
        self.dilation = dilation
        self.layers = nn.Sequential(
            Conv1d(in_channels=in_channels, out_channels=out_channels,
                      kernel_size=5, dilation=dilation),
            nn.LeakyReLU(),
            nn.Conv1d(in_channels=in_channels, out_channels=out_channels,
                      kernel_size=1),
        )

    def forward(self, x):
        
        return x + self.layers(x)
class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride):
        super().__init__()

        self.layers = nn.Sequential(
            ResidualUnit(in_channels=int(in_channels),
                         out_channels=int(in_channels), dilation=3),
            nn.LeakyReLU(),
            Conv1d(in_channels=int(in_channels), out_channels=int(out_channels), 
            kernel_size=2*stride, stride=stride),
        )

    def forward(self, x):
        return self.layers(x)
class Encoder(nn.Module):
    def __init__(self, C=36, D=9):
        super().__init__()
        self.layers = nn.Sequential(
            Conv1d(in_channels=36, out_channels=C, kernel_size=7),
            nn.LeakyReLU(),
            EncoderBlock(in_channels=C ,out_channels=24, stride=1),
            nn.LeakyReLU(),
            EncoderBlock(in_channels=24 ,out_channels=20, stride=2),
            nn.LeakyReLU(),
            EncoderBlock(in_channels=20 ,out_channels=16, stride=2),
            nn.LeakyReLU(),
            EncoderBlock(in_channels=16 ,out_channels=12, stride=2),
            nn.LeakyReLU(),
            Conv1d(in_channels=12, out_channels=D, kernel_size=1), 
        )
    def forward(self, x):
        x = self.layers(x)
        return x
        
if __name__ == "__main__":
      model = Encoder()
      model.train()
      model.qconfig = get_default_qat_qconfig("x86")
      qat_model = prepare_qat(model)
      qat_model_done = convert(qat_model)
      dummy_input = torch.rand(1, 36, 800)
      onnx_file_path = "./qat_export_test.onnx"
      onnx_model = torch.onnx.export(
                                    qat_model_done, 
                                    dummy_input, 
                                    onnx_file_path,
                                    export_params=True
                                 )
```

Error:
```
Traceback (most recent call last):
  File "/media/data/xxx/xxx/prev_code/xxx/data_compression/data_compression/examples/idk.py", line 115, in <module>
    onnx_model = torch.onnx.export(
                 ^^^^^^^^^^^^^^^^^^
  File /media/data/xxx/xxx/prev_code/xxx//prev_code/moritz_martinolli_397/venv/lib/python3.11/site-packages/torch/onnx/__init__.py", line 375, in export
    export(
  File "/media/data/xxx/xxx/prev_code/xxx/prev_code/moritz_martinolli_397/venv/lib/python3.11/site-packages/torch/onnx/utils.py", line 502, in export
    _export(
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/utils.py", line 1564, in _export
    graph, params_dict, torch_out = _model_to_graph(
                                    ^^^^^^^^^^^^^^^^
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/utils.py", line 1117, in _model_to_graph
    graph = _optimize_graph(
            ^^^^^^^^^^^^^^^^
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/utils.py", line 639, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/utils.py", line 1836, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset10.py", line 747, in dequantize
    return symbolic_helper.dequantize_helper(g, input)[0]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/symbolic_helper.py", line 1525, in dequantize_helper
    unpacked_qtensors = _unpack_quantized_tensor(qtensor)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/data/xxx/xxx/prev_code/xxx/venv/lib/python3.11/site-packages/torch/onnx/symbolic_helper.py", line 189, in _unpack_quantized_tensor
    raise errors.SymbolicValueError(
torch.onnx.errors.SymbolicValueError: ONNX symbolic expected the output of `%x.1 : Float(1, 36, 800, strides=[28800, 800, 1], requires_grad=0, device=cpu) = prim::Param()
` to be a quantized tensor. Is this likely due to missing support for quantized `prim::Param`. Please create an issue on https://github.com/pytorch/pytorch/issues  [Caused by the value 'x.1 defined in (%x.1 : Float(1, 36, 800, strides=[28800, 800, 1], requires_grad=0, device=cpu) = prim::Param()
)' (type 'Tensor') in the TorchScript graph. The containing node has kind 'prim::Param'.] 

    Inputs:
        Empty
    Outputs:
        #0: x.1 defined in (%x.1 : Float(1, 36, 800, strides=[28800, 800, 1], requires_grad=0, device=cpu) = prim::Param()
    )  (type 'Tensor')
```

This error arrose while I was doing the whole QuantStub and DeQuantstub workarounds suggested in the FAQ of QAT. 

If no one has time to actually recreate this and debug it, please just tell me if full QAT to ONNX support is planned or not (I've seen contributors saying it is not planned, I've seen posts where there is said that it should be in torch 2.0, so maybe just a clear answer would be good), so I know if I can just disregard this until it eventually becomes an option.

### Versions

```
Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] nvidia-cublas-cu11==11.10.3.66
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu11==11.7.101
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu11==11.7.99
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu11==11.7.99
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu11==8.5.0.96
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu11==10.9.0.58
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu11==10.2.10.91
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu11==11.4.0.1
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu11==11.7.4.91
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu11==2.14.3
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu11==11.7.91
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] onnx==1.17.0
[pip3] onnxruntime==1.20.1
[pip3] onnxscript==0.1.0.dev20241217
[pip3] torch==2.5.1
[pip3] torchaudio==2.5.1
[pip3] torchinfo==1.8.0
[pip3] torchvision==0.20.1
[pip3] triton==3.1.0
```