BMI nf: An Efficient Toolkit for Big Model Inference and Tuning (2022)

Abstract: In recent years, large-scale pre-trained language models (PLMs) containing billions of parameters have achieved promising results on various NLP tasks. Although we can pre-train these big models by stacking computing clusters at any cost, it is impractical to use such huge computing resources to apply big models for each downstream task. To address the computation bottleneck encountered in deploying big models in real-world scenarios, we introduce an open-source toolkit for big model inference and tuning (BMInf), which can support big model inference and tuning at extremely low computation cost. More specifically, at the algorithm level, we introduce model quantization and parameter-efficient tuning for efficient model inference and tuning. At the implementation level, we apply model offloading, model checkpointing, and CPU-GPU scheduling optimization to further reduce the computation and memory cost of big models. Based on above efforts, we can efficiently perform big model inference and tuning with a single GPU (even a consumer-level GPU like GTX 1060) instead of computing clusters, which is difficult for existing distributed learning toolkits for PLMs. BMInf is publicly released at https://github.com/OpenBMB/BMInf .

## introduction recent years have witnessed the great success of pre-trained language models (plms) @xcite in the nlp community. various techniques of plms enable us to train big models containing billions of parameters from large-scale unlabeled corpora in a self-supervised fashion. up to now, these big models (with billions of parameters like gpt-3 @xcite ) have achieved promising results on various nlp tasks and gained extensive attention from researchers. despite the success figure 1 : the overall framework of bminf. to make bminf convenient for users, the underlying implementation and the hardware adaptation will not be exposed to users, and these modules can be automatically executed. as shown in figure 1 , bminf is built based on a four-level framework, the most important part of which lies in its algorithm level and implementation level. at the algorithm level, we introduce model quantization to compress big models from high-bit floating-point parameters to low-bit fixed-point ones, which can significantly reduce the memory cost of big models. the faster computation speed of low-bit numbers can also accelerate the computation of big models. besides model quantization, we also introduce parameter-efficient tuning methods @xcite , which freeze the parameters of big models to reduce the computation and memory cost. by inserting additional learnable modules into big models, parameter-efficient tuning can tune these additional modules to help big models handle specific tasks. some recent works @xcite @xcite have shown that applying parameter-efficient tuning on big models can achieve results comparable to finetuning all model weights. at the implementation level, we implement model offloading and model checkpointing, which can make full use of cpu memory to store massive parameters of big models. moreover, model offloading and checkpointing can drop parameters and computation graphs during both the forward and backward propagation, which can further save gpu memory to operate more data. for the underlying arithmetic operators, we reimplement the mixed-precision cuda arithmetic operators, which can better utilize the tensor cores of gpus to further speed up the computation, especially accelerating the mixed-precision computation in model quantization. considering model offloading and model checkpointing bring extra cpu-gpu communication to load offloaded model weights, we perform cpu-gpu scheduling optimization to synchronously execute weight loading and model computation. this cpu-gpu scheduling optimization can alleviate the time waiting for weight loading. all of model offloading, model checkpointing, and parameter-efficient tuning can benefit from the scheduling optimization. due to the algorithm-level and implementationlevel efficiencies, bminf can work on various gpus at the hardware level, including both powerful gpus (e.g. tesla v100 and tesla a100) and consumer gpus (e.g. gtx 1060 and gtx 1080ti). in section 4, we will show that bminf can run models with more than 10 billion parameters on a consumer gpu gtx 1060, which is quite difficult for existing plm-related distributed toolkits such as megatron @xcite and deep-speed @xcite . at the model level, bminf supports various possible architectures @xmath0.4 ⇥ 10 38 < l a t e x i t s h a 1 _ b a s e 6 4 = " e h d q u q p l z a u t u p g 6 i y v x h k r l e f @xmath1.4 ⇥ 10 38 < l a t e x i t s h a 1 _ b a s e 6 4 = " c p y n f 2 x x d a s 7 u u e j 5 3 u t u b m l c w @xmath2 figure transformer-based plms, and users can choose their own model architectures for inference and tuning. to make bminf more convenient for users, the underlying implementation and the hardware adaptation are automatically executed and will not be exposed to users. in the following sections, we will show more details about bminf, especially at the algorithm level and implementation level. in this section, we briefly introduce how bminf supports big model inference and tuning in an efficient manner at the algorithm level, including model quantization and parameter-efficient tuning. ## implementations to reduce the computation and memory cost in this section, we briefly introduce how bminf reduces the computation and memory cost at the implementation level, including model offloading and model checkpointing, as well as cpu-gpu scheduling optimization. in fact, we also reimplement efficient mixed-precision arithmetic operators to better utilize gpu tensor cores. since the reimplementation of cuda operators is too detailed to be described in words, we will not show it here and recommend our readers refer to the source code. ## evaluation in this section, we present some evaluation results of big model inference and tuning to show the efficiency of our toolkit bminf. the following results are based on the model cpm-2 @xcite . cpm-2 is a chinese plm with over 10 billion parameters. since cpm-2 has an encoder-decoder architecture, it can be used for both text understanding and text generation. the original cpm-2 is implemented with the distributed toolkits deep-speed and megatron, which are currently the most efficient open-source tools for running big models. ## conclusion and future work in this paper, we introduce an efficient toolkit bminf to provide a way to use large-scale plms. by applying model quantization, parameterefficient tuning, model offloading, model checkpointing, cpu-gpu scheduling optimization, as well as the reimplementation of mixed-precision arithmetic operators, bminf can perform big model inference and tuning with less than 1/30 of the gpu memory and 10 times speedup, as compared with existing open-source distributed toolkits for pretraining and fine-tuning plms. in the future, our work to improve bminf will focus on the following three directions: we hope this toolkit can help researchers utilize big models for their own works and advance the adaption of big models in the nlp community.