GKD: A General Knowledge Distillation Framework for Large-scale Pre-trained Language Model
Year: 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 memorylimited 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. 1

heading: 1 Introduction
text: Pre-trained language models, such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), 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 (Hinton et al., 2015), 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
*This work was done when the author visited Zhipu.AI. †Corresponding authors. The other authors also include Yang Yang, Hongyin Tang, Keqing He, Jiahao Liu, and Jingang Wang from Meituan. 1The code is available at https://github.com/aitsc/ GLMKD.
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 (Liu et al., 2022; Wu et al., 2022) and a comprehensive toolkit (Yang et al., 2020b). However, despite these rich research outcomes, there are still major challenges in deploying knowledge distillation systems for real-world industrial-strength applications, including:
• Obstacles to Distilling Ultra-large-scale PLMs. Contrary to distillation in controlled laboratory settings aimed at models with billions of parameters, many industrial-strength applications (Yu et al., 2022) rely on ultralarge-scale PLMs (on the order of 10B or even larger). The training of ultra-large-scale PLMs is already challenging, and the distillation process requires simultaneous training of both large and small models, leading directly to difficulties in distillation of ultra-large-scale PLMs. Furthermore, there are also methods (Wu et al., 2021a; Yuan et al., 2021) for distilling multiple large models into a single small model, which pose significant challenges in memory-constrained GPU environments.
• Obstacles to Switching Distillation Methods. Deploying a knowledge distillation system requires the implementation of numerous distillation methods to meet different requirements, but due to the differences in implementation of these methods, it is difficult to switch and combine them easily within a framework. It is important to have an architecture that accommodates a range of distillation meth-
134
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:
• Larger-scale Model Distillation. We propose a teacher-student parallel strategy based on advanced memory optimization methods, addressing the challenge of distilling ultralarge-scale PLMs (over 10B) due to memory constraints. The proposed strategy supports distillation of at least 100B-scale PLMs on 8 NVIDIA A100 (40GB) GPUs.
• More Compatible Method Architecture. 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.
• Easy-to-use Open Source Toolkit. We have open-sourced the required toolkit for GKD, which provides a command-line interface for 25 distillation methods, facilitating developers to deploy knowledge distillation systems for ultra-large-scale PLMs.

heading: 2 Related work
text: 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 (Jiao et al., 2020), attention layers (Wang et al., 2021), soft labels (Jafari et al., 2021), and hard labels (Jafari et al., 2022). These works validated their methods with PLMs of hundreds of millions of parameters, such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), XLNet (Yang et al., 2019), 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 (Liu et al., 2021) 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 (Shoeybi et al., 2019) based on model parallelism and Zero Redundancy Optimizer (ZeRO) (Rajbhandari et al., 2020) based on data parallelism.
While some code for knowledge distillation methods focused on language models was made public (Sanh et al., 2019; Jiao et al., 2020; Sun et al., 2020), there was a lack of a general framework for deploying knowledge distillation systems. TextBrewer (Yang et al., 2020b) 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 (Wu et al., 2022) requiring dynamic intervention of the intermediate layer computation in the model; SID (Aguilar et al., 2020) changing the intermediate layer features during training; Continuation-KD (Jafari et al., 2022) 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.

heading: 3 GKD
text: 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.

heading: 3.1 Overview Framework
text: Figure 1 shows the overview framework of GKD, which consists of six main processes:
(1) User Requirements: This process begins with the user specifying their requirements and forming a configuration file, which includes the choice of training task, distillation method, teacher model, student model, etc.
(2) Model Building: This process addresses the obstacles to distilling ultra-large-scale PLMs by implementing a teacher-student parallel strategy that combines Megatron-LM (Shoeybi et al., 2019) and ZeRO (Rajbhandari et al., 2020). The process involves selecting and executing parameter initialization strategies for the student model, such as initializing the student model with a pre-trained student, a truncated parameter teacher, random initialization methods, or other distilled students. It also includes initializing the training data with a tokenizer.
(3) Model Training: This process addresses the obstacles to switching distillation methods by implementing an efficient adaptive architecture that is compatible with various methods. This process includes the initialization of methods to extract and compute different model features based on the requirements of different methods at different iteration numbers.
(4) Multiple Training: This process is utilized for methods that require multiple training, such as task-specific methods (Jiao et al., 2020) that necessitate distillation in the task-specific stage after distillation in the pre-training stage.
(5) Analysis: This process confirms the compliance of the distilled student model with deployment requirements through analysis, such as examining the performance on the test set and other phenomena that can be utilized to enhance the model.
(6) Deployment: This process deploys the student model on the corresponding device, such as low-computing mobile devices or services with higher load deployment under equal computing power.
These six processes are performed in sequence to form the workflow of the knowledge distillation system. The greatest contribution of GKD lies in the design of the model building and training, as the other processes do not pose a challenge to the deployment of the knowledge distillation system. In the following sections, we will provide a detailed description of how GKD enables largerscale model distillation in the model building and more compatible method architectures in the model training.

