智源打造基于 Triton 的大模型算子库，助力 AI 芯片软硬件生态建设

2024-07-09
北京
本文字数：2782 字
阅读完需：约 9 分钟

2024 年大模型进入了新的发展阶段，AI 全领域开启了更为迅猛的量变积累。一方面，模型突破了模态的隔离，文本、语音、视觉等各种形式之间产生的丰富的结合，大大增加了模态的多样性；同时，模型参数量从百亿、千亿级膨胀到万亿级，训练数据量从 TB 级达到 PB 级，上下文令牌数量也从几千增长到百万级，计算规模空前庞大；此外，算法结构的创新也带来了 MoE、模型量化、定制算子等更加复杂的计算需求。

而在硬件设备的层面，由英伟达主导的 CUDA 生态以 SIMT 编程模型和 CUDA 编程语言为核心，从高性能算子库、开发工具链和 GPU 驱动等各个层次全面协同，建立了一套完整的体系，并长期占据着高性能计算的领先地位。虽然有生态竞争者——诸如 OpenCL 和 ROCm 等——在试图挑战和替代 CUDA，但是无论在推广程度、使用体验还是计算性能方面都存在一定差距。新兴的各种 AI 芯片硬件架构不一、指令集不一、各自有自己的 AI 编译器，算子库也各自实现，整体呈现十分割裂且难以强大的生态。哪怕有部分 AI 芯片厂商不断模仿、跟进和对齐 CUDA，然而受限于芯片架构的差异和底层的封闭属性，厂商的生态适配仍旧面临开发难度大、任务重以及各自为战的难点。

算力需求增长而资源供应紧张的局面下，多元芯片的开源共建无疑是撼动 CUDA 地位、打造全新格局的机遇。而作为软硬件衔接的关键环节，编译技术这一层次则成为构建统一生态的入口。

为降低大模型新算法的开发门槛、加速芯片架构的创新，智源研究院在今年的智源大会上发布了使用 Triton 语言实现的算子库 FlagGems，以大模型的计算需求为导向，面向多元芯片，借助 Triton 编译器的开源和轻量级优势，提供了一套易适配、高性能的算子实现，以推动基于 Triron 的统一、开源的软硬件生态。

算子库技术选型

在编译技术的多条路线中，统一的中间语言、统一的算子接口、统一的开源算子库都经过了一定的探索和实践。

统一的中间语言：提供更自由的表达能力和更灵活的优化空间，但其设计研发显著依赖硬件架构相关的底层信息，要求厂商的深度参与和多元芯片的高度协作
统一的算子接口：对上层框架能够保持良好的一致性，对下层则要求芯片厂商各自开发算子库，无法确保一致的算子特性，也无法配合使用框架的图优化技术
统一的开源算子库：能够在厂商之间做到源码共享，省去重复开发的成本且保障一致的算子实现，并且能够由厂商对编译器的个性化适配来发挥硬件相关的性能优势

FlagGems 采用开源算子库的技术路线，基于 OpenAI 推出的 Triton 编程语言及编译器，以 eager 模式接入 PyTorch 框架。相比使用 Triton 的 Inductor 技术，FlagGems 不仅具有同等的简便易用的优点，而且能够更加细致深入地提升算子性能，从而提高模型的训练和推理吞吐。

算子库现状

FlagGems 以 Llama2、Llava2 等十余个国内外热门开源大模型为范例，通过 profiler 工具监测和分析模型的算子调用情况，汇总到一百余个算子。模型包含了语言、视觉、语音、多模态等各种类型，支持前向推理和反向训练的不同场景，基本能够代表主流大模型的应用实践，覆盖大部分的算子调用需求。

FlagGems 以 PyTorch 框架在 eager 模式下调用的原生 CUDA 算子为基线，展现出了强势的性能和可观的优化空间。在访存密集型算子上，FlagGems 有部分算子性能超越 CUDA 算子；在计算密集型算子上，FlagGems 基本追平 CUDA 算子；在融合算子上，FlagGems 全面优于 CUDA 算子。

FlagGems 目前已经支持 Bert 和 Llama2 两个模型的单机单卡推理，支持 Aquila 模型的多卡分布式训练。

算子库创新点

3.1 运行时优化

在模型运行的过程中，影响性能的不止有 GPU 上核函数的执行效率，还有 CPU 端侧的运行时开销。Triton 语言在定义算子功能时会包装 Autotuner 或 Heuristics 装饰器，其中 Autotuner 使用自动调优的方法为核函数选择编译参数，而 Heuristics 则使用启发式的方法计算编译参数。编译器在执行函数时需要调用所有装饰器的执行序列，复杂的包装会为运行时带来较大的时间开销。为了尽可能压缩运行时的过程，FlagGems 设计了 LibEntry 机制，使用独立维护的缓存来获取和调用核函数，绕过其他运行时装饰器的多层调用。在算子首次运行时，仍然会沿着 Triton 的原调用路径进行自动调参和启发式计算来取得编译参数，获得的核函数加入缓存；而当算子再次调用到键值相同的核函数时，运行时机制可以直接从缓存取得核函数，快速进行调用运行。

LibEntry 优化技术能够将 CPU 端的运行时开销降低约 70%，在运行时开销为主要瓶颈的小规模算子上获得良好的收益，达到与 PyTorch ATen 持平的 CPU 端性能。

3.2 代码生成

Triton 算子实现以 PyTorch 框架定义的张量作为输入和输出，依托 Python 原生特性无缝融入 PyTorch 生态。然而 Triton 函数的操作对象是张量指针，PyTorch 的输入输出本质是张量视图，二者的差异给 FlagGems 的代码编写带来了挑战。

Triton 函数中的张量指针与张量在存储上的布局密切相关，使用偏移量访问存储和读写数据；而张量视图则是存储的变形，更关心张量的形状尺寸，对于张量布局和连续性感知不强。在 Triton 函数处理连续张量时，数据的寻址和访存都比较简单，可以无视指针与视图的差异；而当 PyTorch 向算子输入非连续张量时，Triton 就不能以与连续存储相同的方式来获取数据，强制连续化可能增加的拷贝开销，使得张量步长成为函数中必要的输入信息。

对于高维张量输入，步长的数量是不固定的，鉴于 Triton 无法支持变长元组参数，算子库需要对所有维度分别编写函数代码，这对开发者造成了一定的负担。FlagGems 利用 pointwise 运算的相似性以及偏移量计算代码的模式化特点，使用自动代码生成的方式处理高维张量的非连续访存问题。

自动代码生成机制能够使用模板生成外层包装函数和 Triton 核心函数，自动地处理张量广播、索引空间计算、并行参数计算、地址偏移计算等任务。对于经过转置、切片、跨步等处理的非连续张量，自动代码生成的 GPU 性能可以显著超过强制连续化的方法、达到与 PyTorch ATen 持平的表现。

下面以两个张量相加为例，分别展示 CUDA 写法、Triton 写法和自动代码生成写法，其中自动代码生成写法，通过简单加法逻辑表示，即可生成 Triton 核心函数。

算子库规划

FlagGems 计划分期实现主流大模型的算子，包括基本数学运算、线性代数、科学计算、张量处理等多种类型的单算子，预计于 2024 年末完成，到目前为止已经完成了超过 50%。

FlagGems 算子库完成后，大模型的开发者和使用者可以仅用一行代码将 ATen 算子替换为 FlagGems，便捷地部署到英伟达 GPU 或其他 AI 芯片上，而无需考虑代码修改或后端适配等问题。FlagGems 的持续优化也会带来更高的推理与训练速度，更加充分地利用芯片算力，辅助和推动大模型进一步的发展。