# Efficient Training of Large Language Models on Distributed Infrastructures:
A Survey

Jiangfei Duan1, Shuo Zhang1, Zerui Wang1, Lijuan Jiang, Wenwen Qu, Qinghao Hu,
Guoteng Wang,
Qizhen Weng, Hang Yan, Xingcheng Zhang, Xipeng Qiu, Dahua Lin, Yonggang Wen,
Xin Jin,
Tianwei Zhang and Peng Sun${}^{\text{\Letter}}$ Jiangfei Duan and Dahua Lin
are with Shanghai AI Laboratory and the Chinese University of Hong Kong.
E-mail: {dj021<EMAIL_ADDRESS>Shuo Zhang and Xipeng Qiu are with
Shanghai AI Laboratory and Fudan University.
E-mail: {zhangshuo<EMAIL_ADDRESS>Zerui Wang is with Shanghai AI
Laboratory and Shanghai Jiao Tong University.
E-mail<EMAIL_ADDRESS>Qinghao Hu, Yonggang Wen and Tianwei Zhang are
with Nanyang Technological University.
E-mail: {qinghao.hu, ygwen<EMAIL_ADDRESS>Xin Jin is with School
of Computer Science, Peking University.
E-mail<EMAIL_ADDRESS>Lijuan Jiang, Wenwen Qu, Guoteng Wang, Qizhen
Weng, Hang Yan, Xingcheng Zhang and Peng Sun are with Shanghai AI Laboratory.
E-mail: {jianglijuan, quwenwen, wangguoteng, wengqizhen, yanhang,
zhangxingcheng<EMAIL_ADDRESS>1 Equal Contribution.

###### Abstract

Large Language Models (LLMs) like GPT and LLaMA are revolutionizing the AI
industry with their sophisticated capabilities. Training these models requires
vast GPU clusters and significant computing time, posing major challenges in
terms of scalability, efficiency, and reliability. This survey explores recent
advancements in training systems for LLMs, including innovations in training
infrastructure with AI accelerators, networking, storage, and scheduling.
Additionally, the survey covers parallelism strategies, as well as
optimizations for computation, communication, and memory in distributed LLM
training. It also includes approaches of maintaining system reliability over
extended training periods. By examining current innovations and future
directions, this survey aims to provide valuable insights towards improving
LLM training systems and tackling ongoing challenges. Furthermore, traditional
digital circuit-based computing systems face significant constraints in
meeting the computational demands of LLMs, highlighting the need for
innovative solutions such as optical computing and optical networks.

###### Index Terms:

Large Language Models; Distributed Training; Machine Learning Systems.

††publicationid: pubid:

##  1 Introduction

Large Language Models (LLMs) are transforming the AI industry, demonstrating
remarkable capabilities across a wide range of tasks and applications,
including personal assistants [1], code copilot [2], chip design [3] and
scientific discovery [4]. The success of this revolution is built on the
unprecedented scale of transformer-based LLMs like GPT [5], LLaMA [6], Gemini
[7], etc. Moreover, it is evidenced that the scaling of LLMs has not plateaued
[8]. This trend has significantly shifted the design of underlying training
systems and infrastructures, since LLM typically follows a relatively fixed
architecture and its training exclusively occupies huge GPU clusters over
extended periods. For example, the pretraining of LLaMA-3 takes approximately
54 days with 16K H100-80GB GPU on Meta’s production cluster [9].

LLM training highlights significant challenges for today’s training system and
infrastructure in terms of “SER”, i.e., Scalability, Efficiency, and
Reliability. Scalability demands that both infrastructure and systems adapt
seamlessly to massive clusters of tens of thousands of GPUs or AI
accelerators, while preserving training correctness and model accuracy. This
requires innovative solutions in hardware configuration, networking, and
training framework. Efficiency focuses on maximizing resource utilization
across the entire cluster, often measured by Model FLOPs Utilization (MFU).
Achieving high MFU involves optimizing computation, minimizing communication
overhead, and efficiently managing memory at an unprecedented scale.
Reliability are critical given the prolonged duration of LLM training, usually
lasting weeks to months. The system must maintain consistent performance and
be resilient to various types of failures, including hardware malfunctions,
network issues, and software errors. It should be capable of swiftly detecting
and recovering from these failures without significant loss of progress or
training quality. These interrelated challenges necessitate a holistic
approach to system and infrastructure design, pushing the boundaries of large-
scale distributed computing and opening new avenues for research and
innovation in high-performance machine learning systems.

This survey paper aims to provide a comprehensive overview of the advancements
in LLM training systems and infrastructure, addressing the aforementioned
challenges. This survey spans from the distributed training infrastructure to
training systems. We examine innovative approaches to infrastructure design,
covering advancements in GPU clusters, high-performance networking, and
distributed storage systems tailored for LLM workloads. We also explore the
key aspects of distributed training systems, including parallelism strategies,
computation, communication and memory optimizations that enhance the
scalability and efficiency. We also delve into fault tolerance mechanisms for
improving training reliability. By synthesizing recent advancements and
identifying future research directions, this survey aims to provide
researchers and practitioners with insights into the most promising avenues
for improving LLM training systems. Our goal is to offer a valuable resource
that not only addresses current challenges but also paves the way for future
innovations in large-scale machine learning infrastructure.

Figure 1: Overall structure of this survey.

Organization. Fig. 1 shows the organization of this survey. Section 2
discusses the background information about LLM architecture, characteristics
and challenges of LLM training. In Section 3, we summarize the key aspects of
the training infrastructure, including AI accelerators, network infrastructure
and storage systems. In Section 4, we investigate the parallelism schemes for
distributed LLM training. In Section 5, we discuss computation optimizations
to harness the unprecedented computational capabilities. In Section 6, we
discuss techniques to optimize the memory footprint in LLM training. In
Section 7, we introduce communication optimizations to minimize the
communication overhead. In Section 8, we first present a failure analysis, the
cover approaches to enable fast failure detection and recovery. Finally, we
conclude the survey in Section 9.

##  2 Background

###  2.1 Transformer-based LLMs

The current state-of-the-art Large Language Models (LLMs) are predominantly
transformer-based. Their core architecture is built around the attention
mechanism [10], which allows the model to dynamically weigh the importance of
different words in a sentence. Fig. 2 depicts the typical architecture of a
transformer layer [10], which can be stacked multiple times to construct an
LLM. The input text is first tokenized into individual tokens, which are then
converted into a token vector $X$ via an embedding layer. To preserve the
sequential nature of the text, the token vector is embedded with positional
information. The resulting token vector is then fed into the transformer
layer, which consists of an Attention block and a Feed-Forward Neural Network
(FFN) block.

Suppose the input token vector is $X=[x_{1},x_{2},\cdots,x_{n}]$. These tokens
are first transformed into query $Q$, key $K$ and value $V$ tensors via linear
transformation. The attention mechanism computes the attention output as
follows:

$\text{Attention}(Q,K,V)=\texttt{softmax}\left(\frac{QK^{T}}{\sqrt{d}}\right)V$
(1)

where $d$ is the dimension of the key tensor. This formula ensures that the
LLM can focus on relevant parts of the input sequence by calculating a
weighted sum of the values, where the weights are derived from the similarity
between the queries and keys. After the attention layer, the output is passed
through a FFN for further processing.

Nowadays, LLMs generally follow the original decoder-only transformer
architecture but incorporate modifications to the attention mechanism and FFN
to enhance efficiency and performance. The original attention mechanism, known
as Multi-Head Attention (MHA) [10], suffers from quadratic computational
complexity and high memory consumption due to the key-value cache. To address
these issues, several variants such as Multi-Query Attention (MQA) [11],
Group-Query Attention (GQA) [12] and Multi-Latent Attention (MLA) [13] have
been proposed. A notable advancement of the FFN component is the Mixture-of-
Experts (MoE) [14, 15] architecture, which employs a sparsely activated FFN.
In MoE, only a subset of the FFN layers (or experts) are activated for each
input, significantly reducing the computational load while maintaining high
model capacity.

###  2.2 LLM Training Workloads Characteristics

Figure 2: A typical Transformer layer contains an Attention block and a Feed-
Forward Network (FFN) block.

The characteristics of LLM training workloads diverge significantly from
traditional deep learning workloads, primarily due to their complexity and
scale. These unique characteristics affect training system design,
performance, scalability, and resource utilization. Here, we highlight key
differences and requirements for LLMs.

(1) Homogeneous Model Architecture. Unlike prior DL workloads that employed
various model architectures (e.g., LSTM [16], CNN [17]) for different tasks,
LLMs predominantly use the Transformer architecture [10]. Models like GPT [5],
LLaMA [6], InternLM[18] and MOSS[19] all share this common foundation. This
architectural uniformity presents significant potential for optimizing system
performance for a particular model architecture.

(2) Unprecedented Scale and Training Duration. LLM training operates at an
unparalleled scale, typically updating models with hundreds of billions of
parameters with terabyte-scale training datasets. Such magnitude necessitates
distributed training across huge GPU clusters and presents challenges in
maintaining high efficiency. Besides, the training of LLMs can span weeks or
months, demanding robust fault tolerance mechanisms and efficient
checkpointing strategies to safeguard against data loss and facilitate the
resumption of interrupted training sessions.

(3) Specialized Software Optimization. To accommodate the enormous model size
of LLMs, specialized systems implement advanced techniques to optimize
execution. For example, Megatron [20] and Alpa [21] accelerate training
through hybrid parallelism. DeepSpeed [22] reduces memory consumption by
integrating state-sharding optimizers.

(4) Shift in Training Paradigm. Traditional DL workloads follow a task-
specific paradigm, training models on domain-specific data for particular
tasks, such as translation. In contrast, LLMs adopt a self-supervised training
approach on extensive datasets to create foundation models, which are then
adapted for various downstream tasks. This paradigm shift represents a
substantial change in the model development pipeline, including pretraining
and alignment phases, and results in distinct workload characteristics
compared to previous DL workloads. From the datacenter perspective, LLM
development involves numerous small-scale workloads that are associated with
pretraining, including alignment (i.e., fine-tuning) and periodical evaluation
workloads [23].

###  2.3 LLM Training Challenges

The unique characteristics of LLM training workloads give rise to significant
challenges in developing efficient training systems and infrastructure. These
challenges primarily manifest in three critical areas: scalability,
efficiency, and reliability. Each of these challenges stems directly from the
massive scale of LLMs and the complexity of their training processes,
requiring innovative solutions that push the boundaries of distributed
computing and machine learning systems. Below, we detail these challenges and
their implications for LLM training:

(1) Scalablity. The success of LLMs is largely attributed to their scale, with
performance often improving as LLMs grow larger [8]. However, the scaling of
model size introduces substantial scalability challenges, since training LLMs
requires increasingly large clusters of GPUs or specialized AI accelerators.
First, building scalable infrastructure that can provide massive computations
and memory capacity is essential. This involves designing and deploying large
clusters of GPUs or specialized AI accelerators, high-performance networking
to connect these devices, and distributed storage systems capable of handling
enormous datasets and model checkpoints. The challenge lies in ensuring that
these components work together efficiently at scale, managing heat
dissipation, power consumption, and hardware failures in large-scale
deployments. Second, designing scalable training systems that can effectively
utilize massive accelerators in parallel is crucial. This includes designing
parallelization strategies and communication algorithms that can achieve near-
linear scalability with thousands of accelerators while maintaining consistent
LLM accuracy.

(2) Efficiency. The enormous computational requirements of LLM training
translate to high training costs, making it imperative to maximize the
efficiency of hardware and software systems. The efficiency can be measured
with MFU (Model FLOPs Utilization), which quantifies how effectively the
system uses available computing resources. However, achieving high efficiency
at scale remains a significant challenge. For instance, LLaMA3 achieves only
$38\%$ to $41\%$ MFU on 16K GPUs [9], highlighting the difficulty in
maintaining high utilization as systems scale. Maximizing efficiency demands
the optimization in parallelism, computation, communication and memory. First,
the parallelism of distributed LLM training demands careful design to minimize
communication demands. Second, optimized computation operators and lower
precision arithmetic are crucial to achieve high GPU FLOPS utilization. Third,
communication overhead needs to minimized to reduce GPU idle time. Finally,
efficient memory optimizations are required to hold LLMs in existing hardwares
and reduce FLOPs waste of recomputation.

(3) Reliability. Ensuring the reliability of LLM training over extended
periods is paramount. As training jobs can span weeks or months on large
clusters of tens of thousands of GPUs, the probability of training failures
increases, necessitating fast failure detection and recovery mechanism for
resilient LLM training. First, LLM training jobs may crash due to various
errors, making it hard to identify the exact fault reason across tens of
thousands GPUs quickly. Second, the hang of LLM training jobs results in all
GPUs becoming idle due to the synchronous nature of training, resulting in
significant waste. Moreover, some intricate anomalies, like redundant link
failures or stragglers, may not cause immediate crashes but can lead to
training slowdowns. This instability can result in reduced training
efficiency. To tackle these challenges, robust anomaly detection systems
capable of detecting both catastrophic failures and performance degradations
are essential. Additionally, implementing fault-tolerant training frameworks
that can seamlessly handle node failures and network issues is crucial.

###  2.4 Related Survey

This work focuses on the efficient training system and infrastructure of
transformer-based LLMs, including the design of underlying distributed
infrastructure, paradigm of parallelism, optimizations of computation and
communication, efficient memory management and resilience of training system.
We also investigate the efficient training systems of emerging workloads such
as MoE, a promising efficient LLM variant, and fine-tuning, a necessary stage
to align the capabilities of LLMs. However, this work does not cover the
evolution of promising LLM architectures [24, 25] and the algorithms for
training [26], instruction tuning [27] and alignment [28] towards powerful and
safe LLMs. While previous works [29, 30, 31] have discussed some aspects of
LLM training systems, their primary focus was not on the design of efficient
training systems and infrastructure. Wan et al. [29] aims to provide a
holistic view of efficient LLM advances in model- and data-centric methods.
Liu et al. [30] covers the training and inference deployment techniques of
LLMs. Xu et al. [31] targets on discussing the resource-efficient strategies
for LLM development in both algorithmic and systemic aspects. This work also
discuss the approaches for quantized LLM training and efficient LLM fine-
tuning, but we focus on the systemic approaches. The algorithmic approaches to
compress and fine-tune LLMs are discussed by Zhu et al. [32] and Han et al.
[33]. The discussion scope of this work does not include advanced optimization
algorithms [34] and distributed DNN training systems [35]. While Liang et al.
[36] have extensively reviewed auto parallelism approaches, their focus is on
general DNNs rather than LLMs specifically.

Figure 3: Infrastructure overview for distributed LLM training.

##  3 Infrastructure for LLM Training

forked edges, for tree= grow=east, reversed=true, anchor=base west, parent
anchor=east, child anchor=west, base=center, font=, rectangle, draw=black,
rounded corners, align=left, text centered, minimum width=4em, edge+=black,
line width=1pt, s sep=3pt, inner xsep=2pt, inner ysep=3pt, line width=0.8pt,
ver/.style=rotate=90, child anchor=north, parent anchor=south, anchor=center,
, where level=1text width=12em,font=,, where level=2text width=12em,font=,, [
Infrastructure for LLM Training, ver [ AI Accelerators, fill=blue!10 [ NVIDIA
GPUs, fill=yellow!10 [ Ampere [37], Hopper [38], Blackwell [39] , leaf, text
width= 32em ] ] [ Other AI Accelerators, fill=yellow!10 [ AMD GPU [40], GAUDI
[41], TPU [42], Graphcore IPU [43],
Cerebras CS-2 [44] , leaf, text width= 32em ] ] ] [ Network Infrastructure,
fill=blue!10 [ Chip-to-Chip, fill=yellow!10 [ Cube-Mesh Topology: NVLink-1.0
[45]
Fully-Connected Topology: NVSwitch [46], Infinity Fabric [47]
2D/3D-Torus Topology: TPUv2 [48], TPUv3 [49], TPUv4 [42] , leaf, text width=
32em ] ] [ Node-to-Node, fill=yellow!10 [ GPUDirect-RDMA [50], InfiniBand
EDR/HDR/NDR [51],
RoCE-v1 [52], RoCE-v2 [53], iWARP [54] , leaf, text width= 32em ] ] [ Network
Topology, fill=yellow!10 [ HPC Topology: Clos [55], BCube [56], DCell [57],
Jellyfish [58],
Torus [59], Dragonfly [60], Dragonfly+ [61]
Training-Optimized Topology: Rail-Optimized [62], HPN [63],
Rail-Only [64], BiGraph [63], HammingMesh [65]
Reconfiguable Topology: SiP-ML [66], TopoOpt [67], TPUv4 [42] , leaf, text
width= 32em ] ] [ Load Balancing & CC, fill=yellow!10 [ Load Balancing:  ECMP
[68], Multi-Flow with Enhanced-ECMP [9],
Packet spraying [69], Ethereal [70], HPN [63], MegaScale [71]
Congestion Control (CC): PFC [72], TIMELY [73], Swift [74],
DCQCN [75, 76], HPCC [77], EQDS [78], RoCC [79],
MLTCP [80], CASSINI [81], MLT [82] , leaf, text width= 32em ] ] ] [ Storage
Systems, fill=blue!10 [ Checkpoint Storage, fill=yellow!10 [ Tectonic [83],
HDFS [84], Ceph Object Storage [85] , leaf, text width= 32em ] ] [ Training
Data Storage, fill=yellow!10 [ Lustre [86], GPFS [87], BeeGFS [88], Alluxio
[89], JuiceFS [90],
Quiver [91], Fluid [92] , leaf, text width= 32em ] ] ] [ Scheduling Systems,
fill=blue!10 [ Workload Scheduling, fill=yellow!10 [ Tiresias [93], THEMIS
[94], ElasticFlow [95], Gavel [96],
Gandiva${}_{\text{fair}}$ [97], FGD [98], Lucid [99], Pollux [100], Sia [101],
Crius [102], Hydro [103], Acme [23] , leaf, text width= 32em ] ] [ Resource
Scheduling, fill=yellow!10 [ Cassini [81], HIRE [104], SiloD [105], Synergy
[106], EnvPipe [107],
Zeus [108], Perseus [109] , leaf, text width= 32em ] ] ] ]

Figure 4: Studies on infrastructure optimizations for distributed LLM
training.

In this section, we explore the infrastructure design for training LLMs,
encompassing accelerators, networks, storage, and scheduling systems (Fig. 3).

###  3.1 AI Accelerators

The rapid progress of LLMs is significantly benefiting from the evolution of
GPUs and AI accelerators, which are crucial for improving the model training
performance.

####  3.1.1 NVIDIA Graphics Processing Unit (GPU)

NVIDIA GPUs have become an essential component for distributed LLM training
due to their superior ability to handle parallel computations. These
processors are built with a multitude of compact, high-efficiency cores that
can execute numerous tasks simultaneously. The design of GPUs is ideally
matched for the matrix and vector operations in LLM training. They offer
support for various numerical precision formats such as FP32, TF32, FP16,
BF16, FP8, INT8, and even FP4. This allows researchers to well balance the
training speed and accuracy, making LLM training more efficient [110].
NVIDIA’s GPU programming language (i.e. CUDA) makes it easier to manage how
tasks are split and processed in parallel on GPUs. This helps researchers
harness the full power of GPUs for training advanced LLMs.

A typical GPU comprises an array of Streaming Multiprocessors (SMs), with each
SM housing several cores that share an instruction unit but are capable of
executing distinct threads in parallel. The shared memory within each SM
allows for effective data exchange and synchronization among threads, which is
essential for optimizing the memory access patterns required for LLM
computations. Furthermore, GPUs are furnished with high-bandwidth memory
(HBM), which accelerates data transfer and mitigates memory access bottlenecks
in computationally intensive tasks. The latest GPU architectures, such as
NVIDIA’s Ampere [37], Hopper [38] and Blackwell [39], are continuously pushing
the boundaries of LLM computation. They offer enhanced memory bandwidth and
capacity, increased floating-point operations per second (FLOPS), and
specialized mixed-precision computing units like Tensor Cores. Notably,
NVIDIA’s Hopper architecture introduces a significant advancement with the
Transformer Engine [111], a feature that leverages mixed FP8 and FP16
precisions to expedite the training of Transformer-based LLMs.

####  3.1.2 Other AI Accelerators

Distributed LLM training on AMD GPUs has become a reality, particularly on
Frontier [112], the world’s first exascale supercomputer. Each Frontier node
is equipped with 8 MI250X [40] AMD GPUs, each with 64 GB of HBM and a
theoretical FP16 peak performance of 191.5 TFLOPS. This configuration presents
an unparalleled opportunity for training trillion-parameter models
efficiently. The key to unlocking this potential lies in adapting existing
CUDA-based tools and frameworks to the ROCm platform [113, 114]. Notably,
ROCm-enabled versions of FlashAttention [115] and FlashAttention2 [116] have
been developed, allowing for the efficient execution of attention.

Various AI accelerators with powerful computing and software optimizations
have been developed to train LLMs. GAUDI [41] offers a heterogeneous compute
architecture comprising two Matrix Multiplication Engines and a cluster of
fully programmable tensor processor cores, enabling efficient handling of LLM
training operations. This processor can support the training of a GPT-3 model
with 175 billion parameters using 384 GAUDI2 cards [117]. Google TPUv4 [42]
supercomputer has 4096 chips, supporting LLM training at an average of
approximately $60\%$ of peak FLOPS. Graphcore Bow Pod64 [43], a one-rack setup
with 64 Bow-class IPUs, achieves 22 petaFLOPS. It supports GPT-3 model
training with 256 IPUs. Cerebras CS-2[44] is a wafer-scale deep learning
accelerator comprising 850,000 processing cores, each providing 48KB of
dedicated SRAM memory. It is used to train Cerebras-GPT, a family of open
compute-optimal language models [118].

###  3.2 Network Infrastructure

Communication overhead is a major obstacle to scaling LLM training [119, 120].
For example, reducing model gradients during training can lead to over 90% of
the training time being spent on communication [121]. To address this, the
research community has focused on improving communication infrastructure for
LLM training.

Figure 5: Five chip-to-chip topologies: tree topology, cube-mesh topology,
switch-based fully-connected topology, P2P-based fully-connected topology, and
2D-torus topology.

####  3.2.1 Chip-to-Chip Communications

Chip-to-chip communication is pivotal for data transfer between AI
accelerators within a node, significantly impacting the efficiency of LLM
training. Traditionally, this communication has relied on PCI Express (PCIe)
[122], which employs a tree topology—a hierarchical structure where multiple
devices connect to a single root complex. Over the years, PCIe has improved
its bandwidth: PCIe 3.0 offers approximately 1 GB/s per lane, totaling about
16 GB/s for a configuration with 16 lanes; PCIe 4.0 doubles the bandwidth to 2
GB/s per lane, while PCIe 5.0 further increases it to 4 GB/s per lane. Despite
these enhancements, PCIe’s inherent limitations in bandwidth, latency, and
scalability render it suboptimal for LLM training [123]. To address these
limitations, specialized chip-to-chip interconnects like NVLink [45] are
increasingly preferred for LLM training. Compared to PCIe, these advanced
interconnects provide significantly higher bandwidth and lower latency by
utilizing various topologies: cube-mesh, fully-connected, and 3D-torus.
Additionally, shared memory models, specialized communication protocols, and
synchronization mechanisms also play crucial roles.

Cube-Mesh Topology. NVLink-1.0 [45] offers a direct and high-speed connection
between GPUs, with each link providing 160 GB/s of bidirectional bandwidth.
This architecture enables the formation of planar mesh structures for four
GPUs and a cube-mesh topology for eight GPUs, which can be configured into a
DGX-1 server. This cube-mesh configuration, although not an all-to-all
connection, significantly enhances data communication efficiency and training
performance on GPUs.

Fully-Connected Topology. Many interconnects utilize either switch-based or
P2P-based fully connected topologies to improve chip-to-chip communication
performance. NVIDIA uses NVSwitch [46] to achieve switch-based all-to-all
interconnections among GPUs. In the DGX-2 [124] system, six NVSwitches fully
connect each of the sixteen GPUs to all others, providing a bidirectional
bandwidth of 300 GB/s between any two GPUs. This bandwidth increased to 600
GB/s with NVSwitch 2.0 and further to 900 GB/s with NVSwitch 3.0. Intel, AMD,
and Huawei Ascend utilize P2P-based fully-connected topology for their
accelerators, where each chip directly connects to every other chip within the
same node using Ethernet or Infinity Fabric [47]. Compared to the switch-based
topology, the bandwidth between two GPUs in a P2P-based topology is limited by
the bandwidth of the directly connected link.

2D/3D-Torus Topology. Google’s TPU systems utilizes Torus network topology
[59] for chip-to-chip communication. It establishes connectivity by linking
each TPU chip to its four adjacent neighbors in a grid, with wraparound edges,
thus forming a toroidal structure. This architectural design ensures low
latency and high bandwidth due to the availability of multiple direct paths
between chips. Specifically, the TPUv2 [48] supercomputer employs a 16x16 2D
torus configuration, encompassing 256 chips, interconnected via high-speed
Inter-Chip Interconnect (ICI) links. The TPUv3 [49] supercomputer utilizes a
32x32 2D torus, comprising 1024 chips. Advancing from the 2D torus design, the
TPUv4 [42] supercomputer organizes compute resources into multi-machine cubes
with a 3D torus topology. Each TPU machine contains four chips arranged in a
2x2x1 mesh, interconnected through the ICI links. Sixteen of these TPU
machines are combined to form a data center rack, where ICI links within the
rack interconnect to create a 4x4x4 mesh, resulting in a 3D torus structure.
This advanced configuration significantly enhances communication efficiency
and scalability, particularly beneficial for LLM training.

####  3.2.2 Node-to-Node Communications

Remote Direct Memory Access (RDMA)[54] enables high-speed and low-latency data
transfer between nodes. RDMA allows direct memory access from the memory of
one computer to another without involving either node’s operating system.
GPUDirect-RDMA[50] enhances this process by enabling direct communication
between GPUs across different nodes, bypassing the CPU entirely. This
technology is particularly beneficial for LLM training, as it accelerates the
synchronization of model parameters and gradients. The two most prevalent RDMA
technologies are InfiniBand [51] and RDMA over Converged Ethernet (RoCE)[52].

InfiniBand is a high-speed, low-latency networking technology widely utilized
in HPC (High Performance Computing) environments, such as the Eagle
supercomputer [125]. This technology necessitates a dedicated network
infrastructure, reflecting its design focus on delivering superior
performance. Over the years, InfiniBand has significantly evolved in terms of
bandwidth capabilities, advancing from EDR (Enhanced Data Rate) at 100 Gbps to
HDR (High Dynamic Range) at 200 Gbps, and more recently to NDR (Next Data
Rate) at 400 Gbps per link [126]. RoCE leverages the existing Ethernet
infrastructure to deliver RDMA capabilities. This approach offers a more cost-
effective and easier-to-deploy solution, particularly in data centers that
already utilize Ethernet. RoCE is available in two versions: RoCE-v1 [52],
which operates as an Ethernet link layer protocol, and RoCE-v2 [53], which
operates over UDP. Industry leaders such as ByteDance and Meta have employed
these technologies to scale LLM training. Another RDMA protocol, the Internet
Wide Area RDMA Protocol (iWARP) [54], enables RDMA over TCP/IP networks.
However, due to its comparatively limited performance, iWARP is not commonly
used for distributed LLM training[127].

####  3.2.3 Network Topology

Figure 6: Four typical network topologies in large-scale GPU clusters: Clos
topology, Dragonfly+ topology, rail-optimization topology, and rail-only
topology.

In LLM training clusters, the network architecture is structured into frontend
and backend components (Fig. 3). The frontend network handles a variety of
traffic, such as job management, model inference, and storage activities,
while the backend network is dedicated to the high-volume traffic generated
during the training process. Our primary focus in optimizing LLM training lies
in improving the performance and efficiency of this backend network, so as to
scale AI accelerators to tens of thousands.

HPC Network Topology. Conventional topologies for HPC environments can also be
used in AI clusters for distributed training, such as Clos [55], BCube[56],
DCell[57], Jellyfish[58], Torus [59], Dragonfly [60] and Dragonfly+ [61]. The
Clos network architecture, commonly known as a Fat-Tree topology, is widely
used in LLM training clusters. In a Clos-based cluster, each server, equipped
with one or more NICs, is organized into racks connected to leaf switches.
These leaf switches link to spine switches, providing inter-rack connectivity
and forming a pod. The pods are further interconnected with core switches,
facilitating any-to-any communication across servers within the cluster. For
example, Meta’s last-generation GPU cluster architecture, supporting up to
24,000 GPUs, consists of eight pods with full-fat bandwidth between them and
use a 7:1 oversubscription ratio on the core layer [9]. Meta uses 24,000 GPUs
to this cluster for training Llama 3 405B.

Training-Optimized Topology. Many network topologies are co-designed with
distributed training algorithms. The rail-optimized topology [62] enhances the
connection from GPUs to leaf switches. Within each rail, GPUs that share the
same index across different servers are interconnected via the same leaf
switch. This configuration improves collective communication performance by
reducing network interference between data flows. The SuperPod architecture
utilizes a rail-optimized network and is capable of connecting over 16,000
GPUs [128]. ByteDance employs a three-layer rail-optimized network to connect
more than 10,000 GPUs in its MegaScale system design [71]. However, rail-
optimized network designs could be less efficient because they necessitate
connecting GPUs to distant switches, which requires costly and power-hungry
optical transceivers. These optical components increase power consumption and
heat, leading to higher rates of network failures, which is significant for
distributed LLM training. Alibaba further optimizes the rail-optimized
topology with a 2-tier, dual-plane architecture named HPN [63]. This
architecture employs the latest 51.2Tbps single-chip switch and supports 1,000
GPUs in a tier1 network and up to 15,000 GPUs within one pod.

The network traffic analysis of GPT/OPT-175B model training reveals that 99%
of GPU pairs do not carry any traffic, and less than 0.25% of GPU pairs handle
pipeline/tensor parallel and data parallel traffic [64]. Based on these
findings, the rail-only topology [64] eliminates the connection between
different rails on rail-optimized network. Each rail is connected by a
dedicated but separate Clos network. Communication across GPUs on different
rails is managed by forwarding the data through internal chip-to-chip
interconnects. This method could effectively reduce costs while maintaining
performance. HammingMesh [65] organizes GPUs into groups with a 2D-torus
topology and connects these 2D-torus groups via sparsely connected switches.
This design aims to save costs without compromising training performance.
Given GPUs that are only connected by PCIe, BiGraph [129] proposes a new
network architecture that exports intra-node GPU communication to outside the
node, bypassing PCIe bandwidth bottlenecks. It features a two-layer network
interconnected via a Clos architecture, with unique shortest paths for
communication that support application-controlled traffic routing.

Reconfiguable Topology. Reconfigurable network can be dynamically adjusted to
optimize communication patterns for improving the training performance. They
typically utilize optical switching and customized configurations to improve
bandwidth utilization, flexibility, and scalability of the network
infrastructure. Driven by Silicon Photonic (SiP) interfaces, SiP-ML [66]
advances two principal architectures: SiP-OCS and SiP-Ring. SiP-OCS
incorporates a fully connected configuration that maximizes the bandwidth by
employing commercially accessible optical circuit switches, linking GPUs to
all switches via Tbps SiP interfaces. Conversely, the SiP-Ring utilizes a
switchless ring configuration, reducing the reconfiguration latency by
integrating micro-ring resonators within the SiP interfaces. Wang et al.
proposed TopoOpt [67] for co-optimizing the network topology and
parallelization strategies in distributed training. This approach not only
optimizes the computational and communication demands but also addresses the
physical layer of network topology. TPUv4 [42] features Optical Circuit
Switches (OCS), allowing for dynamic reconfiguration of the 3D-Torus based
interconnect topology, optimizing the data flow for the varied and intensive
communication patterns that characterize LLM training. For instance, with 512
chips, TPUv4 offers flexibility in 3D-Torus topologies such as 4x4x32 or
8x8x8.

####  3.2.4 Load Balancing & Congestion Control

Load Balancing. The network traffic of LLM training is characterized by a
small number of elephant flows. Specifically, LLM training demonstrates
periodic bursts of network traffic due to gradient synchronization [63]. Each
burst demands significant network bandwidth. Moreover, each compute node
involved in LLM training generates very few connections [63]. The conventional
load-balancing technique, ECMP (Equal-Cost Multi-Path routing) [68], uses hash
algorithms to evenly distribute traffic across equivalent paths, such as those
from leaf switches to spine switches in a Clos topology. However, this hash-
based scheme is inefficient for handling LLM training traffic, which consists
of a small number of elephant flows. When multiple elephant flows are routed
to the same link, it can result in congestion and high latency.

Various strategies have been developed to address the load balancing
challenges in large-scale GPU clusters. During the Llama 3 405B training, the
collective library establishes 16 network flows between two GPUs, rather than
a single flow, thereby reducing traffic per flow and enhancing load balancing
opportunities [9]. Additionally, the Enhanced-ECMP (E-ECMP) protocol
effectively distributes these 16 flows across different network paths by
hashing on additional fields in the RoCE header of packets. Packet spraying
[69] distributes packets from a flow across all available parallel links,
which can lead to out-of-order packets. NICs need to process out-of-order RDMA
packets. Based on the traffic pattern of LLM training, Ethereal [70]
demonstrates that greedily assigning paths to each flow can uniformly
distribute the load across all network paths and resolve the ECMP hash
conflict problem. In a large-scale GPUs cluster, HPN [63] achieves efficient
load balancing by identifying precise disjoint equal paths and balancing the
load within the collective communication library. MegaScale [71] shows that a
rail-optimized topology can also mitigate ECMP hashing conflicts.

Congestion Control. Lossless transmission is crucial in RDMA clusters.
Priority-based Flow Control (PFC) [72] is a flow control mechanism that
prevents packet loss. When congestion occurs in a PFC-enabled queue on a
downstream device, the device instructs the upstream device to halt traffic in
the queue, thereby ensuring zero packet loss. Since PFC is a coarse-grained
mechanism, it can lead to head-of-line blocking [130], which significantly
diminishes network throughput. To address these challenges, various general-
purpose congestion control schemes have been developed. These techniques
include TIMELY [73], Data Center Quantized Congestion Notification (DCQCN)
[75, 76], Swift [74], High Precision Congestion Control (HPCC) [77], Edge-
Queued Datagram Service (EQDS) [78], and Robust Congestion Control (RoCC)
[79]. These schemes monitor network congestion, adjust data rate to alleviate
congestion, and recover the rate to minimize throughput reduction.

When there are concurrent training jobs, many congestion control schemes
leverage the bursty and periodic traffic patterns to interleave the network
traffic effectively. MLTCP [80] interleaves the communication phases of jobs
that compete for bandwidth based on a key insight: training flows should
adjust their congestion window size based on the number of bytes sent in each
training iteration. CASSINI [81] optimizes job placement on network links by
considering the communication patterns of different jobs. MLT [82] leverages
the characteristics of LLM training, where the gradients of earlier layers are
less important than those of later layers, and larger gradients are more
significant than smaller ones. Consequently, in the event of communication
congestion, MLT prioritizes queuing or discarding packets based on the
importance of the gradients within them at the switch level to mitigate
communication congestion issues.

###  3.3 Storage

The storage system plays a critical role in distributed LLM training and need
to meet several key requirements. Firstly, it should align with the computing
power of GPUs to maximize their utilization and avoid resource wastage caused
by storage bottlenecks. Secondly, it should support the storage of large-scale
structured and unstructured training datasets and be scalable in distributed
processing environments. Additionally, the storage and retrieval of model
checkpoints present challenges in LLM training, requiring the system to meet
the writing and reading bandwidth dictated by the model size and training
duration. Lastly, the storage system should satisfy traditional enterprise-
level requirements, such as data protection, high availability, and security.

####  3.3.1 Storage Systems for Checkpoint

The model checkpoint size is enormous in LLM training. As the number of
parameters increases, so does the volume of data that needs to be written,
demanding greater write bandwidth from the storage system. For example, the
checkpoint size is 980GB for an LLM with 70B parameters. Numerous storage
systems have been deployed in large-scale GPU data centers to manage model
checkpoints. Meta’s distributed filesystem, Tectonic [83], enables thousands
of GPUs to save and load model checkpoints simultaneously, providing efficient
and scalable storage solutions for extensive training operations [131]. At
ByteDance, HDFS [84] is employed for centralized model checkpoint maintenance,
ensuring consistency and reliability at scale [71]. To mitigate bandwidth
bottlenecks during checkpoint recovery, a common approach designates a single
worker to read the checkpoint partition from HDFS and then broadcast it to
other workers sharing the same data. Distributed object stores, such as Ceph
Object Storage [85], offer easier scalability. This advantage stems from their
lack of a hierarchical directory tree or namespace, simplifying consistency
maintenance. Due to these benefits, object stores have become widely adopted
for model checkpoint storage.

####  3.3.2 Storage Systems for Training Data

The raw dataset for LLM training is substantial. LLaMA 3 was trained on over
15 trillion tokens, which is more than seven times larger than LLaMA 2’s
dataset [6]. Each token requires about 2 bytes, equaling roughly 30 TB of
data. Preparing datasets for training involves extensive preprocessing steps,
including data crawling and cleaning, requiring significant experimentation.
Typically, the data processed during these steps exceeds 100 times the final
size of the training dataset [132]. For instance, the WanJuan-CC dataset [132]
selectively extracts approximately 68 billion documents, generating around 1
trillion high-quality tokens, equivalent to a data size of 2 TB, after
discarding $99\%$ of the raw data. Therefore, the total data volume for LLM
training is expected to exceed tens of PB.

Parallel file systems such as Lustre [86], GPFS [87], and BeeGFS [88] are
frequently deployed on leading high-performance computing systems to ensure
efficient I/O, persistent storage, and scalable performance. These systems are
also widely used in training clusters for data loading, providing the
necessary infrastructure to handle large-scale training data efficiently.
Additionally, it is crucial for file systems to enable engineers to perform
interactive debugging on jobs utilizing thousands of GPUs, as code changes
need to be immediately accessible to all nodes [131].

During the training of most LLMs, each token is typically encountered only
once. However, employing data caching remains crucial to mitigate I/O
bottlenecks during data loading. This strategy involves prefetching training
data from slower backend storage to faster cache storage. Alluxio [89] and
JuiceFS [90] enhance LLM training by efficiently caching training data from
underlying storage systems such as HDFS or object storage. Quiver [91]
supports the transparent reuse of cached data across multiple jobs and users
operating on the same dataset. Fluid [92] leverages Alluxio for data caching,
incorporating a mechanism that enables on-the-fly autoscaling of the cache
based on I/O conditions.

###  3.4 Scheduling

LLM training workloads generally run on large-scale multi-tenant
infrastructures (e.g., GPU clusters, public clouds) where users share cluster
resources. Effective scheduling mechanisms are crucial for managing these
workloads, ensuring efficient resource utilization and task execution [133].
Unlike _task-level_ scheduling (e.g., pipeline scheduling [134, 135, 136]),
which focuses on fine-grained optimization of single-job execution (§4.1.3),
_cluster-level_ scheduling aims to optimize resource allocation and task
scheduling across the entire cluster. We categorize existing _cluster-level_
scheduling systems into two types, workload scheduling and resource
scheduling, according to their primary optimization aspects.

####  3.4.1 Workload Scheduling

Schedulers tailored for DL training workloads have been actively explored in
recent years [137, 138, 139, 140, 141, 93, 94, 95]. To enhance resource
utilization, three advanced features are commonly implemented: (1)
_heterogeneous-aware_ schedulers (e.g., Gavel [96], Gandiva${}_{\text{fair}}$
[97]) focus on optimizing job allocation across different GPU generations; (2)
_job-packing_ schedulers (e.g., FGD [98], Lucid [99]) enable fine-grained GPU
sharing to fully facilitate hardware capability; (3) _adaptive-scaling_
schedulers (e.g., Pollux [100], Sia [101]) dynamically adjust the number of
GPUs as well as the training hyperparameters to accelerate training progress.
However, these schedulers are designed for general DL workloads and may not be
directly applicable to LLMs due to the unique characteristics of LLM workloads
[23].

To better manage LLM workloads, several recent studies have proposed systems
tailored to LLMs. Crius [102] jointly considers hybrid parallelism (§4.1) and
hardware-affinity within heterogeneous clusters. It investigates the workflow
efficiency of integrating adaptive parallelism configuration at the cluster
scheduling level, offering significant opportunities to improve the efficiency
of training multiple LLMs concurrently. To achieve highly efficient
hyperparameter tuning for LLMs, Hydro [103] scales down the model to a smaller
surrogate model for hyperparameter search, and then fuses multiple models into
a single entity to enhance the hardware utilization. Additionally, Hydro
extends the resources for tuning workloads by interleaving them with pipeline-
enabled LLM pretraining tasks, effectively utilizing the pipeline bubbles.
Acme [23] further characterizes the workload mixture of the LLM development
workflow and proposes a system to efficiently schedule associated jobs related
to LLM training, including decoupled evaluation scheduling for timely model
quality feedback as well as LLM-involved failure diagnosis and automatic
recovery.

####  3.4.2 Resource Scheduling

In addition to workload scheduling, associated resource scheduling (e.g., CPU,
memory, and networking) is another critical aspect of cluster-level
management. For networking, Cassini [81] enables the interleaving of bandwidth
demands during the up and down phases of different jobs by using an affinity
graph to determine time-shift values for adjusting communication phases. HIRE
[104] introduces an innovative in-network computing scheduling system for
datacenter switches, significantly reducing the network detours and tail
placement latency. For storage, SiloD [105] treats data cache and remote I/O
as first-class resources for joint allocation, showing significant throughput
improvements. For CPU and memory, Synergy [106] enhances training efficiency
by optimizing CPU core allocations instead of relying on GPU-proportional
allocation. Additionally, some works focus on energy conservation. EnvPipe
[107] leverages the bubble time in the pipeline parallelism, stretching the
execution time of pipeline units by lowering the SM frequency to save energy.
Zeus [108] automatically configures the batch size and GPU power limit for
improving energy efficiency during training. Perseus[109] introduces an
efficient graph cut-based iterative algorithm to obtain the iteration time-
energy Pareto front for large model training job.

##  4 Parallelism Schemes for LLM Training

forked edges, for tree= grow=east, reversed=true, anchor=base west, parent
anchor=east, child anchor=west, base=center, font=, rectangle, draw=black,
rounded corners, align=left, text centered, minimum width=4em, edge+=black,
line width=1pt, s sep=3pt, inner xsep=2pt, inner ysep=3pt, line width=0.8pt,
ver/.style=rotate=90, child anchor=north, parent anchor=south, anchor=center,
, where level=1text width=12em,font=,, where level=2text width=12em,font=,, [
Parallelism Schemes for LLM Training, ver [ Hybrid Parallelism, fill=blue!10 [
Data Parallelism, fill=yellow!10 [ PyTorch-DDP [142], Horovod [143], Sharded
Weight Update [144],
ZeRO-3 [145], FSDP [146] MiCS [147], leaf, text width= 32em ] ] [ Tensor
Parallelism, fill=yellow!10 [ Megatron-TP [20], Optimus [148], Tesseract
[149], 3-D TP [150], leaf, text width= 32em ] ] [ Pipeline Parallelism,
fill=yellow!10 [ Pipeline Bubble: GPipe [151], PipeDream [134, 135],
DAPPLE [152], Interleaved 1F1B [153], TeraPipe [154], Tessel [155],
Zero Bubble [156], Chimera [136], Hanayo [157], Seq1F1B [158],
Breadth-First [159], DynaPipe [160], DISTMM [161], GraphPipe [162]
Memory Imbalance: BPipe [163], MPress [164], mCAP [165],
Chimera [136], Hanayo [157], V-Shape [166], Varuna [167],
AdaPipe [168], leaf, text width= 32em ] ] [ Sequence Parallelism,
fill=yellow!10 [ Ring Self-Attention[169], Megatron-SP [170], USP [171],
DeepSpeed-Ulysses [172], Context Parallel [173], LoongTrain [174],
Blockwise Ring Attention [175], DistFlashAttn [176], DSP [177],
Striped Attention [178], BurstAttention [179], WallFacer [180], leaf, text
width= 32em ] ] [ Expert Parallelism, fill=yellow!10 [ Sparse Activation:
Gshard[15], Switch transformer[181], Tutel[182],
Deepspeed-MoE[183], Deepspeed-TED[184], Megablock[185],
ScatterMoE[184]
Communication Optimization: PipeMoE[186], ScheMoE[187],
Lina[188], Janus[189], TA-MoE[190]
Load Balance: FasterMoE[191, 192], SmartMoE[193],
FlexMoE[194], Prophet[195], leaf, text width= 32em ] ] ] [ Auto Parallelism,
fill=blue!10 [ General Framework, fill=yellow!10 [ Mesh-TensorFlow [196],
GSPMD [197], OptCNN [198], Alpa [21],
FlexFlow [199], Tofu [200], HyPar [201], AutoMap [202], Slapo [203],
TensorOpt [204], PipeDream [134], DAPPLE [152], AutoPipe [205],
Device Placement [206, 207, 208], Piper [209], Unity [210], Aceso [211],
PartIR [212], nnScaler [213], OneFlow [214], AutoDDL [215], leaf, text width=
32em ] ] [ Transformer-Specific, fill=yellow!10 [ DeepSpeed-Autotuning [216],
Galvatron [217], Merak [218],
Colossal-Auto [219], Galvatron-BMW [220], leaf, text width= 32em ] ] ] [
Heterogeneous Parallelism, fill=blue!10 [ Heterogeneous Hardware,
fill=yellow!10 [ HetPipe [221], AccPar [222], Whale [223], AMP [224],
Pathways [225], HPH [226], SDPIPE [227], HAP [228],
PipePar [229], Yuan et al. [230], SWARM [231], FusionAI [232], leaf, text
width= 32em ] ] [ Heterogeneous Model, fill=yellow!10 [ DeepSpeed-Chat [233],
HuggingFace TRL [234], OpenRLHF [235],
Adpative Placement and Parallelism [236], ReaLHF[237],
PUZZLE[238] , leaf, text width= 32em ] ] ] ]

Figure 7: Studies on parallelism schemes for distributed LLM training.

The continual growing scales of LLMs demand substantial computational
resources and memory capacity. Distributed training, leveraging large-scale
HPC clusters, has emerged as a crucial approach to efficiently train these
models. In this section, we investigate various parallelism schemes proposed
to enhance the utilization of HPC clusters for LLM training. We categorize
these approaches into three main groups: Hybrid Parallelism, Auto Parallelism,
and Heterogeneous Parallelism. Hybrid Parallelism combines multiple
handcrafted parallelization strategies, such as data parallelism, tensor
parallelism, pipeline parallelism, sequence parallelism, and expert
parallelism. Auto Parallelism automatically determines the optimal
parallelization strategy based on the model and hardware characteristics.
Heterogeneous Parallelism exploits the heterogeneity in hardware or model for
efficient training. This includes techniques that leverage different types of
accelerators or leverage the heterogeneity within a single model (e.g., RLHF
training) to improve the overall training efficiency on HPC clusters.

Most of today’s state-of-the-art parallelization strategies adopt a Single
Program Multiple Data (SPMD) programming model, akin to the MPI paradigm
[239], where the same program runs across multiple processors, each working on
different pieces of data [225]. For example, data, model and sequence
parallelism utilizes the SPMD programming model. This approach ensures
uniformity and consistency in operations, making it well-suited for large-
scale, distributed training environments. Some strategies explore to break the
restriction of SPMD and further improve the resource utilization with the
Multiple Program Multiple Data (MPMD) model, where different programs (or
different parts of a program) run on different processors, handling different
parts of the data or model [225]. For example, pipeline parallelism runs
different parts of an LLM on different devices. In addition, auto parallelism
and heterogeneous parallelism can leverage both SPMD and MPMD models to
increase the resource utilization. Therefore, we discuss these approaches
according to the dimensions along which parallelism occurs and whether the
computing devices employed are homogeneous or heterogeneous, rather than
focusing on the underlying programming models.

###  4.1 Hybrid Parallelism

Hybrid parallelism typically combines multiple handcrafted parallelization
strategies to partition different parallelizable dimensions of an LLM. These
strategies include data parallelism, tensor parallelism, pipeline parallelism
and sequence parallelism, as illustrated in Fig. 8. The combination of data
parallelism, tensor parallelism, and pipeline parallelism is also referred to
as 3D parallelism.

####  4.1.1 Data Parallelism

Data parallelism is the most commonly used parallelization strategy for
distributed training due to its high scalability and ease of implementation.
It follows the Single Program Multiple Data (SPMD) programming model. Data
parallelism partitions the input training data along the batch dimension,
where each GPU processes its assigned segment of data, as shown in Fig. 8(a).
Throughout the training process, the data first undergoes forward computation
with the full model weights layer by layer, and then performs backward
computation in the reverse order. Each layer produces gradients that will be
aggregated across all GPUs using collective communication operations for
optimizer updates.

Data parallelism incorporates various sharding strategies, which significantly
influence the memory footprint and communication overhead. Supposing the
global world size is $W$ (i.e. the number of devices), a sharding factor $F$
is introduced to control the sharding strategy used [146], defined as the
number of devices across which parameters are partitioned ($1\leq F\leq W$).
We have the following scenarios.

Full replication ($F=1$): this sharding strategy is simplified as vanilla data
parallelism. Pytorch-DDP [240] and Horovod [143] fully replicate the model
across all devices and use All-Reduce for gradient aggregation. They also
divide the gradients into small buckets to overlap the gradient communication
with backward computation.

Full sharding ($F=W$). This sharding strategy comes with the lowest memory
consumption but the most communication overhead ($1.5\times$ over vanilla data
parallelism). Full sharding strategy fully shards the model, where each device
holds only $\frac{1}{W}$ of the model parameters. The full weights and
gradients are communicated and recovered on-demand before the computation, and
immediately discarded afterwards. ZeRO-3 [145] employs per-parameter sharding
to shard the full model and utilizes All-Gather and Reduce-Scatter for
unsharding and sharding communication, respectively. Sharded Weight Update
[144] also employs per-parameter sharding but focuses more on sharding the
redundant parameter update computation across all the devices. FSDP (Fully
Sharded Data Parallelism)[146] achieves the same functionality by sharding
model parameters in the grain of module units and provides more user-friendly
API.

Hybrid sharding ($1<F<W$). In this strategy [146], all the devices are divided
into a $N\times M$ device mesh. The model parameters are sharded along the $N$
dimension of the mesh, and replicated along the $M$ dimension. MiCS [147]
invokes All-Gather collective to gather the sharded parameters and All-Reduce
to aggregate the gradients. FSDP [146] replaces the All-Reduce with Reduce-
Scatter to reduce the memory and communication overhead. Compared to full
replication and full sharding, hybrid sharding is more flexible to provide a
trade-off between memory consumption and communication overhead via adjsuting
$F$ based on the model architecture and hardware constraints.

Figure 8: An example of 3D-parallelism with data parallelism, tensor
parallelism, and pipeline parallelism.

####  4.1.2 Tensor parallelism

Tensor parallelism (Fig. 8(b)), also known as intra-layer model parallelism,
is a technique proposed to enable the training of LLMs across multiple GPUs.
It partitions the parameter tensors of each layer along multiple dimensions,
effectively distributing the model parameters across the available GPUs.
Tensor parallelism communicates intermediate activation tensors, the size of
which is much smaller than that of parameters and gradients communicated in
data parallelism, except for long context LLM training scenarios. However, in
tensor parallelism, it is challenging to overlap the communication with
computation, necessitating the use of high-bandwidth connections.
Consequently, tensor parallelism is more commonly employed in a single GPU
node.

Tensor parallelism can be divided into 1-D [20], 2-D [148], 2.5-D [149] and
3-D [150] parallelism according to the dimensionality of the partition. There
are two parameter matrices in both the MLP and self-attention module of
transformer-based LLMs. Megatron-LM [20] first adopts 1-D tensor parallelism
to partition the first parameter matrix along its column, and the second
parameter matrix along its row. It replicates the input and output tensors of
each partitioned module and introduces two All-Reduce collective communication
across all GPUs to fit an LLM into multiple GPUs. Inspired by Scalable
Universal Matrix Multiplication Algorithm (SUMMA) [241] and Cannon’s algorithm
[242] for 2-D parallel matrix multiplication, Optimus [148] further partitions
the input and parameter tensors in 2 dimensions to improve the communication
and memory efficiency of 1-D tensor parallelism. Tesseract [149] extends the
2.5-D matrix multiplication method [243], which is proposed to improve the
efficiency of Cannon’s algorithm, to LLM training and proposes the 2.5-D
tensor parallelism to overcome the abundance of unnecessary communication
resulting from the increasing model size. 3-D tensor parallelism [150] adopts
and improves the 3-D parallel matrix multiplication algorithm [244] for linear
operations and achieves a perfect load balance across multiple devices for LLM
training.

####  4.1.3 Pipeline Parallelism

Pipeline parallelism (Fig. 8(c)) [151], also known as inter-layer model
parallelism, is proposed to accommodate large models across multiple GPUs,
particularly across different nodes. Pipeline parallelism partitions the
layers of a model into multiple stages, where each stage consists of a
consecutive set of layers in the model and is mapped to a set of GPUs. Unlike
tensor parallelism, which typically demands high-bandwidth connections like
NVLink for communication, pipeline parallelism only necessitates the exchange
of intermediate tensors at designated cutting points, resulting in less
frequent communication requirements. Therefore, pipeline parallelism is
suitable for scaling up the LLM training across multiple GPU nodes that are
connected with small bandwidth. For example, Strati et al. [245] adopt
pipeline parallelism to fully utilize geo-distributed resources to overcome
the shortage of GPUs. Due to the data dependency of different stages, pipeline
parallelism typically splits the input data into multiple micro-batches for
pipelining to enable the efficient training of giant models. However, it comes
with two significant problems. First, the pipeline bubble problem reduces the
utilization of GPUs due to the time spent on waiting for the output of the
previous stage. Second, there exists memory consumption imbalance across
different stages since the former stages need to hold more active micro-
batches than the latter for better pipelining and higher utilization. We
detail each problem below.

Pipeline Bubble. Efficient micro-batch scheduling algorithms are proposed to
reduce pipeline bubbles. GPipe [151] introduces a fill-drain schedule that
injects all micro-batches at once for forward pass execution, followed by
backward passes. Gpipe incurs significant pipeline bubbles due to the warm up
and cool down of both forward and backward passes. PipeDream [134, 135]
introduces a 1F1B (1 Forward 1 Backward) schedule, which executes the backward
pass of a micro-batch as soon as the corresponding forward pass has completed,
to reduce the pipeline bubbles in the asynchronous scenario. DAPPLE [152]
employs an early backward schedule to first inject a fixed number of micro-
batches at the beginning of each stage and then interleave forward and
backward passes with round robin. Interleaved 1F1B [153] adapts the 1F1B
schedule but assigns multiple stages to each GPU (i.e. looping pipeline
placement). The pipeline bubble is reduced at the cost of higher communication
and peak memory consumption. Chimera [136] introduces a bidirectional pipeline
to reduce bubbles with weight duplication. Hanayo [157] further proposes a
wave-like pipeline that assigns multiple symmetrical stages to one GPU to
improve the pipeline utilization. Zero bubble [156] splits the backward
computation into two parts: activation and parameter gradient computation. It
schedules the forward and activation gradient computation with 1F1B and then
fills the bubbles with parameter gradient computation, which reduce bubbles
with higher peak memory consumption. Breadth-First [159] runs all micro-
batches at once in looping pipeline placement to reduce the communication
overhead when combined with sharded data parallelism. TeraPipe [154] splits
the micro-batch along the sequence dimension and exploits more fine-grained
token parallelism to reduce pipeline bubbles. However, The memory overhead of
TeraPipe is large since it is based on the GPipe schedule. Seq1F1B [158]
splits the sequence into chunks and utilizes the 1F1B schedule to reduce the
peak memory consumption while achieving low pipeline bubble rates. DynaPipe
[160] uses a dynamic micro-batching approach to the multi-task training of
LLMs with variable-length input. It introduces a memory-aware adaptive
scheduling algorithm and ahead-of-time communication planning to further
reduce the pipeline bubble rates. Tessel [155] is a two-phase approach,
including repetitive pattern construction and schedule completion, to
automatically search for the efficient pipeline schedule for a specified
partition strategy. DISTMM [161] launches doubled micro-batches to bypass the
dependency barrier caused by the large batch requirement of multi-modal
training, thus reducing idle cycles. GraphPipe [162] preserves the DNN graph
topology and partitions it into stages that can be concurrently executed to
improve pipeline utilization and reduce memory consumption.

Memory Imbalance. Pipeline parallelism typically injects more micro-batches
into the beginning stages to improve the pipeline utilization, resulting in
higher activation memory consumption in these stages. To resolve this problem,
BPipe [163] and MPress [164] employ D2D (device to device) transfer to swap
intermediate activation tensors from high-load GPU to light-load GPU during
runtime. MPress also combines the activation recomputation to reduce the
memory footprint. Chimera [136] introduces a bidirectional pipeline that
combines two pipelines in different directions together to achieve more
balanced memory consumption. Each GPU holds two symmetric stages, leading to
weight duplication. Hanayo [157] turns the bidirectional pipeline into two
data parallel pipelines to remove the weight duplication and achieves a
balanced memory consumption by assigning multiple stages to one GPU
symmetrically. V-Shape [166] partitions the model into stages twice the number
of devices, where the two halves of stages are placed in reverse order. By
varying the offsets between stages, V-Shape makes a trade-off between peak
memory consumption and bubble utilization. mCAP [165] utilizes an incremental-
profiling approach to partition models evenly across GPUs with respect to peak
memory usage.

Peak memory consumption limits the number of active micro-batches in pipeline
parallelism, thus its efficiency. Activation recomputation can be employed to
reduce the peak memory consumption effectively. Varuna [167] combines pipeline
parallelism and activation recomputation to achieve this goal. It designs a
static rule-based schedule enumerated for a given pipeline with an
opportunistic policy to hide jitter and reduce bubbles. The static schedule is
generated based on constraints including activation recomputation timing,
activation memory management, and backward computation prioritization. To
resolve the memory imbalance with low recomputation overhead, AdaPipe [168]
adopts adaptive recomputation to support different recomputation strategies
for different stages, and adaptive partitioning based on the 1F1B schedule to
balance the computation of each stage.

####  4.1.4 Sequence parallelism

The context window of today’s LLMs grows rapidly, and the most powerful LLM
can support millions of tokens [7]. Such ultra long sequence leads to
significant memory and computation requirements for LLM training: linearly
increasing memory footprint of activations and quadratic complexity of
attention mechanism. Recomputing activations in the backward can reduce the
peak memory consumption but also introduce significant overheads (30% with
full recompute). Large tensor parallelism degree incurs significant
communication overhead. Sequence parallelism (Fig. 8(d)) [169, 170] is
proposed to accommodate the long sequence training and distribute the
computation in multiple GPUs efficiently within the memory capacity. It
divides the input data into multiple chunks along the sequence dimension and
each chunk is fed to one GPU for computation. Since sequence parallelism
replicates the model parameters, it is typically combined with tensor and
pipeline parallelism to scale up the LLM training. When used together with
tensor parallelism, sequence parallelism distributes the memory and
computation of attention on multiple GPUs, but incurs redundant memory
consumption and computation in the non-tensor parallel regions of a
transformer layer. Megatron-SP [170] splits these computations along the
sequence dimension to reduce redundant the activation computation and memory
consumption without increasing the communication.

Although sequence parallelism partitions the memory, computation and
communication across multiple GPUs, the quadratic causal attention still
presents remarkable challenges in training efficiency, including the key-value
tensor communication overhead, IO-awareness attention computation overhead and
load imbalance among GPUs due to the causal attention mask. Most sequence
parallelism approaches for attention are ring-based [169, 173, 175, 176, 178,
179]. Ring Self-Attention [169] leverages sequence parallelism and calculates
the self-attention with ring-style communication to scale up the context
window of LLM training. It first transmits the key tensors among GPUs to
calculate the attention scores in a circular fashion, and then calculates the
self-attention output based on the attention scores and value tensors
transmitted in a similar fashion. DistFlashAttn [176] transmits the key-value
tensor chunks concurrently to utilize IO-awareness FlashAttention [115] kernel
and balance the computation of different GPUs by filling the idle cycles of
earlier tokens with later tokens. Megatron Context Parallel [173] also
leverages the FlashAttention kernel and removes the unnecessary computation
resulted from low-triangle causal masking. It further balances the computation
among GPUs by exchanging half of the chunk with the symmetric GPU.
DistFlashAttn and Context Parallel also overlap the key-value tensor
communication and attention computation with separate CUDA streams. Striped
Attention [178] resolves the imbalance by assigning each GPU a subset of
tokens distributed uniformly throughout the sequence, instead of contiguous
chunks. BurstAttention [179] computes the attention with FlashAttention on
each GPU and utilizes double buffers to overlap the communication and
computation. Blockwise Ring Attention [175] extends Ring Self-Attention [169]
to blockwise attention, which computes the attention in small blocks to reduce
the memory footprint. Inspired from N-body simulation, WallFacer [180] first
divides GPUs into subgroups and replicates query and key-value tensors in each
subgroup via asynchronous AllGather. The attention computation leverages
multiple ring-style P2P communication to enhance efficiency. A final
asynchronous ReduceScatter is needed to distribute the attention output.

DeepSpeed-Ulysses [172] differs from previous ring-based approaches by
splitting the head dimension instead of sequence dimension and leverages All-
to-All to shift the partition dimension from sequence to head. DeepSpeed-
Ulysses can be seamlessly combined with existing attention implementation,
e.g., FlashAttention, and the workload among GPUs is naturally balanced.
However, the parallelism degree of DeepSpeed-Ulysses is restricted by the
number of heads, especially for LLMs using MQA [11] and GQA [12]. LoongTrain
[174] and USP [171] are concurrent work that integrate the advantages of
DeepSpeed-Ulysses and Ring Attention. They organize the GPUs into 2-dimension
mesh, forming hybrid ulysses- and ring-style process groups. During training,
they first perform All-to-All among ulysses groups to switch partition from
sequence to head dimension, and then perform attention computation with Ring-
Attention among ring groups. LoongTrain further proposes Double-Ring-Attention
to fully utilize available bandwidth for inter-node communication and overlap
communication with computation. DSP [177] dynamically switches the parallelism
dimension according to the computation stage in multi-dimensional
transformers, like DiT [246].

####  4.1.5 Expert Parallelism

The Mixture-of-Experts (MoE) is currently the most popular sparse model in
LLMs. While significantly increasing the number of parameters in LLMs, MoE
does not greatly increase the computational cost with conditional
computations[247]. The basic framework of MoE, as shown in Fig. 9, consists of
multiple expert networks that handle different subsets of training data and a
gate network that applies routing algorithm to assign the input tokens to
different experts networks. MoE enables the training of large models with
parameters beyond the trillion scale and is claimed to be used in popular LLM
models such as Mixtral 8x7B [248] and DeepSeek2 [13].

Figure 9: Expert parallelism. The dotted line highlights the MoE components
within the transformer model, where each device maintains one expert for
expert parallelism and collaborate based on All-to-All communication.

Sparse Activation. With the increasing model size, all experts cannot be
accommodated and trained on a single device. Therefore, GShard [15] extends
the idea of MoE to Transformers in distributed settings, where experts are
distributed across different workers and collaborates with All-to-All
communication, as shown in Fig. 9. Subsequent research for expert parallelism
generally follows the same paradigm. For example, the Switch Transformer [181]
incorporates the design of distributed MoE training on the T5 model. But
unlike the top-2 routing algorithm used in GShard, the Switch Transformer
routes each token to only the top-1 expert to maximize computational
efficiency. Additionally, DeepSpeed-MoE [183] proposes a new distributed MoE
architecture that applies shared experts in each worker and places more
experts in deeper layers to balance communication costs with training
accuracy.

Expert parallelism can be effectively integrated with conventional 3D
parallelism. For example, GShard, Switch Transformer, and DeepSpeed-MoE all
treat expert parallelism as an orthogonal dimension of hybrid parallelism. For
efficient hybrid training, DeepSpeed-TED [249] presents a hybrid parallel
algorithm that combines data, tensor, and expert parallelism to enable the
training of MoE models. The authors partition the MoE parameters into “tiles”
of a predefined size to avoid high optimizer memory spikes and propose
communication optimizations like Duplicate Token Dropping (DTD) and activation
checkpointing to eliminate duplicate data in All-to-All communication.
However, it is challenging to choose the optimal hybrid parallelism plan due
to the dynamic nature of MoE, and switching between different parallelism
strategies during runtime also incurs substantial overhead. Therefore, some
research like Tutel [182] designs an adaptive parallelism switching algorithm
that applies the same distribution model layout for all possibly optimal
strategies, and can dynamically switch the parallelism strategy at every
iteration without any extra overhead.

Since General Matrix Multiplications (GeMMs) require the size of all experts’
inputs to be consistent, existing MoE training frameworks often perform token
dropping and padding to match the same expert capacity, which wastes
computation. Megablocks [185] optimizes grouped GeMMs by implementing Block
Sparse Matrix Multiplication and supports different batch sizes for expert
computations in a single kernel to avoid unnecessary token dropping in MoE
training. Another framework that supports grouped GeMMs is ScatterMoE [184],
which implements the ParallelLinear kernel that fuses grouped GeMMs and
scattered read and write operations to reduce the memory footprint for top-k
($k\geq 2$) gating.

Communication Optimization. All-to-all communication in expert parallelism can
seriously affect the training efficiency of MoE, especially in poor network
environments. Existing distributed training systems try to optimize the
performance of MoE by overlapping communication tasks with computing tasks so
that some communication costs can be hidden. For example, Tutel [182] divides
the input tensors into groups along the expert capacity dimension and overlaps
computation and communication among different groups to hide All-to-All
overhead. FasterMoE [191, 192] uses a similar strategy to Tutel but splits the
tensor along the expert dimension. Additionally, Tutel [182] also optimizes
the All-to-All kernel implementation by aggregating small messages into a
single large chunk inside the nodes before exchanging data among different
nodes. This optimization is also used in FasterMoE and ScheMoe [187]. Based on
the overlap strategy in Tutel, PipeMoE [186] models the execution time of both
communication and computation tasks based on the workloads and designs an
adaptive algorithm to find the optimal number of partitions to minimize
training time. ScheMoE [187] considers data compression approaches before All-
to-All communication and modularizes time-consuming operations, including data
compression, collective communication, and expert computation. ScheMoE then
proposes an adaptive optimal scheduling algorithm to pipeline communication
and computing operations to improve training efficiency.

Expert parallelism usually interacts with other parallel strategies in MoE
training. It is possible to reduce communication overhead by fine-grained task
scheduling. For example, Lina [188] analyzes the All-to-All overhead of MoE
during distributed training and inference systematically and finds that All-
to-All latency is prolonged when it overlaps with AllReduce operations. Lina
proposes prioritizing All-to-All over AllReduce to improve its bandwidth and
reduce its blocking period in distributed training. Additionally, Lina
incorporates tensor partitioning and pipelining to perform micro-op scheduling
similar to Tutel. Lina also dynamically schedules resources based on expert
popularity during inference to minimize overhead. Janus [189] designs a data-
centric paradigm that keeps data in place and moves experts between GPUs based
on a Parameter Server. The data-centric paradigm uses fine-grained
asynchronous communication and allows experts to move between GPUs using non-
blocking communication primitives such as pull. Janus implements a topology-
aware strategy to effectively pull experts between nodes and supports expert
prefetching to pull all external experts to local CPU memory.

There are some research optimizes MoE training from model-system co-design
perspective. For example, TA-MoE [190] proposes a topology-aware routing
strategy for large-scale MoE training. TA-MoE abstracts the dispatch problem
into an optimization objective to obtain the target dispatch pattern under
different topologies and designs a topology-aware auxiliary loss based on the
dispatch pattern. This approach adaptively routes the data to fit the
underlying topology without sacrificing model accuracy.

Load Balance. Due to the sparse and conditional computing nature of MoE, a
popular expert may receive more tokens than others in expert parallelism
(usually caused by a poor routing algorithm), leading to serious load
imbalance and affecting the training efficiency of MoE. FasterMoE [192]
proposes the shadowing experts approach, which dynamically broadcasts the
parameters of popular experts to all other GPUs based on the workload of
previous iterations. By spreading the workload of popular experts across
different devices, the shadowing experts approach reduces the impact of skewed
expert popularity. SmartMoE [193] adopts a two-stage approach to search for
the optimal parallel plan for load balance. First, SmartMoE designs a data-
sensitive performance model that divides parallel plans into pools, where the
cost of switching parallel modes within a pool is relatively low. Then,
SmartMoE can switch to the appropriate parallelism (referred to as expert
placement in SmartMoE) to keep load balance during the online stage. FlexMoE
[194] found that the distribution of expert-to-device mapping does not shift
significantly over a short period, so it introduces fine-grained replicated
expert parallelism that duplicates heavy experts across multiple devices.
FlexMoE monitors data workload and uses three placement adjustment primitives
(i.e., expand, shrink, migrate) to generate optimal placement solutions if the
balance ratio is exceeded. Prophet [195] presents a systematic, fine-grained,
and efficient load balancing training method for large-scale MoE models.
Taking the MoE model, device pool, and token distribution as inputs, Prophet’s
planner iteratively searches and evaluates expert placements and finally
outputs a well-balanced expert placement. Additionally, Prophet hides the
overhead of these resource allocation operations using a layer-wise scheduling
strategy.

###  4.2 Auto Parallelism

Given an arbitrary DNN model and a GPU cluster, there exists a vast array of
options for parallelism, encompassing the partitioning of individual layers
and their partitioning degrees. It is a time-consuming and knowledge-intensive
process to design handcrafted hybrid parallelism approaches that can maximize
the training efficiency, requiring expert understanding of the model
architecture, hardware characteristics, and intricate trade-offs involved in
parallelization strategies. Moreover, the efficient implementation of optimal
parallelization strategies often necessitates substantial human efforts. To
address these challenges, auto parallelism emerges as a promising solution,
which seeks to automatically determine the most effective parallelization
strategy for a given DNN model on a specific GPU cluster. By leveraging
sophisticated algorithms and heuristics, auto parallelism systems can analyze
the model architecture, hardware specifications, and performance
characteristics to identify the optimal combination of parallelism techniques,
such as data, tensor, and pipeline parallelism. This approach streamlines the
process of optimizing the distributed training across various models and
infrastructures, enhancing the overall efficiency and reducing the manual
effort. Furthermore, auto parallelism can adapt to changing hardware
configurations and model architectures, automatically adjusting the
parallelization strategy to maintain the optimal performance. In the
following, we categorize existing auto parallelism systems into general and
transformer-specific frameworks, according to the targeted model architecture.

####  4.2.1 General Framework

General auto parallelism frameworks focus on automatically parallelizing
general DNNs on a specific computation cluster. These frameworks typically
follow a three-step process: (1) defining the search space of parallelization
strategies; (2) developing performance models to measure the training
efficiency of different strategies; (3) designing algorithms to efficiently
identify the optimal parallelization strategy. Below we investigate different
approaches according to the search space they cover.

Some works have explored the search space of hybrid data and pipeline
parallelism strategies for DNN training optimization. These approaches focus
on partitioning DNNs automatically and designing pipeline schedules to improve
the pipeline utilization. PipeDream [134] measures the efficiency of pipeline
partitions with the execution time of the slowest stage and develops a dynamic
programming algorithm to partition the DNN evenly by minimizing the slowest
stage. DAPPLE [152] builds an analytical model to estimate the execution time
of one partition strategy and uses dynamic programming to determine the
optimal pipeline partition. AutoPipe [205] constructs a simulator to simulate
the pipeline execution and proposes a heuristic algorithm to obtain the
balanced partition. AutoPipe also automatically splits the micro-batch to
reduce the latency of the warm-up stage. Some device placement approaches
[206, 207, 208] use reinforcement learning to predict the optimal operator
placement for pipeline parallelism.

Researchers also explore the automated data and model parallelism by
partitioning operators along different dimensions. OptCNN [198] partitions
operators along all divisible dimensions in their output tensor and utilizes
an analytical performance model to pick the optimal parallelization strategy,
including the parallelizable dimensions and degree of parallelism, which
defines how to parallelize an individual layer across different devices.
FlexFlow [199] further extends the search space to Sample-Operator-Attribute-
Parameter (SOAP), which includes almost all the divisible dimensions in input
and output tensors, and introduces a novel execution simulator for accurate
performance modeling. FlexFlow efficiently finds an optimal parallelization
strategy with MCMC sampling. Tofu [200] and HyPar [201] develop dynamic
programming algorithms that minimize the total communication cost rather than
the end-to-end performance, to identify the optimal partition for each
operator in the hybrid data and model parallelism space. TensorOpt [204]
optimizes the parallelization strategy under a given memory budget with a
frontier tracking algorithm. AutoMap [202] employs Monte Carlo Tree Search
(MCTS) to select a sequence of partitioning rules defined by PartIR [212] for
a set of selected important operators via a learned scorer. The whole
parallelization strategy is propagated from the strategy via the selected
operators.

Recent works also design approaches for automated data, model and pipeline
parallelism. Piper [209] designs a two-level dynamic programming approach to
find the optimal hybrid data, tensor and pipeline parallelism combined with
activation recomputation. It first divides the model into small partitions for
the pipeline and then splits operators within each partition. Alpa [21]
formulates a comprehensive space by viewing parallelisms as two hierarchical
levels: inter-operator and intra-operator parallelism. Then it automatically
derives an efficient parallel execution plan at each parallelism level. Unity
[210] jointly optimizes the parallelization and algebraic transformations by
representing them as substitutions on a unified parallel computation graph.
Aceso [211] proposes an iterative bottleneck alleviation approach to
significantly reduce the search time. It identifies a performance bottleneck
at every step and adjusts the strategy to mitigate the bottleneck until
convergence. nnScaler [213] introduces three primitives to enable the
composition of the search space with arbitrary partitioning and spatial-
temporal scheduling of the partitioned model. Domain experts can apply
constraints to the primitives to build effective and small search spaces,
which can be automatically explored with low overheads. AutoDDL [215]
customizes a coordinate descent algorithm by iteratively updating the SBP
[214] distributions for each layer and quickly discover an optimal strategy
with near-optimal communication cost.

General auto parallelism frameworks demand efficient system support for
various parallelization strategies, in addition to fast optimization
algorithms for optimal parallelization strategy discovery. This is because
parallelism often involves complex computation and communication operators,
especially for model parallelism that partitions operators. Prior works have
developed efficient systems that enable a wide range of parallelization
strategies, either by building upon modern DL frameworks [21, 213] or
implementation from scratch [199]. Mesh-TensorFlow [196] observes the
intrinsic complexity of implementing a parallelization strategy, and first
proposes to abstract the device cluster into a multi-dimensional mesh, and
abstract parallelism into partitioning the iteration space (i.e. tensor
dimensions). By mapping the tensor and mesh dimensions, a hybrid data and
model parallelism strategy can be easily implemented with high performance.
For example, data and model parallelisms split the batch and hidden
dimensions, respectively. GSPMD [197] further provides a unified way to
achieve various general parallelism schemes with simple tensor sharding
annotations based on JAX [250] and XLA [251]. OneFlow [214] proposes SBP
(Split, Broadcast, Partial-value) abstraction for partition and allows users
to specify the placement and SBP signature for tensors to implement different
parallelization strategies. PartIR [212] decouples the model from its
partitioning and designs a compiler stack for users to compose SPMD sharding
strategies incrementally via a schedule. Similar to TVM [252], Slapo [203]
defines a comprehensive set of schedule primitives for parallelization and
subgraph optimization like operator fusion and activation checkpointing. These
schedules are decoupled from execution and preserves the original model
structure for progressive optimization.

####  4.2.2 Transformer-Specific Framework

As LLMs are based on the transformer architecture, recent works tailor
automated systems for transformers. DeepSpeed-Autotuning [216] automatically
tunes the system knobs to figure out good performance-relevant configurations
in the user-defined tuning space, including the degree of parallelism.
Galvatron [217] designs a dynamic programming algorithm to generate the most
efficient hybrid data, tensor and pipeline parallelism strategy. Merak [218]
introduces an automatic model partitioner for non-intrusive automatic
parallelism and a high-performance 3D parallel runtime engine to enhance the
utilization of available resources. Colossal-AI [219, 253] provides a unified
interface for modular usage of hybrid data, tensor, sequence and pipeline
parallelism. Galvatron-BMW [220] extends the space of Galvatron to include
sharded data parallelism and activation recomputation, and searches for the
optimal strategy considering both the memory consumption and computation while
maximizing the hardware utilization.

###  4.3 Heterogeneous Parallelism

The escalating computational demands of LLM training have spurred advancements
in heterogeneous hardware, which harnesses diverse computing resources and
globally distributed devices. This heterogeneity is also reflected in model
architectures, particularly with Reinforcement Learning from Human Feedback
(RLHF). Utilizing heterogeneous hardware and diverse model architectures has
become essential for the efficient training of LLMs.

####  4.3.1 Heterogeneous Hardware

The massive computational requirements of LLM training have driven the
evolution of accelerators, leading to clusters with mixed device types and
uneven interconnect bandwidths. Additionally, modern data and computing
clusters are often distributed globally due to factors such as power
shortages. These phenomena have motivated the adoption of heterogeneous
parallelisms, which leverage diverse computing resources and geographically
distributed devices to accelerate LLM training.

Some works leverage heterogeneous computing resources, such as CPUs, GPUs, and
specialized accelerators, to enhance the performance of LLMs. The distinct
computation, memory capacity and interconnect bandwidth of these devices
introduce challenges for efficient LLM pretraining. HetPipe [221] partitions
the heterogeneous cluster into multiple virtual works. Each virtual work
processes mini-batches with the pipeline parallelism, and different virtual
works employ asynchronous data parallelisms to improve the throughput. AccPar
[222] proposes flexible tensor partitioning to balance the computation of
different accelerators and uses dynamic programming to automatically decide
tensor partitioning among heterogeneous devices for DNNs. Whale [223] proposes
a unified abstraction to ease the efforts for parallel training of giant
models on heterogeneous clusters. It seamlessly adapts to heterogeneous GPUs
through automatic graph optimizations and balances the workloads with hardware
information. AMP [224] utilizes a heterogeneous-aware performance model to
find the optimal hybrid data, tensor and pipeline parallelism strategy. HPH
[226] arranges different GPUs into stages according to the compute-
communication ratio in descending order and formulates the model partitioning
as an integer programming problem to minimize the iteration time. Pathways
[225] employs a sharded dataflow model and asynchronous gang-scheduling to
efficiently execute ML models on heterogeneous cluster. SDPIPE [227]
introduces a semi-decentralized scheme that decentralizes the communication
model synchronization and centralizes the process of group scheduling for
pipeline parallelism to utilize heterogeneous devices. HAP [228] uses an
A∗-based search algorithm to generate the optimal tensor sharding strategy,
sharding ratio across heterogeneous devices and the communication methods for
distributed training. PipePar [229] proposes a dynamic programming algorithm
to partition the model into stages for pipeline considering both the
heterogeneity of GPUs and network bandwidths.

Some other works explore geo-distributed devices, featuring the low network
bandwidth, to enhance the training efficiency. Yuan et al. [230] partition the
LLMs into computational tasklets and propose a novel scheduling algorithm to
efficiently utilize a group of heterogeneous devices connected by a slow
heterogeneous network for hybrid data and pipeline parallelism. SWARM
parallelism [231] partitions the model into equal-sized stages and prioritizes
routing inputs to stable peers with lower latency for workload balance. It
also adaptively moves devices across stages to maximize the training
throughput. FusionAI [232] splits the training computation graph (DAG) into
subgraphs (sub-DAG) and generates a load balanced task schedule to utilize
heterogeneous consumer GPUs connected with low bandwidth for pipeline
training. Communication compression approaches, like CocktailSGD [254], can
also be leveraged to train LLMs efficiently in low-bandwidth clusters.

####  4.3.2 Heterogeneous Model

Figure 10: An example of RLHF. Inference process: 1The actor model generates a
response from a given query. 2The critic model, reward model, and reference
model use the query and response pairs to generate the value, score, and KL
divergence required for training through inference. Training process: 3The
actor model and critic model use the data collected in the inference process
to update their weights through gradient descent.

During the LLM training process, heterogeneity is not only reflected in the
hardware, but also in the model. Training may involve the interaction of
several different models. A specific example is Reinforcement Learning from
Human Feedback (RLHF). RLHF is a training method that aims to align AI systems
more closely with human preferences [255], leveraging human’s advantages in
judging appropriate behavior rather than demonstrating. This method has
received widespread attention, especially for fine-tuning large language
models. However, due to the particularity of the Proximal Policy Optimization
(PPO)[256] algorithm, the model heterogeneity is introduced into the RLHF
training, making the training process of RLHF very different from pre-training
and supervised fine-tuning.

In principle, RLHF consists of three different stages: the stage 1 is
supervised fine-tuning, the stage 2 is the training of the reward model, and
the stage 3 is PPO training. Model heterogeneity is presented in stage 3, as
shown in Fig. 10. The PPO training stage consists of two different processes,
namely the inference process that generates data, and the training process
that updates the weights of the actor model and critic model. PPO training is
performed via the collaboration of these two processes. Moreover, the training
stage introduces higher memory cost, as we need to serve several copies of
auto-regressive generation models and reward models at the same time, and more
time costs, because we must wait for the experience generation to be completed
before updating the weights.

Many frameworks have been proposed for RLHF training. For instance, DeepSpeed-
Chat [233] uses Hybrid Engine to seamlessly switch model partitioning between
training and inference, such as using tensor parallelism to improve throughput
during inference and using ZeRO [145] or LoRA[257] to improve memory
utilization during training, providing outstanding system efficiency for RLHF
training. HuggingFace TRL [234] can make full use of various parameter-
efficient fine-tuning (PEFT) methods, such as LoRA or QLoRA[258], to save
memory cost, and use a dedicated kernel designed by unsloth[259] to increase
the training speed of RLHF. ColossalAI-Chat [253] is another end-to-end RLHF
training framework that also supports LoRA and supports the use of ZeRO [145]
to reduce memory redundancy.

However, the above work adopts a flattening strategy for model placement, that
is, placing the four models in RLHF on the same device, and then using methods
such as ZeRO or LoRA to minimize memory cost. But using only ZeRO will lead to
memory bottlenecks when training larger models, while using efficient
parameter fine-tuning strategies such as LoRA will damage model performance.
To solve this problem, OpenRLHF [235] uses Ray [260] and vLLM [261] to
distribute the reward models to different devices, avoiding placing all four
models in PPO on the same device. Similarly, Adpative Placement and
Parallelism (APP) framework [236] proposed two other model placement
strategies, namely Interleaving Strategy and Separation Strategy. It captures
the fact that the generation part and the training part can run independently
during PPO training, and some serialization can be eliminated by placing them
on different devices, which introduces additional communication but can
overlap well with computing.

Meanwhile, there are some works that apply the parallel strategies in the
first two stages to the stage 3 of RLHF in a fine-grained scheduling manner.
For example, ReaLHF[237] switches the most suitable parallel mode for
different sub-stages in stage 3 by redistributing parameters, which greatly
increases the optimization space. PUZZLE[238] reschedules the order of task
execution according to the affinity of different stages, so that stages with
better affinity can effectively cover execution and improve training
efficiency.

##  5 Computation Optimizations

Today’s AI accelerators offer unprecedented computational capabilities in
terms of FLOPs. However, effectively utilizing these FLOPs to their full
potential requires sophisticated optimization techniques. This section
introduces systems and techniques of computation optimizations to effectively
utilize GPU FLOPs. We first elaborate operator optimizations including the
core attention operator optimizations and automatic optimizations via
compilers. Remarkable performance for operator and computing graphs is gained
based on exploiting massive parallelism and efficient multi-level memory
access concerning the underlying hardware features. Second, mixed-precision
training is detailed where computations are accelerated benefiting from
reduced precision. 16-Bit floating point mixed training has been the de facto
method in most training systems. Low-bit fixed points as low as 1-bit have
been studied and employed for high training efficiency.

forked edges, for tree= grow=east, reversed=true, anchor=base west, parent
anchor=east, child anchor=west, base=center, font=, rectangle, draw=black,
rounded corners, align=left, text centered, minimum width=4em, edge+=black,
line width=1pt, s sep=3pt, inner xsep=2pt, inner ysep=3pt, line width=0.8pt,
ver/.style=rotate=90, child anchor=north, parent anchor=south, anchor=center,
, where level=1text width=15em,font=,, where level=2text width=15em,font=,, [
Computation Optimizations for LLM Training, ver [ Operator Optimizations,
fill=blue!10 [ Manual Optimizations, fill=yellow!10 [ FlashAttention [115],
FlashAttention-2 [116],
FlashAttention-3 [262] , BPT [263], SWattention [264],
Bikshand et al. [265] , ByteTransformer [266] , leaf, text width= 26em ] ] [
Automatic Optimizations, fill=yellow!10 [ Kernel-level: Halide [267], TVM
[252], Roller [268],
Triton [269], ALCOP [270]
Graph-level: Chimera [271], Welder [272], Slapo [203],
TorchDynamo & TorchInductor [273], JIT-Q [274] , leaf, text width= 26em ] ] ]
[ Mixed-precision Training, fill=blue!10 [ 16-Bit Floating Point,
fill=yellow!10 [ FP16 Mixed-Precision Training [275], Campo [276],
BF16 Mixed-Precision Training [277], THC [278] , leaf, text width= 26em ] ] [
Sub-8-Bit Floating Point, fill=yellow!10 [ Wang et al. [279], Sun et al.
[280], FP8-LM [281] ,
Rouhani et al. [282] , leaf, text width= 26em ] ] [ Low-Bit Fixed Point,
fill=yellow!10 [ INT8: Jetfire [283]
INT4: Xi et al. [284]
1-Bit: BitNet [285], BitNet b1.58 [286] , leaf, text width= 26em ] ] ] ]

Figure 11: Studies on computation optimizations for distributed LLM training.

###  5.1 Operator Optimizations

Operator optimizations can be categorized into manual and automatic
optimizations. Manual optimizations mainly focus on the attention operator,
while automatic optimizations are applied more broadly.

####  5.1.1 Manually Optimized Attention Operator

Attention, as the core of transformer architectures, plays a crucial role in
the training efficiency of LLMs. Given a query $q$, and lists of keys
$k_{1},k_{2},...,k_{n}$ and values $v_{1},v_{2},...,v_{n}$, where
$q,k_{i},v_{i}\in\mathbb{R}^{d}$, the attention is computed as follows,

$s_{i}=dot(q,k_{i}),\enspace
s^{{}^{\prime}}_{i}=softmax(s_{i})=\frac{e^{s_{i}}}{\sum_{j}{e^{s_{j}}}},\enspace
o_{i}=\sum_{i}v_{i}s^{{}^{\prime}}_{i}.$

The self-attention exhibits quadratic time and memory complexity relative to
sequence length. The substantial memory consumption and frequent access to
high-bandwidth memory (HBM) imposed by self-attention constrain both the
performance and the context length of transformer models. Extensive work is
presented to optimize self-attention. We focus on exact attention
optimizations while lossy optimizations, like linear attention, are out of our
scope.

Memory-efficient attention is primarily proposed to mitigate the large memory
cost. Rabe et al. [287] prove that self-attention needs $O(logn)$ memory
complexity instead of $O(n^{2})$. By employing lazy softmax, the division by
$\sum_{j}{e^{s_{j}}}$ in softmax can be delayed to the very end of the
attention operation. Thus the summation could be processed incrementally which
requires just a scalar (i.e. $O(1)$) to maintain the intermediate result but
not change the output. The self-attention requires extra $O(logn)$ memory
complexity to keep the additional index into the list of queries to compute
the results to all queries sequentially.

The FlashAttention series further demonstrate fast and memory-efficient exact
attention with IO-awareness, high parallelism, and balanced workloads on GPU.
In FlashAttention [115], an IO-aware tiling algorithm is proposed to reduce
the number of memory reads/writes between slow HBM and fast on-chip SRAM based
on the online softmax. More specifically, the softmax could be calculated one
block at a time by tracking the normalization statistics including the maximum
score and the sums of exponentiated scores. The tiling algorithm thus fuses
all the computation operation chain in self-attention including matrix
multiply, softmax, matrix multiply, etc, in one cuda kernel for reduced HBM
access. FlashAttention-2 [116] further improves the low-occupancy and
unnecessary shared memory reads/writes in FlashAttention with additional
parallelism in sequence length dimension and improved warp-level scheduling
for data sharing inside a thread block. Besides, the popular training systems
[174] generally employ FlashAttention-2 for high performance. FlashAttention-3
[262] speeds up attention on H100 GPU by excavating the newly presented
hardware capabilities as the former FlashAttention implementations are based
on A100 GPU. An interleaved block-wise GEMM and softmax algorithm is
redesigned based on FlashAttention-2 to hide the non-GEMM operations in
softmax with the asynchronous WGMMA instructions for GEMM. Besides, by
leveraging the asynchrony of the Tensor Cores and Tensor Memory Accelerator
(TMA), overall computation is overlapped with data movement via a warp-
specialized software pipelining scheme. Blockwise Parallel Transformer (BPT)
[263] further reduces the substantial memory requirements by extending the
tiling algorithm in FlashAttention to fuse the feedforward network.

The attention operation is also optimized on various architectures by
leveraging hardware-specific features. For instance, SWattention [264] designs
a two-level blocking attention algorithm to exploit the underlying hardware of
the new Sunway architecture, building upon FlashAttention. Similarly, Bikshand
et al. [265] implement FlashAttention-2 on the H100 GPU using the Cutlass
library. They utilize the TMA and WarpGroup Matrix-Multiply-Accumulate (WGMMA)
instructions to optimize data copying and GEMM operations, respectively.
Additionally, tensor layout transformations and software pipelining of data
copying and computations between the two GEMMs are carefully designed based on
the Cutlass library.

Attention mechanisms are also optimized for variable-length sequences, which
are common in distributed LLM training. These variable-length sequences can
incur significant memory and computation costs if padded to the maximum
length. FlashAttention-2 efficiently handles variable-length inputs by
parallelizing the sequence length dimension inseparably. ByteTransformer [266]
focuses on padding-free transformers for variable-length inputs, maintaining a
position array during computation. This array records the mapping relationship
of valid tokens between the original tensor and the intermediate packed
tensor. The fused Multihead Attention algorithm for long sequences employs
optimized grouped GEMM for unpadded tensors. This optimization reduces the
memory and computation overhead associated with padding, thereby enhancing
performance.

####  5.1.2 Automatic Optimizations via Compilers

DNN compilers play an important role in optimizing key computations in LLM
training. Highly efficient kernels of operators are generated automatically
which mitigates the burden of library-based kernel optimizations on diverse
hardware vendors to a great extent. Operator fusion is performed by analyzing
the computation graphs automatically in the training process.

Efficient Operator Kernel Generation. Halide [267] and TVM [252] generate
high-performance operator implementations automatically, relying on multiple
effective schedule primitives that exploit parallelism and data locality on
various backends. Furthermore, Roller [268] optimizes the cost of searching
for optimal alternatives in the large search space of kernel implementations.
It primarily generates a tile kernel consisting of Load, Store, and Compute
interfaces, following which the complete operator kernel is constructed by a
scale-up-then-scale-out approach. Triton [269] provides a C-based language and
compiler that facilitates expressing and optimizing tile tensor programs for
competitive performance. In particular, effective optimizations such as
hierarchical tiling and shared memory allocation are supported via machine-
dependent compiling passes. ALCOP [270] performs automatic load-compute
pipelining to overlap the high-latency memory access with computations for
operators on GPUs. Multi-stage pipelining is utilized by pipeline buffer
detection as well as sophisticated index analysis and substitution in
complicated loop structures.

Graph-level Optimizations for Operator Fusion. With the disparity of speed of
computing cores and memory bandwidth enlarging, modern DNNs are restricted by
memory access. Data reuse among inter-operators is excavated via operator
fusion using compilers. Plenty of compiler works [288, 289, 290, 291] performs
operator fusion by setting expert rules. Particularly, Chimera [271] works on
optimizing compute-intensive operator chains. The operator chain is firstly
decomposed into a series of computation blocks and the optimal block execution
order is then selected to maximize data reuse according to an analytical
model. In addition, replaceable microkernels are designed to leverage
hardware-specific intra-block optimizations. Welder [272] lowers the computing
graph into a tile-level data-flow graph whose nodes are operator tiles and
edges are marked with the memory level of the tensor data reused by the
connected nodes. Operator fusion combinations that maximize data reuse across
different levels of memory hierarchies are searched at the tile level.

Pytorch2 [273] presents two extensions, i.e. a Python-level JIT compiler
TorchDynamo and the corresponding compiler backend TorchInductor, to enable
more robust graph compilation on various back-ends for remarkable performance
improvement without sacrificing the flexibility of Python. Slapo [203]
proposes a schedule language to decouple model execution from definition.
Declaring a set of schedule primitives, users could convert the model for
high-performance kernels. JIT-Q [274] proposes just-in-time quantization for
weights which enables storing only a high-precision copy of weights during
training and creates low-precision weight copies based on the in-memory ALU
augmentations of the commercial PIM (processing-in-memory) solutions.

###  5.2 Mixed-precision Training

Low-precision training is an effective methodology to reduce the computation,
storage, and communication costs in training large-scale models. Nowadays LLM
training generally leverages FP16 and BF16 data types. In particular, BF16 can
represent the same range of values as that of FP32. BF16 training is utilized
in models such as BLOOM [292] since the loss slowly diverges when the loss
scalar becomes too low in FP16 [293]. However, fast bfloat16 support is only
available on TPUs, or GPUs developed with or after the NVIDIA Ampere series.
Furthermore, mixed-precision training and techniques such as loss scaling are
exploited to ensure numerical stability due to the limited dynamic range
represented by reduced precision. 8-Bit or even lower-bit training is also
becoming the focus of quantitative research.

####  5.2.1 16-Bit Floating Point

Popular training systems often employ FP16/BF16 mixed-precision strategies to
reduce precision during training, as highlighted by works like Megatron-LM
[20] and Colossal-AI [253]. The FP16 mixed-precision training scheme [275]
utilizes the IEEE half-precision format to store weights, activations, and
gradients for forward and backward arithmetic operations. To maintain model
accuracy at reduced precision, a single-precision copy of weights is kept for
accumulation at each optimizer step. Loss scaling is also applied to preserve
the values of small-magnitude gradients. Campo [276] optimizes the casting
cost incurred by conversions between FP32 and FP16 through automatic graph
rewriting. This is crucial since the casting cost can sometimes negate the
performance benefits of low precision. Campo also employs offline-trained
linear regression models to predict casting costs and execution times for FP32
and FP16 operations. BF16 [277] is also widely used in mixed-precision
training across various fields [294, 295]. It has the same representational
range as FP32 and does not require hyperparameter tuning for convergence. In
addition, THC [278] addresses computational overhead in parameter server
architectures by eliminating the need for decompression and compression. THC
enables direct aggregation of compressed gradient values through the Uniform
Homomorphic Compression property, thus enhancing efficiency.

####  5.2.2 Sub-8-Bit Floating Point

With the newly released chips characterized by lower precision data types such
as FP8, mixed-precision training is designed to train with lower precision.
Newly designed data formats combined with the techniques to ensure numerical
stability are primarily leveraged to enable FP8 training for deep learning
neural networks. Wang et al. [279] use a new FP8 floating point format for
numerical representation of data as well as computations. Chunk-based
computations and stochastic rounding are utilized in the floating point
accumulation and weight update process, respectively, to preserve model
accuracy. Sun et al. [280] propose hybrid 8-bit floating point training across
the whole spectrum of deep learning models without accuracy degradation. The
novel hybrid FP8 formats utilize different exponent bits and mantissa bits for
forward and backward propagation, respectively, since forward and backward
passes have different optimal balances between range and precision. Besides,
the techniques such as loss scaling are used to avoid accuracy degradation.
With the maturation of more accelerators with FP8 data types, an FP8 automatic
mixed-precision framework (FP8-LM) [281] for training LLMs based on NVIDIA
H100 GPU [296] is proposed, where 8-bit gradients, optimizer states, and
distributed parallel training are gradually incorporated and FP8 low-bit
parallelism including tensor, pipeline, and sequence parallelism is specified.
Besides, precision decoupling and automatic scaling are designed to solve the
data underflow or overflow issues due to the narrower dynamic range and
reduced precision. FlashAttention-3 also employs block GEMM quantization and
incoherent processing that exploits hardware support for FP8 low-precision on
H100 GPU. Furthermore, Rouhani et al. [282] train LLMs at sub-8-bit weights,
activations, and gradients with minimal accuracy loss by utilizing micro
scaled data formats that associate scaling factors with fine-grained sub-
blocks of a tensor.

####  5.2.3 Low-Bit Fixed Point

Low-bit fixed point training is also studied for LLM training. Jetfire [283]
maintains an INT8 data flow where inputs and outputs are loaded and stored in
INT8 data formats to accelerate both compute-bound linear operators and
memory-bound non-linear operators. In addition, tiling algorithms are utilized
to excavate shared memory data access with a per-block quantization method,
where higher precision computations are performed, i.e. INT32 for WMMA tensor
core operations for linear operators and FP32 for non-linear operations, to
maintain the accuracy of pretrained transformers. Xi et al. [284] propose a
novel INT4 training algorithm for transformer models. In the forward
propagation, the activation matrix is firstly transformed into a block
diagonal Hadamard matrix to alleviate the accuracy degradation caused by
outliers in the activation, and the transformed matrix is then quantized. In
the backward propagation, bit splitting and leverage score sampling are
exploited to choose informative gradients for quantization based on the
structural sparsity of activation gradients.

Recently, low-precision training for LLMs has advanced to using 1-bit
precision. BitNet [285] employs a novel low-bit precision matrix
multiplication within transformer blocks, utilizing weights that are 1-bit and
activations that are 8-bit. The model weights are centralized around zero to
maximize capacity within the limited numerical range, then binarized to +1 or
-1 using the signnum function. To ensure training stability and accuracy, the
gradients, optimizer states, and a high-precision latent weight copy are
maintained for parameter updates. Building on BitNet, BitNet b1.58 [286]
further enhances modeling capability by lowering model weights to ternary
values {-1, 0, 1}. The weight matrix is scaled by its average absolute value,
and each value is rounded to the nearest integer among -1, 0, and +1.

##  6 Memory Optimizations

forked edges, for tree= grow=east, reversed=true, anchor=base west, parent
anchor=east, child anchor=west, base=center, font=, rectangle, draw=black,
rounded corners, align=left, text centered, minimum width=4em, edge+=black,
line width=1pt, s sep=3pt, inner xsep=2pt, inner ysep=3pt, line width=0.8pt,
ver/.style=rotate=90, child anchor=north, parent anchor=south, anchor=center,
, where level=1text width=15em,font=,, where level=2text width=15em,font=,, [
Memory Optimizations for LLM Training, ver [ Activation Recomputation,
fill=blue!10 [ Dynamic Evicting, fill=yellow!10 [ DTR [297], MegTaiChi [298],
Coop [299], leaf, text width=26em ] ] [ Static Evicting, fill=yellow!10 [
Checkmate [300], LoongTrain [174], Yuan et al. [301],
Selective Checkpointing [170], DistFlashAttn [176]
, leaf, text width=26em ] ] ] [ Redundancy Reduction, fill=blue!10 [ Fully
Sharding, fill=yellow!10 [ ZeRO [145], FSDP [146], leaf, text width=26em ] ] [
Partially Sharding, fill=yellow!10 [ ZeRO++ [302], MiCS [147], PaRO [303], RTP
[304],
AMSP [305], leaf, text width=26em ] ] ] [ Defragmentation, fill=blue!10 [
Tensor-based Defragmentation, fill=yellow!10 [ ROAM [306], ZeRO-R [145],
Imanishi et al. [307],
MegTaiChi [298], Coop [299] , leaf, text width= 26em ] ] [ VMM-based
Defragmentation, fill=yellow!10 [ GMLake [308], Expandable Segments [309] ,
leaf, text width= 26em ] ] ] [ Offloading, fill=blue!10 [ CPU Offloading,
fill=yellow!10 [ Static Offloading: L2L [310], ZeRO-Offload [311],
Elixir [312], Yuan et al. [301]
Dynamic Offloading: TSPLIT [313], PatrickStar [314],
Mobius [315], Harmony [316], TMOF [317],
STRONGHOLD [318] , leaf, text width= 26em ] ] [ SSD Offloading, fill=yellow!10
[ ZeRO-Infinity [319], Angel-PTM [320],
Smart-Infinity [321], Fuyou [322], MoESys[323, 324], leaf, text width= 26em ]
] ] ]

Figure 12: Studies on memory optimizations for distributed LLM training.

Memory consumption during the training of LLMs can be categorized into four
key components: model states, activations, temporary buffers, and memory
fragmentation.

  * •

Model States: Model states encompass the memory consumed by the optimizer
states, gradients, and model parameters. In mixed-precision training [275],
model parameters and activations are stored in 16-bit precision. When training
a model with $\Phi$ parameters, $4\Phi$ bytes are needed to store parameters
and gradients. The 32-bit copies of the parameters, momentum, and variance
each require $4\Phi$ bytes, totaling $12\Phi$ bytes. Therefore, the overall
memory requirement for storing model states is $16\Phi$ bytes.

  * •

Activations: Activations refer to the tensors generated during the forward
pass. These tensors are essential for gradient computation during the backward
phase.

  * •

Temporary Buffers: Temporary buffers are used to store intermediate results.
For example, operations such as gradient AllReduce often fuse gradients in a
bucket into a single flattened buffer before applying the operation to enhance
throughput.

  * •

Memory Fragmentation: Memory fragmentation can lead to scenarios where memory
requests fail despite having a large amount of available memory. This occurs
because usable memory can become fragmented, and there is insufficient
contiguous memory to satisfy the memory request [145].

To address memory constraints of LLM training, various memory-efficient
techniques have been proposed. These include activation recomputation
strategies, which trade increased computation for reduced memory usage;
redundancy reduction methods that minimize data duplication across training
processes; defragmentation techniques that optimize memory allocation and
deallocation to reduce fragmentation and improve memory utilization; and swap
and offload approaches that leverage CPU memory and NVMe SSDs to supplement
GPU memory. Figure 12 outlines the taxonomy of these optimizations for memory-
efficient LLM training.

###  6.1 Activation Recomputation

During the backward phase of model training, activations are essential for
computing gradients. As model sizes increase, the memory required to store
these activations during training can exceed GPU memory capacity, thereby
limiting the scale of the models that can be trained. Activation recomputation
[325] provides a solution by strategically discarding certain activations
during the forward pass and recomputing them as needed during the backward
pass. This approach has become a de facto method for reducing memory
consumption in LLM training. The key to effective activation recomputation is
balancing the memory savings against the additional computational overhead.

We categorize these methods into two primary approaches: static evicting and
dynamic evicting. Static evicting methods typically involve the formulation of
evicting strategies tailored to specific model architectures or modules. In
contrast, dynamic evicting methods make decisions in real-time without prior
knowledge of the model. Although static approaches necessitate modifications
for new models, the structure of the majority of LLMs share similar
architectures, enabling the general application of these strategies during LLM
training. Despite their inherent flexibility, dynamic evicting methods have
not been widely adopted in the training of LLMs. Nevertheless, we still
explore some related works in this section for further reference.

####  6.1.1 Static Evicting

Static evicting involves establishing a fixed plan for discarding activations
during the forward pass and later recomputing them during the backward pass.
Checkmate [300] formulates this activation recomputation problem as a mixed
integer linear program to determine the optimal rematerialization plan for
static deep learning models. However, Checkmate struggles to scale to large
models like LLMs due to the vast search space.

Recently, several works have proposed customized activation recomputation
policies tailored for LLM training. Selective-checkpointing [170] selectively
discards the activations of memory-intensive attention modules. FlashAttention
[115] fuses the attention module into a single kernel, and also employs
selective-checkpointing to reduce memory consumption. DistFlashAttn [176]
addresses the high computation overhead in long sequences caused by the
recomputation of attention modules, employing a rematerialization-aware
gradient checkpointing strategy. Specifically, DistFlashAttn places
checkpoints at the output of the FlashAttention kernel instead of at the
Transformer layer boundary, thereby removing recomputation in the attention
module during the backward pass and only requiring storage of its output.
LoongTrain [174] introduces selective-checkpoint++, which further optimizes
the checkpointing process, particularly for training with long sequences, by
adding attention modules to a whitelist. This method saves the attention
output and softmax statistics (softmax_lse). During the forward pass, it saves
the outputs of the modules in the whitelist, and during the backward pass, it
retrieves these stored outputs instead of recomputing them, continuing the
computation graph and thus reducing the need for recomputing attention.

Unlike recent works that predominantly focus on handcrafted checkpointing
policies on attention modules for LLM training, Yuan et al. [301] carefully
measure the minimum computation cost required to reconstruct each activation
tensor during model training. They derive a Pareto frontier of memory and
computation costs by enumerating all possible checkpointing methods. From this
Pareto frontier, they select a solution that optimally balances computation
and memory costs.

####  6.1.2 Dynamic Evicting

Dynamic evicting makes real-time decisions on which activations to discard and
recompute based on the current state of the training process. DTR [297]
proposes a greedy online algorithm to heuristically evict and rematerialize
tensors at runtime for both static and dynamic models. MegTaiChi [298]
introduces a dynamic tensor evicting that leverages the access patterns of
tensors tracked at runtime. Coop [299] proposes to mitigate the memory
fragmentation issue caused by activation recomputing methods due to evicting
tensors without considering their contiguity. Coop employs an efficient
sliding window algorithm to ensure that only contiguous memory blocks are
evicted, thereby minimizing memory fragmentations.

###  6.2 Redundancy Reduction

Traditional data parallel approaches replicate the entire model state across
all GPUs, which leads to substantial redundant memory usage. Redundancy
reduction techniques are proposed to optimize memory usage by eliminating or
reducing memory redundancies on each device. These techniques often seek to
balance memory efficiency with the induced communication overhead, thereby
facilitating training of larger scale or batch size with acceptable costs.

####  6.2.1 Fully Sharding

The Zero Redundancy Optimizer (ZeRO) [145] optimizes memory redundancies by
fully sharding model states across all GPUs through three stages: ZeRO-1,
ZeRO-2, and ZeRO-3. ZeRO-1 globally distributes optimizer states across all
GPUs. During the training, each GPU conducts independent forward and backward
propagation to compute gradients, which are subsequently synchronized across
all GPUs within the data parallel group using an ReduceScatter operation. Each
GPU is responsible for updating specific shard of the model parameters.
Following this, the updated model parameter shards are collected from other
GPUs using an AllGather operation, ensuring that all GPUs have the latest
model parameters. ZeRO-1 reduces the memory consumption of optimizer states
from $12\Phi$ to ${12\Phi}/{N}$, where $N$ is the size of data parallelism.
Building upon ZeRO-1, ZeRO-2 further shards the gradients across all GPUs and
each GPU only updates its parameter shard, reducing the memory required for
holding gradients from $2\Phi$ to ${2\Phi}/{N}$. ZeRO-3 partitions parameters
in addition to optimizer states and gradients. Each GPU only holds a part of
the parameters. When the parameters from remote GPUs are needed for the
upcoming computation, they are collected by an AllGather operation and
discarded afterward. In ZeRO-3, each GPU holds only the weights, gradients,
and optimizer states corresponding to its specific parameter partition,
reducing the overall memory consumption from $16\Phi$ to ${16\Phi}/{N}$. ZeRO
is widely adopted by numerous frameworks, such as DeepSpeed [183], PyTorch-
FSDP [146], and ColossalAI [253].

####  6.2.2 Partially Sharding

ZeRO faces communication challenges since the latency of collective
communication operations increases with the communication scale. There exists
a trade-off between memory utilization and communication cost in distributed
LLM training. Optimizing communication overhead can be achieved by sharding
the model states across smaller groups of GPUs, which are smaller sets of GPUs
within a large GPU cluster. This approach reduces inter-node communications
and communication scales, though it may lead to higher memory usage due to
increased redundancy of model states. The key is to balance the communication
scale with memory utilization [305].

Several approaches building upon the ZeRO framework have been proposed to
address the communication inefficiencies while improving memory utilization.
ZeRO++ [302] partitions all model states globally across all devices following
ZeRO-3 and further introduces a secondary shard of parameters within subgroups
of GPUs. In the forward phase, it collects parameters leveraging the primary
shard across all GPUs and maintains a secondary shard of parameters within
subgroups, typically within the same node. During the backward phase, it
collects parameters from this secondary shard, reducing communication scale
and inter-node communications. Additionally, ZeRO++ uses quantization to
compress parameters and gradients, effectively diminishing communication
volume with a trade-off in accuracy. MiCS [147] and FSDP [146] shards all
model state components within subgroups and replicates them across subgroups,
thereby reducing communication scale and consequently communication latency,
leading to enhanced training performance. AMSP [305] and PaRO [303]
incorporate three flexible sharding strategies, including Full-Replica, Full-
Sharding, and Partial-Sharding, allowing each component within the model
states to independently choose a sharding strategy. AMSP formulates an
optimization problem to find the optimal sharding strategy that minimizes
communication costs under memory constraints. In addition, AMSP proposes a
customized communication and computation overlap strategy, incorporating these
flexible sharding strategies to achieve optimized training efficiency. RTP
(Rotated Tensor Parallelism)[304] seeks to minimize memory duplication by
strategically sharding activations and rotating weights/gradients.

###  6.3 Defragmentation

GPU memory fragmentation refers to the scattered, unusable chunks of GPU