[Paper Review] 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 Science8 citations

TL;DR

The paper introduces the DeepSpeed4Science initiative, detailing AI system technologies built on DeepSpeed to accelerate large-scale scientific discovery, with two structural biology showcases and a plan for broader science collaborations.

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.

Motivation & Objective

• Motivate the need for AI system technologies tailored to science domains beyond generic LLM acceleration.
• Describe the DeepSpeed4Science approach and its base on DeepSpeed pillars (training, inference, compression).
• Showcase two system challenges in structural biology addressed by DS4Sci (memory explosion in Evoformer attention; long-sequence support for GenSLMs).
• Outline the collaboration model and potential platform for sharing AI system technologies for science.

Proposed method

• Develop customized memory-efficient EvoformerAttention kernels to eliminate memory explosions in Evoformer-centric models.
• Integrate enhanced Megatron-DeepSpeed framework with long-sequence support through memory optimizations for attention masks and position embeddings.
• Enable very-long sequence training/inference for genome-scale foundation models via Megatron-DeepSpeed rebasing and optimizations.
• Fuse kernels and employ tiling, on-the-fly broadcasting, and FP32-safe gradient handling to reduce peak memory while preserving accuracy.
• Leverage sequence parallelism, tensor/pipeline parallelism, and model/data offloading to extend sequence lengths dramatically.

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.

Experimental results

Research questions

• RQ1 How can AI system technologies be specialized to address memory and sequence-length challenges in science-focused models?
• RQ2 Can customized kernels and framework rebasing enable significantly longer context sizes for genome-scale and Evoformer-based models without accuracy loss?
• RQ3 What are the performance/throughput gains when applying DS4Sci optimizations to structural biology and GenSLM-style models?
• RQ4 How can DS4Sci facilitate broader collaboration and sharing of advanced AI system techniques for scientific discovery?

Key findings

• A DS4Sci_EvoformerAttention kernel reduces OpenFold's peak memory by 13X for Evoformer-attention variants without accuracy loss.
• New Megatron-DeepSpeed framework enables training GenSLMs with substantially longer sequences, reporting up to an average of 13x longer sequences and up to 2x throughput in certain cases.
• Megatron-DeepSpeed rebasing adds rotary positional embedding, FlashAttention v1/v2, and new fused kernels to improve long-sequence training and inference.
• Memory optimizations for attention masks and position embeddings, plus sequence parallelism, yield significant extensions of feasible sequence lengths for GenSLMs (e.g., 512K for 25B GenSLM) over prior limits.
• The DS4Sci efforts position DeepSpeed4Science as a platform and repository for sharing advanced AI system technologies for science.

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