企业官网建设流程全解析

DeepSpeed4Science 技术指南：AI4Science 大模型长序列训练与 Evoformer 内存高效注意力内核实战

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

DeepSpeed4Science 是 DeepSpeed 团队发起的面向科学发现场景的 AI 系统技术专项，本篇指南以其概述页 docs/_pages/deepspeed4science.md 为主线，聚焦两项已随项目发布的核心技术：面向 AI4Science 大模型的新版 Megatron-DeepSpeed 长序列训练框架，以及随 DeepSpeed 主仓库发布的内存高效 EvoformerAttention 内核（DS4Sci_EvoformerAttention）。读完本文，你将理解这些技术要解决的科学计算痛点、适用的模型与场景，掌握从安装、构建到在自己 Evoformer 类模型中接入DS4Sci_EvoformerAttention的完整方法，并了解这些技术已落地的科学应用（如结构生物学中的 GenSLMs 与 OpenFold）。配套的完整操作教程见 docs/_tutorials/ds4sci_evoformerattention.md，项目公告见 blogs/deepspeed4science/README.md。

新 Megatron-DeepSpeed 框架下训练 33B GPT 类模型的序列长度能力示意：相比 NVIDIA Megatron-LM 可在不触发显存溢出（OOM）的前提下支持约 9 倍更长的序列。

DeepSpeed 的 EvoformerAttention 内核帮助 OpenFold 将训练峰值显存需求降低约 13 倍。

1. DeepSpeed4Science 是什么：面向科学发现的 AI 系统技术专项

以"通过 AI 系统技术创新，帮助领域专家解开当今最大的科学谜题"为目标，DeepSpeed 团队发起 DeepSpeed4Science 专项，把分布式训练与推理系统能力沉淀为可供科学家按需取用的技术集合。概述页面向两类读者：使用方（科学家）可在此"选购"适合自身模型的技术，并查看每项技术的用途、适用时机、使用入口与科学应用案例；贡献方则可将自己将 DeepSpeed 技术应用于科研的实践补充为展示案例。

如果要引用 DeepSpeed4Science，官方给出的引用方式对应白皮书（arXiv:2310.04610），即：

@article{song2023deepspeed4science, title={DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies}, author={Song, Shuaiwen Leon and Kruft, Bonnie and Zhang, Minjia and Li, Conglong and Chen, Shiyang and Zhang, Chengming and Tanaka, Masahiro and Wu, Xiaoxia and Rasley, Jeff and Awan, Ammar Ahmad and others}, journal={arXiv preprint arXiv:2310.04610}, year={2023} }

该专项首批公布的技术聚焦结构生物学等科研计算中真实存在的系统瓶颈，可归纳为两大方向：其一是在训练框架层支持"极长序列 + 超大基因组/蛋白质语言模型"，其二是在算子层为 Evoformer 这类科学模型特有注意力机制消除显存爆炸。下文分别展开。

2. 技术一：新版 Megatron-DeepSpeed —— AI4Science 大模型长序列训练

2.1 它解决什么问题

AI4Science 领域的基因组规模基础模型（如 GenSLMs）需要远超通用大模型的超长序列支持。在 DeepSpeed 侧，这套能力通过新版 Megatron-DeepSpeed提供：它基于最新版 Megatron-LM 重新 rebase 并使能 DeepSpeed，补齐了长序列支持与多项新能力。值得注意的是，该框架在独立的 Megatron-DeepSpeed 仓库中维护（不在当前 DeepSpeed 仓库内），但其背后复用的 ZeRO、并行策略等系统能力均来自 DeepSpeed 本仓库的实现。

2.2 关键系统能力

在新 Megatron-DeepSpeed 中，长序列训练由多种技术的协同组合支撑：

• ZeRO 风格的数据并行（ZeRO-style data parallelism）与张量并行（tensor parallelism）；
• 序列并行（sequence parallelism）与流水并行（pipeline parallelism）；
• 模型状态卸载（model state offloading），将优化器状态等迁出显存；
• 新增的内存优化手段，例如attention mask 卸载（attention mask offloading）与位置编码切分（position embedding partitioning）。

