文 | rumor醬

編 | YY

---

一提到模型加速，大家首先想到的就是蒸餾、（結(jié)構(gòu)性）剪枝、量化（FP16），然而稀疏矩陣（sparse matrix）運算一直不被大家青睞。原因也很簡單，一是手邊沒有現(xiàn)成的代碼（懶），二是即使用了，速度也不一定有之前的稠密矩陣（dense matrix）快。

不過，框架的開發(fā)者們并沒有停下他們的腳步，就在不久前，HuggingFace開心地宣布，他們可以支持稀疏矩陣運算啦！75%的sparsity換來了1/4的內(nèi)存和2倍的速度提升！

這個消息還是比較令人激動的，首先稀疏矩陣在存儲上省略了0值，另外在計算上，也沒必要計算和0值相關(guān)的結(jié)果。所以稀疏矩陣能顯著提升運算速度，并節(jié)約大量存儲空間。

不過老司機們的第一反應(yīng)肯定是：效率不錯，但效果（精度）怎么樣？

普普通通……（注意上圖高亮的modest，感覺效果的確一般，否則就直接放結(jié)果了=。=）

Anyway，雖然精度有些美中不足，但單從速度上講已經(jīng)很好了。技術(shù)的進步要一步步來，以HuggingFace的效率，之后應(yīng)該還會有更多動作。

細心的同學們看到這里一定很疑惑，為啥壓縮了4倍，但只提升了2倍速呢？

在pytorch_block_sparse^[1]的Github庫中，官方詳細解釋了這個問題：主要是當前使用的CUTLASS庫還不夠快。

在繼續(xù)下文的討論前，先介紹些GPU編程的小知識：

• CUDA(Compute Unified Device Architecture)：Nvidia家的編程平臺，幫大家把C++等程序轉(zhuǎn)換為GPU指令。

• BLAS(Basic Linear Algebra Subprograms)：一個線性代數(shù)計算的API標準。

• cuBLAS：用cuda實現(xiàn)的GPU BLAS計算庫。像我們所用的Pytorch、Tensorflow都是基于一系列的cuda庫開發(fā)的。只用于dense矩陣運算，已經(jīng)配合GPU優(yōu)化得很好了。這也就是為什么之前大家不在意稀疏矩陣，因為這樣就不能用cuBLAS了，同時還得加上更多的邏輯，可能還不如用cuBLAS直接運算dense要快。

• CUTLASS：CUDA Templates for Linear Algebra Subroutines，一個CUDA C++ 模板集，用于在CUDA上實現(xiàn)更多樣的矩陣乘法計算（GEMM)。

HuggingFace為了實現(xiàn)稀疏矩陣，選取了CUTLASS庫，其本身在計算矩陣乘法時就比cuBLAS庫要慢上兩倍。所以即使理論上75%稀疏度應(yīng)該加速4倍，最后測出來也只提升了2倍。

可見如果深入研究出定制化的稀疏矩陣運算庫，速度上可能還會有所提升。

對于想試用的同學，HuggingFace也一如既往地重視“拿來即用”的體驗，提供了兩種使用方法：

自己寫網(wǎng)絡(luò)時，可以直接用BlockSparseLinear替換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)

想轉(zhuǎn)換別人已經(jīng)寫完的網(wǎng)絡(luò)，可以直接轉(zhuǎn)整個模型。可惜不能自動轉(zhuǎn)參數(shù)，需要重新訓練。

from pytorch_block_sparse import BlockSparseModelPatcher # Create a model patcher mp = BlockSparseModelPatcher()# Selecting some layers to sparsify. # This is the "artful" part, as some parts are more prone to be sparsified, other may impact model precision too much.# Match layers using regexp (we escape the ., just because, it's more correct, but it does not change anything here) # the [0-9]+ match any layer number. # We setup a density of 0.5 on these layers, you can test other layers / densities . mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.intermediate\.dense", {"density":0.5}) mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.output\.dense", {"density":0.5}) mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.attention\.output\.dense", {"density":0.5}) mp.patch_model(model)print(f"Final model parameters count={model.num_parameters()}")# => 68 million parameters instead of 84 million parameters (embeddings are taking a lof of space in Roberta)

目前HuggingFace只邁出了一小步，后續(xù)CUTLASS還會繼續(xù)提升，作者也會復(fù)現(xiàn)更多的學術(shù)成果。除了他們之外，OpenAI在20年初也宣布要將Tensorflow的部分計算代碼移植到Pytorch，谷歌和斯坦福在6月的Paper Sparse GPU Kernels for Deep Learning^[2] 也承諾會放出源碼，大家可以把稀疏矩陣的優(yōu)化學習提上日程啦。

---

文末福利
后臺回復(fù)關(guān)鍵詞【入群】
加入賣萌屋NLP/IR/Rec與求職討論群
有頂會審稿人、大廠研究員、知乎大V和妹紙
等你來撩哦~

參考文獻

[1] pytorch_block_sparse:
https://github.com/huggingface/pytorch_block_sparse
[2] Sparse GPU Kernels for Deep Learning:
https://arxiv.org/abs/2006.10901

總結(jié)

以上是生活随笔為你收集整理的HuggingFace又出炼丹神器！稀疏矩阵运算进入平民化时代！的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。