DeepSpeed 模型压缩库

什么是 DeepSpeed 压缩:DeepSpeed 压缩是一个专门构建的库,旨在使研究人员和从业者能够轻松压缩模型,同时提供更快的速度、更小的模型尺寸以及大大降低的压缩成本。

为什么要使用 DeepSpeed 压缩:DeepSpeed 压缩提供新颖的先进压缩技术,以实现更快的模型压缩,并获得更好的模型质量和更低的压缩成本。DeepSpeed 压缩还采用端到端的方法,通过高度优化的推理引擎来提高压缩模型的计算效率。此外,我们的库具有多种内置的先进压缩方法。它支持这些方法和系统优化的协同组合,兼顾两者的优势,同时为高效的 DL 模型推理提供了一个无缝且易于使用的管道。我们强烈建议您阅读我们的博客,以了解更多关于(高级别)我们构建 DeepSpeed 压缩的原因以及它为用户提供的优势。

如何使用 DeepSpeed 压缩:第一部分通用教程将介绍库支持的压缩方法。以下部分将介绍我们关于如何组合不同的压缩方法来执行零成本量化 (ZeroQuant)极端压缩 (XTC) 的研究工作。除非另有说明,否则以下列出的实验结果基于 NVIDIA A100 GPU,并且我们在使用不同的 GPU 硬件时观察到略微不同的结果数字。

1. 通用教程

要使用 DeepSpeed 压缩库,您需要按照安装指南安装 DeepSpeed >= 0.7.0。目前,DeepSpeed 压缩包括七种压缩方法:通过知识蒸馏进行层减少、权重量化、激活量化、稀疏剪枝、行剪枝、头剪枝和通道剪枝。在以下小节中,我们将描述这些方法是什么、何时使用它们以及如何通过我们的库使用它们。

1.1 层减少

什么是层减少

神经网络由输入层、输出层和隐藏层构成。例如,BERT-base 语言模型由嵌入层(输入层)、分类层(输出层)和 12 个隐藏层组成。层减少意味着减少隐藏层的数量,同时保持网络的宽度不变(即,它不减少隐藏层的维度)。无论硬件和/或场景如何,此方法都可以线性减少隐藏层的推理延迟。

何时使用层减少

如果模型非常深,您可以考虑使用此方法。应用知识蒸馏时效果更好。层减少可以应用于预训练和微调阶段。前者生成一个蒸馏的与任务无关的模型,而后者生成一个特定于任务的蒸馏模型。在我们的 XTC 工作(论文教程)中,我们还讨论了何时应用层减少。

如何使用层减少

可以使用 DeepSpeed 配置 JSON 文件启用和配置层减少(配置详细信息)。用户可以通过keep_number_layer自由选择任何深度,并通过teacher_layer选择网络层的任何子集。此外,用户还可以选择是否通过other_module_name从给定模型(教师模型)重新初始化输入/输出层。

为了应用特定于任务的压缩的层减少,我们提供了一个关于如何为 BERT 微调执行此操作的示例。层减少是关于重置网络架构的深度和权重参数的重新初始化,这发生在训练过程之前。该示例包括对客户端代码(DeepSpeedExamples中的compression/bert/run_glue_no_trainer.py)的以下更改

(1) 初始化模型时,模型配置中的层数应与 DeepSpeed 配置 JSON 文件中的keep_number_layer相同。对于 Hugging Face BERT 示例,设置config.num_hidden_layers = ds_config["compression_training"]["layer_reduction"]["keep_number_layer"]

(2) 然后,我们需要使用从deepspeed.compression.compress导入的函数init_compression基于 DeepSpeed JSON 配置重新初始化模型。

(3) 在训练期间,如果未使用 KD,则无需执行任何操作。否则,在计算教师和学生的输出之间的差异时,需要考虑使用teacher_layer JSON 配置应用 KD。

可以通过以下方式在DeepSpeedExamples中运行我们的层减少示例:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/layer_reduction.sh

最终结果为:

Epoch: 18 | Time: 12m 38s
Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8340295466123281/0.8339096826688365

为了应用与任务无关的压缩的层减少,我们提供了一个关于如何在 GPT 预训练阶段执行此操作的示例。

步骤 1:获取最新版本的Megatron-DeepSpeed

步骤 2:进入Megatron-DeepSpeed/examples_deepspeed/compression目录。

