Bert/Transformer模型压缩与优化加速

前言

Bert/Transformer结构及其变体,已成为自然语言处理 (NLP)、语音识别 (ASR)等领域的主流序列建模结构。并且,相比于卷积操作的局部特征提取能力、以及平移不变性,Self-Attention/Cross-Attention的全局Context信息编码或长时依赖建模能力,能够使图像/视觉模型具备更强的相关性建模能力(更好的全局视野)、与动态建模能力(领域或样本适应性)。因此在计算机视觉领域 (CV),Vision Transformer模型结构 (如ViT、Deit、Swin-Transformer、PVT、SegFormer与DETR等)也日益流行、方兴未艾。然而,Transformer模型的高计算复杂度与参数量,限制了模型的推理部署(包括服务端与移动端应用),其计算、存储与运行时内存等资源开销都成为限制因素(存在大量密集计算算子、与访存算子,如BatchMatMul、LayerNorm、Transpose/Reshape等)。例如,从标准Transformer layer的Tensor处理结构来看,MHA的Context特征计算与特征长度的平方、Embedding特征维度成正比:

  • Standard Self-Attention (X=Y) / Cross-Attention (X!=Y):

Bert/Transformer模型压缩与优化加速_第1张图片

  •  Standard FFN:

 本文从以下几个维度举例说明,解析Bert/Transformer模型的加速方法,体现了从建模依赖、数据依赖到硬件依赖的优化层次:

  • 模型结构精简化与知识蒸馏;
  • 模型量化 (Quantization);
  • 网络结构搜索 (NAS: Network Architecture Search);
  • 计算图优化 (计算图等价变换);
  • 推理优化引擎 (Faster Transformer / Torch-TensorRT / AICompiler);

Bert/Transformer模型压缩与优化加速_第2张图片

Bert/Transformer模型压缩与优化加速_第3张图片

模型结构精简化与知识蒸馏

Lite-Transformer

参考文献: https://arxiv.org/abs/2004.11886

Lite Transformer是一种融合卷积与Self-Attention操作的、高效精简的Transformer结构,可应用于NMT、ASR等序列生成任务。其核心是长短距离注意力结构 (LSRA:Long-Short Range Attention),能够有效融合卷积与Attention机制。LSRA将输入Embedding沿Feature维度split成两部分,其中一部分通过GLU (Gate Linear Unit)、一维卷积,用以提取局部Context信息;而另一部分依靠Self-attention,完成全局相关性信息编码。一维卷积的引入,能够减轻计算量与参数量。Lite Transformer核心结构如下,首先将FFN的宽度摊平 (flatten),其次引入LSRA以替换Self-Attention:

Bert/Transformer模型压缩与优化加速_第4张图片

Bert/Transformer模型压缩与优化加速_第5张图片

SAN-M

参考文献: https://arxiv.org/abs/2006.01713

SAN-M表示Self-Attention与DFSMN记忆单元的融合,是一种Transformer ASR模型。DFSMN适合捕获局部信息,Self-Attention模块具备较强的长时依赖建模能力,因此二者存在互补性。SAN-M通过将两个模块的特性融合在一起,实现了优势互补。Biasic Sub-layer表示包含了SAN-M的Self-Attention模块,DFSMN添加在values后面,其输出与Multi-head Attention (MHA)相加:

Bert/Transformer模型压缩与优化加速_第6张图片

MiniLM (知识蒸馏)

参考文献: https://arxiv.org/abs/2002.10957

针对NLP任务,深度自注意力知识蒸馏 (Deep Self-Attention Distillation),通过迁移Teacher model最后一层Self-Attention模块的Attention score信息与Value relation信息,可有效实现Student model的诱导训练。只迁移最后一层的知识,能够直接迁移语义信息,显得简单有效、训练速度更快;而且相比于层间特征迁移,不需要手动设计Teacher-student之间的层对应关系。Attention score信息与Value relation信息的知识迁移如下:

  • Attention score transfer:

  • Value relation transfer:

Bert/Transformer模型压缩与优化加速_第7张图片

Bert/Transformer模型压缩与优化加速_第8张图片

  • 若选择的Self-attention layer为LSRA,除了在Multi-head Attention (MHA)分支迁移Attention score与Value relation;在CNN分支需要迁移Feature map的信息,这里主要计算AT loss:

    式中Fs,j表示学生网络里第j个网络层的特征输出,Ft,j表示Teacher network里第j个group的特征输出。

