【GTC 22】PaiSparse：面向深度学习的稀疏算子库

研究动机

随着越来越多超大语言模型的出现，如何将这些超大语言模型在实际业务中部署落地成为了一个急需解决的问题。这些超大语言模型拥有百亿千亿甚至于万亿的参数，如此巨大的参数使得无法将超大模型直接部署在单张卡上，而需要几个甚至几十个 GPU 参与模型推理，会消耗非常大的计算资源。因此将超大语言模型进行部署的时候，就需要利用模型优化技术将它们进行压缩，从而减少模型参数量并提高推理速度。一些常用的模型优化技术如剪枝和量化虽然能够有效的降低模型大小，但是这些算法却很难在保证精度的情况下达到较大的压缩率。而对于超大语言模型而言，只有达到较高的压缩率，才有可能在单卡上部署起来。因此我们选择了非结构化稀疏来压缩超大语言模型，从而使得他们能够在保证精度的情况下达到较高的压缩率。

目前已有的稀疏算子库 CuSparse 虽然能够有效加速稀疏计算，但是我们发现神经网络训练出来的稀疏分布一般不是均匀的，而是呈现一定的分布特点（比如权重中一些行的稀疏率很高，一些行的稀疏率很低）。然而 CuSparse 是为了通用的高性能计算而开发出来的，一般而言均匀稀疏分布才能在 CuSparse 上得到较好的加速性能，而神经网络这样的分布特点无法在 CuSparse 上获得较好的加速效果。和高性能计算场景不一样的是，在神经网络推理中，我们可以预先根据权重的分布特点对最佳的稀疏算子实现进行搜索，从而生成针对当前权重稀疏分布更优的实现。

PaiSparse

PaiSparse 的设计思路如下图所示：

对于稀疏算子而言，主要可以分为三个计算阶段：

1. Prolog Stage：从 gpu 内存中读取 CSR 格式的数据并存入寄存器中。
2. Reduction Stage：负责进行稀疏计算。
3. Epilog Stage：将计算结果存入 gpu 内存中。

我们预先在上述的三个计算阶段中定义了很多代码片段，并通过集成这些代码片段生成最终的稀疏算子。对于单个稀疏算子，我们的搜索空间包括 launch dim, index data type, thread group 和 reduction method 等。

其中，Index data type 的搜索空间包括 int 16和 int 32。int 16 可以有效减少内存读取的开销，但是会增加bit操作，如何进行选取会根据 weight 的具体分布来决定。Thread group 决定了有多少 thread 参与结果中一行的计算，这个参数可以使得算子的计算更加平衡。

除此之外，我们发现 reduction method 也是影响稀疏算子性能非常关键的一个因素，因此我们提出了四种 reduction method 并通过搜索决定最佳实现。这四种方法分别是：Register reduction，SHFL reduction，Shared memory reduction 和 Atomic reduction。针对一个特定的稀疏权重，我们需要首先计算需要执行多少次 register reduction，然后在此基础上计算执行多少次 SHFL reduction。接下来计算需要执行多少次 shared memory reduction，最后在执行完上述三种 reduction 后，剩余的数据累积到 atomic reduction。通过结合这4种 reduction method，PaiSparse 通过 nctrc 生成所有稀疏算子实现，并通过评估它们的性能决定最优实现。

实验结果

我们首先评估了 PLUG(Pre-training for Language Understanding and Generation)一个中文超大语言模型进行 99% 稀疏的精度结果。该模型拥有 270B 的参数量，通过 1TB以上的训练数据训练而成。

如上图所示，通过 99% 稀疏，PLUG 模型依然能够拥有相比于一些较小的语言模型更优的模型精度。

接下来，我们评估了在 99% 稀疏率下，PaiSparse 的稀疏算子性能。

首先在 SpMV 上，当 batch size=1 的时候，在 FP16 的数据格式上，PaiSparse 相比于CuBlas 有 2.7～22 倍的加速，相比于 CuSparse 有 1.22～1.91 倍的加速。在FP32的数据格式上，PaiSparse 相比于 CuBlas 有 27～37 倍的加速，相比于 CuSparse 有1.15~1.6 倍的加速。

在 SpMM 上，PaiSparse 相比于 CuSparse 拥有 1.8～4.9 倍的加速效果。

最后，比较 99% 稀疏下 PLUG 模型的性能。首先99%的稀疏可以在使用 4 张 GPU 进行推理的情况下，将显存占用从13.5GB 降低到 1.8GB，极大的减少了模型的显存占用，同时相比于原始 PLUG 模型至少需要4张V100进行推理的情况，稀疏后的 PLUG 可以部署在单张 V100 上，减少了部署的成本。比较整个 PLUG 网络端到端的推理速度，集成了 PaiSparse的稀疏网络相比于原始 Dense 网络可以达到 1.6 倍的加速。

总结

稀疏可以在保证模型精度的情况下，有效减少大模型需要的计算资源。我们提出的PaiSparse 通过分析深度学习里权重稀疏分布的特点，实现了搜索稀疏算子的整套方案，从而可以针对具体的稀疏分布生成的最优的稀疏算子实现，达到更好的计算性能。