步骤 3:运行示例 bash 脚本,例如ds_pretrain_gpt_125M_dense_cl_kd.sh。与预训练蒸馏相关的参数为:

(1)--kd,这将启用知识蒸馏。

(2)--kd-beta-ce,这指定了知识蒸馏系数。您通常可以将其设置为默认值 1,但有时调整此超参数会导致更好的蒸馏结果。

(3)--num-layers-teacher—hidden-size-teachernum-attention-heads-teacher,这些参数指定了教师模型的网络配置。请确保它们与检查点中的教师模型维度匹配。

(4)--load-teacher,这是指定教师模型检查点的位置。

(5)--load,这是将要加载的学生模型的初始检查点的位置。默认情况下,它将加载教师模型的底层以进行初始化,但您可以传递您自己的检查点以进行初始化。

除了上述配置之外,您可能还需要修改data_options中的数据路径,以便训练器知道数据位置。为了使事情稍微简单一些,我们提供了一些用于运行不同模型大小的蒸馏的示例脚本,包括 350M(ds_pretrain_gpt_350M_dense_kd.sh)和 1.3B 模型(ds_pretrain_gpt_1.3B_dense_cl_kd.sh)。我们还凭经验发现,分阶段 KD 通常会导致在后续任务上获得更好的预训练蒸馏模型。因此,我们建议一种通过不设置脚本中提供的--kd来提前停止 KD 的简单方法(例如,在剩余的 40% 的训练中禁用 KD)。

步骤 4:蒸馏模型后,还可以选择通过运行脚本125M-L10-Int8-test-64gpu-distilled-group48.sh进一步量化蒸馏模型,该脚本使用 INT8 量化器量化蒸馏模型的权重和激活(权重和激活量化将在以下部分介绍)。请注意,在执行量化时,需要设置-reset-iteration标志。我们在下表中提供了来自 WikiText-2 和 LAMBADA 的零样本困惑度结果。

GPT(125M) 层数 wikitex2 困惑度 LAMBADA
未压缩 12 29.6 39.5
仅量化 12 29.8 39.7
仅蒸馏 10 31.9 39.2
蒸馏 + 量化 10 32.28 38.7

1.2 权重量化

什么是权重量化

权重量化将全精度权重(FP32/FP16)映射到低位权重,例如 INT8 和 INT4。引用自此 Coursera 讲座:“量化涉及将模型转换为使用较低精度的参数和计算的等效表示。这提高了模型的执行性能和效率,但通常会导致模型精度降低”。

何时使用权重量化

一方面,再次引用自此 Coursera 讲座:“移动和嵌入式设备的计算资源有限,因此保持应用程序资源效率非常重要。根据任务的不同,您需要在模型精度和模型复杂度之间进行权衡。如果您的任务需要高精度,那么您可能需要一个大型且复杂的模型。对于需要较低精度的任务,最好使用更小、更简单的模型。”。另一方面,最近的服务器加速器(如 GPU)支持低精度算术。因此,将权重量化与激活量化(在后面的部分介绍)结合起来也可以提供更好的效率。

如何使用权重量化

可以使用 DeepSpeed 配置 JSON 文件启用和配置权重量化(配置详细信息)。我们想要指出的关键配置是:

(1)quantize_groups,组级权重矩阵量化:权重矩阵 W 被划分为多个组,每个组分别进行量化。请参阅本文了解更多详细信息。

(2)quantize_weight_in_forward必须设置为 true 用于 FP32 优化器训练,并设置为 false 用于 FP16。

(3)wq1/wq2,用户可以扩展更多组,例如wq3wq4等。

(4)start_bittarget_bit,为了简化第一个实验,我们建议将它们设置为相同,以便在迭代达到schedule_offset时对目标位应用量化。

客户端代码有两个更改(compression/bert/run_glue_no_trainer.pyDeepSpeedExamples 中)。

(1) 在模型初始化后,使用 DeepSpeed JSON 配置文件对模型应用 init_compression 函数。

(2) 训练完成后,应用 redundancy_clean 函数保存量化权重。

您可以在 DeepSpeedExamples 中运行我们的权重量化示例,方法如下:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/quant_weight.sh

最终结果为:

Epoch: 09 | Time: 27m 10s
Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8414671421293938/0.8422497965825875

1.3 激活量化

什么是激活量化

激活是指每一层的输入。激活量化将输入从全精度/半精度映射到低精度。更多信息请参见 这篇博文

何时使用激活量化

