您的位置: 首页 > 技术杂谈 > 正文

仅用笔记本电脑解析90%蛋白质，单卡推理序列长度破万！Colossal-AI开源方案让AlphaFold推理提速5倍，显存降低75%

2022-10-28 12:00 https://my.oschina.net/u/5650766/blog/5586689 潞晨科技次阅读条评论

蛋白质是生命的物质基础，几乎支持着生命的所有功能。弄清楚蛋白质折叠成什么形状被称为 "蛋白质折叠问题"，在过去的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%的蛋白质结构。

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

展开阅读全文