A Survey on Memory-Efficient Large-Scale Model Training in AI for Science

Protein structure prediction stands as one of the most significant challenges in biochemistry, as high-precision predictions are essential for drug discovery and protein design. AlphaFold 2 [2], developed by DeepMind, is a DL-based tool for predicting protein structures. It can predict the three-dimensional structure of proteins based on their amino acid sequences, representing a major breakthrough in addressing long-standing biological challenges.

The architecture of AlphaFold 2 is built upon variants of Transformer. Its main structure features an input module that receives the amino acid sequence of the target protein and searches for homologs within a sequence database to perform multiple sequence alignment (MSA). The Evoformer module then processes the MSA representation and pair representation, facilitating the exchange and updating of information through Transformer layers. Finally, the structural module converts the abstract representation of proteins into 3D atomic coordinates [73]. AlphaFold 2 has significantly influenced the advancement of biotechnology. However, its training process and datasets are not fully open source, which restricts researchers from applying it to tasks beyond protein structure prediction. Furthermore, despite the AlphaFold 2 model having a relatively small parameter size (about 93M), the memory consumed by intermediate activations is substantial due to its architecture, and peak memory usage increases cubically with the length of the input sequence. Despite using mixed precision and gradient checkpointing, the memory cost of activations in attention module exceeds 20 GB [49]. To facilitate further research, Ahdritz et al. developed OpenFold [48]. As an open-source implementation of AlphaFold 2, it achieves comparable prediction accuracy while runing faster and consuming less memory. It is built based on PyTorch [74], which makes it easier for the community to use and improve.

Although AlphaFold2 and RoseTTAFold [45] achieved breakthroughs in atomic resolution structure prediction, they rely on MSA and similar protein structure templates to achieve better performance. Based on the research of OpenFold, Lin et al. [51] further explored the ability of language model and proposed ESMFold. ESMFold utilizes the internal representation of language model to replace the expensive network module for processing MSA with a transformer module for processing sequences. This simplification greatly improves the speed of ESMFold, which is much higher than models based on MSA. To enable the training of ESMFold, ZeRO stage 3 [75] is used to divide the model parameters, gradients, and optimizer states among GPUs.

The protein language model (PLM) has achieved remarkable success in learning biological information from protein sequences. Most existing models are constrained by autoencoding or autoregressive pre-training objectives, making it challenging to address both protein understanding and generation tasks. Chen et al. [52] explored the possibility of jointly optimizing these two training objectives through an innovative framework and trained a 100B PLM (xTrimoPGLM). Inspired by ESMFold, they combined the folding module with xTrimoPGLM and developed xT-Fold, a high-performance 3D structure prediction tool. The performance of xT-Fold is superior to other PLM-based models on benchmarks such as CAMEO and CASP15. As the model size and the scale of training data are greatly increased, xTrimoPGLM uses various memory optimization methods in the pre-training stage, including 3D-parallelism [76], [77], [75], FP16 mixed precision training [78], and gradient checkpointing [79].

In May 2024, Isomorphic Labs and DeepMind jointly released AlphaFold 3 [47]. The overall architecture of Alphafold 3 is similar to AlphaFold 2, its prediction accuracy exceeds AlphaFold 2, and the model is more efficient. In addition, the application scope of Alphafold 3 has also been expanded. It can predict the structure of a single protein, the complex structure of protein complexes, proteins, nucleic acids, small molecules, etc.

LLMs have demonstrated significant potential in the medical and clinical fields. When trained on extensive medical datasets, these models can produce high-quality outputs, and support decision-making. The applications of such medical models are numerous, including medical knowledge retrieval, assistance in clinical diagnoses, consultations, and personalized medical services.