这些能力的背景知识（ZeRO 各阶段、卸载、流水/张量并行等）可分别参考本仓库中的 zero 教程、zero-offload 教程、pipeline 教程 与 autotp-training 教程。

2.3 效果与科学应用

依据概述页公布的结果，使用新 Megatron-DeepSpeed 训练 33B 参数的 GPT 类模型时，其可支持的序列长度相比 NVIDIA Megatron-LM 提升了约9 倍，且全程不触发显存溢出（见 new-megatron-ds 示意图）。

这一框架已实际应用于基因组规模基础模型 GenSLMs——一个获得 2022 年 ACM Gordon Bell 奖的基因组语言模型项目（来自阿贡国家实验室）。GenSLMs 这类模型在训练与推理阶段都需要远超通用大模型长序列策略的超长序列支持。借助 DeepSpeed4Science 的新 Megatron-DeepSpeed，GenSLMs 团队成功将其 25B 模型训练到512K 序列长度，远超其原先的 42K 序列长度。

3. 技术二：DS4Sci_EvoformerAttention —— 内存高效的 Evoformer 注意力内核

3.1 为什么标准注意力优化方案不管用

Evoformer 是 AlphaFold 等蛋白质结构预测科学模型的核心构件。其MSA（多序列比对）注意力在训练/推理过程中（如蛋白结构预测模型）经常遭遇显存爆炸，根源在于其注意力矩阵规模随序列数 × 残基数急剧增长，前向激活巨大。而 FlashAttention 这类通用方案无法直接支持 Evoformer，因为 Evoformer 使用的是行方向（row-wise）、列方向（column-wise）与三角（triangle）注意力，与标准 Transformer 的自注意力/交叉注意力形态不同，需要定制化优化。

为此，项目发布DS4Sci_EvoformerAttention内核集合，通过减小显存占用、提升训练速度，把 Evoformer 计算扩展到更大的序列数与残基数规模。

3.2 何时使用

当序列数（N_seq）与残基数（N_res）较大时收益最明显：

• 前向（forward）内核面向计算加速优化，在推理时对各种注意力机制使用前向内核都有收益；
• 反向（backward）内核在训练时以少量额外计算换取显存占用的大幅下降，因此训练阶段推荐对显存敏感的操作（如 MSA 行方向注意力、MSA 列方向注意力）使用DS4Sci_EvoformerAttention。

3.3 安装与构建前提

DS4Sci_EvoformerAttention作为 DeepSpeed（版本不低于 0.10.3）的一部分随主仓库发布。它基于CUTLASS实现，编译时需要定位到 CUTLASS。DeepSpeed 的自动探测会依次查找：

1. nvidia-cutlassPython 包（即cutlass_library模块源码目录）；
2. Python 环境与 CMake 前缀（CONDA_PREFIX、VIRTUAL_ENV、CMAKE_PREFIX_PATH、CUDA_HOME等环境变量）；
3. 编译器 include 路径环境变量（CPATH、CPLUS_INCLUDE_PATH、C_INCLUDE_PATH）；
4. 位于 DeepSpeed 旁边或当前工作目录的cutlasscheckout；
5. 常见系统安装前缀（如/usr/local、/opt/cutlass等）。

这套完整的探测顺序可在 op_builder/evoformer_attn.py 的_candidate_cutlass_paths()与include_paths()中看到实现。若探测到的 CUTLASS 是 checkout 形态，构建器还会检查其CHANGELOG.md中是否包含3.1.0，以强制要求CUTLASS >= 3.1.0。

使用中的几个实用要点：

• 最省事的做法是在 DeepSpeed 旁边克隆一份 CUTLASS（NVIDIA/cutlass仓库）：

  git clone https://github.com/NVIDIA/cutlass

  内核会在DS4Sci_EvoformerAttention首次被调用时触发编译。

• 若自动探测没有找到预期安装，可通过CUTLASS_PATH显式指定 CUTLASS checkout 根目录或其include目录。
• 若你的编译环境已妥善配置好 CUTLASS 与 CUDA 编译器（例如在 conda 包中编译），可用CUTLASS_PATH="DS_IGNORE_CUTLASS_DETECTION"跳过 CUTLASS 探测。
• 另有DS_USE_CUTLASS_PYTHON_BINDINGS（对应nvidia-cutlassPython 包源码探测）开关定义于构建器中。

