Skip to content

Hardware-Optimized width Pruning

Choose a tag to compare

View all tags
@peremartra peremartra released this 24 Nov 10:11
v0.2.1
255ac92

🎉 OptiPFair v0.2.1 - Hardware-Optimized Pruning

This release introduces the expansion_divisor parameter for hardware-optimized model pruning, enabling better GPU/TPU performance through aligned tensor dimensions.

✨ What's New

Hardware-Optimized Pruning with expansion_divisor

The new expansion_divisor parameter allows you to round intermediate layer sizes to specific multiples (32, 64, 128, or 256), optimizing pruned models for modern GPU and TPU architectures.

Quick Example:

from optipfair import prune_model

pruned_model = prune_model(
    model=model,
    pruning_percentage=20,
    expansion_divisor=128,  # Round to multiple of 128
    show_progress=True
)

Key Benefits:

  • 🚀 Better GPU performance through optimized memory access patterns
  • ⚡ Improved tensor core efficiency with aligned dimensions
  • 🎯 Flexible integration with both pruning_percentage and expansion_rate
  • 🔧 Simple to use - just one parameter

📚 New Resources

  • Example Notebook: expansion_divisor_example.ipynb - Complete tutorial with comparisons
  • Test Suite: Comprehensive tests in tests/test_expansion_divisor.py
  • Documentation: Updated README, LLM reference manual, and API docs

🔧 Technical Details

New Functions:

  • round_to_divisor(): Utility function for precise rounding to nearest multiple

Modified Functions:

  • prune_model(): Added expansion_divisor parameter
  • prune_model_mlp_glu(): Integrated validation and rounding logic
  • prune_neuron_pairs(): Applies rounding after pruning calculation

Validation:

  • Valid values: None (default), 32, 64, 128, 256
  • Requires either pruning_percentage or expansion_rate
  • Maintains bounds: result always ≥1 and ≤ original size

🔄 Compatibility

  • ✅ Fully backward compatible with v0.2.0
  • ✅ Works with all neuron selection methods (MAW, VOW, PON)
  • ✅ Compatible with both static and data-driven pruning
  • ✅ No breaking changes

📦 Installation

pip install --upgrade optipfair
# or
pip install optipfair==0.2.1

📖 Documentation

🙏 Acknowledgments

Thank you to the community for your feedback and contributions!

Full Changelog: https://github.com/peremartra/optipfair/blob/main/CHANGELOG.md

Assets 2
Loading