Support Pytorch quantization toolkit based QAT sample using FX backend of Torch-TRT

[peri044]

Pytorch quantization toolkit support in FX backend

Goal(s)

Support Pytorch quantization toolkit based QAT sample using FX backend of Torch-TRT

Usecases

Proposed APIs / UX

The usage is similar to the sample here https://github.com/pytorch/TensorRT/blob/main/tests/py/qat/test_qat_trt_accuracy.py

Limitations

No known limitations at this time

Internal Implementation

Design

We need to write converters for quantize and dequantize ops in FX. Converters in TS : https://github.com/pytorch/TensorRT/blob/main/core/conversion/converters/impl/quantization.cpp

Extensions Required to Core API implementations

N/A

Data Structures

N/A

Details specific for TorchScript Support

N/A

Details specific for FX support

Write new converters for quantization ops in FX.

Implementation Phases

Prototype - S

Implement a converters for quantization ops. Once this is done, use the QAT training sample and get a trained nn.Module and compile it with FX backend.

MVP (<1.4.0>) - S

Implement a converters for quantization ops. Once this is done, use the QAT training sample and get a trained nn.Module and compile it with FX backend.

Both prototype and MVP would be the same for this feature. TS converters cannot be reused.

[peri044]

@narendasan to review

[narendasan]

Seems good to me if its just a converter