DeepSpeed4Science 深度解读以 AI 系统优化技术驱动大规模科学发现【免费下载链接】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 团队发起的 AI4Science 专项计划其核心目标不是继续堆叠大模型规模而是围绕结构生物学、气候科学、基因组学等科学领域独有的计算形态如带 Mask/Bias 的特殊注意力、超长序列基因组建模开发专属的 AI 系统优化技术。本文基于仓库内该计划的官方中文资料 blogs/deepspeed4science/chinese/README.md结合本仓库中DS4Sci_EvoformerAttention的 Python 封装、CUDA/C 内核、算子构建器 与 官方教程讲解其设计动机、两大标志性成果以及如何在本地环境中编译、调用与验证这些技术。1. DeepSpeed4Science为科学模型而非通用大模型而生的系统技术随着深度学习在过去十年对自然科学的渗透从药物开发到可再生能源AI 正在逼近通过建模与预测自然现象来加速科学发现的临界点。DeepSpeed 作为面向多类 AI 硬件的开源训练/推理框架为大规模深度学习提供了规模与速度基础DeepSpeed4Science 则在这一基础之上进一步回答一个问题用于加速通用 LLM 的系统技术并不能原样迁移到科学 AI 模型上。科学模型的计算模式往往与 LLM 显著不同例如结构生物学中的 Evoformer 拥有高度特化的注意力变体基因组模型需要远超常规文本模型长度的上下文窗口。通用优化如 FlashAttention在这些场景下常常失效。因此 DeepSpeed4Science 以 DeepSpeed 已有的训练、推理与模型压缩能力作为基础使能技术再针对性开发面向科学发现的 AI 系统技术覆盖气候科学、药物设计、生物学理解、分子动力学模拟、癌症诊断与监测、催化剂/材料发现等领域并采用开放、包容的设计期望最终沉淀为一个分享先进 AI4Science 系统技术的软件平台与统一代码仓库。2. 早期合作版图从微软研究院到美国国家实验室DeepSpeed4Science 的技术验证高度依赖真实科学模型 真实科研团队的联合打磨仓库文档记录了两类合作者。2.1 微软内部合作伙伴科学基础模型SFM微软研究院 AI4Science旨在构建统一的大规模科学基础模型支持药物、材料、生物学、健康等多领域、多输入与多计算任务DeepSpeed4Science 为其生成式 AI 方法如面向分子系统平衡分布预测的 Distributional Graphormer提供训练与推理技术支持。ClimaX微软研究院 AI4Science文档将其定位为第一个可执行多种天气与气候建模任务的基础模型可融合不同变量、不同分辨率的数据集以提升预报精度。挑战在于其训练需要处理数十到数百 PB 级的高分辨率图像数据与长序列DeepSpeed4Science 负责为其预训练/微调更大模型提供系统支持与加速策略。分子动力学与机器学习力场微软研究院 AI4Science用 AI 力场模型逼近第一性原理精度来模拟百万原子级分子系统同时保持经典分子动力学的效率与可扩展性由于需要数百万乃至数十亿步推理对 GNNLLM 混合模型的推理速度提出了很高的系统要求。微软天气微软 WebXT/Bing此前已受益于 DeepSpeed 的多 GPU 训练加速DeepSpeed4Science 继续与 WebXT 天气预报团队合作增强预报服务。2.2 外部合作者计划的首批两大外部模型均来自结构生物学领域哥伦比亚大学的OpenFoldAlphaFold2 的高保真开源复现与阿贡国家实验室的GenSLMs2022 年 ACM 戈登贝尔奖获奖 COVID 基因组模型。此外文档还记载了与阿贡国家实验室在 Aurora 超算上训练万亿参数科学模型、与橡树岭国家实验室和国家癌症研究所NCI合作开展 MOSSAIC 癌症监测从非结构化临床文本中高保真提取与分类信息、以及 Brookhaven 国家实验室用 LLM 构建清洁能源数字孪生模型等扩展工作。3. 展示一DS4Sci_EvoformerAttention 如何消除 Evoformer 内存爆炸3.1 问题93M 参数的模型为何会内存爆炸OpenFold 使研究人员可以在新数据集上训练或微调 AlphaFold2。但官方博客指出从头训练 AlphaFold2 依然极其昂贵模型本身仅有9300 万参数却内置多个需要巨大中间内存的特化注意力变体。在标准 AlphaFold2 训练的微调阶段仅其中一个变体在半精度下就会产生超过 12GB 的张量峰值内存需求远超同等规模的 LLM。即便配合 activation checkpointing 与 DeepSpeed ZeRO 优化内存爆炸仍然严重限制可训练的序列长度与 MSA多序列比对深度而采用近似策略又可能影响模型精度与收敛性。问题根源在于这类带 Bias 的 EvoformerAttention 变体无法被 FlashAttention 这类通用手段高效支持。DeepSpeed4Science 的解法不是近似而是为这类在结构生物学模型中广泛出现的注意力变体设计精确的定制内核——由复杂的融合与矩阵分块策略、动态内存减少方法组成的高内存效率内核DS4Sci_EvoformerAttention作为高质量机器学习模块开放给生物学研究社区。3.2 仓库中的源码级实现该内核在本仓库内有完整的端到端落地可以逐层追踪Python 前端deepspeed/ops/deepspeed4science/evoformer_attn.py 暴露了DS4Sci_EvoformerAttention(Q, K, V, biases)入口支持最多两个 bias 项内部校验 bias 形状bias1形如[B, N, 1, 1, L]bias2形如[B, 1, H, L, L]并断言序列长度大于 16。Autograd 封装同文件中的EvoformerFusedAttention(torch.autograd.Function)在forward中调用_attention得到输出O与用于数值稳定的lselog-sum-exp并存入上下文backward中通过attention_bwd计算dQ/dK/dV以及两个 bias 的梯度可配置是否求 bias 梯度。CUDA 内核csrc/deepspeed4science/evoformer_attn/ 包含attention.cpp绑定、attention_cu.cu前向、attention_back.cu反向以及kernel_forward.h/kernel_backward.h、gemm/、epilogue/、iterators/、transform/等 CUTLASS 风格的分层实现目录可见其基于 CUTLASS 的 GEMM/Attention 融合设计。从代码注释看后向当前将隐藏维度上限约束在 64 以内。构建入口op_builder/evoformer_attn.py 中的EvoformerAttnBuilder定义算子名为evoformer_attn、编译开关为DS_BUILD_EVOFORMER_ATTN构建时只保留计算能力 ≥ 7.0Tensor Core 必需的目标低于 SM 70 会被剔除并会校验 CUTLASS 版本不低于 3.1.0 与 CUDA 11。单元测试tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py 与test_evoformer_attn_builder.py覆盖了内核数值正确性与构建器行为。3.3 何时使用按照仓库内 ds4sci_evoformerattention 教程 的说明当序列数与残基数很大时该内核收益最明显前向内核针对加速计算做了专门优化适合在推理期使用反向内核则用于训练期以计算开销换取内存占用的大幅下降。因此它最适用于训练中内存受限的运算例如MSA 行注意力与 MSA 列注意力。4. 展示二为基因组基础模型GenSLMs提供长序列支持4.1 为什么基因组模型需要超长序列GenSLMs 是阿贡国家实验室获 2022 年 ACM 戈登贝尔奖的基因组模型通过在大规模基因组数据上训练 LLM 来学习 SARS-CoV-2COVID-19基因组的进化目标是从序列中识别与分类新发变异株并泛化到细菌病原体乃至真核生物的预测任务。要达到这一目标模型需要远超通用 LLM 长序列策略如 FlashAttention所能覆盖的上下文窗口用于训练与推理。4.2 系统层面的解决方案与量化收益DeepSpeed4Science 在系统层面发布了对长序列与其他新优化支持的 Megatron-DeepSpeed 框架该框架位于 DeepSpeed 仓库之外本文仅转述官方博客记载的成果其要点是通过注意力掩码异步处理、位置编码切分以及张量并行、流水线并行、序列并行、基于 ZeRO 的数据并行与模型状态异步处理等技术组合让科学家能够以更长的上下文训练 GenSLMs 这类大型科学模型。据博客数据新版本使 GenSLMs 25B 与 33B 模型可支持的最大序列长度分别比旧版 Megatron-DeepSpeed提升约 12 倍与 14 倍在支持的序列长度上也显著超过 NVIDIA Megatron-LM25B/33B 模型分别最高约9.8 倍与 9.1 倍。例如阿贡团队此前在 64 块 GPU 上以 42K 的原始序列长度训练 GenSLMs 25B如今可扩展为512K 核苷酸序列训练且不损失准确性。对于偏好相对位置编码等算法策略的领域科学家该版本也做了相应集成。下图对比了不同框架在不同规模下对两个 GenSLMs 模型支持的最大序列长度每个节点 8 块 40G A100 的 NVIDIA DGX。5. 动手实践在 DeepSpeed 仓库中使用 DS4Sci_EvoformerAttention本仓库内的 ds4sci_evoformerattention 教程 提供了可直接复现的完整流程要点整理如下。5.1 依赖与硬件前提DS4Sci_EvoformerAttention随 DeepSpeed ≥ 0.10.3 发布。内核基于 NVIDIA CUTLASS 实现。DeepSpeed 会自动在以下位置寻找 CUTLASSnvidia-cutlassPython 包、Python 环境与 CMake 前缀、编译器 include 路径环境变量、DeepSpeed 旁或当前工作目录下的cutlass检出以及/usr/local等常见系统安装前缀。若自动探测失败可显式设置CUTLASS_PATH指向 CUTLASS 检出根目录或其include目录。硬件要求 GPU 计算能力≥ 7.0V100 及更新最低 CUDA 11.3官方建议使用 CUDA 11.7 以获得更好性能当前 V100 上反向内核性能不及 A100。构建器会对这些条件做检查若在不带 GPU 的环境交叉编译可设置DS_IGNORE_CUDA_DETECTIONTRUE关闭检查。5.2 多架构编译从构建器实现看Evoformer 构建通过TORCH_CUDA_ARCH_LIST支持混合架构打包低于sm_70的目标会被自动剔除且旧的DS_EVOFORMER_GPU_ARCH环境变量已废弃。教程给出的编译示例为TORCH_CUDA_ARCH_LIST7.0;8.0 \ DS_BUILD_OPS0 DS_BUILD_EVOFORMER_ATTN1 \ pip install -e .内核支持的数据类型与架构族对应关系为SM70Volta与 SM75Turing支持 fp16SM80Ampere/Ada/Hopper同时支持 fp16 与 bf16。输入张量需为torch.float16或torch.bfloat16。5.3 在自己的模型中调用Evoformer 中存在四种注意力机制MSA 行、MSA 列以及两种三角注意力DS4Sci_EvoformerAttention通过传入不同的 Q/K/V 形状与 bias 组合统一覆盖它们。设序列数为N_seq、残基数为N_res、隐藏维度为Dim、头数为Headfrom deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention # (a) MSA 行注意力构建残基对间的注意力权重并把 pair 表示作为附加 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 列注意力让属于同一目标残基的元素交换信息 # Q, K, V: [Batch, N_res, N_seq, Head, Dim] out DS4Sci_EvoformerAttention(Q, K, V, [res_mask]) # (c) 三角自注意力starting node更新 pair 表示ending node 用法类似 # Q, K, V: [Batch, N_res, N_res, Head, Dim] # right_edges: [Batch, 1, Head, N_res, N_res] out DS4Sci_EvoformerAttention(Q, K, V, [res_mask, right_edges])5.4 运行单元测试与基准仓库内提供了单元测试与基准脚本可用如下命令运行pytest -s tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py python tests/benchmarks/DS4Sci_EvoformerAttention_bench.py在 OpenFold 上的应用效果官方博客报告峰值内存最高降低 13 倍且无精度损失说明借助这些定制内核研究者可以用更长的序列、更大更复杂的模型在更广泛的硬件上进行结构生物学实验。6. 总结与路线图DeepSpeed4Science 计划从发布之日起逐步公开其对外合作信息与阶段性技术成果其高层目标是推广能够系统性解决大规模科学发现核心痛点的 AI 系统技术并通过开源软件让全球科学家直接受益。本仓库作为 DeepSpeed 主仓库已经承载了展示一中的核心技术——DS4Sci_EvoformerAttention的 Python 封装、CUDA 实现、构建器与单元测试——意味着这套系统技术不是停留在博客中的概念而是可以直接安装、编译、测试并集成到 OpenFold 等外部科学模型中的可用代码。若需在论文中引用 DeepSpeed4Science官方博客建议引用其技术白皮书arXiv 2310.04610Song et al., 2023。进一步阅读完整的英文原文与配套图片见 blogs/deepspeed4science/chinese/README.md 同目录的英文版本内核使用指南见 docs/_tutorials/ds4sci_evoformerattention.mdPython 调用层见 deepspeed/ops/deepspeed4science/evoformer_attn.py底层实现与构建配置见 csrc/deepspeed4science/evoformer_attn/ 与 op_builder/evoformer_attn.py。【免费下载链接】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），仅供参考 DeepSpeed4Science 深度解读：以 AI 系统优化技术驱动大规模科学发现
DeepSpeed4Science 深度解读以 AI 系统优化技术驱动大规模科学发现【免费下载链接】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 团队发起的 AI4Science 专项计划其核心目标不是继续堆叠大模型规模而是围绕结构生物学、气候科学、基因组学等科学领域独有的计算形态如带 Mask/Bias 的特殊注意力、超长序列基因组建模开发专属的 AI 系统优化技术。本文基于仓库内该计划的官方中文资料 blogs/deepspeed4science/chinese/README.md结合本仓库中DS4Sci_EvoformerAttention的 Python 封装、CUDA/C 内核、算子构建器 与 官方教程讲解其设计动机、两大标志性成果以及如何在本地环境中编译、调用与验证这些技术。1. DeepSpeed4Science为科学模型而非通用大模型而生的系统技术随着深度学习在过去十年对自然科学的渗透从药物开发到可再生能源AI 正在逼近通过建模与预测自然现象来加速科学发现的临界点。DeepSpeed 作为面向多类 AI 硬件的开源训练/推理框架为大规模深度学习提供了规模与速度基础DeepSpeed4Science 则在这一基础之上进一步回答一个问题用于加速通用 LLM 的系统技术并不能原样迁移到科学 AI 模型上。科学模型的计算模式往往与 LLM 显著不同例如结构生物学中的 Evoformer 拥有高度特化的注意力变体基因组模型需要远超常规文本模型长度的上下文窗口。通用优化如 FlashAttention在这些场景下常常失效。因此 DeepSpeed4Science 以 DeepSpeed 已有的训练、推理与模型压缩能力作为基础使能技术再针对性开发面向科学发现的 AI 系统技术覆盖气候科学、药物设计、生物学理解、分子动力学模拟、癌症诊断与监测、催化剂/材料发现等领域并采用开放、包容的设计期望最终沉淀为一个分享先进 AI4Science 系统技术的软件平台与统一代码仓库。2. 早期合作版图从微软研究院到美国国家实验室DeepSpeed4Science 的技术验证高度依赖真实科学模型 真实科研团队的联合打磨仓库文档记录了两类合作者。2.1 微软内部合作伙伴科学基础模型SFM微软研究院 AI4Science旨在构建统一的大规模科学基础模型支持药物、材料、生物学、健康等多领域、多输入与多计算任务DeepSpeed4Science 为其生成式 AI 方法如面向分子系统平衡分布预测的 Distributional Graphormer提供训练与推理技术支持。ClimaX微软研究院 AI4Science文档将其定位为第一个可执行多种天气与气候建模任务的基础模型可融合不同变量、不同分辨率的数据集以提升预报精度。挑战在于其训练需要处理数十到数百 PB 级的高分辨率图像数据与长序列DeepSpeed4Science 负责为其预训练/微调更大模型提供系统支持与加速策略。分子动力学与机器学习力场微软研究院 AI4Science用 AI 力场模型逼近第一性原理精度来模拟百万原子级分子系统同时保持经典分子动力学的效率与可扩展性由于需要数百万乃至数十亿步推理对 GNNLLM 混合模型的推理速度提出了很高的系统要求。微软天气微软 WebXT/Bing此前已受益于 DeepSpeed 的多 GPU 训练加速DeepSpeed4Science 继续与 WebXT 天气预报团队合作增强预报服务。2.2 外部合作者计划的首批两大外部模型均来自结构生物学领域哥伦比亚大学的OpenFoldAlphaFold2 的高保真开源复现与阿贡国家实验室的GenSLMs2022 年 ACM 戈登贝尔奖获奖 COVID 基因组模型。此外文档还记载了与阿贡国家实验室在 Aurora 超算上训练万亿参数科学模型、与橡树岭国家实验室和国家癌症研究所NCI合作开展 MOSSAIC 癌症监测从非结构化临床文本中高保真提取与分类信息、以及 Brookhaven 国家实验室用 LLM 构建清洁能源数字孪生模型等扩展工作。3. 展示一DS4Sci_EvoformerAttention 如何消除 Evoformer 内存爆炸3.1 问题93M 参数的模型为何会内存爆炸OpenFold 使研究人员可以在新数据集上训练或微调 AlphaFold2。但官方博客指出从头训练 AlphaFold2 依然极其昂贵模型本身仅有9300 万参数却内置多个需要巨大中间内存的特化注意力变体。在标准 AlphaFold2 训练的微调阶段仅其中一个变体在半精度下就会产生超过 12GB 的张量峰值内存需求远超同等规模的 LLM。即便配合 activation checkpointing 与 DeepSpeed ZeRO 优化内存爆炸仍然严重限制可训练的序列长度与 MSA多序列比对深度而采用近似策略又可能影响模型精度与收敛性。问题根源在于这类带 Bias 的 EvoformerAttention 变体无法被 FlashAttention 这类通用手段高效支持。DeepSpeed4Science 的解法不是近似而是为这类在结构生物学模型中广泛出现的注意力变体设计精确的定制内核——由复杂的融合与矩阵分块策略、动态内存减少方法组成的高内存效率内核DS4Sci_EvoformerAttention作为高质量机器学习模块开放给生物学研究社区。3.2 仓库中的源码级实现该内核在本仓库内有完整的端到端落地可以逐层追踪Python 前端deepspeed/ops/deepspeed4science/evoformer_attn.py 暴露了DS4Sci_EvoformerAttention(Q, K, V, biases)入口支持最多两个 bias 项内部校验 bias 形状bias1形如[B, N, 1, 1, L]bias2形如[B, 1, H, L, L]并断言序列长度大于 16。Autograd 封装同文件中的EvoformerFusedAttention(torch.autograd.Function)在forward中调用_attention得到输出O与用于数值稳定的lselog-sum-exp并存入上下文backward中通过attention_bwd计算dQ/dK/dV以及两个 bias 的梯度可配置是否求 bias 梯度。CUDA 内核csrc/deepspeed4science/evoformer_attn/ 包含attention.cpp绑定、attention_cu.cu前向、attention_back.cu反向以及kernel_forward.h/kernel_backward.h、gemm/、epilogue/、iterators/、transform/等 CUTLASS 风格的分层实现目录可见其基于 CUTLASS 的 GEMM/Attention 融合设计。从代码注释看后向当前将隐藏维度上限约束在 64 以内。构建入口op_builder/evoformer_attn.py 中的EvoformerAttnBuilder定义算子名为evoformer_attn、编译开关为DS_BUILD_EVOFORMER_ATTN构建时只保留计算能力 ≥ 7.0Tensor Core 必需的目标低于 SM 70 会被剔除并会校验 CUTLASS 版本不低于 3.1.0 与 CUDA 11。单元测试tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py 与test_evoformer_attn_builder.py覆盖了内核数值正确性与构建器行为。3.3 何时使用按照仓库内 ds4sci_evoformerattention 教程 的说明当序列数与残基数很大时该内核收益最明显前向内核针对加速计算做了专门优化适合在推理期使用反向内核则用于训练期以计算开销换取内存占用的大幅下降。因此它最适用于训练中内存受限的运算例如MSA 行注意力与 MSA 列注意力。4. 展示二为基因组基础模型GenSLMs提供长序列支持4.1 为什么基因组模型需要超长序列GenSLMs 是阿贡国家实验室获 2022 年 ACM 戈登贝尔奖的基因组模型通过在大规模基因组数据上训练 LLM 来学习 SARS-CoV-2COVID-19基因组的进化目标是从序列中识别与分类新发变异株并泛化到细菌病原体乃至真核生物的预测任务。要达到这一目标模型需要远超通用 LLM 长序列策略如 FlashAttention所能覆盖的上下文窗口用于训练与推理。4.2 系统层面的解决方案与量化收益DeepSpeed4Science 在系统层面发布了对长序列与其他新优化支持的 Megatron-DeepSpeed 框架该框架位于 DeepSpeed 仓库之外本文仅转述官方博客记载的成果其要点是通过注意力掩码异步处理、位置编码切分以及张量并行、流水线并行、序列并行、基于 ZeRO 的数据并行与模型状态异步处理等技术组合让科学家能够以更长的上下文训练 GenSLMs 这类大型科学模型。据博客数据新版本使 GenSLMs 25B 与 33B 模型可支持的最大序列长度分别比旧版 Megatron-DeepSpeed提升约 12 倍与 14 倍在支持的序列长度上也显著超过 NVIDIA Megatron-LM25B/33B 模型分别最高约9.8 倍与 9.1 倍。例如阿贡团队此前在 64 块 GPU 上以 42K 的原始序列长度训练 GenSLMs 25B如今可扩展为512K 核苷酸序列训练且不损失准确性。对于偏好相对位置编码等算法策略的领域科学家该版本也做了相应集成。下图对比了不同框架在不同规模下对两个 GenSLMs 模型支持的最大序列长度每个节点 8 块 40G A100 的 NVIDIA DGX。5. 动手实践在 DeepSpeed 仓库中使用 DS4Sci_EvoformerAttention本仓库内的 ds4sci_evoformerattention 教程 提供了可直接复现的完整流程要点整理如下。5.1 依赖与硬件前提DS4Sci_EvoformerAttention随 DeepSpeed ≥ 0.10.3 发布。内核基于 NVIDIA CUTLASS 实现。DeepSpeed 会自动在以下位置寻找 CUTLASSnvidia-cutlassPython 包、Python 环境与 CMake 前缀、编译器 include 路径环境变量、DeepSpeed 旁或当前工作目录下的cutlass检出以及/usr/local等常见系统安装前缀。若自动探测失败可显式设置CUTLASS_PATH指向 CUTLASS 检出根目录或其include目录。硬件要求 GPU 计算能力≥ 7.0V100 及更新最低 CUDA 11.3官方建议使用 CUDA 11.7 以获得更好性能当前 V100 上反向内核性能不及 A100。构建器会对这些条件做检查若在不带 GPU 的环境交叉编译可设置DS_IGNORE_CUDA_DETECTIONTRUE关闭检查。5.2 多架构编译从构建器实现看Evoformer 构建通过TORCH_CUDA_ARCH_LIST支持混合架构打包低于sm_70的目标会被自动剔除且旧的DS_EVOFORMER_GPU_ARCH环境变量已废弃。教程给出的编译示例为TORCH_CUDA_ARCH_LIST7.0;8.0 \ DS_BUILD_OPS0 DS_BUILD_EVOFORMER_ATTN1 \ pip install -e .内核支持的数据类型与架构族对应关系为SM70Volta与 SM75Turing支持 fp16SM80Ampere/Ada/Hopper同时支持 fp16 与 bf16。输入张量需为torch.float16或torch.bfloat16。5.3 在自己的模型中调用Evoformer 中存在四种注意力机制MSA 行、MSA 列以及两种三角注意力DS4Sci_EvoformerAttention通过传入不同的 Q/K/V 形状与 bias 组合统一覆盖它们。设序列数为N_seq、残基数为N_res、隐藏维度为Dim、头数为Headfrom deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention # (a) MSA 行注意力构建残基对间的注意力权重并把 pair 表示作为附加 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 列注意力让属于同一目标残基的元素交换信息 # Q, K, V: [Batch, N_res, N_seq, Head, Dim] out DS4Sci_EvoformerAttention(Q, K, V, [res_mask]) # (c) 三角自注意力starting node更新 pair 表示ending node 用法类似 # Q, K, V: [Batch, N_res, N_res, Head, Dim] # right_edges: [Batch, 1, Head, N_res, N_res] out DS4Sci_EvoformerAttention(Q, K, V, [res_mask, right_edges])5.4 运行单元测试与基准仓库内提供了单元测试与基准脚本可用如下命令运行pytest -s tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py python tests/benchmarks/DS4Sci_EvoformerAttention_bench.py在 OpenFold 上的应用效果官方博客报告峰值内存最高降低 13 倍且无精度损失说明借助这些定制内核研究者可以用更长的序列、更大更复杂的模型在更广泛的硬件上进行结构生物学实验。6. 总结与路线图DeepSpeed4Science 计划从发布之日起逐步公开其对外合作信息与阶段性技术成果其高层目标是推广能够系统性解决大规模科学发现核心痛点的 AI 系统技术并通过开源软件让全球科学家直接受益。本仓库作为 DeepSpeed 主仓库已经承载了展示一中的核心技术——DS4Sci_EvoformerAttention的 Python 封装、CUDA 实现、构建器与单元测试——意味着这套系统技术不是停留在博客中的概念而是可以直接安装、编译、测试并集成到 OpenFold 等外部科学模型中的可用代码。若需在论文中引用 DeepSpeed4Science官方博客建议引用其技术白皮书arXiv 2310.04610Song et al., 2023。进一步阅读完整的英文原文与配套图片见 blogs/deepspeed4science/chinese/README.md 同目录的英文版本内核使用指南见 docs/_tutorials/ds4sci_evoformerattention.mdPython 调用层见 deepspeed/ops/deepspeed4science/evoformer_attn.py底层实现与构建配置见 csrc/deepspeed4science/evoformer_attn/ 与 op_builder/evoformer_attn.py。【免费下载链接】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），仅供参考