AI真·煉丹：整整14天，無需人類參與

十三 2024-06-30 14:16:32 來源：量子位

金磊 夢晨 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

現(xiàn)在制藥這事，人類要靠邊站了。

坐標(biāo)蘇州，這是一個1600平的制藥實驗室，它的“打開方式”是這樣的：

門口，沒有人。

走廊，沒有人。

實驗室，也沒有人。

相比以往充斥著科學(xué)家、研究員的實驗室，它更多的是把機(jī)械臂和AI系統(tǒng)塞了進(jìn)去，主打的就是一個全自動化。

或許好奇的小伙伴就要問，這樣的實驗室能干嘛？就是為了自動化而自動化嘛？

事情當(dāng)然沒有那么簡單，你瞧見的只是無人的操作，但在背后，AI做的可遠(yuǎn)遠(yuǎn)不只是替代人工的實驗室操作那么簡單，而是：

14天內(nèi)完成靶點發(fā)現(xiàn)和驗證，還是全自動化干濕實驗閉環(huán)的那種。

要知道，這個過程要放以前，可是需要足足2-3年才能完成……

而且更為精細(xì)化的工作，例如樣本處理、細(xì)胞培養(yǎng)、化合物管理、高通量篩選、新一代測序、高內(nèi)涵成像等等，不論是單一任務(wù)還是“聯(lián)動”任務(wù)，機(jī)器都可以在AI的控制下輕松接手。

△用Echo 650T制備檢測板

△用NovaSeq 6000測序

這便是來自全球AI制藥第一梯隊的“選手”——英矽智能（Insilico Medicine）的第六代智能機(jī)器人實驗室，也是全球首個用AI參與決策的生物學(xué)實驗室。

而在它背后驅(qū)動這一切的AI大腦，則是一個叫做PandaOmics的平臺，可以根據(jù)實驗的進(jìn)程自主做決策、下達(dá)指令。

若是把這個AI平臺單拎出來，它更是囊括了20多種預(yù)測模型和生成生物學(xué)模型，還包含遺傳學(xué)、蛋白質(zhì)組學(xué)、甲基化數(shù)據(jù)、文本文獻(xiàn)和科研基金等海量數(shù)據(jù)，用以支持專業(yè)的靶點識別、分析和排序、適應(yīng)癥探索等生物學(xué)研究。

甚至已經(jīng)有高中生用PandaOmics發(fā)現(xiàn)了藥物新靶點，并且研究成果還登上了國際學(xué)術(shù)期刊！

而且除了PandaOmics之外，英矽智能在人工智能制藥領(lǐng)域擁有端到端的藥物發(fā)現(xiàn)平臺Pharma.AI，其中專注于化學(xué)領(lǐng)域的Chemistry42，還可以針對給定靶點從頭設(shè)計具有特定屬性藥物理化性質(zhì)的新型小分子。

這一切都可以在幾小時到幾十小時內(nèi)完成，且支持并行運行多個任務(wù)。

以及英矽智能還將科技圈最潮的大模型也融入進(jìn)來，在Pharma.AI的架構(gòu)上推出Copilot系統(tǒng)，讓你只要會對話就能使用專業(yè)的AI制藥平臺。

由此可見，現(xiàn)在AI不僅是把制藥這件事變成了“自動駕駛”模式，更是狠狠地把門檻打下去、效率提上來。

AI制藥的流程和工作是方便了，但隨之而來的一個問題便是：如此大的工作量，算力，又是如何解決的呢？

科學(xué)計算與AI，CPU都在發(fā)力

對于上述的問題，包括英矽智能在內(nèi)的AI制藥頭部力量們不約而同的選擇了相似的解決辦法：

充分利用所有可以用、值得用的科學(xué)計算與AI算力平臺。這種平臺可不是你想象的那樣被GPU制霸，相反，其中的CPU用量更大，尤其是英特爾的CPU。

為什么要選擇英特爾？

首要的一個原因，就是英特爾供企業(yè)計算及科學(xué)計算使用的主力CPU，即至強(qiáng)? 可擴(kuò)展處理器系列產(chǎn)品，一直都是物理計算——無論是昔日計算機(jī)輔助制藥，還是今天AI輔助制藥都非常依賴的科學(xué)計算應(yīng)用的關(guān)鍵承載平臺。

