CPU也能速刷AlphaFold2？英特爾：請收下這份23倍通量優(yōu)化指南

芯智訊|芯智訊|Core Zhixun|

2022-10-25 19:49:05 7

攪翻計算生物界的AlphaFold2一開源，各種加速方案就爭相涌現。

妹想到啊，現在居然有了個CPU的推理優(yōu)化版本，不用GPU，效果也出人意料的好——

端到端的通量足足提升到原來的23.11倍。

△ 高達23.11倍的提升不是一蹴而就，但依然驚艷

換個更具體的數據，它直接讓AlphaFold2的通量從每天約4.6個序列提升到了約105.4個。

要知道，由DeepMind開源的AlphaFold2，通過AI算法對蛋白質結構預測實現了接近實驗精度的*預測，可謂是公認的AI for Science標桿。

該領域也一直被認為是最吃AI專用加速芯片，如GPU紅利的前沿方向。

這一最新成果的釋出，就意味著CPU也能在AI for Science領域占有一席之地，并發(fā)揮巨大、而且是出乎大家意料的威力。

同樣，這個成果也意味著AI for Science的入場門檻正在被拉低，對那些想從事基礎科研和創(chuàng)新，但還沒有布局異構IT基礎設施，或沒有大規(guī)模采用AI專用加速芯片的企業(yè)和機構而言也是個難得的福音，意味著他們依靠更常見的IT基礎設施就能開展工作。

那么這個優(yōu)化方案究竟是如何做到的？一起來看。

只用CPU加速AlphaFold2，優(yōu)勢在哪？

這在其預處理、推理、后處理三個部分都有所體現。

預處理階段，由于輸入的氨基酸序列所含信息不多，AlphaFold2一般會先利用已知信息（蛋白質序列、結構模板）來提升預測精度，以此拿到MSA表征（MSA representation）和鄰接表征（pair representation）的三維張量。

這就意味著AlphaFold2屬于大張量模型，在嵌入層的瓶頸不在于并行計算，而是在于內存消耗和異構數據通信。

這正是CPU所擅長的領域。

再看模型推理階段。

該階段通過一個由48個塊（Block）組成的Evoformer網絡進行表征融合。

該網絡的機制是利用Self-Attention來學習蛋白質的三角幾何約束信息，讓兩種表征信息相互影響，從而使得模型能直接推理出相應的三維結構，且要循環(huán)三次。

結構層還會基于不動點注意力機制，對三維結構的每個原子進行預測，最后合成一個高度準確的結果。

這一番動作下來，對算力是個大考驗。

而且原版AlphaFold2會受到顯存限制，導致能夠探索的蛋白質序列長度不足1000aa。但很多蛋白質的序列長度動輒都是2k、3k。

最終在后處理階段，將使用Amber力場分析方法對獲得的三維結構參數優(yōu)化，并輸出最終的蛋白質三維結構。

DeepMind團隊提到，他們用128塊TPUv3從頭訓練一遍AlphaFold2，需要11天的時間。

同時，AlphaFold2代碼是基于JAX的，偏向于專業(yè)從事AI科學計算的研究人員，普通開發(fā)人員部署起來也比較困難。

種種挑戰(zhàn)之下，導致AlphaFold2自開源后，相應的加速方案也接連涌現。不過無論是訓練還是推理，市面上更多見的，還是基于AI專用加速芯片，如GPU的方案。

完全基于CPU的加速方案還是頭一回見，而且一上來就在性能增幅上震驚四座，推理通量可提升到優(yōu)化前的23倍之多。

具體到底是怎么做的？

如何只靠CPU增效達23倍？

提到CPU，你可能已經想到方案的提出者是誰了——

沒錯，就是那個名字，英特爾。

他們基于目前最新的第三代至強可擴展平臺，最終實現了“23.11倍”的通量優(yōu)化成果（相比未優(yōu)化時），其中有5.05倍是靠模型本身的優(yōu)化所帶來，還有4.56倍則是來自傲騰持久內存提供的TB級內存支持。

△ 第三代英特爾^?至強^?可擴展處理器

