from transformers import GPT2LMHeadModel, GPT2TokenizerFast
from transformers.pytorch_utils import Conv1D
from fastai.text.all import *
from fasterai.sparse.all import *Prune Transformers
Prune transformers architecture with fasterai
Note
This example code is taken from the fastai docs
pretrained_weights = 'gpt2'
tokenizer = GPT2TokenizerFast.from_pretrained(pretrained_weights)
model = GPT2LMHeadModel.from_pretrained(pretrained_weights)path = untar_data(URLs.WIKITEXT_TINY)Let’s create our fastai Learner.
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), cbs=[DropOutput], metrics=Perplexity())And let’s try to extend a given prompt with the pretrained model.
prompt = "\n = Unicorn = \n \n A unicorn is a magical creature with a rainbow tail and a horn"preds = learn.model.generate(inp, max_length=40, num_beams=5)tokenizer.decode(preds[0].cpu().numpy())'\n = Unicorn = \n \n A unicorn is a magical creature with a rainbow tail and a horn on its head.\n\nA unicorn is a magical creature with a rainbow tail and a horn'print(learn.validate())[3.695716142654419, 40.27440643310547]learn.fit_one_cycle(1, 1e-4)| epoch | train_loss | valid_loss | perplexity | time |
|---|---|---|---|---|
| 0 | 3.128023 | 2.847160 | 17.238749 | 01:38 |
prompt_ids = tokenizer.encode(prompt)
inp = tensor(prompt_ids)[None]
preds = learn.model.generate(inp.cuda(), max_length=40, num_beams=5)
tokenizer.decode(preds[0].cpu().numpy())'\n = Unicorn = \n \n A unicorn is a magical creature with a rainbow tail and a horn @-@ shaped head . The unicorn is also known as a <unk> or <unk'Make it sparse !
Let’s now retrain, this time introducing sparsity. We reload the pretrained weights first, so that this arm and the dense one above are both one epoch of fine-tuning away from the same starting point.
model = GPT2LMHeadModel.from_pretrained(pretrained_weights) # start again from the pretrained weights
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), cbs=[DropOutput], metrics=Perplexity())Unfortunately, the transformer model uses a custom layer: Conv1D, which is not a part of PyTorch. To overcome this problem, we have to add this layer to our Granularities class, so that it knows what to sparsify.
Here, the Conv1D behaves like a Linear layer: its weight is a matrix of shape (nx, nf), input features by output features.
Conv1D(nf=8, nx=4).weight.shapetorch.Size([4, 8])We can thus add the Conv1D granularity by using the add_granularity method, indicating the target module and the corresponding granularities that it can handle (the same as Linear so we can reuse it)
Granularities.add_granularity(Conv1D, Granularities._granularities_Linear)Let’s now define our SparsifyCallback. We make the model 30% sparse, weight by weight, each Conv1D layer treated independently (context='local'), zeroing the lowest-magnitude weights — this is what large_final selects, as it keeps the largest ones.
sp_cb = SparsifyCallback(sparsity=0.3, granularity='weight', context='local', criteria=large_final, schedule=one_cycle, layer_type=Conv1D)We now only have to pass our callback to fastai
And we can check the predicion to the same prompt as before
prompt_ids = tokenizer.encode(prompt)
inp = tensor(prompt_ids)[None]
preds = learn.model.generate(inp.cuda(), max_length=40, num_beams=5)
tokenizer.decode(preds[0].cpu().numpy())'\n = Unicorn = \n \n A unicorn is a magical creature with a rainbow tail and a horn @-@ like head . It is a member of the <unk> family of unicorns'That’s it! You now have a sparse Transformer. On the wikitext-2 tiny test split, the pretrained GPT-2 scores a perplexity of 40.27; one epoch of fine-tuning brings it to 17.24, and one epoch of fine-tuning from the same pretrained weights, while making every Conv1D layer 30% sparse, gives 17.78. Single run per arm, bs=4, seq_len=256; seed spread not measured. The report counts 25,480,319 zeros over the 84,934,656 Conv1D parameters. The model is neither smaller nor faster as it stands: the zeroed weights are still stored and still multiplied. To turn sparsity into a speed-up, see the Granularity page.
Summary
| Tool | Purpose |
|---|---|
Granularities.add_granularity |
Register a custom layer type (Conv1D) so fasterai knows how to sparsify it |
SparsifyCallback(..., layer_type=Conv1D) |
Sparsify the registered layers during training |
large_final |
Criteria keeping the weights with the largest magnitude |
one_cycle |
Schedule used here for the sparsity ramp |
See Also
- Sparsifier - Core sparsification API
- SparsifyCallback - Callback for training integration
- Granularity - What gets removed, and which granularities can translate to a speed-up
- Schedules - All available sparsification schedules
- Sparsifier Tutorial - Basic sparsification walkthrough