feat: add target modules to bnb quantizers by gsprochette · Pull Request #333 · PrunaAI/pruna

gsprochette · 2025-09-05T16:23:34Z

Description

Add Target Modules to BnB quantizers. To do so, I also introduced a routine to get the ignored modules, and restrict any target modules to only leaf modules (Linear, Conv1d, but not SelfAttention for example). These new functionalities are usefull to apply target modules to HF based quantizers.

Related Issue

Fixes #(issue number)

Type of Change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

How Has This Been Tested?

Added new tests counting the number of quantized modules before and after, + notebook

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Additional Notes

simlang

Looks amazing! Mostly some naming comments, and a general discussion how we should handle defaults, let me know what you think!

simlang · 2025-09-08T08:05:21Z

+    def get_unconstrained_hyperparameter_defaults(
+        self, model: Any, smash_config: SmashConfig | SmashConfigPrefixWrapper
+    ) -> TARGET_MODULES_TYPE:
+        """
+        Get default values for the target_modules based on the model and configuration.
+
+        Parameters
+        ----------
+        model : Any
+            The model to get the default hyperparameters from.
+        smash_config : SmashConfig
+            The SmashConfig object.
+
+        Returns
+        -------
+        TARGET_MODULES_TYPE
+            The default target_modules for the algorithm.
+        """
+        prefix: str
+        if hasattr(model, "transformer"):
+            prefix = "transformer."
+        elif hasattr(model, "unet"):
+            prefix = "unet."
+        else:
+            prefix = ""
+        return {"include": [prefix + "*"], "exclude": [prefix + "lm_head"]}


maybe a more general discussion on the defaults - the defaults here contain some key configuration so the model doesn't break from quantization, which we hide from the user. a user might decide to exclude some specific module for their use case but doesn't know about those defaults are necessary for the model to work and wonder why the model breaks.
Should we join the defaults with the user config instead of overwriting them?

One option I saw is a add_to_default flag, if it's true then you the target modules is dict(include = user_defined_include + include_default, exclude = user_defined_exclude + exclude_default). I have that in mind it's coming in a future PR

simlang

Bugbot caught already a lot, just some renaming-issues.

simlang

Looks good! tysm!
LGTM! 🚀

begumcig

Looks super good to me Gaspar!! Thank you so much for taking the time multiple times to explain everything 💜💜