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 倍，而不会损失准确性。有关该方法的详细信息可在我们的网站上找到。

更新日期：