Workflow

This section describes the supported layer types in full quantization scenarios, and provides the API call sequence and example.

The supported layers are as follows. For details about the examples, see Sample List.

Based on different quantization objects, this section covers two scenarios with different supported quantization types and constraints:

Layers Supported by Full Quantization

If the data type of the original model is float32 (fp32), float16 (fp16), or bfloat16 (bf16), you can perform quantization to convert the data type to HiFloat8 (HiF8), float8 (fp8), MXFP8, or float4. This data format compression enables model lightweighting. Note that:

HIFloat8 (HiF8), FLOAT8 (float8), MXFP8, MXFP4, FP4_E2M1, and FP4_E1M2 are supported only by the Atlas 350 Accelerator Card.

Other product models support only the INT8 and INT4 quantized data types.

Table 1 Layers supported by quantization and their constraints

Supported Layer Type

Source Data Type

Supported Quantized Data Type Combination

Constraints

torch.nn.Linear

float32 (fp32), float16 (fp16), bfloat16 (bf16)

act_type: HIFLOAT8 wts_type: HIFLOAT8

act_type: FLOAT8_E4M3FN wts_type: FLOAT8_E4M3FN

Activations (data) support per-tensor quantization, and weights support both per-tensor and per-channel quantization.

The quantization algorithm is OFMR. The ofmr option must be configured in config.

bfloat16 (bf16)

act_type: MXFP8_E4M3FN wts_type: MXFP8_E4M3FN

2D to 6D input, per-group quantization, and symmetric quantization are supported. The result of cin divided by 32 and rounded up must be an even integer (multiple of 2).

The quantization algorithm is not used. That is, only the mx data type conversion is performed. The mxquant option must be configured in config.

bfloat16 (bf16)

act_type: FLOAT8_E4M3FN wts_type: FLOAT4_E2M1

2D to 6D data input is supported, and bias is set to false.

Activations (data) support the shape (m, k), and weights support the shape (n, k), where k is an integer multiple of 64.

Activations (data) support per-tensor quantization, and weights support per-group quantization.

Both activations and weights support only symmetric quantization.

The MIN-MAX and SmoothQuant algorithms are supported, and the minmax and smoothquant options must be configured in config, respectively.

float16 (fp16), bfloat16 (bf16)

act_type: INT8 wts_type: INT8

  • 2D to 6D data input is supported.
  • Activations (data) support per-tensor quantization, as well as symmetric and asymmetric quantization, with the bias quantized to INT32. Weights support per-tensor and per-channel symmetric quantization.
  • Activations (data) support per-token symmetric quantization, with the bias not quantized. k is a multiple of 16, and n is a multiple of 8. Weights support per-tensor and per-channel symmetric quantization.
  • The MIN-MAX and SmoothQuant algorithms are supported, and the minmax and smoothquant options must be configured in config, respectively.

torch.nn.Conv2d

float32 (fp32), float16 (fp16), bfloat16 (bf16)

act_type: HIFLOAT8 wts_type: HIFLOAT8

act_type: FLOAT8_E4M3FN wts_type: FLOAT8_E4M3FN

Per-tensor and per-channel quantization is supported (padding_mode = zeros).

The quantization algorithm is OFMR. The ofmr option must be configured in config.

Note: The act_type and wts_type parameters indicate the activation and weight quantization types in config, respectively. For details about the quantized data types and concepts such as per-tensor quantization and per-channel quantization, see Compression Concepts.

Layers Supported by Weight-Only Quantization

This section describes the weight-only quantization feature and the corresponding quantization algorithms such as AWQ and GPTQ.

If the data type of the original model is float16 (fp16) or bfloat16 (bf16), you can perform weight-only quantization to convert the data type to HiFloat8 (HiF8), float8 (fp8), MXFP4, or float4. This weight compression enables model lightweighting.

Table 2 Layers supported by quantization and their constraints

Supported Layer Type

Source Data Type

Supported Quantization Type Combination

Constraints

torch.nn.Linear

float16 (fp16), bfloat16 (bf16)

wts_type: HIFLOAT8

wts_type: FLOAT8_E4M3FN

  • Per-tensor and per-channel symmetric quantization types are supported.
  • 2D to 6D data input is supported.

    The OFMR and GPTQ quantization algorithms are supported. The ofmr and gptq options must be configured in config.

wts_type: MXFP4_E2M1

  • Per-group symmetric quantization is supported.
  • 2D to 6D data input is supported.
  • No algorithm needs to be configured. The mxquant option must be configured in config.
  • The GPTQ and AWQ quantization algorithms can be configured. The gptq and awq options must be configured in config.

