关于芯片:解读最佳实践倚天-710-ARM-芯片的-PythonAI-算力优化-龙蜥技术

3次阅读

共计 3851 个字符，预计需要花费 10 分钟才能阅读完成。

编者按：在刚刚完结的 PyCon China 2022 大会上，龙蜥社区开发者朱宏林分享了主题为《ARM 芯片的 Python+AI 算力优化》的技术演讲。本次演讲，作者将向大家介绍他们在倚天 710 ARM 芯片上发展的 Python+AI 优化工作，以及在 ARM 云平台上部署 Python+AI 工作的最佳实际。

以下为本次演讲内容：

（图 / 朱宏林现场演讲）

咱们的场景是 ARM 平台的和 AI 相干的工作，次要的指标是进行性能优化，具体来说咱们首先关注的是深度学习推理工作（inference task），次要起因也是来自于业务需要。

这里说的 ARM 平台不是咱们了解的终端设备，比方手机之类的，而是指服务端平台。在大家印象中，AI 工作，特地是深度学习的程序个别是跑在 GPU 或者 x86 的 CPU 上，出于功耗、老本、性能等因素的思考，云厂商逐渐开始建设 ARM 架构的服务平台，这是一个趋势。当然 ARM 平台还不是很成熟，许多软件还无奈胜利跑起来，更不要说晋升性能了。

咱们想要吸引一部分用户将 AI 利用从原先的 x86 平台上迁徙到 ARM 平台上。这就要求 ARM 平台能提供更好的性能，或者更好的性价比。所以说如何整合 Python+AI 的相干软件使其施展最好的性能成为了咱们关注的重点。

下文的分享整体分为两局部，一部分是介绍咱们进行的优化工作，次要是跟矩阵乘法相干的优化，第二局部是对于 Python AI 利用在 ARM 云平台 - 倚天 710 上的最佳实际。

后面说咱们的优化是和矩阵乘法相干的，那首先须要阐明为什么咱们会关注到这个。

这里有一个绕不开的场景就是深度学习，不论是前几年出名的 AlphaGo，还是以后炽热的 ChatGPT，都用到了大量深度学习的技术，深度学习自身只是 AI 的一个分支，但却影响宽泛，不容忽视。所以咱们从深度学习开始切入，从以后最宽泛应用的深度学习框架，TensorFlow 和 PyTorch 开始。此外，咱们还须要联合硬件场景，即后面说到的 ARM 服务端平台，对于阿里云来说就是联合倚天 710 芯片。

深度学习的实现中蕴含大量的矩阵乘法，甚至有文章间接写出矩阵乘法是深度学习的外围。举个例子，咱们熟知的卷积操作，实际上通过一系列的转换后，输出特色和卷积核会被转换为两个矩阵，而后进行矩阵乘法，输入的后果再解码成特色图，就实现了卷积。除此以外，全连贯层也由矩阵乘法实现，以后风行的 Transformers 构造，被包含 ChatGPT 在内的各类 NLP 模型所应用，也蕴含大量矩阵乘法操作。

咱们能够看一些例子：

能够看到，像 AlexNet、ResNet-50 之类的模型，在进行推理时，大概 90% 计算耗时在执行矩阵乘法。即便对矩阵乘法做一些渺小的优化，影响也是很宽泛的。

咱们后面说的矩阵乘法，更精确的叫法是 GEMM，通用矩阵乘法，其实还蕴含系数和累加操作。然而工夫复杂度依然是 MNK 级别，次要还在于 AB 两个矩阵相乘。直观来看，深度学习波及的矩阵乘法计算量很大，比方常见的卷积操作可能就波及 5000 万次计算，所以优化就显得很有必要，右下图是最奢侈的三层循环迭代法，这种做法通常十分慢，计算机科学家做了许多致力，从优化内存布局和利用向量指令登程，可能将性能晋升 10 倍以上。

内存布局次要分两步，第一步是对矩阵进行分块，即对于一个超大的矩阵，咱们并不是一个一个按程序计算，而是将矩阵切分为一个一个小块，分小块计算。第二步是对分出的小块，外部的元素序列进行重排，例如原来是按行排列的矩阵，那可能第一行四个计算好了，就须要取第二行的前四个，然而要取第二行就须要指针挪动很长的间隔，很容易造成 cache 不命中，于是须要重排，使得他们在内存上间断。优化内存布局次要目标是为了减少 cache 命中率，缩小访存次数。

其次是利用向量化指令，相似 AVX 对于 x86 设施，NEON 对于 ARM 设施。向量化指令实质上是为了同时对多个数据进行计算，例如咱们要对四组数据别离进行乘法，那么惯例状况下须要执行四次，如果将它们对应放入向量寄存器中，只须要一条向量化指令，就能够同时得出四个后果，计算效率失去晋升。当然这个是须要硬件反对。

因为 AI 推理大量应用了矩阵乘法，现在也有许多硬件对矩阵运算进行了减速：

NVIDIA Volta 架构引入了 tensor core，能够高效地以混合精度解决矩阵乘
Intel AMX(Advanced Matrix Extensions) 通过脉动阵列在硬件层面反对矩阵乘
ARM SME(Scalable Matrix Extension) 反对向量外积运算，减速矩阵乘

