Sparse Tensor Core 初見

Deep Learning Aug 2, 2021

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

Tags

Great! You've successfully subscribed.
Great! Next, complete checkout for full access.
Welcome back! You've successfully signed in.
Success! Your account is fully activated, you now have access to all content.