该加速库用于Transformer类模型的神经网络推理。加速库中包含了各种Transformer类模型的高度优化模块,如Encoder和Decoder。
加速库底层算子主要由Ascend C编写,基于Ascend底层高性能算子库/TBE算子库实现Transformer类模型的快速推理。
加速库主要层级划分如图2所示,模型的迁移和应用主要基于以下几层展开。
Operation可以分为两类,加速库本身提供的基础Operation和用户根据自身需求使用Ascend C开发的PluginOperation。
Operation为加速库成图的最小单元,是Transformer类模型可识别的最小单元模块,通常对应PyTorch的Module或者Function,如:ROPE、Self-Attention、FFN等。Operation通常也是基于不同底层算子(Kernel)的组合,例如FFN是matmul/add/activation等算子的组合,是模型之间可复用的模块。
底层Kernel算子由公共算子库提供,对用户不可见。模型适配不涉及Kernel的开发,如果存在不支持的Operation,用户可以给加速库提需求或者用户使用Ascend C开发自定义的PluginOperation。
Layer层为Operation的组合单元,定义为不同模型的通用大颗粒结构,如Transformer_Block/LM_Head层等。
将定义的Layer进行拼接,组成完整的推理流程即为Model层。Model层通常包含了完整的模型推理过程,PyTorch等框架可以直接在对应Python代码侧调用已实现的Model类,实现模型的推理调用。
当前大模型推理加速库的并行方案主要为Tensor Parallelism(张量并行),核心是将模型中的部分计算切分到不同Device上运行,最终通过通信算子合并计算结果。
带有两层Linear Layer(o = W * x + B)的MLP层的张量并行优化示例请参见图3。
第一个Linear的Weight和Bias在dim = 1轴切分,第二个Linear的Weight在dim = 0 轴切分,Bias除以2,得到的两个结果z1和z2,最后通过AllReduce算子相加得到最终输出z。
对于Attention Block,整体逻辑一致,需要对Q/K/V矩阵对应权重在dim = 1切分,然后将最后一个Linear的Weight在dim = 0轴切分,最后通过AllReduce算子相加得到最终输出结果,如图4所示。
上图中的符号说明:
对于Transformer Block而言,其包含Attention Block和MLP层,整体切分思路是上述两个结构的组合,具体策略如下: