SmashConfig
SmashConfig is an essential tool in pruna for configuring parameters to optimize your models. This manual explains how to define and use a SmashConfig.
Defining a simple SmashConfig
Define a SmashConfig using the following snippet:
from pruna import SmashConfig
smash_config = SmashConfig()
After creating an empty SmashConfig, you can set activate a method by adding it to the SmashConfig:
smash_config['quantizers'] = ['hqq']
Additionally, you can overwrite the defaults of the method you have added by setting the hyperparameters in the SmashConfig:
smash_config['quant_hqq_weight_bits'] = 4
You’re done! You created your SmashConfig and can now pass it to the smash function.
Adding a Dataset, Tokenizer or Processor to the SmashConfig
Some methods require a dataset, tokenizer or processor to be passed to the SmashConfig. For example, the gptq quantizer requires a dataset and a tokenizer. We can pass them to the SmashConfig e.g. as follows:
from pruna import SmashConfig
smash_config = SmashConfig()
smash_config.add_tokenizer("facebook/opt-125m")
smash_config.add_data("WikiText")
As you can see in this example, we can add a dataset simply by passing the name of the dataset. However, the add_data() function also supports other input formats. For more information, see the class documentation below.
We can now activate the gptq quantizer by adding it to the SmashConfig:
smash_config['quantizers'] = ['gptq']
Similarly, we can add a processor to the SmashConfig if required, like for example by the c_whisper compiler:
from pruna import SmashConfig
smash_config = SmashConfig()
smash_config.add_processor("openai/whisper-large-v3")
smash_config['compilers'] = ['cwhisper']
If you try to activate a method that requires a dataset, tokenizer or processor and haven’t added them to the SmashConfig, you will recieve an error. Make sure to add them before activating the method! If you want to know which methods require a dataset, tokenizer or processor, you can look at the compression method overview.
SmashConfig Documentation
- class pruna.config.SmashConfig.SmashConfig(max_batch_size: int = 1, device: str = 'cuda', cache_dir_prefix: str = '/home/docs/.cache/pruna', configuration: Configuration | None = None)
Wrapper class to hold a ConfigSpace Configuration object as a Smash configuration.
- Parameters:
max_batch_size (int, optional) – The maximum number of batches to process at once. Default is 1.
device (str, optional) – The device to be used for smashing, e.g., ‘cuda’ or ‘cpu’. Default is ‘cuda’.
cache_dir_prefix (str, optional) – The prefix for the cache directory. If None, a default cache directory will be created.
configuration (Configuration, optional) – The configuration to be used for smashing. If None, a default configuration will be created.
- add_data(arg)
- add_data(dataset_name: str, *args, **kwargs) None
- add_data(datasets: list, collate_fn: str, *args, **kwargs) None
- add_data(datasets: tuple, collate_fn: str, *args, **kwargs) None
- add_data(datamodule: PrunaDataModule) None
Add data to the SmashConfig.
- Parameters:
arg (Any) – The argument to be used.
- add_processor(processor: str | AutoProcessor) None
Add a processor to the SmashConfig.
- Parameters:
processor (str | transformers.AutoProcessor) – The processor to be added to the SmashConfig.
- Returns:
This method updates the internal configuration state but does not return a value.
- Return type:
None
- add_tokenizer(tokenizer: str | AutoTokenizer) None
Add a tokenizer to the SmashConfig.
- Parameters:
tokenizer (str | transformers.AutoTokenizer) – The tokenizer to be added to the SmashConfig.
- Returns:
This method updates the internal configuration state but does not return a value.
- Return type:
None