蛋白质是生命的物质基础,几乎支持着生命的所有功能。弄清楚蛋白质折叠成什么形状被称为 "蛋白质折叠问题",在过去的50年里一直是生物学的一个巨大挑战。AlphaFold将Transformer引入了蛋白质结构预测领域,首次获得了原子级别精度, 如今被广泛应用于生物研究与药物开发领域。但是实际部署AlphaFold模型也带来了诸多挑战:
1)AlphaFold推理对GPU显存的需求极大,使用原版AlphaFold推理长度为3千的氨基酸序列结构,至少需要60G显存。
2)一个长序列的推理时间对于AlphaFold模型来说甚至长达几个小时。
3)整体流程中,因为需要对若干数据库进行MSA检索,预处理部分耗时巨大。
今年三月份,开源项目Colossal-AI开源子系统FastFold对这些挑战做出了初步回应,在单张GPU上相比AlphaFold推理速度提升2.01~4.05倍。现在,新版FastFold来了!
新版FastFold通过细粒度显存管理优化将推理含有1200个残基的蛋白质结构所需的16GB显存降到5GB(单精度),远低于NVIDIA RTX3050的8GB显存,也就是说,你 可以使用自己的笔记本推理90%的蛋白质结构 ,显著降低部署硬件成本!在单张NVIDIA A100(80GB)上,FastFold将单张GPU可推理的最大氨基酸序列长度突破至一万,覆盖了99.9999%的蛋白质。同时,FastFold通过GPU Kernel优化,动态轴并行,模型推理速度在上一版优化基础上再次提速25%,达到5倍加速,并使用全流程并行进行前处理加速,大大降低了推理所需的绝对时间。
开源地址:https://github.com/hpcaitech/ColossalAI#Biomedicine
技术路线
-
细粒度显存管理优化
AlphaFold虽然已经使用了chunk分块技术,但是它的推理仍然对于显存有非常高的要求,在实际应用是一个巨大的挑战,如果我们想要获得3000长序列的精确结构,则至少需要60G的显存,只有最顶级的专业显卡,例如价值约10万元的英伟达A100,才能满足这样的显存要求,这就让众多研究者与机构对与AlphaFold望而却步。
FastFold对分块计算技术进行了重构与优化,对各模块进行了针对化调整,优化分块的逻辑与效率;引入局部chunk或是重计算来减少中间变量的显存消耗;将分块技术扩展到了整个模型推理的过程中,降低了40%长序列推理时计算模块的峰值显存占用。
同时,我们提出了显存共享技术,通过在推理过程中构造张量可变结构,让每个操作都能够inplace地对张量进行修改,避免了在函数递归调用过程中显存复制的产生,零成本地降低了推理过程中的显存消耗,最多能够降低50%的显存开销。
在细粒度的显存管理优化下,我们首次突破了一万长序列的难关,使用半精度bf16在一张A100上即可进行推理,覆盖了99.99%的蛋白质,使用单精度fp32也能够推理7K长的氨基酸序列。
同时,90%的蛋白质结构长度都小于1200,通过显存管理技术,我们推理序列长度为1200的蛋白质仅需5GB显存,远低于NVIDIA RTX3050的8GB显存,让每一位研究者无需专业设备,只用自己笔记本上的消费级显卡,就能推理90%的蛋白质结构。
![](https://i-blog.csdnimg.cn/blog_migrate/e5d94d05c231230db0a1fde0310ce42e.png)
Chunk 示意图
-
GPU Kernel优化
FastFold使用了算子融合等计算优化技术,并根据对AlphaFold模型的性能特性,重新实现了LayerNorm和Fused Softmax,大幅度的提高了AlphaFold在GPU平台上的计算效率。针对原版AlphaFold有2到4倍的性能提升,大幅降低了AlphaFold在实际场景部署的成本。在新版本中,我们基于Triton实现了更优的算子,在原版FastFold的基础上继续实现了约25%的提速。
-
动态轴并行
FastFold根据AlphaFold的计算特征创新性地提出了动态轴并行技术。不同于传统的张量并行,动态轴并行选择在AlphaFold的特征的序列方向上进行数据划分,并使用All_to_All进行通信。动态轴并行对比张量并行有几个优势:1)支持Evoformer中的所有计算模块;2)所需的通信量比张量并行小得多;3)显存消耗比张量并行低;4)DAP给通信优化提供了更多的空间,如计算通信重叠。并且我们在新版本中增添了ExtraMSA和TemplatePairStack的并行化。在实际部署中,动态轴并行可以将计算分布到多个GPU上,从而大幅度的降低长序列模型的推理时间,对比原版AlphaFold有9到11倍的性能提升(使用8卡)。
-
全流程并行加速
FastFold通过使用Ray作为分布式计算引擎,支持异步的分布式调用。通过Ray,我们可以实现AlphaFold的推理流程中各个阶段的调度,并使得数据预处理部分并行化。对于单氨基酸序列的蛋白质结构推理,我们实现了约3倍前处理加速;对于N条氨基酸序列组成的蛋白质复合物结构推理,我们实现了约3*N倍前处理加速,大大降低了部署AlphaFold的成本。
性能测试
-
显存测试
自然界中90%的蛋白质残基数量小于1200,新版本FastFold只需要5GB显存就能完成推理,在消费级显卡NVIDIA RTX3050上也能完成。
下图展示了FastFold以单精度(fp32)推理不同长度序列时的显存情况,在单张NVIDIA A100 GPU(80GB)上,FastFold能够以单精度高效地推理长达7K的氨基酸序列,使用bf16则可以将能够推理的最大序列长度拓展到一万以上,提升显著。
-
速度测试
下图展示了FastFold以单精度推理不同长度序列的用时情况和当序列长度固定在2K时,GPU数量对于推理速度的影响。可以看到,在NVIDIA A100 GPU(80GB)硬件平台上,新版本FastFold单卡推理2K长氨基酸序列仅需不到10分钟。如果进一步将硬件平台拓展到4*A100,得益于动态轴并行,FastFold可以在两分钟半左右完成推理,实现几乎线性比加速。
下图展示了FastFold单层Evoformer的前向计算和反向传播的速度,得益于进一步优化的GPU Kernel,相较于0.1版本,新版FastFold在单层Evoformer表现上平均提速约25%,这意味着新版FastFold不但在推理上性能优异,在训练上也会有更好的表现。
One More Thing
FastFold上述突破基于面向大模型时代的通用深度学习系统 Colossal-AI,通过高效多维自动并行、异构内存管理、大规模优化库、自适应任务调度等实现高效快速部署AI大模型训练和推理,降低AI大模型应用成本。
Colossal-AI相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用,广受好评。
Colossal-AI注重开源社区建设,提供中文教程,开放用户社群及论坛,对于用户反馈进行高效交流与迭代更新,不断添加PaLM、AlphaFold、OPT等前沿应用。
自然开源以来,Colossal-AI已经多次在GitHub及Papers With Code热榜位列世界第一,与众多已有数万star的明星开源项目一起受到海内外关注!
项目开源地址: https://github.com/hpcaitech/ColossalAI#Biomedicine
参考链接:
[1] Tiessen, A., Pérez-Rodríguez, P. & Delaye-Arredondo, L.J. Mathematical modeling and comparison of protein size distribution in different plant, animal, fungal and microbial species reveals a negative correlation between protein size and protein number, thus providing insight into the evolution of proteomes. BMC Res Notes 5, 85 (2012). https://doi.org/10.1186/1756-0500-5-85