PTQ Calibration in FX

[peri044]

PTQ Calibration in FX

Goal(s)

Support PTQ in FX. Both DataLoader Calibrator and Cache Calibrator options in Torchscript should be supported via FX as well.

Usecases

Proposed APIs / UX

The usage is similar to how users use PTQ with torchscript backend. There should be no difference.

Example Workflow

self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
            self.testing_dataloader,
            cache_file="./calibration.cache",
            use_cache=False,
            algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
            device=torch.device("cuda:0"),
        )

 trt_mod = torchtrt.compile(self.model, calibrator=self.calibrator)

Limitations

No known limitations at this time

Internal Implementation

Design

The internals of the feature are already present in ptq.py in the torch_tensorrt.py directory. We need to expose this support via calibrator keyword argument in FX API

Extensions Required to Core API implementations

N/A

Data Structures

N/A

Details specific for TorchScript Support

N/A

Details specific for FX support

See above.

Implementation Phases

Prototype - S

Implement a keyword argument and pass it to internal FX compile API. Implement a python example test in FX similar to how it is present in tests/py/ptq for TS backend

MVP (<1.4.0>) - S

Implement a keyword argument and pass it to internal FX compile API. Implement a python example test in FX similar to how it is present in tests/py/ptq for TS backend

Both prototype and MVP would be the same as this feature is not too involved. Most of the calibrator work for TS can be reused.

[peri044]

@narendasan to review