面向超大规模模型的蒸馏优化方法、装置、介质及设备与流程

本发明属于知识蒸馏，具体而言涉及面向超大规模模型的蒸馏优化方法、装置、介质及设备。

背景技术：

1、通过知识蒸馏，大规模预训练语言模型的参数规模被大大减少，这极大地促进了其在各种设备上应用。然而，现实世界中的工业级应用与受控实验环境不同，用户需要在规模更大的语言模型上使用这些复杂的蒸馏方法，在内存受限的显卡上和多种方法的切换上部署会收到阻碍。

2、近年来，研究者对压缩大规模语言模型的知识蒸馏方法进行了广泛的研究和探讨。这些研究大多关注如何更好地利用语言模型特征将知识从大型教师模型转移到较小的学生模型，包括隐藏层、注意层、软标签和硬标签。这些研究使用了数亿个参数的预训练语言模型，例如bert、roberta、xlnet等，验证了其方法的有效性。然而，依赖于超大规模预训练语言模型（10b甚至更大）的蒸馏系统在gpu上的部署受到了限制。一种离线蒸馏方法，在训练学生之前保存教师特征，减少了内存压力，但受到特征规模较小且没有教师学生交互的方法的限制。

3、此外，虽然一些针对语言模型的知识蒸馏方法的代码已经公开，但缺乏一个通用框架来部署知识蒸馏系统。有一些已有的工作在统一的蒸馏工具上进行了尝试：textbrewer提供了一些抽象和简单的蒸馏过程和损失函数，但缺乏许多方法的实现，难以适应越来越复杂的蒸馏方法。这些方法在实现上存在显著差异，例如diito需要动态干预模型中间层的计算；sid在训练过程中更改中间层特征；continuationkd随着纪元增加而改变损失计算方法等等。这些实现上的差异使它们难以在一个单一框架中轻松切换和组合，从而阻碍了各种先进方法在知识蒸馏系统中的应用。

4、综上，现有技术中的知识蒸馏方法无法兼容多种多样的蒸馏方法以至于应用时难以选择当前应用场景最适当的蒸馏方法。

技术实现思路

1、鉴于上述的分析，本发明实施例旨在提供一种面向超大规模模型的蒸馏优化方法、装置、介质及设备，用以解决现有技术中的知识蒸馏方法无法兼容多种多样的蒸馏方法以至于应用时难以选择当前应用场景最适当的蒸馏方法的问题。

2、本发明第一方面实施例提供一种面向超大规模模型的蒸馏优化方法，包括以下步骤：

3、获取需求并根据所述需求形成蒸馏任务配置文件，所述蒸馏任务配置文件包括教师模型和学生模型；

4、基于所述教师模型和所述学生模型，根据教师学生并行策略进行模型构建以形成教师学生并行模型；

5、通过自适应架构对所述教师学生并行模型进行训练以获得学生模型；

6、所述自适应架构包括操作钩子、抽取钩子以及辅助模型；

7、所述通过自适应架构对所述教师学生并行模型进行训练以获得学生模型，包括：

8、依据所述蒸馏任务配置文件的蒸馏算法之一，通过所述操作钩子对所述教师学生并行模型的推理过程进行修改；

9、通过所述抽取钩子分别提取所述教师模型特征和所述学生模型特征；

10、将所述教师模型特征和所述学生模型特征以及所述抽取钩子均输入所述辅助模型；

11、依据所述抽取钩子确定损失函数，以确定所述教师模型特征和所述学生模型特征之间的差异进而获取损失，并依据所述损失对所述学生模型进行训练。

12、在一些实施例中，所述蒸馏任务配置文件包括训练任务、蒸馏算法集合、教师模型和学生模型。

13、在一些实施例中，所述基于所述配置文件，根据教师学生并行策略进行模型构建以形成教师学生并行模型，包括：

14、所述教师模型至少包括顺序排列的第一教师子模型和第二教师子模型，采用数据并行将所述第一教师子模型的各第一参数矩阵部署于第一gpu组，将所述第二教师子模型的各第二参数矩阵部署于第二gpu组；

15、所述学生模型至少包括顺序排列的第一学生子模型和第二学生子模型，所述第一学生子模型和所述第一教师子模型相对应，所述第二学生子模型和所述第二教师子模型相对应，采用数据并行将所述第一学生子模型的各第三参数矩阵部署于所述第一gpu组，将所述第二学生子模型的各第四参数矩阵部署于第二gpu组。

