随着深度学习模型的规模大幅增长，以及人工智能被越来越多地用于广泛的应用之中，降低运行这些模型计算成本的动机从未像现在这样强烈。
稀疏性是降低这些成本的一种方法，并且极具前景。稀疏性通常是指通过剪枝过程在模型权重中诱导出尽可能多的零，然后在进行计算时跳过这些值的方法。
虽然在使用稀疏性来减少FLOPs和参数计数等理论上的成本措施方面已经取得了很大的成功，但在将任务性能保持在可接受的水平上的同时，实现实际速度的提高要困难得多，特别是在通用加速器上，例如使用低精度数字格式的GPU。
为了在实践中实现以上理论上的益处，并在Graphcore拟未的IPU上进一步研究稀疏技术，我们引入了PopSparse。PopSparse是一个库。通过利用IPU独特的硬件特性和数据中定义的块状结构，它可以支持实现快速稀疏操作。我们的目标是两种不同类型的稀疏性：静态的和动态的。静态的稀疏性模式在编译时是固定的，动态的稀疏性模式可以在每次模型运行时改变。我们对IPU上两种模式的稀疏密集矩阵乘法进行了基准测试。结果表明，在具有较大的矩阵尺寸和块尺寸的稀疏级别范围内，在IPU上PopSparse的实现比密集矩阵乘法更快。此外，静态稀疏性通常优于动态稀疏性。虽然以前在GPU上的工作只显示了非常高的稀疏度（通常为99%及以上）的加速，但我们证明了拟未的静态稀疏实现在较低的稀疏度（约90%）下优于FP16中的等效密集计算。
这些成果在线发表在arXiv①上，并被ICLR 2023稀疏神经网络研讨会②接收。
用于深度神经网络中
稀疏密集矩阵乘法的拟未IPU
深度学习中的稀疏性概念通常是指将模型权重稀疏化，以减少相关的存储和计算成本。这通常是通过训练结束时的单一剪枝步骤（Zhu和Gupta，2017）或可能通过一些稀疏的训练制度（Evci等人，2019）来完成。这些方法通常可以实现模型权重减少90至99%，同时将准确性保持在可接受的水平（Hoefler等人，2021）。
虽然可以很容易地实现模型尺寸优势，但由于所产生的非结构化稀疏模式与高度并行的深度学习加速器里的矢量指令集不一致，计算效率的提高往往更难实现（Qin等人，2022）。因此，一种提高稀疏计算效率的方法是对稀疏模式加以某种程度的结构。这促进了广泛的结构化剪枝技术，如神经元（Ebrahimi和Klabjan，2021）、通道（He等人，2017）、块（Gray等人，2017）、2：4（Mishra等人，2021）。然而，与同等的非结构化方法相比，这些方法通常会造成任务性能下降。
我们带来了PopSparse，这是一个在拟未IPU上有效加速稀疏密集矩阵乘法（SpMM）的库。IPU具备多种有助于加速稀疏运算的架构特征：
  • 大型、高带宽的片上SRAM，提高通信量大、算术效率低的操作的性能。
  • 细粒度的MIMD处理模型，将工作划分到1472个独立的计算单元（tile）上，实现了高度的并行性。
  • 提前编译模型，支持在编译期已知稀疏结构时采用进一步的优化。