Singhal et al. [3] found that the benchmark used to evaluate the clinical knowledge of the model is limited. To solve it, they proposed MultiMedQA, a benchmark composed of seven medical Q&A datasets. Further, they fine-tuned the command of PaLM [34] and got Flan-PaLM [32], which shows excellent performance in multiple-choice questions. Moreover, they suggested instruction prompt tuning, and the obtained Med-PaLM model can further adapt to medical field. The pre-training model of Med-PaLM, PaLM, uses ZeRO stage 3 and TP to shard model parameters across all TPUs during the training process and uses gradient checkpointing to reduce the memory overhead of activations.

Tu et al. [5] aim to further explore the potential of generalist models (those capable of performing multiple tasks without fine-tuning, such as GPT-3, PaLM, and GPT-4) in medical field. Utilizing the language and multimodal foundation models PaLM 540B and PaLM-E 562B [55], the researchers develop Med-PaLM M. They also designed a multimodal biomedical benchmark called MultiMedBench. As a large-scale biomedical generalist model, Med-PaLM M demonstrates strong performance across all tasks on MultiMedBench, matching or exceeding the capabilities of task-specific models.

Currently, most LLMs utilized in medical field are either closed source or have relatively small model sizes. Chen et al. [58] believe that while models designed for generalist tasks typically perform well across tasks, their performance in certain tasks falls short compared to models tailored for those specific purposes. They developed Meditron-70B through continued pretraining of Llama-2 and fine-tuned it for particular tasks, resulting in superior performance compared to all open-source generalist and medical LLMs across all medical benchmarks. Furthermore, to enable training, they extended Nvidia’s Megatron-LM [76] and developed Megatron-LLM²²2https://github.com/epfLLM/megatron-LLM, a distributed training library that accommodates various model architectures. They also implemented 3D parallelism to optimize the training efficiency of models at this scale.

HuatuoGPT [59] aims to assist doctors with diagnosis and treatment in Chinese scenarios, it combines the dialogue capabilities of ChatGPT with the expertise of real-world physicians. By employing instruction fine-tuning and reinforcement learning techniques, it exploits distilled instruction data generated by ChatGPT alongside actual response data from doctors. The upgraded version, HuatuoGPT-II [61], successfully passed the Chinese National Pharmacist Licensure Examination in October 2023 and nearly all medical qualification exams, highlighting its robust capabilities within the Chinese medical landscape. ZeRO data parallelism is enabled during the training of HuatuoGPT and HuatuoGPT-II.

The emergence and rapid development of LLM heralded a new era in the fields of biopharmaceuticals and chemical sciences, providing innovative methods for drug discovery, chemical synthesis and optimization, and elucidation of complex biological pathways.

Drug labels possess distinct characteristics that differentiate them from texts in other domains, making it ineffective to apply generic models directly to this specific field. ValizadehAslani et al. [64] leveraged drug label text corpora and conducted pre-training on BERT base checkpoints using masked language modeling (MLM) as the pre-training task. This effort resulted in PharmBERT, a domain-specific BERT model for drug labeling. Their work showcased exceptional performance by evaluating PharmBERT on three downstream tasks.

To tackle the issue of limited knowledge among general-purpose language models in the domains of biopharmaceuticals and chemistry, Chen et al. introduced PharmaGPT [65], a large-scale multilingual language model tailored for these fields. PharmaGPT is pre-trained on high-quality, large-scale, domain-specific datasets and optimized for enhanced performance in biopharmaceuticals and chemistry through fine-tuning and reinforcement learning. Experiments indicate that PharmaGPT outperforms mainstream general models on professional benchmark tests, such as NAPLEX, showcasing its exceptional capability in understanding and generating specialized knowledge.

The application of LLM in the field of chemistry covers multiple aspects, including molecular design, molecular property prediction, and chemical reaction prediction. Additionally, LLMs can be used to analyze and summarize chemical literature, providing researchers with efficient knowledge retrieval and decision-making support. Their powerful sequence processing capabilities enable the handling of complex chemical terminology and diverse data formats, thereby fostering advancements in automation and intelligence in the field of chemistry.

Extracting structured data from chemical literature is a challenging task due to the complexity and heterogeneity of chemical language. Guo et al. [66] broke this task into two cascaded subtasks: product extraction and reaction role labeling. They developed two independent models, ChemBERT and ChemRxnBERT, based on BERT. By pre-training and fine-tuning these models on domain-specific and task-oriented corpora, they enhance the understanding of semantic relationships and contextual features in chemical texts.

In the field of molecular property prediction, Ock et al. introduced CatBERTa [8], which can bypass the requirement of precise atomic coordinates in previous GNN-based methods and predict catalyst properties sorely through molecular text representation. Compared to graph-based methods, using text representation to describe adsorbate-catalyst systems has better interpretability. The prediction accuracy of CatBERTa is comparable to GNN-based methods, highlighting its potential in property prediction tasks.

Irwin et al. pre-trained BART [68] on a dataset containing a large number of simplified molecular line entry system (SMILES) strings to capture the structural relationships and potential properties between chemical molecules, resulting in the Chemformer model [9]. Chemformer is capable of both sequence-to-sequence tasks (e.g. reaction prediction, molecular optimization), and discriminative tasks (e.g. molecular property prediction, biological activity analysis). It achieves or exceeds the performance of state-of-the-art models on multiple benchmarks while significantly reducing convergence time. This effectively addresses the challenges posed by limited computing resources and time constraints.

The objective of molecular design and generation tasks is to create new molecules with specific functions or properties for applications such as drug development, materials science, and catalyst development. Traditional methodologies depend heavily on the expertise of specialists and costly computational simulations. However, the introduction of LLMs has opened new possibilities in this field. Fang et al. introduced a molecular generative pre-trained language model called MolGen [7], which has enhanced its comprehension of chemical structures and semantics through a two-stage pre-training strategy. This strategy encompasses molecular syntax and semantic learning, utilizing self-sacrificing embedded strings (SELFIES) along with domain-agnostic prefix tuning. Furthermore, the study proposed a chemical feedback paradigm that aligns generated molecules with real-world chemical preferences by adjusting their probability distribution, effectively reducing the issue of ”molecular hallucination” (i.e. generated molecules lacking expected functionality in practical applications). Frey et al. [6] conducted an in-depth study on the neural scaling behavior of deep chemical models in relation to dataset and model size, aiming to explore how model performance is enhanced with increased resource investment. They designed and trained a generative pre-trained Transformer model named ChemGPT with over 1B parameters for small molecule generation. By varying both the model and dataset sizes by several orders of magnitude, they demonstrated an improvement in pre-training loss performance as the scale increased.

Weather forecasting represents a crucial application scenario in meteorology, providing the ability to predict future weather changes, which benefits daily life, agriculture, transportation, and other fields. Over the past decade, the swift advancement of high-performance computing devices has led to significant progress in numerical weather prediction (NWP), enhancing the accuracy of forecasts for daily weather, extreme weather events, and climate change. However, as the growth of computing power slows and the complexity of physical models increases, the limitations of traditional methods have become more pronounced. Recently, the rapid progress in deep learning has emerged as a promising avenue for addressing these challenges in the field.

The Pangu-Weather model [10], [11], proposed by Bi et al., employs the 3D Earth-Specific Transformer (3DEST) architecture to process 3D climate data, which is a variant of Swin Transformer [70], [71]. It is the first AI-driven approach to surpass the forecast accuracy of NWP methods. Notably, Pangu-Weather achieves an inference speed of just 1.4 seconds on a single GPU, making it 10000× faster than the previous most accurate NWP method (operational IFS [80]). While the accuracy of 3D models can surpass 2D models, this enhancement comes with a considerable increase in memory requirements, leading to limited network depth. The study employs DP to train on 192 GPUs, with a batch size of 1 assigned to each GPU.

The inherent uncertainty in weather forecasting is inevitable, and the degree of this uncertainty, along with cumulative errors, increase as the forecast period lengthens. In response to this challenge, Chen et al. introduced FuXi [12], a cascaded model architecture based on U-Transformer, which effectively reduce cumulative errors and improves both forecast accuracy and duration. The model pre-training was conducted on a cluster of 8 Nvidia A100 GPUs, utilizing techniques such as FSDP [81], BF16 mixed precision training, and gradient checking to optimize memory usage.