它可以提高计算效率,类似于 权重量化

如何使用激活量化

可以使用 DeepSpeed 配置 JSON 文件启用和配置激活量化(配置详细信息)。一些组件与权重量化相同,例如 schedule_offsetquantization_type。我们想指出的关键配置是

(1)range_calibration,用户可以选择动态或静态。使用“动态”时,激活量化组将自动设置为逐个token(对于基于Transformer的模型)和逐个图像(对于基于CNN的模型)。更多信息请参见 我们的ZeroQuant论文 和代码(deepspeed/compression/basic_layer.pyDeepSpeed 中)。

(2)aq1/aq2,用户可以扩展更多组,例如 aq3aq4 等。

客户端代码更改与 权重量化 相同。

您可以在 DeepSpeedExamples 中运行我们的激活量化示例,方法如下:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/quant_activation.sh

最终结果为:

Epoch: 02 | Time: 28m 50s
Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8375955170657158/0.8422497965825875

1.4 剪枝

什么是剪枝

剪枝旨在通过去除网络连接来减少生成预测所涉及的参数和操作数量。通过剪枝,您可以降低网络的整体参数数量(更多信息请参见 这节Coursera课程)。我们可以将剪枝策略分为两种类型:结构化剪枝和非结构化剪枝(更多信息请参见 这篇论文)。

方法 类型
稀疏剪枝 非结构化和结构化
行剪枝 结构化
头剪枝 结构化
通道剪枝 结构化

1.4.1 稀疏剪枝

什么是稀疏剪枝

稀疏剪枝意味着我们将每个权重矩阵中的一些元素设置为零值。根据用户选择的剪枝方法,零值可能具有结构化模式或非结构化模式。一种执行剪枝的方法是基于权重参数的绝对值,例如 这篇论文。另一种执行剪枝的方法是基于权重被掩盖时对损失函数的影响,例如 这篇论文

何时使用稀疏剪枝

如果您的模型参数过多,您可以考虑使用稀疏剪枝。但是,要看到硬件计算效率的真正好处,密度比(剪枝后保留的权重百分比)必须相当低。

如何使用稀疏剪枝

可以使用 DeepSpeed 配置 JSON 文件启用和配置稀疏剪枝(配置详细信息)。我们想指出的关键配置是

(1)schedule_offset,我们凭经验发现,当使用 method: topk 时,最好将 schedule_offset 设置为总训练步骤的较大值,例如 10%。

(2)method,我们支持 L1 范数、topk 和 snip_momentum 方法。欢迎用户贡献更多方法。

(3)sp1,用户可以扩展更多组,例如 sp2sp3 等。请注意,snip_momentum 方法不需要此配置。

(4)dense_ratio,对于非结构化稀疏剪枝,BRET-base 模型的密度比可以小于 0.1,同时仍能获得良好的准确率。对于 ResNet-50,密度比可以低至 0.3,同时在 ImageNet 上仍具有良好的准确率。对于像 snip_momentum 这样的结构化稀疏剪枝,密度比应在 shared_parameters 中指定,并用于计算全局稀疏率。

(5)frequencyblock_patternschedule_offset_end,它们用于指定剪枝的步长频率、块状剪枝模式(NxM 和 N in M)以及剪枝的结束步长。对于 snip_momentum 方法,这些配置是必须的。

客户端代码更改与 权重量化 相同。

您可以在 DeepSpeedExamples 中运行我们的稀疏剪枝示例,方法如下:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/pruning_sparse.sh

最终结果为:

Epoch: 02 | Time: 26m 14s
Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8416709118695873/0.8447925142392189

1.4.2 行剪枝

什么是行剪枝

行剪枝将权重矩阵某些行的所有元素设置为零值。如果剪枝了一行,则该行中的所有元素都将设置为零。

何时使用行剪枝

行剪枝有利于提高硬件速度,比稀疏剪枝好得多(但与稀疏剪枝相比,可能会导致更大的准确率损失)。它是一个专为两个连续的线性层(例如,Transformer 中的前馈网络)设计的特性。因此,我们建议将行剪枝用于第一个线性层(即 BERT 的 intermediate.dense 层)。减少此矩阵的行维度可以帮助减少后续矩阵的列(即 BERT 的 layer.\\w+.output.dense 层)。行剪枝也适用于其他类型的线性层。

如何使用行剪枝

可以使用 DeepSpeed 配置 JSON 文件启用和配置行剪枝(配置详细信息)。我们想指出的关键配置是

