Prepare Description

Definition of Prepare

Prepare is the process of converting a floating-point model into a pseudo-quantized model. This process involves several key steps:

Computation Graph Capture: Execute the complete forward logic of the model and capture the corresponding computation graph during this process. Subsequently, the computation graph will be pruned based on the positions of QuantStub and DequantStub, retaining only the intermediate parts that require quantization.
Function-type Operator Replacement: Some torch function-type operators (e.g., torch.reciprocal) need to have fake quantization nodes inserted during quantization. Therefore, these operators must be replaced with their corresponding Module-type implementations (e.g., horizon_plugin_pytorch.nn.Reciprocal), so that fake quantization nodes are placed inside this Module. The model remains equivalent before and after replacement. This replacement process is performed based on the computation graph, so operations not in the graph will not be replaced.
Operator Fusion: BPU supports fusing specific computational patterns, where the intermediate results of fused operators are represented with high precision. Therefore, we replace multiple operators to be fused with a single Module to prevent quantizing the intermediate results. The model before and after fusion is also equivalent.
Operator Conversion: Floating-point operators are replaced with QAT (Quantized Awareness Training) operators. According to the configured qconfig, QAT operators will add FakeQuantize nodes at the input/output/weights.

Attention

To keep converted QAT operators work as expected, please ensure that no further modifications are made to the model after calling prepare. For example, converting an unfused BN to a sync BN after prepare may cause the QAT BN to be modified again. The conversion to sync BN should be done before calling prepare.

Model Structure Check: The QAT model is checked, and a check result file is generated.

Note

Principle of Function Operator Replacement

Function operator replacement relies on the computation graph. However, multiple calls to the same line of code in the graph are expanded into multiple nodes and replaced with multiple Modules. If the model contains code blocks that are called different numbers of times during training and inference, it will cause scale misalignment and affect model accuracy.

To solve this problem, we introduce the concept of "Scope". During operator replacement, multiple calls to the same func code within a single Scope will be replaced with the same shared Module. Currently, there are two ways to define a Scope: 1. The forward method of a Module type defines one Scope; 2. A code block marked with the horizon_plugin_pytorch.fx.jit_scheme.dynamic_block interface is one Scope. Examples are as follows:

Multiple calls within one Scope are replaced with different Modules

def sub_one(x):
    return x - 1

class Net(torch.nn.Module):
    def forward(self, x):
        for _ in range(2):
            x += 1  # self._generated_add_0, self._generated_add_1
            x = sub_one(x)  # self._generated_sub_1, self._generated_sub_1
        return x

# The loop is expanded in the graph, equivalent to:

class _Net(torch.nn.Module):
    def forward(self, x):
        x += 1  # self._generated_add_0
        x = sub_one(x)  # self._generated_sub_0
        x += 1  # self._generated_add_1
        x = sub_one(x)  # self._generated_sub_1
        return x

Scope defined via Module

def sub_one(x):
    return x - 1

class Block(torch.nn.Module):
    def forward(self, x):
        x += 1  # self._generated_add_0
        x = sub_one(x)  # self._generated_sub_0
        return x

class Net(torch.nn.Module):
    def __init__(self):
        self.block = Block()

    def forward(self, x):
        x = self.block(x)  # self.block._generated_add_0, self.block._generated_sub_0
        x = self.block(x)  # self.block._generated_add_0, self.block._generated_sub_0
        return x

Scope defined via dynamic_block

def sub_one(x):
    return x - 1

class Net(torch.nn.Module):
    def forward(self, x):
        for _ in range(numpy.random.randint(1, 10)):
            with horizon_plugin_pytorch.fx.jit_scheme.dynamic_block():
                x += 1  # self._generated_add_0
                x = sub_one(x)  # self._generated_sub_1
                x += 1  # self._generated_add_1
        return x

Note: To enable the conversion from function to Module, we implement a special Wrapper subclass for torch.Tensor. The wrapping and unwrapping of this Wrapper are implemented by registering hooks on appropriate Modules in the model. Therefore, do not modify any hooks in the model after prepare is completed, to avoid invalidation of the replacement and consequent errors or accuracy issues.

Note

Accuracy Description

The model after prepare differs from the floating-point model in terms of computational logic in the following aspects:

Pseudo-quantization nodes are added to the model. For a very small number of operators (such as reciprocal) whose outputs may have extremely large values, to adapt to quantization, their outputs will be clipped to a reasonable range by default.

The above operations will cause changes in the numerical values of the model.

The usage of the prepare interface is as follows:

from horizon_plugin_pytorch.quantization import get_qconfig, prepare, PrepareMethod
from horizon_plugin_pytorch.quantization.qconfig_setter import *

# When using templates, example_inputs and qconfig_setter must be provided.
# When method is PrepareMethod.JIT_STRIP, example_inputs must be provided.
# def prepare(
#     model: torch.nn.Module,
#     example_inputs: Any = None,   # used to get model's graph structure, ensuring it can be used to run forward.
#     qconfig_setter: Optional[Union[Tuple[QconfigSetterBase, ...], QconfigSetterBase]] = None,  # qconfig template, supports multiple templates, priority from high to low.
#     method: PrepareMethod = PrepareMethod.JIT_STRIP,  # prepare method
# ) -> torch.nn.Module:

