Hugging Face pytorch_block_sparse 发布

Hugging Face 发布了 pytorch_block_sparse，这是一个旨在通过实现块稀疏矩阵乘法来创建更小、更快的神经网络的扩展。该库解决了现有稀疏代数计算中的效率差距，从而降低语言模型的内存占用并降低生产成本。

块稀疏实现与用法

pytorch_block_sparse 提供了一个 BlockSparseLinear 模块，可作为 torch.nn.Linear 的即插即用替代方案。这允许开发人员通过在初始化期间指定密度参数，以极少的代码更改将稀疏性集成到其模型中。

此外，该库还包含一个 BlockSparseModelPatcher，可以“即时”修改现有模型。这使得模型可以被稀疏化，然后进行常规训练，而无需更改原始模型源代码。

技术基础：NVIDIA CUTLASS

该扩展基于 Yulhwa Kim 的 cutlass tilesparse 概念验证，并利用基于 NVIDIA CUTLASS 的 C++ CUDA 模板进行块稀疏矩阵乘法。

CUTLASS 是一个用于实现高性能 CUDA 内核的 CUDA C++ 模板集合，它能够在不需要汇编语言的情况下实现接近 cuBLAS 的性能水平。该库与 Ampere Tensor Core 原语兼容，在精度损失有限的情况下，可以提供 10 倍或更多的加速。

性能与内存效率

虽然该库中稀疏矩阵的当前性能比经过 cuBLAS 优化的稠密对应物慢大约两倍，但与标准 PyTorch 稀疏矩阵相比，这是一个显著的改进，因为标准 PyTorch 稀疏矩阵比稠密矩阵慢一个数量级。

性能增益随稀疏程度而增加。例如，一个稀疏度为 75% 的矩阵比其稠密等效矩阵快大约 2 倍。内存节省更为显著；在 75% 的稀疏度下，内存消耗减少了 4 倍。

未来开发路线图

pytorch_block_sparse 的未来迭代旨在超越初始化时建立的固定稀疏模式。

计划中的改进包括：

  • Sparsity Pattern Optimization：开发工具来衡量参数的“有用性”以在学习过程中优化稀疏模式。
  • Hardware Acceleration：在块内集成 NVIDIA Ampere 50% 稀疏模式以实现进一步的性能提升。
  • CUTLASS Updates：利用更新版本的 CUTLASS 来增强整体效率。

Sources