飞浆

见证 CPU 加速 AI for Science 普及潜能:基于英特尔® 至强® 的 AI 软硬件组合助百度飞桨优化多项科研模型与技术

  • 为进一步优化飞桨平台 HelixFold 模型和分子动力学模型在蛋白质结构预测、分子动力学计算等科研任务中的效能,百度引入面向单路和双路的第三代英特尔® 至强® 可扩展处理器,充分利用其内置的英特尔® AVX-512,还有搭配的英特尔® 傲腾™ 持久内存和英特尔® oneAPI 工具套件等,有效提升 HelixFold 模型和分子动力学模型的工作效率。

author-image

作者

概述

传统科学范式正面临维度灾难、数据风暴等多方面的挑战,而 AI(Artificial Intelligence,人工智能)技术则凭借其与生俱来的独特优势,开始帮助科学家和研究者突破传统范式,推进并加速科学研究和探索。在这一进程中,面向 AI for Science 的各类应用和方案发挥了关键作用,并迅速成为 AI 技术落地的热点之一,其应用不仅在学科研究层面,还涉及了从微观到介观、再到宏观的数学物理方程多尺度、跨尺度求解,在产业实践层面,也正在作用于智能设计制造与生物医药等产业场景。

作为这一融合创新浪潮中的重要参与者,百度深耕 AI 领域多年,不但拥有全面的 AI 技术布局,更是凭借旗下开源深度学习平台——飞桨平台(https://www.paddlepaddle.org.cn/),对各类硬件实现广泛适配,并直接部署到超算集群,与已有的科学计算生态紧密融合,强力支撑 AI for Science 方案的部署与应用。

为进一步优化飞桨平台面向 AI for Science 方向的重要成果——HelixFold 模型和分子动力学(Molecular Dynamics,MD)模型在蛋白质结构预测、分子动力学计算等科研任务中的效能,百度又与英特尔开展了一系列深度合作,通过引入面向单路和双路的第三代英特尔® 至强® 可扩展平台,充分利用其内置的英特尔® AVX-512,还有搭配的英特尔® 傲腾TM 持久内存和英特尔® oneAPI 工具套件等软硬件产品与技术,在有效提升 HelixFold 模型和分子动力学模型的工作效率的同时,也为超长序列蛋白质结构预测等需求提供了有效的实现途径,让百度飞桨平台在 AI for Science 的普及进程中更具应用优势。

“利用 AI for Science 方法来深入挖掘科学洞见,寻求基本规律并解决实际问题,是 AI 和大数据技术驱动下的全新科研范式。我们的飞桨平台正依托广泛的 AI 技术积累和产业实践优势,助力用户开展科学探索,加速各领域的科学发现。在这一过程中,来自英特尔的一系列软硬件为飞桨平台中 HelixFold 模型、分子动力学模型等的优化和能力提升提供了可观的助益。”

——向辉,百度飞浆 AI for Science 产品负责人

寻求创新范式突破,飞桨平台为科研勾勒 AI for Science新图景

▪ AI for Science 加速科研范式变革

一直以来,不同的科研范式都是科学研究寻求发现与突破的重要手段。例如基于第一性原理的牛顿范式,是通过复杂的科学研究与缜密的计算分析来获得对事物基本原理的认知。分子动力学,就是从这种经典物理的统计力学出发,研究分子间相互作用势函数及运动方程。又比如基于数据驱动的开普勒范式,则是从观察、采集到的数据中揭示事物的基本规律。今天许多生物信息领域的研究就是借助数据分析的方法获得了全新的突破。

诞生于上世纪 40 年代的电子计算机,在过去几十年间以持续高速演进的计算速度和分析精密度,帮助各领域的科学家与研究者获得了上述两种范式所需的庞大计算和数据分析能力,推动了基础学科和应用高速发展,并加速基因工程、材料科学、高速通信等技术实现了从基础理论到大规模产业化应用的跃迁。

科研范围的不断延伸和深入,让基于传统科研范式的计算与数据分析方法也开始遭遇新的挑战。例如在牛顿范式中,从一些自然科学规律,包括薛定谔方程、控制论方程等出发的研究,在面临高维偏微分方程时就可能遇到维度灾难问题(指由于变量或维数的增加,造成计算复杂度急剧增大)。而基于开普勒范式的研究,同样也会遇到因数据风暴而带来的处理成本倍增等问题。

AI 技术的发展及其与科学研究的融合,可帮助人们有效应对上述技术趋势和新兴场景带来的挑战。事实上,AI for Science 因为自身的特性,即广泛地把 AI 技术,更准确地说是深度学习技术用作各学科的科研、技术创新、成果转化等领域的新手段,正被公认是面向未来的科研范式。得益于 AI 技术具备的高通量数据分析、拟合建模等优势,其运用范围目前已延伸至高维偏微分方程求解、加速新材料研发与设计、预测蛋白结构以及在新药研发整个过程中为抗体的筛选、评估、预测、优化提供支持等环节。

▪ 飞桨平台助力勾勒 AI for Science 全景图

AI for Science 与科研场景的融合,同样也需要适宜的硬件平台与软件工具来予以支持,然而这种系统的构建并非易事。复杂的科学问题往往需要多尺度、跨尺度的模拟与多学科的交叉融合,以及基于复杂推演的问题计算等。针对这些需求,作为全球领先的 AI 公司,百度基于飞桨平台所具备的平台优势以及技术积累,开展了大量领先于时代的工作,为诸多科学家和研究者勾勒出了一幅 AI for Science 的全景图。

如图一所示,在核心框架的支持下,百度飞桨平台以 PaddleHelix(螺旋桨)、PaddleScience(赛桨)、Paddle Quantum(量桨)等工具组件,面向一系列包括科学计算、生物计算、量子计算等学科在内的前沿科研方向,提供了丰富的算法、开发接口,以及在不同场景中的算例。

图一  飞桨 AI for Science 全景图

同时,飞桨平台在基础 AI 能力上积累丰厚,其大量的计算机视觉、自然语言处理等领域的技术也能被复用到用户的 AI for Science 任务中。例如,计算流体力学领域的流场仿真,就可以用到计算机视觉领域的 UNet、GAN 等模型;而生物计算领域则可以使用到自然语言处理领域著名的 Transformer 序列标注模型。

基于这一平台,各科研领域的用户既能获得可直接复用的标准案例,也可基于组件对问题进行定制化的复现与分析,并通过数据驱动、与物理机理相结合等不同方法在物理仿真、化合物分子表征、量子纠缠处理等场景中进行探索。值得一提的是,在最新版本的核心飞桨框架(2.4 版本)中,还针对自动微分、编译器等提供了多种创新技术,并已在 3D 定常圆柱绕流等科学计算任务中完成了能力验证。

此外,AI for Science 的推进,同样也离不开飞桨平台为之提供的算力加成与优化。在算力构建与 AI 加速方面一直拥有独到优势的英特尔,以先进的软硬件产品与技术组合为百度与飞桨平台提供了强劲、且经过深度融合优化的算力和 AI 加速能力基座。现在,百度正与英特尔开展一系列技术合作,通过引入面向单路和双路的第三代英特尔® 至强® 可扩展处理器、其内置的英特尔® AVX-512 及其搭配的英特尔® 傲腾TM 持久内存和英特尔® oneAPI 工具套件等,在为平台中蛋白质结构预测任务、分子动力学计算任务提供高性能算力基础的同时,也在透过从指令集到硬件组件、再到软件工具的协同,为模型的定向优化和加速提供全面助力和支撑,使其能力范围获得进一步扩展,从而让飞桨平台在物理、化学、材料、医学生物等领域发挥出更大的作用。

英特尔软硬件助力 HelixFold 高效预测超长序列蛋白质结构

对蛋白质三维结构的解析与预测,一直是探索有机体构成、运行和变化规律的有效手段之一。在过去,科学家和研究者普遍使用 X-射线晶体衍射、冷冻电镜、核磁共振等方法与工具开展解析预测工作。但随着待测序列不断增多,这些方法在速度或效率上已无法满足需求,而基于 AI 技术的新方案则在近年来崭露头角。其中,由 DeepMind 在 2021 年开源的 AlphaFold2 将深度学习等 AI 方法用于蛋白质结构预测,其原子级别的预测精度,为 AI for Science 在生物制药、病理学以及制药学的实践树立了全新里程碑。

图二  AlphaFold2 核心模型框架

AlphaFold2 提供了完整的端到端蛋白质三维结构预测流程,如图二所示,其模型框架大致可分为数据预处理(Preprocessing)、Evoformer 网络模型、结构预测模块(Structure Module)三个部分。从计算原理、计算量和计算复杂度来看,这三个部分对硬件平台的算力和存储性能都有着较高的要求。

▪ 预处理阶段优化

在数据预处理部分,系统需要检索一个非常大的多序列比对(Multiple Sequence Alignment,MSA)数据库,用于获得不同蛋白质之间的共有进化信息,提升预测精度。而在 Evoformer 网络模型和结构预测模块中,系统需要处理多达 192 个模块(48*4),且计算复杂度为 L^3,存储复杂度为 L^2,这就带来了以下两方面的挑战:
 

  • 计算量巨大。尤其当要预测的蛋白序列很长时,需要系统更多采用多实例的并行计算模式来提升处理效率;
  • 内存墙高筑。中间数据的计算过程都需要内存参与,这意味着预测的蛋白序列越长,所需内存容量也就越大。在实践中,当序列长度为 765 时,AlphaFold2 所需内存峰值可达 50GB 左右,而对于长度超过 4,000 的蛋白质序列而言,所需的内存峰值就能超过 1TB。由于传统的动态随机存取存储器(Dynamic Random Access Memory,DRAM)内存较难实现 TB 级的大容量,因此当系统开展长序列解析任务时,推理过程很容易耗尽内存而出现 Out Of Memory (OOM) 的报错。

为应对以上挑战,百度飞桨螺旋桨 PaddleHelix 团队基于飞桨平台,依托高性能并行计算能力及大容量内存,发布了蛋白质结构预测模型 HelixFold。与原版 AlphaFold2 相比,HelixFold 模型通过多项创新技术,包括分支并行与混合并行策略、算子融合优化技术、张量融合低频次访存技术以及多维度内存优化方案等的加入,使预测方案获得了全面的性能优化。在测试中,在 CAMEO 八月份最新的数据集上(60 个蛋白),HelixFold 模型的 TM- Score 指标达到 0.9026,预测精度超过原版 AlphaFold2(0.8912)和基于 PyTorch 的 OpenFold(0.8827),训练效率则提升了 106.97%1

为帮助 HelixFold 模型在蛋白质结构预测赛道中获得更显著的优势,百度与英特尔一起,借助面向单路和双路的第三代英特尔® 至强® 可扩展处理器、英特尔® AVX-512、英特尔® 傲腾TM 持久内存以及英特尔® C++ 编译器等软硬件产品与技术所具备的优势,为 HelixFold 模型预处理、推理与预测等提供了更强的高性能并行计算能力和更大容量内存。

第三代英特尔® 至强® 可扩展处理器一方面具备出色的微架构设计,如更多的核心、线程数量和高速缓存容量,可满足整个预测过程所需,并能凭借多核心优势,为并行计算带来效率上的大幅提升。另一方面,至强® 可扩展处理器内置的英特尔® AVX-512 能以高位宽优势(最大可提供 512 位向量计算能力)提升计算过程中的向量化并行程度,使向量/矩阵运算效率获得有效提升。上述产品与技术的运用,不仅满足了 HelixFold 模型在蛋白质结构预测中的算力所需,也在算力资源挖掘上获得了显著的突破,在多实例并行计算中处理器使用率可达 100%2

借助英特尔® 至强® 可扩展处理器在多核心、高位宽,以及内存访问时的低时延优势,英特尔与百度飞桨 PaddleHelix 团队一起为HelixFold模型预处理阶段的性能提升,加入了三重优化。如图三所示,包括:
 

  • 优化措施 1:进行英特尔® C++ 编译器与 Jemalloc 协同优化。该优化选项可对处理过程中的碎片化内存进行优化,带来显著性能提升;
  • 优化措施 2:面向英特尔® AVX-512 展开优化。通过代码层面的优化,使更多指令可运行于英特尔® AVX-512 上,借助其高位宽优势来提升处理性能;
  • 优化措施 3:进行面向并行化 MSA 的优化。英特尔® 至强® 可扩展处理器能够支持方案并行处理多路 MSA 工作负载,大幅提升处理效率。

▪ 模型推理优化

与此同时,百度和英特尔也在 HelixFold 模型的模型推理阶段打造了行之有效的技术优化方案。HelixFold 模型在原有的模型推理中采用了动态图编译模式,这会面临两方面的性能挑战:
 

  • 对于推理过程中需动态调整的张量(tensor),需在计算图编译过程中一次性全部生成并存入内存,这不但会耗费大量的内存空间,也会大幅延长内存分配时间;
  • 原技术方案只能利用一个或有限个处理器物理核心来进行运算,无法发挥处理器多核性能。

面对上述挑战,英特尔首先与百度飞桨 PaddleHelix 团队一起构建了基于静态图编译模式的技术方案。如图三所示,与动态图模式中需要一次性生成全部张量相比,静态图模式引入 JIT(Just In Time)编译技术,将计算图分成多个子图。每次只需编译一个子图,运行得到结果后再用于下一个子图的编译,并以此类推。这带来的优势显而易见:
 

  • 一方面,逐个编译的模式不再需要生成大规模的张量,可以显著节约内存空间以及内存分配时间;
  • 另一方面,借助英特尔® oneAPI 工具套件中的英特尔® oneDNN,新方案不仅能通过将许多用到的计算基元向量化为英特尔® AVX-512 指令的方式,深度调动英特尔® AVX-512 的高位宽优势来提升编译过程的效率,也能让更多物理核心参与到计算过程中,让第三代英特尔® 至强® 可扩展处理器多核心的优势充分显现。

图三  面向 HelixFold 模型推理流程的优化

▪ 内存优化

而对于 HelixFold 模型所涉及的大规模中间计算过程,英特尔® 傲腾™ 持久内存则为之提供了兼具大容量和高性价比的内存支撑。以英特尔® 傲腾™ 持久内存 200 系列为例,其最高可提供 512GB 的单模组容量,在与面向单路和双路平台的第三代英特尔® 至强® 可扩展处理器搭配后,能在提供3,200MT/S 内存带宽的基础上,理论上实现每路高达 4TB 的英特尔® 傲腾™ 持久内存容量配置,以及每路高达 6TB 的内存总容量(与 DRAM 内存组合后),同时其性能表现也接近 DRAM 内存。由此,新方案在超长序列蛋白质结构预测能力上获得了大幅的提升,能轻松预测推理长度超过 4,000 的蛋白质序列。

联动 DeePMD,英特尔® oneAPI 助分子动力学深度学习方案实现推理加速

基于飞桨平台,百度在另一个 AI for Science 的重要成果,即面向分子动力学计算场景的 DeePMD 模型上,也通过引入英特尔的 AI 软硬件产品组合开展了深度优化,并获得了实质进展。

分子动力学

通过对分子间相互作用势函数及运动方程的求解,分析分子运动的行为规律,模拟体系的动力学演化过程,给出微观量(如:分子的坐标与速度等)与宏观可观测量(如:体系的温度、压强、热容等)之间的关系,从而研究复合体系的平衡态性质和力学性质,是研究材料内部流体行为、通道运输等现象有效的研究手段。

在传统的分子动力学计算任务中,首先需要选取合理的势函数,例如采用第一性原理的方法开展计算,通过电荷密度得到 Hellmann-Feynman 力。这些方法虽然精度较高,但计算代价也相当高昂,需要为之配备足量的计算资源,而这就让计算资源受限成为了分子模拟的尺度难以得到进一步扩展的主因。

近年来,随着 AI 技术在分子动力学计算领域一展所长,上述瓶颈有望得到突破。这其中的佼佼者 DeePMD,正是通过引入深度学习方法,将第一性原理的计算精度带到了更大尺度的分子模拟。基于 DeePMD-kit 工具包,能在无人工干预的前提下,在数小时内以端到端的方式,将用户提供的数据转换成面向多体势能表示和分子动力学的深度势能模型,且能与常用分子动力学模拟软件,如 LAMMPS 等无缝集成。而这可在保持量子力学精度和准确度的基础上,将分子动力学计算任务的速度提高几个数量级。

为了让更多用户从基于 AI 方法的分子动力学计算方案中收益,百度在经过一系列算子优化和编译上的转换后,在中国国内率先完成了与传统分子动力学软件 LAMMPS 以及 DeePMD-kit 工具包的融合,推出了基于飞桨平台的 Paddle_deepmd-kit。

为保证其具有更优的计算能力输出,百度在同样引入面向单路和双路的第三代英特尔® 至强® 可扩展处理器作为其算力核心的同时,还联动英特尔® oneAPI 工具套件,使方案的训练和推理全流程实现了从“0 到 1”的突破。

英特尔® oneAPI 工具套件

一套完整的高级编译器、库以及移植、分析和调试器工具,其利用先进的硬件功能加速计算,可与现有的编程模型和代码库互操作,包括英特尔® oneAPI Base 工具包、英特尔® oneAPI AI Analytics 工具包及 OpenVINO™ 工具套件等不同工具。

其中,英特尔® oneMKL 是一个广泛的科学/工程数据库,其面向基于英特尔® 架构的处理器的特性,尤其是多线程优势进行了专门的优化,在函数计算、特征向量、特征值求解等计算中有着独特的优势。

英特尔与百度飞桨科学计算团队一起,借助英特尔® oneMKL对DeePMD中常用的通用矩阵乘(General Matrix Multiplication,GEMM)函数、矢量加(Vector Add,VADD)函数以及面向JIT的双曲正切(tanh)函数等进行了优化,并基于英特尔® oneMKL,使飞桨平台、LAMMPS和DeePMD的特性得到了更好地融合。同时,第三代英特尔® 至强® 可扩展处理器所内置的英特尔® AVX-512,也对分子动力学模型的优化提供了可靠的算力加速支持。

“在人工智能科学领域,英特尔和百度工程师携手从‘0’到‘1’,在分子动力学领域共同研发,基于英特尔® oneAPI 工具套件和第三代英特尔® 至强® 可扩展处理器,并利用英特尔® AVX-512 等优化分子动力学模型,且支持软件 LAMMPS,调优后能满足研究需要,有效提升用户深度学习应用的工作效率。目前该产品已开放测试,赋能材料研发工作,助力深度学习技术在 AI for Science 领域的应用,并奠定了人工智能框架与传统科学计算相结合的基石。”

——于佃海,百度飞浆总架构师

展望未来

随着 AI for Science 在越来越多的领域中发挥作用,基于百度和英特尔深度技术合作的飞桨 AI 科学计算平台也在科研中产生日益显著的作用。例如,借助 HelixFold 模型使可探索的靶点蛋白范围不断扩大,进而帮助制药企业挖掘更多靶向药物分子,对更多疑难杂症形成针对性的治疗方案;由分子动力学模型推动的分子动力学深度学习方案,也帮助更多用户提升材料模拟计算等方面的效率,加速新材料、新产品的研发速度。

除 HelixFold 模型和分子动力学模型之外,百度和英特尔也正基于飞桨平台,在其它一系列面向 AI for Science 的应用方向上开展广泛合作。例如双方基于英特尔® 架构平台产品与技术,通过对飞桨平台中的特定模型进行优化来更快、更精准地预测海洋气象与潮汐。

面向未来,百度还将与英特尔继续携手,发挥各自在 AI 和计算领域上的优势,为更多科研领域的用户打造高质、高效的 AI 科学计算平台,助力科学探索、加快科学发现,并促进科研成果加速产业落地,使 AI for Science 应用更加普及的同时,也让其使用起来更为便捷和高效。