其整體流程，就是先在預處理階段對模型進行高通量優(yōu)化，然后將模型遷移到PyTorch框架下，接著再在PyTorch版本上進行細節(jié)上的推理優(yōu)化，最后給予TB級內存支持以解決AlphaFold2的內存瓶頸，由此達到不輸專用加速芯片的效果。

更具體點，這些優(yōu)化一共分為五步。

01、預處理階段高通量優(yōu)化

如前文所述，此階段模型在進行蛋白質序列和模版搜索時需要計算平臺執(zhí)行大量的向量/矩陣運算——處理器能不能夠火力全開就顯得尤為重要。

*步優(yōu)化就在此展開，不過這步的優(yōu)化非常簡單，就是借助至強可擴展處理器自帶的多核心、多線程和大容量高速緩存能力直接加速，提升MSA和模板搜索通量。

至強可擴展處理器內置的AVX-512指令集和支持的NUMA ( Non-Uniform Memory Access，非一致存儲訪問) 架構等技術，能以提供最大512位向量計算能力的顯著高位寬優(yōu)勢，來提升計算過程中的向量化并行程度，從而進一步提升預處理階段的整體效率。

這步優(yōu)化支持所有至強可擴展系列CPU，且只需在ICC編譯器中添加一句簡單的配置就OK：

-O3 -no-prec-div -march=icelake-server

02、遷移到面向英特爾^?架構優(yōu)化的 PyTorch

在預處理階段的高通量優(yōu)化完成后，就需要將模型遷至PyTorch了。

因為原始AlphaFold2所基于的JAX庫所提供的加速能力主要針對GPU，且在英特爾^?架構平臺上能夠發(fā)揮的功能有限。

而PyTorch擁有*的動態(tài)圖糾錯方法，與haiku-API有著相似的風格（AlphaFold2一部分也基于haiku-API實現），就更別說還有英特爾^{^?}oneAPI工具套件提供的針對PyTorch的優(yōu)化“利器”：Intel^{^?} Extensions for PyTorch （IPEX）。

因此，為了實現更好的優(yōu)化效果，需要在這里完成 PyTorch版本的遷移。

03、PyTorch JIT

接下來，為了提高模型的推理速度，便于后續(xù)利用IPEX的算子融合等加速手段進行深入優(yōu)化，英特爾又將遷移后的代碼進行了一系列的API改造，在不改變網絡拓撲的前提下，引入PyTorch Just-In-Time (JIT) 圖編譯技術，將網絡最終轉化為靜態(tài)圖。

以上都還只是“熱身動作”，下面才是展現“真正的實力”的時候——

04、切分Attention+算子融合

首先，通過算法設計分析，英特爾發(fā)現，在AlphaFold2模型的嵌入層有一個叫做ExtraMsaStack的模塊，其注意力模塊包含了大量的偏移量計算。

這些運算需要靠張量間的矩陣運算來完成。

其過程就會伴隨著張量的擴張，而張量擴張到一定規(guī)模后，就會讓模型內存需求變得巨大。

比如一個“5120x1x1x64”的張量，其初始內存需求只要1.25MB，擴張過程中就可達到930MB。

這一下子爆出的內存峰值壓力，會讓內存資源在短時間耗盡，繼而可能引發(fā)推理任務的失敗。

同時別忘了，大張量運算所需的海量內存還會帶來不可忽略的內存分配過程，徒增執(zhí)行耗時。

那么，英特爾的第四步優(yōu)化就瞄準這兩個“痛點”，對注意力模塊來了個“大張量切分”的優(yōu)化思路，化大張量為多個小張量。

比如將上述“5120x1x1x64”的張量切分為“320x1x1x64”后，其擴張所需的內存就由930MB降至59.69MB，只占原來的6.4%左右。

沒有了大內存之需，也就不需要進行內存分配了，因此，張量切分后推理速度也上來了。

比如從下圖我們就可以看到，注意力模塊的效率在切片前后有著非常明顯的差別。

△ 注意力模塊切分前后對比

這還沒完。

接著，英特爾利用PyTorch自帶的Profiler對AlphaFold2的Evoformer網絡進行了算子跟蹤分析。

然后他們發(fā)現，有兩種算子（Einsum和Add）的資源占用率很高，且總是連續(xù)同時存在。

因此，英特爾就使用IPEX工具提供的算子融合能力將它倆的計算過程進行融合，以省去中間建立臨時緩存數據結構的時間，提高整體效率。

