Error exporting to onnx qops due to bit overflow

[surajpandey353]

Hi,

I am trying to export QAT model, which has Addition / Multiplication operation with the QuantTensor. These operations result in increased bit_width to tackle the overflow, e.g. multiplying two 8-bit QuantTensor will give 16 bit QuantTensor. This raises an error during export_onnx_qops, which doesn't support bit_width > 8bits.

One solution I tried it to fallback to torch.tensor (QuantTensor.value) during these operation, which seems to work fine with QuantConv layers but not with the quant activation.

Error:

File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/init.py", line 32, in export_onnx_qop
return StdQOpONNXManager.export(*args, **kwargs)
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/onnx/manager.py", line 162, in export
return cls.export_onnx(
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/onnx/standard/manager.py", line 46, in export_onnx
output = super().export_onnx(
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/onnx/manager.py", line 121, in export_onnx
cls.set_export_mode(module, enabled=True)
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/onnx/standard/qoperator/manager.py", line 87, in set_export_mode
_set_layer_export_mode(module, enabled)
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/manager.py", line 115, in _set_layer_export_mode
m.export_mode = enabled
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1225, in setattr
object.setattr(self, name, value)
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/common.py", line 42, in export_mode
self.export_handler.prepare_for_export(self)
File "anaconda3/envs/e2enc/lib/python3.9/site-packages/brevitas/export/onnx/standard/qoperator/handler/act.py", line 51, in prepare_for_export
'input_axis': input_quant_symbolic_kwargs['input_axis']}
KeyError: 'input_axis'

I am feeding quant_input with QuantIdentity(return_quant_tensor=True. Thus I am keeping input_quant = None, return_quant_tensor=True in all the activation layers.

Dependencies:
Torch version 1.10.0+cu111
ONNX version 1.14.0
Brevitas version 0.9.1

[surajpandey353]

Hi,

Alternatively, if elementwise ops QuantEltwiseAdd , QuantCat is supported for q_ops export, this could be a solution.
QuantElementwiseOps