[[390958]]

本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

AI訓練的計算量每年都在大幅增長，最近有研究指出，AI訓練未來的瓶頸不是算力，而是GPU內存。

AI加速器通常會簡化或刪除其他部分，以提高硬件的峰值計算能力，但是卻難以解決在內存和通信上的難題。

無論是芯片內部、芯片間，還是AI加速器之間的通信，都已成為AI訓練的瓶頸。

Transformer模型中的參數數量（紅色）呈現出2年240倍的超指數增長，而單個GPU內存（綠色）僅以每2年2倍的速度擴大。

△多年來SOTA模型的參數數量（紅點）以及AI加速器存儲容量（綠點）的演變

訓練AI模型的內存需求，通常是參數數量的幾倍。因為訓練需要存儲中間激活，通常會比參數（不含嵌入）數量增加3-4倍的內存。

于是，AI訓練不可避免地撞上了“內存墻”（Memory Wall），內存墻不僅是內存容量，也包括內存傳輸帶寬。

在很多情況下，數據傳輸的容量和速度，都沒有觸摸到內存墻。

△訓練不同神經網絡模型所需的內存量

從圖中可以看出，每當GPU內存容量增加時，開發人員就會設計出新模型；

2019年GPT-2所需的內存容量，已經是2012年的AlexNet的7倍以上；

自谷歌團隊在2017年提出Transformer，模型所需的內存容量開始大幅增長。

為什么不能靠多GPU堆顯存

那么，為了擺脫單個硬件的有限內存容量和帶寬，是否可以將訓練擴展到多個AI加速器，使用分布式內存呢？

事實上，這樣做也會面臨內存墻的問題，神經網絡加速器之間移動數據的通信瓶頸，甚至比芯片上的數據移動還慢且低效。

與單系統內存的情況類似，擴展帶寬的技術難題還尚未被攻克。僅在很少的通信和數據傳輸的情況下，橫向擴展才適用于計算密集型問題。

從圖中可以看出，20年間，運算設備的算力提高了90,000倍；

雖然存儲器從DDR發展到GDDR6x，能夠用于顯卡、游戲終端和高性能運算，接口標準也從PCIe1.0a升級到NVLink3.0；

但是和算力的提高幅度相比，通訊帶寬的增長只有30倍，可以說非常緩慢。

由于算力和內存之間的差距越來越大，訓練更大的模型也會更加困難。

怎樣突破“內存墻”

怎樣解決內存限制問題？作者從三個方面進行了分析。

訓練算法的改進

訓練神經網絡模型的一大挑戰，就是要進行蠻力超參數調整。雖然可以通過二階隨機優化方法來實現，不過目前的方法卻增加了3-4倍的內存占用，這一點仍需解決。

微軟的Zero方法（一種萬億級模型參數訓練方法），實現了在相同內存下，通過去除多余的優化狀態變量，來訓練8倍大的模型。

也可以在傳遞過程中只存儲或檢查激活的子集，而不保存所有激活，以此將內存占用減少5倍，不過需要增加20%的計算量。

此外，從單精度算法到半精度（FP16）算法的進展，使硬件計算能力提高了10倍以上，可以進一步研究適合INT8精讀訓練的優化算法。

高效部署

最新的SOTA模型（例如：GPT-3）需要分布式內存部署，這是一個很大的挑戰。可以通過降低精度或刪除其冗余參數，來壓縮這些模型，以進行推理。

在訓練或推理過程中，可以降低至INT4精度，模型占用空間和延遲能夠減少8倍。不過，想要將訓練精度降低到遠低于FP16，仍然很困難。

而刪除冗余參數，則可能導致準確率下降。當前的方法能夠修剪30％的具有結構化稀疏性的神經元，以及80％的具有非結構化稀疏性的神經元，以保證對準確性的影響最小。

AI加速器的設計

雖然很難同時提高存儲帶寬和峰值計算能力，但是可以犧牲峰值計算，以獲得更好的帶寬。

在帶寬受限問題上，CPU的性能要比GPU好得多，但是與相比GPU相比，CPU的峰值計算能力要小一個數量級左右。

因此，可以研究一種在二者之間的另一種架構，實現更高效的緩存。