Switch-Transformer

参考文献:https://arxiv.org/abs/2101.03961

针对内容理解与生成、以及多模态特征表征等AI任务,基于MoE(Mixture of Experts)单元的大模型的参数规模不断扩展(Switch-Transformer是其中的典型代表之一),但大模型对算力的需求、被MoE的稀疏激活(Sparse activation)或动态路由(Dynamic routing)机制有效缓解。MoE单元属于典型的Sample驱动型结构设计,使得MoE-based模型具备良好的动态建模能力(动态卷积、Self-Attention等亦如是),因此MoE能够非常好的适应多任务/多模态、甚至在线实时的场景应用需求。Switch-Transformer的MoE单元结构如下,具体执行过程包括动态路由(Dynamic routing)、数据分派(Data dispatch)、专家计算(Expert computation)与结果合并(Result combine)

Bert/Transformer模型压缩与优化加速_第9张图片

多模态与CLIP应用

具体参考:

随笔记录——不同模态信号、表征与应用_AI Flash-CSDN博客

模型量化 (Quantization)

Transformer ASR模型压缩

参考文献: https://arxiv.org/abs/2104.05784

针对Transformer ASR,文章提出了联合随机稀疏与PTQ量化 (KL量化、ADMM与混合精度设置)的压缩策略,整体实现了10倍压缩,且绝对精度损失约0.5%。总体流程如下:

  • 模型稀疏化,更新Weight重要性,渐进式增加稀疏度:

  • KL量化,计算Activation量化参数:

  • ADMM,优化Weight量化参数:

  • 混合精度量化设置,减少量化误差:

Bert/Transformer模型压缩与优化加速_第10张图片

此外,针对Transformer模型的量化,需要讨论全网络INT8量化与计算的意义:

  • 全网络INT8量化:
    • 同时减少计算密集算子与访存密集算子的开销;
    • 实现模型压缩,INT8模型是FP32模型的1/4;
  • Fully INT8 Attention Path:
    • 基本的全INT8通路:
      • Dense layer的输入与输出均为INT8 Tensor;
      • Reshape与Transpose的操作对象为INT8 Tensor,节省内存开销:
        • Dense+Reshape+Transpose、与Transpose+Reshape+Dense可以实现Op fusion;
      • BatchMatMul、Softmax的操作对象保留为FP32 Tensor,确保模型预测精度;

Bert/Transformer模型压缩与优化加速_第11张图片

  • ​​​​​Encoder的Self-Attention:
    • X=Y;
  • Decoder的Self-Attention:
    • X=Y,Batch_size=Batch_size * Beam_size,Seq_len=1;
    • k_dense与v_dense的输出会添加到Cache,按照INT8类型搬运数据,可节省访存开销:

Bert/Transformer模型压缩与优化加速_第12张图片

PTQ for Vision Transformer

参考文献: https://arxiv.org/abs/2106.14156

针对Vision Transformer的PTQ量化,分别针对FFN与Self-Attention,提出了Similarity-aware Quantization与Ranking-aware Quantization:

  • Similarity-aware Quantization: 基于优化方式确定Weight与Activation量化的最优量化参数,并以Pearson相似度作为Target loss;

Bert/Transformer模型压缩与优化加速_第13张图片

Bert/Transformer模型压缩与优化加速_第14张图片

  • Ranking-aware Quantization: 在相似度Loss基础之上,叠加Attention map的Ranking loss作为约束,以准确感知不同Attention map的重要性排序,确保Self-attention量化的有效性(优先保留重要Attention map的信息);

Bert/Transformer模型压缩与优化加速_第15张图片

PTQ4ViT

参考博客文章:

PTQ4ViT: Vision Transformer的PTQ方法与框架_AI Flash-CSDN博客

Q-Bert

参考文献:https://arxiv.org/abs/1909.05840

针对NLP任务(MNLI、SQuAD等),Q-Bert是首个实现Bert模型混合精度量化的QAT算法。Bert模型是由Embedding layer、Encoder layer以及Task-specific output layer构成的前馈结构模型,且Encoder layer包含了MHA与FFN等子模块。

