用于Transformers高效推理的联合剪枝、量化和蒸馏

2023年10月27日 由 alex 发表 480 0

介绍


预训练的Transformer模型被广泛用于各种自然语言处理任务,如文本分类、问答和生成任务。最近的趋势是模型在不断扩展的同时,性能也在不断提升。然而,Transformer的增长也导致部署所需的计算资源和能量大大增加。模型压缩的目标是在不显著降低准确性的情况下简化模型。剪枝、量化和知识蒸馏是深度学习模型中三种最流行的模型压缩技术。剪枝是一种通过减小模型大小来提高效率或性能的技术。量化可以显著减少存储和计算需求,通过减少表示数据所需的位数。知识蒸馏涉及到训练一个小模型来模仿一个大模型的行为。


9


OpenVINO™神经网络压缩框架(NNCF)开发了联合修剪、量化和蒸馏(JPQD)作为一个单一的联合优化流程,通过在预训练Transformer的迁移学习期间同时进行修剪、量化和蒸馏来改进其推理性能。JPQD减轻了开发者对不同压缩技术的顺序优化的复杂性,从而产生了一个经过优化的模型,显著提高了效率,同时保持了较好的任务准确性。JPQD的输出是一个在OpenVINO IR中结构修剪和量化的模型,可以通过在Intel平台上优化的OpenVINO运行时进行部署。Optimum Intel提供了一个简单的API,将JPQD集成到Hugging Face Transformers的训练流程中。


JPQD与Optimum Intel的BERT-base模型


在本文中,我们介绍如何将JPQD应用于GLUE基准上的BERT-base模型,用于SST-2文本分类任务。


这是一个压缩配置示例,格式遵循NNCF规范。我们在压缩算法和超参数的列表中指定修剪和量化。修剪方法与Movement Pruning(Sanh等人,2020)和Block Pruning For Faster Transformers(Lagunas等人,2021)的工作密切相似,用于非结构化和结构化的运动稀疏性。量化涉及量化感知训练(QAT)。在训练开始时,将使用这个配置对正在优化的模型进行修剪和量化操作。


compression_config = [
    {
        "compression":
        {
        "algorithm":  "movement_sparsity",
        "params": {
            "warmup_start_epoch":  1,
            "warmup_end_epoch":    4,
            "importance_regularization_factor":  0.01,
            "enable_structured_masking":  True
        },
        "sparse_structure_by_scopes": [
            {"mode":  "block",   "sparse_factors": [32, 32], "target_scopes": "{re}.*BertAttention.*"},
            {"mode":  "per_dim", "axis":  0,                 "target_scopes": "{re}.*BertIntermediate.*"},
            {"mode":  "per_dim", "axis":  1,                 "target_scopes": "{re}.*BertOutput.*"},
        ],
        "ignored_scopes": ["{re}.*NNCFEmbedding", "{re}.*pooler.*", "{re}.*LayerNorm.*"]
        }
    },
    {
        "algorithm": "quantization",
        "weights": {"mode": "symmetric"}
        "activations": { "mode": "symmetric"},
    }
]


10


上图显示了BERT-base模型在优化生命周期中的稀疏程度,包括两个主要阶段:


1. 非结构化稀疏化:在第一个阶段中,模型权重按照“sparse_structure_by_scopes”指定的区域逐渐稀疏化。BertAttention层(多头注意力:MHA)将以32x32的块大小进行稀疏化。相反,BertIntermediate和BertOutput层(前馈网络:FFN)将分别按行或列进行稀疏化。第一阶段作为预热阶段由参数“warmup_start_epoch”和“warmup_end_epoch”定义。参数“importance_regularization_factor”定义了权重重要性分数的正则化因子。在预热阶段之前,该因子为零,并在预热期间逐渐增加,最后在预热结束后保持固定值。用户可能需要一些启发式方法来找到稀疏性和任务准确性之间令人满意的权衡。


2. 结构化掩码和微调:第一个预热阶段将生成非结构化稀疏化模型。目前,针对4代英特尔® Xeon® Scalable处理器搭配OpenVINO 2022.3或更高版本,仅支持非结构化稀疏优化的推理加速。有关详细信息,请参阅使用OpenVINO和4th Gen Intel® Xeon® Scalable处理器加速稀疏Transformer模型的推理。但是,可以从模型中完全丢弃一些稀疏结构以节省计算和内存占用。NNCF提供了通过“enable_structured_masking”: true实现结构化掩码的机制,它会自动解决依赖层之间的结构化掩码,并重置不参与任务建模加速的稀疏参数。正如图2所示,“warmup_end_epoch”后稀疏程度下降,模型将继续进行微调。


已知限制:当前仅支持使用结构化剪枝和移动稀疏化对BERT、Wav2Vec2和Swin系列模型进行处理。


对于蒸馏,可以使用transformer API加载教师模型,例如来自Hugging Face Hub的BERT-large预训练模型。OVTrainingArguments通过添加蒸馏超参数(即蒸馏权重和温度)扩展了transformers的TrainingArguments,以便于使用。下面的代码片段展示了如何加载教师模型并创建带有OVTrainingArguments的训练参数。随后,教师模型与实例化的OVConfig和OVTrainingArguments一起提供给OVTrainer。其余的流程与内置的transformers的训练相同,但在内部应用了剪枝、量化和蒸馏。


from optimum. intel import OVConfig, OVTrainer, OVTrainingArguments
# Load teacher model
teacher_model = AutoModelForSequenceClassification.from_pretrained(
    teacher_model_or_path)