目前市面上尚没有能够大规模应用的反对 AMX 或者 SME 的硬件，在这个阶段咱们应该如何优化 CPU 上的 AI 推理算力呢？咱们首先要理解 BF16 数据类型。

BF16(全称 Brain Floating Point)，是由 Google Brain 开发设计的 16 位浮点数格局。

相比传统的 FP16 位浮点数，BF16 领有和 FP32 一样的取值范畴，然而精度较差。但对于深度学习来说，较低的精度并不显著影响后果，而较低的示意范畴则会显著影响模型训练的好坏。

此外，BF16 还具备转换不便的特点，BF16 和 FP32 的互转只须要截断或填充尾数即可。

应用 BF16 还能够节约一半的内存，紧凑的内存示意通常意味着更高的计算吞吐。

最初，咱们也有了硬件指令反对，能够间接对 BF16 数据进行操作。

须要阐明的是 BF16 的扩大蕴含在 ARMv8.6 设施上，当然倚天 710 是 ARMv9 的指令集，同样反对。咱们次要通过 BFMMLA 来进行矩阵乘法计算，例如对于蕴含 128bit 的向量寄存器的设施来说：

输出 A: 大小为 2*4 的 BF16 矩阵，按行存储
输出 B: 大小为 4*2 的 BF16 矩阵，按列存储
输入 C: 大小为 2*2 的 FP32 矩阵

BFMMLA 单指令实现 16 次乘法和 16 次加法，计算吞吐十分高。

当然这时候如果咱们须要 C 是 BF16 类型的话，就须要利用转换指令，例如向量化指令 BFCVT，减速转换过程。

咱们的指标还是给 tensorflow 和 pytorch 用户提供减速，这是整体的流程图，对于一个 AI 推理工作，实际上不论是 TensorFlow 还是 PyTorch 都不会本人间接去计算，而是叫个专门的计算后端，在 ARM 次要是两个，一个是 ARM Compute Library，另一个是 OpenBLAS，他们之间的关系如右图。

TensorFlow 在最近的版本中开始采纳 oneDNN + ACL 作为计算后端，oneDNN 也是一层皮，理论的计算依然是 ACL。用户实际上只须要设置一个环境变量，就能够在不该动代码的状况下取得 BF16 减速。这个改良是由 ARM 公司的研发人员首先实现了。具体操作例子如下：

# 假如 resnet.py 蕴含用户写的模型推理的代码
DNNL_DEFAULT_FPMATH_MODE=BF16 python3 resnet.py

PyTorch 的状况比较复杂，PyTorch 反对 OneDNN + ACL，但无奈很好的施展性能，同时 PyTorch 反对 OpenBLAS 后端，因而能够通过 OpenBLAS 来享受 ARM bf16 扩大带来的性能收益。

OpenBLAS 的 BF16 的 GEMM 优化是由龙蜥社区理事单位阿里巴巴奉献的，于此同时，咱们为了不便用户应用，也在 PyTorch 中退出了一个 API，用户在模型执行前增加一行 torch.set_float32_fast_math_mode(“BF16”)，就能够取得 BF16 减速，不用批改其余代码（须要阐明，这个 api 还没有合入 PyTorch，所以目前要应用咱们提供的 pytorch 镜像才能够取得）。

操作例子如下：

# ...

# 在模型执行前设置 fast math mode
torch.set_float32_fast_math_mode("BF16")
# ...
# 执行模型
pred = model(x)
# ...

之后是一些性能测试的展现，咱们测试了 OpenBLAS 纯矩阵计算的性能比照。别离记录了 GFLOPS 和执行工夫两个指标。

而后测试 TensorFlow 和 PyTorch 的性能比照，在比照中，咱们能够看到，得益于 BF16 扩大，最新的 ECS ARM 平台上的性能优于 x86 平台（g7）。

当初介绍一下在 ARM 平台，特地是倚天 710 的用户，应用 TensorFlow 或 PyTorch 的最佳实际。

要晓得软件版本的抉择非常重要，随便抉择 tensorflow 或者 pytorch 包可能遭逢：

未适配 ARM 架构，装置失败
软件未适配 BF16 扩大或者环境参数有误，无奈施展硬件的全副算力，性能打折
须要精心抉择计算后端，例如目前 pytorch 下 OpenBLAS 较快

在 TensorFlow 上，咱们能够抉择最新的两个官网版本，2.10.1 或者 2.11.0（最新版本），才可能取得 ACL 的 BF16 减速。用户也能够抉择阿里云的镜像，这个和 pip 装置的其实是一样的，没有区别。

对于 PyTorch 用户，官网版本只有在最新的 1.13.0 才可能取得 ACL 减速，然而正如后面所说的，理论性能并不突出。阿里云则提供了带最新 OpenBLAS 的 PyTorch，在 docker 拉取时标注 torch_openblas 就能够取得。此外，咱们也提供了 modelzoo 镜像，蕴含模型的测试代码和验证代码。

目前咱们依然在进行相干的工作，期待后续能为大家提供更加欠缺的镜像。欢送大家入群一起摸索相干技术。