CPU也能速刷AlphaFold2？英特爾：請收下這份23倍通量優化指南
豐色 明敏 發自 凹非寺
量子位 | 公眾號 QbitAI
攪翻計算生物界的AlphaFold2一開源，各種加速方案就爭相湧現。
妹想到啊，現在居然有了個CPU的推理優化版本，不用GPU，效果也出人意料的好——
端到端的通量足足提升到原來的23.11倍。
△ 高達23.11倍的提升不是一蹴而就，但依然驚艷
換個更具體的數據，它直接讓AlphaFold2的通量從每天約4.6個序列提升到了約105.4個。
要知道，由DeepMind開源的AlphaFold2，通過AI演算法對蛋白質結構預測實現了接近實驗精度的精準預測，可謂是公認的AI for Science標杆。
該領域也一直被認為是最吃AI專用加速晶元，如GPU紅利的前沿方向。
這一最新成果的釋出，就意味著CPU也能在AI for Science領域佔有一席之地，並發揮巨大、而且是出乎大家意料的威力。
同樣，這個成果也意味著AI for Science的入場門檻正在被拉低，對那些想從事基礎科研和創新，但還沒有布局異構IT基礎設施，或沒有大規模採用AI專用加速晶元的企業和機構而言也是個難得的福音，意味著他們依靠更常見的IT基礎設施就能開展工作。
那麼這個優化方案究竟是如何做到的？一起來看。
只用CPU加速AlphaFold2，優勢在哪？
這在其預處理、推理、後處理三個部分都有所體現。
預處理階段，由於輸入的氨基酸序列所含信息不多，AlphaFold2一般會先利用已知信息（蛋白質序列、結構模板）來提升預測精度，以此拿到MSA表徵（MSA representation）和鄰接表徵（pair representation）的三維張量。
這就意味著AlphaFold2屬於大張量模型，在嵌入層的瓶頸不在於並行計算，而是在於內存消耗和異構數據通信。
這正是CPU所擅長的領域。
再看模型推理階段。
該階段通過一個由48個塊（Block）組成的Evoformer網路進行表徵融合。
該網路的機制是利用Self-Attention來學習蛋白質的三角幾何約束信息，讓兩種表徵信息相互影響，從而使得模型能直接推理出相應的三維結構，且要循環三次。
結構層還會基於不動點注意力機制，對三維結構的每個原子進行預測，最後合成一個高度準確的結果。
這一番動作下來，對算力是個大考驗。
而且原版AlphaFold2會受到顯存限制，導致能夠探索的蛋白質序列長度不足1000aa。但很多蛋白質的序列長度動輒都是2k、3k。
最終在後處理階段，將使用Amber力場分析方法對獲得的三維結構參數優化，並輸出最終的蛋白質三維結構。
DeepMind團隊提到，他們用128塊TPUv3從頭訓練一遍AlphaFold2，需要11天的時間。
同時，AlphaFold2代碼是基於JAX的，偏向於專業從事AI科學計算的研究人員，普通開發人員部署起來也比較困難。
種種挑戰之下，導致AlphaFold2自開源後，相應的加速方案也接連湧現。不過無論是訓練還是推理，市面上更多見的，還是基於AI專用加速晶元，如GPU的方案。
完全基於CPU的加速方案還是頭一回見，而且一上來就在性能增幅上震驚四座，推理通量可提升到優化前的23倍之多。
具體到底是怎麼做的？
如何只靠CPU增效達23倍？
提到CPU，你可能已經想到方案的提出者是誰了——
沒錯，就是那個名字，英特爾。
他們基於目前最新的第三代至強可擴展平台，最終實現了「23.11倍」的通量優化成果（相比未優化時），其中有5.05倍是靠模型本身的優化所帶來，還有4.56倍則是來自傲騰持久內存提供的TB級內存支持。
△ 第三代英特爾? 至強? 可擴展處理器
其整體流程，就是先在預處理階段對模型進行高通量優化，然後將模型遷移到PyTorch框架下，接著再在PyTorch版本上進行細節上的推理優化，最後給予TB級內存支持以解決AlphaFold2的內存瓶頸，由此達到不輸專用加速晶元的效果。
更具體點，這些優化一共分為五步。
01、預處理階段高通量優化
如前文所述，此階段模型在進行蛋白質序列和模版搜索時需要計算平台執行大量的向量/矩陣運算——處理器能不能夠火力全開就顯得尤為重要。
第一步優化就在此展開，不過這步的優化非常簡單，就是藉助至強可擴展處理器自帶的多核心、多線程和大容量高速緩存能力直接加速，提升MSA和模板搜索通量。
至強可擴展處理器內置的AVX-512指令集和支持的NUMA ( Non-Uniform Memory Access，非一致存儲訪問) 架構等技術，能以提供最大512位向量計算能力的顯著高位寬優勢，來提升計算過程中的向量化並行程度，從而進一步提升預處理階段的整體效率。
這步優化支持所有至強可擴展系列CPU，且只需在ICC編譯器中添加一句簡單的配置就OK：
-O3 -no-prec-div -march=icelake-server
02、遷移到面向英特爾? 架構優化的 PyTorch
在預處理階段的高通量優化完成後，就需要將模型遷至PyTorch了。
因為原始AlphaFold2所基於的JAX庫所提供的加速能力主要針對GPU，且在英特爾? 架構平台上能夠發揮的功能有限。
而PyTorch擁有良好的動態圖糾錯方法，與haiku-API有著相似的風格（AlphaFold2一部分也基於haiku-API實現），就更別說還有英特爾? oneAPI工具套件提供的針對PyTorch的優化「利器」：Intel? Extensions for PyTorch （IPEX）。
因此，為了實現更好的優化效果，需要在這裡完成 PyTorch版本的遷移。
03、PyTorch JIT
接下來，為了提高模型的推理速度，便於後續利用IPEX的運算元融合等加速手段進行深入優化，英特爾又將遷移後的代碼進行了一系列的API改造，在不改變網路拓撲的前提下，引入PyTorch Just-In-Time (JIT) 圖編譯技術，將網路最終轉化為靜態圖。
以上都還只是「熱身動作」，下面才是展現「真正的實力」的時候——
04、切分Attention 運算元融合
首先，通過演算法設計分析，英特爾發現，在AlphaFold2模型的嵌入層有一個叫做ExtraMsaStack的模塊，其注意力模塊包含了大量的偏移量計算。
這些運算需要靠張量間的矩陣運算來完成。
其過程就會伴隨著張量的擴張，而張量擴張到一定規模後，就會讓模型內存需求變得巨大。
比如一個「5120x1x1x64」的張量，其初始內存需求只要1.25MB，擴張過程中就可達到930MB。
這一下子爆出的內存峰值壓力，會讓內存資源在短時間耗盡，繼而可能引發推理任務的失敗。
同時別忘了，大張量運算所需的海量內存還會帶來不可忽略的內存分配過程，徒增執行耗時。
那麼，英特爾的第四步優化就瞄準這兩個「痛點」，對注意力模塊來了個「大張量切分」的優化思路，化大張量為多個小張量。
比如將上述「5120x1x1x64」的張量切分為「320x1x1x64」後，其擴張所需的內存就由930MB降至59.69MB，只佔原來的6.4%左右。
沒有了大內存之需，也就不需要進行內存分配了，因此，張量切分後推理速度也上來了。
比如從下圖我們就可以看到，注意力模塊的效率在切片前後有著非常明顯的差別。
△ 注意力模塊切分前後對比
這還沒完。
接著，英特爾利用PyTorch自帶的Profiler對AlphaFold2的Evoformer網路進行了運算元跟蹤分析。
然後他們發現，有兩種運算元（Einsum和Add）的資源佔用率很高，且總是連續同時存在。
因此，英特爾就使用IPEX工具提供的運算元融合能力將它倆的計算過程進行融合，以省去中間建立臨時緩存數據結構的時間，提高整體效率。
從下圖我們可以看到，兩運算元融合後光是在單元檢測中的計算效率就提升到了原來的6倍。
△ 運算元 Einsum Add 融合效果圖
至此，經過以上幾個步驟的優化，AlphaFold2在CPU上的總體性能已經得到了大約5倍的提升。
還差最後一步：
05、多實例並行優化
在這一步，英特爾先利用至強可擴展平台上基於NUMA架構的核心綁定技術，讓每個推理工作負載都能穩定地在同一組核心上執行，並優先訪問對應的近端內存，從而提供更優、也更穩定的並行算力輸出。
然後引入英特爾? MPI庫幫助模型在多實例並行推理計算時實現更優的時延、帶寬和可擴展性。
但這些動作還不足以破解限制AlphaFold2發揮潛能的一個重要因素：內存瓶頸。
眾所周知，在面向不同蛋白質的結構測序工作中，序列長度越長，推理計算複雜度就越大。
而在我們對模型進行了並行計算能力的優化後，更多計算實例的加入還會進一步凸顯這一問題。
英特爾用「星際探索」這一比喻對這種現象做了非常形象的說明。
△ 不同序列長度的蛋白質所代表的計算複雜度示意圖
可以看到，當蛋白質序列長度達到4700aa時，此時內存需求就已經大於1.3TB，計算複雜度對於系統來說就相當於「探索宇宙級別」了。
如果再加上64個實例並行執行，內存容量的需求就會衝到一個令人驚嘆的量級。
那麼，一個具備超大容量內存支持的平台就顯得尤為重要。
給至強可擴展處理器配上傲騰持久內存便可滿足這一需求。
△ 英特爾? 傲騰? 持久內存
提到傲騰持久內存，我們並不陌生。今年6月在《用CPU方案打破內存牆？學Paypal堆傲騰擴容量，漏查欺詐交易量可降至1/30》一文中已經小結了它對於AI應用的關鍵作用，即提供更大容量的內存子系統來滿足那些內存敏感型AI應用將更多數據貼近算力的需求。
目前最新一代的傲騰持久內存200系列，可以在提供接近主流DRAM內存性能的基礎上，實現每路高達4TB的容量，或者說，與DRAM內存組合時可提供每路高達6TB的內存總容量。
有了它，我們甚至能夠實現10000aa序列長度的蛋白結構預測。
用了它，英特爾這個方案的優化就基本完成，模型的總體性能也可在之前優化步驟的基礎上再次得到4.56倍的提升。
最後，附上一份用於這個優化方案的英特爾官方推薦配置。
△ AlphaFold2 高通量推理實現所需的系統——英特爾推薦配置
如果你想看更多細節中的細節，也可以訪問英特爾公開分享的白皮書《通量提升達23.11倍！至強可擴展平台助力AlphaFold2端到端優化》。
或者觀看英特爾聯合國際學術期刊《Science》推出的「架構師成長計劃」第二季第八期課程《AI驅動的生命科學創新範式之變》。
其中不但有量子位總編輯李根、晶泰科技首席研發科學家楊明俊博士和英特爾人工智慧架構師楊威博士圍繞這個主題的精彩討論，還有晶泰科技在AI製藥領域的領先實踐分享，以及英特爾這個AlphaFold2優化方案更為細緻且可視化的呈現。
經驗和成果在業界已開始擴散
光說不練假把式，英特爾這個方案看起來很美，但要真正實現落地才能讓人信服。
事實上，在對這個方案進行摸索和開發的過程中，英特爾與相關領域合作夥伴或用戶的協作與交流就一直沒停過，不但吸收了各方的經驗，實現了博採眾長和互相借鑒、啟發。
在成型後，也第一時間廣泛分享給到夥伴和用戶們，讓他們能夠根據自身特定的環境、應用狀況和需求，開展實戰驗證和推進更進一步的探索。
例如國內某高校就曾嘗試在數百台基於至強可擴展處理器的伺服器上，採用該方案提供的經驗和方法來進行測試，並取得了一舉兩得的結果——
其順利實踐了短序列高通量的、面向蛋白質組學級別的批量化結構預測，既降低了蛋白質組的AlphaFold2預測成本，又提高了集群推理的總通量。
上文提及的國內明星AI製藥公司晶泰科技，也在自身的研發中，將自主研發的AI演算法與AlphaFold2結合，從而驗證靶點、精確解析活性構象，為後續的藥物發現打下良好基礎。
通過充分利用CPU的TB級內存支持，在公有雲上部署英特爾版AlphaFold2優化方案後，科學家可以實現針對短序列的單節點高通量推理優化，從而加快蛋白組學結構分析進程，並預測序列長度超過4700aa的蛋白質序列。
這不光能拓展AlphaFold2在研發中的探索範圍，也能以更低的成本，讓更高精度的演算法工具作用於更早的研發環節，進一步加速藥物發現。
蛋白質結構預測，只是AI for Science的序幕
不久前，引爆了AI for Science的AlphaFold2又公開了新進展——
它已經成功預測出包括植物、細菌、真菌在內的100萬個物種的2.14億個蛋白質結構，並將數據集對外開源。
而這些，都還只是AI for Science的序幕。
在蛋白質結構預測、生物計算、藥物開發之外，AI在物理、天文、化學等領域也開始逐漸展露頭角。
前沿領域，今年《Nature》上的一項「改寫物理教科書」的研究，正是通過AI開展的。
歐洲核子研究組織（CERN）的科學家利用機器學習，發現了質子內部存在5個夸克的有力證據，這一成果顛覆了一直以來質子只有三個夸克的理論。
應用落地層面，前面提到的晶泰科技已經構建了一套智能計算、自動化實驗、專家經驗結合的三位一體的研發模式，提供一站式小分子藥物發現、大分子藥物發現，藥物固體形態研發，以及化學合成服務。
晶泰科技目前已建成數千平的自動化實驗室，與智能演算法「乾濕結合」，形成實驗數據與演算法預測間的交互閉環，保證AI演算法的產業落地和交付能力。
△ 晶泰科技官網
當下行業的大勢所趨，就是利用AI，從生產工具、生產關係等不同維度突破科學探索瓶頸。
英特爾所提出的CPU版AlphaFold2加速方案在生命科學領域中發揮出的巨大價值就是最有力的佐證。
而這或許還只是其技術在AI for Science領域釋出的一個起點。
實際上，當AI成為科學家的工具後，一種科研新範式已在應運而生。
它不同於亞里士多德時代的演繹法，不是基於經驗的試錯，不是將探索發現寄托在偶然的正確之上。
它也不是牛頓愛因斯坦時代的「假設再驗證」，不再依賴於人類群體中極少數天才的靈光一現。
它有一點像計算機時代下，結合理論進行數值模擬分析（如天氣預報），但仍舊有本質區別——
它將探索未知的基礎，第一次不再純粹基於人類群體的認知。
當海量客觀存在的數據成為「初始反應物」，在以深度學習為代表的技術驅動下，科研探索的邊界或許將發生前所未有的改變。
當下AlphaFold2一系列成果震驚世人，便是對此的最佳詮釋。
接下來，我們可能會看到生物、醫藥領域之外的更多科學家，樂意將AI作為科研探索的生產力、助手，更多突破人類想像的科研成果將會湧現……
光是想想就非常期待了~
最後，歡迎掃下方海報中的二維碼觀看課程，或者戳閱讀原文翻閱白皮書哦～
— 完 —
量子位 QbitAI · 頭條號簽約
關注我們，第一時間獲知前沿科技動態