wts_type: FLOAT4_E2M1

  • Per-group symmetric quantization is supported.
  • 2D to 6D data input is supported.
  • The MIN-MAX, GPTQ, and AWQ quantization algorithms are supported. The minmax, gptq, and awq options must be configured in config.

wts_type: INT8

  • Per-tensor, per-channel, and per-group quantization, as well as symmetric and asymmetric quantization, are supported.
  • The MIN-MAX, GPTQ, and AWQ quantization algorithms are supported, and the minmax, gptq, and awq options must be configured in config, respectively.
  • When wts_type is set to INT8, the K and N axes of the original model weight must be 32-element aligned. When wts_type is set to INT4, the K and N axes must be 64-element aligned.

float32 (fp32), float16 (fp16), bfloat16 (bf16)

wts_type: INT4

  • Per-tensor, per-channel, and per-group quantization, as well as symmetric and asymmetric quantization, are supported.
  • The MIN-MAX, GPTQ, and AWQ quantization algorithms are supported, and the minmax, gptq, and awq options must be configured in config, respectively.
  • When wts_type is set to INT4, the K and N axes must be 64-element aligned.

API Call Sequence

Figure 1 API call sequence
The user implements the operations in blue, while those in gray are implemented by using AMCT APIs.
  1. First, construct the original PyTorch model (ensure that it can perform inference normally) and the quantization configuration. Then call the quantize API to modify the original model and generate a quantization calibration model. Different data types use different quantization algorithms during modification.
    • Full quantization scenario
      • For HIFLOAT8 and FLOAT8_E4M3FN: Replace the operators to be quantized with OFMRQuant operators in the model to generate a quantization calibration model. For details about the OFMR algorithm, see OFMR Algorithm.
      • For MXFP8_E4M3FN and MXFP4_E2M1: No quantization algorithm is used. The conversion is directly performed based on the conversion formula.
      • For FLOAT8_E4M3FN (configured for activations) and FLOAT4_E2M1 (configured for weights): If the SmoothQuant algorithm is not configured, replace the operators to be quantized with MinMaxCalibrator operators in the model to generate a quantization calibration model. For details about the MIN-MAX quantization algorithm, see MIN-MAX Algorithm. If the SmoothQuant algorithm is configured, replace the operators to be quantized with SmoothQuant operators in the model to generate a quantization calibration model. For details about the SmoothQuant quantization algorithm, see SmoothQuant Algorithm.
      • For INT8 (configured for activations) and INT8 (configured for weights): If the SmoothQuant algorithm is not configured, replace the operators to be quantized with MinMaxCalibrator operators in the model to generate a quantization calibration model. If the SmoothQuant algorithm is configured, replace the operators to be quantized with SmoothQuant operators in the model to generate a quantization calibration model.
    • Weight-only quantization scenario
      • For HIFLOAT8, FLOAT8_E4M3FN, MXFP4_E2M1, FLOAT4_E2M1, FLOAT4_E1M2, INT8, and INT4: If no quantization algorithm is configured, replace the operators to be quantized with MinMaxCalibrator operators in the model to generate a quantization calibration model. For details about the MIN-MAX quantization algorithm, see MIN-MAX Algorithm. If the GPTQ algorithm is configured, replace the Linear operators to be quantized in the model with GPTQuant operators to generate a quantization calibration model.
      • For MXFP4_E2M1, FLOAT4_E2M1, FLOAT4_E1M2, INT8, and INT4: If the AWQ algorithm is configured, replace the Linear operators to be quantized in the model with LinearAWQuant operators to generate a quantization calibration model.
  2. Call the convert API to convert the quantization calibration operator into the quantized deployable operator corresponding to the NPU and output the quantized deployable model (of the torch.nn.Module type).

    If the generated deployable model needs to run inference, torch_npu must be used. For details, see PyTorch Graph Mode User Guide (TorchAir).

Example

The following code snippet demonstrates key steps for reference only. Do not copy, build, or run it directly.

# 1. Import the AMCT package.
import amct_pytorch as amct

# 2. Call AMCT to quantize the model.
# 2.1 Generate a quantization calibration model.
# Build a graph of the network for quantization.
ori_model = build_model()
model = copy.deepcopy(ori_model)
# Perform quantization configuration.
cfg = {
        'batch_num': 1,
        'quant_cfg': {
            'weights': {
                'type': 'int8',
                'symmetric': True,
                'strategy': 'tensor',
            },
        },
        'algorithm': {'minmax'},
        }
# Call the quantization API to generate a quantization calibration model.
amct.quantize(model, cfg)

# 2.2 Generate a quantized deployable model.
# Call the API to convert the quantization calibration model into a quantized deployable model.
amct.convert(model)