A Compression-Compilation Framework for On-mobile Real-time BERT Applications

1 Introduction

Pre-trained large-scale language models such as BERT [Devlin et al., 2019], XLNet [Yang et al., 2019], RoBERTa [Liu et al., 2019] have substantially advanced the state-of-the-art across a wide spectrum of NLP tasks. With the increasing popularity of mobile AI applications and the concerns of information security and privacy, it is desirable to deploy these well-trained models on edge devices, and furthermore, to meet real-time requirements. However, these models often consist of hundreds (or even thousands) of computation layers and hundreds of millions of parameters. Therefore, how to accommodate the large and extremely deep models, such as BERT to edge device becomes an imminent problem.

There have been some efforts to compress the BERT model while maintaining the accuracy for downstream NLP tasks. MobileBERT [Sun et al., 2020] is able to reduce the memory requirement, but there is still a considerable execution overhead due to a large number of computation units, thus leading to high inference latency. Moreover, the large number of model layers also brings challenges in compiling models to mobile devices. To the best of our knowledge, only TensorFlow-Lite (TFLite) [Abadi et al., 2015] supports deploying BERT models on mobile CPU (not on mobile GPU), while no other frameworks can even support BERT models on mobile CPU.

∗These authors contributed equally

In this paper, we propose a compression-compilation codesign framework to optimize the structures of BERT variants for mobile devices. This is the first framework that involves compiler optimizations in the architecture search loop, aiming to co-optimize the model accuracy and computation resource usage. We also propose a highly effective layer fusion method to reduce intermediate results to achieve lower latency on both mobile CPU and GPU. Our framework outperforms the state-of-the-art framework, TFLite, by up to 7.8× speedup. Thus achieving the least latency while executing on mobile devices. We will release our model and framework.

Computational graph

a

b

Figure 2: a. Fusion example demonstrated in Figure 4. b. Sample fusion candidates for a computational graph section with an input (marked with A). Each layer has an input either from the previous layer/layers or from its weights, marked with other alphabets. Each number (from 1 to 4) denotes a fusion candidate (or fused block) based on mathematical properties.

2 Framework Design

There are two processes in CANAO: training and compiler code generation (as shown in Figure 3). The training process includes a controller and a trainer. The controller predicts/generates the model hyperparameters (i.e., network architecture); the trainer trains the predicted model and (quickly) evaluates its accuracy by fine-tuning the model to downstream tasks. The compiler code generation process takes the predicted model and returns execution information (e.g. number of fused layers, latency, CPU/GPU utilization).

The execution information together with the model accuracy from the training process will be feedback to the controller to improve the prediction of neural architectures. After the compiler-aware NAS, the generated codes by our optimized compiler will be deployed for mobile CPU/GPU executions.

For the training process, the controller generates the architectural hyperparameters of neural networks. This includes two phases: 1) The determination of the number of transformer blocks; 2) The optimization of size for each layer. We find that layer number affects the accuracy the most for BERT related models, thus it should be the first thing we determine when searching the optimized model architecture. Then we optimize the layer size by considering both inference latency and model accuracy, which are set as reward signals to feedback to the controller. The controller serves to find the optimal architecture by maximizing the expected reward.

The compiler code generation process includes three steps: 1) Generate a computational graph from the controllergenerated model and apply multiple optimizations on this graph. 2) Employ a novel compiler-based layer fusion optimization to further improve execution performance. This plays a key role in achieving better hardware efficiency. 3) Employ code generation and optimization to generate and further optimize the inference code. The generated inference code is tested on mobile devices. According to the feedback from the device side, the controller makes a better tradeoff between model accuracy and latency.

2.1 Controller Architecture Search

Our search space includes the number of layers, hidden layer size, and intermediate embedding size of the feedforward layers. We apply the recurrent neural network for searching the model architecture in the Controller. The recurrent network can be trained with a policy gradient method to maximize the expected reward of the sampled architectures. The accuracy and latency are used as the reward signal to feedback to the

1)

controller, which is trained by using the reinforcement learning method to explore the architecture. Our framework can search for a desirable model that achieves a good balance between accuracy and latency, preventing from searching the architecture manually.

2.2 Compiler Code Generation

This section introduces our compiler optimizations that optimize the latency reward for the feedback. There are two phases: Lightweight Polynomial-based Layer Fusion (LPFusion) and Polyhedral-based Code Generation.

LP-Fusion We identify all fusion candidates in a model based on two kinds of properties in the polynomial calculation: computation laws (i.e., associative, commutative, and distributive) and data access patterns. Fig. 2b shows four fusion candidates (or fused blocks) for a computational graph. Layer fusion reduces not only the memory consumption of intermediate results, but also the number of operators. Take Fig. 2b 3 for example, without layer fusion, the computation function is defined as:(⋆+F)⊙G+(⋆+F)⊙H. The layer and computation count numbers are 4 and 5, respectively. After fusion, the computation function is simplified as:(⋆+F)⊙ (G+H). Where layer and computation count numbers become 1 and 3, respectively. This process can significantly reduce the operator number and computation overhead. Compared with prior work on loop fusion [Ashari et al., 2015; Bezanson et al., 2017; Boehm et al., 2018], the novelty of this approach is that we exploit a restricted domain of DNN execution. Thus, we can enable more aggressive optimizations without very expensive exploration.

