技术解读倚天 ECS 实例 — Arm 芯片的 Python-AI 算力优化

深度学习技术在图像识别、搜索推荐等领域得到了广泛应用。近年来各大 CPU 厂商也逐渐把 AI 算力纳入了重点发展方向，通过《Arm 芯片 Python-AI 算力优化》我们将看到龙蜥社区 Arm 架构 SIG（Special Interest Group）利用最新的 Arm 指令集优化 Python-AI 推理 workload 的性能。

倚天ECS实例的AI推理软件优化

阿里云推出的倚天Arm ECS实例，拥有针对AI场景的推理加速能力，我们将了解加速的原理以及以及相关的软件生态适配。

卷积神经网络(CNN)在图像和语音领域使用广泛，神经网络算法相比传统的算法消耗了更多算力。为了探索对计算的优化，我们进一步看到AlexNet模型(一种CNN)的推理过程的各个层的计算资源消耗占比。

可以看到名为conv[1-5]的5个卷积层消耗了90%的计算资源，因此优化CNN推理的关键就是优化卷积层的计算。

我们进一步来看如何对图像应用卷积核：

使用im2col根据卷积核尺寸，将图像转化为若干块(patch)
将多个卷积核展开成若干向量
对由图像块组成的矩阵和由多个卷积核展开组成的矩阵应用矩阵乘法

上面一页的计算应用了矩阵乘法操作，为什么我们不采用更加直接的迭代计算方式，而是采用需要额外内存的矩阵乘法呢？这里有两个关键因素：

深度学习的卷积计算量很大，典型计算需要涉及5000万次乘法和加法操作，因此对计算的优化十分重要
计算机科学家们已经深入探索了矩阵乘法操作，矩阵乘法操作可以被优化得非常快。

在fortran世界中，GEMM(general matrix multiplication)已经成为一个通用操作:

该操作通过对数据重新排列，精心设计计算过程，利用多线程和向量指令，可以比自己实现的朴素版本快十倍以上。因此使用矩阵运算带来的收益相比额外的开销是值得的。

因为AI推理大量使用了矩阵乘法，如今也有许多硬件对矩阵运算进行了加速：

NVIDIA Volta架构引入了tensor core，可以高效地以混合精度处理矩阵乘
Intel AMX(Advanced Matrix Extensions) 通过脉动阵列在硬件层面支持矩阵乘
ARM SME(Scalable Matrix Extension) 支持向量外积运算，加速矩阵乘

虽然在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推理

Tensorflow下可以通过OneDNN + ACL(Arm Compute Library)来使用BFMMLA加速
Pytorch已经支持了OneDNN + ACL，但是目前还在试验状态，无法很好地发挥性能。但是Pytorch同时支持OpenBLAS作为其计算后端，因此可以通过OpenBLAS来享受ARM bf16扩展带来的性能收益

可以看到相比默认的eigen实现，开启OneDNN + ACL后，perf获得的计算热点已经从fmla(向量乘加)转换到了bfmmla，算力显著提升。

从workload角度评测，上图对比了两种机型:

g7：Intel IceLake实例
g8m：倚天ARM服务器

左边柱状图中蓝色柱子表示算力对比，橙色柱子表示考虑性价比后使用倚天处理器获得的收益。可以看到在Resnet50和BERT-Large模型的推理场景下，软件优化后的倚天处理器皆可获得一倍左右的性价比收益。

在上文中，我们看到使用倚天处理器若想获得较高收益，软件版本的选择十分重要。随意选择tensorflow或者pytorch包可能遭遇：

未适配arm架构，安装失败
软件未适配bf16扩展或者环境参数有误，无法发挥硬件的全部算力，性能打折
需要精心选择计算后端，例如目前pytorch下OpenBLAS较快

因此我们提供了Docker镜像，帮助云上的用户充分使用倚天ECS实例的AI推理性能：

accc-registry.cn-hangzhou.cr.aliyuncs.com/tensorflow/tensorflow
accc-registry.cn-hangzhou.cr.aliyuncs.com/pytorch/pytorch

通过Serverless能力充分释放算力

除了使能更多的硬件指令，另一种充分释放硬件算力的方式就是通过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。

这种框架下有两个难点:

Python的code_object是散落在heap的各处且不连续的，因此mmap复制整个heap是行不通的。我们采用的方式是以code_object为根，遍历对象图，对感兴趣的内容复制并紧凑排布
Python的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，修改环境变量运行和使用cdsAPI做dump即可对现有的应用启动进行加速了。

code-data-share-for-python是一个新项目，需要大家的参与和反馈，欢迎通过以下链接了解和使用: