架构及方案设计

架构设计

基于算子加速库模式的大模型推理架构实现如图1所示。
图1 算子加速库推理架构

LLM类大模型推理采用分层架构设计,基于Transformer结构实现大模型推理。
  • CANN提供异构计算基础能力,包括Runtime、ACLNN、HCCL、DVPP等底层硬件相关的通用能力。
  • Operation基于二进制算子提供单算子下发能力。
  • Layer基于Operation和ACLNN单算子构图,支持Attention、MLP、Decoder-Layer模块化,也支持多卡Tensor并行及8bit量化。
  • Model基于GraphOp积木拼接模型,实现模型流程打通。
  • 基于推理引擎统一接口对接Hugging Face,快速实现模型迁移。

框架介绍

该加速库用于Transformer类模型的神经网络推理。加速库中包含了各种Transformer类模型的高度优化模块,如EncoderDecoder

加速库底层算子主要由Ascend C编写,基于Ascend底层高性能算子库/TBE算子库实现Transformer类模型的快速推理。

加速库主要层级划分如图2所示,模型的迁移和应用主要基于以下几层展开。

图2 加速库层级
  1. Operation

    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。

  2. Layer

    Layer层为Operation的组合单元,定义为不同模型的通用大颗粒结构,如Transformer_Block/LM_Head层等。

  3. Model

    将定义的Layer进行拼接,组成完整的推理流程即为Model层。Model层通常包含了完整的模型推理过程,PyTorch等框架可以直接在对应Python代码侧调用已实现的Model类,实现模型的推理调用。

方案介绍

当前大模型推理加速库的并行方案主要为Tensor Parallelism(张量并行),核心是将模型中的部分计算切分到不同Device上运行,最终通过通信算子合并计算结果。

带有两层Linear Layer(o = W * x + B)的MLP层的张量并行优化示例请参见图3

图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所示。

图4 Attention Block张量并行

上图中的符号说明:

  • @:表示常规的数学上定义的矩阵相乘。
  • x:表示q/k/v输入。
  • W:表示q/k/v权重。

对于Transformer Block而言,其包含Attention Block和MLP层,整体切分思路是上述两个结构的组合,具体策略如下:

  1. Attention Block内:Attention中Q、K、V的权重在head轴竖切,最后一个linear的权重横切。
  2. MLP中gate_proj和up_proj的权重在head轴竖切,down_proj(最后一个)的权重横切。
  3. 插入通信算子:单层一共两个AllReduce,分别在Attention和Linear后,Add(残差)之前。
图5 Transformer Block张量并行