Polyhedral-based Code Generation LP-Fusion supports grouping multiple layers with varied output shapes, i.e., in the code-level, the nested loop structures of these layers may be different. Traditional compilers cannot support this kind of loop fusion well, mainly due to the complexity of such loop analysis. As shown in Fig. 2a and Figure 4, there are three operators: Mul-1, Mul-2, and Add. Mul-1 and Mul-2 take

func mul1: T *in0, T *in1, int row, int col, T *out for i = 0 to i < row for j = 0 to j < col let idx = i * col + j out[idx] = in0[idx] * in1[idx] func mul2: T *in0, T *in1, int col, T *out for j = 0 to j < col out[j] = in0[j] * in1[j] func fuse_add: T* in0, T* in1, T* in2, T* in3, int row, int col, T* out for i = 0 to i < row for j = 0 to j < col let idx = i * col + j out[idx] = in0[idx] * in1[idx] + in2[j] * in3[j] func fuse_add': T* in0, T* in1, T* in2, T* in3, int row, int col, T* out for j = 0 to j < col let temp = in2[j] * in3[j] for i = 0 to i < row let idx = i * col + j out[idx] = in0[idx] * in1[idx] + temp

matrices A and B as their input, respectively. Add takes the output of Mul-1 and Mul-2 as input to generate the result. The shape of matrix A is M × N while the shape of B is 1 × N. We have two options to perform loop fusion:fuse add, and fuse add'. In case fuse add, in2 and in3 are single dimensional arrays, so in2[j] * in3[j] incurs redundant computation for each outer-loop iteration (except the first one). The case fuse add' resolves this redundant computation with a proper loop permutation; however, it degrades the data locality because memory access for both matrices in0 and in1 becomes column-major, inconsistent with their memory storage. To address this complexity, our compiler extends the polyhedral analysis model [Wilde, 1993] to generate both versions and employs auto-tuning to dynamically select the optimal version. Moreover, our compiler also employs an extended polyhedral analysis model to analyze the loop structure and data dependency, transform indices, and generate optimal fused code for other loop fusion cases.

3 Experiments and Demonstrations

3.1 Training and Evaluation Setup

Our models are trained on a server with 16× NVIDIA Tesla V100 (Volta) GPUs. We use English Wikipedia [Devlin et al., 2019] and BooksCorpus [Zhu et al., 2015] to train the models and finetune on GLUE benchmark [Wang et al., 2018]: MNLI [Williams et al., 2018], SST-2 [Socher et al., 2013], MRPC [Dolan and Brockett, 2005], STS-B [Cer et al., 2017], RTE [Wang et al., 2018], and CoLA [Warstadt et al., 2019]. The sequence length is 128. We evaluate our framework on a Samsung Galaxy S20 cell phone with Qualcomm Snapdragon 865. For each model, we run our framework and TFLite 100 times with 8 threads on CPU and all pipelines on GPU.

3.2 Demonstration onMobile

Figure 1 shows the interface of our real-time BERT application. Figure 1 left is the Question Answering task. Type a random question that is related to the paragraph, it will automatically highlight the answer in the test. Figure 1 right is the Text Generation task. Given a starting sentence, it can automatically generate new sentences by word.

3.3 Evaluation Results

We compare the accuracy and latency of four models: BERTBASE [Devlin et al., 2019], MobileBERT [Sun et al., 2020], DistilBERT [Sanh et al., 2019], and CANAOBERT. Table 2 shows the accuracy and latency results. We manage to significantly reduce latency compared to BERTBASE, DistilBERT, and MobileBERT on both CPU and GPU. Compared with BERTBASE, our model is 5.2× faster on CPU and 4.1× faster on GPU with 0.5-2% accuracy loss. Compared with MobileBERT, our model is 1.49 × faster on CPU and 1.53× faster on GPU with only 0.4-1% accuracy decrease.

3.4 Effectiveness of Compiler Optimizations

We compare with a state-of-the-art framework, TFlite. Table 1 shows inference latency comparison results. TFLite only supports mobile CPU execution, and other frameworks do not support BERT models on mobile devices. And GPU performance is unusually slower than CPU (only 0.6× speedup for CANAOBERT over TFLite on CPU). The fully optimized framework can achieve up to 2.0× speedup on CPU, and 2.4× on GPU, over TFLite's CPU execution. Notably, comparing to BERTBASE on TFLite (352ms on CPU), our overall model and framework (45ms on GPU) can achieve up to 7.8× speedup.

4 Conclusion

We introduced a novel compression-compilation co-design framework to optimize the structures of BERT variants for mobile devices. We presented two BERT applications on mobile devices: Question Answering and Text Generation. Both can be executed in real-time with latency as low as 45ms.