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

发自凹非寺量子位|公众号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进行初步的测试。用于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