为RTP-LLM提供Arm CPU后端,助力Arm
RTP LLM 是一个大语言模型推理和加速的引擎,主要是用 C++ 和 Python 编写的,它在阿里巴巴的电商平台得到了生产环境的部署和应用。在性能上,最原始的后端是基于 GPU 扩大的高性能内核,可以达到一个业界领先的性能水平。在灵活性和应用性上面,前端使用的是类似于 Hugging Face 的 Library ,去加载模型和权重,所以在模型的支持上和群众的利用上也是比较容易使用的。在后端的选择上面,最开始支持 GPU 的设备类型,目前在异构的计算平台上有了更多后端加入,比如一些硬件加速的处理板卡或者是像通用 CPU 的设备类型都是可以在这个框架上得到一个应用的支持。最后会有一个 Learning Paths ,是提交的 RTP LLM 的入门教程,感兴趣想实践的话,可以通过这个教程做一个入门的学习。
关于如何使用 RTP LLM ,在硬件要求上,选择的是阿里云上以前的 710 平台,基于 Arm Neoverse N2 的服务器处理器,操作系统目前选的是 22.04 ,其他的像阿里 Linux 或者龙蜥的平台,也是没有问题的。最少的核数和内存的要求取决于要推理的模型的大小,对硬件的要求也会有相应的提高。如果要下载或者是构建源码,加上 Configure 等于 Arm 的构建参数来生成一个 Maga Transformer 的一个包。 Maga 是 Make Alibaba Great Again的意思,比较巧合的是龙蜥和 Arm 的代号都是以 A 开头,所以如果想不到更好的代号,可以考虑用 Maga 。编译完用 Keep In Store 来本地安装,就可以通过 Paisoimport 包来做推理的服务部署。中间必不可少的是要下载一个需要推理的大语言模型的模型文件,这里以 Hugging Face Cli 来从实例,以 0.5b 这样一个比较小的方便测试的模型,做我们后续推理的模型准备。
大概介绍一下 RTP LLM 整体的架构分层,由 Python 做前端,负责去加载权重,开始推理的服务,一些基本的前端处理都在 Python 端完成。中间是 C++ 层,会定义不同模型的推理实现的逻辑,就是 Model Forward 的 Operation 里面一些具体的算子中间组合的逻辑。最底下是不同的硬件平台主要实现的硬件的算子,比如 CPU 的一些实现,在底层的 Device Ops 里面会要求开发者做一个实现。
如果想要增加一个新的硬件知识,比如 Arm 的 CPU 以这个为入口实现一些基本算子,在框架上会有一个罗列。在开发过程中发现一些算子 API 传入参数的定义,更适配于调用的逻辑,所以在我们 CPU 的实现上会有一些要适配的东西,所以要理解输入输出参数的一些定义。
还需要对函数的逻辑做验证,在这个框架里面有一些基础的功能函数的单元测试的用例,中间的话会有一个 GPT Model 除去上层 Python 逻辑之后的 C++ 的逻辑验证,会根据 Token ID 得到一些输出的 Score ,选出最大的一个可能的 Token 的分值,跟匹配的参考平台做比较。还有一个验证就是从 Python 端开始,输入一个有益的语句,走 Tokenizer ,最后的输出这样的结果,我们做对比去验证这个功能是不是完善且常见的。
中间一层是 GPT Model Forward Workflow ,这是最开始谷歌这边提出来的一个常见的 Incode 或者 Decode 的逻辑。这里更多的是 Decode Only 的逻辑,通过 GPT Model Forward 的流程把前面串起来。最开始有一个 Tokenizer 之后的 Token ID ,会经过 Embedding Lookup ,通过一个 ID 查到一组数组,作为一个 Hidden 的变量传到后续的 Layer里面反复的迭代,根据模型的定义,比如 24 层就做 24 次的循环。
中间的每一个 Attention Layer 可以看到它有具体的更多的实现。比如做 Gemm 是一个通用的矩阵的乘法,后续有 Context Attention 和 Decoder Attention 等的实现。 Context Attention 可以进一步展开如上图,这是我们在实现上需要重点去优化和关注的一些算子。
具体 Profile 大语言模型的推理过程,我们会根据不同模型的大小,不同的阶段进行矩阵乘,比如在最开始的处理 Prompt 的阶段,还有最后的 Decode 阶段会有不同矩阵大小去做矩阵乘,也可以看到不同的耗时分布。在 Decode 的阶段,更多的还是 Memory Bound 的处理瓶颈,需要大量内存的搬移,受限于内存带宽是大语言模型推理一个比较显著的特点。在这块做优化需要更多考虑量化或者一些常用的压缩技巧,来减少对内存带宽的消耗。
以上是一个关于矩阵的乘法的示例,矩阵乘一般是有一个 M×K 的 A 矩阵和一个 K×N 的 B 矩阵相乘得到一个 C 矩阵的结果,它的纬度就是 M×N 。最容易想到的办法就是通用的 3 层循环,但是计算机计算的 Cache,如果 B 矩阵不多做任何处理,就会有很多的 Cache Miss ,计算效率是比较低的。
常见的优化就是首先把 B 矩阵做一个转置,做一些向量化的处理,利用 CPU 或者其他硬件品牌的一些硬件加速的指令,进一步加快矩阵乘法,这是矩阵乘的一个基本原理。
有很多不同硬件平台已经有自己的一些优化的实现, Arm Compute Library 会提供 Gemm 的一个算子, Dash-infer 是 Model Scope 提供的推理的框架,会针对移迁的平台提供 BF16 的 Gemm 的算子。 Openblas 也有自己的一些实现,他们各自有不同的特点,集成到大语言模型的框架里面会有自己的好处和痛点。
评估不同实现的基础的性能,为后续在我们的这种语言推理模型上哪种有更好的实现或者应用做一个选择。首先要衡量它的效果,当拿到语言推理的时候,用同样的接口,不同的实现方式来做一些单元测试,去衡量不同的实现的效率,为我们后续做集成打基础。
前面提到这些开源实现有各自不同的特点,如果集成到 RTP LLM 里面,会有一些自己的好处,比如 ACL 集成调用起来比较简单,把相应的 A、B 矩阵的指针传进去,直接做一些通用的 API 的 Configure 和 Run 就可以。但是它还会带来一些问题,比如 Memory Allocation 或者是右边矩阵重复的 Packing ,这样会导致计算效率不高,因为有内存分配和重复的 Packing 在每次矩阵的计算里面。所以我们就想要一个不需要记录状态,实现更加简单的API的集成,比如 Dash-infer 有这种特点,也利用到了 Arm 的一些硬件加速的指令。还有 Arm 新推出的 Kleidi AI ,它也是有 Stateless 的一些特点,但是在最开始那版集成的时候,有一些像BF16 的支持还没有准备,所以选择了 Dash-infer 。
这一版目前在开源里面利用了 BF16 矩阵乘的 Kernel 。接下来要考虑利用量化,比如从 16 比特权重的数据类型压缩成 4 比特的权重。 GPU 有 GPTQ 的量化格式,已经在业界广泛的应用到了,我们做了一些调研,希望通过直接加载 GPTQ 的权重模式在 CPU 上做一个推理。在计算数据类型的选择上,还是利用 BF16 的 Computer Tap ,或者用 I8MM 硬件加速指令集来做 Tap 的选择。
如果要去集成 Kleidi AI 加速引擎的话,首先了解一下它的特点,它更像 ACL 的一个底层加速库,没有更多内存的动态分配、 Scheduling的处理,需要自己去做一些类似的预分配的处理。在硬件的加速决定集上,同一套接口不变的情况下,底层又通过不同的 Dotprod 或者 I8MM ,可以通过同一套 API 实现在不同硬件平台同等效果的加速。
再具体使用上,就是把文件导入进来,通过拷贝或者第 3 方,在我们自己的 Project 里面做一个构建和编译。具体的使用上更多是要考虑的,比如右边的矩阵权重不变的情况下,做了一次的拍品,左边的矩阵和真正的 Matmul 这种 Code Operation 在具体的 Gemm API里面做一个实现。
总结一下 RTP LLM 使能工作要做的事情,主要就是实现 Gemm 适配不同的数据类型的要求,比如 P32 或者 P16,实现自己 Context Attention 的逻辑,还有一些其它的逻辑,就可以自己在自己的硬件平台上做大语言模型的使能。最开始我们选择了 ACL 做 P16 或者 P32 的支持,但是我们发现了 ACL 的一些痛点,效率上不是很高,后续和阿里巴巴一起合作做了一些优化,最终实现的效果就是可以平台上实现 End-to-End 推理,目前在生产线上也得到了部署。
上图是具体一些性能指标的测试,分别是 0.5B 和 7B 的模型, Gemm 的计算类型选用的 BF16 的精度。在不同的 Input Length 和 Output Length 情况下, Context Time 就是处理前面 Prompt 需要的总时间,是毫秒单位; Generate Time 就是每个 Token 之间的时延,单位也是毫秒;最后的 TPS ,通过上图下方扇形图观察在两个不同阶段算子的 Profile 的耗时,为不同的阶段去做进一步的优化,后续的优化基于这些 Profile 数据做进一步的提高。
我们关于未来的计划包含以下四方面,第一就是 Upstream KleidiAI 的支持,第二是加入量化的支持,让 GPTQ 这样 Int4 的格式可以在 Arm 平台上做推理,还有加入更多的像 Bert 一些类型的模型支持,最后就是提高推理的性能,包括 Throughput 和Latency。