Inspired by the concept of foundational models [82], Nguyen et al. designed and trained a basic model named ClimaX [13], aimed at effectively adapting to general tasks related to the Earth’s atmosphere. ClimaX can be trained on heterogeneous datasets to solve diverse climate and weather tasks. Remarkably, even with pre-training conducted at lower resolutions and with limited computational resources, ClimaX surpasses previous data-driven models in weather and climate prediction benchmarks. The pre-training used 80 NVIDIA V100 32GB GPUs and enabled FP16 mixed precision training. Subsequently, Bodnar et al. introduced Aurora [14], a large-scale atmospheric foundation model designed to excel in various prediction tasks, particularly in regions with limited data or extreme weather conditions. After training on diverse weather and climate datasets, Aurora can effectively adapt to new atmospheric prediction tasks by learning a universal representation of atmospheric dynamics. The model was trained on 32 GPUs with 1 batch size per GPU. To meet memory constraints, BF16 mixed precision training and gradient checkpointing were also used.

Despite the progress achieved in applying large-scale models across various domains, the continuous increase in memory overhead during model training remains a pressing concern that requires attention. We identify the underlying causes of this issue from the following key factors:

• •

  The number of model parameters continues to increase. In recent years, there has been a notable increase in the scale of models across various fields, as shown in Fig. 1. This trend is driven by the demand for processing complex and large-scale data, advancements in computing resources, and the pursuit of enhanced performance and accuracy in downstream tasks. Consequently, this growth has also resulted in higher memory requirements.

• •

  The necessity to process long sequence inputs. Long sequence data is prevalent across various scientific domains, such as DNA and protein sequences in biology, molecular structures in chemistry, and time-series weather data in meteorology. As the length of input data increases, the memory required to store activations also grows significantly. This challenge is exacerbated by the quadratic computational complexity inherent in transformer models. The situation may become worse when employing a more complex architecture.

• •

  The demand for handling diverse data structures. In scientific fields, the use of high-resolution images and multimodal data significantly contributes to increased memory consumption during large model training. High-resolution images contain a substantial amount of pixel information, thereby amplifying the demands of data processing and storage. Additionally, multimodal data requires the handling of diverse data types, such as images, text, and audio. Each modality necessitates an independent processing pathway and fusion mechanism, which further elevates the complexity and memory requirements of the model.

Data parallelism (DP) is a commonly used distributed training technique, particularly in environments with multiple GPUs and nodes. DP can significantly accelerate training process and improve overall training throughput. It involves replicating the model across computing devices. During training, each device is assigned a subset of the batched training data, known as a mini-batch, which is used as input for the model. Due to the different input data across model replicas, it is necessary to aggregate the gradients from all computing devices before updating model parameters. In terms of implementation, DP architectures can be categorized into parameter server architecture (Fig. 2a) and decentralized architecture (Fig. 2b). In parameter server architecture [83], devices are divided into parameter servers and workers. The parameter server stores the latest parameters of the model and aggregates gradients collected from workers to perform parameter updates. Each worker communicates with the parameter server during training to retrieve the latest parameters or upload gradients calculated using mini-batches.

Parameter servers are likely to become a communication bottleneck, as all workers need to communicate with servers frequently, especially in large-scale training scenarios. In contrast, decentralized architecture avoids this issue. In such a setup, computing device peers perform gradient aggregation through collective communications such as all-reduce, which enhances scalability. Moreover, the implementation of PyTorch DDP [84] leverages the overlap between computation and communication to further increase training throughput.

DP divides the input training data along the batch dimension, which reduces the size of intermediate activations on each device during training and greatly improves training efficiency. However, its problem is also obvious. DP requires each device to hold a model replica, which is infeasible for large models that a single device cannot store.

The memory overhead during model training primarily comes from model states and intermediate activations. The model states can be further divided into optimizer states, gradients, and parameters. Assuming the number of parameters in the model is denoted as $P$, and utilizing Adam optimizer with FP16 mixed precision during training, the memory overhead of model parameters is $2P+4P=6P$ bytes. The memory overhead of gradients is $2P$ bytes, and optimizer states consume $4P+4P=8P$ bytes. Thus, the total memory overhead is $16P$ bytes. If training GPT-3 (175B), would require approximately 2.8TB of memory. To address the redundancy of model states, Rajbhandari et al. introduced Zero Redundancy Optimizer (ZeRO) [75], which shards redundant model states across devices, reducing memory overhead up to $\frac{1}{N_{d}}$, where $N_{d}$ represents the number of devices. The core principle of ZeRO is that each device maintains only a fraction of the model states instead of a complete copy. The implementation of ZeRO has been integrated into Microsoft’s open-source distributed training framework, DeepSpeed [85], making it easily accessible for researchers and developers.

Inspired by ZeRO, PyTorch developed Fully Sharded Data Parallel (FSDP) [81]. A report from Hugging Face³³3https://huggingface.co/blog/deepspeed-to-fsdp-and-back indicates that the performance of PyTorch FSDP is comparable to that of DeepSpeed ZeRO. Since PyTorch v2.4⁴⁴4https://github.com/pytorch/pytorch/releases/tag/v2.4.0, the PyTorch team has restructured FSDP and introduced FSDP2, adding additional optimizations to enhance performance.

The rapid increase in model parameters has rendered the memory and computing resources of a single device insufficient to support the training of the entire model. Although ZeRO can partition the model across computing devices, it does not reduce computational workloads on each device. For large models with vast parameters and high computational demands, relying solely on DP becomes inadequate for meeting the training needs. Model parallelism (MP) has emerged as a crucial technique for further improving the scale of deep learning models. MP primarily includes two techniques: tensor parallelism (TP) and pipeline parallelism (PP), which will be discussed in this section and the following section. TP, also referred to as intra-layer parallelism, finely divides the computation within each model layer into smaller tensor operations that can be executed concurrently across multiple devices, thereby effectively distributing both memory overhead and computational workload.

Matrix multiplication can be divided into smaller operations by row or column, allowing the results to be calculated separately and then aggregated to obtain the final outcome, as shown in Fig. 3. In line with this idea, Megatron-LM [76] has developed a tensor parallelism approach specifically designed for the transformer architecture. Taking the MLP layer as an example, it consists of two parameter matrices: $W_{1}$ with a shape of $(h,4h)$ and $W_{2}$ with a shape of $(4h,h)$, here we denote hidden size as $h$. Initially, $W_{1}$ is partitioned column-wise, while $W_{2}$ is partitioned row-wise. This results in the parameter matrices stored on each device having the shapes: $W_{1}^{\prime}(h,4h/N_{d})$ and $W_{2}^{\prime}(4h/N_{d},h)$. After forward pass on each device is completed, an all-reduce is performed to obtain the final result.

Megatron-LM TP divides tensors into one-dimensional segments. While this approach reduces model parameters stored on each device to $\frac{1}{N_{d}}$, it still needs to store complete intermediate activations on each device. Additionally, forward and backward computations of each transformer layer involve 4 all-reduce, causing significant communication overhead. To address this, researchers have proposed multidimensional tensor partitioning methods [86], [87], [88], which further reduce memory overhead on each device and mitigate the frequent communication issues of 1D TP [89]. In practice, TP is typically applied within nodes rather than across nodes to prevent communication from becoming a bottleneck for model training efficiency [90].

Pipeline parallelism, also known as inter-layer parallelism, involves partitioning the model layers horizontally into multiple stages. Each stage consists of consecutive layers, which are assigned to different devices for sequential execution. Unlike TP, which divides intra-layer tensor calculations, PP partitions the model into stages, where each device transfers activations and gradients between adjacent stages during forward and backward passes. PP not only reduces memory demands and computational workloads on each device but also has less communication overhead. However, a major challenge associated with PP is device idle time (often referred to as pipeline bubbles). While asynchronous PP schemes [91], [92] can eliminate pipeline bubbles, they have no guarantee of model convergence [93]. As synchronous PP schemes continue to improve, their performance and efficiency are comparable with asynchronous PP, making them a more favorable option. In the following paragraphs, we will discuss in detail several common PP schemes.

In vanilla PP, data can be passed to the next device only after the preceding device completes its calculation, due to the sequential dependence in data flow. This leads to pipeline bubbles, resulting in a waste of resources and diminishing overall efficiency of the system.

GPipe [94] splits each input batch into several micro-batches, which are input to the model sequentially. After completing the processing of each micro-batch, each device immediately passes its results to the next device and begins to process the next micro-batch. This approach significantly reduces the pipeline bubbles, and finally, each device uses the accumulated gradients to perform synchronous parameter updates. To further enhance memory efficiency, GPipe incorporates gradient checkpointing [79], which retains only a subset of intermediate activations and recalculates the remaining activations during backpropagation. However, GPipe still faces the issue of high memory consumption for intermediate activations.

The GPipe scheduling involves completing the forward passes for all micro-batches before executing the backward passes (F-then-B). Let $N_{s}$ denote the number of stages in the model, $N_{d}$ denote the number of devices, and $N_{m}$ denote the number of micro-batches. In the GPipe schedule, each device must retain intermediate activations of $N_{m}$ micro-batches. To reduce memory consumption, researchers introduced a one forward pass followed by one backward pass pattern (1F1B) [95], [96]. This scheduling strategy effectively reduces the number of on-the-fly micro-batches by executing the backward pass in advance. With 1F1B schedule, the pipeline bubble is equal to that of GPipe, and the activations of micro-batches required to store on each device decrease from $N_{m}$ to $N_{d}$ (where $N_{m}$ is typically greater than $N_{d}$).

The aforementioned PP schemes assume that the number of pipeline stages $N_{s}$ and the number of devices $N_{d}$ are equal, with each device possessing one stage. Narayanan et al. [95] split the stages based on the 1F1B schedule, allowing each device to have multiple stages. For instance, in Fig. 4, an 8-layer model is partitioned into 4 stages, where each stage consists of two consecutive layers, and each device handles one stage (device 1 possesses L1 and L2, device 2 possesses L3 and L4, etc.). If the model is divided into 8 stages, each device will be responsible for the computation of 2 stages. This interleaved 1F1B schedule performs a finer-grained scheduling of forward and backward computations, reducing pipeline bubbles and enhancing efficiency. However, the trade-off for this reduction in pipeline bubbles includes increased communication and higher activation memory usage.

In training iteration, the computational cost of backpropagation is greater than that of forward propagation (generally estimated to be about twice as much), resulting in longer computation time for backpropagation. This is the primary contributor to the existence of pipeline bubbles. Qi et al. [97] suggest splitting the backward pass into the calculation of the gradients of activations (denoted as $B$), and the calculation of the gradients of parameters (denoted as $W$). They found that the $B$ pass is independent of the $W$ pass within the current layer and only depends on the $B$ pass from the previous layer. Theoretically, one could complete $B$ passes of all layers before performing $W$ passes for each layer. This allows for the design of more efficient pipeline schedules. One of the proposed pipeline schedules (ZB-H1) is similar to 1F1B, but it offers the flexibility to schedule $W$ after the corresponding $B$ within the same stage, effectively filling bubbles in the schedule. Additionally, in terms of activation memory usage, ZB-H1 ensures that peak memory usage on all devices does not exceed 1F1B. If the memory limit is adjusted to accommodate peak memory usage exceeding 1F1B, along with the optimizer post-validation strategy proposed by the author, zero bubble scheduling (ZB-H2) can be achieved. ZB-H2 fills additional $F$ during the warm-up phase to fill bubbles prior to the initial $B$, followed by a rearrangement of $W$ passes at the tail to completely eliminate all bubbles in pipeline schedule.