qat_model = prepare(
    float_model,
    example_inputs=example_inputs,
    qconfig_setter=QconfigSetter(
        get_qconfig(),
        templates=[
            ModuleNameTemplate({"": qint8}),
            ConvDtypeTemplate(),
            SensitivityTemplate(table, ratio=0.2),
        ],
    ),
)

PrepareMethod

The prepare method includes JIT_STRIP and EAGER. JIT_STRIP belongs to Graph Mode, while EAGER belongs to PrepareMethod.EAGER. Their comparison is as follows:

method	Principle	Advantages	Disadvantages
PrepareMethod.JIT_STRIP(Graph Mode)	Use hooks and subclass tensor to get the graph structure, performing operator replacement/operator fusion on the original forward.	Fully automatic, minimal code modification, hides many detail issues, easy to debug.	Dynamic code blocks need special handling.
PrepareMethod.EAGER(Eager Mode)	Does not sense the graph structure. operator replacement/operator fusion needs to be done manually.	Flexible usage, controllable process, easy to debug and handle various special needs.	Requires more manual operations, more code modifications, high learning cost.

Currently, JIT_STRIP is our recommended method. The JIT_STRIP will identify and skip pre-process and post-process based on the positions of QuantStub and DequantStub in the model.

Example

import copy

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.quantization import DeQuantStub, QuantStub

from horizon_plugin_pytorch import March, set_march
from horizon_plugin_pytorch.fx.jit_scheme import Tracer
from horizon_plugin_pytorch.quantization import (
    FakeQuantState,
    get_qconfig,
    PrepareMethod,
    prepare,
    set_fake_quantize,
)

class Net(torch.nn.Module):
    def __init__(self, input_size, class_num) -> None:
        super().__init__()
        self.quant0 = QuantStub()
        self.quant1 = QuantStub()
        self.dequant = DeQuantStub()
        self.conv = nn.Conv2d(3, 3, 1)
        self.bn = nn.BatchNorm2d(3)
        self.classifier = nn.Conv2d(3, class_num, input_size)
        self.loss = nn.CrossEntropyLoss()

    def forward(self, input, other, target=None):
        # Preprocess that does not need quantization. Use JIT_STRIP to exclude these operations from the computational graph.
        input = (input - 128) / 128.0

        x = self.quant0(input)
        y = self.quant1(other)

        n = np.random.randint(1, 5)
        m = np.random.randint(1, 5)

        # Since the python code is not regenerated, this dynamic loop is retained in the QAT model.
        for _ in range(n):
            for _ in range(m):
                # Dynamic code blocks involving operator replacement or fusion must be marked.
                # The marked part refers to the logic that requires operator replacement or fusion, rather than the for loop.
                with Tracer.dynamic_block(self, "ConvBnAdd"):
                    x = self.conv(x)
                    x = self.bn(x)
                    x = x + y

        x = self.classifier(x).squeeze()

        # Since the python code is not regenerated, this dynamic control flow is retained in the QAT model
        if self.training:
            assert target is not None
            x = self.dequant(x)
            return F.cross_entropy(torch.softmax(x, dim=1), target)
        else:
            return torch.argmax(x, dim=1)

model = Net(6, 2)
train_example_input = (
    torch.rand(2, 3, 6, 6) * 256,
    torch.rand(2, 3, 6, 6),
    torch.tensor([[0.0, 1.0], [1.0, 0.0]]),
)
eval_example_input = train_example_input[:2]

model.eval()
set_march("nash-e")

model.qconfig = get_qconfig()
qat_model = prepare(
    model,
    example_inputs=copy.deepcopy(eval_example_input),
    method=PrepareMethod.JIT_STRIP,
)

qat_model.graph.print_tabular()

# opcode         name              target                                                     args                              kwargs
# -------------  ----------------  ---------------------------------------------------------  --------------------------------  ----------
# placeholder    input_0           input_0                                                    ()                                {}
# call_module    quant0            quant0                                                     (input_0,)                        {}
# placeholder    input_1           input_1                                                    ()                                {}
# call_module    quant1            quant1                                                     (input_1,)                        {}
# call_module    conv              conv                                                       (quant0,)                         {}
# call_module    bn                bn                                                         (conv,)                           {}
# get_attr       _generated_add_0  _generated_add_0                                           ()                                {}
# call_method    add_2             add                                                        (_generated_add_0, bn, quant1)    {}
# scope_end is automatically inserted during the trace process to mark the boundaries of sub-modules or dynamic code blocks, not corresponding to any calculations
# call_function  scope_end         <function Tracer.scope_end at 0x7f65d90e5e50>              ('_dynamic_block_ConvBnAdd',)     {}
# call_module    conv_1            conv                                                       (add_2,)                          {}
# call_module    bn_1              bn                                                         (conv_1,)                         {}
# get_attr       _generated_add_1  _generated_add_0                                           ()                                {}
# call_method    add_3             add                                                        (_generated_add_1, bn_1, quant1)  {}
# call_function  scope_end_1       <function Tracer.scope_end at 0x7f65d90e5e50>              ('_dynamic_block_ConvBnAdd',)     {}
# call_module    classifier        classifier                                                 (add_3,)                          {}
# call_function  squeeze           <method 'squeeze' of 'torch._C._TensorBase' objects>       (classifier,)                     {}
# call_function  argmax            <built-in method argmax of type object at 0x7f66f04cf820>  (squeeze,)                        {'dim': 1}
# call_function  scope_end_2       <function Tracer.scope_end at 0x7f65d90e5e50>              ('',)                             {}
# output         output            output                                                     ((argmax,),)                      {}

print(qat_model)

# GraphModuleImpl(
#   (quant0): QuantStub(
#     (activation_post_process): FakeQuantize(
#       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=qint8, qscheme=torch.per_tensor_symmetric, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
#       (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
#     )
#   )
#   (quant1): QuantStub(
#     (activation_post_process): FakeQuantize(
#       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=qint8, qscheme=torch.per_tensor_symmetric, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
#       (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
#     )
#   )
#   (dequant): DeQuantStub()
#   (conv): Identity()  # Since the forward code remains unchanged, conv and bn will still be executed, so after fusion, Conv and Bn must be replaced with Identity
#   (bn): Identity()
#   (classifier): Conv2d(
#     3, 2, kernel_size=(6, 6), stride=(1, 1)
#     (activation_post_process): FakeQuantize(
#       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=qint8, qscheme=torch.per_tensor_symmetric, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
#       (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
#     )
#     (weight_fake_quant): FakeQuantize(
#       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=qint8, qscheme=torch.per_channel_symmetric, ch_axis=0,         scale=tensor([1., 1.]), zero_point=tensor([0, 0])
#       (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
#     )
#   )
#   (loss): CrossEntropyLoss()
#   (_generated_add_0): ConvAdd2d(  # Automatically replace '+' with Module form, and fuse Conv and Bn into it
#     3, 3, kernel_size=(1, 1), stride=(1, 1)
#     (activation_post_process): FakeQuantize(
#       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=qint8, qscheme=torch.per_tensor_symmetric, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
#       (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
#     )
#     (weight_fake_quant): FakeQuantize(
#       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=qint8, qscheme=torch.per_channel_symmetric, ch_axis=0,         scale=tensor([1., 1., 1.]), zero_point=tensor([0, 0, 0])
#       (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
#     )
#   )
# )

qat_model.train()
set_fake_quantize(qat_model, FakeQuantState.QAT)

for _ in range(3):
    ret = qat_model(*train_example_input)
    ret.backward()

Attention

When dynamic code blocks involve operator replacement or fusion, they must be marked with Tracer.dynamic_block. Otherwise, it will lead to quantization information confusion or forward errors.
Parts of the model where the call count changes (sub-modules or dynamic blocks), if only executed once during the trace, may get fused with non-dynamic parts, leading to forward errors.

When using the graph-based prepare method, it's recommended that the model only contains deployment logic.

# The execution path is controlled by `training` or other flags, which may not match the final deployment state.
# This can easily lead to misplaced quant and dequant nodes.
def forward(self, input, gt):
    conv_out = self.conv(input)
    if self.training:
        return self.loss(conv_out, gt)
    else:
        return self.sigmoid(conv_out)

# Extract deployment logic and only apply prepare to this part.
def forward_infer(self, input, gt):
    conv_out = self.conv(input)
    return self.sigmoid(conv_out), conv_out

# Non-deployment logic is handled externally and excluded from prepare.
def forward(self, input, gt):
    sig_out, conv_out = self.forward_infer(input, gt)
    if self.training:
        return self.loss(conv_out, gt)
    else:
        return sig_out

If it’s not feasible to extract clean deployment logic due to readability or maintainability concerns, then the following checks should be performed:

Check for missing or unexpected keys when loading the checkpoint. Missing or unexpected quantization parameters may indicate that some deployment logic was not quantized or some non-deployment logic was incorrectly quantized. Missing or unexpected model parameters may indicate a mismatch between the forward logic used during prepare and the one used when the checkpoint was generated.
Check whether the FX graph generated by multiple prepare runs is consistent. Inconsistent FX graphs suggest that the forward logic differs across runs, and it should be verified whether this is intended.

Model Check

When example_inputs is provided, prepare will perform a model structure check by default. If the check completes, a model_check_result.txt file can be found in the running directory. If the check fails, you need to modify the model based on the warning prompts or call horizon_plugin_pytorch.utils.check_model.check_qat_model separately to check the model. The check process is the same as check_qat_model in the debug tool, and the analysis of the result file is detailed in the check_qat_model related documentation.

#Prepare Description

#Definition of Prepare

#PrepareMethod

#Example

#Model Check

Prepare Description

Definition of Prepare

PrepareMethod

Example

Model Check