AWQ Quantization

This module provides AWQ (Activation-aware Weight Quantization) quantization using llmcompressor to compress language models while maintaining high inference quality.

Overview

AWQ is a post-training quantization method that reduces model size and improves inference speed by quantizing weights to lower precision while keeping activations at higher precision. This implementation supports various quantization schemes optimized for different use cases.

Installation

Make sure you have installed the required dependencies:

pip install -r requirements.txt

Quick Start

Basic Usage

Run AWQ quantization with minimal configuration:

python quantization/awq_quantize.py --model_path /path/to/your/model

This creates an awq-4bit subfolder within your model directory containing the quantized model, ready for inference.

Advanced Usage

For more control over the quantization process:

python quantization/awq_quantize.py \
    --model_path /path/to/your/model \
    --output_suffix "awq-4bit" \
    --group_size 128 \
    --max_seq_length 4096 \
    --num_calibration_samples 512 \
    --dataset /path/to/your/calibration/dataset

Parameters

Parameter	Description	Default	Required
`--model_path`	Path to the model to quantize	-	Yes
`--output_suffix`	Name of the output subfolder	`awq-4bit`	No
`--scheme`	Quantization scheme (see below)	`W4A16_ASYM`	No
`--group_size`	Quantization group size	`128`	No
`--max_seq_length`	Maximum sequence length for calibration	`4096`	No
`--num_calibration_samples`	Number of calibration samples	`512`	No
`--dataset`	Custom calibration dataset path	`open_platypus`	No
`--ignore_layers`	Layer names to ignore for quantization	`["lm_head"]`	No
`--device`	Computing device (auto/cuda/cpu)	`auto`	No
`--log_level`	Log level (DEBUG/INFO/WARNING/ERROR)	`INFO`	No
`--dry_run`	Show config only, don't quantize	`False`	No

Quantization Schemes

Scheme	Description	Use Case
W4A16_ASYM	4-bit weights, 16-bit activations, asymmetric	Recommended: High compression with good quality
W4A16_SYM	4-bit weights, 16-bit activations, symmetric	Alternative 4-bit option
W8A16	8-bit weights, 16-bit activations	Conservative: Higher precision, larger size

Output

The script creates a quantized model in a subfolder within your specified model path:

/path/to/your/model/
├── original_model_files...
└── awq-4bit/  # or your custom suffix
    ├── quantized_model_files...
    └── quantization_config.json

The quantized model is immediately ready for inference and can be loaded using standard model loading procedures.

Requirements

Python 3.10+
Sufficient RAM for model loading
CUDA-compatible GPU (recommended for larger models)

Notes

Quantization time depends on model size and number of calibration samples
The calibration dataset significantly impacts quantization quality
Monitor GPU memory usage during quantization of large models
Use --dry_run to preview the quantization configuration before running
The --ignore_layers parameter helps preserve important layers like output heads

Troubleshooting

Issue: Out of memory errors during quantization Solution:

Reduce --num_calibration_samples (e.g., from 512 to 256)
Use --device cpu for CPU-only quantization (slower but uses less GPU memory)
Ensure sufficient system RAM is available

Issue: Poor quality after quantization Solution:

Try using W8A16 scheme for higher precision
Increase --num_calibration_samples for better calibration
Use a representative calibration dataset similar to your target use case

Issue: Specific layers causing issues Solution: Add problematic layer names to --ignore_layers to exclude them from quantization