搅翻计算生物界的AlphaFold2一开源，各种加速方案就争相涌现。
妹想到啊，现在居然有了个CPU的推理优化版本，不用GPU，效果也出人意料的好——
端到端的通量足足提升到原来的23.11倍。
△ 第三代英特尔? 至强? 可扩展处理器
其整体流程，就是先在预处理阶段对模型进行高通量优化，然后将模型迁移到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%左右。
没有了大内存之需，也就不需要进行内存分配了，因此，张量切分后推理速度也上来了。
比如从下图我们就可以看到，注意力模块的效率在切片前后有着非常明显的差别。
△ 算子 Einsum+Add 融合效果图
至此，经过以上几个步骤的优化，AlphaFold2在CPU上的总体性能已经得到了大约5倍的提升。
还差最后一步：
05、多实例并行优化
在这一步，英特尔先利用至强可扩展平台上基于NUMA架构的核心绑定技术，让每个推理工作负载都能稳定地在同一组核心上执行，并优先访问对应的近端内存，从而提供更优、也更稳定的并行算力输出。
然后引入英特尔? MPI库帮助模型在多实例并行推理计算时实现更优的时延、带宽和可扩展性。
但这些动作还不足以破解限制AlphaFold2发挥潜能的一个重要因素：内存瓶颈。
众所周知，在面向不同蛋白质的结构测序工作中，序列长度越长，推理计算复杂度就越大。
而在我们对模型进行了并行计算能力的优化后，更多计算实例的加入还会进一步凸显这一问题。
英特尔用“星际探索”这一比喻对这种现象做了非常形象的说明。
△ 英特尔? 傲腾? 持久内存
提到傲腾持久内存，我们并不陌生。今年6月在《用CPU方案打破内存墙Paypal堆傲腾扩容量，漏查欺诈交易量可降至1/30》一文中已经小结了它对于AI应用的关键作用，即提供更大容量的内存子系统来满足那些内存敏感型AI应用将更多数据贴近算力的需求。
目前最新一代的傲腾持久内存200系列，可以在提供接近主流DRAM内存性能的基础上，实现每路高达4TB的容量，或者说，与DRAM内存组合时可提供每路高达6TB的内存总容量。
有了它，我们甚至能够实现10000aa序列长度的蛋白结构预测。
用了它，英特尔这个方案的优化就基本完成，模型的总体性能也可在之前优化步骤的基础上再次得到4.56倍的提升。
最后，附上一份用于这个优化方案的英特尔官方推荐配置。
蛋白质结构预测，只是AI for Science的序幕
不久前，引爆了AI for Science的AlphaFold2又公开了新进展——
它已经成功预测出包括植物、细菌、真菌在内的100万个物种的2.14亿个蛋白质结构，并将数据集对外开源。
而这些，都还只是AI for Science的序幕。
在蛋白质结构预测、生物计算、药物开发之外，AI在物理、天文、化学等领域也开始逐渐展露头角。
前沿领域，今年《Nature》上的一项“改写物理教科书”的研究，正是通过AI开展的。
欧洲核子研究组织（CERN）的科学家利用机器学习，发现了质子内部存在5个夸克的有力证据，这一成果颠覆了一直以来质子只有三个夸克的理论。
应用落地层面，前面提到的晶泰科技已经构建了一套智能计算、自动化实验、专家经验结合的三位一体的研发模式，提供一站式小分子药物发现、大分子药物发现，药物固体形态研发，以及化学合成服务。
晶泰科技目前已建成数千平的自动化实验室，与智能算法“干湿结合”，形成实验数据与算法预测间的交互闭环，保证AI算法的产业落地和交付能力。
— 完 —
文章知识点与官方知识档案匹配，可进一步学习相关知识算法技能树首页概览33917 人正在系统学习中
声明：本站部分文章及图片源自用户投稿，如本站任何资料有侵权请您尽早请联系jinwei@zod.com.cn进行处理,非常感谢！