硬件与软件下限：需要计算能力 7.0 及以上的 GPU（NVIDIA V100 或更新），最低CUDA 11.3；官方建议使用 CUDA 11.7+ 以获得更好性能。此外，V100 上反向内核的性能目前不如 A100 理想。扩展会在编译时同时校验上述条件（见is_compatible()），任一不满足即失败；如在无 GPU 环境做交叉编译，可设置DS_IGNORE_CUDA_DETECTION=TRUE跳过校验。在 op_builder/evoformer_attn.py 中可看到：低于sm_70的架构会被filter_ccs()剔除（Evoformer 内核依赖 Tensor Core）。

3.4 多架构打包构建（Multi-Arch Build）

Evoformer 现已支持直接通过TORCH_CUDA_ARCH_LIST进行混合架构打包，示例：

TORCH_CUDA_ARCH_LIST='7.0;8.0' \ DS_BUILD_OPS=0 DS_BUILD_EVOFORMER_ATTN=1 \ pip install -e .

说明与约束：

• TORCH_CUDA_ARCH_LIST控制生成的 CUDA 切片（各架构顺序无关）；
• 低于sm_70的目标会被自动剔除，因为 Evoformer 内核依赖 Tensor Core；
• DS_EVOFORMER_GPU_ARCH已弃用，Evoformer 构建会忽略它，请改用TORCH_CUDA_ARCH_LIST（构建器会在检测到该变量时打印弃用警告）。

不同架构家族支持的数据类型矩阵如下：

3.5 单元测试与基准

仓库内置了完整的功能测试与性能基准，可直接运行验证：

pytest -s tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py python tests/benchmarks/DS4Sci_EvoformerAttention_bench.py

• test_DS4Sci_EvoformerAttention.py 使用一个标准注意力参考实现（attention_reference：QK^T 乘缩放因子 → 叠加 bias → softmax → 与 V 相乘）做逐元素对比，覆盖 fp16/bf16 两种数据类型以及多种张量形状（如(1, 256, 256, 4, 32)、(1, 512, 256, 8, 8)），同时校验前向输出与 Q/K/V/bias 的梯度；测试文件中还以skip_on_arch依据数据类型跳过不支持 bf16 的低算力架构。
• test_evoformer_attn_builder.py 针对 OpBuilder 的兼容性判断做测试。
• DS4Sci_EvoformerAttention_bench.py 用 CUDA event 计时，分别采集定制内核与基线（PyTorch 参考实现）的前向/反向耗时，用于性能对比。

3.6 在自己模型中接入：核心 API 与四种注意力用法

在你的模型中使用时，从deepspeed.ops.deepspeed4science导入即可：

from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention

DS4Sci_EvoformerAttention通过统一的调用签名支持 Evoformer 的四种注意力机制（MSA 行方向、MSA 列方向，以及两种三角注意力）。其中以N_seq表示序列数、N_res表示残基数；隐藏维度Dim与头数Head因注意力类型而异。注意：输入张量必须是torch.float16或torch.bfloat16。

(a) MSA 行方向注意力（MSA row-wise attention）：为残基对构建注意力权重，并把 pair representation 的信息作为额外 bias 项融合进来：

# 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 列方向注意力（MSA column-wise attention）：让属于同一目标残基的元素交换信息：

# 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) 三角自注意力（Triangular self-attention）：更新 pair representation，分为 around-starting 与 around-ending 节点两种，下面给出 around-starting 节点的示例，around-ending 与之类似：

# 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])

3.7 接口约定与底层实现要点

封装实现位于 deepspeed/ops/deepspeed4science/evoformer_attn.py，理解它有助于规避使用陷阱：

