AI加速器與機器學習算法:協同設計與進化
作者 | Shashank Prasanna
翻譯 | 胡燕君
此刻,你應該是在電腦或手機上看這篇文章。不管怎樣,這些機器都屬于現代計算機,它們都有中央處理器(CPU)和其他為特定功能服務的專用芯片,例如顯卡、聲卡、網卡、傳感器融合等。處理特定任務時,專用處理器往往比通用CPU更快更高效。計算機發展早期,CPU都會和專用處理器配合使用。1970年代的8位和16位CPU需要依賴軟件來模擬浮點指令,因此執行浮點運算非常慢。而由于計算機輔助設計(CAD)和工程模擬等應用對浮點運算的速度要求較高,于是人們開始用數學協處理器(math coprocessor)輔助CPU,讓數學協處理器分擔所有浮點運算任務,它的浮點運算速度和效率都比CPU更高。這就是專用處理器的一個例子。
關注AI和半導體行業的朋友近來可能聽說過“機器學習(ML)專用處理器”(即AI加速器)。最常見的AI加速器莫過于NVIDIA GPU,此外還有Intel的Habana Gaudi處理器、Graphcore的Bow IPU、Google的TPU、AWS的Trainium和Inferentia芯片等。為什么如今有這么多AI加速器可供選擇?它們和CPU有什么不同?算法如何改變才能適應這些硬件?硬件又該如何發展才能支持最新的算法?本文將一一解答。本文主要內容包括:
- 為什么需要專用AI加速器?
- ML硬件的分類:CPU、GPU、AI加速器、FPGA和ASIC
- “硬件感知(Hardware-aware)”的算法和“算法感知(Algorithms-aware)”的硬件
- AI加速器與高效ML算法的協同進化
- 針對推理的AI加速器與高效算法
- 針對訓練的AI加速器與高效算法
- AI加速器的未來
構建ML專用處理器有三個方面的原因:能效、性能、模型大小及復雜度。近來,要提高模型準確率,通常做法是擴大模型參數量,并用更大型的數據集訓練模型。計算機視覺、自然語言處理和推薦系統都采用這種做法。語言模型方面,前幾年誕生的GPT-3有1750億參數,當時被視為大模型的“天花板”,但后來又出現了GLaM和NVIDIA MT-NLG,參數量分別達到1.2萬億和5300億。按照歷史規律,模型將越來越大,而現有處理器的算力將無法滿足大模型在訓練時間和推理延遲方面的要求。不過,構建AI專用加速器的最重要原因還是能效,開發AI專用芯片可節省巨大的能源,可覆蓋研發投入有余。為什么需要高能效的處理器?ML模型越大,需要執行的內存訪問操作就越多。與內存訪問相比,矩陣-矩陣運算和矩陣-向量運算的能效高很多。根據斯坦福大學韓松博士的論文(https://arxiv.org/pdf/1506.02626v3.pdf),讀取內存的能耗比加/乘運算操作的能耗高出好幾個數量級。大型神經網絡由于無法片上存儲,需要執行更多DRAM讀取操作,因此能耗還要更高。圖源:https://arxiv.org/pdf/1506.02626v3.pdf
用通用處理器運行深度神經網絡時,擴大處理器規模可以帶來性能提升,但即便只是小幅提升也須以大量能耗和巨額設備成本為代價。
CPU等通用處理器犧牲低能耗換取通用性,AI加速器等專用處理器則犧牲通用性換取低能耗。
使用AI加速器則不一樣。AI加速器通過改進設計,可以減少內存訪問,提供更大的片上緩存,還可以具備特定的硬件功能(如加速矩陣-矩陣計算)。由于AI加速器是基于特定而構建的設備,可根據算法進行適配改進,因此其運行效率會比通用處理器更高。
接下來我們談談加速器的不同種類,以及它們的通用和專用程度。 如上圖所示,最具通用性的是CPU,可以運行任意代碼。專用芯片可執行的任務,CPU也能執行,如圖像處理、語音處理、機器學習等。然而,CPU的性能和能效都比較低。
通過以上分類,我們對各種處理器有了大致認識。下面我們來談談這些處理器如何與軟件互相配合。通用計算模型有兩個組成部分:(1)軟件與算法;(2)運行軟件的硬件處理器。一般而言,這兩部分互相獨立——編寫軟件時很少會考慮軟件會在什么硬件上運行;而硬件設計的出發點則是讓硬件盡可能支持更多種類的軟件。
通用計算模型逐漸演進,需要應用于深度學習等要求高性能和高能效的領域,因此需要加強算法設計和硬件設計之間的聯系。現代ML計算模型同樣包含兩個組成部分:(1)ML算法與軟件框架;(2)與AI加速器配合使用的通用處理器。不同于通用計算模型的開發,ML研發人員會針對特定硬件設計ML算法。比如,編寫代碼時可以充分利用硬件的特定功能(如支持多種精度:INT8、FP16、BF16、FP32)實現特定的芯片功能(混合精度、結構化稀疏)。然后,用戶就可以通過常見的ML軟件框架使用這些功能。同理,硬件設計師會針對特定算法構建AI加速器。比如,為加速ML矩陣計算而設計專門的芯片功能(如NVIDIA在Volta GPU架構中引入Tensor Core)。這就是AI加速器和ML算法的協同進化。硬件設計師為AI加速器增加ML算法可以使用的功能,而ML研發人員則利用AI加速器的硬件功能量身設計新的算法。
硬件和軟件的協同可以帶來更好的性能和更高的能效。
AI加速器分為兩類:(1)用于訓練的AI加速器;(2)用于推理的AI加速器。由于訓練和推理的目標不同,而AI加速器是針對特定工作負載的專用處理器,因此有必要為不同類型的工作負載分別設計處理器。用于訓練的AI加速器的目標是減少訓練時間,而且應具備能配合訓練算法的硬件特點。因此,AI訓練加速器的功率通常較大,內存空間也較大,以滿足較高的吞吐量(每秒處理的數據)要求。由于AI訓練加速器注重吞吐量,因此提高吞吐量和利用率有助于降低能耗成本(即通過擴大吞吐量降低“能耗/吞吐量”之間的比率)。AI訓練加速器還支持混合精度訓練,使用較低精度以加快計算速度,使用高精度累積計算結果,從而實現比通用處理器更高的能效。(后文還將詳談AI加速器的混合精度訓練。) AI推理加速器的目標是,在運行大量獨立數據批次時降低預測延遲,因此需要具備高能效特點,需要降低“能耗/預測”之間的比率。雖然也可將訓練加速器用于推理加速(畢竟訓練中的前向傳播過程本質上即是一種推理任務),但使用訓練加速器時,“能耗/推理”之比會大很多,因為訓練加速器處理小型數據批次的推理請求時利用率較低。訓練加速器就像公交車,只有保持乘客滿員才能實現高能效(油耗/乘客數之比低)。如果偌大的公交車每次只載一個人,其油耗/乘客數之比將變得極高。而推理加速器就像跑車,其速度比公交車快,只載一人時能效比公交車高(跑車的單名乘客油耗比公交車低)。但如果想用跑車一次載50個人,它就會跑得極慢(何況超載違法)。下文將分別談論訓練和推理的工作流,以及AI加速器和軟件應具備什么特點才能在訓練和推理中實現高性能和高能效。
ML推理即是根據新的數據使用訓練好的模型以輸出預測結果。本節將討論AI加速器上運行的可提升推理性能和效率的算法。
訓練后量化(PTQ):使用一個訓練好的、以FP32格式計算的模型,確定比例因子,然后將FP32映射為INT8。確定比例因子的方法是:TensorRT衡量每個神經網絡層中激活函數輸出的分布,然后找到一個使參考分布(reference distribution)和量化分布(quantized distribution)之間信息損失(KL散度)最小的比例因子。
量化感知訓練(QAT):在訓練中計算比例因子,使模型可以適應信息損失并將信息損失降到最低。
可見,硬件不斷發展,具備更多可提升效率的功能(如降低精度)。同時,算法也不斷進化,可以更好地利用硬件的功能。我的另一篇文章提供了NVIDIA TensorRT在GPU上實行量化的代碼示例(https://towardsdatascience.com/a-complete-guide-to-ai-accelerators-for-deep-learning-inference-gpus-aws-inferentia-and-amazon-7a5d6804ef1c)。使用AWS Inferentia芯片進行量化NVIDIA GPU設計之初用于圖像處理加速,后來才演變成強大的AI加速器,而AWS Inferentia芯片一開始即是為機器學習推理而生。每塊AWS Inferentia芯片含4個NeuronCore。NeuronCore是基于脈動陣列的矩陣相乘引擎,有兩級存儲層次結構和極大的片上緩存空間。AWS Inferentia芯片支持FP16、BF16和INT8數據類型,不支持更高精度的格式——畢竟AWS Inferentia是一種推理專用處理器,推理時無須用到更高的精度。正如NVIDIA為GPU推出了TensorRT編譯器,AWS也推出了AWS Neuron SDK和AWS Neuron編譯器,該編譯器支持量化和優化,可提高推理效率。盡管AWS Inferentia芯片支持INT8格式,但截至本文撰寫時,AWS Neuron編譯器只支持量化到FP16和BF16格式。用FP32格式訓練的模型會在編譯過程中自動被轉化為BF16格式。如果在使用AWS Neuron編譯器之前人工將FP32格式的權重量化為FP16,那么編譯器就會保留FP16精度用于推理。與GPU相比,AWS Inferentia芯片不可編程,專用性比GPU更強,更接近ASIC。如果模型中包含的操作均為AWS Inferentia所支持,那么對于特定的模型和批次規模(batch size)而言,使用Inferentia就比使用GPU更能提高模型的能效。然而,如果模型含有Inferentia不支持的操作,AWS Neuron編譯器會自動將相應操作置于主機CPU上,這就導致CPU和加速器之間需要進行數據搬運,進而降低性能和效率。
ML訓練即利用訓練數據優化模型參數,以提高模型的預測準確度。本節將討論AI加速器上運行的算法如何提升推理性能和能效。 接下來我們依舊會討論精度,不過這次是從訓練工作流的角度。如前所述,訓練時,模型權重和激活函數都以FP32格式存儲,FP32遵循早在深度學習之前就誕生的IEEE 754浮點數標準。FP32之所以被選為機器學習默認的浮點數表示形式,是因為訓練時FP16可表示的信息量不夠大,而FP64可表示的信息量則太大,而且也不必用到這么高的精度。機器學習需要一種精度處在FP16 和FP64之間的表示格式,但當時的硬件并不支持。換言之,當時的硬件并不能滿足ML算法的需求,并未成為“算法感知”的硬件。如果當時的ML研究人員有更好的選擇,他們應該會選一種不同于FP32的格式,或者使用混合精度來提升性能和效率。混合精度也正是目前AI加速器的發展方向。實現混合精度,需要硬件和算法的協同設計。
混合精度訓練提升性能與效率
矩陣乘法運算是神經網絡訓練和推理的基本操作。AI加速器的主要工作即為在神經網絡的不同層中將輸入數據和權重的大型矩陣相乘。混合精度訓練背后的思想是,訓練時的矩陣乘法發生在較低精度表示(FP16、BF16、TF32),因此它們更快和能效更高,然后用FP32格式累積運算結果,以降低信息損失,從而提升訓練速度和能效。使用NVIDIA GPU進行混合精度訓練2017年,NVIDIA宣布推出Volta GPU架構,其中包含專門用于機器學習的Tensor Core運算單元。Tensor Core通過FP16運算和FP32累積結果實現混合精度訓練(https://arxiv.org/abs/1710.03740)。NVIDIA的新一代新架構還支持更多低精度格式(BF16、TF32)。在芯片層次,Tensor Core執行低精度(reduced-precision)融合乘加(FMA)運算,用FP32累積結果。每一代NVIDIA架構的進步都體現了硬件和算法之間的協同設計和協同發展。
- NVIDIA Volta架構(2017)引入第一代Tensor Core,當時僅支持FP16運算和FP32累積結果。
- NVIDIA Turing架構(2018)的Tensor Core支持更低精度的INT8和INT4(主要可以加速推理,而非加速訓練)。
- NVIDIA Ampere架構(2020)的Tensor Core還支持BF16和TF32,也就是說,它可以執行FP16、BF16和TF32運算,并且用FP32累積結果,以實現混合精度
混合精度訓練的一大難點是軟件層面的實現。用戶必須在訓練時執行額外的操作,比如將權重轉化為FP16格式,但同時會保留權重的FP32副本和損失縮放(loss scaling)。盡管NVIDIA可以讓深度學習框架在只需修改少量代碼的情況下執行這些操作,它對用戶的要求依然很高,不像使用FP32訓練那么簡單。
NVIDIA的Ampere架構支持TF32,可以有效解決這一用戶體驗難題。TF32格式的好處在于,它結合了FP32的動態范圍和FP16的精度,因此深度學習框架無需轉換格式和保留副本等額外操作即可直接支持TF32格式。然而,在為開發者減少麻煩的情況下,使用TF32可實現比FP32更好的性能,但NVIDIA依然推薦使用FP16或BF16格式進行混合精度訓練,以便獲得最快的訓練性能。使用其他AI加速器進行混合精度訓練Intel Habana Gaudi處理器Habana Gaudi加速器支持混合精度訓練的方式與NVIDIA GPU類似——通過一個附加工具配合深度學習框架,使用格式轉換和副本保存功能。若想體驗Intel Habana Gaudi AI加速器的功能,可以通過AWS云服務器啟動Amazon EC2 DL1實例,該實例配備8個Gaudi加速器。AWS Tranium芯片AWS在2021年re:Invent大會上宣布推出Tranium芯片,該芯片由AWS的Annapurna實驗室研發,用于AI加速。目前,Tranium 芯片尚未得到大規模應用。AWS在大會上介紹稱,Tranium將支持FP16、TF32、BF16、INT8,以及一種稱為cFP8(定制8位浮點數)的全新格式。
如今,ML算法研究和硬件設計都在蓬勃發展。AI加速器也將在性能和能效方面持續進步,逐漸可以像通用處理器一樣無縫使用。現代的AI加速器已具備理想中的硬件功能,例如支持INT1和INT4,這兩種精度類型尚未被用于訓練和推理,但或許它們的存在可以催生新的ML算法。AI加速器之間的互聯也漸見革新。隨著模型規模越來越大,我們需要更大的計算集群,將更多AI加速器連接起來,從而支持更大的工作負載。為此,NVIDIA推出了高帶寬的NVLink和NVSwitch,用于GPU之間的互聯;Intel的Habana Gaudi處理器則在片上集成了基于以太網的RoCE RDMA。未來AI應用將更加廣泛, AI加速器也將成為現代計算環境的中流砥柱。希望未來的AI加速器帶來更好的用戶和開發者體驗。如今的異構計算模型需要協調多個CPU和AI加速器,對大部分數據科學家和開發人員而言,掌握它們的聯網和存儲設置難度頗高。使用Amazon SageMaker等云托管服務可省去管理基礎設施的麻煩,可以方便地擴大機器學習規模,然而,開源框架仍希望用戶對底層硬件、精度類型、編譯器選擇和聯網原語等有較深的了解。未來,開發人員可以登入遠程IDE,然后使用開源ML框架運行代碼,而不必考慮代碼在何種設備上以何種方式運行。他們唯一需要思考的只是成本和速度之間的權衡——想獲得高速度就多花錢,想省錢就在速度上妥協。我是個樂觀的人,我認為距離這樣的未來已經不遠了。(本文經授權后編譯發布。原文:https://towardsdatascience.com/ai-accelerators-machine-learning-algorithms-and-their-co-design-and-evolution-2676efd47179)
頭圖源自brookhaven, Pixabay
*博客內容為網友個人發布,僅代表博主個人觀點,如有侵權請聯系工作人員刪除。
數字通信相關文章:數字通信原理
通信相關文章:通信原理







