[论文解读] DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies

Shuaiwen Leon Song, Bonnie Kruft+90 more|arXiv (Cornell University)|Oct 6, 2023

Machine Learning in Materials Science被引用 8

一句话总结

本文介绍 DeepSpeed4Science 计划， detailing 基于 DeepSpeed 的 AI 系统技术以加速大规模科学发现，包含两场结构生物学展示以及面向更广科学协作的计划。

ABSTRACT

In the upcoming decade, deep learning may revolutionize the natural sciences, enhancing our capacity to model and predict natural occurrences. This could herald a new era of scientific exploration, bringing significant advancements across sectors from drug development to renewable energy. To answer this call, we present DeepSpeed4Science initiative (deepspeed4science.ai) which aims to build unique capabilities through AI system technology innovations to help domain experts to unlock today's biggest science mysteries. By leveraging DeepSpeed's current technology pillars (training, inference and compression) as base technology enablers, DeepSpeed4Science will create a new set of AI system technologies tailored for accelerating scientific discoveries by addressing their unique complexity beyond the common technical approaches used for accelerating generic large language models (LLMs). In this paper, we showcase the early progress we made with DeepSpeed4Science in addressing two of the critical system challenges in structural biology research.

研究动机与目标

• 为科学领域定制的 AI 系统技术的需求提供动机，超越通用的大模型加速。
• 描述 DeepSpeed4Science 的方法及其基于 DeepSpeed 的支柱（训练、推理、压缩）。
• 展示 DS4Sci 针对结构生物学的两个系统挑战（Evoformer 注意力的内存爆炸；GenSLMs 的长序列支持）。
• 概述协作模型以及用于分享科学领域 AI 系统技术的平台潜力。

提出的方法

• 开发定制的内存高效 EvoformerAttention 内核，以消除以 Evoformer 为中心的模型的内存爆炸。
• 通过内存优化实现对注意力掩码和位置嵌入的长序列支持，集成增强的 Megatron-DeepSpeed 框架。
• 通过 Megatron-DeepSpeed 重构和优化，使基因组规模的基础模型能够进行非常长序列的训练/推理。
• 融合内核并采用分块、在线广播，以及 FP32 安全梯度处理，在保持精度的同时降低峰值内存。
• 利用序列并行、张量/流水线并行以及模型/数据卸载，显著扩展序列长度。

Figure 1: DeepSpeed4Science approach: developing a new set of AI system technologies that are beyond generic large language model support, tailored for accelerating scientific discoveries and addressing their complexity.

实验结果

研究问题

• RQ1 AI 系统技术如何专门应对科学焦点模型的内存和序列长度挑战？
• RQ2 定制内核和框架重构能否在不损失精度的前提下，实现基因组规模和 Evoformer 基模型的显著更长的上下文？
• RQ3 将 DS4Sci 的优化应用于结构生物学和 GenSLM 风格模型时，性能/吞吐量的提升是什么？
• RQ4 DS4Sci 如何促进更广泛的协作与分享用于科学发现的先进 AI 系统技术？

主要发现

• 一个 DS4Sci_EvoformerAttention 内核使 OpenFold 的峰值内存在 Evoformer-注意力变体上降低 13 倍且不损失精度。
• 新的 Megatron-DeepSpeed 框架使 GenSLMs 的训练具备显著更长的序列，并在某些场景中平均序列长度提升最多 13 倍，吞吐量提高至 2 倍。
• Megatron-DeepSpeed 重构增加了旋转位置嵌入、FlashAttention v1/v2 以及新的融合内核以改进长序列训练和推理。
• 对注意力掩码和位置嵌入的内存优化，加上序列并行，显著扩展 GenSLMs 的可行序列长度（例如 25B GenSLM 可达 512K），超越此前极限。
• DS4Sci 的努力将 DeepSpeed4Science 定位为一个分享科学领域先进 AI 系统技术的平台与知识库。

Figure 2: Peak memory requirement for training variants of the MSA attention kernels (with bias) with the maximum possible training sample dimension in OpenFold. (Left) The original OpenFold implementation with EvoformerAttention used in AlphaFold2. The memory explosion problems in training/inferenc