• Tutorials >
  • (prototype) FX Graph Mode Quantization User Guide
Shortcuts

(prototype) FX Graph Mode Quantization User Guide

Author: Jerry Zhang

FX Graph Mode Quantization requires a symbolically traceable model. We use the FX framework (TODO: link) 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. Data dependent control flow (if statements / for loops etc 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

1.a. 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:

after:

quantization code:

qconfig_dict = {"": qconfig}
model_fp32.traceable_submodule = \
  prepare_fx(model_fp32.traceable_submodule, qconfig_dict)

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

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_dict = {"": 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_dict,
    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:

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)

quantization code:

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

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:
    ...
  1. 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_dict,
    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_dict,
    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.


이 튜토리얼이 어떠셨나요?

© Copyright 2022, PyTorch & 파이토치 한국 사용자 모임(PyTorch Korea User Group).

Built with Sphinx using a theme provided by Read the Docs.

PyTorchKorea @ GitHub

파이토치 한국 사용자 모임을 GitHub에서 만나보세요.

GitHub로 이동

한국어 튜토리얼

한국어로 번역 중인 PyTorch 튜토리얼입니다.

튜토리얼로 이동

커뮤니티

다른 사용자들과 의견을 나누고, 도와주세요!

커뮤니티로 이동