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模型，如GenSLMs，使用更长的序列。

new Megatron-DeepSpeed

该图展示了系统能力，即在训练一个33B参数的类GPT模型时，使用我们新的Megatron-DeepSpeed框架能够支持更长的序列长度。结果显示，新的Megatron-DeepSpeed能够支持比NVIDIA的Megatron-LM长9倍的序列长度，而不会触发内存不足错误。

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

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