不同的Bert layer体现了不同的量化敏感度(Quantization Sensitivity),配以不同的量化比特数,直接决定了量化后模型的预测精度。通常Bert的Embedding layer与Output layer的量化敏感度相对较高,故Embedding layer会配以较高的量化比特数(如8bits),Output layer则保留为FP32浮点实现。另外,尽管不同Encoder layer的结构相一致,但由于对Context信息编码的贡献有所区别,因此表现出不同的量化敏感性。

  • 量化敏感度与混合精度量化:Q-Bert采用线性均匀量化方式(min-max量化),并通过Dense layer的Hessian信息(二阶梯度)定义其量化敏感度:

Bert/Transformer模型压缩与优化加速_第16张图片

式中Hessian矩阵特征值(Top eigenvalues)的均值,反映了针对同一个下游任务(相同训练数据分布),不同Dense layer的量化敏感度;Hessian矩阵特征值的标准差,反映了针对不同的下游任务(不同训练数据分布),Dense layer量化敏感度的占比。然后,基于不同Dense layer的量化敏感度排序,可确定Bert模型的混合精度量化配置,并进一步执行QAT微调训练。

  • Group-wise量化:为了进一步提高量化后模型的精度,将MHA的Q、K、V与O权重矩阵按Attention heads分成不同的Group,每个Group定义各自的量化范围(min-max范围),从而更精细的保证不同Group的量化分辨率。而每个Group又可以根据实际需求、进一步划分为不同的Sub-group,以满足更为苛刻的部署精度需求: 

    Bert/Transformer模型压缩与优化加速_第17张图片

网络结构搜索 (NAS)

AdaBert

参考文献: https://arxiv.org/abs/2001.04246

针对NLP任务,AdaBert通过可微分搜索、与知识蒸馏,实现卷积类型的NLP模型搜索:

  • 搜索空间:类似于DARTS的Cell设计,实现搜索空间的构造;区别于Vanilla CNN,DARTS的Cell单元结构体现出了非规整性,可能不利于实际推理引擎的加速;

    Bert/Transformer模型压缩与优化加速_第18张图片

  • 搜索策略:基于Gumbel Softmax实现网络结构的随机采样(类似于FBNet)、与可微分搜索,并且采样概率随着训练的进行、逐渐锐化,逼近Argmax采样;另外,基于FLOPS与Model size构造了Efficiency-aware Loss,作为搜索训练的资源约束,实现Hardware-aware搜索;
  • 知识蒸馏:多层次、任务相关的知识迁移;

    Bert/Transformer模型压缩与优化加速_第19张图片

DynaBert

参考文献: https://arxiv.org/abs/2004.04037

针对NLP任务,DynaBert模型压缩策略可实现多个维度的压缩搜索(主要是Width与Depth维度)。DynaBert是一种动态可伸缩性质的One-shot NAS:首先基于预训练Bert模型(或其变体),进行多维度正则化训练;正则化训练之后,按照不同的资源约束需求,能够进行网络参数的裁剪、以获得子网络,并进行子网络的微调训练;从正则训练、到子网络微调,可协同知识蒸馏训练,进一步保持子网络的预测精度。主要步骤描述如下:

  • 第一步:基于Neuron与Attention head的重要性,执行参数重排(Weight Re-wiring):

Bert/Transformer模型压缩与优化加速_第20张图片

  • 第二步:多阶段微调训练,实现宽度、与深度方向的自适应正则化:

Bert/Transformer模型压缩与优化加速_第21张图片

NAS-Bert

参考博客文章:

NAS-Bert——确保One-shot与Task-agnostic_AI Flash-CSDN博客

Evolved Transformer

参考文献: https://arxiv.org/abs/2004.11886

针对序列生成任务,基于NAS搜索获得的Transformer结构:

  • 搜索空间:包括两个Stackable cell,分别包含在Transformer encoder与Transformer decoder。每个Cell由NAS-style block组成, 可通过左右两个Block转换输入Embedding、再聚合获得新的Embedding,进一步输入到Self-Attention模块。
  • 搜索策略:基于EA (Evolutional Aligorithm)的搜索策略;

网络结构如下,融合了一维卷积与Attention的特点:

Bert/Transformer模型压缩与优化加速_第22张图片

HAT: Hardware-aware Transformer

参考文献: https://arxiv.org/abs/2005.14187

对Transformer模型进行网络结构搜索时,通常会面临两个关键问题:

  • FLOPS不能完全代表Transformer模型的计算速度/实际执行效率,即只能间接反映硬件平台特性;
  • 不同硬件设备上,与计算硬件或计算库相适配的最优Transformer结构不尽相同;

