DS4Sci_EvoformerAttention 解决了扩展以 Evoformer 为中心的结构生物学模型的内存爆炸问题
1. 什么是 DS4Sci_EvoformerAttention
DS4Sci_EvoformerAttention
是一系列内核的集合,旨在通过减少内存占用和提高训练速度,将 Evoformer 计算扩展到更多的序列和残基。
2. 何时使用 DS4Sci_EvoformerAttention
DS4Sci_EvoformerAttention
在序列和残基数量很大时最为有益。前向内核经过优化以加速计算。在推理过程中使用前向内核对于各种注意力机制都是有益的。相关的反向内核可以在训练期间使用,以减少内存占用,但会牺牲一些计算。因此,在训练中,对于内存受限的操作,例如 MSA 行向注意力(MSA row-wise attention)和 MSA 列向注意力(MSA column-wise attention),使用 DS4Sci_EvoformerAttention
是有益的。
3. 如何使用 DS4Sci_EvoformerAttention
3.1 安装
DS4Sci_EvoformerAttention
作为 DeepSpeed >= 0.10.3 的一部分发布。DS4Sci_EvoformerAttention
是基于 CUTLASS 实现的。您需要克隆 CUTLASS 仓库并在环境变量 CUTLASS_PATH
中指定其路径。
git clone https://github.com/NVIDIA/cutlass
export CUTLASS_PATH=/path/to/cutlass
当首次调用 DS4Sci_EvoformerAttention
时,内核将被编译。
DS4Sci_EvoformerAttention
需要计算能力 7.0 或更高(NVIDIA V100 或更高版本 GPU)的 GPU,并且最低 CUDA 版本为 11.3。建议使用 CUDA 11.7 或更高版本以获得更好的性能。此外,目前 V100 内核上的反向内核性能不如 A100。
3.2 单元测试和基准
单元测试和基准测试可在 DeepSpeed 仓库的 tests
文件夹中找到。您可以使用以下命令运行单元测试和基准测试。
pytest -s tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py
python tests/benchmarks/DS4Sci_EvoformerAttention_bench.py
3.3 将 DS4Sci_EvoformerAttention 应用到您的模型
要在用户的模型中使用 DS4Sci_EvoformerAttention
,您需要从 deepspeed.ops.deepspeed4science
导入 DS4Sci_EvoformerAttention
。
from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention
DS4Sci_EvoformerAttention
通过使用以下示例中所示的不同输入,支持 Evoformer 中的四种注意力机制(MSA 行向、MSA 列向以及两种三角注意力)。在这些示例中,我们将序列数量表示为 N_seq
,残基数量表示为 N_res
。隐藏状态的维度 Dim
和头数 Head
在不同的注意力机制中有所不同。请注意,DS4Sci_EvoformerAttention
要求输入张量为 torch.float16
或 torch.bfloat16
数据类型。
(a) MSA 行向注意力 为残基对构建注意力权重,并将对表示中的信息作为附加偏置项进行整合。
# Q, K, V: [Batch, N_seq, N_res, Head, Dim]
# res_mask: [Batch, N_seq, 1, 1, N_res]
# pair_bias: [Batch, 1, Head, N_res, N_res]
out = DS4Sci_EvoformerAttention(Q, K, V, [res_mask, pair_bias])
(b) MSA 列向注意力 允许属于同一目标残基的元素交换信息。
# Q, K, V: [Batch, N_res, N_seq, Head, Dim]
# res_mask: [Batch, N_seq, 1, 1, N_res]
out = DS4Sci_EvoformerAttention(Q, K, V, [res_mask])
(c) 三角自注意力 更新对表示。三角自注意力有两种:围绕起始节点和围绕结束节点。下面是围绕起始节点的三角自注意力的示例。围绕结束节点的三角自注意力类似。
# Q, K, V: [Batch, N_res, N_res, Head, Dim]
# res_mask: [Batch, N_res, 1, 1, N_res]
# right_edges: [Batch, 1, Head, N_res, N_res]
out = DS4Sci_EvoformerAttention(Q, K, V, [res_mask, right_edges])
4. DS4Sci_EvoformerAttention 科学应用
4.1 DS4Sci_EvoformerAttention 解决了 OpenFold 中扩展以 Evoformer 为中心的结构生物学模型的内存爆炸问题
OpenFold 是 DeepMind AlphaFold2 的社区复现,使得在新的数据集上训练或微调 AlphaFold2 成为可能。训练 AlphaFold2 会导致内存爆炸问题,因为它包含几种自定义的 Evoformer 注意力变体,这些变体产生了异常大的激活。通过利用 DeepSpeed4Science 的 DS4Sci_EvoformerAttention 内核,OpenFold 团队能够将峰值内存需求降低 13 倍,而不会损失准确性。有关该方法的详细信息可在我们的网站上找到。