TW Community

Sparse Tensor Core 初見

NVIDIA 介紹了 A100 有搭載 Sparse Tensor Core[1],一開始看到名字,還以為有些函數可以因此受益並改進了,但大略看過後, Sparse Tensor Core 的主要用途應該還在 deep learning (深度學習) 上面。

首先,Sparse Tensor Core 並不是可以計算 Sparse Matrix (CSR, ELL, 等在前面[2]有介紹的) ,他是能算一種特殊存法的矩陣,(如果從矩陣觀點來看我不怎麼會叫他 Sparse 就是了)

給定一個矩陣,每四個元素只有兩個元素能儲存值,因此會類似於 ELL 除了儲存值之外,也要額外儲存 col_idx ,但這邊可以注意的是,因為限制每四個儲存兩個,所以只需要儲存它的位移值就好,例如用 4bit on/off 來代表那四個中有哪兩個存值,或者個別以 2 bit 來表示0~3 ,以這篇[3]的 3.1 章節來看應該使以個別 2 bit 來儲存。
(0-based, 最開始是 0),例如第零行的第二組假設他存的是[0, 3],那我們就可以利用第二組 -> 前面已有八個元素再配合位移值 8 + [0, 3] -> 可以得出存的位置是在 [8, 11]

sparse tensor core from [3]

這邊我們也可以來做個小計算,如果我們拿 ELL 來存的話要耗多少空間,假設矩陣 n x n 且用 double/float 跟 int (4 byte) 來記錄

  • Dense: n * n * 8 (double) | 4n^2 (float)
  • Ell: n * n/2 * 8 + n * n/2 * 4 = 6n^2 (double) | 4n^2 (float)
  • 以剛剛假設的存法: n * n/2 * 8 + n * n/2 * (1/2) = 4.25n^2 (double) | 2.25n^2 (float)

可以看到 ELL 在 double 時, 只節省 1/4 使用空間,但卻損失了一半的資料,而在 float 時,完全是多此一舉了;而在剛剛假設的狀況下,基本上可以省下一半的使用空間。因此在 sparse tensor core 至少可以有兩倍的效能提升 (假設在同等的計算晶片)。

在 ELL 的情況下,就可以知道為甚麼從矩陣角度來看,我會覺得他不是 Sparse 了,因為這樣存並沒有辦法帶來給我們效益。且通常下 Sparse Matrix 處理的資料大約落在 O(n) , 但這裡的矩陣還是 O(n^2)。

(我對於 deep learning 也還在學習中,如果有錯誤或者可以補充的地方,麻煩在底下留言了)
這一篇[3]的做法 (4.1) 是先以正常的方式訓練,然後再將每四個權重中的兩個比較小的丟掉(丟掉的第二輪訓練就會永遠是零),然後再以這樣的權重再去訓練,此流程也有實作在 Automatic SParsity (ASP) for PyTorch 中。
以下擷取了部分結果

classification accuracy from [3]

最左行是用 half precision (半精準) 的原先模型,第二行是將權重的部分砍掉一半,最後是改用 int 8 來存權重,可以看到其實這樣訓練下來他們的準確率是差不多的,但這樣可以加速最後使用模型來推演的效率。在這篇[3]後面也還有 Image Segmentation 或者 Detection 的結果,有興趣的可以在點進去詳細了解。
除了 Nvidia 這篇外,也有其他論文有用演算法來做類似的事情,但不像 Nvidia 這篇是還有利用到硬體來幫忙加速,例如這一篇 Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks[4] 整理了很多篇論文如何利用 sparse 來加速 deep learning,例如會動態地訓練時將某些邊去掉或加回去,這樣子來訓練模型等。

結語:如果注意最後 deep learning 推演的速度的話可以多考慮 sparse tensor core 或者其他篇利用 sparse 的方式。(就我感覺,應該也可以幫助減緩 overfitting?)

參考資料:

  1. NVIDIA Ampere Architecture In-Depth
  2. 稀疏矩陣 (Sparse Matrix) 介紹
  3. Accelerating Sparse Deep Neural Network
  4. Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks