LLaMA 2端到端推理打通！来自中国团队

qin123456

Buddy Compiler 端到端 LLaMA2-7B 推理示例已经合并到buddy-mlir仓库［1］主线。


我们在 Buddy Compiler 的前端部分实现了面向 TorchDynamo 的第三方编译器，从而结合了 MLIR 和 PyTorch 的编译生态。


目前，前端部分可以覆盖 LLaMA 计算图，转换到 MLIR 后我们集成了部分向量化和并行优化，并在 AVX512 平台上进行了测试。整个推理过程可以跑通但还需要大量优化。以下是相关链接和现状：

• ［E2E］ Buddy Compiler 端到端LLaMA2-7B 推理示例［2］

•［E2E］上述端到端推理示例目的是展示编译栈设计，并非完备的 LLaMA问答工具

• [Frontend] Buddy Dynamo Compiler[3]

• ［Midend］集成面向矩阵乘法的向量化以及面向循环的并行优化

• ［Backend］端到端示例在 X86 AVX512 机器上进行测试（Ubuntu 22.04）

• ［WIP］ 开发并集成各种优化（现在速度太慢）

• ［WIP］ 开发并集成各种优化（现在速度太慢）

• ［WIP］在多种 CPU 后端进行测试(Arm Neon, RVV, etc.• ［WIP］対接多种硬件平台（GPU，Gemmini Accelerator, etc.)
• ［WIP］増加前端覆盖程度（Stable Diffusion, CLIP, Whisper, etc.)



概述

AI大模型的爆发为软硬件设计带来了新的抓手和机会。随着模型的规模、类型、模态的增加和发散，单一软硬件技术栈覆盖各种场景的能力越来越受限。这种趋势也加深了我对软硬件生态重要性的认识。在我看来，整个生态的最重要的三条设计原则如下（重要程度由高到低）：


• 技术路线标准化（生态的立足之本）
• 上手门槛低（足够多的贡献者和玩家是生态繁荣的关键）
• 优化上限高（決定了生态的发展潜力）



我们基于以上思考打造了 BuddyCompiler，致力于实现软硬件协同设计生态。我们的目标是实现从领域特定编程语言（DSL）到领域特定硬件架构（DSA）的编译流程和协同设计。



本文将介绍如何使用 Buddy Compiler 完成 LLaMA 2 的端到端推理。同时，我们也会分享 Buddy Compiler 的整体设计理念和未来的规划。具体的构建流程请参阅此处的文档［4］（大模型的构建需要十足的耐心和一台性能不错的机器）我们基于PyTorch 和 MLIR 打通了 LLaMA 端到端的推理通路，但是尚未进行完整的性能分析和优化。


目前，我们只应用了针对矩阵乘法的向量化优化，以及针对循环的并行计算优化。


优化和调优的策略仍在开发阶段，因此目前的性能还处于较低水平。我们的初步目标并不是追求极致的性能，而是建立一个标准的端到端通路，在此基础上再做各层级性能优化和调优。


技术路线



技术路线的标准化不仅是我们努力追求的核心原则，而且我们坚信这能够吸引更多的贡献者和合作伙伴，同时还能够有效降低用户的学习门槛和后续的维护成本。


如图所示，我们选择 PyTorch 生态对接各种 AI 模型，选择 MLIR 生态作为 Buddy Compiler 的中同表示。


我们将 Buddy Compiler 作为Torch Dynamo 的自定义编译器组成整个编译栈，从而希望实现将各种AI模型映射到多种硬件架构上的愿景。以下是一些设计点：


使用TorchDynamo作为Trace工具对接Al模型


TorchDynamo 使用 CPython 的 Frame Evaluation API 特性能够做到更加准确的PyTorch 图的捕捉，详情可以参阅PyTorch 的文档［5］。除此之外，PyTorch2.× 提供了全面的编译支持，尤其是提供了非常友好的自定义编译器对接方案［6］。
因此，PyTorch 2.x TorchDynamo 成力了我们对接 AI模型的不二选择。

选择Aten IR作为对接层级


根据 PyTorch 的文档［5］，Core Aten IR是服务于对接后端的算子集合，它相比Prime IR 抽象级别也更高。我们倾向于使用較高的抽象級别映射到 MLIR 从而有更多的空间和信息来进行多层优化，因此我们选择 Aten IR 作为对接层级，可以很方便地映射到 Linalg Dialect 和 TOSADialect 的各种操作。

使用MLIR Python Bindings实现INDyn amo Compiler生成TOSA/Linalg Ops

Buddy Compiler 前端中的 Dynamo Compiler（或者叫做 Dynamo Importer）的作用是将PyTorch 的 Aten IR 转换到 MLIR，同时也负责将模型参数进行处理和打包。

Dynamo Compiler 遍历从 TorchDynamo传入的 FX Graph，并且针对每个 Aten IR 节点进行转换，转换过程使用 MLIRPython Bindings 生成目标 MLIR 代码。


Dynamo Compiler 支持対接 Dialect 的优先级设定，即可以选择 Linalg Dialect或者 TOSA Dialect 作为首选项进行转换。


最终，Dynamo Compiler 负责输出转换后的 MLIR Module 和权重参数。


使用Buddy Compiler工具链进行优化和下降


我们的整个编译通路目前并没有完全使用Python 脚本完成，而是使用 CMake 将前、中、后端集成起来。

这一定程度上简化了并解耦了前、中、后端的开发流程。编译优化和下降的所有Pass 注册到 buddy-opt 工具中，它包括所有的上游 Pass 和 Buddy Compiler 中的优化 Pass，因此 buddy-opt 是上游mlir-opt 的超集。面向通用硬件的编译流程我们使用 MLIR Core Dialect 进行实现，从而达成最大化的复用，我们的工具也因此和所有 LLVM/MLIR 的工具兼容，例如mlir-translate, IIc 等等。


目前我们针对循环采用并行计算的优化，其中相关中间表示和优化 Pass 完全来自上游的 OMP Dialect，可以直接复用并带来不错的优化效果。从此也可以看出统一生态的优势。此外，我们针对粗颗粒度的Operations 没汁了向量化算法迸行編代化，使用 Vector Dialect 也可以实现跨SIMD/Vector 平台的效果。