Selective Smashing with Target Modules

Some algorithms let you target specific modules in your model via the target_modules parameter. It specifies exactly which parts of the model the algorithm should operate on.

TARGET_MODULES_TYPE = Dict[Literal["include", "exclude"], List[str]]

The parameter is a dictionary with two keys, include and exclude, each mapping to a list of pattern strings to match module paths.

A module is targeted if its path in the model matches at least one include pattern and does not match any exclude pattern.

Check out this tutorial notebook to learn more about how to use target_modules.

Pattern Format

Each of the include and exclude lists contains glob patterns, allowing you to match module paths like you would in a file search:

• * to match any number of characters (e.g., attention.* matches attention.to_q, attention.to_k, etc.)

• ? to match exactly one character

• [abc] to match any single character from the set (e.g., to_[qk] matches to_q and to_k)

Default Values

If target_modules is not provided (i.e., None), default values are inferred automatically from the model, configuration and algorithm used.

If a target_modules dictionary is provided but missing either the include or exclude key:

• Missing include: defaults to ["*"] (considering all modules)

• Missing exclude: defaults to [] (excluding no modules)

Usage Example target_modules

The following example shows how to use target_modules with the quanto quantizer to target your model’s transformer, excluding the embedding layers.

from pruna import SmashConfig

smash_config = SmashConfig()
smash_config["quantizer"] = "quanto"
smash_config["quanto_target_modules"] = {
    "include": ["transformer.*"],
    "exclude": ["*embed*"]
}

Previewing the Targeted Modules

You can preview the targeted modules by using the expand_list_of_targeted_paths function as shown in the example below:

from transformers import AutoModelForCausalLM
from pruna.config.target_modules import expand_list_of_targeted_paths

model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M")
target_modules = {
    "include": ["model.layers.[01].*attn.*"],
    "exclude": ["*v_proj"]
}
print(expand_list_of_targeted_paths(target_modules, model))

This will return the list of module paths that match the include and exclude patterns. In this example, the output contains the first two attention modules (model.layers.0.self_attn and model.layers.1.self_attn) and the q_proj, k_proj and o_proj layers inside them.

Note that this will list all modules that match the patterns, although some algorithms may only apply to the linear layers among those.