(1)method,目前仅支持 topk 方法。欢迎用户贡献更多方法。

(2)rp1,用户可以扩展更多组,例如 rp2rp3 等。

(3)related_modules,如“何时使用行剪枝”中所述,如果我们进行行剪枝,后续矩阵将受到影响。因此,需要了解模块之间的连接关系。

客户端代码更改与 权重量化 相同。

您可以在 DeepSpeedExamples 中运行我们的行剪枝示例,方法如下:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/pruning_row.sh

最终结果为:

Epoch: 02 | Time: 27m 43s
Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8440142638818136/0.8425549227013832

1.4.3 头剪枝

什么是头剪枝

头剪枝专为具有多头注意力的网络设计,例如基于Transformer的模型(更多信息请参见 这篇博文)。例如,BERT-base(BERT-large)模型具有 12 个头(24 个头)。

何时使用头剪枝

头剪枝有利于提高硬件速度。此外,如 这篇博文 所述:“在 这篇论文 中做出了令人惊讶的观察,即使在正常训练模型(使用所有头)之后,也可以在测试时移除许多头,而不会显着影响 BLEU 分数,事实上,在某些情况下,移除少量头会导致 BLEU 分数提高。”。

注意:头剪枝是专为注意力层(例如,Transformer 中的多头注意力)设计的特性。目前,它只能应用于 Transformer 的输出矩阵(即 BERT 中的 attention.output.dense)。剪枝输出矩阵也会导致 Query/Key/Value 矩阵被剪枝。

如何使用头剪枝

可以使用 DeepSpeed 配置 JSON 文件启用和配置头剪枝(配置详细信息)。我们想指出的关键配置是

(1)num_heads:用户需要为其模型提供正确的头部数量。

(2)modules:模块 attention.output.dense 是为 Hugging Face BERT 模型专门设计的。目前,当 Query/Key/Values 是分开的矩阵并紧随 attention.output.dense 时,我们仅支持这种情况。我们很乐意协助并欢迎对注意力模型的变体做出贡献。

(3)related_modules:如“何时使用头剪枝”中所述,剪枝注意力输出矩阵也会导致 QKV 矩阵被剪枝。因此,此处的输入为 [“self.query”, “self.key”, “self.value”]。

客户端代码更改与 权重量化 相同。

您可以在 DeepSpeedExamples 中运行我们的头剪枝示例,方法如下:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/pruning_head.sh

最终结果为:

Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8397350993377484/0.8377746135069162

1.4.4 通道剪枝

什么是通道剪枝

通道剪枝专为卷积层和计算机视觉设计。根据 wikipedia.org,“图像的颜色数据存储在三个值数组中,称为通道”。例如,一个具有三个通道的图像通过 ResNet-18 在第一层后会产生 64 个通道。

何时使用通道剪枝

通道剪枝是专为两个连续的 CONV2d 层(例如,ResNet 中的残差连接)设计的特性。因此,我们建议将通道剪枝用于第一个 CONV2d 层。减少此层的输出通道数可以帮助减少下一层的输入通道数。通道剪枝也适用于其他类型的 CONV2d 层。

如何使用通道剪枝

可以使用 DeepSpeed 配置 JSON 文件启用和配置通道剪枝(配置详细信息)。

您可以在 DeepSpeedExamples 中运行我们的通道剪枝示例,方法如下:

pip install torch torchvision
DeepSpeedExamples/compression/cifar$ bash run_compress.sh

最终结果为:

after_clean
epoch 10 testing_correct: 0.7664

请注意,以上结果是在“ResNet”模型中不使用批归一化(BN)时得到的。如果您对模型使用 BN 并应用通道剪枝,则清理模型后的验证结果将与清理前的模型不同。对于此类情况,我们建议用户在应用 redundancy_clean 后进一步微调模型。

2. ZeroQuant 教程:高效且经济实惠的训练后量化

在本节中,我们将介绍如何应用 DS-Compression 执行免费的 INT8 量化和轻量级的 INT4/INT8 混合精度量化。有关更多详细信息,请参阅 我们的论文

什么是 ZeroQuant

ZeroQuant 是一种高效的训练后量化方法,包括 (1) 针对权重和激活的细粒度硬件友好量化方案,可以显著减少量化误差;(2) 一种新颖且经济的逐层知识蒸馏算法 (LKD),即使无法访问原始训练数据;(3) 高度优化的量化系统后端支持,以消除量化/反量化开销。通过这些技术,ZeroQuant 能够 (1) 将模型量化为 INT8 而无需任何成本,以及 (2) 将模型量化为 INT4/INT8 混合精度量化,且资源需求极低(例如,BERT-base 量化需要 31 秒)。

何时使用 ZeroQuant

当您想要将基于Transformer的模型量化为 INT8 或 INT4/INT8 格式时,始终建议首先尝试 ZeroQuant,尤其是在模型对资源(GPU 和/或时间)要求很高的情况下进行量化感知训练,或者在无法访问原始训练数据时。

如何使用 ZeroQuant

您可以在 DeepSpeedExamples 中运行我们的 BERT 示例,方法如下:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt
DeepSpeedExamples/compression/bert$ bash bash_script/ZeroQuant/zero_quant.sh

最终结果为:

Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8427916454406521/0.8453010577705452

您可以运行我们的 GPT 示例,方法如下:

DeepSpeedExamples/compression/gpt2$ pip install -r requirements.txt
DeepSpeedExamples/compression/gpt2$ bash bash_script/run_zero_quant.sh

最终结果为:

Before converting the module COVN1D to linear and init_compression: 19.371443732303174
Before cleaning, Epoch at 0 with Perplexity: 19.47031304212775
After cleaning with Perplexity: 19.47031304212775

注意:目前,我们仅支持零成本量化。敬请关注 ZeroQuant 论文中提出的逐层知识蒸馏代码发布。

3. XTC 教程:用于极端压缩的简单而有效的压缩管道

在本节中,我们将介绍如何应用 DeepSpeed Compression 库来进行轻量级层减少和超低比特精度(二进制/三进制)量化。特别是,我们将指导您实现XTC 方法,具体包括:

(1) 获取具有 8 位激活量化的 1 位或 2 位 BERT-base(12 层)。

(2) 将 12 层的 Bert-base 减少到 5 层,然后获取其 1 位或 2 位对应版本。

什么是 XTC

XTC(eXTreme Compression 的缩写)是我们一种简单而高效的新方法,它通过轻量级层减少和鲁棒的二值化将模型压缩到极限。XTC 通过简单而有效的二值化技术将模型大小减少了 32 倍,同时 GLUE 任务的平均得分几乎没有损失。通过结合极端量化和轻量级层减少,我们可以进一步改进二值化模型,在保持 97% 准确率的情况下实现 50 倍的模型大小减少。有关更多详细信息,请参阅我们在论文中如何推导出我们的方法,我们在其中对当前用于极端压缩的各种技术的的影响进行了系统研究。

何时使用 XTC

如果您想在保持竞争性能的同时显著压缩您的模型,XTC 可能是理想的选择。它是一种简单且对超参数调整友好的方法。

如何使用 XTC

安装:BERT 模型的 XTC 极端压缩示例位于 DeepSpeedExamples 中的 compression/bert/bash_script/XTC。您需要通过以下方式安装依赖项:

DeepSpeedExamples/compression/bert$ pip install -r requirements.txt

XTC 方法的实现:为了适应没有经过微调的模型或特定于任务的压缩模型的用户,使用参数 --model_name_or_path yoshitomo-matsubara/bert-base-uncased-${TASK_NAME},我们的 Python 脚本 run_glue_no_trainer.py 会自动从 Hugging Face 下载模型。用户还可以使用自己的模型,这些模型具有更高的准确率,作为教师模型和学生模型的初始化。

3.1 使用 8 位激活量化的 1 位或 2 位 BERT-base(12 层)

有关配置,请参阅 DeepSpeedExamples 中的 compression/bert/config/XTC/ds_config_W1A8_Qgroup1_fp32.json。在我们的论文中,我们使用 FP32("fp16": {"enabled": false})进行训练,同时直接将 8 位量化("bits": 8)应用于激活,并将 1 位量化("start_bits": 1, "target_bits": 1)应用于注意力(查询、键、值)和前馈权重矩阵("modules": ["attention.self", "intermediate", "output.dense"])在训练开始时("schedule_offset": 0)。此外,我们还将 1 位量化应用于 word_embeddings 作为权重量化。

可以通过以下命令运行此示例:

DeepSpeedExamples/compression/bert$ bash bash_script/XTC/quant_1bit.sh

最终结果为:

Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8293428425878757/0.8396053702196908

