DeepSpeed4Science 概述与教程

秉承微软解决人类最紧迫挑战的使命，微软 DeepSpeed 团队正通过一项名为 DeepSpeed4Science 的新计划来应对这一机遇，旨在通过人工智能系统技术创新构建独特的能力，帮助领域专家解锁当今最大的科学奥秘。本页面将作为 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}
}

新的 Megatron-DeepSpeed，用于大规模科学人工智能模型训练

我们很荣幸推出新的 Megatron-DeepSpeed，这是一个用于大规模模型训练的更新框架。我们重新基于并使用最新的 Megatron-LM 增强了 DeepSpeed，以支持长序列和许多其他功能。借助新的 Megatron-DeepSpeed，用户现在可以通过 ZeRO 风格的数据并行、张量并行、序列并行、流水线并行、模型状态卸载以及注意力掩码卸载和位置嵌入分区等多种新添加的内存优化技术的协同组合，训练像 GenSLMs 这样具有更长序列的大型科学人工智能模型。

new Megatron-DeepSpeed

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

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

同时，我们新的 Megatron-DeepSpeed 已应用于基因组尺度基础模型GenSLMs，这是来自阿贡国家实验室（Argonne National Lab）并荣获 2022 年 ACM 戈登贝尔奖的基因组尺度语言模型。为了实现其科学目标，GenSLMs 及类似模型在训练和推理时需要超长序列支持，这超出了通用大型语言模型的长序列策略。通过利用 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 倍。