IPU上PopSparse库里的SpMM
SpMM可以写作：
其中，⊙表示元素乘法，*表示内积，Y ∈ Rᵐ ˣⁿ和X ∈ Rᵏ ˣⁿ 分别为密集输出和输入矩阵。稀疏权重矩阵(M⊙W)通过M ∈ Bᵐ ˣᵏ (B = {0,1})定义，它是代表稀疏模式的掩码，其本身来自块掩码M̂ ∈ B⁽ᵐᐟᵇ⁾ˣ⁽ᵏᐟᵇ⁾，W ∈ Rᵐ ˣᵏ定义权值。
在这个公式中，(M⊙W)具有块稀疏结构，其中形状为b×b的连续正方形权重块包含所有非零元素，给定的块尺寸为b。尺寸m和k被称为输出和输入特征尺寸，n被称为批尺寸，对应于它们在权重稀疏神经网络计算中的典型作用。
我们定义密度，d=∑ᵢⱼMᵢⱼ/(m⋅k)，其中标准项稀疏性对应于(1 — d)。我们使用浮点运算（FLOP）计数来量化SpMM中有用的算术工作量，即：2⋅m⋅k⋅n⋅d。请注意，这只计算对非零值进行的运算，并不取决于块尺寸b。PopSparse库支持非结构化（块尺寸b = 1）或结构化的稀疏性（b = 4、8、16）。
静态稀疏性
对于静态稀疏性，在编译期间，分区器已知稀疏模式和矩阵形状，它将稀疏矩阵的非零元素（具体值尚不清楚）跨k维度拆分为qₖ分区，并将密集矩阵跨n维度拆分为qₙ分区。图1a说明了静态稀疏情况下稀疏矩阵的分布、计算和缩减。在k维度上的拆分不必大小均匀，并且会被拣选，以确保非零元素的平衡分布（图中用彩色小块表示）。模型被编译后，在执行期间，将提供W的特定非零值。由于事先已知是稀疏模式，因此不需要额外的数据交换。然后可以执行与密集矩阵的局部点积计算，最后跨tile进行缩减以给出密集输出矩阵。
动态稀疏
在动态稀疏性方面，稀疏操作数的非零元素的最大数量是固定的，而稀疏性模式在执行过程中更新。为了最大限度地减少计算周期，我们在编译时采用了一个规划器来优化分区，同时适应执行过程中可以定义的全部稀疏模式。与静态稀疏性相比，由于每次运行的稀疏性模式都可能发生变化，因此还有一个额外的步骤，即通过分区器分配稀疏矩阵的索引和非零值。这在图1b中得到了证明。此外，由于需要在m和k维度上进行固定大小的分区，每个分区的非零元素的数量可能不平衡。在图1b中，k维的第二个分区的B元素不适合在Tile 1上。其中三个B元素溢出到Tile 2。然而，Tile 2没有与B元素对应的所有分区信息（m、k、n的片段）来计算结果。因此，需要引入额外的传播步骤来完成计算。
总的来说，与静态情况相比，我们的动态稀疏实现有几个额外的成本：
  • 编译后的额外控制流，这产生了一些成本开销。
  • 我们必须为尽可能大的tile间通信量进行编译，因此平均来说效率较低。
  • 当数据不平衡时，需要额外的传播和计算阶段来重新平衡数据。所需步骤的数量取决于不平衡的程度。
图1：输入特征维度k的静态和动态稀疏分区的说明。图片由作者提供。
基准测试
我们的微观基准实验在IPU和GPU上探索单一稀疏密集矩阵乘法SpMM：Y=(M⊙W)*X。用于IPU实验和GPU基线的API见表1。表2详细列出了基准测试参数的范围。
表1：对IPU和GPU进行基准测试时使用的API。
IPU：该操作在Bow-2000机箱中的单个Bow IPU上执行一次，随机产生稀疏模式和值。我们提取周期计数信息，并在1.85GHz的恒定时钟下将这些周期计数转换成TFLOP/秒值。主机传输被排除在测量时间之外。
GPU：GPU上的基准稀疏实现由cuSPARSE库v11.6.2（NVIDIA，2022）提供。这些都是根据使用的稀疏格式划分的，没有静态和动态稀疏模式之间的区别。针对非结构化和块稀疏，我们分别考虑压缩稀疏行（CSR）和块稀疏行（BSR）。我们生成一个随机的稀疏模式和值，复制到设备上，执行25次迭代操作，并将结果复制到主机上进行验证。我们在5次迭代后开始计时，用cudaEventRecord测量壁钟时间。所有实验都是在4个A100-SXM4-40G机箱中运行的单个A100 GPU上进行的。
表2：为基准扫描的参数范围。图片由作者提供。
这项工作的主要结果显示在图2中。这些图显示了具有固定特征尺寸的计算FLOP/秒（对数比例）与密度的关系，每条线表示同一操作的不同实现方式。通过构造，密集实现显示出密度的线性比例，因为在FLOP/s的计算中不计算零。而另一方面，对于稀疏实现，完美缩放会随着密度的变化，预测恒定的FLOP/s。
图2：IPU FP16和GPU块稀疏矩阵乘法性能随着不同块大小b的密度变化，正方形特征大小m=k=4096，最好超过批大小n。图片由作者提供。
对于密集基线，虽然我们在FP16中大致看到IPU和GPU之间的芯片对芯片的奇偶校验，但必须强调的是，在云端使用的IPU的价格大约是一半③，突出了它们在这些基本操作中的成本效率。此外，对于稀疏操作，我们看到使用IPU的好处甚至更大。
具体来说，对于IPU（图2a），我们看到使用静态稀疏操作和块稀疏性都分别提高了性能。例如，在块尺寸为16的动态情况下，与密集基线相比，在所考虑的所有密度下都有提升，在低密度（<5%）情况下，非结构化（b=1）的静态稀疏性优于密集。此外，当同时使用静态和块状稀疏性方法时，我们看到吞吐量比密集基线多出5倍。
对于GPU（图2b），我们看到稀疏操作的性能要差得多，所测试的情况都没有超过FP16密集基线的性能。然而，我们确实发现，FP32块状稀疏（BSR）的情况随着密度的降低而扩展得非常好，看起来在非常低的密度情况下可能会超过密集基线。
比较这两个平台，我们可以看到，与密集的GPU基线相比，IPU实现了多达5倍的吞吐量优势，当只考虑稀疏操作时，吞吐量提高了25倍。由于有额外的价格优势，我们相信这为在IPU上研究稀疏应用提供了一个颇具说服力的例证。
然而，应该注意的是，即使在IPU上，稀疏操作也只是在某些条件下优于密集基线。更广泛的结果可以在我们的论文④中找到，但是作为近似指南，稀疏性通常会在以下情况下带来IPU（FP16，两者的较大的批尺寸）上密集基线加速：
  • 静态稀疏性，块尺寸b=1，特征（m=k）≥4096，密度d≤1/32。
  • 静态稀疏性，块尺寸b≥4：特征（m=k）≥4096，密度d≤1/8。
  • 动态稀疏性，块尺寸b≥8，特征（m=k）≥4096，密度d≤1/32。