我们想提到的另一个重要特性是 weight_quantization 中的 quantize_groups,这里将其设置为 1 以匹配我们 XTC 论文的 FP32 训练设置。我们发现,在 FP16 训练下,较少的量化组数(例如,1 或 2)可能导致训练不稳定。因此,我们建议在 FP16 下使用较多的组数(例如,64)。DeepSpeedExamples 中的 compression/bert/config/ds_config_W1A8_Qgroup64_fp16.json 是 FP16 示例配置,其中 "fp16": {"enabled": true}"weight_quantization": {"shared_parameters": {"quantize_weight_in_forward": false}} 与 FP32 情况不同。

使用此配置,我们对从 Hugging Face 下载的现有微调模型进行量化。对于 2 位权重量化,用户需要更新 ds_config JSON 文件。为了了解下载模型的压缩性能与我们论文的对比,我们在下表中收集了结果(在 MNLI 和 QQP 上使用 18 个训练 epoch 的 1/2 位 BERT)。本教程与论文之间的差异是由于它们使用了不同的检查点。在TinyBERT 中引入的数据增强对于较小的任务(如 mrpc、rte、sst-b 和 cola)将会有显著的帮助。有关更多详细信息,请参阅我们的论文

XTC quantization results

3.2 将 12 层 BERT-base 压缩到 1 位或 2 位 6/5 层 BERT

本节包含两部分:(a) 我们首先进行轻量级层减少,以及 (b) 基于 (a) 中的模型,我们进行 1 位或 2 位量化。

3.2.1 轻量级层减少

DeepSpeedExamples 中的 compression/bert/config/XTC/ds_config_layer_reduction_fp16.json 是将 12 层 BERT-base 减少到 6 层的示例配置。学生层的初始化来自教师的第 i 层,其中 i= [1, 3 ,5 ,7 ,9 ,11](注意层从 0 开始),这在我们的 XTC 论文中称为 Skip-BERT_5。此外,学生的模块(包括嵌入、池化器和分类器)也从教师处初始化。对于 5 层层减少,需要将 ds_config_layer_reduction_fp16.json 中的配置更改为 "keep_number_layer": 5"teacher_layer": [2, 4 ,6, 8, 10](如 compression/bert/config/ds_config_TEMPLATE.json 中所示)。

可以通过以下命令运行此示例:

DeepSpeedExamples/compression/bert$ bash bash_script/XTC/layer_reduction.sh

最终结果为:

Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8377992868059093/0.8365541090317331

值得注意的是,当使用单阶段知识蒸馏(--distill_method one_stage)时,教师模型和学生模型输出之间的差异(att_loss 和 rep_loss)也需要与初始化保持一致。请参阅 compression/bert/util.pyforward_loss 下的 _kd_function 函数。

对于 mnli/qqp,我们设置 --num_train_epochs 36--learning_rate 5e-5,并使用上述 JSON 配置。结果如下所示(我们还包括了 fp16 训练结果)。使用 fp32 明显比 fp16 产生更稳定的性能,尽管 fp16 可以加快训练时间。

XTC layer reduction results

3.2.2 6 层(5 层)BERT 的 1 位或 2 位量化

在上述层减少模型准备就绪后,我们现在继续使用 1/2 位量化压缩模型。DeepSpeedExamples 中的 compression/bert/config/XTC/ds_config_layer_reduction_W1Q8_fp32.json 是示例配置,其中我们在 compression/bert/config/XTC/ds_config_W1A8_Qgroup1_fp32.json 的基础上将层减少设置为 true。除了配置之外,我们还需要使用脚本 compression/bert/bash_script/XTC/layer_reduction_1bit.sh 中的 --pretrained_dir_student 更新学生模型的路径。用户可以通过添加 --pretrained_dir_teacher 使用不同的教师模型进行训练。

可以通过以下命令运行此示例:

DeepSpeedExamples/compression/bert$ bash bash_script/XTC/layer_reduction_1bit.sh

最终结果为:

Epoch: 18 | Time: 18m 11s
Clean the best model, and the accuracy of the clean model is acc/mm-acc:0.8140601120733572/0.8199755899104963

使用上述命令,现在可以获得 1 位 6 层模型的结果。现在,我们在下表中列出了更多关于 2/1 位 6/5 层模型的结果。请注意,我们用于以下压缩的检查点来自 3.2.1 节中的上表。

XTC 6-layer and quantization

XTC 5-layer and quantization

更新: