本文发表于2023年7月10日至12日举行的第61届计算语言学协会年会（Association for Computational Linguistics）的行业分会会议录中，卷5，页码134-148。研究题目为“GKD：一个用于大规模预训练语言模型的通用知识蒸馏框架”。主要作者包括*谭世成（安徽大学）、谭文岚（Zhipu.ai）、王元春（中国人民大学）、龚文文（清华大学）、赵姝†（安徽大学）、张鹏†（Zhipu.ai）、唐杰†（清华大学）。其他作者还包括来自美团的杨扬、汤鸿音、何克清、刘佳豪和王金刚。
这项研究属于自然语言处理领域，特别是模型压缩与知识蒸馏方向。其背景是，以BERT、RoBERTa为代表的大规模预训练语言模型虽在各项NLP任务中取得了巨大成功，但其庞大的参数量（动辄数亿甚至上百亿）导致其在计算资源受限的现实应用场景（如移动设备、高负载服务）中部署困难。知识蒸馏作为一种模型压缩技术，通过将大型“教师”模型的知识迁移到小型“学生”模型中，从而减小模型规模、降低推理成本，受到广泛关注。
然而，将知识蒸馏系统应用于真实的工业级场景面临两大严峻挑战。首先，超大规模模型蒸馏的障碍：工业应用常依赖百亿（10B）甚至更大规模的超大型预训练模型，而蒸馏过程需要同时训练庞大的教师模型和较小的学生模型，在GPU内存有限的条件下极难实现。其次，蒸馏方法切换的障碍：部署一个蒸馏系统往往需要实现众多不同的蒸馏方法来满足不同需求，但这些方法的实现差异巨大，难以在一个统一的框架内轻松切换和组合。此外，许多复杂方法需要动态干预模型中间层的计算或特征，容易造成内存浪费。
为解决这些挑战，研究团队提出了GKD，一个面向大规模预训练语言模型的通用知识蒸馏框架。其主要目标是：1）支持在内存有限的GPU上对超大规模PLMs（如100B级）进行蒸馏；2）提供一种兼容性高的架构，能够在一个框架内轻松切换、组合和使用多种不同的蒸馏方法，同时保证训练效率。
GKD的整体工作流程包含六个主要过程，其核心贡献在于“模型构建”和“模型训练”这两个过程的设计。
第三部分：详细工作流程 GKD的工作流程从用户需求到模型部署，依次包括以下六个过程： 1. 用户需求：用户通过配置文件指定需求，包括训练任务、蒸馏方法、教师模型、学生模型等选择。 2. 模型构建：此过程旨在解决超大规模模型蒸馏的内存挑战。GKD提出了一种师生并行策略。传统的策略是将所有模型（一个或多个教师和一个学生）的参数完整加载到每个GPU上，这严重限制了可训练的模型规模。GKD的师生并行策略灵感来源于训练单个超大模型的优化技术，它结合了基于模型并行（Model Parallelism）的Megatron-LM和基于数据并行（Data Parallelism）的ZeRO优化器。该策略的核心思想是将教师模型和学生模型的参数均匀地拆分并分配到不同的GPU上，每个GPU上存放的是教师和学生模型相对应的部分参数。例如，在图2示例中，通过将模型并行度和数据并行度均设为2，可以将两个6层Transformer教师和一个4层Transformer学生的内存占用至少减少一半。如果进一步利用ZeRO-Offload技术，还可以将优化器状态存储在CPU内存中，从而进一步降低GPU内存占用。通过这种策略，GKD能够在有限的内存资源下，构建起包含超大规模教师和学生的蒸馏环境。 3. 模型训练：此过程旨在解决多种蒸馏方法难以切换和组合的挑战。GKD设计了一种高效的自适应架构。该架构通过动态钩子机制和辅助模型来实现。具体工作流程（见图3）是：在每个训练迭代中，动态钩子机制会根据配置文件（类似JSON格式）构建两类钩子——提取钩子用于提取模型各层的特征（如隐藏状态、注意力分数等），操作钩子用于修改模型的推理过程（例如替换某些层，如Theseus方法）。配置文件只需记录方法所需的操作，这些钩子会在模型前向传播过程中被触发。随后，一个独立的辅助模型会接收这些钩子提取到的特征，并根据不同蒸馏方法的需求，计算相应的损失函数（如均方误差MSE、KL散度等）。这种设计的优势在于，它将方法的具体操作与模型的主体训练流程解耦，使得新增或切换方法只需修改配置文件即可，无需改动核心训练代码。同时，它避免了为兼容所有方法而预先提取全部中间特征所造成的内存浪费，实现了按需、高效的特征提取与计算。 4. 多次训练：某些蒸馏方法（如TinyBERT）需要分阶段训练，例如先在预训练阶段进行通用知识蒸馏，再在特定任务微调阶段进行任务特定知识蒸馏。此过程负责管理这类需要多轮训练的方法流程。 5. 分析：此过程对蒸馏得到的学生模型进行分析，例如在测试集上评估性能，或利用GKD记录的特征距离等信息分析蒸馏过程，以确认模型是否符合部署要求或为进一步优化提供 insights。 6. 部署：将最终满足要求的学生模型部署到目标设备上，如移动设备或服务器。
在实验部分，研究团队对GKD的两个核心能力进行了验证。
第四部分：主要结果 1. 更兼容的方法架构验证：为验证自适应架构的兼容性，研究者在GKD上成功实现了25种主流蒸馏方法，包括22种专门为语言模型设计的方法以及3种来自计算机视觉的经典方法（如KD、TAKD、DGKD）。如表2所示，所有方法均能在GKD框架下运行，并在SuperGLUE基准验证集上取得了预期的结果。一个有趣的发现是，在其实验设置下，TinyBERT（未使用数据增强）的性能超过了所有最新的方法。这表明，最新的方法不一定总是最有效的，实际应用中需要根据具体需求选择合适的方法，而GKD恰好为此提供了便捷的测试和切换平台。 2. 更大规模模型蒸馏验证：为验证师生并行策略的有效性，研究者测试了在不同策略下蒸馏不同规模模型时的内存和时间消耗（见表3）。结果表明，传统策略在蒸馏6B参数规模的教师模型时即发生GPU内存溢出。而采用GKD的师生并行策略，并结合模型并行（MP）、数据并行（DP）、ZeRO优化器状态分割以及ZeRO-Offload（将优化器状态卸载到CPU）等技术，可以在8块NVIDIA A100 (40GB) GPU上成功支持高达100B参数规模教师模型的蒸馏（从100B教师蒸馏到20B学生）。表3的第9-11行具体展示了通过分割模型参数、分割优化器状态、或将优化器状态存储于CPU内存，可以逐步降低单GPU的内存占用。研究还通过附录B.2对内存与时间消耗的权衡进行了详细分析，为使用者提供了配置建议：在内存不足时，应优先考虑使用ZeRO分割优化器状态和梯度，其次增加模型并行度，最后再考虑使用Offload到CPU。
进一步探索：得益于动态钩子机制，GKD可以轻松组合不同方法。研究者在附录A.1中尝试了数十种不同模型特征的组合方式，并发现了一个名为bestc的组合方法（结合了TinyBERT、MiniLMv2和软标签的特征），其在SuperGLUE上的平均性能超过了单一的TinyBERT方法1.24%。此外，得益于自适应架构，GKD能够记录训练过程中教师与学生模型所有特征的距离。附录A.2利用这一特性，分析了特征距离与训练损失、以及特征距离与任务性能之间的相关性，增强了蒸馏过程的可解释性。例如，分析发现，在TinyBERT的预训练蒸馏阶段，虽然只显式地蒸馏了嵌入层、注意力分数和隐藏状态，但其他许多特征（如值矩阵V）的距离也随着训练损失下降而减小。同时，预训练任务的性能与下游任务性能并非总是正相关，这为设计更高效的蒸馏目标提供了线索。
第五部分：结论与意义 本论文提出的GKD框架，是首次针对工业场景下大规模预训练语言模型知识蒸馏的系统性探索。其核心贡献和价值在于： * 科学价值与应用价值：GKD通过创新的师生并行策略和高效的自适应架构，解决了工业部署中大规模模型蒸馏和多种方法兼容的两大核心痛点。它使得在有限硬件资源下蒸馏百亿级模型成为可能，并极大地降低了尝试和组合不同先进蒸馏方法的工程门槛。 * 促进研究与生产：GKD作为一个开源工具包，提供了25种蒸馏方法的命令行接口，方便开发者和研究人员快速部署针对超大规模PLMs的知识蒸馏系统，有望推动知识蒸馏研究从实验室环境走向真实世界的规模化应用，促进高效学生模型的大规模生产与部署。
第六部分：研究亮点 1. 支持超大规模蒸馏：首次提出并实现了面向超大规模PLMs（>10B）的师生并行蒸馏策略，在8块A100 GPU上验证了支持至少100B规模模型的蒸馏能力。 2. 高度兼容与灵活：提出基于动态钩子机制和辅助模型的自适应架构，在一个框架内无缝支持了至少25种主流蒸馏方法，并能轻松进行方法组合，这在现有工具中是独一无二的。 3. 开源与易用性：研究团队开源了完整的GKD工具包，提供了简洁的命令行接口，极大地方便了工业界和学术界的使用与二次开发。 4. 增强可解释性：框架内置的特征记录与分析功能，为理解知识蒸馏的内部机制、探索更优的蒸馏目标提供了有力工具。
第七部分：其他有价值内容 除了核心框架，论文中展示的实验结果也为领域提供了有价值的参考。例如，不同蒸馏方法在相同设置下的性能对比（表2），为方法选择提供了实证依据；对不同并行和内存优化技术的详细性能剖析（表3、表5、附录B.2），为在实践中配置分布式训练提供了实用指南；对方法组合和特征相关性的探索（附录A），则为未来研究开辟了新的方向。