從下圖我們可以看到，兩算子融合后光是在單元檢測中的計算效率就提升到了原來的6倍。

△ 算子 Einsum+Add 融合效果圖

至此，經過以上幾個步驟的優(yōu)化，AlphaFold2在CPU上的總體性能已經得到了大約5倍的提升。

還差最后一步：

05、多實例并行優(yōu)化

在這一步，英特爾先利用至強可擴展平臺上基于NUMA架構的核心綁定技術，讓每個推理工作負載都能穩(wěn)定地在同一組核心上執(zhí)行，并優(yōu)先訪問對應的近端內存，從而提供更優(yōu)、也更穩(wěn)定的并行算力輸出。

然后引入英特爾^?MPI庫幫助模型在多實例并行推理計算時實現更優(yōu)的時延、帶寬和可擴展性。

但這些動作還不足以破解限制AlphaFold2發(fā)揮潛能的一個重要因素：內存瓶頸。

眾所周知，在面向不同蛋白質的結構測序工作中，序列長度越長，推理計算復雜度就越大。

而在我們對模型進行了并行計算能力的優(yōu)化后，更多計算實例的加入還會進一步凸顯這一問題。

英特爾用“星際探索”這一比喻對這種現象做了非常形象的說明。

△ 不同序列長度的蛋白質所代表的計算復雜度示意圖

可以看到，當蛋白質序列長度達到4700aa時，此時內存需求就已經大于1.3TB，計算復雜度對于系統(tǒng)來說就相當于“探索宇宙級別”了。

如果再加上64個實例并行執(zhí)行，內存容量的需求就會沖到一個令人驚嘆的量級。

那么，一個具備超大容量內存支持的平臺就顯得尤為重要。

給至強可擴展處理器配上傲騰持久內存便可滿足這一需求。

△ 英特爾^?傲騰? 持久內存

提到傲騰持久內存，我們并不陌生。今年6月在《用CPU方案打破內存墻？學Paypal堆傲騰擴容量，漏查欺詐交易量可降至1/30》一文中已經小結了它對于AI應用的關鍵作用，即提供更大容量的內存子系統(tǒng)來滿足那些內存敏感型AI應用將更多數據貼近算力的需求。

目前最新一代的傲騰持久內存200系列，可以在提供接近主流DRAM內存性能的基礎上，實現每路高達4TB的容量，或者說，與DRAM內存組合時可提供每路高達6TB的內存總容量。

有了它，我們甚至能夠實現10000aa序列長度的蛋白結構預測。

用了它，英特爾這個方案的優(yōu)化就基本完成，模型的總體性能也可在之前優(yōu)化步驟的基礎上再次得到4.56倍的提升。

最后，附上一份用于這個優(yōu)化方案的英特爾官方推薦配置。

△ AlphaFold2 高通量推理實現所需的系統(tǒng)——英特爾推薦配置

如果你想看更多細節(jié)中的細節(jié)，也可以訪問英特爾公開分享的白皮書《通量提升達23.11倍！至強可擴展平臺助力AlphaFold2端到端優(yōu)化》。

或者觀看英特爾聯(lián)合國際學術期刊《Science》推出的“架構師成長計劃”第二季第八期課程《AI驅動的生命科學創(chuàng)新范式之變》。

其中不但有量子位總編輯李根、晶泰科技首席研發(fā)科學家楊明俊博士和英特爾人工智能架構師楊威博士圍繞這個主題的精彩討論，還有晶泰科技在AI制藥領域的*實踐分享，以及英特爾這個AlphaFold2優(yōu)化方案更為細致且可視化的呈現。

經驗和成果在業(yè)界已開始擴散

光說不練假把式，英特爾這個方案看起來很美，但要真正實現落地才能讓人信服。

事實上，在對這個方案進行摸索和開發(fā)的過程中，英特爾與相關領域合作伙伴或用戶的協(xié)作與交流就一直沒停過，不但吸收了各方的經驗，實現了博采眾長和互相借鑒、啟發(fā)。

在成型后，也*時間廣泛分享給到伙伴和用戶們，讓他們能夠根據自身特定的環(huán)境、應用狀況和需求，開展實戰(zhàn)驗證和推進更進一步的探索。

例如國內某高校就曾嘗試在數百臺基于至強可擴展處理器的服務器上，采用該方案提供的經驗和方法來進行測試，并取得了一舉兩得的結果——

其順利實踐了短序列高通量的、面向蛋白質組學級別的批量化結構預測，既降低了蛋白質組的AlphaFold2預測成本，又提高了集群推理的總通量。

上文提及的國內明星AI制藥公司晶泰科技，也在自身的研發(fā)中，將自主研發(fā)的AI算法與AlphaFold2結合，從而驗證靶點、精確解析活性構象，為后續(xù)的藥物發(fā)現打下*基礎。

通過充分利用CPU的TB級內存支持，在公有云上部署英特爾版AlphaFold2優(yōu)化方案后，科學家可以實現針對短序列的單節(jié)點高通量推理優(yōu)化，從而加快蛋白組學結構分析進程，并預測序列長度超過4700aa的蛋白質序列。

這不光能拓展AlphaFold2在研發(fā)中的探索范圍，也能以更低的成本，讓更高精度的算法工具作用于更早的研發(fā)環(huán)節(jié)，進一步加速藥物發(fā)現。

蛋白質結構預測，只是AI for Science的序幕

不久前，引爆了AI for Science的AlphaFold2又公開了新進展——

它已經成功預測出包括植物、細菌、真菌在內的100萬個物種的2.14億個蛋白質結構，并將數據集對外開源。

而這些，都還只是AI for Science的序幕。

在蛋白質結構預測、生物計算、藥物開發(fā)之外，AI在物理、天文、化學等領域也開始逐漸展露頭角。

前沿領域，今年《Nature》上的一項“改寫物理教科書”的研究，正是通過AI開展的。

歐洲核子研究組織（CERN）的科學家利用機器學習，發(fā)現了質子內部存在5個夸克的有力證據，這一成果顛覆了一直以來質子只有三個夸克的理論。

應用落地層面，前面提到的晶泰科技已經構建了一套智能計算、自動化實驗、專家經驗結合的三位一體的研發(fā)模式，提供一站式小分子藥物發(fā)現、大分子藥物發(fā)現，藥物固體形態(tài)研發(fā)，以及化學合成服務。

晶泰科技目前已建成數千平的自動化實驗室，與智能算法“干濕結合”，形成實驗數據與算法預測間的交互閉環(huán)，保證AI算法的產業(yè)落地和交付能力。

△ 晶泰科技官網

當下行業(yè)的大勢所趨，就是利用AI，從生產工具、生產關系等不同維度突破科學探索瓶頸。

英特爾所提出的CPU版AlphaFold2加速方案在生命科學領域中發(fā)揮出的巨大價值就是最有力的佐證。

而這或許還只是其技術在AI for Science領域釋出的一個起點。

實際上，當AI成為科學家的工具后，一種科研新范式已在應運而生。

它不同于亞里士多德時代的演繹法，不是基于經驗的試錯，不是將探索發(fā)現寄托在偶然的正確之上。

它也不是牛頓愛因斯坦時代的“假設再驗證”，不再依賴于人類群體中極少數天才的靈光一現。

它將探索未知的基礎，*次不再純粹基于人類群體的認知。

當海量客觀存在的數據成為“初始反應物”，在以深度學習為代表的技術驅動下，科研探索的邊界或許將發(fā)生前所未有的改變。

當下AlphaFold2一系列成果震驚世人，便是對此的最佳詮釋。

接下來，我們可能會看到生物、醫(yī)藥領域之外的更多科學家，樂意將AI作為科研探索的生產力、助手，更多突破人類想象的科研成果將會涌現……

光是想想就非常期待了~

免責聲明：本文采摘自“芯智訊”，本文僅代表作者個人觀點，不代表薩科微及行業(yè)觀點，只為轉載與分享，支持保護知識產權，轉載請注明原出處及作者，如有侵權請聯(lián)系我們刪除。

上一篇：: 薩科微金航標公司喬遷之喜

下一篇：: 全球封裝基板供應商列表

CPU也能速刷AlphaFold2？英特爾：請收下這份23倍通量優(yōu)化指南

客服QQ2635709730

客服QQ
2635709730