用于更小、更快的语言模型的块稀疏矩阵

一次一个零，节省空间和时间

在之前的博客文章中，我们介绍了稀疏矩阵以及它们对改进神经网络的作用。

基本假设是全密度层通常是多余的，并且可以在不显著损失精度的情况下进行剪枝。在某些情况下，稀疏线性层甚至可以提高精度或/和泛化能力。

主要问题是，当前可用的支持稀疏代数计算的代码严重缺乏效率。我们仍在等待官方 PyTorch 支持。

这就是我们失去耐心并在今年夏天花了一些时间解决这个“空白”的原因。今天，我们很高兴发布扩展pytorch_block_sparse。

该库本身，或者与蒸馏和量化等其他方法结合使用时，可以实现更小、更快的神经网络，Hugging Face 认为这对于让任何人都能以低成本在生产中使用神经网络并改善最终用户体验至关重要。

用法

提供的 BlockSparseLinear 模块是 torch.nn.Linear 的直接替换，在您的模型中使用它非常简单。

# from torch.nn import Linear
from pytorch_block_sparse import BlockSparseLinear

...

# self.fc = nn.Linear(1024, 256)
self.fc = BlockSparseLinear(1024, 256, density=0.1)

该扩展还提供了一个 BlockSparseModelPatcher，允许“即时”修改现有模型，如本示例笔记本中所示。这样的模型可以像往常一样进行训练，而无需更改模型源代码。

NVIDIA CUTLASS

这个扩展基于 Yulhwa Kim 的 cutlass tilesparse 概念验证。

它使用基于CUTLASS 的C++ CUDA 模板进行块稀疏矩阵乘法。

CUTLASS 是用于实现高性能 CUDA 内核的 CUDA C++ 模板集合。通过 CUTLASS，可以在不使用汇编语言代码的情况下，实现接近 cuBLAS 性能的自定义内核。

最新版本包含所有Ampere Tensor Core 原语，在精度损失有限的情况下提供10 倍或更高的加速。pytorch_block_sparse 的后续版本将利用这些原语，因为块稀疏性与 Tensor Cores 的要求 100% 兼容。

性能

在库的当前阶段，稀疏矩阵的性能比其 cuBLAS 优化的密集对应物慢大约两倍，我们相信未来可以改进这一点。

这对于 PyTorch 稀疏矩阵来说是一个巨大的改进：它们当前的实现比密集实现慢一个数量级。

但更重要的一点是，使用稀疏矩阵的性能增益随着稀疏性而增长，因此75% 稀疏的矩阵比密集等效矩阵大约快 2 倍。

内存节省更加显著：对于75% 的稀疏性，内存消耗如您所料减少 4 倍。

未来工作

能够高效地训练块稀疏线性层只是第一步。稀疏模式目前在初始化时是固定的，当然在学习过程中对其进行优化将带来巨大的改进。

因此，在未来的版本中，您可以期待能够优化稀疏模式的工具，以衡量参数的“有用性”。块内NVIDIA Ampere 50% 稀疏模式可能会带来另一个显著的性能提升，就像升级到更新的 CUTLASS 版本一样。

所以，敬请期待在不久的将来，更多的稀疏性优化！