DS4Sci_EvoformerAttention 消除了扩展以 Evoformer 为中心的结构生物学模型的内存爆炸问题
1. 什么是 DS4Sci_EvoformerAttention
DS4Sci_EvoformerAttention
是一个内核集合,旨在通过减少内存占用并提高训练速度,将 Evoformer 计算扩展到更多数量的序列和残基。
2. 何时使用 DS4Sci_EvoformerAttention
DS4Sci_EvoformerAttention
在序列和残基数量很大时最有效。前向内核经过优化以加速计算。在各种注意力机制的推理过程中使用前向内核是有益的。关联的反向内核可以在训练期间用于减少内存占用,但会以牺牲一些计算为代价。因此,在训练中针对 MSA 行注意力和 MSA 列注意力等内存受限操作使用 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 列和 2 种三角形),如下面的示例所示。在示例中,我们将序列数表示为 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 倍,而不会损失准确性。有关方法的详细信息,请参阅 我们的网站。