豐色 明敏 發自 凹非寺
量子位 | 公衆号 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 · 頭條号簽約
關注我們，第一時間獲知前沿科技動态