• 微信公众号:美女很有趣。 工作之余,放松一下,关注即送10G+美女照片!

Bert/Transformer模型加速

互联网 diligentman 2小时前 2次浏览

前言

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

  • Standard Self-Attention:

Bert/Transformer模型加速

  •  Standard FFN:

Bert/Transformer模型加速

 本文从以下几个维度举例说明,解析Bert/Transformer模型的加速方法:

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

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

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模型加速

Bert/Transformer模型加速

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 layer,DFSMN添加在values后面,其输出与Multi-head Attention (MHA)相加:

Bert/Transformer模型加速

Bert/Transformer模型加速

MiniLM (知识蒸馏)

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

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

  • Attention score transfer:

Bert/Transformer模型加速

  • Value relation transfer:

Bert/Transformer模型加速

Bert/Transformer模型加速

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

    Bert/Transformer模型加速

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

模型量化 (Quantization)

Transformer ASR模型压缩

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

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

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

Bert/Transformer模型加速

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

Bert/Transformer模型加速

Bert/Transformer模型加速

Bert/Transformer模型加速

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

Bert/Transformer模型加速

Bert/Transformer模型加速

此外,针对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模型加速

  • ​​​​​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模型加速

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模型加速

Bert/Transformer模型加速

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

Bert/Transformer模型加速

Bert/Transformer模型加速

Bert/Transformer模型加速

网络结构搜索 (NAS)

AdaBert

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

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

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

    Bert/Transformer模型加速

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

    Bert/Transformer模型加速

DynaBert

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

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

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

Bert/Transformer模型加速

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

Bert/Transformer模型加速

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 layer。
  • 搜索策略:基于EA (Evolutional Aligorithm)的搜索策略;

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

Bert/Transformer模型加速

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,用于预测不同子网络结构、在指定硬件平台上的执行速度 (Latency)。搜索预测的好处:一方面可直接、有效获取硬件平台特性,作为超网络预训练的资源约束;另一方面,无需在硬件平台上测量推理延迟,有利于跨平台模型结构搜索;

Bert/Transformer模型加速

Bert/Transformer模型加速

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模型加速

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

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

Bert/Transformer模型加速

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

Bert/Transformer模型加速

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

Bert/Transformer模型加速

Bert/Transformer模型加速

Bert/Transformer模型加速

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

计算图优化

LINM

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

Bert/Transformer模型加速

推理优化引擎

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模型加速

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

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

Bert/Transformer模型加速


程序员灯塔
转载请注明原文链接:Bert/Transformer模型加速
喜欢 (0)