SparseGPT

repository·master·Indexed 21 days ago

https://github.com/ist-daslab/sparsegpt

A tool for one-shot pruning of massive language models, including OPT, BLOOM, and LLaMA. It supports unstructured, n:m (structured), and sparse + quantized compression. The library provides scripts for performing various compression tasks, such as uniform sparsity and magnitude pruning, and allows for evaluation using datasets like c4, ptb, and wikitext.

Tokens
1.5K
Snippets
7
Records
8
Agent score
25%

What's inside SparseGPT

  1. Run SparseGPT on OPT models

    master

    Use opt.py to perform various compression tasks on OPT models. You can specify the HuggingFace model name and the evaluation dataset (e.g., c4).

    Common Tasks:

    • Dense Baseline: Run the model without pruning.
    • Magnitude Baseline: Run pruning using the Magnitude Pruning (GMP) method.
    • Uniform Sparsity: Prune to a specific sparsity level (e.g., 50%).
    • N:M Sparsity: Prune to a structured 2:4 sparsity pattern.
    • Sparse + Quantized: Combine sparsity with weight quantization (e.g., 4-bit).
    # Run dense baseline
    python opt.py facebook/opt-125m c4
    
    # Run magnitude baseline
    python opt.py facebook/opt-125m c4 --sparsity .5 --gmp
    
    # Prune to 50% uniform sparsity with SparseGPT
    python opt.py facebook/opt-125m c4 --sparsity .5
    
    # Prune to full 2:4 sparsity with SparseGPT
    python opt.py facebook/opt-125m c4 --prunen 2 --prunem 4
    
    # Prune to 50% + 4-bit with SparseGPT
    python opt.py facebook/opt-125m c4 --sparsity .5 --wbits 4
  2. Install SparseGPT dependencies

    master

    To use SparseGPT, ensure you have the following Python packages installed. The implementation has been tested with these specific versions:

    • torch: v1.10.1+cu111 or higher
    • transformers: v4.21.2
    • datasets: v1.17.0
    pip install torch==1.10.1+cu111 transformers==4.21.2 datasets==1.17.0
  3. Prune an OPT model using SparseGPT

    master

    You can prune OPT models using the opt.py script. The script requires a model identifier, a calibration dataset, and various configuration flags to define the pruning strategy.

    Supported OPT Models (Colab friendly)

    • facebook/opt-125m
    • facebook/opt-350m
    • facebook/opt-1.3b (Larger models like 6.7b and 13b can be used with bitsandbytes.)

    Calibration Datasets

    Select one of the following for calibration:

    • c4 (Default)
    • ptb
    • wikitext

    Pruning Configuration Flags

    • --sparsity <float>: Specifies unstructured sparsity as a floating point number in [0, 1].
    • --prunen <int> and --prunem <int>: Specifies semistructured N:M pruning patterns.
    • --gmp: Enables Magnitude Pruning instead of SparseGPT.
    • --wbits <int>: Applies quantization on top of sparsity.
    • --save <path>: Directory to save the pruned model.

    After execution, the script prints perplexity scores on wikitext2, ptb, and c4 benchmarks.

    # Example: Prune facebook/opt-125m to 0.5 unstructured sparsity via SparseGPT
    python opt.py facebook/opt-125m c4 --sparsity 0.5 --save sparse_opt/opt-125m
  4. Install SparseGPT dependencies and setup

    master

    To use the SparseGPT demo, you need to install the datasets and transformers libraries and clone the repository. This is typically done in a Colab or Jupyter environment.

    !pip install -q datasets
    !pip install -q transformers
    !git clone https://github.com/IST-DASLab/sparsegpt
  5. Configure SparseGPT output and logging

    master

    When running the pruning scripts, you can use the following flags to manage outputs:

    • --save <path>: Specify the path where the sparsified model checkpoint should be saved.
    • --log_wandb: Optionally log evaluation results to Weights & Biases (W&B).
  6. Compare generations between dense and sparse models

    master

    Once a model is pruned, you can compare the text generation quality of the original (dense) model against the pruned (sparse) model using the transformers library.

    from transformers import AutoTokenizer, OPTForCausalLM
    
    device = 'cuda'
    
    # Load dense model
    model_dn = OPTForCausalLM.from_pretrained('facebook/opt-125m', torch_dtype='auto').to(device)
    
    # Load sparse model (from your saved path)
    model_sp = OPTForCausalLM.from_pretrained('sparse_opt/opt-125m', torch_dtype='auto').to(device)
    
    # Initialize tokenizer
    tokenizer = AutoTokenizer.from_pretrained('facebook/opt-125m')
    
    input_text = "It takes a great deal of bravery"
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
    
    # Generate with dense model
    output_ids_dn = model_dn.generate(input_ids)
    print("Dense:", tokenizer.decode(output_ids_dn[0].cpu(), skip_special_tokens=True))
    
    # Generate with sparse model
    output_ids_sp = model_sp.generate(input_ids)
    print("Sparse:", tokenizer.decode(output_ids_sp[0].cpu(), skip_special_tokens=True))