针对序列生成任务,HAT (Hardware-aware Transformer)从搜索空间、搜索策略与搜索预测三方面加以分析:

  • 搜索空间:Encoder-decoder Attention的任意连接方式,以及Transformer layer内部结构(网络深度、Attention head数目、Width等);
  • 搜索策略:训练了SuperTransformer作为超网络,并在超网络预训练之后,结合资源约束,通过进化搜索算法(EA: Evolutional Algorithm)寻找最优子网络;不同于权重共享型NAS (如FBNet、SPOS等),HAT是一种动态可伸缩类型的One-shot NAS;
  • 搜索预测:单独、离线训练了MLP回归模型作为Cost model或Predicter,用于预测不同子网络结构(Architecture embedding)、在指定硬件平台上的执行速度 (GPU或CPU latency),并且回归预测效果良好(以相关性系数作为Metric)。搜索预测的好处在于:一方面可直接、有效获取硬件平台特性,作为超网络预训练的资源约束;另一方面,在训练或搜索阶段,无需在硬件平台上测量推理延迟,提升搜索效率,从而有利于跨平台模型结构搜索;

Bert/Transformer模型压缩与优化加速_第23张图片

Bert/Transformer模型压缩与优化加速_第24张图片

AutoFormer

参考文献: https://arxiv.org/abs/2107.00651

针对Vision Transformer的结构搜索,AutoFormer提出的Weight entanglement,在不额外增加Block choice的前提下,通过通道宽度、网络深度、Attention head数目等多个维度的调整,实现了Vision Transformer模型的动态可伸缩预训练与结构搜索。Weight entanglement的做法,类似于BigNAS、FBNet-v2的通道搜索,都不会额外增加通道维度的权重参数量。相比于手工设计的CNN模型(ResNet、ResNext、DenseNet)与Vision Transformer模型(ViT、Deit),AutoFormer模型在相同资源开销条件下、能够获得最好的识别精度。

Bert/Transformer模型压缩与优化加速_第25张图片

AutoFormer总体思路如上图所示,是一种基于Weight entanglement的动态可伸缩搜索方法,其搜索维度包括Attention heads、通道宽度与网络深度。下面从搜索空间、搜索策略与搜索效率这些维度加以分析:

  • 搜索空间:Embedding dimension、Q-K-V dimension、Attention heads、MLP ratio与Network depth;根据不同的资源约束,分别设置Supernet-tiny、Supernet-small与Supernet-base三个基本的Template model;

Bert/Transformer模型压缩与优化加速_第26张图片

  • 搜索策略:基于Weight entanglement的One-shot NAS,完成超网络预训练之后,通过进化算法执行子网络搜索。包含l个网络层的子网络结构可按下式表示,表示第i层的Block结构,表示相应的权重参数:

        而每个Block结构,都是按照Weight entanglement原则从超网络采样获得,下式表示n个动态选择范围:

Bert/Transformer模型压缩与优化加速_第27张图片

  • 搜索效率:按照权重共享形式、完成超网络预训练,训练效率高、且收敛速度快;另外,由于没有引入额外的Block choice,因此训练时的Memory cost较低;

AutoFormerV2

具体参考:

ViT结构优化——Searching the Search Space (S3 NAS)_AI Flash-CSDN博客

计算图优化

LINM

LINM (Loop-invariant Node Motion)是一种计算图等价变换技术,通过将Transformer模型涉及的自回归解码 (Auto-regressive Decoding)的重复计算逻辑 (Encoder-decoder Attention的k/v计算、cache gather等)移至While-loop之外,在确保计算功能不变的情况下,实现计算效率的提升:

Bert/Transformer模型压缩与优化加速_第28张图片

计算图等价变换是在Graph-level执行的、确保计算结果不变的性能优化策略,其总体原则包括:

  • 首先,寻找更具效率的等价计算形式,比如Flops更低、或对后端引擎/硬件更为友好;

  • 其次,在计算量(如Flops)相当的情况下,寻找访存交换量、或Kernel launch数量更低的等价计算形式,以提升计算强度(参考Roofline model的定义);

MatMul替换为卷积

如果推理引擎的卷积算子具备很高的计算效率,并且为了避免一些不必要的访存算子开销,可以按如下计算图等价变换,将Dense layer替换为Conv1D layer:

