深度学习技术在图像识别、搜索推荐等领域得到了广泛应用。近年来各大 CPU 厂商也逐渐把 AI 算力纳入了重点发展方向,通过《Arm 芯片 Python-AI 算力优化》我们将看到龙蜥社区 Arm 架构 SIG(Special Interest Group) 利用最新的 Arm 指令集优化 Python-AI 推理 workload 的性能。
阿里云推出的倚天Arm ECS实例,拥有针对AI场景的推理加速能力,我们将了解加速的原理以及以及相关的软件生态适配。
卷积神经网络(CNN)在图像和语音领域使用广泛,神经网络算法相比传统的算法消耗了更多算力。为了探索对计算的优化,我们进一步看到AlexNet模型(一种CNN)的推理过程的各个层的计算资源消耗占比。
可以看到名为conv[1-5]
的5个卷积层消耗了90%的计算资源,因此优化CNN推理的关键就是优化卷积层的计算。
我们进一步来看如何对图像应用卷积核:
im2col
根据卷积核尺寸,将图像转化为若干块(patch
)上面一页的计算应用了矩阵乘法操作,为什么我们不采用更加直接的迭代计算方式,而是采用需要额外内存的矩阵乘法呢?这里有两个关键因素:
在fortran
世界中,GEMM(general matrix multiplication)已经成为一个通用操作:
该操作通过对数据重新排列,精心设计计算过程,利用多线程和向量指令,可以比自己实现的朴素版本快十倍以上。因此使用矩阵运算带来的收益相比额外的开销是值得的。
因为AI推理大量使用了矩阵乘法,如今也有许多硬件对矩阵运算进行了加速:
虽然在AI算力上GPU要远高于CPU,但是CPU因为其部署方便,且无需在主机-设备间拷贝内存,在AI推理场景占有一席之地。目前市面上尚没有可以大规模使用的支持AMX或者SME的硬件,在这个阶段我们应该如何优化CPU上的AI推理算力呢?我们首先要了解BF16数据类型。
BF16(Brain Float 16)是由Google Brain 开发设计的16位浮点数格式。相比传统的IEEE16位浮点数,BF16拥有和IEEE单精度浮点数(FP32)一样的取值范围,但是精度较差。研究人员发现,在AI训练和推理中,使用BF16可以节约一半的内存,获得和单精度浮点数接近的准确率。
根据右图,BF16指数的位数和FP32是一致的,因此BF16和FP32的相互转换只要截断尾数即可,左下角图上便是tensorflow源码中的转换实现。
引入BF16的一大价值是如今的很多硬件计算的瓶颈在寄存器宽度或者访问内存的速度上,更紧凑的内存表示往往可以获得更高的计算吞吐,在理想情况下,BF16相比FP32可以提高一倍的吞吐(FLOPS)。
如今我们虽然无法大规模使用到支持AMX/SME的硬件,但是Armv8.6-A提供了bf16扩展,该扩展利用了有限的128bit向量寄存器,通过BFMMLA
指令执行矩阵乘法运算:
A
: 大小为2*4的BF16矩阵,按行存储B
: 大小为4*2的BF16矩阵,按列存储C
: 大小为2*2的FP32矩阵该指令单次执行进行了16次浮点数乘法和16次浮点数加法运算,计算吞吐非常高。
阿里巴巴向OpenBLAS项目贡献了sbgemm
(s表示返回单精度,b表示输入bf16)的硬件加速实现,从GEMM吞吐上看,BF16相比FP32 GEMM吞吐提升超过100%。
倚天ECS实例是市面上少数可以支持bf16指令扩展的ARM服务器。目前已经支持了Tensorflow和Pytorch两种框架的AI推理
BFMMLA
加速可以看到相比默认的eigen实现,开启OneDNN + ACL后,perf获得的计算热点已经从fmla
(向量乘加)转换到了bfmmla
,算力显著提升。
从workload角度评测,上图对比了两种机型:
左边柱状图中蓝色柱子表示算力对比,橙色柱子表示考虑性价比后使用倚天处理器获得的收益。可以看到在Resnet50
和BERT-Large
模型的推理场景下,软件优化后的倚天处理器皆可获得一倍左右的性价比收益。
在上文中,我们看到使用倚天处理器若想获得较高收益,软件版本的选择十分重要。随意选择tensorflow或者pytorch包可能遭遇:
因此我们提供了Docker镜像,帮助云上的用户充分使用倚天ECS实例的AI推理性能:
accc-registry.cn-hangzhou.cr.aliyuncs.com/tensorflow/tensorflow
accc-registry.cn-hangzhou.cr.aliyuncs.com/pytorch/pytorch
除了使能更多的硬件指令,另一种充分释放硬件算力的方式就是通过Serverless架构提高CPU利用率。Python作为动态语言,其模块是动态导入的,因此启动速度不是Python的强项,这也制约了Python workload在Serverless场景的普及。
Python应用启动的主要耗时在模块导入,Python模块导入步骤为:
code_object
其中的第二步在首次加载模块时,要对.py
文件进行编译,获得code_object
, 为了降低将来加载的开销,Python解释器会序列化并缓存code_object
到.pyc
文件。
即便模块导入过程已经通过缓存机制优化过了,但是读取.pyc
文件并反序列化依旧比较耗时。
在这里我们借助了OpenJDK的AppCDS的思路:将heap上的code_object
复制到内存映射文件中(mmap)。在下次加载模块时,直接使用mmap中的code_object
。
这种框架下有两个难点:
code_object
是散落在heap的各处且不连续的,因此mmap复制整个heap是行不通的。我们采用的方式是以code_object
为根,遍历对象图,对感兴趣的内容复制并紧凑排布code_object
会引用.data
段的变量,在Linux的随机地址安全机制下,.data
段的数据的地址在每次运行时都会随机变化,这样mmap中的指针就失效了。我们的解决方式是遍历所有对象,针对.data
段的指针进行偏移量修复因为该项目共享了python的code_object
,因此名字是code-data-share-for-python
,简称pycds
。
我们测试了bota3
、numpy
、flask
等常用的python苦,平均可以节省20%的模块导入耗时。
对于现有的python应用可以轻易地使用pycds,且无需修改任何代码:
# 安装pycds
pip install code-data-share # 安装pycds
# 生成模块列表
PYCDSMODE=TRACE PYCDSLIST=mod.lst python -c 'import numpy’
# 生成 archive
python -c 'import cds.dump; cds.dump.run_dump("mod.lst", "mod.img")’
# 使用archive
time PYCDSMODE=SHARE PYCDSARCHIVE=mod.img python -c 'import numpy'
real 0m0.090s
user 0m0.180s
sys 0m0.339s
# 对比基线
time python -c 'import numpy'
real 0m0.105s
user 0m0.216s
sys 0m0.476s
我们仅仅通过安装PyPI,修改环境变量运行和使用cds
API做dump即可对现有的应用启动进行加速了。
code-data-share-for-python
是一个新项目,需要大家的参与和反馈,欢迎通过以下链接了解和使用:
ARM 架构 SIG链接地址:https://openanolis.cn/sig/ARM_ARCH_SIG
原文链接
本文为阿里云原创内容,未经允许不得转载。