預測超長蛋白質這事，CPU贏了

十三 2023-06-30 19:51:42 來源：量子位

金磊 楊凈 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

AI模型的推理在CPU上完成加速和優(yōu)化，竟然不輸傳統(tǒng)方案？

至少在生命科學和醫(yī)療制藥方向，已經透露出這種信號。

例如在處理AlphaFold2這類大型模型這件事上，大眾普遍的認知可能就是堆GPU來進行大規(guī)模計算。

但其實從去年開始，CPU便開始苦練內功，使端到端的通量足足提升到了原來的23.11倍。

而現(xiàn)如今，CPU讓這個數(shù)值great again——再次提升3.02倍！

不論是像抗菌肽這種較短的氨基酸序列，還是像亨氏綜合征蛋白這樣超長的序列，都可以輕松hold住。

而且所有的預測任務，在不考慮最高通量、僅僅是順序執(zhí)行，8個小時就能全部搞定。

甚至國內已經有云服務提供商做了類似的優(yōu)化方案：

相比于GPU，基于CPU的加速方案在性價比上更為理想，而且在特定的情況下（超過300或400氨基酸），幾乎只有CPU能把它算完，而GPU的失敗率會很高。

要知道，像AlphaFold2這類任務，可以說是公認的AI for Science標桿。

從上述的種種跡象表明，CPU不再是“你以為的你以為”，而是以一種新勢力進軍于此，并發(fā)揮著前所未有的威力。

CPU，正在大步邁進新時代。

英特爾自己刷新自己

事實上，此次備受關注的CPU加速方案，背后不是別人，正是發(fā)明了CPU的英特爾。

2022年，英特爾以第三代至強???可擴展處理器為硬件基座，使AlphaFold2通量優(yōu)化提升達23.11倍。一年后，他們在此基礎上，再次實現(xiàn)自我刷新。

2022年，英特爾基于第三代至強???可擴展平臺，針對AlphaFold2的設計特點，在預處理、模型推理、后處理三階段實現(xiàn)了端到端優(yōu)化。

如今，原有的五大端到端基礎步驟之上，第四代至強???可擴展處理器的加入，再次給AlphaFold2帶來整體推理性能的提升。

此次優(yōu)化方案主要圍繞預處理和模型推理兩個方面，基本劃分為五個步驟。

第一步：預處理階段，借助第三代或第四代至強???可擴展處理器的多核優(yōu)勢及其內置AVX-512技術，實現(xiàn)針對性的高通量優(yōu)化。

第二步到第五步模型推理階段的優(yōu)化，與2022年方案類似。

第二步，將深度學習模型遷移至面向英特爾???架構優(yōu)化的PyTorch，并逐模塊地從JAX/haiku完成代碼遷移。

第三步，引入JIT圖編譯技術，將網格轉化為靜態(tài)圖，以提高模型推理速度。

第四步，切分注意力模塊和算子融合，即對注意力模塊進行大張量切分的優(yōu)化思路；與此同時，使用IPEX（英特爾???擴展優(yōu)化框架，建議版本為IPEX-2.0.100+cpu或更高）對Einsum和Add兩種算子進行融合。

第五步，借助至強???可擴展平臺的計算和存儲優(yōu)勢實施針對性優(yōu)化。比如基于NUMA架構技術，挖掘多核心優(yōu)勢，破解多實例運算過程中的計算和內存瓶頸。

不過除了提供更強的基礎算力，第四代至強???可擴展平臺還帶來了諸多針對AI工作負載的優(yōu)化加速技術。

具體可以拆分為四項：（詳細優(yōu)化方案可點擊閱讀原文獲取）

一、TPP技術降低推理過程中的內存消耗

TPP（Tensor Processing Primitives）相當于是一種虛擬的張量指令集架構，能讓英特爾???AVX-512等物理指令集予以抽象，生成經過優(yōu)化的平臺代碼。

具體到計算執(zhí)行上，TPP能實現(xiàn)兩種優(yōu)化方式：以單指令多數(shù)據(jù)方式處理數(shù)據(jù)；優(yōu)化內存訪問模式，提升緩存命中率來提高數(shù)值計算和訪存效率。

這樣一來，狹長矩陣乘法的空間復雜度從 O (n^2) 降為 O (n) ，運算內存峰值也將大幅降低，更有助于處理長序列蛋白結構預測的問題。

二、支持DDR5內存與大容量緩存帶來張量吞吐提升

AlphaFold2中大量的矩陣計算過程需要內存來支撐，因此內存性能影響著整個模型運行性能。

第四代至強?? 可擴展處理器帶來兩種解決思路——支持DDR5內存，以及大容量末級緩存：

