optipfair 0.1.4

A library for structured pruning & Bias visualization of large language models

OptiPFair

optiPfair

The Python library for making LLMs both efficient (via pruning) and fair (via bias analysis).

Documentation · Report Bug · Request Feature

---

🚀 Interactive Demos: Try OptiPFair NOW

Experience the power of OptiPFair directly in your browser. No installation, no setup.

Live Bias Visualization (Recommended Demo)	Pruning Compatibility Check	Bias Compatibility Check
Analyze any compatible model from Hugging Face with a full UI.	Check if your model's architecture can be pruned by OptiPFair.	The coder's alternative to our live demo.
🚀 Launch the Live Demo on HF Spaces

---

✅ Why OptiPFair?

OptiPFair is more than just another pruning library. It's a toolkit designed for the modern AI developer who cares about both performance and responsibility.

• Efficiency & Fairness in One Place: Stop juggling tools. OptiPFair is the only library designed to integrate structured pruning with powerful, intuitive bias visualization and analysis.

• Optimized for Modern Architectures: We focus on what works now. The library is specialized for GLU-based models like LLaMA, Mistral, Gemma, and Qwen, ensuring relevant and effective pruning.

• Go Beyond Numbers with Bias Visualization: Don't just get a bias score. Our visualization tools (PCA, heatmaps, mean differences) help you understand how and where your model encodes bias, enabling more effective mitigation.

• 🤖 AI-Assisted Development: Accelerate your workflow using the included LLM Reference Manual. Provide it to your favorite LLM (ChatGPT, Claude) to get expert-level help and generate integration code instantly.

• 🔬 Backed by Research: Our methods aren't arbitrary. They are built upon and validated by ongoing applied research in model optimization and fairness analysis.

---

⚙️ Installation

Choose the installation method that best suits your needs. For bias visualization features, you'll need the [viz] extra. Standard Installation For core pruning functionality:

pip install optipfair

Full Installation (with Bias Visualization) To use the bias analysis and visualization tools, install with the [viz] extra dependencies:

pip install "optipfair[viz]"

Developer Installation To install from the source for contributing or development:

git clone https://github.com/peremartra/optipfair.git
cd optipfair
pip install -e .

---

⚡ Quick Start

See how to use OptiPFair's core features in just a few lines of code.

Pruning with the Python API

Prune 20% of the MLP neurons from a model using the Maximum Absolute Weight (MAW) method.

from transformers import AutoModelForCausalLM
from optipfair import prune_model

# Load a pre-trained model
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")

# Prune 20% of neurons from MLP layers
pruned_model, stats = prune_model(
    model=model,
    pruning_type="MLP_GLU",
    neuron_selection_method="MAW",
    pruning_percentage=20,
    show_progress=True,
    return_stats=True
)

# Print pruning statistics
print(f"Original parameters: {stats['original_parameters']:,}")
print(f"Pruned parameters: {stats['pruned_parameters']:,}")
print(f"Reduction: {stats['reduction']:,} parameters ({stats['percentage_reduction']:.2f}%)")

# Save the pruned model
pruned_model.save_pretrained("./pruned-llama-model")

The pruning process yields tangible results in model size and performance. Here's a sample comparison for Llama-3.2-1B after pruning 20% of its MLP neurons:

Metric	Original Model	Pruned Model	Improvement
Total Parameters	1.24B	1.07B	-13.03%
Inference Speed	Benchmark in progress	Benchmark in progress	Coming soon
MMLU Score	Benchmark in progress	Benchmark in progress	Minimal change expected

Results based on the MAW pruning method. Full benchmark results will be published shortly.

Visualizing Bias with the Python API

Generate visualizations to analyze how a model's activations differ when processing prompts with varying demographic attributes.

from transformers import AutoModelForCausalLM, AutoTokenizer
from optipfair.bias import visualize_bias

# Load model and tokenizer
model_name = "meta-llama/Llama-3.2-1B"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define prompt pairs that differ in demographic attributes
prompt_pairs = [
    ("The white man walked through the neighborhood. The police officer thought he", 
     "The Black man walked through the neighborhood. The police officer thought he")
]

# Generate visualizations and get bias metrics
_, metrics = visualize_bias(
    model, 
    tokenizer,
    prompt_pairs=prompt_pairs,
    visualization_types=["mean_diff", "pca"],
    layers="first_middle_last",
    output_dir="./bias_analysis"
)

# Print overall bias metrics
overall = metrics["pair_1"]["metrics"]["overall_metrics"]
print(f"Mean activation difference: {overall['mean_difference']:.6f}")

The code above generates the following types of visualizations, allowing for a deep dive into model fairness:

---

Dive Deeper: Features & Concepts

For those who want to understand the nuts and bolts of OptiPFair.

Supported Models

OptiPFair is designed to work with transformer-based language models that use GLU architecture in their MLP layers. This includes most modern LLMs:

• LLaMA family (LLaMA, LLaMA-2, LLaMA-3)
• Mistral models
• Gemma models
• Qwen models
• Phi models
• ... and other models with a similar GLU architecture.

Neuron Selection Methods

OptiPFair provides multiple methods for calculating neuron importance, allowing you to tailor the pruning strategy to your specific needs:

1. MAW (Maximum Absolute Weight): The default and typically most effective method for GLU architectures. It identifies influential neurons based on the magnitude of their connections.
2. VOW (Variance of Weights): Identifies neurons based on the variance of their weights.
3. PON (Product of Norms): Uses the product of L1 norms to identify important neurons.

Expansion Rate vs. Pruning Percentage

You can define the pruning target in two intuitive ways:

1. Pruning Percentage: Directly specify the percentage of neurons to remove (e.g., 20%).
2. Expansion Rate: Specify the target MLP expansion rate as a percentage (e.g., target a 140% expansion rate, down from a model's original 400%). This is often more comparable across different model sizes.

---

🗺️ Roadmap & Community

The OptiPFair project is actively developed. Here's what's planned for the future.

Future Roadmap

Our goal is to make OptiPFair the go-to toolkit for efficient and fair model optimization. Key upcoming features include:

• Attention Pruning: Implementing pruning methods for attention heads and layers.
• Depth Pruning: Adding support for removing entire layers or blocks.
• Advanced Benchmarks: Integrating more comprehensive performance and evaluation benchmarks.
• GPU Optimizations: Creating a v2.0 with significant GPU-specific optimizations for faster execution.
• Large-Scale Model Support: Adding compatibility for DeepSpeed and FSDP to handle 70B+ models efficiently.

🤝 Contributing

Contributions are welcome! Whether it's bug reports, feature requests, or code contributions, please check out our contributing guidelines to get started.

Citation

If you use OptiPFair in your research or projects, please cite the library:

@software{optipfair,
  author = {Pere Martra},
  title = {OptiPFair: A Library for Structured Pruning and Bias Visualization of Large Language Models},
  year = {2024},
  url = {https://github.com/peremartra/optipfair}
}

License

This project is licensed under the Apache 2.0 License. See the LICENSE file for details.