编者按：在刚刚结束的PyConChina2022大会上，龙蜥社区开发者朱洪林分享了主题的技术演讲《ARM 芯片的 Python+AI 算力优化》。在本次演讲中，笔者将介绍他们在倚天710ARM芯片上的Python+AI优化工作，以及在ARM云平台上部署Python+AI任务的最佳实践。以下为本次演讲内容：（图/朱洪林现场演讲）我们的场景是ARM平台上的AI相关任务。主要目标是优化性能。具体来说，我们首先关注深度学习的推理任务（inferencetask），主要原因也是来自于业务需求。这里所说的ARM平台并不是我们理解的终端设备，比如手机，而是服务器平台。在大家的印象中，AI任务，尤其是深度学习程序，一般运行在GPU或者x86CPU上。考虑到功耗、成本、性能等因素，云厂商逐渐开始构建基于ARM架构的服务平台。是一种趋势。当然，ARM平台还不是很成熟，很多软件都不能顺利运行，更谈不上提升性能了。我们想吸引一些用户将AI应用从原来的x86平台迁移到ARM平台。这就需要ARM平台提供更好的性能或者更好的性价比。因此，如何整合Python+AI相关软件，使其发挥最佳性能，成为大家关注的焦点。下面的分享整体分为两部分。一部分是介绍我们的优化工作，主要涉及矩阵乘法。第二部分是关于PythonAI在ARM云平台-倚天710上应用的最佳实践。一、优化工作简介前面说了我们的优化是和矩阵乘法相关的，所以首先需要说明一下为什么要关注这。这里一个无法回避的场景就是深度学习。无论是几年前大名鼎鼎的AlphaGo，还是现在大热的ChatGPT，都使用了大量的深度学习技术。深度学习本身只是人工智能的一个分支，但其影响却十分广泛。不容忽视。所以我们从深度学习入手，从目前使用最广泛的深度学习框架TensorFlow和PyTorch入手。另外，我们还需要结合硬件场景，就是上面提到的ARM服务器平台，对于阿里云来说，是结合倚天710芯片。深度学习的实现包括大量的矩阵乘法，甚至有文章直接写到矩阵乘法是深度学习的核心。比如大家熟知的卷积运算，其实是经过了一系列的变换，将输入的特征和卷积核转化为两个矩阵，然后进行矩阵乘法，输出结果解码为特征图，它是完成的卷积。另外，全连接层也是通过矩阵乘法来实现的。目前流行的Transformers结构被包括ChatGPT在内的各种NLP模型所采用，同时也包含了大量的矩阵乘法运算。我们可以看一些例子：你可以看到像AlexNet、ResNet-50这样的模型，在做推理的时候，大约90%的计算时间都花在了矩阵乘法上。即使是对矩阵乘法的小优化也会产生广泛的影响。我们前面说的矩阵乘法，更准确的叫GEMM，generalmatrixmultiplication，其实包括系数和累加运算。但是时间复杂度还是在MNK级别，主要是AB的两个矩阵相乘。直观上，深度学习涉及大量的矩阵乘法计算。例如，常见的卷积运算可能涉及5000万次计算，因此需要进行优化。右下图是最简单的三层循环迭代方式。这种方法通常很慢，计算机科学家从优化内存布局和使用向量指令开始，为将性能提高10倍以上做出了很多努力。内存布局主要分为两步。第一步是将矩阵分成块。即对于一个非常大的矩阵，我们不是按顺序一个一个计算，而是把矩阵分成小块，小块计算。第二步，重新排列分离出来的小块的内部元素序列。比如矩阵是按行排列的，可能会计算前四行，需要取第二行的前四行。但是取第二行，指针需要移动很远的距离，很容易造成cachemiss，所以需要重新排列，让它们在内存中连续。优化内存布局的主要目的是提高缓存命中率，减少内存访问次数。第二种是使用向量化指令，类似于x86设备的AVX和ARM设备的NEON。矢量化指令的本质是同时计算多个数据。比如我们要将四组数据相乘，一般情况下需要执行四次。如果将它们对应地放入向量寄存器，则只需要一条向量化指令。说明，可以同时得到四个结果，提高了计算效率。当然，这需要硬件支持。由于AI推理广泛使用矩阵乘法，现在很多硬件加速矩阵运算：NVIDIAVolta架构引入tensorcore，可以通过脉动阵列在硬件层面高效处理混合精度IntelAMX（AdvancedMatrixExtensions）的矩阵乘法支持矩阵乘法ARMSME(ScalableMatrixExtension)支持向量外积运算，加速矩阵乘法目前市场上还没有支持AMX或SME的硬件可以大规模使用。现阶段我们应该如何优化CPU上的AI推理能力？羊毛布？我们首先需要了解BF16数据类型。BF16（全称BrainFloatingPoint）是由GoogleBrain开发设计的一种16位浮点数格式。BF16与传统的FP16位浮点数相比，取值范围与FP32相同，但精度较差。但对于深度学习，较低的精度不会显着影响结果，而较低的表示范围会显着影响模型训练的质量。此外，BF16还具有转换方便的特点。BF16和FP32之间的转换只需要截尾或补尾即可。使用BF16还节省了一半的内存，紧凑的内存表示通常意味着更高的计算吞吐量。最后，我们还有硬件指令支持，可以直接对BF16数据进行操作。需要注意的是BF16的扩展是包含在ARMv8.6设备中的，当然倚天710是ARMv9的指令集，也是支持的。我们主要使用BFMMLA进行矩阵乘法计算。例如，对于含有128位向量寄存器的设备：输入A：BF16矩阵，大小为2*4，按行存储输入B：BF16矩阵，大小为4*2，按列存储输出C：FP32矩阵BFMMLA大小为2*2。一条指令完成16次乘法和16次加法，计算吞吐量非常高。当然，此时如果我们需要C是BF16类型的，就需要应用转换指令，比如向量化指令BFCVT，来加快转换过程。我们的目标是为tensorflow和pytorch用户提供加速。这是整体流程图。对于AI推理任务，TensorFlow和PyTorch都不会自己直接计算，而是调用专门的计算后端。ARM中主要有两个，一个是ARMComputeLibrary，一个是OpenBLAS，它们之间的关系如右图所示。TensorFlow在最近的版本中已经开始使用oneDNN+ACL作为计算后端。oneDNN也是一层皮，实际计算的还是ACL。用户实际上只需要设置一个环境变量就可以在不接触代码的情况下获得BF16加速。这项改进首先由ARM公司的研发人员完成。具体操作示例如下：#假设resnet.py包含用户编写的模型推理代码DNNL_DEFAULT_FPMATH_MODE=BF16python3resnet.pyPyTorch比较复杂。PyTorch支持OneDNN+ACL，但性能不佳，而且PyTorch支持OpenBLAS后端，因此可以通过OpenBLAS享受ARMbf16扩展带来的性能优势。OpenBLASBF16的GEMM优化由龙蜥社区理事单位阿里巴巴贡献。同时，为了方便用户，我们还为PyTorch添加了一个API。用户添加一行torch.set_float32_fast_math_mode("BF16")，不需要修改其他代码就可以得到BF16加速（需要注意的是这个api还没有集成到PyTorch中，所以只能通过我们使用的pytorch镜像来获取）提供）。操作示例如下：#...#在模型执行前设置快速数学模式torch.set_float32_fast_math_mode("BF16")#...#执行模型pred=model(x)#...之后，还有一些性能测试演示，我们测试了OpenBLAS纯矩阵计算的性能对比。GFLOPS和执行时间这两个指标分别记录。然后测试TensorFlow和PyTorch的性能对比。在对比中我们可以看到，得益于BF16扩展，在最新的ECSARM平台上的性能优于x86平台（g7）。2、在ARM云平台上使用PythonAI的最佳实践——倚天710下面介绍一下在ARM平台上使用TensorFlow或PyTorch的最佳实践，尤其是倚天710的用户，软件版本的选择一定要知道非常重要。如果随意选择tensorflow或者pytorch包，可能会遇到：ARM架构不适配，安装失败。选择计算后端。例如，pytorch下的OpenBLAS比TensorFlow更快。我们可以选择官方最新的两个版本，2.10.1或者2.11.0（最新版）来获取ACL的BF16加速。用户也可以选择阿里云的镜像，其实和pip安装的是一样的，没有区别。对于PyTorch用户来说，正式版只能在最新的1.13.0中获得ACL加速，但如前所述，实际表现并不突出。阿里云为PyTorch提供了最新的OpenBLAS，可以在dockerpull时标记torch_openblas获取。另外，我们还提供了modelzoo镜像，里面包含了模型的测试代码和验证代码。目前我们还在进行相关工作，期待未来能为大家提供更完整的形象。欢迎大家进群一起探讨相关技术。AISIG主页地址：https://openanolis.cn/sig/AI_SIG原文链接本文为阿里云原创内容，未经允许不得转载。
