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_128")

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:

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')

Wrapper class to hold a ConfigSpace Configuration object as a Smash configuration.

add_data(data: str | DataLoader | LightningDataModule | Dataset) None

Add data to the SmashConfig.

Parameters:

data (str | torch.utils.data.DataLoader | pytorch_lightning.LightningDataModule | torch.utils.data.Dataset) – The data to be added to the SmashConfig. String Options: MNIST, ImageNet, COCO, common_voice, pubchem_selfies, LAION256, inpaint, img2img, Polyglot, WikiText, OpenAssistant, C4, AIPodcast

Returns:

This method updates the internal configuration state but does not return a value.

Return type:

None

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

load_from_json(path: str) None

Load a SmashConfig from a JSON file.

Parameters:

path (str) – The file path to the JSON file containing the configuration.

Returns:

This method updates the internal configuration state but does not return a value.

Return type:

None

save_to_json(path: str) None

Save the SmashConfig to a JSON file, including additional keys.

Parameters:

path (str) – The file path where the JSON file will be saved.

Returns:

This method updates the file system but does not return a value.

Return type:

None