heading: 3.2 Model Building
text: The challenge in building models lies in allocating ultra-large-scale PLMs, consisting of a student and one or more teacher models, on a GPU with only several tens of GB of memory. To address this challenge, we propose a teacher-student parallel strategy that splits the model parameters to different GPUs while preserving the feature distance computation between the teacher and student models. This strategy is inspired by the optimization of single ultra-large-scale PLMs, including MegatronLM (Shoeybi et al., 2019) which splits each parameter matrix in the transformer across multiple GPUs, and ZeRO (Rajbhandari et al., 2020) which partitions each layer of transformers sequentially across multiple GPUs.
As shown in Figure 2, we demonstrate the comparison between the previous strategy and our proposed teacher-student parallel strategy using an example. The example includes the allocation of
two 6-layer transformer teacher models and one 4- layer transformer student model on the GPU. The current methods allocate all the model parameters on each GPU, severely limiting the training of ultralarge-scale PLMs and multiple models. To reduce the memory usage on each GPU without compromising the interaction between the teacher and the student, our teacher-student parallel strategy evenly distributes the parameters of the teacher and student on different GPUs, with each GPU corresponding to the matching parameters of the teacher and student. With the model parallel and data parallel count being 2, the memory usage can be reduced by at least half. If utilizing ZeRO-Offload (Ren et al., 2021), the optimizer states can further be stored in CPU memory to reduce the utilization of GPU memory.

heading: 3.3 Model Training
text: The challenge in training models lies in how to easily switch and use different distillation methods within a single framework. To address this challenge, we propose an efficient adaptive architecture that is compatible with various methods. It implements the operation of different methods and the calculation of features through a dynamic hook
mechanism and an auxiliary model, respectively. As shown in the workflow in Figure 3, the dynamic hook mechanism constructs extraction hooks for extracting model features and operation hooks for modifying the model inference process during each iteration. These hooks are described by a configuration file similar to JSON, which only requires recording the operations required by the method and playing a role during the model inference process. The auxiliary model calculates the loss function based on these hooks and the returned model features. Table 1 describes the features that this architecture can adapt to existing methods.
It is worth noting that GKD can achieve method combination by integrating hooks from different methods. GKD can also record all model features through extraction hooks and save the distance of teacher and student features in the auxiliary model for later analysis of the correlation between feature distance and task performance in the distillation process.

heading: 4 Experiments
text: 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.

heading: 4.1 Experimental Setup
text: Datasets All methods that require distillation in the pre-training stage use BooksCorpus (Zhu et al., 2015) and English Wikipedia as training data (19GB). For the task-specific stage (fine-tuning), we evaluate different distillation methods using the more challenging SuperGLUE benchmark (Wang
et al., 2019).
Methods We tested 22 distillation methods specifically designed for language models, as well as three classic methods (KD, TAKD, and DGKD) from computer vision, which are listed in Tables 1 and 2. The implementation of the teacherstudent parallel strategy was carried out using the Megatron-LM (Shoeybi et al., 2019) and DeepSpeed (Rasley et al., 2020) framework.
Models The commonly used BERT (Devlin et al., 2019) lacks open-source ultra-large-scale PLMs, so we employed a more advanced GLM (Du et al., 2022), which boasts open-source models of 10Bscale or even 130B-scale (Zeng et al., 2023), significantly reducing the deployment cost of the knowledge distillation system. The scale of teachers and students are presented in Tables 2 and 3.
Refer to Appendix C for more implementation details.

