AWQ › Implementation
Language AI HandbookAll chapters AWQ › ImplementationOne section · Sections in this chapter5 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's walk through a practical AWQ quantization using the AutoAWQ library, which implements the full AWQ algorithm with optimized CUDA kernels for inference.
# First, install the required packages
!uv pip install autoawq transformers torch accelerate --quietWe'll show AWQ quantization on a small model to show the workflow. The process involves loading the model, configuring quantization parameters, and running the algorithm on a calibration dataset.
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
# Load model and tokenizer
model_path = "facebook/opt-125m" # Small model for demonstration
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)AWQ requires a calibration dataset to compute activation statistics. The calibration set should be representative of the model's intended use case. For general-purpose language models, a diverse set of text samples works well. You need enough variety for the activation statistics to reflect the model's typical behavior rather than a narrow slice.
# Prepare calibration data
calibration_texts = [
"The quick brown fox jumps over the lazy dog.",
"Machine learning models require careful optimization to reach good performance.",
"Natural language processing has transformed how computers understand human text.",
"Quantization reduces model memory requirements while maintaining accuracy.",
"Large language models show remarkable capabilities in text generation.",
]The AWQ quantization configuration specifies the bit width, group size, and other parameters. Group size controls how many weights share a quantization scale, with smaller groups giving better accuracy at the cost of additional metadata storage overhead.
# Configure AWQ quantization
quant_config = {
"zero_point": True, # Use asymmetric quantization with a zero point
"q_group_size": 128, # Number of weights that share one scale and zero point
"w_bit": 4, # 4-bit weight quantization target
"version": "GEMM", # Optimized GEMM kernel variant for throughput
}
# Run AWQ quantization
model.quantize(
tokenizer,
quant_config=quant_config,
calib_data=calibration_texts,
)The quantize method runs the full AWQ algorithm: it collects activation statistics from the calibration data, computes optimal per-channel scaling factors using grid search, turns and quantizes the weights, and fuses the inverse scaling into the preceding layers. This all happens automatically inside the library call, with the resulting model stored in the quantized format ready for saving.
# Save the quantized model
output_path = "./opt-125m-awq"
model.save_quantized(output_path)
tokenizer.save_pretrained(output_path)The output confirms the successful creation of the quantized model directory, which now contains the compressed weights and configuration files needed for inference. The saved files include the INT4-packed weight tensors, quantization scales and zero points, and a configuration file describing the quantization parameters.
Loading and Using AWQ Models
Once quantized, AWQ models load efficiently for inference. The AutoAWQ library gives optimized CUDA kernels that perform dequantization on-the-fly during matrix multiplication, avoiding the need to fully dequantize weights into memory.
# Load quantized model for inference
quant_model = AutoAWQForCausalLM.from_quantized(
output_path,
fuse_layers=False, # Disable fusion to allow layer inspection
)
tokenizer = AutoTokenizer.from_pretrained(output_path)
# Generate text
prompt = "The future of artificial intelligence"
tokens = tokenizer(prompt, return_tensors="pt").to("cuda")
output = quant_model.generate(**tokens, max_new_tokens=50)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)The generated text shows that the 4-bit quantized model retains the linguistic capabilities of the original, creating coherent and contextually appropriate output despite the 4x compression. The fuse_layers=True option lets additional optimizations that combine sequential operations like dequantization and matrix multiply into single fused CUDA kernels, reducing memory bandwidth requirements and improving throughput further.
Understanding the Quantization Structure
AWQ stores quantized weights in a packed format where multiple 4-bit values share INT32 storage. Let's examine the structure of a quantized layer to understand what is stored on disk and in memory.
# Inspect a specific layer (e.g., first attention projection)
layer_name = "model.decoder.layers.0.self_attn.q_proj"
layer = quant_model.model.decoder.layers[0].self_attn.q_proj
# Get parameter shapes
qweight_shape = layer.qweight.shape
scales_shape = layer.scales.shape
qzeros_shape = layer.qzeros.shape if hasattr(layer, "qzeros") else None
# Calculate compression ratio
# qweight stores 8 x 4-bit weights per int32 (4 bytes)
# Original FP16 weights would be 2 bytes each
num_weights = layer.qweight.numel() * 8
fp16_size = num_weights * 2 # bytes
int4_size = layer.qweight.numel() * 4 # bytes
compression_ratio = fp16_size / int4_sizeThe inspection reveals the internal structure of AWQ quantization: weights are packed eight-per-INT32 into the qweight tensor, while scales are kept in FP16 precision to preserve accurate dequantization. The compression ratio confirms the expected 4x reduction in weight memory usage compared to FP16. The scales and zero points occupy very little additional space because they are shared across groups of 128 weights.