一方面，與上個方案DDR4內存帶寬25.6GBps?(3,200MHz）相比，DDR5內存帶寬提升了超50%，達到38.4GBps?(4,800MHz）以上 ；另一方面，末級緩存也由上一代的最高 60MB提升至現(xiàn)在最高112.5MB，幅度87.5%。

三、內置AI加速引擎AMX

英特爾在第四代至強???可擴展處理器中創(chuàng)新內置了AI加速器——英特爾???AMX，類似GPU里的張量核心，加速深度學習推理過程并減少存儲空間。

它支持INT8、BF16等低精度數(shù)據(jù)類型，尤其BF16數(shù)據(jù)類型在精度上的表現(xiàn)不遜于FP32數(shù)據(jù)類型，AlphaFold2使用AMX_BF16后，推理時間可縮短數(shù)倍之多。

四、高帶寬內存HBM2e增加訪存通量

每個英特爾???至強???CPU Max系列，都擁有4個基于第二代增強型高帶寬內存?(HBM2e)?的堆棧，總容量為64GB?(每個堆棧的容量為16GB）。

由于能同時訪問多個DRAM芯片，它可提供高達1TB/s的帶寬。而且配置更靈活，有三種不同模式與DDR5內存一起協(xié)同工作：HBM Only、HBM Flat以及HBM Cache。

綜上，第四代英特爾???至強???可擴展處理器所帶來的四種優(yōu)化技術讓AlphaFold2的端到端通量得到了再進一步提升，與第三代相比實現(xiàn)了高達3.02倍的多實例通量提升。

當然，除了CPU之外，英特爾在探索驗證AlphaFold2優(yōu)化方案、步驟和經驗過程中，同樣也能提供其他AI加速芯片，給產業(yè)鏈上的生態(tài)伙伴提供強勁支持。

甚至已經給出了行業(yè)備受認可的解決方案。

就在前段時間，英特爾聯(lián)合Github上知名的AI+科學計算的開源項目——Colossal-AI的團隊潞晨科技，成功優(yōu)化了AlphaFold2蛋白質結構預測的性能，并將其方案開源。

基于AI專用加速芯片Habana???Gaudi??，他們成功將端到端推理速度最高提升3.86倍（相較于此前使用的方案），應用成本相較于GPU方案最多降低39%。

醫(yī)藥和生命科學領域，AI還有何作為？

大模型，毋庸置疑是近來科技圈最為火爆的技術之一。

它憑借自身強算法、多數(shù)據(jù)、大算力的結合所帶來的泛用性，在醫(yī)藥和生命科學領域同樣大步發(fā)展著。

這一過程，AI宛如從破解人類的自然語言，躍進到了破解生命的自然語言：

• 人類自然語言大模型：從26個字母，到詞/句/段。
• 生命自然語言大模型：從21個氨基酸字母，到蛋白質/細胞/生命體。

那么具體而言，現(xiàn)在AI大模型可能會讓醫(yī)藥和生命科學領域產生怎樣的變革？

我們不妨以百圖生科推出的，世界首個AI大模型驅動的AI生成蛋白平臺AIGP（AI Generated Protein）為例來了解一番。

AIGP背后所依靠的，是一個千億參數(shù)的跨模態(tài)生命科學大模型，通過“挖掘公開數(shù)據(jù)和獨特自產數(shù)據(jù)”、“跨模態(tài)預訓練和科學計算”，以及“蛋白質讀寫系統(tǒng)和細胞讀寫系統(tǒng)”，三大步驟實現(xiàn)對蛋白質空間及生命體的建模。

如此大模型能力之下，百圖生科便具備了一系列給定Protein（抗原），設計與之以特定方式結合的Protein（抗體）的能力。

也因此參與到了一系列前沿藥物的開發(fā)，包括高性能免疫調控彈頭設計、難成藥靶點精準設計、定表位抗體彈頭設計、可溶性TCR設計等。

除此之外，百圖生科也具備對給定細胞/細胞組合，發(fā)現(xiàn)調控細胞的有效蛋白靶點/組合，并繼而快速設計調控蛋白的能力。

這就為多種疾病的靶點發(fā)現(xiàn)、耐藥/不響應患者改善、靶點科學線索轉化帶來新的可能。

不過有一說一，百圖生科的例子也是只是AI之于醫(yī)藥、生命科學領域變革的一隅。

但今年生物醫(yī)學領域的著名獎項（加拿大蓋爾德納獎）史無前例地頒給了人工智能科學家、DeepMind創(chuàng)始人Demis Hassabis等人。

這也從側面反映了生命科學、醫(yī)藥領域對于AI的認可，以及更多的期待。