专栏名称: 英特尔中国
关注我们,了解潮流IT资讯!
目录
相关文章推荐
嗨推  ·  私域直播,危险了呢! ·  2 天前  
安利云学堂  ·  涟漪丨快乐跳操+营养打底 ... ·  3 天前  
跟宇宙结婚  ·  日常念叨:上饿了么搜【跟宇宙结婚】领红包哟 ·  3 天前  
51好读  ›  专栏  ›  英特尔中国

基于英特尔® 架构的AlphaFold2解决方案

英特尔中国  · 公众号  ·  · 2024-09-06 18:35

正文

概 述



对蛋白质三维结构的预测解析,已被广泛用于医学、制药、生物等领域的研究与创新。AlphaFold2的到来,及其为这一应用领域带来的令人惊艳的预测精度与效率,则是让它受到众多科研机构、高校实验室和生物、医药企业的关注,成为近年来“AI for Science”领域最具明星效应的创新技术。


但AlphaFold2的落地,一直都有较高的部署和应用门槛,如面临着构建成本、高通量计算和大规模推理以及大内存需求等挑战。 为了让这项技术能真正普惠相关行业和企业,一直活跃在“AI for Science”创新前沿的英特尔结合自身优势,以内置AI加速能力的英特尔 ® 至强 ® 处理器为硬件基座,并结合端到端优化方案来应对上述挑战。 通过一系列优化举措,AlphaFold2不仅在英特尔 ® 架构CPU平台上实现了比GPU平台更为出色的性能表现,同时对内存的“增量降耗”措施,也让用户能够突破性地在该平台上开展更长蛋白质序列的预测。而CPU平台在构建成本、易获取性以及生态支持方面的优势,也有助于更多用户跨越AlphaFold2的落地门槛。


AlphaFold2正为产业创新拓展更多可能性



得益于AI技术的高速发展,基于深度学习的蛋白质结构预测方法近年来异军突起,并卓有成效地解决了传统的、基于实验方法的蛋白质结构解析工具,包括X-射线晶体衍射、冷冻电镜、核磁共振等在解析速度上的短板。这其中表现最令人刮目相看的,当属DeepMind提出的AlphaFold系列方案,例如AlphaFold2就曾以惊人的92.4分(GDT_TS分数)的表现实现了原子级别的预测精度,被认为“已可替代传统实验方法” 1


图一 AlphaFold2已成为AI制药领域重要的基础架构


在一战成名后,越来越多的科研机构、高校实验室和生物、医药企业都借助AlphaFold2方案开展蛋白质结构预测,并融入基于AI的新药研发、新诊疗方法探索、生物学新研究以及生物信息学、合成生物学等一系列科研和创新方向中。例如在一项研究中,科学家利用AlphaFold2为sigma-2和5-HT2A两种蛋白结构预测出数以亿计的潜在药物组合,并有着相当高的成功率,能为阿尔茨海默病、精神分裂症等疾病药物的研发提供更高效的途径 2


AlphaFold2在落地与实践中面临的挑战



随着更多新药物研制、诊疗新方法探索正借助AlphaFold2来实现加速,最终用户们在构建和使用AlphaFold2的实战中,也开始面临更多和更为严峻的挑战,包括:


  • 构建平台时的成本挑战: AlphaFold2任务离不开高质量强算力硬件基础设施的支持,虽然GPU一直在多种AI任务中担纲主力,但其昂贵的市场价格以及难以获取的状况,也让用户不得不面临高昂的构建成本。同时用户在既有IT基础设施之外构建额外的AI算力平台,也容易带来算力调度不易、重复投资等问题;


  • 高通量计算和大规模推理压力: AlphaFold2在架构上的多项创新设计,例如在结构网络层引入不变点注意力(Invariant Point Attention)机制等,在带来更优预测效果之余,也使整个方案从预处理阶段至推理任务都需要面对高通量的计算压力。同时为加速成果的问世,用户也亟需提升推理效率来缩短预测时长;


  • 大内存需求问题: 在面向不同蛋白质的结构预测工作中,序列长度越长,张量运算规模与推理计算复杂度就越大,再结合更多序列的并行计算,其所需的内存容量也就越高。一旦承载平台的内存或显存容量不足,就可能限制更长序列的预测工作。

