架构及方案设计
架构设计
- CANN提供异构计算基础能力，包括Runtime、ACLNN、HCCL、DVPP等底层硬件相关的通用能力。
- Operation基于二进制算子提供单算子下发能力。
- Layer基于Operation和ACLNN单算子构图，支持Attention、MLP、Decoder-Layer模块化，也支持多卡Tensor并行及8bit量化。
- Model基于GraphOp积木拼接模型，实现模型流程打通。
- 基于推理引擎统一接口对接Hugging Face，快速实现模型迁移。
框架介绍
该加速库用于Transformer类模型的神经网络推理。加速库中包含了各种Transformer类模型的高度优化模块，如Encoder和Decoder。
加速库底层算子主要由Ascend C编写，基于Ascend底层高性能算子库/TBE算子库实现Transformer类模型的快速推理。
加速库主要层级划分如图2所示，模型的迁移和应用主要基于以下几层展开。
- 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。
- Layer
Layer层为Operation的组合单元，定义为不同模型的通用大颗粒结构，如Transformer_Block/LM_Head层等。
- Model
将定义的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所示。
上图中的符号说明：
- @：表示常规的数学上定义的矩阵相乘。
- x：表示q/k/v输入。
- W：表示q/k/v权重。
对于Transformer Block而言，其包含Attention Block和MLP层，整体切分思路是上述两个结构的组合，具体策略如下：
- Attention Block内：Attention中Q、K、V的权重在head轴竖切，最后一个linear的权重横切。
- MLP中gate_proj和up_proj的权重在head轴竖切，down_proj（最后一个）的权重横切。
- 插入通信算子：单层一共两个AllReduce，分别在Attention和Linear后，Add（残差）之前。
