<jupyter_start><jupyter_text>How to sparsify a Pytorch model<jupyter_code>from transformers import RobertaConfig
from transformers import RobertaForMaskedLM
from pytorch_block_sparse import BlockSparseModelPatcher
import re

config = RobertaConfig(
    vocab_size=52_000,
    max_position_embeddings=514,
    num_attention_heads=12,
    num_hidden_layers=6,
    type_vocab_size=1,
)

model = RobertaForMaskedLM(config=config).cuda()

# =>84 million parameters
print(f"Initial model parameters count={model.num_parameters()}")
 
# Create a model patcher
mp = BlockSparseModelPatcher()

# Show names that can be used: this returns a list of all names in the network that are patchable.
# These names are escaped to be used as regexps in mp.add_pattern()
patchables = mp.get_patchable_layers(model)

dedup_layers = []

# Pretty print the regexps: replace layer number with regexp matching numbers, and dedup them
# This is a bit specific to Roberta, but should work for most transformers, it's just for ease of reading.
for patchable in patchables:
    r = patchable["regexp"]
    r = re.sub(r'[0-9]+', '[0-9]+', r)
    if r not in dedup_layers:
        dedup_layers.append(r)
        layer = patchable['layer']
        print(f"{r}\n   => {layer.in_features}x{layer.out_features}, bias={layer.bias is not None}")


# Selecting some layers to sparsify.
# This is the "artful" part, as some parts are more prone to be sparsified, other may impact model precision too much.

# Match layers using regexp (we escape the ., just because, it's more correct, but it does not change anything here)
# the [0-9]+ match any layer number.
# We setup a density of 0.5 on these layers, you can test other layers / densities .
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.intermediate\.dense", {"density":0.5})
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.output\.dense", {"density":0.5})
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.attention\.output\.dense", {"density":0.5})
mp.patch_model(model)

print(f"Final model parameters count={model.num_parameters()}")

# => 68 million parameters instead of 84 million parameters (embeddings are taking a lof space in Roberta)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11252847380410023, 'fraction_numerical': 0.017312072892938495, 'mean_word_length': 4.674418604651163, 'pattern_counts': {'":': 3, '<': 5, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '140529', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 730, 'n_tokens_neox': 696, 'n_words': 224, 'path': 'doc/notebooks/ModelSparsification.ipynb', 'repo_name': 'owensgroup/pytorch_block_sparse', 'size': '2195'}