CPU 也能速刷 AlphaFold2？英特尔：请收下这份 23 倍通量优化指南

搅翻计算生物界的 AlphaFold2 一开源，各种加速方案就争相涌现。

妹想到啊，现在居然有了个 CPU 的推理优化版本，不用 GPU，效果也出人意料的好——

端到端的通量足足提升到原来的 23.11 倍。

高达 23.11 倍的提升不是一蹴而就，但依然惊艳

AlphaFold2 逐步通量提升*² (每日样本数量，序列长度=765)

换个更具体的数据，它直接让 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 倍。

至此，经过以上几个步骤的优化，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 倍的提升。

最后，附上一份用于这个优化方案的英特尔官方推荐配置。

如果你想看更多细节中的细节，也可以访问英特尔公开分享的白皮书《通量提升达 23.11 倍！至强可扩展平台助力 AlphaFold2 端到端优化》。

或者观看英特尔联合国际学术期刊《Science》推出的 “架构师成长计划” 第二季第八期课程《AI 驱动的生命科学创新范式之变》。

其中不但有量子位总编辑李根、晶泰科技首席研发科学家杨明俊博士和英特尔人工智能架构师杨威博士围绕这个主题的精彩讨论，还有晶泰科技在 AI 制药领域的领先实践分享，以及英特尔这个 AlphaFold2 优化方案更为细致且可视化的呈现。

光说不练假把式，英特尔这个方案看起来很美，但要真正实现落地才能让人信服。

事实上，在对这个方案进行摸索和开发的过程中，英特尔与相关领域合作伙伴或用户的协作与交流就一直没停过，不但吸收了各方的经验，实现了博采众长和互相借鉴、启发。

在成型后，也第一时间广泛分享给到伙伴和用户们，让他们能够根据自身特定的环境、应用状况和需求，开展实战验证和推进更进一步的探索。

例如国内某高校就曾尝试在数百台基于至强可扩展处理器的服务器上，采用该方案提供的经验和方法来进行测试，并取得了一举两得的结果——

其顺利实践了短序列高通量的、面向蛋白质组学级别的批量化结构预测，既降低了蛋白质组的 AlphaFold2 预测成本，又提高了集群推理的总通量。

上文提及的国内明星 AI 制药公司晶泰科技，也在自身的研发中，将自主研发的 AI 算法与 AlphaFold2 结合，从而验证靶点、精确解析活性构象，为后续的药物发现打下良好基础。

通过充分利用 CPU 的 TB 级内存支持，在公有云上部署英特尔版 AlphaFold2 优化方案后，科学家可以实现针对短序列的单节点高通量推理优化，从而加快蛋白组学结构分析进程，并预测序列长度超过 4700aa 的蛋白质序列。

这不光能拓展 AlphaFold2 在研发中的探索范围，也能以更低的成本，让更高精度的算法工具作用于更早的研发环节，进一步加速药物发现。

不久前，引爆了 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 作为科研探索的生产力、助手，更多突破人类想象的科研成果将会涌现……

光是想想就非常期待了~