PP typically has the issue of high activation memory usage, and the 1F1B pipeline schedules mentioned above have the problem of memory imbalance. Regarding this, Qi et al. [93] proposed an analytic framework to decompose existing pipeline schedules into building blocks and found that peak memory usage of a pipeline schedule is highly related to the lifespan of the building blocks. Based on this, a series of V-Shape building blocks are proposed in the article, which have more balanced peak memory on each device compared to interleaved 1F1B.

In the distributed training of large-scale deep learning models, PP is a commonly used technique. Compared to TP, it only communicates between adjacent stages and has lower communication overhead. Therefore, it is usually applied between nodes to scale up model size [95].

Supporting long sequence lengths in large-scale models is crucial for applying AI to scientific endeavors. Many scientific challenges inherently involve complex analysis of large-scale, high-dimensional data, frequently presented as long sequences. Long-sequence issues are prevalent in fields including structural biology, chemistry, drug development, and atmospheric science. For instance, many proteins consist of hundreds or even thousands of amino acid residues, where the sequence context of each residue can be essential to its 3D structure. Long-sequence support enhances the ability to capture non-local interactions among residues, thereby improving the quality of predictions. Additionally, complex molecules or chemical reactions often necessitate long sequence representations, including molecular formulas, SMILES, reaction pathways, etc. Climate prediction models require analyzing long-term historical and global spatial data, represented in time series or grid formats. Despite the importance of long-sequence support, its implementation presents challenges. Processing long sequences demands significant increases in computational power and memory, with the memory usage of attention mechanisms rising quadratically with sequence length.

In traditional distributed training techniques such as DP and TP, model parameters or input data are partitioned and assigned to multiple devices for processing. However, when the model needs to handle long inputs, these methods may lose efficiency due to memory bottlenecks or computational overhead. Sequential parallel (SP) is relatively novel and differs from DP and MP methods in that it aims to reduce the memory overhead caused by activations during model training. Although SP strategies divide input along sequence dimension, the specific implementation varies.

Li et al. [98] considered solving the problem from a system perspective by combining ring-style communication with self-attention, allowing the input sequence to be distributed to each device and calculating attention scores across devices. Compared with Megatron-LM TP [76], the proposed method is not limited by the number of attention heads, as long as the sequence length is divisible by the degree of SP.

In Megatron-LM TP, some activations within the transformer layer are not distributed across devices, leading to increased memory overhead. The SP method proposed by Li et al. [98] can help mitigate this issue; however, it requires replicating model parameters and optimizer states across all devices, making it impractical for large-scale model training. Korthikanti et al. [90] found that the LaterNorm and Dropout operations within the transformer layer operate independently along the sequence dimension. Consequently, modifications were implemented based on Megatron-LM to effectively distribute the computational workloads and activations associated with LaterNorm and Dropout, without incurring additional communication overhead.

Jacobs et al. proposed Ulysses [99], a method that divides the input along sequence dimension onto all devices and conducts all-to-all communication on the Query (Q), Key (K), and Value (V) matrices before performing attention computation. This ensures that each device operates with disjoint attention heads, allowing for distributed attention calculation. Once the calculations are complete, another all-to-all is performed to get the results and restore the state that inputs were divided along sequence dimension. Ulysses SP and Megatron-LM TP share similarities in the distributed calculation of self-attention, both concerning limitations on the number of attention heads. A key advantage of Ulysses is that the communication volume for a single device decreases linearly as SP degree increases, while in Megatron-LM, it is independent of TP degree.

Liu et al. [100] developed a ring-style distributed attention computation method that optimizes resource usage by conserving a portion of the Q, K, and V blocks. During calculations, each device exchanges K and V blocks with its neighboring devices, enabling it to perform computations using blockwise self-attention and feedforward networks (FFN) [101]. In contrast to the approach taken by Li et al. [98], this method can eliminate communication overhead by picking an proper chunk size and overlapping communication with computation.