另一方面，就算是把應(yīng)用的主題從相對傳統(tǒng)的制藥相關(guān)的科學(xué)計算任務(wù)，切換到更偏AI的應(yīng)用上，英特爾也算是頗有建樹，這一點從它對以AlphaFold2為代表的開源蛋白質(zhì)預(yù)測模型的支持上就可見一斑。

△AlphaFold2基本架構(gòu)

?先，AlphaFold2整個端到端的處理過程，涉及?量復(fù)雜多樣的計算類型。從早期的數(shù)據(jù)收集、特征提取等預(yù)處理階段，到基于深度學(xué)習(xí)的蛋?質(zhì)結(jié)構(gòu)預(yù)測，再到后續(xù)的結(jié)果分析，這是?個?度異構(gòu)的?作負(fù)載。

?英特爾? ?強(qiáng)? 可擴(kuò)展處理器可以輕松勝任這一系列多樣化的任務(wù)。以?強(qiáng)? CPU Max系列處理器為例，它采?全新微架構(gòu)、更多內(nèi)核（最?達(dá)56個），能以更?頻率和更?緩存，去應(yīng)對?通量的預(yù)處理和后處理?作。

它在內(nèi)存和輸入/輸出（I/O）子系統(tǒng)性能上有著顯著的增強(qiáng)，還結(jié)合大容量末級緩存使AlphaFold2推理過程中關(guān)鍵的張量吞吐獲得了大幅提升。

△英特爾? 至強(qiáng)? CPU Max 系列處理器

其次，由于AlphaFold2所采?的深度學(xué)習(xí)模型規(guī)模巨?，推理過程中的張量運算不僅量?，且維度極?。這就要求承載平臺具備強(qiáng)?的AI運算加速能?。

在這?點上，新款?強(qiáng)? 系列處理器內(nèi)置的英特爾? AMX（?級矩陣擴(kuò)展）技術(shù)，可以顯著加速?規(guī)模矩陣乘法運算。

在FP32/BF16混合精度計算下，其理論峰值可達(dá)每時鐘周期1024次乘加操作。針對AlphaFold2推理任務(wù)中所需的大量矩陣運算操作，AMX_BF16能在保持較高精度的同時，提高計算速度并減少存儲空間。

△AMX_BF16推理優(yōu)化帶來更低內(nèi)存占用和更大輸入長度

另???，AlphaFold2因其?維張量運算和?序列并?計算，在推理過程中常?臨超?內(nèi)存需求，不光影響推理速度，還會限制更?蛋?質(zhì)序列的預(yù)測。

為此英特爾從軟硬協(xié)同的方式給出完整解決方案。

一面是提升內(nèi)存容量和帶寬。解決方案中，英特爾? ?強(qiáng)? CPU Max系列處理器除支持DDR5內(nèi)存外，還集成了HBM（?帶寬內(nèi)存）。單顆處理器的HBM容量?達(dá)64GB，且具有高達(dá)460GB/s帶寬。

另一面是提供了多種降低內(nèi)存的軟件優(yōu)化方法。如面向PyTorch對張量計算原語（Tensor Processing Primitives，TPP）技術(shù)進(jìn)行擴(kuò)展，以及切分Attention模塊和算子融合的推理優(yōu)化方案，幫助AlphaFold2在通用矩陣乘法計算中所需的內(nèi)存峰值大幅降低。

△熱點算子與融合效果

經(jīng)過一系列加強(qiáng)和優(yōu)化后，最終效果如何呢？

如圖所示，在基于至強(qiáng)? CPU Max系列處理器的優(yōu)化流程中，每個優(yōu)化步驟獲得的提升累積后，獲得了相對于基線性能（對比組1，基于第三代至強(qiáng)? 可擴(kuò)展處理器，未實施優(yōu)化）高達(dá)33.97倍的通量提升。

根據(jù)測算，性能提升中的74%源自預(yù)處理階段的高通量優(yōu)化，26%要歸功于對推理過程的優(yōu)化。

此外，在同樣開啟IPEX（面向PyTorch的英特爾? 擴(kuò)展優(yōu)化框架）的情況下，相比對比組2（基于第三代至強(qiáng)? 可擴(kuò)展處理器，但實施過優(yōu)化），方案在升級使用至強(qiáng)? CPU Max 系列處理器后，其內(nèi)置的HBM內(nèi)存、英特爾? AMX的加成，則帶來了48.3%的性能提升。

△切分Attention模塊和算子融合的推理優(yōu)化方案

