(prototype) FX Graph Mode Quantization User Guide

Author: Jerry Zhang

FX Graph Mode Quantization requires a symbolically traceable model. We use the FX framework to convert a symbolically traceable nn.Module instance to IR, and we operate on the IR to execute the quantization passes. Please post your question about symbolically tracing your model in PyTorch Discussion Forum

Quantization will only work on the symbolically traceable parts of your model. The data dependent control flow-if statements / for loops, and so on using symbolically traced values-are one common pattern which is not supported. If your model is not symbolically traceable end to end, you have a couple of options to enable FX Graph Mode Quantization only on a part of the model. You can use any combination of these options:

1. Non traceable code doesn’t need to be quantized

   1. Symbolically trace only the code that needs to be quantized

   2. Skip symbolic tracing the non-traceable code

2. Non traceable code needs to be quantized

   1. Refactor your code to make it symbolically traceable

   2. Write your own observed and quantized submodule

If the code that is not symbolically traceable does not need to be quantized, we have the following two options to run FX Graph Mode Quantization:

Symbolically trace only the code that needs to be quantized

When the whole model is not symbolically traceable but the submodule we want to quantize is symbolically traceable, we can run quantization only on that submodule.

before:

class M(nn.Module):
    def forward(self, x):
        x = non_traceable_code_1(x)
        x = traceable_code(x)
        x = non_traceable_code_2(x)
        return x

after:

class FP32Traceable(nn.Module):
    def forward(self, x):
        x = traceable_code(x)
        return x

class M(nn.Module):
    def __init__(self):
        self.traceable_submodule = FP32Traceable(...)
    def forward(self, x):
        x = self.traceable_code_1(x)
        # We'll only symbolic trace/quantize this submodule
        x = self.traceable_submodule(x)
        x = self.traceable_code_2(x)
        return x

quantization code:

qconfig_mapping = QConfigMapping().set_global(qconfig)
model_fp32.traceable_submodule = \
  prepare_fx(model_fp32.traceable_submodule, qconfig_mapping, example_inputs)

Note if original model needs to be preserved, you will have to copy it yourself before calling the quantization APIs.

Skip symbolically trace the non-traceable code

When we have some non-traceable code in the module, and this part of code doesn’t need to be quantized, we can factor out this part of the code into a submodule and skip symbolically trace that submodule.

before

class M(nn.Module):

    def forward(self, x):
        x = self.traceable_code_1(x)
        x = non_traceable_code(x)
        x = self.traceable_code_2(x)
        return x

after, non-traceable parts moved to a module and marked as a leaf

class FP32NonTraceable(nn.Module):

    def forward(self, x):
        x = non_traceable_code(x)
        return x

class M(nn.Module):

    def __init__(self):
        ...
        self.non_traceable_submodule = FP32NonTraceable(...)

    def forward(self, x):
        x = self.traceable_code_1(x)
        # we will configure the quantization call to not trace through
        # this submodule
        x = self.non_traceable_submodule(x)
        x = self.traceable_code_2(x)
        return x

quantization code:

qconfig_mapping = QConfigMapping.set_global(qconfig)

prepare_custom_config_dict = {
    # option 1
    "non_traceable_module_name": "non_traceable_submodule",
    # option 2
    "non_traceable_module_class": [MNonTraceable],
}
model_prepared = prepare_fx(
    model_fp32,
    qconfig_mapping,
    example_inputs,
    prepare_custom_config_dict=prepare_custom_config_dict,
)

If the code that is not symbolically traceable needs to be quantized, we have the following two options:

Refactor your code to make it symbolically traceable

If it is easy to refactor the code and make the code symbolically traceable, we can refactor the code and remove the use of non-traceable constructs in python.

More information about symbolic tracing support can be found in: (TODO: link)

before:

def transpose_for_scores(self, x):
    new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
    x = x.view(*new_x_shape)
    return x.permute(0, 2, 1, 3)

This is not symbolically traceable because in x.view(*new_x_shape) unpacking is not supported, however, it is easy to remove the unpacking since x.view also supports list input.

after:

def transpose_for_scores(self, x):
    new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
    x = x.view(new_x_shape)
    return x.permute(0, 2, 1, 3)

This can be combined with other approaches and the quantization code depends on the model.

Write your own observed and quantized submodule

If the non-traceable code can’t be refactored to be symbolically traceable, for example it has some loops that can’t be eliminated, like nn.LSTM, we’ll need to factor out the non-traceable code to a submodule (we call it CustomModule in fx graph mode quantization) and define the observed and quantized version of the submodule (in post training static quantization or quantization aware training for static quantization) or define the quantized version (in post training dynamic and weight only quantization)

before:

class M(nn.Module):

    def forward(self, x):
        x = traceable_code_1(x)
        x = non_traceable_code(x)
        x = traceable_code_1(x)
        return x

after:

1. Factor out non_traceable_code to FP32NonTraceable non-traceable logic, wrapped in a module

class FP32NonTraceable:
    ...

2. Define observed version of FP32NonTraceable

class ObservedNonTraceable:

    @classmethod
    def from_float(cls, ...):
        ...

3. Define statically quantized version of FP32NonTraceable and a class method “from_observed” to convert from ObservedNonTraceable to StaticQuantNonTraceable

class StaticQuantNonTraceable:

    @classmethod
    def from_observed(cls, ...):
        ...

# refactor parent class to call FP32NonTraceable
class M(nn.Module):

   def __init__(self):
        ...
        self.non_traceable_submodule = FP32NonTraceable(...)

    def forward(self, x):
        x = self.traceable_code_1(x)
        # this part will be quantized manually
        x = self.non_traceable_submodule(x)
        x = self.traceable_code_1(x)
        return x

quantization code:

# post training static quantization or
# quantization aware training (that produces a statically quantized module)v
prepare_custom_config_dict = {
    "float_to_observed_custom_module_class": {
        "static": {
            FP32NonTraceable: ObservedNonTraceable,
        }
    },
}

model_prepared = prepare_fx(
    model_fp32,
    qconfig_mapping,
    example_inputs,
    prepare_custom_config_dict=prepare_custom_config_dict)

calibrate / train (not shown)

convert_custom_config_dict = {
    "observed_to_quantized_custom_module_class": {
        "static": {
            ObservedNonTraceable: StaticQuantNonTraceable,
        }
    },
}
model_quantized = convert_fx(
    model_prepared,
    convert_custom_config_dict)

post training dynamic/weight only quantization in these two modes we don’t need to observe the original model, so we only need to define thee quantized model

class DynamicQuantNonTraceable: # or WeightOnlyQuantMNonTraceable
   ...
   @classmethod
   def from_observed(cls, ...):
       ...

   prepare_custom_config_dict = {
       "non_traceable_module_class": [
           FP32NonTraceable
       ]
   }

# The example is for post training quantization
model_fp32.eval()
model_prepared = prepare_fx(
    model_fp32,
    qconfig_mapping,
    example_inputs,
    prepare_custom_config_dict=prepare_custom_config_dict)

convert_custom_config_dict = {
    "observed_to_quantized_custom_module_class": {
        "dynamic": {
            FP32NonTraceable: DynamicQuantNonTraceable,
        }
    },
}
model_quantized = convert_fx(
    model_prepared,
    convert_custom_config_dict)

You can also find examples for custom modules in test test_custom_module_class in torch/test/quantization/test_quantize_fx.py.