LLaMA2端到端推理打通!来自中国团队编译器

作者: 小郑 2023-11-30 12:47:52
阅读(68)
发自凹非寺量子位|公众号QbitAIBuddyCompiler端到端LLaMA2-7B推理示例已经合并到buddy-mlir仓库[1]主线。我们在BuddyCompiler的前端部分实现了面向TorchDynamo的第三方编译器,从而结合了MLIR和PyTorch的编译生态。目前,前端部分可以覆盖LLaMA计算图,转换到MLIR后我们集成了部分向量化和并行优化,并在AVX512平台上进行了测试。整个推理过程可以跑通但还需要大量优化。以下是相关链接和现状:[E2E]BuddyCompiler端到端LLaMA2-7B推理示例[2][E2E]上述端到端推理示例目的是展示编译栈设计,并非完备的LLaMA问答工具[Frontend]BuddyDynamoCompiler[3][Midend]集成面向矩阵乘法的向量化以及面向循环的并行优化[Backend]端到端示例在X86AVX512机器上进行测试(Ubuntu22.04)[WIP]开发并集成各种优化(现在速度太慢)[WIP]在多种CPU后端进行测试(ArmNeon,RVV,etc.)[WIP]对接多种硬件平台(GPU,GemminiAccelerator,etc.)[WIP]增加前端覆盖程度(StableDiffusion,CLIP,Whisper,etc.)概述AI大模型的爆发为软硬件设计带来了新的抓手和机会。随着模型的规模、类型、模态的增加和发散,单一软硬件技术栈覆盖各种场景的能力越来越受限。这种趋势也加深了我对软硬件生态重要性的认识。在我看来,整个生态的最重要的三条设计原则如下(重要程度由高到低):技术路线标准化(生态的立足之本)上手门槛低(足够多的贡献者和玩家是生态繁荣的关键)优化上限高(决定了生态的发展潜力)我们基于以上思考打造了BuddyCompiler,致力于实现软硬件协同设计生态。我们的目标是实现从领域特定编程语言(DSL)到领域特定硬件架构(DSA)的编译流程和协同设计。本文将介绍如何使用BuddyCompiler完成LLaMA2的端到端推理。同时,我们也会分享BuddyCompiler的整体设计理念和未来的规划。具体的构建流程请参阅此处的文档[4](大模型的构建需要十足的耐心和一台性能不错的机器)我们基于PyTorch和MLIR打通了LLaMA端到端的推理通路,但是尚未进行完整的性能分析和优化。目前,我们只应用了针对矩阵乘法的向量化优化,以及针对循环的并行计算优化。优化和调优的策略仍在开发阶段,因此目前的性能还处于较低水平。我们的初步目标并不是追求极致的性能,而是建立一个标准的端到端通路,在此基础上再做各层级性能优化和调优。技术路线技术路线的标准化不仅是我们努力追求的核心原则,而且我们坚信这能够吸引更多的贡献者和合作伙伴,同时还能够有效降低用户的学习门槛和后续的维护成本。如图所示,我们选择PyTorch生态对接各种AI模型,选择MLIR生态作为BuddyCompiler的中间表示。我们将BuddyCompiler作为TorchDynamo的自定义编译器组成整个编译栈,从而希望实现将各种AI模型映射到多种硬件架构上的愿景。以下是一些设计点:使用TorchDynamo作为Trace工具对接AI模型TorchDynamo使用CPython的FrameEvaluationAPI特性能够做到更加准确的PyTorch图的捕捉,详情可以参阅PyTorch的文档[5]。除此之外,PyTorch2.x提供了全面的编译支持,尤其是提供了非常友好的自定义编译器对接方案[6]。因此,PyTorch2.xTorchDynamo成为了我们对接AI模型的不二选择。选择AtenIR作为对接层级根据PyTorch的文档[5],CoreAtenIR是服务于对接后端的算子集合,它相比PrimeIR抽象级别也更高。我们倾向于使用较高的抽象级别映射到MLIR从而有更多的空间和信息来进行多层优化,因此我们选择AtenIR作为对接层级,可以很方便地映射到LinalgDialect和TOSADialect的各种操作。使用MLIRPythonBindings实现DynamoCompiler生成TOSA/LinalgOpsBuddyCompiler前端中的DynamoCompiler(或者叫做DynamoImporter)的作用是将PyTorch的AtenIR转换到MLIR,同时也负责将模型参数进行处理和打包。DynamoCompiler遍历从TorchDynamo传入的FXGraph,并且针对每个AtenIR节点进行转换,转换过程使用MLIRPythonBindings生成目标MLIR代码。DynamoCompiler支持对接Dialect的优先级设定,即可以选择LinalgDialect或者TOSADialect作为首选项进行转换。最终,DynamoCompiler负责输出转换后的MLIRModule和权重参数。使用BuddyCompiler工具链进行优化和下降我们的整个编译通路目前并没有完全使用Python脚本完成,而是使用CMake将前、中、后端集成起来。这一定程度上简化了并解耦了前、中、后端的开发流程。编译优化和下降的所有Pass注册到buddy-opt工具中,它包括所有的上游Pass和BuddyCompiler中的优化Pass,因此buddy-opt是上游mlir-opt的超集。面向通用硬件的编译流程我们使用MLIRCoreDialect进行实现,从而达成最大化的复用,我们的工具也因此和所有LLVM/MLIR的工具兼容,例如mlir-translate,llc等等。目前我们针对循环采用并行计算的优化,其中相关中间表示和优化Pass完全来自上游的OMPDialect,可以直接复用并带来不错的优化效果。从此也可以看出统一生态的优势。此外,我们针对粗颗粒度的Operations设计了向量化算法进行编译优化,使用VectorDialect也可以实现跨SIMD/Vector平台的效果。如果希望面向特定加速器(例如GemminiAccelerator)生成代码,也可以使用buddy-translate和buddy-llc生成到特定于硬件的LLVMIR,从而最终生成加速器的硬件指令。Forward函数搭配BuddyCompilerTextContainer完成端到端推理在完成编译优化和下降之后,模型的Forward函数将会被构建为共享库。由于我们很多场景需要在模拟器和开发平台上测试,因此我们没有选择将执行的流程交给Python生态,而是进行AOT编译,生成可执行文件。为了配合从MLIR构建出来的Forward函数实现端到端的推理,我们提供了C++版本的TextContainer和MemRefContainer来作为文本输入输出的Tokenizer和数据容器。最终,在C++的main函数中将输入的文本和权重参数加载到数据容器,然后调用Forward函数进行推理,输出的Token再交由TextContainer进行后处理即可得到最终的文本。未来工作我们目前打通了LLaMA到CPUSIMD/Vector平台的通路,使用X86AVX512进行初步的测试。LLaMA2端到端推理打通!来自中国团队编译器用于VectorDialect的跨平台性,ArmNeon和RISC-VVectorExtesion也是可以天然支持的,我们正在进行广泛测试。同时我们也在支持尝试将BuddyCompiler中的Gemmini加速器支持对接到大模型推理的通路上。此外,GPU的优化Pass也在开发中。在模型层面,我们希望测试并对接更多的多模态大模型,进一步提升BuddyDynamoCompiler前端的覆盖程度。在完成上述工作后,下一个里程碑是将多模态大模型编译到多种硬件平台上。总的来说,接下来前、中、后端将会相对解耦地进行开发,最终进行对接。前端:对接更多模型完善算子覆盖程度。中端:进行更详细的性能分析,面向各种计算负载开发不同层级的优化。后端:对CPU各SIMD/Vector平台进行测试,完成面向GPU/Gemmini加速器的端到端通路。三个部分均相对完备的时候就考虑用Python接口将所有工具链包装起来,形成一个更为优雅的使用流程。总结如今大模型推理的软件栈也层出不穷,技术路线也各不相同。我们使用的Torch2.x+BuddyCompiler的编译栈设计策略实际上是希望融合PyTorch和MLIR两大生态进行AI计算。当然,在MLIR生态里面做大模型推理,我们认为Torch-MLIR+IREE目前是相对比较完备的解决方案。nod.ai的SHARKTurbine[7]就使用了这种技术路线。相比于Torch-MLIR搭配IREE的组合,BuddyCompiler更强调SimpleButPowerful设计,采用极致复用策略和完全代码生成策略。相比于Torch-MLIR的TorchDialect层级,BuddyCompiler更偏向直接复用TOSA/Linalg对接AtenIR;相比于IREE覆盖一切后端的Runtime和ExecutionEngine的设计,BuddyCompiler更偏向进行完全代码生成。SharkTurbine封面[8]是一个飞机涡轮发动机,这和他们的技术路线非常契合,TorchDynamo+Torch-MLIR+IREE是一个极其精密且重型的编译栈,这样的“发动机“理论上可以带着他们飞跃任何高山沟壑。相比而言,BuddyCompiler更像是电动汽车的三电平台,可以以此为基础打造各种性格的电动汽车。对我们来说,LLaMA的支持不是起点也不是终点,是在探索路上偶遇的一座高山,我们希望翻过它,看看山那边的世界,尤其是开着自己造的车!致谢感谢所有BuddyCompiler的贡献者,特别感谢一起努力跑通LLaMA的伙伴:zhanghb97,weilinquan,xTayEx,EllisLambda,Lester-1,LAJIidea,SForeKeeper,LHY-24,xlinsist,qingqing12138.同时感谢OSPP组委会提供的开源项目席位。关于IntelligentComputingIntelligentComputing是由之江实验室和美国科学促进会(AAAS)共同创办,是《科学》合作期刊框架中智能计算领域的第一本开放获取(OpenAccess)国际期刊。期刊以“面向智能的计算、智能驱动的计算”以及“智能、数据与计算驱动的科学发现”为主题,主要刊载原创研究论文、综述论文和观点论文。[1]https://github.com/buddy-compiler/buddy-mlir[2]https://github.com/buddy-compiler/buddy-mlir/tree/main/examples/BuddyLlama#buddy-compiler-llama-example[3]https://github.com/buddy-compiler/buddy-mlir/tree/main/frontend/Python[4]https://github.com/buddy-compiler/buddy-mlir/tree/main/examples/BuddyLlama[5]https://pytorch.org/docs/stable/torch.compiler.html#torch-compiler[6]https://pytorch.org/docs/stable/torch.compiler_custom_backends.html#custom-backends[7]https://github.com/nod-ai/SHARK-Turbine#shark-turbine[8]https://github.com/nod-ai/SHARK-Turbine#shark-turbineg