BMC ook: A Task-agnostic Compression Toolkit for Big Models (2022)

Abstract: Recently, pre-trained language models (PLMs) have achieved great success on various NLP tasks and have shown a trend of exponential growth in model size. To alleviate the unaffordable computational costs brought by the size growth, model compression has been widely explored. Existing efforts have achieved promising results in compressing medium-sized models for specific tasks, while task-agnostic compression for big models with over billions of parameters is rarely studied. Task-agnostic compression can provide an efficient and versatile big model for both prompting and delta tuning, leading to a more general impact than task-specific compression. Hence, we introduce a task-agnostic compression toolkit BMCook for big models. In BMCook, we implement four representative compression methods, including quantization, pruning, distillation, and MoEfication. Developers can easily combine these methods towards better efficiency. To evaluate BMCook, we apply it to compress T5-3B (a PLM with 3 billion parameters). We achieve nearly 12x efficiency improvement while maintaining over 97% of the original T5-3B performance on three typical NLP benchmarks. Moreover, the final compressed model also significantly outperforms T5-base (a PLM with 220 million parameters), which has a similar computational cost. BMCook is publicly available at https://github.com/OpenBMB/BMCook .

## introduction as the sizes of pre-trained language models (plms) increase, especially after reaching 10 billion parameters @xcite @xcite @xcite , powerful intelligence capabilities emerge in these big models, supporting plms to accomplish tasks that previous smaller † corresponding authors q p d m tensorflow @xcite pytorch @xcite textpruner @xcite textbrewer @xcite bmcook (this work) to improve the efficiency of plms, model compression is an essential solution. there are several compression techniques, including model distillation @xcite , model quantization @xcite , and model pruning @xcite . based on these techniques, practitioners can conduct task-specific compression during finetuning @xcite and task-agnostic compression during pre-training @xcite . previous studies mainly focus on applying taskspecific compression for medium-sized plms with around one-hundred million parameters, such as bert base @xcite @xcite , while compressing large-scale plms with over billions of parameters is rarely studied. in this work, we focus on the task-agnostic compression of big models because it enables developers to utilize the powerful intelligence of big models with fewer computation resources for both prompting @xcite and delta tuning (aka parameter-efficient tuning) @xcite . both prompting and delta tuning are the current core approaches to drive big models. there exist two challenges for the taskagnostic compression of big models. first, big models require high compression rates to achieve affordable costs while existing compression toolkits only support one or two techniques as shown in in this work, we introduce bmcook, a taskagnostic compression toolkit for big models. bm-cook has three main characteristics: (1) zeroredundancy training. bmcook is developed on an efficient training toolkit, bmtrain foot_0 , which supports the zero-redundancy optimizer with offloading @xcite to handle the memory challenge. (2) flexible combination. to achieve better efficiency, we make bmcook flexible to support arbitrary combinations of different compression techniques. to this end, we implement four popular compression techniques and distribute each technique to different parts of one unified training life-cycle. (3) runtime model modification. since some compression techniques require to access the inner hidden states of plms, developers have to modify the code of model implementation provided by a third-party package. to make the compression easier to operate, bmcook implements runtime modification by monkey patch to get rid of modifying the source code of plms. we evaluate the effectiveness of bmcook on t5-3b @xcite , a t5 model with 3 billion parameters. experimental results show that bm-cook achieves nearly 12x compression efficiency by combining all four techniques while maintaining over 97% original performance on three typical nlp benchmarks, including sst-2 @xcite , mnli @xcite , and squad @xcite . besides, the compressed model significantly outperforms t5-base, which has similar computation costs. bmcook is supported by open lab for big model base (openbmb) foot_1 . we hope bmcook can help researchers explore better compression methods for large-scale plms in the future and help practitioners to improve their model efficiency in real-world applications. ## design and implementation as mentioned in the introduction, we implement three main characteristics in bmcook, zeroredundancy training, flexible combination, and runtime model modification. in this section, we will introduce the design and the implementation details of these three characteristics. ## evaluation to validate the effectiveness of bmcook, we study task-agnostic compression on t5-3b @xcite , which has 3 billion parameters. since task-agnostic compression would benefit various downstream tasks, we evaluate the performance of adapter tuning @xcite of t5-3b and its compressed variants. we also study t5-base and t5-large, which have 220 million and 770 million parameters, respectively. training and evaluation data. we use the pile dataset @xcite for task-agnostic compression training, which is a large-scale corpus for pre-training language models. the training objective is masked language modeling used by t5. note that we turn on distillation during the compression training in all experiments because we find knowledge distillation with mse loss can improve model performance in our pilot experiments. besides, we choose three downstream datasets for evaluation: sst-2 @xcite , a representative singlesentence classification dataset, mnli @xcite , a representative sentence-pair classification dataset, squad v1.1 @xcite , a representative question-answer dataset. for the first two datasets, we use accuracy as the evaluation metric. for the third dataset, we use both exact match and f1 score as the evaluation metrics. we evaluate model performance on their development sets. we adopt the same task templates and label words of the original t5 paper @xcite . hyper-parameters. the learning rate of taskagnostic compression training is 1e-4 while that of adapter tuning ranges from 1e-6 to 1e-5. the batch size of task-agnostic compression training and adapter tuning is 32. we use 4 nvidia a100 gpus in the experiments. the training step of taskagnostic compression training ranges from 10k to 50k according to the compression methods. the training epoch of adapter tuning ranges from 3 to 5. to fairly compare the efficiency of t5-3b and its variants, we define a new metric, named activated model size, because @xcite mentioned that the computation of transformer is linear in the model size, which excludes the embedding layer. hence, we consider the parameters of self-attention networks and ffns. for the original model, the activated model size is equal to its original model size. although it is intuitive to directly compare the speedup of compressed models, there is no inference toolkit supporting all the compressed methods. hence, we focus on the theoretical computational cost in this work. in the evaluation, we set the pruning sparsity to 50%, i.e., we prune 50% of the parameters and reduce half of the activated model size. besides, we quantize the parameters to 8 bits, which reduces three-fourths of the activated model size compared to the floating-point version. for moefication, we dynamically involve 50% of parameters in ffns for specific input. therefore, the activated model size of the modified ffns is about half of the original ffns. note that transformer consists of both attention layers and ffns and the model size of ffns are about 70%. the final activated model size of the modified transformer is about 66% of the original one. if we combine all three techniques, we can achieve a compressed model with about one-twelfth of the original activated size. we report the evaluation results in ## conclusion and future work in this paper, we introduce a task-agnostic compression toolkit for big models, named bmcook. this toolkit contains four popular techniques and is designed to be flexible to support arbitrary combinations. users can easily compress a plm by adding several lines to its pre-training code and specifying the strategy in a configuration file. in the future, there are three directions to further improve bmcook. first, we will enrich the options of existing compression techniques, such as knowledge distillation on attention matrices @xcite and extreme low-bit quantization @xcite . second, there are some other compression techniques that are not covered by bm-cook, such as weight sharing @xcite and low-rank decomposition @xcite . third, we will explore the automatic search for better compression strategies or configurations. given a specific computation budget, we want to find the compression strategy that achieves the best model performance, which is similar to neural architecture search @xcite . meanwhile, we will also plan to enrich the inference toolkits to support different compression techniques. although compression techniques have been fast developed, the inference toolkits are still lagging behind. recently, there are some efforts to support compressed models in inference, such as bminf @xcite and @xcite , while they are still limited to specific compression techniques.