ov_config = OVConfig(compression=compression_config)
trainer = OVTrainer(
    model=model,
    teacher_model=teacher_model,
    args=OVTrainingArguments(save_dir, num_train_epochs=1.0, do_train=True,
                             do_eval=True, distillation_temperature=3, distillation_weight=0.9),
    train_dataset=dataset["train"].select(range(300)),
    eval_dataset=dataset["validation"],
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
    data_collator=default_data_collator,
    ov_config=ov_config,
    task="text-classification",
)
# Train the model like usual, internally the training is applied with pruning, quantization, and distillation
train_result = trainer.train()
metrics = trainer.evaluate()
# Export the quantized model to OpenVINO IR format and save it
trainer.save_model()


此外,NNCF还提供了其他任务的JPQD示例,例如问答。


基于BERT的演示到端到端JPQD


设置具有必要依赖项的Python环境。


conda create -n optimum_intel python=3.8
conda activate optimum_intel
python -m pip install --upgrade pip
python -m pip install "optimum-intel[openvino,nncf]"@git+https://github.com/huggingface/optimum-intel.git"optimum-intel[openvino,nncf]"@git+https://github.com/huggingface/optimum-intel.git
git clone https://github.com/huggingface/optimum-intel.git
cd optimum-intel/examples/openvino/text-classification
python -m pip install -r requirements.txt


在GLUE上使用JPQD的BERT运行文本分类示例。


TASK_NAME=sst2
python run_glue.py \
    --model_name_or_path bert-base-uncased \--model_name_or_path bert-base-uncased \
    --task_name $TASK_NAME \
    --teacher_model_name_or_path yoshitomo-matsubara/bert-large-uncased-sst2 \
    --nncf_compression_config ./configs/bert-base-jpqd.json \
    --distillation_weight 0.9 \
    --output_dir ./jpqd-bert-base-ft-$TASK_NAME \
    --overwrite_output_dir \
    --do_train \
    --do_eval \
    --max_seq_length 128 \
    --per_device_train_batch_size 32 \
    --learning_rate 2e-5 \
    --optim adamw_torch \
    --num_train_epochs 5 \
    --logging_steps 10 \
    --evaluation_strategy steps \
    --eval_steps 250 \
    --save_strategy epoch \
    --save_total_limit 3 \
    --fp16 \
    --seed 42


所有JPQD的配置和结果都保存在./jpqd-bert-base-ft-$TASK_NAME目录中。优化的OpenVINO IR用于在英特尔平台上进行高效推断。


对Xeon进行BERT-base性能评估和准确性验证


11


显示了基于BERT-base模型对文本分类任务性能评估和4th Gen Intel® Xeon® Scalable处理器上的准确度验证结果。BERT-base FP32模型作为基准。BERT-base INT8 (QAT)指的是仅使用8位量化方法进行优化的模型。BERT-base INT8 (JPQD)指的是通过剪枝、量化和蒸馏方法进行优化的模型。


在这里,我们使用性能提示“吞吐量”来评估输入序列长度为128的模型性能。


结果显示,与基准相比,BERT-base INT8 (QAT)已经可以达到2.39倍的压缩比和3.17倍的性能提升,而在SST-2上的准确度下降不明显(1.3%)。BERT-base INT8 (JPQD)可以进一步提高压缩比至5.24倍,从而实现4.19倍的性能改进,同时在SST-2上的准确度下降很小(<1%)。


12


通过适当的微调,JPQD甚至可以在提高性能的同时提高模型的准确性。上图显示了在第四代Intel® Xeon® Scalable处理器上对问答任务性能评估和准确性验证结果的BERT-base模型。与FP32基准相比,BERT-base INT8 (JPQD)可以将压缩率提高到5.15倍,同时提高SQuAD上的Exact Match (1.35%)和F1 score (1.15%)指标的性能至4.25倍。


13


上图显示了JPQD优化过的BERT-base模型在文本分类任务中每个层的参数计数可视化。你可以发现全连接层实际上是“密集”的,而大多数(多头注意力)MHA层与原始模型相比要稀疏得多。


14


上图显示了经JPQD优化的BERT-base模型在文本分类任务中每层的MHA头数,其中active(蓝色)表示保留的MHA头数,而pruned(灰色)表示删除的MHA头数。与在所有转换层中均匀剪枝不同,我们观察到JPQD倾向于在保留较低层的权重的同时大量剪枝最高层,这与Movement Pruning(Sanh等,2020)的实验结果类似。


结论


在本文中,我们介绍了一种联合剪枝、量化和蒸馏(JPQD)的方法,以加速英特尔平台上的转换器推断。以下是三个关键要点:


1. Optimum Intel提供了一个简单的API,可以将JPQD集成到训练流程中,以便在预训练转换器的迁移学习过程中同时进行剪枝、量化和蒸馏。将生成优化的OpenVINO IR以实现在英特尔架构上高效推断。

2. 用于文本分类任务的BERT-base INT8(JPQD)模型可以达到5.24倍的压缩比,在第四代Intel® Xeon® Scalable处理器上实现4.19倍的性能提升,并且与BERT-base FP32模型相比,在SST-2上保持了最小的准确率下降(<1%)。

3. 用于问答任务的BERT-base INT8(JPQD)模型可以达到5.15倍的压缩比,在第四代Intel® Xeon® Scalable处理器上实现4.25倍的性能提升,并且在SQuAD上改进了Exact Match(1.35%)和F1评分(1.15%),与BERT-base FP32模型相比。

文章来源:https://medium.com/openvino-toolkit/joint-pruning-quantization-and-distillation-for-efficient-inference-of-transformers-21333481f2ad
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消