预测超长蛋白质这事，CPU 赢了

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 的优化方案，请点击【阅读原文】。

— 完 —

点这里关注我，记得标星哦～

一键三连「分享」、「点赞」和「在看」

科技前沿进展日日相见 ~