打破內存牆的新利器成行業「熱搜」!持久內存讓打工人也能玩轉海量數據+高維模型

  • 2022 年 6 月 23 日
  • AI

不得不說,Colossal-AI訓練系統這個開源項目的漲星速度是真快。

在「沒十幾塊顯卡玩不起大模型」的當下,它硬是只用一張消費級顯卡,成功單挑了180億參數的大模型。

難怪每逢新版本發佈前後,都會連續好幾天霸榜GitHub熱門第一。

使用github-star-history製圖

之前我們也介紹過,Colossal-AI的一個重點就是打破了內存牆限制,如訓練GPT-2與英偉達自己的Megatron-LM,相比GPU顯存最高能節省91.2%。

隨着AI模型參數量的不斷增長,內存不夠的問題逐漸凸顯,一句CUDA out of memory讓不少從業者頭疼。

甚至伯克利AI實驗室學者Amir Gholami一年前曾發出預言[1]來內存牆將是比算力更大的瓶頸:

內存容量上,GPU單卡顯存容量每兩年才翻倍,需要支撐的模型參數卻接近指數級增長。

傳輸帶寬上,過去20年才增長30倍,更是遠遠比不上算力20年增長9萬倍的速度。

因此,從芯片內部到芯片之間,甚至是AI加速器之間的數據通信,都阻礙着AI進一步發展和落地。

為了搞定這個問題,全行業都在從不同角度想辦法。


1

為了打破內存牆,業界做出哪些努力?

首先,從模型算法本身入手減少內存使用量。

比如斯坦福&紐約州立大學布法羅分校團隊提出的FlashAttention,給注意力算法加上IO感知能力,速度比PyTorch標準Attention快了2-4倍,所需內存也僅是其5%-20%。

論文鏈接:arxiv.org/abs/2205.14135

又比如,東京大學&商湯&悉尼大學團隊提出將分層ViT與掩碼圖像建模整合在一起的新方法。內存使用量比之前方法減少了70%。

論文鏈接:arxiv.org/abs/2205.13515

同類研究其實層出不窮,就先列舉最近發表的這兩個成果。

這些單獨的方法雖然有效但應用面較窄,需要根據不同算法和任務做針對性的設計,不太能泛化。

接下來,被寄予厚望能解決內存牆問題的還有存算一體芯片。

這種新型芯片架構在存儲單元中嵌入計算能力,以此消除數據搬運的時延和功耗,來突破馮諾依曼瓶頸。

存算一體芯片以憶阻器技術為代表,這種電路元件阻值會隨着通過的電流改變,如果電流停止,電阻會停留在當前值,相當於「記住」了電流量。

如果把高阻值定義為1,低阻值定義為0,憶阻器就可以同時實現二進制的計算和存儲。

來自doi:10.1038/s41586-021-03748-0

不過存算一體芯片行業還在起步階段,需要材料學的進步來推動。一方面,能做到量產的就不多,另一方面也缺少對應的編譯器等軟件基礎設施支持,所以離真正大規模應用還有一段距離。

當下,基於現有軟硬件框架做優化就成了比較務實的選項。

如前面提到的Colossal-AI,用多維並行的方式減少多GPU並行時相互之間的通信次數,又通過向CPU「借內存」的方法讓GPU單卡也能訓練大模型。

具體來說,是根據動態查詢到的內存使用情況,不斷動態轉換張量狀態、調整張量位置,高效利用GPU+CPU異構內存。

這樣一來,當AI訓練出現算力足夠但內存不夠的情況時,只需加錢添購DRAM內存即可,這聽起來可比買GPU划算多了。

然而,這裡又面臨一個新的問題。

GPU平台直接「借內存」,並不是一種很高效的選擇(不然大伙兒都去堆內存條了)——

與CPU相比,GPU平台的內存可擴展性其實沒那麼高、也不具備L1-L3高速緩存。數據在CPU與GPU之間交換走的PCIe接口效率也要低一些。

對於那些對時延更敏感的AI應用場景來說,是否存在一種更合適的解決方案?



2

用CPU方案打破內存牆,可行嗎?

要問行不行,還得先看有沒有。

從業界來看,確實已經有不少公司開始基於CPU平台搭建一些AI項目,其中一些如個性化推薦、基於AI的實時決策系統等,都屬於「對時延非常敏感」的決策型AI。

而決策型AI,正是深受內存牆困擾的「受害者」之一——

不是因為模型參數量大,而是因為模型對數據庫的要求高。

與其他訓練完直接投入使用的AI不同,決策型AI必須每天從現實環境中獲取新鮮數據,將決策變得更「精準」,這需要大量的低時延數據交互。

因此,背後的數據庫也需要具備大規模並發讀寫、實時性強、可擴展等特性。

