LLaMA 利来老牌国际官网2端到端推理打通来自中国团队

发布时间：2024-09-26 01:35:47 丨 浏览次数：

　　

　　此外，GPU 的优化 Pass 也在开发中★■■◆◆。在模型层面，我们希望测试并对接更多的多模态大模型◆■★■◆，进一步提升 Buddy Dynamo Compiler 前端的覆盖程度。在完成上述工作后◆◆■★★，下一个里程碑是将多模态大模型编译到多种硬件平台上。

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

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

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

　　同时我们也在支持尝试将 Buddy Compiler 中的 Gemmini 加速器支持对接到大模型推理的通路上★■◆★◆。

　　Intelligent Computing是由之江实验室和美国科学促进会（AAAS）共同创办，是《科学》合作期刊框架中智能计算领域的第一本开放获取（Open Access）国际期刊。期刊以“面向智能的计算、智能驱动的计算”以及“智能■★■■◆、数据与计算驱动的科学发现”为主题，主要刊载原创研究论文、综述论文和观点论文◆◆■。

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

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

　　由于我们很多场景需要在模拟器和开发平台上测试◆◆■★◆◆，因此我们没有选择将执行的流程交给 Python 生态◆★，而是进行 AOT 编译，生成可执行文件。

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

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

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

　　最终，在 C++ 的 main 函数中将输入的文本和权重参数加载到数据容器，然后调用 Forward 函数进行推理，输出的 Token 再交由 Text Container 进行后处理即可得到最终的文本★■■。

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

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

　　在完成编译优化和下降之后，模型的 Forward 函数将会被构建为共享库。

　　本文为澎湃号作者或机构在澎湃新闻上传并发布◆■★，仅代表该作者或机构观点，不代表澎湃新闻的观点或立场◆■，澎湃新闻仅提供信息发布平台。申请澎湃号请用电脑访问。

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

　　三个部分均相对完备的时候就考虑用 Python 接口将所有工具链包装起来，形成一个更为优雅的使用流程。

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

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

　　Shark Turbine 封面[8]是一个飞机涡轮发动机利来老牌国际官网，这和他们的技术路线非常契合利来老牌国际官网，TorchDynamo + Torch-MLIR + IREE 是一个极其精密且重型的编译栈■◆，这样的★★★“发动机■★■◆★★“理论上可以带着他们飞跃任何高山沟壑。相比而言■■◆★★■，Buddy Compiler 更像是电动汽车的三电平台，可以以此为基础打造各种性格的电动汽车。对我们来说，LLaMA 的支持不是起点也不是终点，是在探索路上偶遇的一座高山，我们希望翻过它，看看山那边的世界，尤其是开着自己造的车！

　　如今大模型推理的软件栈也层出不穷，技术路线也各不相同。我们使用的 Torch 2.x + Buddy Compiler 的编译栈设计策略实际上是希望融合 PyTorch 和 MLIR 两大生态进行 AI 计算。当然，在 MLIR 生态里面做大模型推理★★，我们认为 Torch-MLIR + IREE 目前是相对比较完备的解决方案■■。nod■◆★.ai 的 SHARK Turbine[7] 就使用了这种技术路线。