概述
由DeepMind在2021年发布的AlphaFold2，凭借自身在蛋白质结构预测上的高可信度，以及远优于传统实验方法的效率和成本表现，树起了一座“AI for Science”的全新里程碑。它不仅在生命科学领域掀起了颠覆式的革新，也成为了AI在生物学、医学和药学等领域落地的核心发力点。
随着AlphaFold2项目在产、学、研各细分领域中的启动与落地，其技术管线对于推理的高通量和高性能的需求也是与日剧增。一直活跃在“AI for Science”创新前沿的英特尔结合自身优势，以内置AI加速能力的产品技术，特别是第三代和目前最新的第四代至强® 可扩展平台为硬件基座，对AlphaFold2实施了端到端的高通量优化，并在实践中实现了比专用AI加速芯片更为出色的表现。其中，第三代英特尔® 至强® 可扩展处理器上的优化，可使通量提升至优化前的23.11倍1，而第四代英特尔® 至强® 可扩展处理器则可在此基础上使通量再获高达3.02倍的提升2。
如此显著的优化成效，基于英特尔® 架构的软硬件协作功不可没：
- 硬件支撑：英特尔® 至强® 可扩展平台的核心产品和技术特性，包括第三代和第四代英特尔® 至强® 可扩展处理器在算力输出上的越来越出色的表现，及其内置的AI加速引擎，如英特尔® 高级矢量扩展512（英特尔® AVX-512）和英特尔® 高级矩阵扩展（Intel® Advanced Matrix Extensions，英特尔® AMX）等技术带来的预处理、推理计算优化，以及高带宽内存 (High Bandwidth Memory, HBM)、全新DDR5内存等特性对张量吞吐、数据访存通量的明显提升；
- 软件加成：软件是充分利用或释放硬件加速潜能的“钥匙”，例如在模型推理阶段，序列长度为n的情况下，推理时间复杂度为O(n2)，此时原始AlphaFold2在CPU上的推理时长是难以接受的。英特尔为此采取了一系列软件调优举措，包括对注意力模块（attention unit）开展大张量切分（tensor slicing），以及使用英特尔® oneAPI工具套件实施算子融合等优化方法，解决了AlphaFold2在CPU平台上面临的计算效率低和处理器利用率不足等难题，同时也缓解了调优方案执行各环节中面临的内存瓶颈等问题。
本文的核心任务，就是要介绍上述基于英特尔® 架构、致力于在CPU平台上加速AI应用的软硬件产品技术组合在AlphaFold2端到端优化中扮演的关键角色，并详细分享对它们进行配置、调优以求持续提升AlphaFold2应用性能表现的核心经验和技巧，从而为所有计划开展或正在推进类似探索、实践的合作伙伴及最终用户们提供一些关键的参考和建议，让整个产业界能够进一步加速相关应用的落地并尽可能提升其收益。
图一 基于英特尔® 至强® 可扩展平台的AlphaFold2 推理优化路线图及其实现的性能提升3
蛋白质结构解析任务繁重，AlphaFold2勇担重任
如生物学中心法则（Central Dogma）所揭示的，脱氧核糖核酸（DNA）、核糖核酸（RNA）和蛋白质（包括多肽4）之间“转录-翻译”的关系，清晰呈现了有机体内的信息传递路径，也让人们认识到：对蛋白质三维结构开展有效解析与预测，就能对有机体的构成，及其运行和变化的规律实施更深层次的诠释和探究，进而可为生物学、医学、药学乃至农业、畜牧业等行业和领域的未来研究与发展提供高质量的生物学假设。
虽然许多基于传统实验方法的蛋白质结构解析工具，包括X-射线晶体衍射、冷冻电镜、核磁共振等已获普遍运用，但通过传统实验方法进行结构解析的速度，远赶不上氨基酸序列的增加速度，这就造成海量待测样品/序列可能会在实验室中等待数月乃至数年才能得到解析。以UniProtKB/Swiss-Prot数据库搜集和整理的数据为例，单从实验获得的已知蛋白序列就已高达57万条之多5。
AI技术的高速发展，则为破解上述效率问题提供了新的思路--人们开始将深度学习等方法运用于蛋白质结构预测，其中由DeepMind在2020年CASP 146上提出的AlphaFold2方案尤其令人瞩目，它以惊人的92.4分（GDT_TS分数）的表现实现了原子级别的预测精度，被认为“已可替代传统实验方法”7。
AlphaFold2端到端预测：三个阶段协作增效
与以往多是间接预测蛋白质结构的AI方法不同，AlphaFold2提供了完整的端到端蛋白质三维结构预测流程。如图三所示，其工作流程大致可分为预处理（Preprocessing）、深度学习模型推理（DL Model Inference）以及后处理（Postprocessing）三个阶段，各阶段执行的功能如下：
- 预处理
由于初始输入的氨基酸序列所含信息往往较少，因此AlphaFold2在预处理阶段会先利用已知信息（包括蛋白质序列、结构模板）来提升预测精度。包括借助一些蛋白质搜索工具在特定序列数据库中使用多序列比对（MSA）方法，以及在特定结构数据库中进行模板搜索，从而获得不同蛋白质之间的共有进化信息；
- 深度学习模型推理
在该阶段中，AlphaFold2首先会借助嵌入（Embedding）过程，将来自预处理阶段的模板MSA信息、MSA和目标构成MSA表征（MSA representation）的三维张量，同时也将模板邻接信息和额外的MSA构成邻接表征（pair representation）的三维张量，随后两种表征信息会通过一个由48个块（Block）组成的Evoformer网络进行表征融合。在这一进程中，模型将通过一种Self-Attention机制来学习蛋白质的三角几何约束信息，并让两种表征信息相互影响来使模型推理出相应的三维结构，且循环三次；
- 后处理
这一阶段，AlphaFold2将使用Amber力场分析方法对获得的三维结构参数优化，并输出最终的蛋白质三维结构。
AlphaFold2在预测精度上取得的优势，源于四点全新的设计思路：
- 在预处理阶段通过MSA方法等，将模板蛋白质结构和序列保守性信息融入预测特征；
- 在特征嵌入阶段，将保守性最高的MSA特征单独取出，压缩其余的Extra MSA，并与模板特征交互；
- 在模型推理阶段，采用独特的双轨注意力模块和深层Transformer架构，并引入循环回收机制；
- 在结构网络层引入不变点注意力（Invariant Point Attention）机制。但这也意味着AlphaFold2从执行之初，直至整个推理过程都需要面对高通量的计算压力。
五大步骤：至强® 可扩展平台助AlphaFold2实现端到端优化
随着越来越多的科研机构、实验室和企业开始借助AlphaFold2进行蛋白质结构预测，各行业和领域内的使用者也开始遇到越来越多、也渐趋严峻的挑战。例如结构预测各环节面临着庞大的计算量，使用者需要更加充分地挖掘硬件的计算潜力来提升执行效率；为缩短结构预测时间，他们还需要利用更多计算节点来构建效率更高的并行计算方案等。
基于英特尔® 至强® 可扩展平台提供的内置AI加速能力，对于运算和存储性能的均衡设计，以及对硬件和软件协同优化能力的兼顾，英特尔着手对AlphaFold2进行了端到端的全面优化，以帮助生物学等应用领域的使用者们应对以上挑战。针对AlphaFold2的设计特点，该优化方案主要聚焦在预处理和模型推理两个层面，在第三代英特尔® 至强® 可扩展处理器和第四代英特尔® 至强® 可扩展处理器上，其基本可划分为以下五个步骤。
- 第一步：预处理阶段-高通量优化
预处理阶段的高通量计算需求，使方案在执行时面临非常明显的并行计算压力。借助第三代或第四代英特尔® 至强® 可扩展处理器的多核优势及其内置的英特尔® AVX-512技术，方案能够实现针对预处理阶段的高通量优化。
如前文所述，AlphaFold2会在预处理阶段对特定序列数据库和结构数据库中的已知序列/模板信息进行搜索，包括使用jackhmmer等蛋白质搜索工具来执行MSA方法，即从数据库中抽取和输入与氨基酸序列相近的序列并进行对齐的过程，其目的是找出同源的序列/模板组成表征信息来为后续推理过程提供输入，由此提高预测精度。
这一过程中，计算平台需要执行大量的向量/矩阵运算。以模板搜索为例，其本质为计算两个隐马尔可夫模型（Hidden Markov Model，HMM）间的距离。当输入的氨基酸序列很长（例如执行中输入长度达数百的氨基酸序列）且需并行执行大量实例时，如果无法让处理器的算力“火力全开”去提升平台的并行计算效率，那么整个预处理过程的效率就会变得乏善可陈。
在实践中，由英特尔® 至强® 可扩展处理器内置的英特尔® AVX-512以及使用英特尔® C++ 编译器 配置 jemalloc的优化方法，为方案提供了更进一步的性能调优空间。其中，针对序列/模板搜索所需的大量向量/矩阵运算需求，英特尔® AVX-512技术，能以显著的高位宽优势（最大可提供512位向量计算能力）来提升计算过程中的向量化并行程度，从而有效提升向量/矩阵运算效率。
- 第二步：模型推理阶段-将深度学习模型迁移至面向英特尔® 架构优化的PyTorch
原始版本的AlphaFold2是基于DeepMind的JAX和haiku-API做的网络实现，但目前JAX上还没有面向英特尔® 架构平台的优化工具。而PyTorch拥有良好的动态图纠错方法，与haiku-API有着相似的风格，并可以采用面向PyTorch的英特尔® 扩展优化框架（Intel® Extension for PyTorch，IPEX，可由英特尔® oneAPI AI工具套件提供）。为实现更好的优化效果，方案选择将深度学习模型迁移至面向英特尔® 架构优化的PyTorch，并最终逐模块地从JAX/haiku上完成了代码迁移。
- 第三步：模型推理阶段-PyTorch JIT
为提高模型的推理速度，便于利用IPEX的算子融合等加速手段，优化方案中还对迁移后的代码进行了一系列的API改造，在不改变网络拓扑的前提下，引入PyTorch Just-In-Time (JIT) 图编译技术，将网络最终转化为静态图。
- 第四步：模型推理阶段-切分Attention模块和算子融合
AlphaFold2的嵌入过程是构成MSA表征张量和邻接表征张量来作为Evoformer网络输入的关键步骤。从其算法设计可以获知，其注意力模块中包含了大量的偏移量（bias）计算。
这种偏移量计算是通过张量间的矩阵运算来完成的，因此运算过程中会伴随张量的扩张。当张量达到一定规模后，扩张过程对内存容量的需求就会变得巨大。这就使AlphaFold2在嵌入过程中面临两个问题：一方面是巨大的内存峰值压力，其需求量会使内存资源在短时间耗尽，尤其是内存峰值在相互叠加之后，进而可能造成推理任务的失败；另一方面，大张量运算所需的海量内存也会带来不可忽略的内存分配过程，从而增加执行耗时。
为此，英特尔提出了”对注意力模块进行大张量切分”的优化思路，即，将大张量切分为多个较小的张量，来降低扩张中的内存需求。例如将上述“5120 x 1 x 1 x 64”的张量切分为“320 x 1 x 1 x 64”后，其扩张所需的内存就由930MB降至59.69MB，仅为未进行张量切分时的6.4%左右，有效消减了内存峰值压力。
同时，英特尔也发现利用PyTorch自带的Profiler对AlphaFold2的Evoformer网络进行算子跟踪分析时， Einsum和Add这两种算子占用了大部分的算力资源。因此，英特尔就考虑使用IPEX（建议版本为IPEX-2.0.100 + cpu或更高）提供的算子融合能力来实现上述两种计算过程的融合。
传统的深度学习计算过程都是逐一操作：例如Einsum计算过程结束后，函数返回值需要在Python进程中建立一个临时缓存，然后通过调用Add算子，再次进入oneDNN完成第二个函数的运算，这中间来回折返的过程时间消耗不可忽略。如图五所示，算子融合带来的优势就在于，在前一操作结束后可以马上执行后一操作，节省了中间建立临时缓存数据结构的时间。同时从时间轴上不难看出，经过融合后，两个连续的算子合并为一个，用时也显著缩短。
- 第五步：模型推理阶段-破解多实例运算过程中的计算和内存瓶颈
为了让推理性能在多实例进程中获得更接近线性的增长表现，优化方案也借助英特尔® 至强® 可扩展平台提供的高效且更为均衡的计算和存储优势实施了有针对性的优化。
这些优势包括借助基于NUMA架构的核心绑定技术，来充分挖掘至强® 可扩展处理器的多核心优势。这一技术可对处理器节点以及访问本地内存进程予以精确控制，让每个推理工作负载都能稳定地在同一组核心上执行，并优先访问对应的近端内存，从而提供更优、也更稳定的并行算力输出。
四剂“强芯针”：新一代至强® 可扩展处理器为AlphaFold2推理带来 “推背感”
英特尔硬件与加速器的持续更新，正为AlphaFold2带来更大的优化空间。随着全新第四代英特尔® 至强® 可扩展处理器被逐步引入AlphaFold2的工作环境，其不仅提供了更强的基础算力，也带来多项针对AI工作负载的优化加速技术。在基于第四代英特尔® 至强® 可扩展处理器的优化工作中，英特尔基于以下四剂“强芯针”，让AlphaFold2的推理性能获得更为显著的提升。
1. 借助TPP技术，降低推理过程中的内存消耗
在深度学习系统开发中，诸如算子（Operators）、算法概念（Algorithmic Concepts）以及计算模式（Computational Motifs）等编程范式（Programming Paradigm）通常会面向特定平台进行调优，这会对系统的构建便利性、性能调优以及可移植性造成障碍。为此，张量计算原语（Tensor Processing Primitives，TPP）技术是在2D张量上定义了一组低层级的基本算子，通过有效、且可移植的张量级算子来应对这一问题。TPP可被看成是一种虚拟的张量指令集架构，能将英特尔® AVX-512等物理指令集予以抽象，并生成经优化的平台代码。
根据自身软硬件特性，英特尔面向PyTorch对TPP进行了扩展。面向PyTorch的英特尔® TPP扩展（Intel® Tensor Processing Primitives Extension for PyTorch）不仅能让开发者直接使用TPP调用英特尔® oneAPI等库来生成优化代码，也可利用面向PyTorch的TPP作为构建块来表示底层张量计算。
在优化方案中引入TPP技术能让AlphaFold2在通用矩阵乘法（General Matrix Multiplication，GEMM）等计算中获得优势，降低内存消耗并更好地利用全新第四代英特尔® 至强® 可扩展处理器所具备的大容量末级缓存优势。例如在Evoformer模块中需要进行大量的狭长矩阵乘法运算。对于在处理器上执行的矩阵乘法计算，一般会采用两种重要的优化方式：
- 以单指令多数据（Single Instruction Multiple Data，SIMD）方式处理数据；
- 优化内存访问模式，提升缓存命中率来提高数值计算和访存效率。
通过引入面向PyTorch的英特尔® TPP扩展，英特尔在AlphaFold2实现了以上两种优化。如图六所示，一方面由libxsmm(小矩阵乘法函数库)构建起来的TPP能借助BRGEMM（Batch Reduce General Matrix Multiplication）最大化利用第四代英特尔® 至强® 可扩展处理器所内置的SIMD运算单元，同时小矩阵乘法也能有效提高缓存命中率，使处理器的大容量末级缓存优势在计算过程中获得更充分的利用。TPP的引入，令狭长矩阵乘法的空间复杂度从O(n^2)降为O(n)，这使得运算过程中所需的内存峰值大幅降低，有效缓解了长序列蛋白质结构预测工作中面临的“序列长度天花板”问题。
2. 支持DDR5内存与大容量缓存带来张量吞吐提升
通过对算法架构的解析可知，AlphaFold2中大量的矩阵运算过程都需要内存予以支撑，因此内存性能是影响AlphaFold2性能的重要因素。而随着预测序列长度的增加，计算中所需的内存也会成倍增加，内存性能，尤其是内存带宽对系统整体性能的影响也会更为明显。
与此同时，更优的缓存策略也能让AlphaFold2进一步发挥潜能。由于张量间的矩阵运算会涉及大量的内存数据访存，而更靠近处理器运算单元末级缓存在延迟性能上比内存高出一个数量级。因此在复杂的矩阵运算中，更多的热数据通过末级缓存而非内存来访存，可以带来显著的性能提升。
第四代英特尔® 至强® 可扩展处理器对DDR5内存的支持，以及所具备的大容量末级缓存，为张量吞吐量的提升提供了更佳途径。新一代DDR5内存不仅频率更高、工作电压更低，还具有远超DDR4内存的带宽速度。与DDR4内存25.6GBps（3,200MHz）的带宽相比，DDR5内存带宽达到了38.4GBps（4,800MHz）以上，提升幅度超过了50%。同时，新处理器的末级缓存也由上一代的最高60MB提升至本代的最高112.5MB，提升幅度达到了87.5%8。性能更高的内存与容量更大的末级缓存，使AlphaFold2推理过程中关键的张量吞吐获得了显著提升。
3. 全新英特尔® AMX_BF16在保证精度的前提下加速推理过程
第四代英特尔® 至强® 可扩展处理器面向深度学习应用推出的“杀手锏”之一就是其创新的内置AI加速引擎，即英特尔® AMX。作为矩阵相关的加速器，英特尔® AMX能显著加速基于CPU平台的深度学习推理和训练，提升AI整体性能。英特尔® AMX对INT8、BF16等低精度数据类型都有着良好的支持（通过AMX_INT8、AMX_BF16等不同指令集执行操作），其中BF16数据类型在精度上有着不逊于FP32数据类型的表现。
针对AlphaFold2推理过程所需的大量矩阵运算，AMX_BF16能在保持较高精度的同时，提高计算速度并减少存储空间。究其原因，是因为英特尔® AMX在解决矩阵乘法问题时，直接采用了分块矩阵乘法的方式。其内部所定义的Tile矩阵乘法 (Tile Matrix Multiply Unit，TMUL)加速模块，能直接对矩阵寄存器中的数据实施矩阵运算操作，由此运算效率可得到大幅提升。实践数据表明，AlphaFold2在推理过程中使用AMX_BF16后，推理时间可缩短数倍之多。
4. 高带宽内存HBM2e增加访存通量
与第四代英特尔® 至强® 可扩展处理器一同发布、采用了相同微架构的英特尔® 至强® CPU Max系列中，还加入了对HBM的支持，这也能让运行在其上的AlphaFold2推理负载更进一步。作为一种采用3D堆叠技术的全新内存技术，HBM能为AI应用场景所需的各类计算负载提供更大的内存带宽支持。
- 每个英特尔® 至强® CPU Max 系列都拥有4个基于第二代增强型高带宽内存（HBM2e）的堆栈，总容量为64GB（每个堆栈的容量为16GB）；
- 由于能同时访问多个DRAM芯片，因此HBM在带宽方面相较DDR技术更具优势，其中HBM2e可提供高达1TB/s的带宽；
- HBM内存可根据工作负载特性，以“HBM Only”、“HBM Flat”以及“HBM Cache”三种不同的模式，通过灵活的配置与DDR5内存一起协同工作。
在实践中，HBM2e内存被证明能有效缓解AlphaFold2推理负载对高带宽内存的需求，带来访存通量的大幅提升，从而降低整体的推理时长。
多个优化步骤实施后的总体性能表现
基于英特尔® 至强® 可扩展平台开展的AlphaFold2端到端优化，包括一系列并行计算能力优化举措的引入，使得整个AlphaFold2端到端处理过程的性能获得了质的提升。如图七所示，在基于第三代英特尔® 至强® 可扩展处理器的优化流程中，每个优化步骤获得的提升累积后，最后相比优化前通量提升可达23.11倍9。
图七 基于第三代英特尔® 至强® 可扩展处理器的优化流程中多种优化措施带来的累计性能提升10
而来自第四代英特尔® 至强® 可扩展处理器的优化加持，则使AlphaFold2的端到端通量获得再进一步的提升，如图八所示，与第三代英特尔® 至强® 可扩展处理器相比，融合AMX_BF16、HBM内存等技术的第四代英特尔® 至强® 可扩展平台能实现高达3.02倍的多实例通量提升11。
图八 第四代英特尔® 至强® 可扩展处理器带来多实例通量提升11
在探索和验证上述端到端AlphaFold2优化方案、步骤和经验的过程中，英特尔扮演的角色并非“独行侠”，而是与同在寻求相关解决方案的，专攻医药和生命科学研究和创新的产、学、研领域用户及合作伙伴们积极开展了广泛及深入的协作，这些协作起到了博采众长的效果，也为方案的普适性带来了助益。
同样，在优化方案基本定型，并展现了显著的通量提升效果以及能够担起更长序列蛋白质结构预测重任的能力后，众多合作伙伴与用户也第一时间参考和借鉴了方案中的方法、经验与技巧，并结合自身特定的环境、应用现状和需求，开展了实战验证和更进一步的探索。
总结与展望
得益于AI技术的高速发展和演进，它与科学前沿研究的结合正在快速地改变世界并造福人们的生活。始终走在AI应用创新与落地一线的英特尔，也在这一过程中借助至强® 可扩展平台，包括硬件层面的第三代英特尔® 至强® 可扩展处理器和第四代英特尔® 至强® 可扩展处理器，以及其软件层面的英特尔® oneAPI 工具套件等，基于这些软硬件之间的无缝组合与高效协作，以及多样化的AI 优化方法，为AlphaFold2提供了持续改良的端到端高通量计算优化方案。
面向未来，英特尔还将继续携手科学前沿领域的合作伙伴，推进更多英特尔产品、技术与AlphaFold2等新技术开展交互与融合，在更多层面助力和加速“AI for Science”的技术创新，让AI应用为各类前沿科学研究和探索带来更多加速、助力与收获。