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/DeepSpeedDeepSpeed4Science 是 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-DeepSpeedGenSLMs 团队成功将其 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 的自动探测会依次查找nvidia-cutlassPython 包即cutlass_library模块源码目录Python 环境与 CMake 前缀CONDA_PREFIX、VIRTUAL_ENV、CMAKE_PREFIX_PATH、CUDA_HOME等环境变量编译器 include 路径环境变量CPATH、CPLUS_INCLUDE_PATH、C_INCLUDE_PATH位于 DeepSpeed 旁边或当前工作目录的cutlasscheckout常见系统安装前缀如/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 旁边克隆一份 CUTLASSNVIDIA/cutlass仓库git clone https://github.com/NVIDIA/cutlass内核会在DS4Sci_EvoformerAttention首次被调用时触发编译。若自动探测没有找到预期安装可通过CUTLASS_PATH显式指定 CUTLASS checkout 根目录或其include目录。若你的编译环境已妥善配置好 CUTLASS 与 CUDA 编译器例如在 conda 包中编译可用CUTLASS_PATHDS_IGNORE_CUTLASS_DETECTION跳过 CUTLASS 探测。另有DS_USE_CUTLASS_PYTHON_BINDINGS对应nvidia-cutlassPython 包源码探测开关定义于构建器中。硬件与软件下限需要计算能力 7.0 及以上的 GPUNVIDIA V100 或更新最低CUDA 11.3官方建议使用 CUDA 11.7 以获得更好性能。此外V100 上反向内核的性能目前不如 A100 理想。扩展会在编译时同时校验上述条件见is_compatible()任一不满足即失败如在无 GPU 环境做交叉编译可设置DS_IGNORE_CUDA_DETECTIONTRUE跳过校验。在 op_builder/evoformer_attn.py 中可看到低于sm_70的架构会被filter_ccs()剔除Evoformer 内核依赖 Tensor Core。3.4 多架构打包构建Multi-Arch BuildEvoformer 现已支持直接通过TORCH_CUDA_ARCH_LIST进行混合架构打包示例TORCH_CUDA_ARCH_LIST7.0;8.0 \ DS_BUILD_OPS0 DS_BUILD_EVOFORMER_ATTN1 \ pip install -e .说明与约束TORCH_CUDA_ARCH_LIST控制生成的 CUDA 切片各架构顺序无关低于sm_70的目标会被自动剔除因为 Evoformer 内核依赖 Tensor CoreDS_EVOFORMER_GPU_ARCH已弃用Evoformer 构建会忽略它请改用TORCH_CUDA_ARCH_LIST构建器会在检测到该变量时打印弃用警告。不同架构家族支持的数据类型矩阵如下架构家族fp16bf16Sm70Volta支持不支持Sm75Turing支持不支持Sm80Ampere/Ada/Hopper支持支持3.5 单元测试与基准仓库内置了完整的功能测试与性能基准可直接运行验证pytest -s tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py python tests/benchmarks/DS4Sci_EvoformerAttention_bench.pytest_DS4Sci_EvoformerAttention.py 使用一个标准注意力参考实现attention_referenceQK^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_EvoformerAttentionDS4Sci_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.FunctionEvoformerFusedAttention实现前向保存 Q/K/V、输出 O 与以fp32 保存的lselog-sum-exp供反向使用反向中lse先算出delta再求 dQ/dK/dV 以及两个 bias 的梯度bias 梯度默认以 fp32 累积后转回输入 dtype前向要求序列长度 16seq_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.pyEvoformerAttnBuilder构建开关为DS_BUILD_EVOFORMER_ATTN它会额外链接-lcurandROCm PyTorch 环境除外。3.8 科学应用OpenFold 中消除训练显存爆炸OpenFold 是 DeepMind AlphaFold2 的社区复现实现允许在新数据集上训练/微调 AlphaFold2。AlphaFold2 训练存在显存爆炸问题——其包含多个自定义 Evoformer 注意力变体会产生异常巨大的激活值。借助 DeepSpeed4Science 的DS4Sci_EvoformerAttention内核OpenFold 团队在无精度损失的前提下将训练峰值显存需求降低了约 13 倍。4. 技术选型与落地建议需求推荐技术应用位置 / 前置条件AI4Science 大模型如基因组/蛋白质语言模型需要远超通用大模型的长序列训练新版 Megatron-DeepSpeed独立 Megatron-DeepSpeed 框架组合 ZeRO 式数据并行、张量/序列/流水并行、状态卸载及 attention mask 卸载、位置编码切分等可参考 33B 模型 9 倍序列长度与 GenSLMs 25B/512K 序列的应用案例Evoformer 类模型结构生物学MSA/三角注意力显存爆炸DS4Sci_EvoformerAttention随 DeepSpeed 主仓库发布要求 GPU 计算能力 ≥ 7.0、CUDA ≥ 11.3训练推荐配合反向内核用于 MSA 行/列方向注意力Evoformer 注意力推理加速DS4Sci_EvoformerAttention前向内核支持四种 Evoformer 注意力形态输入需 fp16/bf16接入DS4Sci_EvoformerAttention的最小步骤可归纳为① 确认环境满足 GPU/CUDA/CUTLASS 前提必要时按上文方式构建含多架构打包② 从deepspeed.ops.deepspeed4science导入③ 按四种注意力机制之一整理 Q/K/V 与 bias注意形状断言、dtype 与序列数/残基数布局随注意力类型改变这两点④ 用 单元测试 验证正确性、用 基准脚本 量化收益。5. 深入阅读指引若需进一步查阅本仓库中的第一手材料可沿以下路径展开总览页docs/_pages/deepspeed4science.md本篇文章依据的主体文档EvoformerAttention 完整操作教程docs/_tutorials/ds4sci_evoformerattention.mdPython 调用封装与形状/前置断言deepspeed/ops/deepspeed4science/evoformer_attn.pyCUTLASS 探测与编译兼容性逻辑op_builder/evoformer_attn.pyCUDA/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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考