INC ONNX Runtime 3.x API design

[mengniwang95]

INC ONNX Runtime 3.x API Design

Target

• Design 3.x API for onnxrt framework and keep consistent with ORT user API

Main principles

• autotune is the exposed user interface API, which requires a set of configurations.
• _quantize is an internal API.
• Each algorithm has its own configuration class, such as GPTQConfig and autotune will use a set of configurations.

Repo Architecture

• common
  • config # basic configuration definition.
  • utils
    • logger # logger is a common used tool.
• onnxrt
  • algorithms # make sure all algos here have user interface.
    • smooth_quant
      • algo_entry.py # entry of smooth quant algo
      • ...
    • weight_only
      • rtn.py
      • gptq.py
      • ...
      • algo_entry.py # entry of weight-only algo
  • quantization
    • init.py # make sure all configurations and autotune are imported here.
    • config.py # contains all configurations, such as SmoothQuantConfig, get_default_config.
    • quantize.py # the internal interface of quantization.
  • utils
    • utility.py # saving common used functions, such as calculate_scale_zp.
    • onnx_model.py # contains all functions of onnx model.

Previous Design

StaticQuant & SmoothQuant

# Tuning space
# 'weight': {
#             'dtype': ['int8'],
#             'scheme': ['sym'],
#             'granularity': ['per_channel'],
#             'algorithm': ['minmax']
#             },
# 'activation': {
#             'dtype': ['uint8'],
#             'scheme': ['asym', 'sym'],
#             'granularity': ['per_tensor'],
#             'algorithm': ['minmax', 'kl']
#             },
#         },
conf = PostTrainingQuantConfig(
    backend="default",
    op_name_dict=op_name_dict,
    op_type_dict=op_type_dict,
    recipes={
        "smooth_quant": True,
        "smooth_quant_args": {"folding": folding, "alpha": np.arange(0.1, 0.4, 0.05).tolist()},
    }, # without smoothquant recipes, it's static quantization
)
calib_dataloader = Dataloader()
int8_model = quantization.fit(
    model,
    conf,
    calib_dataloader=calib_dataloader,
)
int8_model.save(model_path)

Weight-only Quantization

# Tuning space
# 'MatMul': &cap_weight_only_integer_linear {  # only MatMul now
#     'weight': {
#                 'dtype': ['int'],
#                 'bits': [4, 3, 8], # [1-8], # 4
#                 # group_size=-1 means per-channel, others means per-group
#                 'group_size': [32, -1, 1, 16, 64, 128, 256, 512, 1024], # [1-inf], # 32
#                 'scheme': ['sym', 'asym'], # sym, no ZP
#                 'algorithm': ['RTN', 'AWQ', 'GPTQ']
#                 },
conf = PostTrainingQuantConfig(
    approach="weight_only",
    op_type_dict={
        ".*": {  # re.match
            "weight": {
                "bits": 8,  # 1-8 bit
                "group_size": -1,  # -1 (per-channel)
                "scheme": "sym",
                "algorithm": "RTN", # select from ['RTN', 'AWQ', 'GPTQ']
            },
        },
    },
    recipes={
        'rtn_args':{'accuracy_level': 0},
        'gptq_args':{'percdamp': 0.01, 'actorder':True, 'block_size': 128, 'actorder': False, 'mse': False, 'perchannel': True, 'accuracy_level': 0},
        'awq_args':{'enable_auto_scale': True, 'enable_mse_search': True, 'accuracy_level': 0},
    },
)

New Design

StaticQuant & SmoothQuant

Configuration

The argument to config is data or a list of data. If the parameters can be assembled into different configurations, the returned obj will be a list of configurations used for autotuning.

# default configuration, this APIs can be updated later
conf = get_sq_default_config()

# customized configuration
fp32config = FP32Config()

conf = SmoothQuantConfig(
    alpha = [0.5, 1.0],
    folding = [True, False],
    scale_sharing=[True, False], # whether share the same scale for layers with the same input.
    auto_alpha_args = {
        "init_alpha": 0.5,
        "alpha_min": 0.0,
        "alpha_max": 1.0,
        "alpha_step": 0.1,
    }
)
# By default, conf is a global configuration
conf.set_local(["lm_head", "Linear"], fp32config) # op_name, op_type fallback

Quantize Interface

from neural_compressor.onnxrt.quantization import (
    SmoothQuantConfig,
    TuningConfig,
    autotune,
)
# _quantize is only for internal use
from neural_compressor.onnxrt.quantization.quantize import _quantize

model = _quantize(
    model,
    conf,
    calibration_data_reader , # calibration function
) # int8 model here is an onnx model, not an INC model

conf = TuningConfig(quant_configs=[SmoothQuantConfig(alpha=[0.5, 1.0])])
# SmoothQuantConfig(alpha=0.5), SmoothQuantConfig(alpha=1.0) == SmoothQuantConfig.expend(SmoothQuantConfig(alpha=[0.5, 1.0]))

model = autotune(
    model,
    conf,
    calibration_data_reader , # calibration function
)

onnx.save(model, 'model.onnx')

Weight-only Quantization

Configuration

# default configuration
# by default, last matmul is quantized
conf = get_rtn_default_config() 
conf = get_gptq_default_config()
conf = get_awq_default_config()

# customized configuration
fp32config = FP32Config()
conf = RTNConfig(
    dtype=['int'],
    group_size=['32', '128'],
    use_sym=[True, False],
)
conf = GPTQConfig(
    dtype=['int'],
    group_size=['32', '128'],
    use_sym=[True, False],
    use_actorder=[True, False],
)
conf = AWQConfig(
    dtype=['int'],
    group_size=['32', '128'],
    use_sym=[True, False],
    folding=[True, False],
)

# By default, conf is a global configuration
conf.set_local(["lm_head"], fp32config) # op_name, op_type fallback

Quantize Interface

from neural_compressor.onnxrt.quantization import (
    RTNConfig,
    GPTQConfig,
    AWQConfig,
    TuningConfig,
    autotune,
)
# _quantize is only for internal use
from neural_compressor.onnxrt.quantization.quantize import _quantize

fp32config = FP32Config()
conf = RTNConfig(dtype=['int'])
conf = GPTQConfig(dtype=['int'])

conf.set_local(["lm_head"], fp32config) # op_name, op_type fallback
conf.set_local(["*.fc1"], RTNconfig) # specific configuration for op_name, op_type.

model = _quantize(
    model,
    conf,
    calibration_data_reader, # calibration function
) # int8 model here is an onnx model, not an INC model

# below config will try three weight-only quantization configuration one-by-one
conf = TuningConfig(quant_configs=[RTNConfig(), GPTQConfig(), AWQConfig()])
# below config will try different group_size with fp32 lm_head one-by-one
conf = TuningConfig(quant_configs=[GPTQConfig(group_size=['32', '128'].set_local("lm_head", fp32config)])

model = autotune(
    model,
    conf,
    calibration_data_reader , # calibration function
)

onnx.save(model, 'model.onnx')