GKD : A General Knowledge Distillation Framework for Large-scale Pre-trained Language Model (2023)

Abstract: Currently, the reduction in the parameter scale of large-scale pre-trained language models (PLMs) through knowledge distillation has greatly facilitated their widespread deployment on various devices. However, the deployment of knowledge distillation systems faces great challenges in real-world industrial-strength applications, which require the use of complex distillation methods on even larger-scale PLMs (over 10B), limited by memory on GPUs and the switching of methods. To overcome these challenges, we propose GKD, a general knowledge distillation framework that supports distillation on larger-scale PLMs using various distillation methods. With GKD, developers can build larger distillation models on memory-limited GPUs and easily switch and combine different distillation methods within a single framework. Experimental results show that GKD can support the distillation of at least 100B-scale PLMs and 25 mainstream methods on 8 NVIDIA A100 (40GB) GPUs.

## introduction pre-trained language models, such as bert @xcite , roberta @xcite , and their variants, have achieved excellent success in natural language processing (nlp) tasks when they usually have hundreds of millions of parameters. considering computationally expensive resource constraints, a wide range of real-world applications are often impeded. knowledge distillation @xcite , as a method for compressing large-scale pre-trained language models, is attracting more and more attention. as large-scale plms continue to grow in scale, and with advancements in knowledge distillation methods, it becomes increasingly pressing to apply knowledge distillation research in controlled laboratory settings to the real world. the field of knowledge distillation for language models has witnessed a phenomenal progress in recent years, particularly with regards to the reduction of model size, leading to the development of a plethora of sophisticated distillation techniques @xcite and a comprehensive toolkit @xcite . however, despite these rich research outcomes, there are still major challenges in deploying knowledge distillation systems for real-world industrial-strength applications, including: it is important to have an architecture that accommodates a range of distillation meth-ods while ensuring efficient training, such as avoiding excessive extraction of intermediate features that lead to memory waste. thus, a compatible and efficient architecture is crucial for successful deployment of knowledge distillation systems. to overcome these challenges, we present a general knowledge distillation framework (gkd) for deploying knowledge distillation systems that support various scale plms and methods. to overcome the obstacles to distilling ultra-large-scale plms, gkd leverages the techniques of training large transformer models to the distillation process that requires training multiple large (teacher) and small (student) models simultaneously, incorporating the latest model and data parallel strategies. to overcome the obstacles to switching distillation methods, gkd employs a dynamic hook mechanism and auxiliary model to extract and operate intermediate layer features and inference process of models in each iteration. while being compatible with various methods, it avoids the waste of memory caused by extracting all intermediate layers. gkd presents the first exploration of knowledge distillation for language models in industrial scenarios. specifically, our main contribution lies in: we propose an efficient adaptive architecture compatible with various methods, addressing the challenge of switching and using different distillation methods within a single framework with difficulty. the proposed architecture supports at least 25 model distillation methods. ## related work in recent years, knowledge distillation for compressing plms has gained increased attention. these works studied ways of better utilizing language model features for transferring knowledge from large teacher models to a smaller student model, involving hidden layers @xcite , attention layers @xcite , soft labels @xcite , and hard labels @xcite . these works validated their methods with plms of hundreds of millions of parameters, such as bert @xcite , roberta @xcite , xlnet @xcite , etc. however, deployment of the distillation system on gpus with limited memory has been hindered by the reliance on ultra-large-scale plms (10b or even larger). an offline distillation method @xcite that saved teacher features before training the student individually reduced memory pressure, but was limited to methods with smaller feature scales and without teacher-student interaction. in this work, gkd was compatible with ultra-large-scale plms distillation via the introduction of megatron-lm @xcite based on model parallelism and zero redundancy optimizer (zero) @xcite based on data parallelism. while some code for knowledge distillation methods focused on language models was made public @xcite @xcite , there was a lack of a general framework for deploying knowledge distillation systems. textbrewer @xcite packaged some abstract and simple distillation processes and loss functions, but lacked implementation of many methods and was difficult to adapt to increasingly complex distillation methods. there were significant differences in the implementation of these methods, such as diito @xcite requiring dynamic intervention of the intermediate layer computation in the model; sid @xcite changing the intermediate layer features during training; continuation-kd @xcite altering the loss calculation method as the epoch increased, and so on. these differences in implementation made it difficult for them to be easily switched and combined within a single framework, hindering the application of various advanced methods in knowledge distillation systems. in this work, gkd accommodated various advanced knowledge distillation methods through a dynamic hook mechanism and auxiliary models. ## gkd in this section, we first introduce the overview framework of the proposed gkd, then delve into the details of how gkd implements larger-scale model distillation and a more compatible method architecture, from the perspective of model building and training. ## experiments in this section, we verified that gkd, which is used for distillation of language models, can support at least 100b-scale parameters and 25 mainstream methods on 8 nvidia a100 (40gb) gpus. ## conclusions in this paper, we propose a general knowledge distillation framework, gkd, for deploying knowledge distillation systems targeting large-scale plms. gkd satisfies the demands of real-world applications by employing a parallel strategy and