KV Cache Quantization

This section describes what KV cache is, why KV cache quantization is required, which layers are supported by KV cache quantization, as well as how KV cache quantization is called.

Overview

Key-value cache (KV cache) is a cache technology that reuses computation results by storing key-value pairs to improve performance and reduce memory consumption. In large-scale training and inference, the continuous increase of batch_size and sequence_length leads to surging memory overhead occupied by KV cache, which even exceeds that of the model itself. INT8 KV cache quantization can greatly reduce the memory consumption during model running and reduce the model deployment cost.

The layer that supports KV cache quantization is torch.nn.Linear, whose source data types are float32 and float16. For details about the KV cache quantization sample, see Sample List.

This feature is supported only by the following products:

Atlas A2 training product/Atlas A2 inference product

Atlas inference product

Atlas A3 training product/Atlas A3 inference product

Atlas 350 Accelerator Card

API Call Sequence

Figure 1 shows the API call sequence for KV cache quantization.

Figure 1 API call sequence for KV cache quantization

The user implements the operations in blue, while those in gray are implemented by using AMCT APIs.

Build an original PyTorch model and then generate a quantization configuration file by using the create_quant_cali_config API.
Find the Linear operator to be quantized based on the PyTorch model and quantization configuration file. Then, call the create_quant_cali_model API to modify the original PyTorch model and add the IFMR/HFMG quantization operator after the Linear output.
Use the calibration dataset to perform forward pass in the PyTorch environment, call the IFMR/HFMG quantization algorithm to calibrate the output based on the quantization algorithm configuration, and output the result to the quantization factor record file in the corresponding format.

To support more operator types or customize other operations, you can call the QuantCalibrationOp API to construct a graph, perform quantization calibration, and output the quantization factor record file.

Example

Take the following steps to get started. Update the sample code based on your situation.
Tweak the arguments passed to AMCT API calls as required.

Import the AMCT package and set the log level using the environment variable in "AMCT (PyTorch)" in Post-installation Actions.
1
import amct_pytorch as amct
(Optional) Validate the inference script and environment setup in the source PyTorch environment. (Update the sample code based on your situation.)
You are advised to run inference on the original model in the PyTorch environment based on the test dataset to validate the environment setup and inference script.

This step is recommended as it guarantees a properly functioning original model for inference with acceptable accuracy. You can use a subset from the test dataset to improve the efficiency.
1
user_do_inference_torch(ori_model, test_data, test_iterations)

Run AMCT to quantize the model.

Parse quant.cfg to generate a quantization configuration file.

config_file = './tmp/config.json'
amct.create_quant_cali_config(config_file=config_file,
			         model=ori_model,
			         quant_layers=None,
			         config_defination="./configs/quant.cfg")

Modify the graph by inserting quantization operators for quantization parameter calculation.
Find the Linear operator to be quantized, insert the IFMR/HFMG quantization operator after the Linear output, and calibrate and calculate the quantization factors online, so that the output can be quantized into INT8.
1 2 3 4
record_file = './tmp/record.txt' calibration_model = amct.create_quant_cali_model(config_file=config_file, record_file=record_file, model=ori_model)
Run inference on the modified model (calibration_model) in the PyTorch environment based on the calibration dataset (calibration_data) to determine the quantization factors. (Update the sample code based on your situation.)
Pay attention to the following points:
1. Ensure that the calibration dataset and the preprocessed data match the model to preserve the accuracy.
2. Ensure that the number of forward passes (specified by batch_num) is large enough. Otherwise, the subsequent process will fail.
If you get the error "[IFMR]: Do layer xxx data calibration failed!" during the calibration, rectify the fault by referring to Why Do I See "[IFMR]: Do layer xxx data calibration failed!" During Calibration?
1
user_do_inference_torch(calibration_model, calibration_data, batch_num)
Output the quantization factor record file.

Parent topic: Graph-based Quantization