QuantizationModifier refactor

[bfineran]

Main feature branch for the QuantizationModifier refactor project. Goal of the project is to create a new interface that provides more flexibility (via the QuantizationScheme config) while including default targets that protect more against future torch default changes.

UX

Users have full control over what gets quantized through a default scheme, overriding defaults for module types or submodules with scheme_overrides, and skipping module types or submodules with ignore

sample config:

!QuantizationModifier
    start_epoch: 0.0
    scheme:
        input_activations:
            num_bits: 8
            symmetric: False
        weights:
            num_bits: 8
            symmetric: True
    scheme_overrides:
        feature_extractor: "default"
        classifier:
            input_activations:
                num_bits: 8
                symmetric: False
            weights: null
        Conv2d:
            input_activations:
                num_bits: 8
                symmetric: True
    ignore: ["ReLU", "input"]
    disable_quantization_observer_epoch: 2.0
    freeze_bn_stats_epoch: 3.0
    model_fuse_fn_name: 'fuse_module'
    strict: True

Tasks

• [QuantizationModifier] refactor base - move deprecated code to legacy file, add object routing for yaml load #1059
• [QuantizationModifier] pydantic classes for defining quantization schemes to generate QConfigs #1061
• [QuantizationModifier] base refactor flow - quantize entire module from QuantizationScheme #1185
• [QuantizationModifier] improved quantization flow - control of propagation with schemes and stronger testing #1198
• [QuantizationModifier] exclude_module_types implementation #1199
• [QuantizationModifier] submodule_schemes property impl - target specific submodules by scheme #1201
• [QuantizationModifier] module_type_schemes - override quantization scheme by layer type #1202
• [QuantizationModifier] target hardware support #1203
• [QuantizationModifier] freeze bn stats and disable observers for QAT finetuning support #1206
• [QuantizationModifier] num_calibration_steps support (PTQ) #1208
• [QuantizationModifier] override params for model fuse step #1209
• [QuantizationModifier] refactor QuantizationScheme to its own file #1223
• [QuantizationModifier] QATWrapper support #1226
• [QuantizationModifier] logging support #1231
• [QuantizationModifier] potentially re-load quantization schemes on qconfig load #1236
• [QuantizationModifier] UX refactor - submodule_overrides and ignore #1239
• [QuantizationModifier] strict mode - raise if unmatched submodules or types #1241
• [QuantizationModifier] take ownership of add_observers_, unit test fixes #1261

Sample section of ResNet-50 quantized with new default flow

Test Plan

• Testing included with each feature PR, lifecycle tests cover these features
• yaml testing included
• planned to work with research team on generating practical recipes with the updated modifier
• sample resnet50 onnx export in description

[KSGulin]

Nice!!

[corey-nm]

🚀 🎉