基于英特尔 ® 架构的AphaFold2解决方案



为让AlphaFold2的应用能够真正普及,助力到更多科研机构、高校实验室以及生物制药企业的创新事业中, 英特尔启动了针对自身架构的AlphaFold2方案构建及优化工作。 它以至强 ® 处理器,例如至强 ® CPU Max系列、第四代至强 ® 可扩展处理器等CPU产品作为算力基座,并借助它们内置的英特尔 ® 高级矩阵扩展(Intel ® Advanced Matrix Extensions,英特尔 ® AMX)技术、所集成的高带宽内存(High Bandwidth Memory, HBM),并配以软硬结合的端到端优化方案,为用户打造出了一条通途,让他们能够在更易于获取和使用的CPU平台上,以更优的性价比,更加可用且便捷的方式部署和应用AlphaFold2,从而能让更多生物、制药领域的用户搭上AlphaFold2的“快车”。


英特尔这一方案不仅能满足在多实例、长序列预测场景下的高通量计算性能和效率要求,也能以大内存优势帮助用户实现对更长蛋白质序列的预测。此外,该方案还进一步地为用户提供了面向蛋白质复合体预测(AlphaFold2 Multimer)的实现方式,将单一蛋白质预测进一步扩展到复合体范围,使之具有更强的可用性。


如欲了解更多优化细节与结果,请参阅:

《CPU加速Alpha-Fold2更上一层楼!第四代至强 ® 可扩展平台带来3.02倍通量提升》 https://www.intel.cn/content/www/cn/zh/customer-spotlight/cases/cpu-acceleration-alphafold2-next-level-throughput.html


1、兼具性能、经济性和可用性的CPU方案


如图二所示,AlphaFold2的基本架构分为预处理、深度学习模型推理以及后处理三个主要部分,因此其 计算负载是一种混合计算负载。 而传统上,收集数据、数据预处理等所需的通用计算任务本来就适于部署在CPU平台上,并能借助丰富的优化工具和方法提升性能。而随着更多AI加速技术“嵌入”CPU,以及围绕CPU平台的AI加速生态逐步完善,已有充分数据显示, 以英特尔 ® 至强 ® 可扩展处理器为代表的CPU算力平台可完全满足AlphaFold2推理所需的算力。


图二 AlphaFold2基本架构


而在满足AlphaFold2任务所需的性能之余,基于CPU平台构建的方案还能为用户带来两方面的收益: 一方面 是能让用户在既有硬件平台上即可完成算力平台的搭建,并方便地将算力在其它计算任务与AlphaFold2之间调度,在充分调度空闲算力之余也能避免额外建设费用; 另一方面, 方案也帮助用户避免了GPU方案所面临的成本、采购以及算力调度等方面的难题,也对降低TCO有益。


同时,英特尔 ® 架构CPU平台庞大的市场占有率、广泛的部署基础,及其基于高效且稳定的x86架构所培育的丰富而友好的生态环境(包括对各类主流软件框架的广泛支持、借助插件对开源框架的优化,以及一系列用于提升AI性能的开源工具链等),也让CPU算力平台的搭建和获取都变得便捷且平滑。无论是采用私有化部署方式,还是依托公有云服务,AlphaFold2任务在任何时候,包括高密度使用场景下都能获得稳定且高效的CPU平台算力支持。


2、软硬件协同应对高通量计算与推理压力


在硬件上, 基于英特尔 ® 架构的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)技术,以显著的高位宽优势来提升计算过程中的向量化并行程度,从而有效提升向量/矩阵运算效率。


3、多种方法降低内存消耗,提升可预测序列长度


在解决方案中,英特尔 ® 至强 ® 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的英特尔解决方案也已与多个用户展开实践合作。


可靠性能,已经验证!






请到「今天看啥」查看全文