一文了解AOL算子加速库
过去一年,随着ChatGPT的发布与快速迭代,基于大数据量、大参数量、大算力的预训练大模型已成为人工智能产业的主要路线。大模型的普及与发展不仅依靠模型本身的创新,更依赖于算力底座的支撑以及软件生态的繁荣,需要伙伴和开发者的积极参与。为加速大模型算力释放,昇腾提供AOL(Ascend Operator Library)算子加速库,不仅覆盖Softmax、MatMul等基础算子,也包括了大模型结构泛化的Flash Attention等高性能融合算子,开发者可以直接使用昇腾内置的算子加速库使能大模型创新与应用。
1 为什么需要算子加速库
算子实际上是面向深度学习任务高度优化的函数,它是网络中层或者节点的计算逻辑,开发者在实现这些计算逻辑的时候,不仅要实现功能逻辑,还要考虑硬件指令的适配,如何支持不同大小/类型的输入,以及如何针对不同的输入Shape进行切分,并保证算子的运行性能。可见,一个算子的代码如此复杂,如果神经网络中的每一个算子都由开发人员实现一遍,那么可想而知,神经网络构建的工作量有多大,难度有多高,特别是面对当前日益复杂的大模型。能否针对某种计算逻辑开发一个通用实现?这样,开发者便像搭房子一样构建网络而无需关注具体的算子实现。
为此,昇腾算子加速库提供了丰富的经过深度优化的高性能算子,且更亲和昇腾AI处理器,可供开发者直接调用,使能大模型极致性能优化,正是这些算子组成了能够让AI任务在昇腾硬件上飞速执行的弹药库。
2 昇腾基础算子加速库
昇腾基础算子加速库包括1400+种高性能算子,其中包括常用深度学习算法计算类型的NN(NeuralNetwork)算子库、能够提供高性能媒体数据处理能力的DVPP(Digital Vision Pre-Processing)算子库、分布式训练中提供卡间高效数据传输的HCCL(Huawei Collective Communication Library)算子库等。
- NN(NeuralNetwork)算子库:覆盖了包括TensorFlow、Pytorch、MindSpore、ONNX等框架在内的常用深度学习算法的计算类型,包括Softmax、MatMul、Convolution等典型算子,在算子库中占有最大比重。
- BLAS(Basic Linear Algebra Subprograms)算子库:BLAS为基础线性代数程序集,是进行向量和矩阵等基本线性代数操作的数值计算库,CANN支持通用的矩阵乘和基础的Max、Min、Sum、乘加等基础数学运算。
- DVPP(Digital Video Pre-Processor)算子库:DVPP是昇腾AI处理器内置的图像处理单元,支持高性能视频编解码、图片编解码、图像裁剪缩放等预处理能力,通过DVPP算子可快速体验媒体处理硬加速能力。
- AIPP(AI Pre-Processing)算子库:AIPP可用于在AI Core上完成图像预处理,包括图像尺寸更改、色域转换(转换图像格式)、减均值/乘系数(图像归一化),并与模型推理过程融合,满足推理输入要求。
- HCCL(Huawei Collective Communication Library)算子库:HCCL是基于昇腾硬件的高性能集合通信库,提供单机多卡以及多机多卡间的数据并行、模型并行集合通信方案。支持AllReduce、Broadcast、Allgather、ReduceScatter、AlltoAll等通信原语,Ring、Mesh、HD等通信算法,在HCCS、RoCE和PCIe高速链路实现集合通信。
3 昇腾融合算子加速库
在主流大模型网络模型中,大量使用典型的Multi-Head Attention结构,带来了巨大的计算和内存开销。其运行过程中,矩阵乘和softmax结果存放在片上内存会带来巨大的内存消耗,访存性能严重下降,甚至会导致模型无法正常运行,同时网络中的矩阵和向量计算串行执行,也会导致硬件算力发挥受限。
斯坦福的Tri DAO提出了FlashAttention融合算子,其原理是对attention处理过程进行切分和计算等价,使得attention的多个步骤在一个算子中完成,并且通过多重循环、每次处理一小部分数据,以近似流式的方式访问片上内存,减少了片上内存访问的总数据量,并能够将计算和数据搬运更好的重叠隐藏。
CANN针对昇腾AI处理器的片上内存和缓存大小,以及数据搬运通路,优化实现FlashAttention融合算子,充分利用片上缓存,提升Attention处理性能。根据实测,在一些典型场景中,CANN的FlashAttention算子相比融合前的小算子取得了5倍以上的性能提升。开发者可直接调用相关算子接口使能大模型极致性能优化。
Flash Attention算子融合机制:
由此可知,算子融合(Operator Fusion)基本操作就是将多个连续的操作或算子合并成一个单一的算子,常用于深度学习和其他计算密集型任务中,以减少计算和内存开销。使用算子融合有助于提高执行效率,减少数据传输和临时存储,以及提高缓存利用率。
除了Flash Attention类算子,昇腾还提供了MoE结构下的高性能融合算子、针对分布式场景提供MC²(Merged Compute and Communication)通算融合类算子等。
MoE (Mixture-of-Experts,混合专家系统)结构将稠密网络的FFN层扩展成具有相同结构的专家网络,通过路由或门控网络决定激活哪些expert进行计算,从而能够实现在低成本下训练更大的模型。昇腾针对MoE结构提供了诸多高性能融合算子,比如GroupedMatMul融合算子可支持专家分组能力、使能多专家多核并行计算;MoeGatingTopKSoftmax将gating部分的softmax以及topk计算融合;MoEInitRouting、MoEFinalizeRouting对routing进行加速处理以提升MoE结构在NPU上的执行性能。
MC²通算融合类算子的工作原理是将通信和计算的操作过程融合,使得计算和通信可以流水并行提升性能,典型算子包括AllGatherMatMul、MatMulReduceScatter、MatMulAllReduce等,这些算子将AllGather、ReduceScatter、AllGather等通信过程和MatMul计算融合,使SP、TP等并行场景下的通信和矩阵计算(Linear)流水并行,可将GPT-3的整网性能提升10%左右。
目前,昇腾已开放部分融合算子代码样例,开发者可以访问https://gitee.com/ascend/cann-ops-adv进行体验和使用。这意味着,开发者不仅可以直接使用昇腾融合算子加速库使能大模型极致性能,也可以从自有的大模型算法和应用出发,基于开放的融合算子代码样例和Ascend C算子编程API开展定制优化,从而提升大模型运行性能或优化资源占用,构建差异化算法竞争力。
4 更多介绍
关于昇腾算子加速库的相关介绍就到这里,欢迎大家关注后续技术分享。如需获取更多学习资源请登录昇腾社区。
更多CANN开放能力介绍:
《Ascend C 2.0全新升级,支撑大模型融合算子高效开发》
《深度解读昇腾CANN小shape算子计算优化技术,进一步减少调度开销》
- 点赞
- 收藏
- 关注作者
评论(0)