Pegasus Model Weights Compression/Pruning

I have trained pegasus for condition generation on a custom dataset. And have been successfully using it as well. My project aim is text summarization deployed on an heroku app. My model weights are over 2.5GB and heroku can’t support the same.

I trained pegasus on in legacy dir in transformers. I was thinking if pruning would help in reducing the model capacity significantly but still keeping the accuracy and performance intact.

I couldn’t try keras pruning due to the complexity of the trainer present in and I do not have any idea on how that works.

Any way you guys could help would be much appreciated.

hey @SteveMama, have you tried quantizing your model first? this might be a “quick win” to reduce the size of your model by 2-3x.

if that’s not good enough, then my suggestion would be to try movement pruning which is designed for fine-tuning language models and is implemented in the nn_pruning library: GitHub - huggingface/nn_pruning: Prune a model while finetuning or training.

(i’m not sure if pegasus is supported in the library, but you can post and issue to find out!)

Thanks for getting back @lewtun
the model is stored as a PyTorch model. so, I’ve been going through the docs of Pytorch, correct me if Im wrong but, is it possible to prune or reduce model size post-training and train again?

i think the answer strongly depends on the type of pruning method (structured vs unstructured), but it is possible in principle to iteratively prune a neural network - this was essentially the strategy in a now-famous paper by han et al:

now, when it comes to transformers, my current understanding is that you can either prune the attention heads post-training (see e.g. here), but if you want to keep training you’ll probably need to use something like movement pruning that i linked to above.

my personal experience is that pruning is still quite tricky to get right and requires some care to preserve the performance of the original model. i would suggest trying quantization or optimizations with frameworks like ONNX Runtime before devoting a lot of time into pruning

@lewtun Thanks a ton. Quantization did help me reduce the model by 3x. however, I can’t load it as a PegasusForConditionalGeneration model since the weights and layers no.s do not match and aren’t initialized. Is there a way, I can load the quantised model and avoid the problem.

@lewtun essentially, when I am training the load the quantized model I saved in my drive. This warning occurs. and post inference I find the summary generated in the form of numbers and not readable text.
How can this be fixed

hmm this is a bit odd :grimacing:

how did you do the quantization? i just remembered that someone already asked about quantizing pegasus here so maybe you can check whether you can dynamically quantize the model in a similar to how i described it there and then try generating some outputs with the same model in memory (i.e. don’t save and reload)

if that works, then my guess is that from_pretrained doesn’t support loading quantized models (i can have a look) and you might need to do the loading in native pytorch

@lewtun I chose dynamic quantization approach. And I dont think from_pretrained support loading quantised models.
I’d be happy if you can take a look. Can we hop on a google meet call and you can help me out.

hey @SteveMama, i had a closer look at from_pretrained and indeed it does not support loading quantized models because quantization changes the model’s state_dict (for example by introducing scale and zero_point parameters).

however, i think there is a work around that involves the following steps:

1. quantize your model

import torch
from transformers import PegasusTokenizer, PegasusForConditionalGeneration

# load fine-tuned model and tokenizer
model_ckpt = "google/pegasus-cnn_dailymail"
model = PegasusForConditionalGeneration.from_pretrained(model_ckpt)
# quantize model
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8

2. save the quantized model’s state_dict and config to disk

# save config
# save state dict
quantized_state_dict = quantized_model.state_dict(), "")

3. in your heroku app, create a dummy model using the saved config

from transformers import AutoConfig
# load config and dummy model
config = AutoConfig.from_pretrained("pegasus-quantized-config")
dummy_model = PegasusForConditionalGeneration(model.config)

4. quantize dummy model and load state dict

reconstructed_quantized_model = torch.quantization.quantize_dynamic(
    dummy_model, {torch.nn.Linear}, dtype=torch.qint8

from here you should be able to run reconstructed_quantized_model.generate and produce coherent outputs - let me know if you cannot.

an alternative is to use torchscript for the serialization as done in this pytorch tutorial (link), but i am not very familiar with torchscript and it seems somewhat complicated because you need to trace out the forward pass with some dummy inputs (seems to work for BERT but less sure about seq2seq models like pegasus)

hope that helps!