日前，深势科技研发团队将最新注意力机制加速技术Flash-Attention与Uni-Fold深度结合，进一步优化了模型的显存利用和计算效率。本次优化后，Uni-Fold端到端训练速度再增18%，将现有蛋白折叠模型（Jumper et al, 2021）的端到端训练时间由11天降至4天。此外，该技术大幅降低了模型推理的显存需求，在不使用模型并行与分块计算技术时，支持的最大序列长度增至2倍。相关实现已开源至深势科技Github仓库。
· 开源网址：
https://github.com/dptech-corp/flash-attention
图1 Flash-Attention让Uni-Fold快马加鞭。
图为端到端训练的总GPU小时数
作为困扰了结构生物学界50年的难题，蛋白质结构预测问题对生命科学和药物研发至关重要。为进一步推动领域发展，深势科技于2021年12月推出Uni-Fold，完整复现了目前学术界已知的最好的蛋白折叠算法（Jumper et al, 2021），并开源了训练代码。今年8月，深势科技开源Uni-Fold v2.0.0，进一步提高精度和效率，是目前效率最高的蛋白质折叠模型。日前，深势科技赓续技术迭代，历经月余打磨，将最新的注意力机制（Attention Mechanism）加速技术Flash-Attention融合进Uni-Fold中，让模型快马加鞭。
Flash-Attention借鉴NVIDIA Apex扩展包中的FMHA（Fast Multi-Head Attention）的实现，整体融合注意力算子，并对softmax运算进行分块和增量计算（Dao et al, 2022; Milakov et al, 2018），通过重计算而非重加载访存，以避免频繁显存访问，实现了运算效率的提高并降低峰值显存占用。然而，目前开源的Flash-Attention支持的场景尚为有限，并不能直接适用于Uni-Fold中的Evoformer模块中的计算。
图2 Flash-Attention算法示意图
（图Dao et al, 2022）
为解决这一问题，通过深入Flash-Attention的CUDA实现，适配数据流的多级重排和多级内存控制等，深势科技实现了该技术对Evoformer的支持。使用Flash-Attention后，Uni-Fold的单层Evoformer计算速度提升13%，显存消耗下降25%；端到端训练速度提升了18%，相较原始蛋白结构预测模型（Jumper et al, 2021）提升2.6倍，由11天降至4.1天（见图1）。同时，得益于显存的优化，在不开启分块计算（chunking）与模型并行的情况下，单次预测能支持的最大残基数目提升2倍，由1400左右提升至3000左右。
图3 Uni-Fold+FA与Uni-Fold的推理速度和显存使用大小
（使用40G显存GPU，精度为bfloat16）
通过融合Flash-Attention技术，Uni-Fold实现了进一步的加速与显存优化。深势科技将继续关注领域最新进展，致力于算法与产品打磨，为用户提供更加优质的服务，探索创造更多可能。
[1] Dao T, Fu D Y, Ermon S, et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135, 2022.
[2] Milakov M, Gimelshein N. Online normalizer calculation for softmax. arXiv:1805.02867, 2018.
[3] Jumper, J., Evans, R., Pritzel, A. et al. Highly accurate protein structure prediction with AlphaFold. Nature 596, 583–589 (2021).
深势科技是“AI for Science”科学研究范式的引领者和践行者，致力于运用人工智能和分子模拟算法，结合先进计算手段求解重要科学问题，为人类文明最基础的生物医药、能源、材料和信息科学与工程研究打造新一代微尺度工业设计和仿真平台。
我们开创性地提出了「多尺度建模+机器学习+高性能计算」的革命性科学研究新范式，并推出了Hermite®药物计算设计平台、Bohrium微尺度科学计算云平台等微尺度工业设计基础设施，颠覆了现有研发模式，打造“计算指导实验、实验反馈设计”的全新范式，为药物、材料领域带来极具突破性的计算模拟及设计工具。
深势科技是国家高新技术企业、北京市“专精特新”中小企业，总部位于北京，并在上海、深圳、海口等城市布局研发中心。科研技术团队由中国科学院院士领衔，汇集了近百名数学、物理、化学、生物、材料、计算机等多个领域的优秀青年科学家和工程师，其中公司的博士及博士后占比超过35%。核心成员获得过2020年全球计算机高性能计算领域的最高奖项“戈登贝尔奖”，相关工作当选2020年中国十大科技进展和全球AI领域十大技术突破。
深势诚邀有志之士，加入我们的“星途”家族：
你将与古往今来的“万千星辰”交相呼应，共创星辉；
你将与我们一同探索新理论、发现新技术、创造新产品；
在“AI for Science”的时代，深势希望助力科学“探险者”们见微知著、洞见未来。
未来虽远，前行必达，青春欲启，梦想莫待。
关于深势科技的更多信息，请访问：https://dp.tech
深势科技社招官网：https://dptechnology.jobs.feishu.cn/index/m/?spread=1FMM2VH