heading: 4.2 Results
text: More Compatible Method Architecture To verify the proposed adaptive architecture can effectively be compatible with various methods, we tested 25 mainstream distillation methods and present the results in Table 2. The results demonstrate that these methods can be easily switched and utilized in GKD. It is worth noting that TinyBERT (without data augmentation) outperformed all the latest methods in our setup. This suggests that the latest methods may not necessarily be the most effective, and different requirements may necessitate different methods. Additionally, the reliability of GKD is further validated from the perspective of loss function values in Appendix B.1.
Larger-scale Model Distillation To verify the proposed teacher-student parallel strategy can support distillation of 100B-scale model on 8 NVIDIA A100 (40GB) GPUs, we present the memory and time consumption of different strategies for distilling models of varying scale in Table 3. The results indicate that previous strategies encountered GPU memory overflow when distilling 6B-scale models, whereas our strategy is capable of supporting the distillation of 100B-scale models. The results in rows 9, 10, and 11 respectively demonstrate that GPU memory consumption can be reduced through splitting the model parameters, optimizer states, or storing the optimizer states in CPU memory. If not limited to 8 GPUs, our strategy has the potential to distill even larger models. Appendix B.2 further examines the trade-off between memory and time consumption.

heading: 4.3 Further Exploration
text: In addition to compatibility with various methods, GKD also allows for effortless combination of different methods. In Appendix A.1, we have discovered a method that achieves SOTA results by combining the advantages of different distillation methods. Appendix A.2 presents a tool that analyzes the correlation between feature distance and task performance through GKD, enhancing the interpretability of the distillation process.

heading: 5 Conclusions
text: 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
adaptive architecture, allowing for the distillation of ultra-large scale PLMs (over 10B) and the switch of various advanced distillation methods. In the future, we plan to launch our knowledge distillation system for facilitating the mass production and deployment of student models.

heading: Acknowledgements
text: This work is supported by Technology and Innovation Major Project of the Ministry of Science and Technology of China under Grant 2020AAA0108400 and 2020AAA0108402, the Natural Science Foundation of China under Grant No. 61836013, the Major Program of the National Social Science Foundation of China under Grant No. 18ZDA032, and funds from CCFZhipu.AI and Beijing Academy of Artificial Intelligence (BAAI). The GPUs used are sponsored by Zhipu.AI.

