用於更小、更快的語言模型的塊稀疏矩陣

一次一個零，節省空間和時間

在之前的部落格文章中，我們介紹了稀疏矩陣以及它們對改進神經網路的作用。

基本假設是全密度層通常是多餘的，並且可以在不顯著損失精度的情況下進行剪枝。在某些情況下，稀疏線性層甚至可以提高精度或/和泛化能力。

主要問題是，當前可用的支援稀疏代數計算的程式碼嚴重缺乏效率。我們仍在等待官方 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 版本一樣。

所以，敬請期待在不久的將來，更多的稀疏性最佳化！