重要的是，在当前的权重稀疏训练和推理的背景下使用这些范围设置具有一定的挑战性。主要原因在于，很难找到不降低性能的块稀疏模式，而且根据理论效率指标（每个非零FLOP），非结构化稀疏性总是表现得与块稀疏性一样好或更好。我们希望这项工作能够证明，作为实现稀疏模型实际加速的方法，块稀疏性颇具前景，并将激发研究界对有效块稀疏剪枝算法的进一步研究。
您可以在Paperspace上的IPU上亲自尝试PopSparse操作⑤，并尝试将它们集成到您自己的模型中。如果您正在积极开展稀疏性研究，而这些研究又可以受益于IPU的支持，您可以考虑参加我们的学术计划⑥，该计划支持研究人员使用IPU取得了很多优秀的成果。
动手尝试
在IPU上尝试SpMM：在Paperspace上运行⑦。
参考资料
[1] Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression, 2017. URL https://arxiv.org/abs/1710.01878.
[2] Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners, 2019. URL https://arxiv.org/abs/1911.11134.
[3] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks, 2021. URL https://arxiv.org/abs/2102.00554.
[4] Eric Qin, Raveesh Garg, Abhimanyu Bambhaniya, Michael Pellauer, Angshuman Parashar, Sivasankaran Rajamanickam, Cong Hao, and Tushar Krishna. Enabling flexibility for sparse tensor acceleration via heterogeneity, 2022. URL https://arxiv.org/abs/2201.08916.
[5] Abdolghani Ebrahimi and Diego Klabjan. Neuron-based pruning of deep neural networks with better generalization using kronecker factored curvature approximation, 2021. URL https://arxiv.org/abs/2111.08577.
[6] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks, 2017. URL https://arxiv.org/abs/1707.06168.
[7] Scott Gray, Alec Radford, and Diederik P. Kingma. GPU kernels for block-sparse weights. 2017.
[8] Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelerating sparse deep neural networks, 2021. URL https://arxiv.org/abs/2104.08378.
[9] NVIDIA. API reference guide for cuSPARSE, 2022. URL https://docs.nvidia.com/cuda/cusparse/.
①https://arxiv.org/abs/2303.16999
②https://www.sparseneural.net/?
③https://docs.paperspace.com/gradient/machines/
④https://arxiv.org/abs/2303.16999
⑤https://console.paperspace.com/github/graphcore-research/notebooks?container=graphcore%2Fpytorch-jupyter%3A3.2.0-ubuntu-20.04&machine=Free-IPU-POD4&file=%2Fsparsity_benchmarks%2FSpMM.ipynb
⑥https://www.graphcore.cn/academic
⑦https://console.paperspace.com/github/graphcore-research/notebooks?container=graphcore%2Fpytorch-jupyter%3A3.2.0-ubuntu-20.04&machine=Free-IPU-POD4&file=%2Fsparsity_benchmarks%2FSpMM.ipynb
⑧https://medium.com/u/b23ad1d05ffe?source=post_page-----3791284a2127--------------------------------
本篇博客作者：
AI工程师李祉颐、研究科学家Douglas Orr、软件工程师Valeriu Ohan和研究科学家Dominic Masters⑧
获取更多Graphcore资讯，阅读深度技术文章，并与其他创新者们一起交流，请至中国官网graphcore.cn，以及关注Graphcore微信、微博和知乎创新社区。
Graphcore中国官网
Graphcore官方微信
Graphcore微博创新社区
Graphcore知乎创新社区
点击阅读原文，查看论文。
继续阅读
阅读原文