智源打造基于 Triton 的大模型算子库,助力 AI 芯片软硬件生态建设
![智源打造基于Triton的大模型算子库,助力AI芯片软硬件生态建设](https://static001.geekbang.org/infoq/6b/6b6ed6aca11cc1dab4f1deb6bac3ba6b.png)
2024 年大模型进入了新的发展阶段,AI 全领域开启了更为迅猛的量变积累。一方面,模型突破了模态的隔离,文本、语音、视觉等各种形式之间产生的丰富的结合,大大增加了模态的多样性;同时,模型参数量从百亿、千亿级膨胀到万亿级,训练数据量从 TB 级达到 PB 级,上下文令牌数量也从几千增长到百万级,计算规模空前庞大;此外,算法结构的创新也带来了 MoE、模型量化、定制算子等更加复杂的计算需求。
而在硬件设备的层面,由英伟达主导的 CUDA 生态以 SIMT 编程模型和 CUDA 编程语言为核心,从高性能算子库、开发工具链和 GPU 驱动等各个层次全面协同,建立了一套完整的体系,并长期占据着高性能计算的领先地位。虽然有生态竞争者——诸如 OpenCL 和 ROCm 等——在试图挑战和替代 CUDA,但是无论在推广程度、使用体验还是计算性能方面都存在一定差距。新兴的各种 AI 芯片硬件架构不一、指令集不一、各自有自己的 AI 编译器,算子库也各自实现,整体呈现十分割裂且难以强大的生态。哪怕有部分 AI 芯片厂商不断模仿、跟进和对齐 CUDA,然而受限于芯片架构的差异和底层的封闭属性,厂商的生态适配仍旧面临开发难度大、任务重以及各自为战的难点。
算力需求增长而资源供应紧张的局面下,多元芯片的开源共建无疑是撼动 CUDA 地位、打造全新格局的机遇。而作为软硬件衔接的关键环节,编译技术这一层次则成为构建统一生态的入口。
为降低大模型新算法的开发门槛、加速芯片架构的创新,智源研究院在今年的智源大会上发布了使用 Triton 语言实现的算子库 FlagGems,以大模型的计算需求为导向,面向多元芯片,借助 Triton 编译器的开源和轻量级优势,提供了一套易适配、高性能的算子实现,以推动基于 Triron 的统一、开源的软硬件生态。
![](https://static001.geekbang.org/infoq/db/db6b3475d145af22beecd7738acb6269.jpeg?x-oss-process=image%2Fresize%2Cp_80%2Fauto-orient%2C1)
算子库技术选型
在编译技术的多条路线中,统一的中间语言、统一的算子接口、统一的开源算子库都经过了一定的探索和实践。
![](https://static001.geekbang.org/infoq/7a/7a6d42967b6fedcd0f2bbf881ddb8b17.png)
统一的中间语言:提供更自由的表达能力和更灵活的优化空间,但其设计研发显著依赖硬件架构相关的底层信息,要求厂商的深度参与和多元芯片的高度协作
统一的算子接口:对上层框架能够保持良好的一致性,对下层则要求芯片厂商各自开发算子库,无法确保一致的算子特性,也无法配合使用框架的图优化技术
统一的开源算子库:能够在厂商之间做到源码共享,省去重复开发的成本且保障一致的算子实现,并且能够由厂商对编译器的个性化适配来发挥硬件相关的性能优势
![](https://static001.geekbang.org/infoq/5b/5b8b101e967eaf589755d34431d57c0c.png)
FlagGems 采用开源算子库的技术路线,基于 OpenAI 推出的 Triton 编程语言及编译器,以 eager 模式接入 PyTorch 框架。相比使用 Triton 的 Inductor 技术,FlagGems 不仅具有同等的简便易用的优点,而且能够更加细致深入地提升算子性能,从而提高模型的训练和推理吞吐。
算子库现状
FlagGems 以 Llama2、Llava2 等十余个国内外热门开源大模型为范例,通过 profiler 工具监测和分析模型的算子调用情况,汇总到一百余个算子。模型包含了语言、视觉、语音、多模态等各种类型,支持前向推理和反向训练的不同场景,基本能够代表主流大模型的应用实践,覆盖大部分的算子调用需求。
FlagGems 以 PyTorch 框架在 eager 模式下调用的原生 CUDA 算子为基线,展现出了强势的性能和可观的优化空间。在访存密集型算子上,FlagGems 有部分算子性能超越 CUDA 算子;在计算密集型算子上,FlagGems 基本追平 CUDA 算子;在融合算子上,FlagGems 全面优于 CUDA 算子。
![](https://static001.geekbang.org/infoq/67/6767b5cc32a92ed27b5b0919ecd4ea3c.png)
FlagGems 目前已经支持 Bert 和 Llama2 两个模型的单机单卡推理,支持 Aquila 模型的多卡分布式训练。
算子库创新点
3.1 运行时优化
在模型运行的过程中,影响性能的不止有 GPU 上核函数的执行效率,还有 CPU 端侧的运行时开销。Triton 语言在定义算子功能时会包装 Autotuner 或 Heuristics 装饰器,其中 Autotuner 使用自动调优的方法为核函数选择编译参数,而 Heuristics 则使用启发式的方法计算编译参数。编译器在执行函数时需要调用所有装饰器的执行序列,复杂的包装会为运行时带来较大的时间开销。 为了尽可能压缩运行时的过程,FlagGems 设计了 LibEntry 机制,使用独立维护的缓存来获取和调用核函数,绕过其他运行时装饰器的多层调用。在算子首次运行时,仍然会沿着 Triton 的原调用路径进行自动调参和启发式计算来取得编译参数,获得的核函数加入缓存;而当算子再次调用到键值相同的核函数时,运行时机制可以直接从缓存取得核函数,快速进行调用运行。
![](https://static001.geekbang.org/infoq/d0/d0fe8fe524a89efc7dc5e49fa0f7e053.png)
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 核心函数。
![](https://static001.geekbang.org/infoq/c0/c00507d44513caa2b7d091f3b107d0ea.png)
![](https://static001.geekbang.org/infoq/47/475915c86af51032c8f47792e46c138e.png)
算子库规划
FlagGems 计划分期实现主流大模型的算子,包括基本数学运算、线性代数、科学计算、张量处理等多种类型的单算子,预计于 2024 年末完成,到目前为止已经完成了超过 50%。
FlagGems 算子库完成后,大模型的开发者和使用者可以仅用一行代码将 ATen 算子替换为 FlagGems,便捷地部署到英伟达 GPU 或其他 AI 芯片上,而无需考虑代码修改或后端适配等问题。FlagGems 的持续优化也会带来更高的推理与训练速度,更加充分地利用芯片算力,辅助和推动大模型进一步的发展。
生态共建
FlagGems 本着生态共建的宗旨,已经得到多家公司参与算子开发,如中科加禾、硅基流动;也得到多家芯片企业适配芯片编译器,如天数智芯、沐曦、昆仑芯等;同时,也与新华三等整机厂商联合推广。
智源研究院定期组织 Triton 中国社区系列活动,推动更多开发者参与。
欢迎体验 FlagGems: https://github.com/FlagOpen/FlagGems
评论