Bert/Transformer模型压缩与优化加速_第29张图片

QKV计算合并

MHA (Multi-head Attention)模块在计算Attention map与Context编码特征之前,需要基于输入Tensor计算QKV矩阵、并Split成多份Heads。可通过计算图等价变换,将QKV计算合并到同一个Dense layer,能够充分利用NPU或GPU的并发计算特性:

B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.sample_num_heads, -1).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]

推理优化引擎

Faster Transformer

GitHub链接:https://github.com/NVIDIA/FasterTransformer

Faster Transformer是NVIDIA针对标准BERT/Transformer做的推理优化库,其发布时间线:

  • 1.0版本:
    • 2019年7月,开源了FasterTransformer 1.0,针对 BERT中的Transformer Encoder进行优化和加速;
    • 面向BERT分类、自然语言理解场景;
    • 底层由CUDA和cuBLAS实现,支持FP16和FP32计算,FP16可充分利用Volta和Turing架构的Tensor Core计算单元;
    • 提供C++ API、TF Op与TensorRT Plugin三种接口;
    • 参考资料;
  • 2.0版本:
    • 2020年2月,新增对Transformer decoder的优化和加速,包括decoder与decoding两种加速模式;
    • 面向生成式场景,如NMT、文本内容生成与ASR等;
    • 底层由CUDA和cuBLAS实现,支持FP16和FP32计算模,FP16可充分利用Volta和Turing架构的 Tensor Core计算单元;
    • 提供C++ API、TF Op与TensorRT Plugin三种接口;
    • 参考资料;
  • 2.1版本:
    • 2020年6月,引入Effective Transformer优化;通过remove_padding的支持,提高计算与访存效率;
    • 并新增PyTorch Op接口;
    • 参考资料;
  • 3.0版本:
    • 2020年9月,新增BERT encoder的INT8量化加速支持;
    • 仅支持Turing架构GPU;
    • 同时支持PTQ与QAT方法,提供了TF量化工具;
    • 相比于FP16计算,约20~30%加速,但存在精度损失风险;
  • 3.1版本:
    • 2020年12月,新增对PyTorch使用INT8推理的支援;
    • 在Turing以后的GPU上,FP16的性能比3.0提升了 10% ~ 20%;
    • INT8的性能比3.0最多提升了70%;
  • 4.0版本:
    • 2021年4月,新增对GPT-3等百亿/千亿级参数规模模型的多机多卡推理加速能力;
    • 新增FP16 fused MHA算子,同时支持Volta与Turing架构的GPU;
    • 以及对解码端Kernel的优化,可以省略已完成语句的计算,节省计算资源;
    • 参考资料;
  • 支持矩阵:

Bert/Transformer模型压缩与优化加速_第30张图片

Transformer计算图表示,主要包含三个部分的表示:

  • Encoder layer;
  • Decoder layer;
  • Decoding logic;

Bert/Transformer模型压缩与优化加速_第31张图片

Torch-TensorRT

参考博客文章:

深度学习模型压缩与优化加速(Model Compression and Acceleration Overview)_AI Flash-CSDN博客_深度学习模型压缩

AICompiler

基于AICompiler(AI编译优化),可通过Op fusion方式 (算子融合),提升Transformer模型的执行效率。Op fusion具备的好处主要是:1)首先,将多个零碎算子合并成计算功能等价的一个大算子,可减少密集的Kernel launch开销 (CUDA核函数启动);2)其次,由于相邻算子之间存在内存读写开销,Op fusion可有效减少密集的内存访问开销;3)最终通过算子实现的自动寻优、与代码生成 (CodeGen),可自动生成优化的可执行代码。 

总结

深度学习领域的高性能计算,总结为如下公式:

高性能计算 = 高效率算法 + 模型压缩 + 系统/硬件优化

可以单独的强调其中一种优化策略,或者实现多种选项的联合优化(例如Hardware-aware NAS,压缩与编译联合优化等),以满足实际应用场景的部署需求。在工业界,包括阿里淘系MNN、阿里云PAI、华为昇腾计算、Open AI Lab的Tengine、百度EasyDL等平台,在该领域均构建了相对完善的软硬件技术栈,为开放、普惠的AI应用提供了坚实基础。

你可能感兴趣的:(自动语音识别,Transformer,模型压缩,自然语言处理,语音识别,计算机视觉,模型压缩,深度学习)