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，正在大步邁進新時代。

1、英特爾自己刷新自己

事實上，此次備受關注的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%。

圖片

2、醫藥和生命科學領域，AI還有何作為？

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

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

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

（1）人類自然語言大模型：從26個字母，到詞/句/段。

（2）生命自然語言大模型：從21個氨基酸字母，到蛋白質/細胞/生命體。

圖片

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

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

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

圖片

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

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

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

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

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

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