想当年,其实估摸着也就大半年前,多多同学还在实验室瞪大眼睛盯着一种叫做xilinx系列的板子,调试着一种叫做VHDL的语言,还记得那个写代码的工具叫做Vivado,不知道大家听说过没有?那个时候,我想实现一个复杂的公式,涉及的计算稍微复杂点(比如来个开方)就要写一大串代码(虽然常用的复杂函数是有IP核可以调的),同时调试过程十分麻烦,甚至要具体到clock对齐。总而言之,十分难忘。那个时候业余时间写下一行Python代码解决一个问题,简直可以直呼“爽啊”。当然,硬件代码虽然难写,但毕竟计算速度、能耗比、并行优势一直很好,所以即便不好写,还是依旧使用广泛。
那么回到今天的主题,在这样只能进行整型数计算的设备上,我们的无敌的BERT如何部署上去呢?当多多看到这篇文章的时候,立刻就来分享给大家了。以下内容主要参考文章:I-BERT: Integer-only BERT Quantization。代码已开源,HuggingFace Git 仓库也有。
相关背景
- 最近两年,以BERT为代表的预训练模型大火,在自然语言处理/最近计算机视觉超越各种达到新SOTA,当然预训练模型虽然效果一个比一个好,参数大小也是越来越大。
- 想要把这些大模型高效(存储空间小,运算速度快)跑起来大致有这么几个办法: A. Pruning,去掉不太重要的参数; B. 知识蒸馏,找个小一点的模型仿大模型; C. 更快更好的计算框架设计; D. 结合计算硬件对模型进行针对性的优化设计; E. Quantization,量化,比如32bit浮点数变成16bit浮点数甚至8bit整型数。
- 目前已有很多压缩Transformer的研究,其中Quantization方面的也不少,但这些Quantization的方法大多不是完完全全的整型数量化(GELU、Softmax、LayerNorm部分都还是用浮点数计算的)。
- 新的Integer-only优秀硬件不断更新,比如ARM处理器。
I-BERT主要贡献
- 针对激活函数GELU,Softmax函数,I-BERT提出了一种高效并且正确的integer-only的估算方法,该方法为基于整型计算实现的轻量级二阶多项式。
- 针对Transformer里的LayerNorm,使用integer-only迭代算法对平方根进行估算。
- 基于GELU、Softmax、LayerNorm的估算,将Transformer进行了量化。 具体的Transformer量化为:用INT8乘法和INT32累加器处理Embedding和Matrix multiplication;然后基于INT32的累加结果计算GELU、Softmax和LayerNorm,然后再将INT32量化回INT8。整个计算图中的参数和激活函数计算都是用Intergers。图1展示了I-BERT和其他工作的区别。
- I-BERT在基于Roberta-base/large进行了实现,并在GLUE上取得了和全精度模型相近的效果; 和32位浮点计算相比,将I-BERT基于TensorRT部署到T4 GPU上,有4倍的加速效果。
图1 对比I-BERT和之前的工作
I-BERT实现
I-BERT全文都在传递一个思想,其实和我上一篇写的贝叶斯优化很相似:如果一个函数的计算解决不了,就找个近似函数,如果近似函数不好找,就限定自变量范围来找!
1
基本量化方法
这里简单介绍以基本的量化概念和方法,在uniform symmetric quantization框架下(对称、均匀的量化),一个实数x会被映射成一个整数q,q的范围是[-2^(b-1), 2^(b-1) -1],b为量化的bit数,正式定义为:
简单来讲为:
- 如果一个数x不在[-a,a]范围内则直接用x=a或者x=-a,如果在范围内则保持x。
- 得到x除以2^(b-1)-1得到的整数商作为量化结果。
把量化的结果转化为原始值也就是用q乘以上式中的分母,可以看到量化是损失的,一般来说,量化之后再反量化过程无法得到和原来一摸一样的数字。除了均匀量化,还有很多非均匀、不对称的量化方法,这里不再深入介绍。
1
Integer-Only的非线性函数GELU计算
我们先放出非线性GELU的表达式,然后思考下这么复杂的函数如何用整数计算进行近似估计,orz我感觉挺难的反正:
I-BERT是这么解决这个问题的?
由于整型计算可以计算2阶多项式,比如a(x b)^2 c。那么先用2阶多项式估计GELU函数,再用整型计算2阶多项式。为什么是2阶不是更高阶呢?因为阶数约高,估算函数也越复杂,作者发现2阶多项式足够估计GELU了。
2阶多项式估算GLUE又可以转化为优化如下目标函数,主要是对erf函数进行估计:找到一个2阶多项式,系数为a,b,c,能让GELU函数与2阶多项式误差最小。
erf函数有这么两个特点:对于输入很大的值而言,erf结果基本趋近于1或者-1,并且erf函数还是一个奇函数,因此只考虑优化正实数域的优化。在这两个特点下,作者找到一个相对较优的2阶多项式:
其中a=-0.2888,b=-1.769。sgn为符号函数。最终GELU的近视表达为:
整个GELU的估算和效果如图所示:
图2 GELU算法和效果
1
Integer-Only的softmax计算
个人觉得这几个估算十分精妙!
softmax函数的表达式是和特点:
然后可以发现x=x-x_max之后都会变成非正数,然后任何一个非正数可以表达为x=(-ln2)z p,其中z是一个非负整数,p是一个在[-ln2,0]之间的实数,因此x的指数可以表示为:
>>是移位运算,在整型逻辑计算硬件中十分常见。最终将估计问题变化为简单的:估计[-ln2,0]之间的p的指数值。同样使用2阶多项式进行估计,最终得到:
到这里,整个integer-only的softmax计算逻辑如图所示:
图3 Softmax估计算法
1
Integer-Only的Layer Norm
Layer Norm在Transformer中大量使用,并且使用了非线性函数来计算方差和均值。
而且,注意LayerNorm的计算是在模型运行的是时候动态计算的,所以这些开方、平方运算也需要跟着一起算。不过这些平方、开方运算在整型运算里面还算常见,常见的方法就是迭代求解。如下图所示,便不再详细介绍啦:
图4 Layer Norm中的平方根估算方法
1
I-BERT的效果
虽然说I-BERT能完成用整型计算得到Transformer的结果了,但也不能损失太多的效果。总体上I-BERT的效果是不错的。如下图所示:
图5 GLUE任务上的效果
同时I-BERT在硬件上的运算速度也很好,相比于浮点数计算有2-4倍的加速。
图6 硬件运行加速对比
总结
本文对I-BERT进行了介绍,整体上I-BERT能完全基于整型计算进行,并且对于GELU、Softmax、LayerNorm的估算非常有借鉴意义。笔者能从这几个估算过程感受懂啊数学之美啊!!!!!