Plugin-based Quantization

Overview

Plugin-based quantization is model-agnostic, providing support for language and vision models. It enables fast quantization for Transformer models.

Overall Process

There are two quantization parameter optimization strategies: post-training quantization (PTQ) and quantization aware training (QAT). PTQ completes quantization calibration using the inference project, while QAT requires integration with the training project.

API Usage Description

import os
import sys
sys.path.append('path/to/dopt_torch_py3')
####  Import from API.
from dopt.dopt_lm.do_opt import (
    generate_config_file,   ## Generate a quantization configuration file.
    optimize_model,         ## Convert the floating-point nn.module into the nn.module with quantization operators inserted using the configured quantization configuration file.
    set_quant_state,        ## Enable quantization inference separately for activations and weights.
    set_calibrate_state,    ## Set the updatable states of quantization parameters.
    generate_quant_params,  ## Export quantization parameters.
)
#### Insert quantization operators. A quantization configuration file is generated when this API is called for the first time. Once the file is configured, execute PTQ or QAT.
def get_quant_model(model, dopt_config):
    ## **model**: nn.module instance object, which is of the floating-point type.
    ## **dopt_config**: quantization configuration file path. The value is **path/to/config.json**.
    if not os.path.exists(dopt_config):
        generate_config_file(model, dopt_config)
        exit()
    model = optimize_model(model, dopt_config)
    return model

Quantization strategy configuration:

{
    "layer_strategy": {
        "node_name": {
            "type": "<class 'torch.nn.modules.linear.Linear'>",
            "quant_strategy": "need to set",
            "weight": {
                "bit": 4,
                "group_size": 128,
                "weight_algo": "group_min_max",
            },
            "input" : {
                "bit" : 8,
                "input_algo":  "min_max",
                "unsigned_quant": True
            }
        }
    }
}

Configuration description:

Each strategy (specified by quant_strategy) has its own default settings. Configure the strategies one by one. Table 1 describes the parameters.

Table 1 Configuration parameter description

Parameter	Configuration Description	Option	Option Description
quant_strategy	N/A	Quant_aigc_ptq Quant_aigc_qat Quant_act_weight_eco Quant_lm_head Quant_Embed_MinMax	PTQ quantization strategy. The default value is U8S8. The quantizer is a minmax. QAT quantization strategy. The default value is U8S8. The quantizer is omni_minmax. 16-4 group-linear quantization strategy. The default group_size is set to 128. This strategy is for LLMs only. The QAT quantizer is used. 16-4 group-linear quantization strategy. The default group_size is set to 128. This strategy is for LLMs only. The QAT quantizer is used. 8-bit embedding quantization.
weight	bit	4, 8	Weight quantization bit width.
weight	group_size	128 , 256	With a weight bit of 4, the model is considered low-bit quantization. In this case, set this key-value pair only when grouplinear quantization is used.
weight	weight_algo	min_max group_min_max omni_min_max omni_group_min_max	Per-channel quantization. The min-max statistics are used for PTQ. Per-group quantization. The min-max statistics are used for PTQ. QAT quantizer, which uses the per-channel quantization granularity. QAT quantizer, which uses the per-group quantization granularity.
input	bit	8, 16	Activation quantization bit width.
input	input_algo	min_max ema omni_min_max	Per-tensor quantization. The min-max statistics are used for PTQ, resulting in minimal clipping error. Per-tensor quantization. The moving average statistics are used for PTQ, resulting in minimal rounding error. Per-tensor quantization. The QAT quantizer is used.
input	unsigned_quant	True False	Asymmetric quantization. Symmetric quantization:

PTQ for Quantization Parameter Optimization

In the generated quant_config, the float strategy is used for each layer. Change the strategy to Quant_aigc_ptq. See the preceding table for advanced parameter settings.

Use the inference project and select a proper amount of data to perform forward inference on the model returned by the get_quant_model API. This process optimizes the model's quantization. The sample code is provided below:

## A configuration file is generated upon the first call. You need to manually configure the quantization strategy in the file.
model = get_quant_model(model, quant_config)
model.eval()
## Enable the quantizer.
set_quant_state(model, weight_state=True, input_state=True)
## Enable the quantization parameter calibration status.
set_calibrate_state(model, True)
## Use real inference data to perform quantization calibration.
for data in datasets:
    model(data)
set_calibrate_state(model, False)
torch.save(model.state_dict(), 'pth_save_path')

QAT for Quantization Parameter Optimization

In the generated quant_config, the float strategy is used for each layer. Change the strategy to Quant_aigc_qat. See the preceding table for advanced parameter settings.

Use the training project to train the model with quantization operators.

model = get_quant_model(model, quant_config)
model.train()
set_quant_state(model, weight_state=True, input_state=True)
set_calibrate_state(model, True)
"""
do training
"""
set_calibrate_state(model, False)
torch.save(model.state_dict(), 'pth_save_path')

Quantization Parameter Extraction and Export

from dopt.dopt_lm.do_opt import generate_quant_params, optimize_model
## **model**: floating-point nn.module instance object; **quant_config**: quantization configuration file.
model = optimize_model(model, quant_config)
## Load the .pth file saved after quantization calibration.
model.load_state_dict(
    torch.load('pth_save_path', map_location=torch.device('cpu')),
    strict=True,
)
## Disable quantization parameter update and export quantization parameters.
## Tips: Parameter export can be executed immediately upon the completion of calibration. However, these steps are separated to enhance process integrity, mitigating the risk of errors requiring the repetition of prior stages.
set_quant_state(model, weight_state=True, input_state=True)
set_calibrate_state(model, False)
generate_quant_params(
    model,
    output_dir,
    quant_param_2=False,
    embedding_separate=True,
)
## The **fake_quant_weight.pth** file and quantization parameter file are generated in **output_dir**. **fake_quant_weight.pth** is used only to replace weights when the ONNX graph is exported.

ONNX export:

Use torch.onnx.export to export the floating-point model (fake_quant_weight.pth needs to be loaded). If necessary, simplify the ONNX model, and ensure that the names of convolution operators and linear operators remain consistent with those defined by torch.