戴尔PowerEdgeXE8545助力企业应对AI时代浪潮人体免疫系统的抗体蛋白、胶原蛋白、抗冻蛋白、核糖体等,每一种蛋白都有其独特的三维蛋白结构。蛋白质的研究对生命科学和药物开发具有重要价值。人类的很多疾病都是由蛋白质错误折叠引起的,比如帕金森病、阿尔茨海默病、亨廷顿舞蹈病等。药物的原理类似于一把钥匙,而锁就是疾病的目标,通常可以看作是一种蛋白质。在锁孔中加入“关键小分子”,与蛋白质结合,抑制蛋白质的正常功能,或激活蛋白质的某些功能。因此,了解和预测蛋白质的形状将有助于科学家设计新的更有效的疾病治疗方法,有助于发现新药,并降低实验成本。蛋白质由氨基酸序列组成,但真正决定蛋白质功能的是它的三维结构,即氨基酸序列如何折叠,如下例所示。蛋白质结构发现的主要方法有X射线晶体衍射、核磁共振、冷冻电镜三维重建等,从2013年开始流行。但冷冻电镜的采购成本高,且图像重建需要大量的计算能力,解决一个新的蛋白质3D结构往往需要很长时间。1972年诺贝尔化学奖获得者、美国生物化学家克里斯蒂安·安芬森提出,给定一个氨基酸序列,可以从理论上预测蛋白质的三维结构。五十年来,为了验证这一理论,科学家们尝试了各种模型来预测蛋白质结构,但在全球蛋白质结构预测领域最著名的CASP竞赛中,到2018年的预测准确率仅为40%左右.蛋白质越大,模型就越复杂和困难,因为需要考虑更多氨基酸之间的相互作用。据统计,枚举蛋白质的可能构型平均搜索空间为10的300次方。使用分子动力学结构预测等传统计算方案需要极高的计算能力和较长的计算时间。在过去的50年中,只有17%的人类蛋白质组得到了结构解析。谷歌DeepMind在2020年推出的AlphaFold2改变了这一切。在2020年12月发布的CASP14成绩单中,AlphaFold2将CASP蛋白质结构预测得分提高到92.4分(满分100分),距离真实蛋白质结构仅差一个原子。2021年7月,Alphafold2模型结构和训练过程在Nature发表,蛋白质结构数据库和推理代码开源。Alphafold2可以预测98.5%的人类蛋白质结构,60%的结构位置预测是可靠的。《科学》杂志将AlphaFold2评选为2021年十大科学发现。与第一代AlphaFold的卷积神经网络相比,Alphafold2利用多序列比对(MSA)将蛋白质结构和生物信息整合到深度学习算法中,主要包括神经网络EvoFormer和结构模块StructureModule:EvoFormer主要将图网络和多序列比对相结合,完成结构预测。图网络将蛋白质相关信息构建成图,表示不同氨基酸之间的距离;tripleAttention自注意力机制用于处理氨基酸之间的关系图。结构模块主要将EvoFormer获得的信息转化为蛋白质的3D结构。AlphaFold2是一个端到端的神经网络,它反复将最终的损失应用于输出,然后在输出上递归以不断逼近正确的结果。那么,训练AlphaFold2并使用AlphaFold2进行蛋白质结构预测的推理计算需要什么样的计算支持呢?戴尔科技中国研究院和戴尔数据中心事业部解决方案团队从GitHub上下载了AlphaFold2模型代码,部署在戴尔PowerEdgeXE8545服务器上,使用NVIDIAA100GPU对AlphaFlod2进行了68-2750个氨基酸残基不同大小的测试。蛋白质的3D结构预测,AlphaFold2的计算性能和特性评估。DellPowerEdgeXE8545是戴尔科技最新推出的15G服务器家族中专为AIGPU计算设计优化的加速服务器。4U空间可支持4块A100GPU加速卡,通过NVLink可实现GPU间600GB/s的pear-to-pear高速直接通信。测试环境软硬件配置如下:AMDEPYC771364核处理器×21024GB内存NvidiaA100GPU×4,80GB/500WCentOSLinux7.0Python3.8.0,TensorFlow2.5.0CUDA11.5,ThecuDNN8.3XE8545推理对由68-2750个氨基酸残基组成的蛋白质的3D结构预测所消耗的计算时间如下表所示(Top1模型,即推荐置信度最好的模型),计算时间为单次A100推理从19.3分钟到2个半小时。根据DeepMind论文的Top5模型,XE8545单卡A100的推理计算时间如下:通过性能日志分析,我们可以清楚的看到AlphaFold2在推理过程中是CPU和GPU交替计算的过程,第一阶段是同源序列搜索和模板搜索。而特征构建,以及最后阶段3D结构生成的计算过程主要由CPU计算;中间第二阶段的Evoformer神经网络和结构模块的计算主要由GPU计算。XE8545提供的强大GPU计算能力和AMD多核CPU计算能力(128核),可以保证AlphaFold2在规定时间内完成大型蛋白质3D结构的预测计算。我们还比较了不同GPU对AlphaFold2推理计算性能的影响。我们选用戴尔7750工作站,配备NVIDIARTX5000显卡,对比蛋白质结构预测(Top1机型)的计算性能。性能与消费级GPU性能相差无几;蛋白质越大,使用A100结构预测的加速性能越明显。预测1511个残基的蛋白质3D结构,XE8545+A100花费RTX500065%的时间;预测2000个残基的蛋白质三维结构,XE8545只需要RTX5000的50%时间。我们可以看到,RTX5000在预测2800个残基的蛋白质结构时,由于内存容量和计算能力的限制,无法完成结构预测工作,而XE8545依然在数小时内成功完成了相同规模的蛋白质结构预测。从模型训练来看,Alphafold2及后续类似的蛋白质结构预测模型,由于使用了Transformer机制,模型训练需要非常高的计算能力,通常需要64-512个GPU组成一个计算集群,采用分布式训练机制,以便在较短的时间内实现模型收敛。DeepMind在论文中提到,训练AlphaFold2模型使用了128个谷歌TPU芯片,完成模型训练需要将近2周的时间。2022年3月,上海交通大学联合绿辰科技发布的FastFold模型,使用256块A100GPU进行初始训练,512块A100GPU进行Fine-tuning,用时2.81天完成模型训练。DellTechnologies的AIGPU分布式训练解决方案可以提供高速GPU计算、快速读写小文件IO(蛋白质数据库中有大量小文件)、高带宽和低延迟的网络通信,帮助用户在深度学习框架下实现分布式训练的自动化实现和性能优化,轻松应对AI时代的浪潮。此外,2021年发布的《戴尔科技AI GPU分布式训练技术白皮书》还可以为用户AI大模型GPU分布式训练提供基础设施解决方案、参考架构和优化建议。公众号后台回复关键字“白皮书”,轻松获取~
