背景介绍
在具有上亿级别参数的大模型推理的过程中,由于模型的权重过大或者是输入数据过大导致无法加载到一张卡上或一台机器上,此时就需要将推理计算任务分布到多张卡或多个计算机上,当前常用的分布式并行策略有数据并行(Data Parallelism),张量并行(Tensor Parallel,TP),流水线并行(Pipeline Parallelism)等。当前大模型主要以transformers及其变体作为基本结构,其中一个transformer层由一个前面的self attention块和后面的MLP层组成,如图1所示,有一种张量并行策略就是在一个transformer层内对self attention块和MLP(Multilayer Perceptron)块分别进行张量并行。
以MLP层举例,对MLP的两个GEMM(通用矩阵乘)计算的右矩阵A和B,将A以列切的方式并行,在得到GELU层输出后,再将B以行切的方式并行,最后在Dropout之前进行一次Allreduce操作。
图2 MLP的Tensor并行

对于具有非常多transformer层的大模型来说,频繁的通信操作使得通信时延成为分布式推理的一个关键因素。加速库通信算子介绍引入加速库的四类通信算子,集成了高性能集合通信库、小包通信机制,提供了单机多卡、多机多卡集合通信原语,实现分布式推理,此外加速库还提供了通信并行算子,提升了推理性能。
父主题: 分布式推理