DeepSpeed4Science 概述和教程

为了配合微软解决人类最紧迫挑战的使命，微软的 DeepSpeed 团队正在响应这一机会，推出了名为 DeepSpeed4Science 的新计划，旨在通过 AI 系统技术创新构建独特的能力，帮助领域专家解开当今最大的科学谜团。本页面作为 DeepSpeed4Science 计划中所有已发布（或将来将要发布）的技术的概述页面，方便科学家查找所需的技术。DeepSpeed4Science 计划的详细信息可以在 我们的网站 上找到。对于每项技术，我们将介绍其用途、何时使用它、如何使用它的链接以及现有技术在科学应用中的应用（我们欢迎用户在您将我们的技术应用于科学研究时贡献更多展示）。

要引用 DeepSpeed4Science，请引用我们的 白皮书

@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}
}

用于大规模 AI4Science 模型训练的新 Megatron-DeepSpeed

我们很荣幸地介绍 新 Megatron-DeepSpeed，它是一个用于大规模模型训练的更新框架。我们将 DeepSpeed 与最新的 Megatron-LM 重新整合，以支持长序列和其他许多功能。使用新的 Megatron-DeepSpeed，用户现在可以通过 ZeRO 风格的数据并行、张量并行、序列并行、流水线并行、模型状态卸载以及一些新添加的内存优化技术（如注意力掩码卸载和位置嵌入分区）的协同组合，训练其大型 AI4Science 模型（如 GenSLM） ，这些模型具有更长的序列。

new Megatron-DeepSpeed

该图描绘了使用我们的新 Megatron-DeepSpeed 框架训练具有 33B 个参数的 GPT 类模型，在支持长序列长度方面的系统能力。结果表明，新的 Megatron-DeepSpeed 使得序列长度比 NVIDIA 的 Megatron-LM 长 9 倍，而不会触发内存不足错误。

要查看新的 Megatron-DeepSpeed 如何帮助实现新的系统功能，例如训练具有超长序列长度的模型，请阅读我们的 教程。

同时，我们的新 Megatron-DeepSpeed 已经应用于基因组规模基础模型 GenSLM，它是一个 2022 年 ACM Gordon Bell 奖 获奖的基因组规模语言模型，来自阿贡国家实验室。为了实现他们的科学目标，GenSLM 和类似模型需要非常长的序列支持，无论是训练还是推理，都超出了通用 LLM 的长序列策略。通过利用 DeepSpeed4Science 的新 Megatron-DeepSpeed，GenSLM 团队能够以 512K 的序列长度训练他们的 25B 模型，远高于他们最初的 42K 序列长度。有关方法的详细信息可以在 我们的网站 上找到。GenSLM 团队还在他们的 示例 中展示了如何在 GenSLM 库中使用 DeepSpeed4Science。

内存高效的 EvoformerAttention 内核

Evoformer 是 DeepMind 的 AlphaFold 等科学模型的关键构建块。然而，EvoFormer 的多序列比对 (MSA) 注意力在训练/推理过程中经常遇到内存爆炸问题，例如在蛋白质结构预测模型中。现有的技术（如 FlashAttention）无法有效地支持 Evoformer，因为 EvoFormerAttention 使用的是行级/列级/三角形注意力，与需要自定义优化的标准 Transformer 自注意力和交叉注意力不同。为了缓解内存爆炸问题，我们引入了 DS4Sci_EvoformerAttention 内核，它是一组内核，可以提高 EvoFormer 变体的内存效率。 DS4Sci_EvoformerAttention 易于使用。要查看如何使用它，请参阅我们的 教程。

DS4Sci_EvoformerAttention 已经应用于 OpenFold，它是 DeepMind 的 AlphaFold2 的社区复现，它使得能够在新的数据集上训练或微调 AlphaFold2。借助 DS4Sci_EvoformerAttention 内核，OpenFold 团队能够将峰值内存需求降低 13 倍，而不会损失准确性。有关方法的详细信息可以在 我们的网站 上找到。

DS4Sci_EvoformerAttention

该图显示了 DeepSpeed 的 EvoFormerAttention 内核如何帮助将 OpenFold 训练的峰值内存需求降低 13 倍。