GitHub - intel/neural-compressor: SOTA low-bit LLM quantization (INT8/FP8/INT4/FP4/NF4) & sparsity; leading model compression techniques on TensorFlow, PyTorch, and ONNX Runtime (original) (raw)

Intel® Neural Compressor aims to provide popular model compression techniques such as quantization, pruning (sparsity), distillation, and neural architecture search on mainstream frameworks such as TensorFlow, PyTorch, and ONNX Runtime, as well as Intel extensions such as Intel Extension for TensorFlow and Intel Extension for PyTorch. In particular, the tool provides the key features, typical examples, and open collaborations as below:

What's New

Installation

Choose the necessary framework dependencies to install based on your deploy environment.

Install Framework

Install Neural Compressor from pypi

# Install 2.X API + Framework extension API + PyTorch dependency
pip install neural-compressor[pt]
# Install 2.X API + Framework extension API + TensorFlow dependency
pip install neural-compressor[tf]

Note: Further installation methods can be found under Installation Guide. check out our FAQ for more details.

Getting Started

After successfully installing these packages, try your first quantization program. Following example code demonstrates FP8 Quantization, it is supported by Intel Gaudi2 AI Accelerator.
To try on Intel Gaudi2, docker image with Gaudi Software Stack is recommended, please refer to following script for environment setup. More details can be found in Gaudi Guide.

Run a container with an interactive shell, more info

docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu24.04/habanalabs/pytorch-installer-2.6.0:latest

Run the example,

from neural_compressor.torch.quantization import ( FP8Config, prepare, convert, )

import torch import torchvision.models as models

model = models.resnet18() qconfig = FP8Config(fp8_config="E4M3") model = prepare(model, qconfig)

Customer defined calibration. Below is a dummy calibration

model(torch.randn(1, 3, 224, 224).to("hpu"))

model = convert(model)

output = model(torch.randn(1, 3, 224, 224).to("hpu")).to("cpu") print(output.shape)

More FP8 quantization doc.

Following example code demonstrates weight-only large language model loading on Intel Gaudi2 AI Accelerator.

from neural_compressor.torch.quantization import load

model_name = "TheBloke/Llama-2-7B-GPTQ" model = load( model_name_or_path=model_name, format="huggingface", device="hpu", torch_dtype=torch.bfloat16, )

Note: Intel Neural Compressor will convert the model format from auto-gptq to hpu format on the first load and save hpu_model.safetensors to the local cache directory for the next load. So it may take a while to load for the first time.

Documentation

Overview
Architecture Workflow APIs LLMs Recipes Examples
PyTorch Extension APIs
Overview Dynamic Quantization Static Quantization Smooth Quantization
Weight-Only Quantization FP8 Quantization MX Quantization Mixed Precision
Tensorflow Extension APIs
Overview Static Quantization Smooth Quantization
Transformers-like APIs
Overview
Other Modules
Auto Tune Benchmark

Note: From 3.0 release, we recommend to use 3.X API. Compression techniques during training such as QAT, Pruning, Distillation only available in 2.X API currently.

Selected Publications/Events

Note: View Full Publication List.

Additional Content

Communication