Skip to main content
`Ctrl`+`K`

# Quantizing LLMs for ONNX Runtime GenAI

This document provides examples of quantizing large language models (LLMs) to UINT4 using the AWQ algorithm via the Quark API, and exporting them to ONNX format using the ONNX Runtime Gen AI Model Builder.

ONNX Runtime GenAI offers an end-to-end pipeline for working with ONNX models, including inference using ONNX Runtime, logits processing, search and sampling, and key-value (KV) cache management. For detailed documentation, visit the ONNX Runtime Gen AI Documentation. The tool includes a Model Builder that facilitates exporting models to the ONNX format.

Note

For large models, it is recommended to run this workflow on a data center GPU, as it is memory-intensive. Most laptop and desktop GPUs will not have enough memory to process larger parameter models efficiently. If you haven’t done so already, and have a GPU, we suggest you create a fresh Quark environment with a PyTorch ROCm or CUDA install for your platform https://pytorch.org/get-started/locally/.

## Preparation

### Model Preparation

To use Llama2 models, download the HF Llama2 checkpoint. Access to these checkpoints requires a permission request to Meta. For more information, refer to the Llama2 page on Hugging Face. Once permission is granted, download the checkpoint and save it to the `<llama checkpoint folder>`.

### Installation

We will use a script from the `examples/` directory for Quark. This directory is found in the Quark `.zip` that can be downloaded at 📥amd_quark.zip release_version.

The `quantize_quark.py` script, that we will use, is found in a sub-directory of `examples/`. It requires some additional dependencies be installed:

    cd examples/torch/language_modeling/llm_ptq/
    pip install -r requirements.txt

## Quark UINT4 Quantization with AWQ

Quantization Configuration: AWQ / Group 128 / Asymmetric / FP16 input_tensors

Use the following command to quantize the model:

    python3 quantize_quark.py --model_dir <llama checkpoint folder> \
                              --output_dir <quantized safetensor output dir> \
                              --quant_scheme uint4_wo_128 \
                              --num_calib_data 128 \
                              --quant_algo awq \
                              --dataset pileval_for_awq_benchmark \
                              --seq_len 512 \
                              --model_export hf_format \
                              --data_type float16 \
                              --custom_mode awq

This will generate a directory containing the safe tensors at the specified `<quantized safetensor output dir>`.

Note

About ``–exclude_layers``:

- By default, certain layers (like `lm_head`) are automatically excluded from quantization based on the model type.

- To include all layers in quantization (exclude nothing), use `--exclude_layers` without any arguments.

- To specify custom layers to exclude, use `--exclude_layers "layer1" "layer2"` (supports wildcards like `"*down_proj*"`).

Example: `--exclude_layers` (includes lm_head in quantization)

Note

To quantize the model for BF16 input_tensors, use the `--data_type bfloat16` flag.

Note

To specify a group size other than 128, such as 32, use `--quant_scheme uint4_wo_32` instead of `--quant_scheme uint4_wo_128`. Available group sizes are 32, 64, and 128 (e.g., `uint4_wo_32`, `uint4_wo_64`, `uint4_wo_128`).

### (Optional) Quark UINT4 Quantization with Different Group Sizes per Layer

Quark supports quantizing layers with different group sizes, providing finer-grained control over the quantization process. This allows you to better balance performance and accuracy. For example, to quantize the model with 32 group size for lm_head, while 128 group size for the rest, use the following command:

    python3 quantize_quark.py --model_dir <llama checkpoint folder> \
                              --output_dir <quantized safetensor output dir> \
                              --quant_scheme uint4_wo_128 \
                              --num_calib_data 128 \
                              --quant_algo awq \
                              --dataset pileval_for_awq_benchmark \
                              --seq_len 512 \
                              --model_export hf_format \
                              --data_type float16 \
                              --exclude_layers \
                              --layer_quant_scheme lm_head uint4_wo_32

Note

This is an advanced feature that is not supported by the standard AWQ model format. As a result, the quantized model is stored in the Quark model format, which does not require the `--custom_mode awq` argument.

Support for the Quark model format in ONNX Runtime GenAI is coming soon in v0.7 release. For early access you can try this feature branch.

## Exporting Using ONNX Runtime Gen AI Model Builder

Install the ONNX Runtime Gen AI tool package using `pip`:

    pip install onnxruntime-genai

To export the quantized model to ONNX format, run the following command:

    python3 -m onnxruntime_genai.models.builder \
            -i <quantized safetensor output dir> \
            -o <quantized onnx model output dir> \
            -p int4 \
            -e dml

Note

The input_tensors data type of the ONNX model depends on the combination of the `-p` (precision) and `-e` (execution provider) flags. For example:

- Using `-p int4 -e dml` will generate an ONNX model with float16 input_tensors prepared for the DirectML execution provider for hybrid (NPU + iGPU) flow.

- To generate an ONNX model with float32 input_tensors for NPU flow, use the `-p int4 -e cpu` flag.