heading: A Further Exploration
text: In this section, we further explore the capabilities of GKD in combining different distillation methods and enhancing the interpretability of the distillation process.
A.1 Method Combination Thanks to the dynamic hook mechanism, GKD is capable of combining methods by integrating hooks from different methods. As shown in Table 4, we demonstrate results from several dozen combinations of different model features. To conserve computational power, we set the batch size to 32 during pre-training and set the sizes of the teacher and student models to 110M and 22M, respectively. In the task-specific stage, the batch size and learning rate were fixed at 16 and 1e-5, respectively, without the use of grid search and seed averaging. Based on the results in Table 4, the following conclusions can be drawn.
(1) We discovered the method BestC which achieves the SOTA, outperforming TinyBERT by 1.24% on average in SuperGLUE. BestC combines
the features of TinyBERT, MiniLMv2, and soft labels. (2) The method that performs distillation in the pre-training stage (row 5) outperforms those using randomly initialized parameters (row 3) or truncated fine-tuned teacher parameters (row 4) in the pre-training stage. (3) The methods using soft labels in the pre-training stage (rows 14 and 17) outperform those not using soft labels (rows 12 and 16). (4) Starting from row 21, we compare
the results of various combinations distilled in the task-specific stage and not distilled (only trained on hard labels). We find that distillation in the taskspecific stage greatly improves the performance of the task.
A.2 Enhanced Interpretability Thanks to the adaptive architecture, GKD can record all model features through extraction hooks and save the distance of teacher and student fea-
calculated by first computing features H ← HHT√
dimensionality . Att1, HS1, Q1, K1, and V1 denote the attention
scores, hidden state, query matrix, key matrix, and value matrix of the first layer transformer, respectively. Soft and Emb denote the soft labels and output of the embedding-layer respectively. KL1, KL5, KL10, KL15, and KL20 denote the KL divergence with temperatures of 1, 5, 10, 15, and 20, respectively.
tures in the auxiliary model for later analysis of the correlation between feature distance and task performance in the distillation process. As an example of TinyBERT’s pre-training stage distillation, we present the Spearman and Pearson correlation coefficients between the feature distance and training loss, and between the feature distance and task performance, respectively, in Figures 4 and 5. The following conclusions can be drawn.
(1) The results shown in Figure 4 indicate that while TinyBERT trains its embedding layer, attention scores, and hidden state, many other features (e.g., the value matrix and features obtained after pair-wise scaled dot-product) also decrease in distance between teacher and student as the training loss decreases. This suggests that we may be able to find a way to automatically have a large number
of student features approach the teacher without having to distill all features, thus reducing the cost of distillation. (2) The results shown in Figure 5 indicate that the distillation in the pre-training stage of TinyBERT actually leads to a decrease in pre-training task performance. This suggests that the performance of pre-training tasks is not necessarily positively correlated with the performance of downstream tasks. It is noteworthy that there are a small number of features (e.g., soft labels) whose distance is related to task performance. Our hypothesis is that distilling features that are related to task performance may further improve task performance, and the third conclusion in Appendix A.1 supports this hypothesis.
KL1 KL5 KL10 KL15 KL20 MSE Spearman
Att1 Att2 Att3 Att4 Att5 Att6 Emb HS1 HS2 HS3 HS4 HS5 HS6
K1 K2 K3 K4 K5 K6 Soft Q1 Q2 Q3 Q4 Q5 Q6 V1 V2 V3 V4 V5 V6 0.16 -0.09 -0.08 -0.09 -0.09 -0.08 -0.19 -0.43 -0.08 -0.06 -0.06 -0.07 -0.07 -0.19 0.23 0.78 0.34 0.39 0.15 0.95 0.11 0.12 0.14 0.41 0.35 0.34 0.38 -0.04 0.05 0.15 0.18 0.15 -0.21 -0.09 -0.08 -0.08 -0.08 -0.08 -0.07 -0.45 -0.08 -0.07 -0.06 -0.07 -0.07 -0.22 0.18 0.46 0.24 0.19 0.0 0.77 0.16 0.1 0.02 0.26 0.29 0.26 0.37 -0.04 0.05 0.15 0.17 0.14 -0.18 -0.08 -0.08 -0.09 -0.08 -0.08 -0.08 -0.45 -0.08 -0.07 -0.06 -0.07 -0.08 -0.22 0.18 0.43 0.23 0.17 0.01 0.71 0.16 0.1 -0.01 0.24 0.29 0.24 0.37 -0.04 0.05 0.15 0.17 0.14 -0.16 -0.08 -0.08 -0.09 -0.08 -0.08 -0.08 -0.45 -0.08 -0.07 -0.06 -0.07 -0.08 -0.22 0.18 0.43 0.24 0.17 0.02 0.69 0.16 0.1 -0.02 0.24 0.29 0.24 0.37 -0.04 0.05 0.15 0.17 0.14 -0.15 -0.08 -0.08 -0.09 -0.08 -0.08 -0.08 -0.45 -0.08 -0.07 -0.06 -0.07 -0.08 -0.22 0.18 0.43 0.24 0.17 0.02 0.69 0.16 0.1 -0.02 0.24 0.29 0.24 0.36 -0.04 0.05 0.15 0.17 0.14 -0.11 -0.08 -0.08 -0.08 -0.08 -0.07 -0.08 -0.45 -0.08 -0.07 -0.06 -0.07 -0.08 -0.23 0.18 0.42 0.24 0.16 0.03 0.59 0.15 0.1 -0.03 0.24 0.29 0.23 0.37 -0.04 0.05 0.15 0.17 0.14
Distance
KL1 KL5 KL10 KL15 KL20 MSE Spearman
Emb
HS1
HS2
HS3
HS4
HS5
HS6
K1
K2
K3
K4
K5
K6
Q1
Q2
Q3
Q4
Q5
Q6
V1
V2
V3
V4
V5
V6
-0.29
0.68
-0.04
-0.09
-0.2
-0.16
-0.18
-0.12
-0.07
-0.06
-0.05
-0.07
-0.07
-0.23
-0.07
-0.06
-0.05
-0.06
-0.08
-0.74
-0.8
-0.62
-0.64
-0.91
-0.58
-0.23
0.52
0.14
-0.08
-0.09
-0.11
-0.09
-0.17
-0.26
-0.43
-0.06
-0.12
-0.15
0.66
-0.08
-0.09
-0.07
-0.07
-0.07
-0.44
-0.54
-0.26
-0.15
-0.56
-0.96
-0.16
0.51
0.04
-0.09
-0.09
-0.1
-0.1
-0.32
-0.28
-0.45
-0.07
-0.17
-0.18
0.8
-0.09
-0.1
-0.07
-0.07
-0.11
-0.37
-0.48
-0.2
-0.1
-0.45
-0.95
-0.13
0.5
0.02
-0.09
-0.09
-0.09
-0.1
-0.48
-0.28
-0.45
-0.06
-0.17
-0.18
0.79
-0.09
-0.1
-0.07
-0.07
-0.15
-0.37
-0.48
-0.18
-0.09
-0.41
-0.95
-0.11
0.49
0.01
-0.1
-0.08
-0.09
-0.1
-0.5
-0.28
-0.44
-0.06
-0.17
-0.18
0.75
-0.09
-0.09
-0.06
-0.07
-0.16
-0.38
-0.48
-0.17
-0.09
-0.39
-0.95
-0.08
0.46
-0.02
-0.1
-0.09
-0.1
-0.12
-0.39
-0.09
-0.47
-0.06
-0.09
0.05
0.04
-0.04
0.17
0.0
-0.05
0.12
-0.41
-0.48
-0.17
-0.13
-0.52
-0.79
Distance after pair-wise scaled dot-product
KL1 KL5 KL10 KL15 KL20 MSE Pearson
Att1 Att2 Att3 Att4 Att5 Att6 Emb HS1 HS2 HS3 HS4 HS5 HS6
K1 K2 K3 K4 K5 K6 Soft Q1 Q2 Q3 Q4 Q5 Q6 V1 V2 V3 V4 V5 V6
-0.1 -0.21 -0.22 -0.3 -0.31 -0.34 -0.38 -0.36 -0.38 -0.38 -0.38 -0.38 -0.38 0.17 0.18 0.58 0.49 0.49 0.33 0.83 0.52 0.4 0.39 0.54 0.47 0.49 0.02 -0.33 -0.08 0.02 0.01 0.01 -0.1 -0.22 -0.27 -0.33 -0.34 -0.36 -0.36 -0.36 -0.38 -0.38 -0.38 -0.39 -0.39 0.11 0.06 0.41 0.37 0.34 0.21 0.78 0.55 0.35 0.13 0.37 0.39 0.37 -0.03 -0.33 -0.08 0.01 0.0 0.0 -0.09 -0.21 -0.27 -0.33 -0.34 -0.36 -0.38 -0.36 -0.38 -0.38 -0.38 -0.39 -0.39 0.1 0.06 0.4 0.36 0.33 0.21 0.72 0.55 0.35 0.11 0.35 0.39 0.36 -0.04 -0.33 -0.08 0.01 0.0 0.0 -0.09 -0.21 -0.27 -0.34 -0.34 -0.36 -0.38 -0.36 -0.38 -0.38 -0.38 -0.39 -0.39 0.1 0.06 0.4 0.36 0.33 0.21 0.7 0.55 0.35 0.1 0.35 0.4 0.35 -0.04 -0.33 -0.08 0.01 0.0 0.0 -0.09 -0.21 -0.27 -0.34 -0.34 -0.36 -0.38 -0.36 -0.38 -0.38 -0.38 -0.39 -0.39 0.1 0.06 0.4 0.36 0.33 0.21 0.7 0.55 0.35 0.1 0.35 0.4 0.35 -0.04 -0.33 -0.08 0.01 0.0 0.0 -0.09 -0.21 -0.27 -0.33 -0.34 -0.35 -0.38 -0.36 -0.38 -0.38 -0.38 -0.39 -0.39 0.09 0.06 0.4 0.36 0.32 0.22 0.64 0.55 0.35 0.09 0.34 0.4 0.35 -0.03 -0.33 -0.08 0.01 0.0 0.0
Distance
KL1 KL5 KL10 KL15 KL20 MSE Pearson
Emb
HS1
HS2
HS3
HS4
HS5
HS6
K1
K2
K3
K4
K5
K6
Q1
Q2
Q3
Q4
Q5
Q6
V1
V2
V3
V4
V5
V6
-0.26
0.69
-0.19
-0.14
-0.13
-0.14
-0.14
-0.5
-0.26
-0.24
-0.27
-0.32
-0.34
-0.22
-0.23
-0.23
-0.27
-0.3
-0.34
-0.71
-0.69
-0.69
-0.66
-0.82
-0.5
-0.17
0.32
-0.12
-0.34
-0.4
-0.41
-0.41
-0.48
-0.33
-0.29
-0.31
-0.36
-0.35
0.07
-0.35
-0.28
-0.3
-0.31
-0.32
-0.43
-0.6
-0.56
-0.44
-0.64
-0.89
-0.26
0.21
-0.12
-0.31
-0.38
-0.39
-0.41
-0.52
-0.35
-0.29
-0.31
-0.36
-0.35
0.36
-0.34
-0.29
-0.3
-0.31
-0.32
-0.34
-0.55
-0.54
-0.42
-0.61
-0.86
-0.33
0.18
-0.13
-0.3
-0.36
-0.37
-0.4
-0.55
-0.35
-0.29
-0.31
-0.35
-0.35
0.29
-0.34
-0.29
-0.3
-0.31
-0.31
-0.31
-0.54
-0.53
-0.42
-0.6
-0.85
-0.36
0.17
-0.13
-0.3
-0.35
-0.36
-0.4
-0.54
-0.36
-0.29
-0.31
-0.35
-0.35
0.21
-0.34
-0.29
-0.3
-0.32
-0.3
-0.31
-0.54
-0.53
-0.42
-0.6
-0.84
-0.29
0.13
-0.07
-0.24
-0.31
-0.32
-0.34
-0.34
-0.51
-0.3
-0.31
-0.23
-0.03
-0.15
-0.52
-0.05
-0.35
-0.16
0.04
-0.25
-0.48
-0.52
-0.51
-0.65
-0.48
Distance after pair-wise scaled dot-product

heading: B Additional Analysis
text: In this section, we further verify the reliability of GKD from the perspective of loss function value, and analyze the balance of memory and time consumption in the teacher-student parallel strategy.
B.1 Are the Loss Values of GKD Normal? In order to further verify the reliability of GKD, we present the loss function values of each method at various distillation stages in Figure 6. The downward trend of all the loss values is consistent with our expectations, with two noteworthy observations: (1) MobileBERT and SID tend to gradually increase the number of distilled layers during training, hence the loss values exhibit an up-and-down trend. (2) The ReCoRD dataset, shown in task-
specific stages, was trained for 5 epochs, therefore some methods may show loss changes in stair-step fashion, such as Annealing-KD and Universal-KD.
B.2 Trade-off between Memory and Time Consumption
In order to speed up the training process while ensuring that the distillation process is not limited by GPU memory, we conducted a full combination of all optimization options to find the best balance between memory and time. Table 5 showcases the resource usage of 5B-scale and 10B-scale teacher models under different MP, DP, ZeRO, and Offload options during distillation. The results of the testing lead us to the following recommendations: In cases of insufficient GPU memory, ZeRO should
be considered first for partitioning the optimizer states and model gradients, followed by increasing the number of model parallelism, and lastly, using ZeRO-Offload to store the optimizer states and model gradients in CPU memory.
C Implementation Details
In this section, we provide further details regarding the hyperparameters and models to facilitate replication by developers.
C.1 Hyperparameters
The batch size, number of iterations, and peak learning rate for the pre-training stage were set to 64, 150000, and 4e-4, respectively. The taskspecific hyperparameters for specific methods were set to the optimal values from their corresponding papers, while other hyperparameters (see Table 6) were kept consistent with the fine-tuning teacher. For single-teacher methods in the taskspecific stage, grid search was used to optimize hyperparameters, including learning rate {5e-6,1e5,2e-5} and batch size {16,32}. Table 7 presents the learning rate and batch size for each method on each dataset in the SuperGLUE benchmark. The results for all methods were averaged over three random seeds.
C.2 Models
Table 8 shows the specific parameters of all the models utilized in this paper. The 110M, 340M, and 10B scale models are from GLM pre-trained models 2. The 293M-scale model with the MobileBERT structure (inverted-bottleneck structure) was obtained by us through a week of pre-training with 16 NVIDIA A100 (40GB) GPUs, and the 25Mscale model is also with the MobileBERT structure.
2https://github.com/THUDM/GLM
When conducting pre-training tasks, the models with the MobileBERT structure require the expansion of the token dimension, thus the actual number of parameters is greater than the scale. The other sized teacher models were tested with randomly initialized parameters to assess resource consumption. All the distillation processes were conducted using half-precision floating-point (fp16) models.

heading: Future Work
text: In the future, we plan to launch our knowledge distillation system for facilitating the mass production and deployment of student models