不得不說，Colossal-AI 訓(xùn)練系統(tǒng)這個開源項目的漲星速度是真快。
在“沒十幾塊顯卡玩不起大模型”的當(dāng)下，它硬是只用一張消費級顯卡，成功單挑了180億參數(shù)的大模型。
難怪每逢新版本發(fā)布前后，都會連續(xù)好幾天霸榜GitHub熱門第一。
使用github-star-history制圖
之前我們也介紹過，Colossal-AI的一個重點就是打破了內(nèi)存墻限制， 如訓(xùn)練GPT-2與英偉達自己的Megatron-LM，相比GPU顯存最高能節(jié)省91.2%。
隨著AI模型參數(shù)量的不斷增長，內(nèi)存不夠的問題逐漸凸顯，一句CUDA out of memory讓不少從業(yè)者頭疼。
甚至伯克利AI實驗室學(xué)者Amir Gholami一年前曾發(fā)出預(yù)言[1]，未來內(nèi)存墻將是比算力更大的瓶頸：
內(nèi)存容量 上，GPU單卡顯存容量每兩年才翻倍，需要支撐的模型參數(shù)卻接近指數(shù)級增長。
傳輸帶寬 上，過去20年才增長30倍，更是遠遠比不上算力20年增長9萬倍的速度。
因此，從芯片內(nèi)部到芯片之間，甚至是AI加速器之間的數(shù)據(jù)通信，都阻礙著AI進一步發(fā)展和落地。
為了搞定這個問題，全行業(yè)都在從不同角度想辦法。
1
為了打破內(nèi)存墻，業(yè)界做出哪些努力？
首先，從模型算法本身入手減少內(nèi)存使用量。
比如斯坦福&紐約州立大學(xué)布法羅分校團隊提出的FlashAttention，給注意力算法加上IO感知能力，速度比PyTorch標(biāo)準(zhǔn)Attention快了2-4倍，所需內(nèi)存也僅是其5%-20%。
論文鏈接：arxiv.org/abs/2205.14135
又比如，東京大學(xué)&商湯&悉尼大學(xué)團隊提出將分層ViT與掩碼圖像建模整合在一起的新方法。內(nèi)存使用量比之前方法減少了70%。
論文鏈接：arxiv.org/abs/2205.13515
同類研究其實層出不窮，就先列舉最近發(fā)表的這兩個成果。
這些單獨的方法雖然有效但應(yīng)用面較窄，需要根據(jù)不同算法和任務(wù)做針對性的設(shè)計，不太能泛化。
接下來，被寄予厚望能解決內(nèi)存墻問題的還有存算一體芯片。
這種新型芯片架構(gòu)在存儲單元中嵌入計算能力，以此消除數(shù)據(jù)搬運的時延和功耗，來突破馮諾依曼瓶頸。
存算一體芯片以憶阻器技術(shù)為代表，這種電路元件阻值會隨著通過的電流改變，如果電流停止，電阻會停留在當(dāng)前值，相當(dāng)于“記住”了電流量。
如果把高阻值定義為1，低阻值定義為0，憶阻器就可以同時實現(xiàn)二進制的計算和存儲。
來自doi：10.1038/s41586-021-03748-0
不過存算一體芯片行業(yè)還在起步階段，需要材料學(xué)的進步來推動。一方面，能做到量產(chǎn)的就不多，另一方面也缺少對應(yīng)的編譯器等軟件基礎(chǔ)設(shè)施支持，所以離真正大規(guī)模應(yīng)用還有一段距離。
當(dāng)下，基于現(xiàn)有軟硬件框架做優(yōu)化就成了比較務(wù)實的選項。
如前面提到的Colossal-AI，用多維并行的方式減少多GPU并行時相互之間的通信次數(shù)，又通過向CPU“借內(nèi)存”的方法讓GPU單卡也能訓(xùn)練大模型。
具體來說，是根據(jù)動態(tài)查詢到的內(nèi)存使用情況，不斷動態(tài)轉(zhuǎn)換張量狀態(tài)、調(diào)整張量位置，高效利用GPU+CPU異構(gòu)內(nèi)存。
這樣一來，當(dāng)AI訓(xùn)練出現(xiàn)算力足夠但內(nèi)存不夠的情況時，只需加錢添購DRAM內(nèi)存即可，這聽起來可比買GPU劃算多了。
然而，這里又面臨一個新的問題。
GPU平臺直接“借內(nèi)存”，并不是一種很高效的選擇(不然大伙兒都去堆內(nèi)存條了)——
與CPU相比，GPU平臺的內(nèi)存可擴展性其實沒那么高、也不具備L1-L3高速緩存。數(shù)據(jù)在CPU與GPU之間交換走的PCIe接口效率也要低一些。
對于那些對時延更敏感的AI應(yīng)用場景來說，是否存在一種更合適的解決方案?
2
用CPU方案打破內(nèi)存墻，可行嗎？
要問行不行，還得先看有沒有。
從業(yè)界來看，確實已經(jīng)有不少公司開始基于CPU平臺搭建一些AI項目，其中一些如個性化推薦、基于AI的實時決策系統(tǒng)等，都屬于“對時延非常敏感”的決策型AI。
而決策型AI，正是深受內(nèi)存墻困擾的“受害者”之一——
不是因為模型參數(shù)量大，而是因為模型對數(shù)據(jù)庫的要求高。
與其他訓(xùn)練完直接投入使用的AI不同，決策型AI必須每天從現(xiàn)實環(huán)境中獲取新鮮數(shù)據(jù)，將決策變得更“精準(zhǔn)”，這需要大量的低時延數(shù)據(jù)交互。
因此，背后的數(shù)據(jù)庫也需要具備大規(guī)模并發(fā)讀寫、實時性強、可擴展等特性。
在這種情況下，如何充分利用內(nèi)存來加速數(shù)據(jù)讀寫，反而成為了比提升算力更加困擾AI的問題。
那么，這些企業(yè)究竟是如何在CPU平臺上解決內(nèi)存墻問題的呢?
以曾經(jīng)在全球引領(lǐng)了在線支付服務(wù)潮流，如今依然處于該領(lǐng)域C位的PayPal為例。
PayPal的業(yè)務(wù)如今已經(jīng)涵蓋了在線轉(zhuǎn)賬、計費和支付，并且客戶規(guī)模已經(jīng)達到了200多個市場的超3.25億消費者和商家，所以它也像傳統(tǒng)的銀行服務(wù)一樣，面臨嚴(yán)峻的欺詐挑戰(zhàn)。
PayPal的應(yīng)對策略，就是打造了一個具備實時識別新出現(xiàn)欺詐模式能力的實時決策系統(tǒng)。
不過欺詐者也在不斷改變欺詐模式，或發(fā)掘新的方式來對抗該系統(tǒng)，因此，PayPal需要不斷提升新型欺詐檢測的準(zhǔn)確性，并且需要盡可能地縮短欺詐檢測時間。
在這種類似貓鼠游戲，比誰反應(yīng)更快、誰能更靈活應(yīng)變的對抗中，起到關(guān)鍵作用的就是數(shù)據(jù)的快速處理及讀寫。
為了實時識別新出現(xiàn)的欺詐模式，PayPal需要更快地處理和分析更多數(shù)據(jù)，就需要將盡可能大體量的數(shù)據(jù)與實時處理做更好的對接。
然而，內(nèi)存墻的問題，在此時也悄然出現(xiàn)了。
PayPal發(fā)現(xiàn)，自己要應(yīng)對的是平臺多年來收集的數(shù)百PB數(shù)據(jù)，隨著其反欺詐決策平臺數(shù)據(jù)量的逐年增長，主索引的規(guī)模也在不斷擴張，以至于幾乎要拖垮其數(shù)據(jù)庫，特別是承載這些數(shù)據(jù)的各節(jié)點的內(nèi)存容量一旦耗盡，反欺詐的效率就會大打折扣，實時性也就無從談起。
于是，PayPal開始考慮采用新的內(nèi)存和存儲技術(shù)，來突破內(nèi)存墻，換言之，提升其數(shù)據(jù)庫方案的整體存儲密度。
恰逢其會，PayPal于2015年開始主要采用來自Aerospike的數(shù)據(jù)庫技術(shù)，而后者正是最早支持英特爾® 傲騰™ 持久內(nèi)存的數(shù)據(jù)庫廠商之一。其創(chuàng)新的混合內(nèi)存架構(gòu)(Hybrid Memory Architecture，HMA)經(jīng)過優(yōu)化，可以幫助PayPal將體量越來越大的主索引存入傲騰持久內(nèi)存而非DRAM中，內(nèi)存墻難題就此破局。
最終的試驗結(jié)果，也驗證了傲騰持久內(nèi)存在打破內(nèi)存墻、提升整個數(shù)據(jù)庫容量和性能方面的價值：
在PayPal現(xiàn)有共計2,000臺Aerospike服務(wù)器中，有200臺已導(dǎo)入了這款持久內(nèi)存，結(jié)果每節(jié)點的存儲空間提升到了原來的約4倍，且保持了應(yīng)用的極速反應(yīng)和低時延。
隨內(nèi)存和存儲容量增大而來的，還有成本上的大幅節(jié)省，據(jù)PayPal和Aerospike進行的基準(zhǔn)測試：
由于單個節(jié)點在數(shù)據(jù)存儲和讀寫上的能力得到了強化，所需服務(wù)器的數(shù)量可以因此減少50%，每集群的成本就可因此降低約30%[2]。
而且，傲騰持久內(nèi)存還有一個BUFF，也在PayPal這個反欺詐應(yīng)用場景里發(fā)揮了令人意想不到的作用，這就是數(shù)據(jù)持久性，能帶來超快的數(shù)據(jù)和應(yīng)用恢復(fù)速度。
相比將主索引存入DRAM，在計劃或非計劃的停機后還需要從存儲設(shè)備中掃描數(shù)據(jù)并重建索引不同，將主索引存入傲騰持久內(nèi)存并做持久化后，不論是意外宕機，還是計劃中的停機，其數(shù)據(jù)都不會因為斷電而消失，整個系統(tǒng)就可以用更快的速度恢復(fù)并重新聯(lián)機。
要問這個速度有多快?PayPal給出的答案是原先需要59分鐘來重建索引，現(xiàn)在只需4分鐘。
PayPal還給出了一些更具整體視角，并從業(yè)務(wù)和最終應(yīng)用功效切入的數(shù)據(jù)來說明它的收益：
它以2015年初步估計的50TB欺詐數(shù)據(jù)量和過去的內(nèi)存系統(tǒng)為基準(zhǔn)，發(fā)現(xiàn)基于傲騰持久內(nèi)存的新方案，可幫助它將服務(wù)級別協(xié)議(SLA)遵守率從98.5%提升到99.95%。
漏查的欺詐交易量，則降到原來的約1/30，整體服務(wù)器的占用空間可降至原來的約1/8(從1024減少到120臺服務(wù)器)，而其整體硬件成本可以降到原來的約1/3。
考慮到預(yù)測的年數(shù)據(jù)增長率約為32%，PayPal的反欺詐系統(tǒng)完全可在新方案上實現(xiàn)經(jīng)濟高效的擴展，并讓它繼續(xù)保持99.95%的欺詐計算SLA遵守率、更短的數(shù)據(jù)恢復(fù)時間、更強的數(shù)據(jù)處理、查詢性能和數(shù)據(jù)一致性以及高達99.99%的可用性。
所以，像這種對數(shù)據(jù)庫性能要求更高的推薦、在線評估類AI應(yīng)用，利用CPU平臺，特別是利用有AI加速能力的CPU+傲騰持久內(nèi)存來打破內(nèi)存墻，加速整體性能表現(xiàn)并降低成本確實是可行，而且也是能夠負擔(dān)得起的。
如前文提及的，除了PayPal這樣的全球型客戶外，國內(nèi)也有不少渴望打破內(nèi)存墻的互聯(lián)網(wǎng)企業(yè)、AI創(chuàng)業(yè)企業(yè)在他們類似的應(yīng)用場景中嘗試了傲騰持久內(nèi)存，結(jié)果也是收獲了內(nèi)存子系統(tǒng)容量大幅擴展+數(shù)據(jù)和應(yīng)用恢復(fù)用時顯著縮短+硬件成本或TCO大降的多重功效。
而且，能用上這套方案的還不止是這些場景。
即使在AI for Science上，目前也有一些科研項目正嘗試充分利用這套方案，來解決內(nèi)存墻的問題。
由DeepMind在2021年發(fā)布的AlphaFold2就算是一例。
得益于加速蛋白質(zhì)三維結(jié)構(gòu)探究的定位，以及預(yù)測的高可信度，AlphaFold2正在生命科學(xué)領(lǐng)域掀起顛覆式的變革，而它的成功秘訣，就在于利用深度學(xué)習(xí)方法進行蛋白質(zhì)結(jié)構(gòu)預(yù)測，這使它在效率、成本等方面遠勝傳統(tǒng)實驗方法(包括X-ray衍射、冷凍電鏡、NMR等)。
因此，幾乎所有生物學(xué)界的從業(yè)者都在著手這一技術(shù)的落地、管線搭建以及性能調(diào)優(yōu)。英特爾也是其中一員。它結(jié)合自身架構(gòu)的軟硬件優(yōu)勢，對AlphaFold2算法進行了在CPU平臺上的端到端高通量優(yōu)化，并實現(xiàn)了比專用AI加速芯片還要出色的性能。
取得這一成績，既得益于第三代英特爾® 至強® 可擴展處理器內(nèi)置的高位寬優(yōu)勢(AVX-512等)，也離不開傲騰持久內(nèi)存對“內(nèi)存墻”的突破。
一方面，在模型推理階段，英特爾專家通過對注意力模塊(attention unit)進行大張量切分(tensor slicing)，以及使用英特爾® oneAPI進行算子融合等優(yōu)化方法提升了算法的計算效率和CPU處理器利用率，加快了并行推理速度，并緩解了算法執(zhí)行中各個環(huán)節(jié)面臨的內(nèi)存瓶頸等問題。
另一方面，傲騰持久內(nèi)存的部署，也提供了TB級內(nèi)存容量的“戰(zhàn)略級”支持，能更輕松地解決多實例并行執(zhí)行時內(nèi)存峰值疊加的內(nèi)存瓶頸。
這個瓶頸有多大?據(jù)英特爾技術(shù)專家介紹：在輸入長度為765aa的條件下，64個實例并行執(zhí)行時，內(nèi)存容量的需求就會突破2TB。在這種情形下，對用戶而言，使用傲騰持久內(nèi)存也是他們目前真正可行的方案。
3
下一步：異構(gòu)芯片，統(tǒng)一內(nèi)存
當(dāng)然，從整個行業(yè)的發(fā)展態(tài)勢來看，CPU搭配大容量持久內(nèi)存的方案，也并非就能一勞永逸地解決“內(nèi)存墻”的問題。
它同樣也只是眾多解決方案中的一種。
那么，是否還有其他針對內(nèi)存墻的方案，既不像存算一體芯片那般遙遠，但又比CPU+持久內(nèi)存的用途更全面、更多樣呢?
答案或許就是異構(gòu)芯片+統(tǒng)一內(nèi)存的路子了。
這里的異構(gòu)芯片，指的可不僅僅是CPU和GPU，還包括有FPGA和ASIC等同樣能為AI計算提供加速的芯片類型。隨著芯粒(Chiplet)技術(shù)的發(fā)展，異構(gòu)計算或許能為打破內(nèi)存墻提供新的可能性。
目前，芯粒互聯(lián)互通的開放標(biāo)準(zhǔn)UCIe(Universal Chiplet Interconnect Express)已獲得大量芯片行業(yè)玩家認可，有望成為主流標(biāo)準(zhǔn)。
這個標(biāo)準(zhǔn)的牽頭者英特爾自己就在積極布局XPU戰(zhàn)略，把標(biāo)量(CPU)、矢量(GPU)、矩陣(ASIC)和空間(FPGA)等不同類型和架構(gòu)芯片的多樣化算力組合在一起。
最近能看到的一項成果便是美國阿貢國家實驗室的下一代超算系統(tǒng)——極光(Aurora)。
極光超算的CPU將采用代號為Sapphire Rapids的第四代英特爾® 至強® 可擴展處理器，并搭配代號為Ponte Vecchio的英特爾® 數(shù)據(jù)中心GPU，雙精度峰值計算性能超過每秒兩百億億次，能支持更準(zhǔn)確的氣候預(yù)測以及發(fā)現(xiàn)應(yīng)對癌癥的新療法等研發(fā)創(chuàng)新活動。
這還是目前可見的進展。在UCIe的支持下，未來還有可能出現(xiàn)不同架構(gòu)、甚至不同工藝制程的IP封裝成為一塊SoC芯片的全新物種。
伴隨異構(gòu)芯片的協(xié)作甚至是異構(gòu)芯粒的整合，不同芯片和芯粒所搭配的內(nèi)存也很可能出現(xiàn)統(tǒng)一或池化的趨勢。
其中一個可能的實現(xiàn)途徑，就是通過光學(xué)I/O來連接不同芯片、芯粒、內(nèi)存等組件，即用光信號代替電信號做芯片間的通信，可以做到更高帶寬、更低時延和更低功率。
例如，光學(xué)I/O方面的創(chuàng)新企業(yè)Ayar Labs，目前已經(jīng)被各大芯片巨頭和高性能計算供應(yīng)商所看好。
在最新一輪1.3億美元的融資中，它的投資方就包括了英特爾、英偉達、格芯和HPE。
或許，距離內(nèi)存“大一統(tǒng)”的時代真的不遠了。
在這種情況下，持久內(nèi)存本身也正在迎來更多的機會。
例如，傲騰持久內(nèi)存目前已實現(xiàn)單條512GB的容量，單條1TB容量的型號也正在籌備中。
如果要真正高效地擴展異構(gòu)系統(tǒng)的統(tǒng)一內(nèi)存池，它所具備的多重優(yōu)勢是不可忽略的。
申請創(chuàng)業(yè)報道，分享創(chuàng)業(yè)好點子。點擊此處，共同探討創(chuàng)業(yè)新機遇！
