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