金磊 楊淨 發自 凹非寺
AI模型的推理在CPU上完成加速和最佳化，竟然不輸傳統方案？
至少在生命科學和醫療製藥方向，已經透露出這種信號。
例如在處理AlphaFold2這類大型模型這件事上，大眾普遍的認知可能就是堆GPU來進行大規模計算。
但其實從去年開始，CPU便開始苦練內功，使端到端的通量足足提升到了原來的23.11倍。
而現如今，CPU讓這個數值great again——再次提升3.02倍！
不論是像抗菌肽這種較短的氨基酸序列，還是像亨氏綜合徵蛋白這樣超長的序列，都可以輕鬆hold住。
而且所有的預測任務，在不考慮最高通量、僅僅是順序執行，8個小時就能全部搞定。
甚至國內已經有云服務提供商做了類似的最佳化方案：
相比於GPU，基於CPU的加速方案在價效比上更為理想，而且在特定的情況下（超過300或400氨基酸），幾乎只有CPU能把它算完，而GPU的失敗率會很高。
要知道，像AlphaFold2這類任務，可以說是公認的AI for Science標杆。
從上述的種種跡象表明，CPU不再是「你以為的你以為」，而是以一種新勢力進軍於此，併發揮著前所未有的威力。
CPU，正在大步邁進新時代。
英特爾自己刷新自己
事實上，此次備受關注的CPU加速方案，背後不是別人，正是發明了CPU的英特爾。
2022年，英特爾以第三代至強®️可擴展處理器為硬體基座，使AlphaFold2通量最佳化提升達23.11倍。一年後，他們在此基礎上，再次實現自我刷新。
2022年，英特爾基於第三代至強®️可擴展平臺，針對AlphaFold2的設計特點，在預處理、模型推理、後處理三階段實現了端到端最佳化。
如今，原有的五大端到端基礎步驟之上，第四代至強®️可擴展處理器的加入，再次給AlphaFold2帶來整體推理性能的提升。
此次最佳化方案主要圍繞預處理和模型推理兩個方面，基本劃分為五個步驟。
第一步：預處理階段，藉助第三代或第四代至強®️可擴展處理器的多核優勢及其內建AVX-512技術，實現針對性的高通量最佳化。
第二步到第五步模型推理階段的最佳化，與2022年方案類似。
第二步，將深度學習模型遷移至面向英特爾®️架構最佳化的PyTorch，並逐模組地從JAX/haiku完成程式碼遷移。
第三步，引入JIT圖編譯技術，將網格轉化為靜態圖，以提高模型推理速度。
第四步，切分注意力模組和運算元融合，即對注意力模組進行大張量切分的最佳化思路；與此同時，使用IPEX（英特爾®️擴展最佳化框架，建議版本為IPEX-2.0.100+cpu或更高）對Einsum和Add兩種運算元進行融合。
第五步，藉助至強®️可擴展平臺的計算和儲存優勢實施針對性最佳化。比如基於NUMA架構技術，挖掘多核心優勢，破解多實例運算過程中的計算和記憶體瓶頸。
不過除了提供更強的基礎算力，第四代至強®️可擴展平臺還帶來了諸多針對AI工作負載的最佳化加速技術。
具體可以拆分為四項：（詳細最佳化方案可點選閱讀原文獲取）
一、TPP技術降低推理過程中的記憶體消耗
TPP（Tensor Processing Primitives）相當於是一種虛擬的張量指令集架構，能讓英特爾®️AVX-512等物理指令集予以抽象，生成經過最佳化的平臺程式碼。
具體到計算執行上，TPP能實現兩種最佳化方式：以單指令多資料方式處理資料；最佳化記憶體訪問模式，提升快取命中率來提高數值計算和訪存效率。
這樣一來，狹長矩陣乘法的空間複雜度從 O (n^2) 降為 O (n) ，運算記憶體峰值也將大幅降低，更有助於處理長序列蛋白結構預測的問題。
二、支持DDR5記憶體與大容量快取帶來張量吞吐提升
AlphaFold2中大量的矩陣計算過程需要記憶體來支撐，因此記憶體性能影響著整個模型運行性能。
第四代至強®️ 可擴展處理器帶來兩種解決思路——支持DDR5記憶體，以及大容量末級快取：
一方面，與上個方案DDR4記憶體頻寬25.6GBps (3,200MHz）相比，DDR5記憶體頻寬提升了超50%，達到38.4GBps (4,800MHz）以上 ；另一方面，末級快取也由上一代的最高 60MB提升至現在最高112.5MB，幅度87.5%。
三、內建AI加速引擎AMX
英特爾在第四代至強®️可擴展處理器中創新內建了AI加速器——英特爾®️AMX，類似GPU裡的張量核心，加速深度學習推理過程並減少儲存空間。
它支持INT8、BF16等低精度資料類型，尤其BF16資料類型在精度上的表現不遜於FP32資料類型，AlphaFold2使用AMX_BF16後，推理時間可縮短數倍之多。
四、高頻寬記憶體HBM2e增加訪存通量
每個英特爾®️至強®️CPU Max系列，都擁有4個基於第二代增強型高頻寬記憶體 (HBM2e) 的堆疊，總容量為64GB (每個堆疊的容量為16GB）。
由於能同時訪問多個DRAM晶片，它可提供高達1TB/s的頻寬。而且配置更靈活，有三種不同模式與DDR5記憶體一起協同工作：HBM Only、HBM Flat以及HBM Cache。
綜上，第四代英特爾®️至強®️可擴展處理器所帶來的四種最佳化技術讓AlphaFold2的端到端通量得到了再進一步提升，與第三代相比實現了高達3.02倍的多實例通量提升。
當然，除了CPU之外，英特爾在探索驗證AlphaFold2最佳化方案、步驟和經驗過程中，同樣也能提供其他AI加速晶片，給產業鏈上的生態夥伴提供強勁支持。
甚至已經給出了行業備受認可的解決方案。
就在前段時間，英特爾聯合Github上知名的AI+科學計算的開源項目——Colossal-AI的團隊潞晨科技，成功最佳化了AlphaFold2蛋白質結構預測的性能，並將其方案開源。
基於AI專用加速晶片Habana®️Gaudi®️，他們成功將端到端推理速度最高提升3.86倍（相較於此前使用的方案），應用成本相較於GPU方案最多降低39%。
醫藥和生命科學領域，AI還有何作為？
大模型，毋庸置疑是近來科技圈最為火爆的技術之一。
它憑藉自身強演算法、多資料、大算力的結合所帶來的泛用性，在醫藥和生命科學領域同樣大步發展著。
這一過程，AI宛如從破解人類的自然語言，躍進到了破解生命的自然語言：
-
人類自然語言大模型：從26個字母，到詞/句/段。
-
生命自然語言大模型：從21個氨基酸字母，到蛋白質/細胞/生命體。
那麼具體而言，現在AI大模型可能會讓醫藥和生命科學領域產生怎樣的變革？
我們不妨以百圖生科推出的，世界首個AI大模型驅動的AI生成蛋白平臺AIGP（AI Generated Protein）為例來了解一番。
AIGP背後所依靠的，是一個千億參數的跨模態生命科學大模型，通過「挖掘公開資料和獨特自產資料」、「跨模態預訓練和科學計算」，以及「蛋白質讀寫系統和細胞讀寫系統」，三大步驟實現對蛋白質空間及生命體的建模。
如此大模型能力之下，百圖生科便具備了一系列給定Protein（抗原），設計與之以特定方式結合的Protein（抗體）的能力。
也因此參與到了一系列前沿藥物的開發，包括高性能免疫調控彈頭設計、難成藥靶點精準設計、定表位抗體彈頭設計、可溶性TCR設計等。
除此之外，百圖生科也具備對給定細胞/細胞組合，發現調控細胞的有效蛋白靶點/組合，並繼而快速設計調控蛋白的能力。
這就為多種疾病的靶點發現、耐藥/不響應患者改善、靶點科學線索轉化帶來新的可能。
不過有一說一，百圖生科的例子也是隻是AI之於醫藥、生命科學領域變革的一隅。
但今年生物醫學領域的著名獎項（加拿大蓋爾德納獎）史無前例地頒給了人工智慧科學家、DeepMind創始人Demis Hassabis等人。
這也從側面反映了生命科學、醫藥領域對於AI的認可，以及更多的期待。
如果您對本文涉及的基於Habana Gaudi與英特爾®️至強®️可擴展處理器對AlphaFold2進行端到端最佳化的技術細節感興趣，如果您也想了解百圖生科在AIGP領域的最新進展，英特爾《至強實戰課》之《AI驅動的生命科學與醫藥創新》將為您帶來更加全面且詳細的真人講解，歡迎大家註冊收看～
以及想要了解更多第四代至強®️可擴展平臺對AlphaFold2的最佳化方案，請點選【閱讀原文】。