• 函数签名实为DS4Sci_EvoformerAttention(Q, K, V, biases)，其中biases为长度不超过 2 的列表；代码会自动把缺失的 bias 槽位补为None（空 bias 场景），因此 MSA 列方向注意力只传[res_mask]也能工作；
• 代码对输入做形状断言：bias1（如res_mask）应为(Batch, N, 1, 1, N_res)形态、bias2（如pair_bias/right_edges）应为(Batch, 1, Head, N, N)形态，不合规会直接报错提示；
• 内部通过torch.autograd.Function（EvoformerFusedAttention）实现：前向保存 Q/K/V、输出 O 与以fp32 保存的lse（log-sum-exp）供反向使用，反向中lse先算出delta，再求 dQ/dK/dV 以及两个 bias 的梯度（bias 梯度默认以 fp32 累积后转回输入 dtype）；
• 前向要求序列长度> 16（seq_len must be greater than 16断言），而 Q 的最后一维会被向上取整到 32 的倍数参与内部计算；
• 反向路径要求隐藏维度不超过 64（对应内核中kMax常量），超过会触发Hidden size is too large断言；
• 所有张量（含 bias）必须位于当前加速器上，内核按需通过EvoformerAttnBuilder().load()首次加载编译。

在底层，对应源码位于 csrc/deepspeed4science/evoformer_attn/：attention.cpp提供算子绑定，attention_cu.cu与attention_back.cu分别承载前向/反向 launch 逻辑，kernel_forward.h与kernel_backward.h为核心 kernel 模板，并配套gemm/、iterators/、epilogue/、transform/等目录存放 CUTLASS 集成所需的 GEMM、迭代器与变换组件。构建入口则在 op_builder/evoformer_attn.py（EvoformerAttnBuilder，构建开关为DS_BUILD_EVOFORMER_ATTN），它会额外链接-lcurand（ROCm PyTorch 环境除外）。

3.8 科学应用：OpenFold 中消除训练显存爆炸

OpenFold 是 DeepMind AlphaFold2 的社区复现实现，允许在新数据集上训练/微调 AlphaFold2。AlphaFold2 训练存在显存爆炸问题——其包含多个自定义 Evoformer 注意力变体，会产生异常巨大的激活值。借助 DeepSpeed4Science 的DS4Sci_EvoformerAttention内核，OpenFold 团队在无精度损失的前提下将训练峰值显存需求降低了约 13 倍。

4. 技术选型与落地建议

接入DS4Sci_EvoformerAttention的最小步骤可归纳为：① 确认环境满足 GPU/CUDA/CUTLASS 前提，必要时按上文方式构建（含多架构打包）；② 从deepspeed.ops.deepspeed4science导入；③ 按四种注意力机制之一整理 Q/K/V 与 bias（注意形状断言、dtype 与"序列数/残基数布局随注意力类型改变"这两点）；④ 用 单元测试 验证正确性、用 基准脚本 量化收益。

5. 深入阅读指引

若需进一步查阅本仓库中的第一手材料，可沿以下路径展开：

• 总览页：docs/_pages/deepspeed4science.md（本篇文章依据的主体文档）；
• EvoformerAttention 完整操作教程：docs/_tutorials/ds4sci_evoformerattention.md；
• Python 调用封装与形状/前置断言：deepspeed/ops/deepspeed4science/evoformer_attn.py；
• CUTLASS 探测与编译兼容性逻辑：op_builder/evoformer_attn.py；
• CUDA/C++ 内核实现：csrc/deepspeed4science/evoformer_attn/；
• 正确性单元测试与构建器测试：test_DS4Sci_EvoformerAttention.py、test_evoformer_attn_builder.py；
• 性能基准：DS4Sci_EvoformerAttention_bench.py；
• 项目启动公告：blogs/deepspeed4science/README.md（含白皮书引用条目）。

综上，DeepSpeed4Science 的两项首批技术分别回应了 AI4Science 训练中的两类现实瓶颈：超长序列导致的大规模训练扩展问题，以及 Evoformer 类注意力带来的显存爆炸问题。前者解决"能不能把序列做得更长"，后者解决"能不能把显存省下来"——两者都已在实际科学模型（GenSLMs、OpenFold）中得到验证，并为后续科学场景的技术发布奠定了可复用的框架与算子基础。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed