Installing SparseML

SparseML is Neural Magic's core toolkit for streamlining the optimization of deep learning models using cutting-edge sparsification methods (e.g., pruning, quantization). This guide covers installation options, including PyPI, GitHub, and specialized applications for various AI domains.

Prerequisites

Hardware Requirements

• CPU: Intel, AMD, or ARM
• RAM: 4GB+ recommended (depends on your models and workflows)
• GPU: NVIDIA (recommended, 16GB+ VRAM); AMD (limited support)

Software Requirements

• OS: Linux (e.g., Ubuntu, CentOS, Red Hat); MacOS (limited support)
• Python: 3.8 - 3.11
• PyTorch: >= 1.1.0, < 2.2

Installation

PyPI

For the most common use case, install the current release version of SparseML using PyPi:

pip install sparseml-nightly

GitHub (Advanced)

For development purposes or advanced use cases, install directly from the GitHub repository:

pip install git+https://github.com/neuralmagic/sparseml.git@main

Or from a locally cloned repository:

git clone https://github.com/neuralmagic/sparseml.git
cd sparseml
git checkout main
pip install -e .[dev]

Specialized Installations

Install SparseML with tailored support with the following extras for domain-specific use cases and tasks.

Generative AI: Hugging Face

This extra provides built-in support for models like Llama, Mistral, MPT, GPT, and others for generative AI, particularly transformer architectures. It enables compatibility of Hugging Face's Transformers pipelines and models to SparseML allowing SOTA pruning, quantization, and other optimizations.

pip install sparseml-nightly[llm]

Object Detection: YOLOv8

For object detection, this extra provides built-in support for YOLOv8 models. It enables compatibility of YOLOv8 models to SparseML, allowing SOTA pruning, quantization, and other optimizations.

pip install sparseml-nightly[yolov8]

Image Classification: TorchVision

For image classification, this extra provides built-in support for TorchVision models. It enables compatibility of TorchVision models to SparseML, allowing SOTA pruning, quantization, and other optimizations.

pip install sparseml-nightly[torchvision]

Natural Language Processing: Hugging Face

This extra provides built-in support for Hugging Face's Transformer models for natural language processing. It enables compatibility of Hugging Face models to SparseML, allowing SOTA pruning, quantization, and other optimizations.

pip install sparseml-nightly[transformers]

Development

For development purposes, this extra provides built-in support for development tools.

pip install sparseml-nightly[dev]