而且值得一提的是，在一項基于某公有云服務(wù)的測試中，基于至強(qiáng)? CPU平臺構(gòu)建的AlphaFold2解決方案還在性能上獲得了遠(yuǎn)優(yōu)于某高端GPU平臺的表現(xiàn)，同時也優(yōu)于由CPU+GPU混合構(gòu)建的方案。

這可是一個非常難得的成績——畢竟過去在很多AI應(yīng)用的測試或?qū)崙?zhàn)中，CPU能有接近或媲美GPU的表現(xiàn)就已經(jīng)算是成功，而AlphaFold2上至強(qiáng)? 平臺則實現(xiàn)了性能+蛋白質(zhì)預(yù)測序列長度的全面反超。

現(xiàn)在還剩下最后一個問題，多個蛋白結(jié)果的解析模型AlphaFold2 Multimer。

也就是從預(yù)測單個蛋白質(zhì)三維結(jié)構(gòu)，發(fā)展到了對多個蛋白質(zhì)分子之間的相互作用及所形成的復(fù)合體結(jié)構(gòu)進(jìn)行預(yù)測。

CPU在這一演變過程中的支持力度如何呢？

答案是不用擔(dān)心！

基于英特爾? 架構(gòu)的AlphaFold2解決方案同樣也面向AlphaFold2 Multimer的管線結(jié)構(gòu)進(jìn)行了優(yōu)化與驗證。

雖然后者的管線結(jié)構(gòu)已根據(jù)蛋白質(zhì)復(fù)合體結(jié)構(gòu)預(yù)測的需求進(jìn)行了調(diào)整，但英特爾AlphaFold2上的優(yōu)化方案，在被用于AlphaFold2 Multimer時同樣有效。

△面向AlphaFold2 Multimer模式的方案實現(xiàn)

CPU加速新藥發(fā)現(xiàn)不是夢

回顧以往，研發(fā)?種新藥動輒需要10年時間，投?20億美元才能起步。

?在AI的助?下，這?成本正?幅降低。以英矽智能為例，它們進(jìn)展最快的項目僅?18個?就找到了治療特發(fā)性肺纖維化（IPF）的潛在全球首創(chuàng)候選藥物并通過實驗驗證，總成本約為280萬美元。

展望未來，隨著AI技術(shù)的進(jìn)?步發(fā)展?jié)B透，它必將重塑制藥業(yè)的創(chuàng)新模式，讓新藥研發(fā)變得更加?效、精準(zhǔn)、經(jīng)濟(jì)。而在這一進(jìn)程中，相關(guān)的科學(xué)計算及AI應(yīng)用任務(wù)，依然需要有強(qiáng)大的算力支撐。

從英矽智能等公司的實踐來看，以?強(qiáng)? 處理器為代表的CPU平臺，正憑借其在性能、成本、?態(tài)等??的獨特優(yōu)勢，成為推動AI時代制藥創(chuàng)新的重要“引擎”。

這也預(yù)?著，在AI改變眾多?業(yè)的當(dāng)下，CPU加速AI應(yīng)用落地，幫助用戶節(jié)支增效以及推進(jìn)其技術(shù)和業(yè)務(wù)創(chuàng)新的腳步從未停止。

AI讓新藥研發(fā)進(jìn)?“?動駕駛”模式，?英特爾? ?強(qiáng)? 處理器則提供了它所需的源源不斷的動?。

在這種合作模式下，AI+制藥還將擦出怎樣的?花，就很值得期待了。

為了科普CPU在AI推理新時代的玩法，量子位開設(shè)了《最“in”AI》專欄，將從技術(shù)科普、行業(yè)案例、實戰(zhàn)優(yōu)化等多個角度全面解讀。

我們希望通過這個專欄，讓更多的人了解英特爾? 架構(gòu)CPU在AI推理加速，甚至是整個AI平臺或全流程加速上的實踐成果，重點就是如何更好地利用CPU來提升AI，包括大模型應(yīng)用的性能和效率。

未來隨著英特爾AI產(chǎn)品技術(shù)組合的進(jìn)一步擴(kuò)展和豐富，我們還將在這里為大家提供更多產(chǎn)品技術(shù)上的優(yōu)秀用例與方案分享，以及技術(shù)應(yīng)用指南。

更多關(guān)于基于英特爾? 架構(gòu)的AlphaFold2解決方案，可點擊文末閱讀原文進(jìn)一步了解。