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

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

概述

  • 为帮助用户在构建和使用 AlphaFold2 面临的挑战:构建平台的成本挑战,高通量计算和大规模推理压力,以及大内存需求问题。英特尔启动了针对自身架构的 AlphaFold2 方案构建及优化工作,基于至强® CPU Max 集成的 HBM 和英特尔® AMX 等软件的支持,新方案在提升了推理效率的同时降低内存消耗,能有效缓解长序列蛋白质结构预测工作中面临的 “序列长度天花板” 问题。结合 CPU 的易获取和成本优势,此方案已开始加速走进科研机构、高校实验室以及生物制药企业。

概述

对蛋白质三维结构的预测解析,已被广泛用于医学、制药、生物等领域的研究与创新。 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 方案开展蛋白质结构预测,并融入基于 AI 的新药研发、新诊疗方法探索、生物学新研究以及生物信息学、合成生物学等一系列科研和创新方向中。例如在一项研究中,科学家利用 AlphaFold2 为 sigma-2 和 5-HT2A 两种蛋白结构预测出数以亿计的潜在药物组合,并有着相当高的成功率,能为阿尔茨海默病、精神分裂症等疾病药物的研发提供更高效的途径2

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

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

▪    构建平台时的成本挑战:AlphaFold2 任务离不开高质量强算力硬件基础设施的支持,虽然 GPU 一直在多种 AI 任务中担纲主力,但其昂贵的市场价格以及难以获取的状况,也让用户不得不面临高昂的构建成本。同时用户在既有 IT 基础设施之外构建额外的 AI 算力平台,也容易带来算力调度不易、重复投资等问题;
▪    高通量计算和大规模推理压力:AlphaFold2 在架构上的多项创新设计,例如在结构网络层引入不变点注意力 (Invariant Point Attention) 机制等,在带来更优预测效果之余,也使整个方案从预处理阶段至推理任务都需要面对高通量的计算压力。同时为加速成果的问世,用户也亟需提升推理效率来缩短预测时长;
▪    大内存需求问题:在面向不同蛋白质的结构预测工作中,序列长度越长,张量运算规模与推理计算复杂度就越大,再结合更多序列的并行计算,其所需的内存容量也就越高。一旦承载平台的内存或显存容量不足,就可能限制更长序列的预测工作。

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

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

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

(如欲了解更多优化细节与结果,请参阅《CPU 加速 AlphaFold2 更上一层楼!第四代至强® 可扩展平台带来 3.02 倍通量提升》https://www.intel.cn/content/www/cn/zh/customer-spotlight/cases/cpu-acceleration-AlphaFold2-next-level-throughput.html


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

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

AlphaFold2 基本架构  

图二 AlphaFold2 基本架构

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

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

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

在硬件上,基于英特尔® 架构的 AlphaFold2 解决方案提供了英特尔® 至强® CPU Max 系列、第四代至强® 可扩展处理器等,来应对 AlphaFold2 面临的高通量计算与推理压力。就以前者为例,它所采用的全新微架构、更多的内核(最多达 56 个性能核)可为各类计算任务提供强有力的性能支撑。它在内存和输入/输出 (I/O) 子系统性能上有着显著的增强,提供了对 DDR5 内存(内存带宽可达 38.4 GBps(4,800MHz)以上)的支持,还结合大容量末级缓存使 AlphaFold2 推理过程中关键的张量吞吐获得了大幅提升。

英特尔® 至强® CPU Max 系列处理器  

图三 英特尔® 至强® 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 等不同指令集执行操作),这在提升推理效率的同时也能显著降低推理时的内存消耗。由于不同的蛋白质结构探索范围所需耗费的内存资源有比较大的差别,英特尔这一解决方案提供的各项为内存增量降耗的硬件方案与优化措施,能有效缓解长序列蛋白质结构预测工作中面临的 “序列长度天花板” 问题。

    面向 AlphaFold2 Multimer 模式的方案实现

与预测单个蛋白质三维结构的 AlphaFold2 Monomer(即 AlphaFold2)相比,AlphaFold2 Multimer 能对多个蛋白质分子之间的相互作用及所形成的复合体结构进行预测。这对研究和探明蛋白质之间的相互作用、信号传导和生物通路等课题有非常重要的意义。借助 AlphaFold2 Multimer,研究者能够更精确地了解蛋白质复合体的结构和动态行为,从而让生物、新药的研究变得更具效率优势。

面向 AlphaFold2 Multimer 模式的方案实现  

图五 面向 AlphaFold2 Multimer 模式的方案实现

基于英特尔® 架构的 AlphaFold2 解决方案同样也面向 AlphaFold2 Multimer 的管线结构进行了优化与验证,虽然其管线结构已根据蛋白质复合体结构预测的需求进行了调整,但英特尔 AlphaFold2 上的优化方案,在被用于 AlphaFold2 Multimer 时同样有效,同样能为之提供充沛且更具性价比的算力支持。目前,基于 AlphaFold2  Multimer 的英特尔解决方案也已与多个用户展开实践合作。

可靠性能,已经验证!

得益于内置 AI 加速引擎、集成 HBM 内存的至强® 系列处理器平台,以及围绕它们开展的端到端优化,整个 AlphaFold2 端到端处理过程的性能,都能在英特尔® 架构 CPU 平台上获得质的提升。如图六所示,在基于至强® CPU Max 系列处理器的优化流程中,每个优化步骤获得的提升累积后,与第三代至强® 可扩展处理器相比获得了高达 33.97 倍的通量提升3。根据测算,性能提升中的 74% 源自预处理阶段的高通量优化,26% 要归功于对推理过程的优化4

在至强® CPU Max 系列处理器上实现多实例通量提升  

图六 在至强® CPU Max 系列处理器上实现多实例通量提升

此外,在同样开启 IPEX 的情况下,相比第三代至强® 可扩展处理器,至强® CPU Max 系列处理器内置的 HBM 内存、英特尔® AMX 的加成,则带来了 48.3% 的性能提升5

在一项基于某公有云服务的测试中,基于英特尔® 架构 CPU 平台构建的解决方案在性能上获得了远优于某高端 GPU 平台的表现,也优于由 CPU + GPU 混合构建的方案。这表明英特尔® 架构 CPU 平台方案在 AlphaFold2 任务中有着显著的优势6

在应对 AlphaFold2 的大内存挑战方面,基于英特尔® 架构的解决方案也有效降低了内存峰值。如图七所示,在加入 TPP 技术后,在预测同样长度的蛋白质序列时,内存峰值仅为不使用 TPP 技术时的几分之一。这也意味着,用户在方案中部署同等容量的内存,可预测更长序列的蛋白质结构。

TPP 技术有效降低内存峰值  

图七 TPP 技术有效降低内存峰值

在解决方案通过 AMX_BF16 来加速 AlphaFold2 的推理过程时,如图八所示,在预测结果没有明显差异的前提下,在不同的输入长度下使用 BF16 数据格式的第四代至强® 可扩展处理器占用内存更少。而在总内存容量一致的情况下,使用 BF16 数据格式的第四代至强® 可扩展处理器可输入长度更大。

AMX-BF16 推理优化带来更低内存占用和更大输入长度  

图八 AMX-BF16 推理优化带来更低内存占用和更大输入长度

拥有上述综合性能优势,并兼具成本、易获取性特点的英特尔® 架构 AlphaFold2 解决方案目前已经开始加速走进科研机构、高校实验室以及生物制药企业。

如果您需要了解更多基于英特尔® 架构的 AlphaFold2 的解决方案详情,请联系:

杨威博士

英特尔中国人工智能架构师

邮箱:wei2.yang@intel.com