在這種情況下,如何充分利用內存來加速數據讀寫,反而成為了比提升算力更加困擾AI的問題。

那麼,這些企業究竟是如何在CPU平台上解決內存牆問題的呢?

以曾經在全球引領了在線支付服務潮流,如今依然處於該領域C位的PayPal為例。

PayPal的業務如今已經涵蓋了在線轉賬、計費和支付,並且客戶規模已經達到了200多個市場的超3.25億消費者和商家,所以它也像傳統的銀行服務一樣,面臨嚴峻的欺詐挑戰。

PayPal的應對策略,就是打造了一個具備實時識別新出現欺詐模式能力的實時決策系統。

不過欺詐者也在不斷改變欺詐模式,或發掘新的方式來對抗該系統,因此,PayPal需要不斷提升新型欺詐檢測的準確性,並且需要儘可能地縮短欺詐檢測時間。

在這種類似貓鼠遊戲,比誰反應更快、誰能更靈活應變的對抗中,起到關鍵作用的就是數據的快速處理及讀寫。

為了實時識別新出現的欺詐模式,PayPal需要更快地處理和分析更多數據,就需要將儘可能大體量的數據與實時處理做更好的對接。

然而,內存牆的問題,在此時也悄然出現了。

PayPal發現,自己要應對的是平台多年來收集的數百PB數據,隨着其反欺詐決策平台數據量的逐年增長,主索引的規模也在不斷擴張,以至於幾乎要拖垮其數據庫,特別是承載這些數據的各節點的內存容量一旦耗盡,反欺詐的效率就會大打折扣,實時性也就無從談起。

於是,PayPal開始考慮採用新的內存和存儲技術,來突破內存牆,換言之,提升其數據庫方案的整體存儲密度。

恰逢其會,PayPal於2015年開始主要採用來自Aerospike的數據庫技術,而後者正是最早支持英特爾® 傲騰™ 持久內存的數據庫廠商之一。其創新的混合內存架構(Hybrid Memory Architecture,HMA)經過優化,可以幫助PayPal將體量越來越大的主索引存入傲騰持久內存而非DRAM中,內存牆難題就此破局。

最終的試驗結果,也驗證了傲騰持久內存在打破內存牆、提升整個數據庫容量和性能方面的價值:

在PayPal現有共計2,000台Aerospike服務器中,有200台已導入了這款持久內存,結果每節點的存儲空間提升到了原來的約4倍,且保持了應用的極速反應和低時延。

隨內存和存儲容量增大而來的,還有成本上的大幅節省,據PayPal和Aerospike進行的基準測試:

由於單個節點在數據存儲和讀寫上的能力得到了強化,所需服務器的數量可以因此減少50%,每集群的成本就可因此降低約30%[2]

而且,傲騰持久內存還有一個BUFF,也在PayPal這個反欺詐應用場景里發揮了令人意想不到的作用,這就是數據持久性,能帶來超快的數據和應用恢復速度。

相比將主索引存入DRAM,在計劃或非計劃的停機後還需要從存儲設備中掃描數據並重建索引不同,將主索引存入傲騰持久內存並做持久化後,不論是意外宕機,還是計劃中的停機,其數據都不會因為斷電而消失,整個系統就可以用更快的速度恢復並重新聯機。

要問這個速度有多快?PayPal給出的答案是原先需要59分鐘來重建索引,現在只需4分鐘。

PayPal還給出了一些更具整體視角,並從業務和最終應用功效切入的數據來說明它的收益:

它以2015年初步估計的50TB欺詐數據量和過去的內存系統為基準,發現基於傲騰持久內存的新方案,可幫助它將服務級別協議(SLA)遵守率從98.5%提升到99.95%。

漏查的欺詐交易量,則降到原來的約1/30,整體服務器的佔用空間可降至原來的約1/8(從1024減少到120台服務器),而其整體硬件成本可以降到原來的約1/3。

考慮到預測的年數據增長率約為32%,PayPal的反欺詐系統完全可在新方案上實現經濟高效的擴展,並讓它繼續保持99.95%的欺詐計算SLA遵守率、更短的數據恢復時間、更強的數據處理、查詢性能和數據一致性以及高達99.99%的可用性。

所以,像這種對數據庫性能要求更高的推薦、在線評估類AI應用,利用CPU平台,特別是利用有AI加速能力的CPU+傲騰持久內存來打破內存牆,加速整體性能表現並降低成本確實是可行,而且也是能夠負擔得起的。

如前文提及的,除了PayPal這樣的全球型客戶外,國內也有不少渴望打破內存牆的互聯網企業、AI創業企業在他們類似的應用場景中嘗試了傲騰持久內存,結果也是收穫了內存子系統容量大幅擴展+數據和應用恢復用時顯著縮短+硬件成本或TCO大降的多重功效。

而且,能用上這套方案的還不止是這些場景。

即使在AI for Science上,目前也有一些科研項目正嘗試充分利用這套方案,來解決內存牆的問題。

由DeepMind在2021年發佈的AlphaFold2就算是一例。

得益於加速蛋白質三維結構探究的定位,以及預測的高可信度,AlphaFold2正在生命科學領域掀起顛覆式的變革,而它的成功秘訣,就在於利用深度學習方法進行蛋白質結構預測,這使它在效率、成本等方面遠勝傳統實驗方法(包括X-ray衍射、冷凍電鏡、NMR等)

因此,幾乎所有生物學界的從業者都在着手這一技術的落地、管線搭建以及性能調優。英特爾也是其中一員。它結合自身架構的軟硬件優勢,對AlphaFold2算法進行了在CPU平台上的端到端高通量優化,並實現了比專用AI加速芯片還要出色的性能。

取得這一成績,既得益於第三代英特爾® 至強® 可擴展處理器內置的高位寬優勢(AVX-512等),也離不開傲騰持久內存對「內存牆」的突破。

一方面,在模型推理階段,英特爾專家通過對注意力模塊(attention unit)進行大張量切分(tensor slicing),以及使用英特爾® oneAPI進行算子融合等優化方法提升了算法的計算效率和CPU處理器利用率,加快了並行推理速度,並緩解了算法執行中各個環節面臨的內存瓶頸等問題。

另一方面,傲騰持久內存的部署,也提供了TB級內存容量的「戰略級」支持,能更輕鬆地解決多實例並行執行時內存峰值疊加的內存瓶頸。

這個瓶頸有多大?據英特爾技術專家介紹:在輸入長度為765aa的條件下,64個實例並行執行時,內存容量的需求就會突破2TB。在這種情形下,對用戶而言,使用傲騰持久內存也是他們目前真正可行的方案。



3

下一步:異構芯片,統一內存

當然,從整個行業的發展態勢來看,CPU搭配大容量持久內存的方案,也並非就能一勞永逸地解決「內存牆」的問題。

它同樣也只是眾多解決方案中的一種。

那麼,是否還有其他針對內存牆的方案,既不像存算一體芯片那般遙遠,但又比CPU+持久內存的用途更全面、更多樣呢?

答案或許就是異構芯片+統一內存的路子了。

這裡的異構芯片,指的可不僅僅是CPU和GPU,還包括有FPGA和ASIC等同樣能為AI計算提供加速的芯片類型。隨着芯粒(Chiplet)技術的發展,異構計算或許能為打破內存牆提供新的可能性。

目前,芯粒互聯互通的開放標準UCIe(Universal Chiplet Interconnect Express)已獲得大量芯片行業玩家認可,有望成為主流標準。

這個標準的牽頭者英特爾自己就在積極布局XPU戰略,把標量(CPU)、矢量(GPU)、矩陣(ASIC)和空間(FPGA)等不同類型和架構芯片的多樣化算力組合在一起。

最近能看到的一項成果便是美國阿貢國家實驗室的下一代超算系統——極光(Aurora)

極光超算的CPU將採用代號為Sapphire Rapids的第四代英特爾® 至強® 可擴展處理器,並搭配代號為Ponte Vecchio的英特爾® 數據中心GPU,雙精度峰值計算性能超過每秒兩百億億次,能支持更準確的氣候預測以及發現應對癌症的新療法等研發創新活動。

這還是目前可見的進展。在UCIe的支持下,未來還有可能出現不同架構、甚至不同工藝製程的IP封裝成為一塊SoC芯片的全新物種。

伴隨異構芯片的協作甚至是異構芯粒的整合,不同芯片和芯粒所搭配的內存也很可能出現統一或池化的趨勢。

其中一個可能的實現途徑,就是通過光學I/O來連接不同芯片、芯粒、內存等組件,即用光信號代替電信號做芯片間的通信,可以做到更高帶寬、更低時延和更低功率。

例如,光學I/O方面的創新企業Ayar Labs,目前已經被各大芯片巨頭和高性能計算供應商所看好。

在最新一輪1.3億美元的融資中,它的投資方就包括了英特爾、英偉達、格芯和HPE。

或許,距離內存「大一統」的時代真的不遠了。

在這種情況下,持久內存本身也正在迎來更多的機會。

例如,傲騰持久內存目前已實現單條512GB的容量,單條1TB容量的型號也正在籌備中。

如果要真正高效地擴展異構系統的統一內存池,它所具備的多重優勢是不可忽略的。

參考鏈接:

[1]//mp.weixin.qq.com/s/kuIsyX0QEIeFHn8tvFE8vw

[2]//www.intel.com/content/www/us/en/customer-spotlight/stories/paypal-customer-story.html

雷峰網