16、在一些实施例中，所述基于所述配置文件，根据教师学生并行策略进行模型构建以形成教师学生并行模型，还包括：

17、所述第一gpu组包括第一gpu和第二gpu，所述第二gpu组包括第三gpu和第四gpu；

18、采用相同的切分方式，将所述第一参数矩阵进行切分以至少获得第一子矩阵和第二子矩阵，将所述第二参数矩阵进行切分以至少获得第三子矩阵和第四子矩阵，将所述第三参数矩阵进行切分以至少获得第五子矩阵和第六子矩阵，将所述第四参数矩阵进行切分以至少获得第七子矩阵和第八子矩阵；

19、将所述第一子矩阵和所述第五子矩阵均部署在所述第一gpu，将所述第二子矩阵和所述第六子矩阵均部署在所述第二gpu，将所述第三子矩阵和所述第七子矩阵均部署在所述第三gpu，将所述第四子矩阵和所述第五子矩阵均部署在所述第四gpu。

20、在一些实施例中，所述基于所述配置文件，根据教师学生并行策略进行模型构建以形成教师学生并行模型，包括：

21、将所述教师模型的各参数矩阵进行切分以至少获得第一教师子矩阵和第二教师子矩阵，采用模型并行将所述第一教师子矩阵部署于第一gpu组，将所述第二教师子矩阵部署于第二gpu组；

22、采用所述教师模型的参数矩阵的切分方式将所述学生模型的各参数矩阵进行切分以至少获得第一学生子矩阵和第二学生子矩阵，采用模型并行将所述第一学生子矩阵部署于第一gpu组，将所述第二学生子矩阵部署于第二gpu组。

23、在一些实施例中，所述蒸馏算法包括ditto，lrc-bert，theseus，sid，tinybert，rail-kd，universal-kd，lrc-bert，annealing-kd，continuation-kd，mobilebert，ckd，mgskd，alp-kd，tmkd，mt-bert，rl-kd，uncertainty，takd，dgkd，kd，pd，pkd，distilbert，minilm，minilmv2中的一种或多种的组合。

24、在一些实施例中，所述优化方法还包括：

25、根据所述学生模型的性能进行评估分析，以判断所述学生模型的是否符合部署要求；

26、若是则将所述学生模型部署在设备上，否则继续对所述学生模型进行优化训练。

27、本发明第二方面实施例提供一种面向超大规模模型的蒸馏优化装置，包括：

28、配置模块，用于获取需求并根据所述需求形成蒸馏任务配置文件，所述蒸馏任务配置文件包括教师模型和学生模型；

29、模型构建模块，基于所述教师模型和所述学生模型，根据教师学生并行策略进行模型构建以形成教师学生并行模型；

30、模型训练模块，通过自适应架构对所述教师学生并行模型进行训练以获得学生模型；

31、所述自适应架构包括操作钩子、抽取钩子以及辅助模型；

32、所述通过自适应架构对所述教师学生并行模型进行训练以获得学生模型，包括：

33、通过所述操作钩子调用所述蒸馏算法，以对所述教师学生并行模型的推理过程进行修改；

34、通过所述抽取钩子分别提取所述教师模型特征和所述学生模型特征；

35、将所述教师模型特征和所述学生模型特征以及所述抽取钩子均输入所述辅助模型；

36、依据所述抽取钩子确定损失函数，依据所述教师模型特征和所述学生模型特征之间的差异确定损失，依据所述损失对所述学生模型进行训练。

37、本发明第三方面实施例提供一种电子设备，包括存储器和处理器，所述存储器存储有计算机程序，所述计算机程序被所述处理器执行时实现如上任一实施例所述的面向超大规模模型的蒸馏优化方法。

38、本发明第四方面实施例提供一种计算机可读存储介质，其上存储有计算机程序，所述计算机程序被处理器执行时实现如上任一实施例所述的面向超大规模模型的蒸馏优化方法。

39、本发明上述实施例至少具有以下有益效果：

40、1、通过基于内存优化的教师学生并行策略使得本发明实施例可以支持超大规模的大模型知识蒸馏。

41、2、通过自适应架构可以实现在蒸馏训练过程中的不同蒸馏算法的切换和组合，便于找到最适应当前蒸馏场景的蒸馏算法。