The NiuTrans System for the WMT 2021 Efficiency Task

Proceedings of the Sixth Conference on Machine Translation (WMT), pages 787–794 November 10–11, 2021. ©2021 Association for Computational Linguistics 787 The NiuTrans System for the WMT21 Efficiency Task Chenglong Wang1, Chi Hu1, Yongyu Mu1, Zhongxiang Yan1, Siming Wu1, Yimin Hu1, Hang Cao1, Bei Li1, Ye Lin1, Tong Xiao1,2 and Jingbo Zhu1,2 1NLP Lab, School of Computer Science and Engineering, Northeastern University 2NiuTrans Research, Shenyang, China clwang1119@gmail.com,huchinlp@gmail.com {xiaotong,zhujingbo}@mail.neu.edu.cn Abstract This paper describes the NiuTrans system for the WMT21 translation efficiency task1. Fol- lowing last year’s work, we explore various techniques to improve the efficiency while maintaining translation quality. We investigate the combinations of lightweight Transformer architectures and knowledge distillation strate- gies. Also, we improve the translation ef- ficiency with graph optimization, low preci- sion, dynamic batching, and parallel pre/post- processing. Putting these together, our system can translate 247,000 words per second on an NVIDIA A100, being 3× faster than our last year’s system. Our system is the fastest and has the lowest memory consumption on the GPU-throughput track. The code, model, and pipeline will be available at NiuTrans.NMT2. 1 Introduction Large and deep Transformer models have dom- inated machine translation (MT) tasks in recent years (Vaswani et al., 2017; Edunov et al., 2018; Wang et al., 2019; Raffel et al., 2020). Despite their high accuracy, these models are inefficient and dif- ficult to deploy (Wang et al., 2020a; Hu et al., 2021; Lin et al., 2021b). Many efforts have been made to improve the translation efficiency, including effi- cient architectures (Li et al., 2021a,b), quantization (Bhandare et al., 2019; Lin et al., 2020), and knowl- edge distillation (Li et al., 2020; Lin et al., 2021a). This work investigates efficient Transformers architectures and optimizations specialized for dif- ferent hardware platforms. In particular, we study deep encoder and shallow decoder Transformer models and optimize them for both GPUs and CPUs. Starting from an ensemble of three deep Transformer teacher models, we train various stu- dent models via sequence-level knowledge distil- 1http://statmt.org/wmt21/ efficiency-task.html 2https://github.com/NiuTrans/NiuTrans. NMT lation (SKD) (Hinton et al., 2015; Li et al., 2021a; Kim and Rush, 2016) and data augmentation (Shen et al., 2020). We find that using a deep encoder (6 layers) and a shallow decoder (1 layer) gives reasonable improvements in speed while maintain- ing high translation quality. We improve the stu- dent model’s efficiency by removing unimportant components, including the FFN sub-layers and multi-head mechanism. We also explore other model-agnostic optimizations, including graph op- timization, dynamic batching, parallel pre/post- processing, 8-bit matrix multiplication on CPUs, and 16-bit computation on GPUs. Section 2 describes the training procedures of the deep teacher models. Then, Section 3 presents various optimizations for reducing the model size, improving model performance and efficiency. Fi- nally, Section 4 details the accuracy and efficiency results of our submissions for the shared efficiency task. 2 Model Overview Following Hu et al. (2020), Li et al. (2021a) and Lin et al. (2021a), we use the SKD method to train our models. Our experiments also show that the SKD method …

Content truncated to 500 words