-
构建平台时的成本挑战:
AlphaFold2任务离不开高质量强算力硬件基础设施的支持,虽然GPU一直在多种AI任务中担纲主力,但其昂贵的市场价格以及难以获取的状况,也让用户不得不面临高昂的构建成本。同时用户在既有IT基础设施之外构建额外的AI算力平台,也容易带来算力调度不易、重复投资等问题;
-
高通量计算和大规模推理压力:
AlphaFold2在架构上的多项创新设计,例如在结构网络层引入不变点注意力(Invariant
Point Attention)机制等,在带来更优预测效果之余,也使整个方案从预处理阶段至推理任务都需要面对高通量的计算压力。同时为加速成果的问世,用户也亟需提升推理效率来缩短预测时长;
-
大内存需求问题:
在面向不同蛋白质的结构预测工作中,序列长度越长,张量运算规模与推理计算复杂度就越大,再结合更多序列的并行计算,其所需的内存容量也就越高。一旦承载平台的内存或显存容量不足,就可能限制更长序列的预测工作。
为让AlphaFold2的应用能够真正普及,助力到更多科研机构、高校实验室以及生物制药企业的创新事业中,
英特尔启动了针对自身架构的AlphaFold2方案构建及优化工作。
它以至强
®
处理器,例如至强
®
CPU Max系列、第四代至强
®
可扩展处理器等CPU产品作为算力基座,并借助它们内置的英特尔
®
高级矩阵扩展(Intel
®
Advanced Matrix Extensions,英特尔
®
AMX)技术、所集成的高带宽内存(High Bandwidth Memory, HBM),并配以软硬结合的端到端优化方案,为用户打造出了一条通途,让他们能够在更易于获取和使用的CPU平台上,以更优的性价比,更加可用且便捷的方式部署和应用AlphaFold2,从而能让更多生物、制药领域的用户搭上AlphaFold2的“快车”。
英特尔这一方案不仅能满足在多实例、长序列预测场景下的高通量计算性能和效率要求,也能以大内存优势帮助用户实现对更长蛋白质序列的预测。此外,该方案还进一步地为用户提供了面向蛋白质复合体预测(AlphaFold2 Multimer)的实现方式,将单一蛋白质预测进一步扩展到复合体范围,使之具有更强的可用性。
如图二所示,AlphaFold2的基本架构分为预处理、深度学习模型推理以及后处理三个主要部分,因此其
计算负载是一种混合计算负载。
而传统上,收集数据、数据预处理等所需的通用计算任务本来就适于部署在CPU平台上,并能借助丰富的优化工具和方法提升性能。而随着更多AI加速技术“嵌入”CPU,以及围绕CPU平台的AI加速生态逐步完善,已有充分数据显示,
以英特尔
®
至强
®
可扩展处理器为代表的CPU算力平台可完全满足AlphaFold2推理所需的算力。
图二 AlphaFold2基本架构
而在满足AlphaFold2任务所需的性能之余,基于CPU平台构建的方案还能为用户带来两方面的收益:
一方面
是能让用户在既有硬件平台上即可完成算力平台的搭建,并方便地将算力在其它计算任务与AlphaFold2之间调度,在充分调度空闲算力之余也能避免额外建设费用;
另一方面,
方案也帮助用户避免了GPU方案所面临的成本、采购以及算力调度等方面的难题,也对降低TCO有益。
同时,英特尔
®
架构CPU平台庞大的市场占有率、广泛的部署基础,及其基于高效且稳定的x86架构所培育的丰富而友好的生态环境(包括对各类主流软件框架的广泛支持、借助插件对开源框架的优化,以及一系列用于提升AI性能的开源工具链等),也让CPU算力平台的搭建和获取都变得便捷且平滑。无论是采用私有化部署方式,还是依托公有云服务,AlphaFold2任务在任何时候,包括高密度使用场景下都能获得稳定且高效的CPU平台算力支持。
在硬件上,
基于英特尔
®
架构的AlphaFold2解决方案提供了英特尔
®
至强
®
CPU Max系列、第四代至强
®
可扩展处理器等,来应对AlphaFold2面临的高通量计算与推理压力。就以前者为例,它所采用的全新微架构、更多的内核(最多达56个性能核)可为各类计算任务提供强有力的性能支撑。它在内存和输入/输出(I/O)子系统性能上有着显著的增强,提供了对DDR5内存(内存带宽可达 38.4GBps(4,800MHz)以上)的支持,还结合大容量末级缓存使AlphaFold2推理过程中关键的张量吞吐获得了大幅提升。
图三 英特尔
®
至强
®
CPU Max系列处理器
这些处理器内置的英特尔
®
AMX加速引擎,则能显著加速基于CPU平台的深度学习推理和训练,提升AI整体性能。英特尔
®
AMX在解决矩阵乘法问题时,直接采用了分块矩阵乘法的方式。其内部所定义的Tile矩阵乘法(Tile Matrix Multiply Unit, TMUL)加速模块,能直接对矩阵寄存器中的数据实施矩阵运算操作,由此可大幅提升运算效率。针对AlphaFold2推理任务中所需的大量矩阵运算操作,AMX_BF16能在保持较高精度的同时,提高计算速度并减少存储空间。
在软件优化层面,
英特尔一方面为AlphaFold2提供了面向英特尔
®
架构优化的PyTorch和PyTorch Just-In-Time(JIT)图编译技术。用户通过将深度学习模型迁移至面向英特尔
®
架构优化的PyTorch能实现更好的动态图纠错方法,而PyTorch Just-In-Time(JIT)图编译技术则能将网络转化为静态图,有利于推理速度的提高。
另一方面,针对AlphaFold2预处理阶段的高通量计算压力,英特尔也借助至强
®
系列处理器的多核优势及内置的英特尔
®
高级矢量扩展512(英特尔
®
AVX - 512)技术,以显著的高位宽优势来提升计算过程中的向量化并行程度,从而有效提升向量/矩阵运算效率。
在解决方案中,英特尔
®
至强
®
CPU Max系列处理器还有一项优势表现突出,那就是它集成的HBM内存,HBM能为AI应用场景所需的各类计算负载提供更大的内存容量、更高的带宽以及更低的功耗。
每个至强
®
CPU Max系列处理器都拥有总容量高达64GB的HBM内存,且具有高达460GB/s的带宽。
HBM内存可根据工作负载特性,以
“单一HBM模式”、“HBM平坦模式”
以及
“HBM缓存模式”
三种不同的模式,通过灵活的配置与DDR5内存一起协同工作。在面向AlphaFold2的实践中,HBM内存能有效缓解推理任务中大张量运算带来的海量内存需求,实现对更长蛋白质序列的预测。并以高带宽特性带来显著的访存通量提升,从而降低整体的推理时长。
与此同步,英特尔还在解决方案中
提供了多种降低内存的软件优化方法。
例如英特尔面向PyTorch对张量计算原语(Tensor Processing Primitives,TPP)技术进行了扩展,不仅能让开发者直接使用TPP调用英特尔
®
oneAPI等库来生成优化代码,也可利用面向PyTorch的TPP作为构建块来表示底层张量计算。这可
帮助AlphaFold2在通用矩阵乘法(General Matrix Multiplication,GEMM)等计算中获得优势,使得运算过程中所需的内存峰值大幅降低。
其次,切分Attention模块和算子融合的推理优化方案,也可通过将大张量切分为多个较小张量的方案来降低AlphaFold2执行时的张量扩张带来的巨大内存峰值方面的压力。同时方案也借助面向PyTorch的英特尔
®
扩展优化框架(Intel
®
Extensions for PyTorch,IPEX)对算力资源消耗较大的热点算子进行融合,提升计算效率。
此外,英特尔
®
AMX对INT8、BF16等低精度数据格式都有着良好的支持(通过AMX_INT8、AMX_BF16等不同指令集执行操作),这在
提升推理效率的同时也能显著降低推理时的内存消耗。
图四 热点算子与融合效果
由于不同的蛋白质结构探索范围所需耗费的内存资源有比较大的差别 , 英特尔这一解决方案提供的各项为内存增量降耗的硬件方案与优化措施,能有效缓解长序列蛋白质结构预测工作中面临的“序列长度天花板”问题。
4、面向AlphaFold2 Multimer模式的方案实现
与预测单个蛋白质三维结构的AlphaFold2 Monomer(即Al-phaFold2)相比,
AlphaFold2 Multimer能对多个蛋白质分子之间的相互作用及所形成的复合体结构进行预测。
这对研究和探明蛋白质之间的相互作用、信号传导和生物通路等课题有非常重要的意义。借助AlphaFold2 Multimer,研究者能够更精确地了解蛋白质复合体的结构和动态行为,从而让生物、新药的研究变得更具效率优势。
图五 面向AlphaFold2 Multimer模式的方案实现
基于英特尔
®
架构的AlphaFold2解决方案同样也面向Alpha-Fold2 Multimer的管线结构进行了优化与验证,虽然其管线结构已根据蛋白质复合体结构预测的需求进行了调整,但英特尔AlphaFold2上的优化方案,在被用于AlphaFold2 Multimer时同样有效,同样能
为之提供充沛且更具性价比的算力支持。
目前,基于AlphaFold2 Multimer的英特尔解决方案也已与多个用户展开实践合作。