# Performance Modeling and Workload Analysis of Distributed Large Language
Model Training and Inference

Joyjit Kundu†∗, Wenzhe Guo†∗, Ali BanaGozar†∗, Udari De Alwis∗, Sourav
Sengupta∗, Puneet Gupta‡ and
Arindam Mallik∗ ∗Interuniversity Microelectronics Centre (IMEC), Leuven,
Belgium
‡Department of Electrical and Computer Engineering, University of California,
Los Angeles
$\dagger${joyjit.kundu<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Aligning future system design with the ever-increasing compute needs of large
language models (LLMs) is undoubtedly an important problem in today’s world.
Here, we propose a general performance modeling methodology and workload
analysis of distributed LLM training and inference through an analytical
framework that accurately considers compute, memory sub-system, network, and
various parallelization strategies (model parallel, data parallel, pipeline
parallel, and sequence parallel). We validate our performance predictions with
published data from literature and relevant industry vendors (e.g., NVIDIA).
For distributed training, we investigate the memory footprint of LLMs for
different activation re-computation methods, dissect the key factors behind
the massive performance gain from A100 to B200 ($\sim$ 35x speed-up closely
following NVIDIA’s scaling trend), and further run a design space exploration
at different technology nodes (12 nm to 1 nm) to study the impact of logic,
memory, and network scaling on the performance. For inference, we analyze the
compute versus memory boundedness of different operations at a matrix-multiply
level for different GPU systems and further explore the impact of DRAM memory
technology scaling on inference latency. Utilizing our modeling framework, we
reveal the evolution of performance bottlenecks for both LLM training and
inference with technology scaling, thus, providing insights to design future
systems for LLM training and inference.

$\dagger$$\dagger$footnotetext: These authors contributed equally to this work

##  1 Introduction & Background

Transformer architecture has emerged as one of the most widely used neural
network architectures in various artificial intelligence (AI) applications
domain. The whole zoo of large language models (LLMs) (e.g., class of GPT
models and Llamma variants) with their ever-increasing model size are examples
of transformer architecture, playing a dominant role in today’s natural
language processing and image classification[15]. Training a large language
model requires a huge amount of data and compute time, resulting in
significant carbon emissions along with an estimated cost of tens of thousands
to millions of dollars. For instance, training a GPT-3 transformer model costs
around $10M [11, 33]. However, the cost in the long run can be dominated by
inference when the same model is deployed to serve multiple users for a long
period of time. Thus, it is important to understand the impact of the trend in
LLMs and underlying system architecture on the performance per total cost of
operation (TCO) in the context of both large-scale training and inference.
Detailed analysis of the performance per TCO would help us identify the pain
points and invest in required areas or components while designing future
compute systems or models.

###  1.1 Transformers

The decoder-based transformer architecture is quite regular and each layer of
it primarily consists of a multi-head attention block (MHA) and a multi-layer
perceptron block (MLP)[32]. This structural regularity and almost static
nature of the data flow at a higher abstraction level allow analytical
modeling of the performance of LLMs at a data center level[16]. The attention
mechanism is at the heart of the transformer architecture. The prediction
accuracy of LLMs depends on the sequence length or the context of the
model[30]. Unfortunately, the execution time and memory complexity of
attention grows quadratically with sequence length[7]. An important challenge
in the LLM community is scaling the performance of transformer models with
long sequences. There are primarily three aspects that determine the
performance: the number of floating point operations (FLOPs), memory accesses,
and communication over the network. For instance, the FlashAttention [7, 6]
work addresses this problem for LLM training by focusing on the memory access
to and from DRAM at the cost of FLOPs. Similarly, the implementation of KV-
cache is crucial for scaling the performance of inference. However, the
performance bottleneck is not static and often shifts with the evolution of
LLMs, compute architectures, or technology. A recent analysis suggests
communication will have a significant overhead (40-75%) in runtime as models
and hardware evolve [24]. Thus, a generic framework that can expose different
tradeoffs in the performance of LLM training and inference workloads with a
connection to the technology is essential for hardware-software co-design.

###  1.2 Performance Bottlenecks

The primary operations in Transformer architecture can be categorized into
three groups or kernels: tensor contractions (general matrix-matrix multiply –
GEMM or matrix-vector multiply – GEMV), normalization (e.g., softmax, layer-
norm), and element-wise operations (e.g., non-linearities, biases, and
dropout) [11]. Among these, GEMM or GEMV (depending on training or inference)
is the most critical operation that dictates the overall performance of a
transformer. The performance at a GEMM level can be characterized by studying
the balance between pure computation and memory accesses. If the time taken by
an operation is dominated by the count of arithmetic operations not the memory
access time, it is categorized as compute-bound. For memory-bound operations,
the execution time is primarily determined by the memory accesses while the
time spent in actual computation is negligible. The arithmetic intensity is a
metric that shows the compute or memory-boundedness of a kernel by capturing
the number of arithmetic operations per byte transferred to and from the
memory. Tensor contractions in distributed training are generally compute
intensive since they involve fat GEMMs ($m\approx n\approx k>>1$),
characterized by matrices that are either square or closer to square in shape
due to the substantial batch and sequence dimensions; however, in the auto-
regressive generation phase of inference, they are mostly memory-bound since
the inherent sequential nature of token generation leads to skinny GEMMs
-characterized by matrices that are long rectangles in shape or GEMVs. The
other two types of operations are in general memory-bound as well. Kernel-
fusion is one of the techniques commonly used to improve the arithmetic
intensity of such operations [1]. The above issues concern the performance at
a single accelerator level. The other important aspect of LLM training or
inference at scale is the data transfer through the network. The communication
overhead becomes quite important for large-scale training and inference on
advanced multi-GPU systems when compute is relatively fast[24, 1].

###  1.3 Parallelization

Distributed training involves different parallelism strategies: Data
parallelism (DP), Tensor model parallelism (TP), Pipeline parallelism, and
Sequence parallelism (SP)[4, 18, 14]. In DP, each GPU processes a portion of
the data but shares the same model parameters to compute the gradients
locally. The gradients across the devices are then reduced (all-reduce) to
finally update the model parameters synchronously. The memory footprint due to
DP at each GPU is dependent on the minibatch size ($=$ batch size/DP-degree),
sequence length, model dimension, model weights, gradients, and optimizer
states. Tensor model parallelism alleviates the memory requirement of model
related parameters. Essentially, TP partitions a tensor operation across
multiple devices (e.g., the model weight matrix is split across rows or
columns). Depending on the partitioning, each device might end up having the
partial sums due to tensor contraction that need to be further reduced before
going to the next stage of the computation, consequently causing network
overhead. Here, the same data is copied to all participating devices. Using
Megatron-LM’s TP parallelization strategy [28], which is explained in Section
3.2, the communication is of the all-reduce type, and the overhead is
minimized. PP is a type of model parallelism that involves distributing the
layers across multiple devices. Each device processes a set of layers and
passes the activations to the next device. In PP, the minibatch is further
divided into multiple microbatches that are passed in a pipeline fashion. This
approach can introduce idle times, known as pipeline bubbles. This bubble time
can be reduced using techniques like interleaved pipeline schedules as
discussed in [28]. TP distributes the attention and MLP blocks, but not the
Dropout and Layer-norm following them. Despite being computationally
inexpensive, the Dropout and Layer-norm contribute to a considerable amount of
activation memory. Sequence parallelism was proposed to parallelize these
blocks along the sequence dimension to proportionally reduce their memory
footprint without incurring communication overhead. Since training is data
intensive, usually combinations of DP, TP, SP, and PP are used to scale the
performance across several hundreds or thousands of nodes. While for inference
the data is limited and thus, most implementations involve only TP across a
few devices within a node. Of course, batched inference with LLMs may require
more devices to fit the model and data into the device memory.

We consider all the above-discussed issues starting from the characteristics
of the workloads in terms of task graphs, parallelization, mapping of that
onto a system architecture, and modeling the performance of every kernel at
each device and system level to derive key insights on performance. Our main
contributions in this paper are given below:

  * •

We construct a general framework to model and analyze the performance of
distributed LLM training and inference workloads that thoroughly considers the
impacts of compute, memory sub-system, and network communication.

  * •

Upon extensive validation (GEMM, GEMV, training, and inference) and
performance analysis, we bring in the insights behind the massive performance
gain achieved by NVIDIA from A100 to B200.

  * •

We run design space exploration at different technology nodes and for
different DRAM technologies to investigate the importance of compute vs
memory.

  * •

We analyze the impact of off-chip memory technology scaling and network on LLM
training and inference. Also, explore the GEMM level analysis of memory and
compute boundedness for inference.

##  2 Related Work, Trends and Gap areas

Numerous prior works [21, 35, 13, 9, 36, 8] focus on design space exploration
for conventional Deep Neural Networks (DNNs). However, there are relatively
few studies in the LLM space [16, 34, 37, 3, 38]. In [21], the authors propose
an analytical model to assess the impact of different memory hierarchies on
performance and energy consumption for convolution kernels. However, the model
does not support end-to-end DNN performance modeling and focuses only on
optimizing a single layer. Ref. [35], presents a DNN framework that can be
simulated in a cycle-level SoC simulator, which incorporates data
transformation, movement cost, and software framework overheads. In [13], the
authors propose a domain-specific genetic algorithm-based method for efficient
DNN mappings considering a comprehensive map space including computation
order, tile sizes, and parallelization. The mapping proposed in [9], further
improved the DNN computation scheduling problem by mapping the multiple
scheduling problems into a single constrained optimization problem that can be
solved directly without incurring the high cost of iterative scheduling. The
model proposed in [8], utilizes a gradient-decent based algorithm to search
the accelerator-cost function space in order to determine the optimum mapping.

LLMs, nevertheless, have their own unique challenges that are different from
DNNs. To optimize hardware utilization for LLM execution, different mapping
strategies for compute-bound training and memory-bound inference are
essential. AMPeD [16] is an analytical model for end-to-end performance
modeling of distributed transformer training that explores different
parallelization strategies, tunes accelerator and system specifications.
However, it neglects the modeling of the memory subsystem, technology
implications, and network behavior. In [34], the authors emphasize on the
communication network and scale-out modeling for server grade architecture.
However, the authors do not study different mappings, architectural
organizations, and technologies. The analytical framework, in [37], offers a
comprehensive review of efficient LLM inference and introduces a roofline-
based analytical tool for systematic analysis. This framework identifies
hardware bottlenecks and provides insights into memory and compute
requirements for each layer of LLMs. However, it focuses on existing hardware
and does not include customization of components like bandwidth, memory, and
compute. Additionally, visibility into the GEMM operations within a layer is
still a work in progress. DeepFlow as proposed in [3], integrates micro
architecture generation, compute graph transformation, device mapping, and
performance prediction engine. It adopts a gradient descent technique to
explore the design space with the knowledge of hardware specifications. The
model characterizes only the LSTM workloads and validates with the existing
GPUs. In the paper [38], the authors introduces a hardware evaluation
framework for LLM inference workloads. It is fast, accurate, and versatile,
allowing for detailed evaluation of various hardware designs. The model breaks
down GPU/TPU hardware to the systolic array level, offering greater design
flexibility. It includes an automatic mapper for performance-optimal
scheduling and an area-based cost model for architectural decision-making.
However, it primarily focuses on inference and does not address LLM training.

To tackle the challenges of end-to-end performance modeling and workload
analysis of LLM training and inference workloads, we propose a general
analytical methodology or platform, _Optimus_ \dagger\dagger$\dagger$Optimus
is an in-house proprietary platform to perform early design space exploration
of LLM workloads. We conduct performance scaling analysis, explore the impact
of logic and memory scaling on performance, and perform design space
exploration at different semiconductor technology nodes to identify optimized
system architectures.

##  3 Methodology

We build on DeepFlow’s [3] strong foundation and extend the framework
extensively to model state-of-the-art LLM workloads, and architectures like
advanced GPUs, TPUs, and future custom designs, enabling us to investigate the
interplay between software and hardware to expose performance bottlenecks.
DeepFlow is a cross-stack pathfinding framework based on analytical
performance modeling that integrates technology parameters and system-level
architecture with workload characteristics like compute graphs and
parallelization strategies. At the core, it applies a hierarchical roofline
model with a memory-sub system aware tiling to predict the performance of
GEMMs. However, DeepFlow is currently tailored to older machine learning
workloads, such as LSTMs, which can be reduced to a single bulky GEMM
operation. Additionally, the framework’s evaluation and validation have been
conducted solely on old-generation GPU architectures, specifically P4 and
V100. Significant efforts are needed to construct the configuration file for a
new architecture as DeepFlow requires tedious low-level technology parameter
specifications to derive important quantities like area per cell, energy per
flip, etc. This approach makes it difficult to model new-generation commercial
architectures, like NVIDIA GPUs, since the technology details are generally
not revealed.

Figure 1: Overview of our performance modeling framework: $\mu$Arch engine
generates a microarchitecture from the inputs. The architecture abstraction
layer constructs a high-level representation of the underlying architecture.
Given an LLM workload, the framework builds a task graph and parallelizes
across multiple devices based on mapping. The performance prediction engine
predicts the execution time.

###  3.1 Framework overview

Fig. 1 illustrates the high-level overview of our methodology. We start with
the task graph of LLM training or inference and map that onto the system
architecture based on the chosen parallelization strategy and mapping. We
adopt the optimized mapping strategy proposed in the Megatron scheme [28]. At
every device level for GEMMs or GEMVs, we follow a hierarchical Roofline model
based on DeepFlow [3]. Different activation recomputation techniques are also
incorporated to minimize model memory footprint, offering more optimization
options as discussed later.

On the system architecture side, we introduce an intermediate architecture
abstraction layer between the microarchitecture engine and the performance
prediction engine. This abstraction layer extracts the high-level performance
drivers, such as compute throughput and memory sub-system bandwidths, derived
from the underlying microarchitecture engine. It can also directly receive a
high-level system description from external inputs, which avoids tedious
microarchicture parameter calibration and greatly eases the process of
constructing the system configuration description for any new architectures
without compromising prediction accuracy. This enables us to extend the
studied processors to include modern GPUs (e.g., A100, H100, B100, and B200),
TPUs, and custom architectures.

###  3.2 Mapping & Parallelization strategy

As discussed in Section 1.3, two primary parallelization strategies for
scaling DNN training are data parallelism and model parallelism, with model
parallelism further classified into tensor, pipeline and sequence parallelism
[28, 14]. We adopt the model parallelization strategy outlined in the Megatron
paper [28] for training LLMs. This approach focuses on minimizing
synchronization requirements between processing devices by partitioning
matrices effectively, thus reducing the need for communication. We also model
the sequence parallelism proposed in [14] for distributing Dropout and Layer-
norm layers. For simplicity, we focus on explaining the concept in relation to
the MLP block of a transformer layer, but as depicted in Figure 2, it applies
to the MHA block as well.

Similar to any other MLP, a GEMM operation between the input matrix (_I_) and
the weight matrix (_Wi_) is followed by a non-linear function, such as GELU.
To parallelize the GEMM while minimizing synchronization needs, _Wi_ can be
partitioned along its columns. This way, when the input is multiplied by each
partition, the non-linear function can be applied independently without
requiring synchronization. However, the resulting output (_O_), which serves
as the input to the second MLP layer, is also partitioned along its columns.
Consequently, in the second MLP layer, the corresponding weight matrix (_Wo_)
should be divided along its rows in the same proportion as the previous
layer’s output. This ensures that corresponding elements from both matrices
are on the same device and can be multiplied to produce partial results. The
partial results of the second GEMM are then reduced across the processing
devices before being passed to the dropout layer. This approach splits both
GEMMs in the MLP block across processing devices and requires only a single
all-reduce operation in the forward pass per MLP block. It is worth mentioning
that in the MHA block, computation of _Key_ , _Query_ , and _Value_ for
different attention heads are processed in parallel across multiple devices
since they are independent. The rest follows a similar parallelization
strategy as the MLP block. Besides, we also adopt various PP schedules,
including GPipe [10], PipeDream-Flush [17], and Interleaved 1F1B [18].

Figure 2: The model parallelism strategy proposed in Megatron-LM paper [28]
effectively reduces the need for synchronization and communication.

###  3.3 Activation recomputation

Training LLMs demands a large amount of memory. The total required memory
mainly consists of model parameters, optimizer states, and activations. Since
training dictates that the activations in all layers need to be saved for
backward gradient computation, activations become the critical bottleneck for
scaling LLMs. We have implemented two mainstream activation recomputation
techniques to alleviate the storage issue, namely full recomputation [5] and
selective recomputation [14]. Full recomputation checkpoints LLM layers and
recompute all the activations by executing the forward pass again. Despite
saving substantial memory space, it doubles the forward pass time. The
required activation memory size, $A_{\rm full}$, can be expressed as

$A_{\rm full}=N_{\rm ckp}A_{\rm inp}+L/N_{\rm ckp}(A_{\rm tot}-A_{\rm inp})$
(1)

where $N_{\rm ckp}$ is the number of checkpoints, $A_{\rm inp}$ is the input
activation size of a transformer layer, $A_{\rm tot}$ is the total activation
size of a transformer layer, and $L$ is the number of layers.

Selective recomputation selects the memory-intensive parts of LLM layers, such
as softmax and dropout outputs, which are not computationally expensive for
recomputation. The required activation memory size, $A_{\rm sel}$, can be
expressed as,

$A_{\rm sel}=L(A_{\rm tot}-(A_{\rm sm}+A_{\rm do\\_\rm mask}+A_{\rm do\\_\rm
out}))$ (2)

where $A_{\rm sm}$ is the size of the input activation to a Softmax layer,
$A_{\rm do\\_\rm mask}$ is the mask size of the dropout layer after the
Softmax layer, $A_{\rm do\\_\rm out}$ is the size of the output from the
Dropout layer.

###  3.4 Modeling All-to-All communication

Both training and inference use all-gather or all-reduce communication
collectives where all devices within a node or the full system communicate
with each other for gathering or reducing data. The communication involves
fetching the data from the respective device memory and sending it in a
pipeline fashion to hide the latency. We model ring-topology [25] and Double
binary trees topology for global communications [27, 12]. Ring all-reduce is a
bandwidth-optimal algorithm, suitable for data-intensive communication when
the impact of latency is negligible [22]. The communication cost is determined
by the slowest connection between processors, independent of the number of
processors. The algorithm consists of two stages: scatter-reduce and all
gather. All the processors are arranged in a logical ring. In the scatter-
reduce stage, each processor sends a chunk of data to its right neighbor and
reduces the data received from its left neighbor. Each processor ends up with
one chunk of the final result. An all-gather operation is then performed
across all the processors. The data transfer communication time, $T_{r}$, can
be expressed as,

$T_{r}=\frac{2K}{N\times\rm BW}(N-1)+2~{}l~{}(N-1)$ (3)

where $K$ is the data volume to be transferred, $N$ is the number of
processors, $\rm BW$ is the network bandwidth between processors, $l$ is the
network latency. While during training the latency term is negligible, for
inference, its contribution can be non-negligible due to low data volume.
Thus, we also model double binary trees-based communication that is both
bandwidth and latency optimal [27]. The revised communication time is
expressed as [31],

$T_{r}=\frac{2K}{N\times\rm BW}(N-1)+2~{}l~{}\log_{2}(N)$ (4)

The second term in the above reduces the impact of latency and helps scale
inference up to 8 GPUs. It is worth noting that for inference, the data volume
is generally low and the network bandwidth is underutilized. We apply a
utilization factor to derive the actual bandwidth.

###  3.5 KV-cache modeling

Unlike training, GenAI inference processes one token at a time (in the auto-
regressive generation or decoding phase) sequentially depending on previously
generated tokens [26]. This sequential nature complicates the use of
parallelization techniques, such as flash-attention [6], that are effective
during training. As mentioned in Section 1.1, _Key_ s and _Value_ s are used
to calculate the scaled dot-product attention for each _Query_. During the
decoding phase of an LLM, the attention calculation for previous tokens is
repeated at each generation step.

KV-caching optimizes this process by focusing solely on calculating attention
for the new token while caching the previously computed _Key_ s and _Value_ s.
The required total KV-cache size is given by $(2\times{\rm
batch~{}size}\times{\rm context}\times{\rm
precision}\times{\rm\\#layers}\times{\rm embedding~{}dimesnion})$. The first
factor of 2 is due to the _Key_ and _Value_ -matrices. Without KV-cache the
_Key_ s and _Value_ s for all the previous tokens need to be recomputed. The
trade-off for this approach is the increased memory and bandwidth required to
store and load the _Key_ and _Value_ states.

###  3.6 Design space exploration framework

Design Space Exploration (DSE) relies on connecting technology parameters to
an architecture template following a system topology to create a micro-
architecture that serves as a blueprint for the underlying device. With a
given budget and allocation of hardware resources (i.e., area, power, and chip
perimeter) on the micro-architecture, the DSE framework derives the essential
coarse-grained quantities, such as compute throughput, memory capacities,
memory bandwidths, and network bandwidths (both inter-node and intra-node).
The analytical performance model estimates the performance (i.e., execution
time) of a given workload based on these high-level quantities that define the
system. Thus, the DSE framework solves a constrained optimization problem. The
search space contains all possible choices of area, power, and perimeter
fractions for each component in the micro-architecture. The constraint is a
given resource budget. A gradient-descent search algorithm is employed to find
the optimal design point that minimizes the execution time.

##  4 Validation

Extensive validations of the proposed methodology have been conducted through
experimental measurements and comparison with published data across various
platforms. Below we present detailed results for GPUs.

###  4.1 Distributed GEMM and GEMV validation

The primitive component of LLM workloads is the GEMM kernel. Thus, it is
essential to validate its implementation in different scenarios. The training
phase generally involves fat GEMMs. Whereas, skinny GEMMs such as GEMV,
dominate computation in the inference phase due to the autoregressive token
generation phase. This section covers the validation of both types.

Fat GEMM kernels are generally compute-bound kernels due to high arithmetic
intensity– they were intensively validated on the old-generation GPUs in
DeepFlow. We extend the validation studies and verified its prediction
accuracy on advanced GPUs (A100s). On the other hand, GEMV kernels are
typically bounded by the memory bandwidth in GPUs. Since GEMV kernels move
small volumes of data between memory levels, DRAM bandwidth is usually
underutilized. The utilization depends largely on the matrix/vector dimension.
In order to capture the effect of the underutilized memory bandwidth, a memory
bandwidth utilization factor is introduced to adjust the roofline model-based
time calculation. Matrix/vector dimensions were selected to cover a wide range
of kernel types used in the LLMs. The selected GEMV kernels are profiled on
A100 GPUs, their memory utilization in multilevel memory as well as compute
utilization are recorded and clustering techniques are used to quantify the
DRAM utilization factors. This process results in reducing the absolute
percentage error between the measurements and the prediction to 5.4% (depicted
in blue in Fig. 3). We also simplify the modeling process by applying a
constant DRAM utilization factor to all the GEMV kernels, resulting in
negligible errors for large matrices (indicated by orange points). For smaller
sizes, the software overhead has a non-negligible impact. Worth mentioning
that we have extended our modeling framework to accommodate TPUs and custom
architectures.

Figure 3: Correlation between GPU runtime and our prediction for GEMV
validation on a single A100 GPU.

###  4.2 LLM training validation for GPUs

Our framework provides training time predictions for decoder-based transformer
models, such as GPTs. Running large-scale LLM training is not feasible for us,
thus, we have validated our framework and methodology against the published
data [28, 14]. We model the impact of various parallelism strategies,
including TP, DP, PP, and SP. Due to high memory consumption by activations,
different activation recomputation techniques are considered, including full
recomputation and selective recomputation. Table 1 shows the comparisons
between the reported training time per batch for different systems of A100
GPUs (8 to 3072 GPUs) and our predictions for different parallelism settings
across various GPT models (22 B - 1 T parameters models). Among all the cases,
the relative errors are mostly well below 10%. It is worth noting that TP and
SP are always implemented within a node due to their higher communication
overhead. DP and PP are usually implemented across nodes. For PP, the
interleaved scheduling strategy assigns multiple pipeline stages in each
device and allows for a large number of microbatches with a small memory
footprint, which minimizes the pipeline bubbles. Our framework accurately
models every aspect of different parallelisms and can also produce accurate
memory breakdown.

Table 1: Validation of training time per batch for LLMs on systems of A100 GPUs with varying choices of parallelism strategies for different GPT models. Model | # GPUs | Batch size | DP-TP-PP-SP | Activation recomputation |  $t_{ref}$ in s from [28][14] |  $t_{pred}$ in s |  $\delta E$ (%)
---|---|---|---|---|---|---|---
Only TP and PP
GPT-22B | 8 | 4 | 1-8-8-1 | full | 1.4 | 1.4 | 2.1
GPT-175B | 64 | 64 | 1-8-8-1 | full | 18.1 | 16.9 | 6.9
GPT-530B | 280 | 280 | 1-8-35-1 | full | 49.1 | 46.8 | 4.6
GPT-1008B | 512 | 512 | 1-8-64-1 | full | 94.4 | 87.9 | 6.9
TP, PP and SP
GPT-22B | 8 | 4 | 1-8-8-8 | selective | 1.1 | 1.1 | 0.0
GPT-175B | 64 | 64 | 1-8-8-8 | selective | 13.8 | 12.9 | 5.9
GPT-530B | 280 | 280 | 1-8-35-8 | selective | 37.8 | 35.5 | 6.2
GPT-1008B | 512 | 512 | 1-8-64-8 | selective | 71.5 | 69.1 | 3.4
DP, TP and PP
GPT-310B | 1920 | 2160 | 15-8-16-1 | full | 37.6 | 34.1 | 9.5
GPT-530B | 2520 | 2520 | 9-8-35-1 | full | 54.2 | 51.2 | 5.5
GPT-1008B | 3072 | 3072 | 6-8-64-1 | full | 102.4 | 100.7 | 1.6
Table 2: Validation (with [19]) of inference latency on systems of A100 and H100 GPUs with varying degree of TP for three different Llama models. Here, the batch size is set to 1, summarization and generation stages involve 200 tokens. Model | # GPUs | TP | tnvidia in ms (A100) | tpred in ms (A100) |  $\delta E$ (%) | tnvidia in ms (H100) | tpred in ms (H100) |  $\delta E$ (%)
---|---|---|---|---|---|---|---|---
Llama2-70B | 8 | 8 | 4735 | 4284 | 9.5 | 3202 | 3147 | 1.7
Llama2-70B | 4 | 4 | 6403 | 6019 | 6.0 | 4116 | 3986 | 3.2
Llama2-70B | 2 | 2 | 10500 | 10042 | 4.4 | 6267 | 6186 | 1.3
Llama2-13B | 8 | 8 | 1693 | 1514 | 10.6 | 1201 | 1209 | 0.7
Llama2-13B | 4 | 4 | 1894 | 1748 | 7.7 | 1431 | 1258 | 12.1
Llama2-13B | 2 | 2 | 2499 | 2492 | 0.2 | 1717 | 1617 | 5.8
Llama2-13B | 1 | 1 | 3884 | 4263 | 9.7 | 2396 | 2599 | 8.5
Llama2-7B | 8 | 8 | 1187 | 1096 | 7.7 | 828 | 899 | 8.6
Llama2-7B | 4 | 4 | 1280 | 1166 | 8.9 | 924 | 869 | 5.6
Llama2-7B | 2 | 2 | 1544 | 1526 | 1.2 | 1143 | 1016 | 11.1
Llama2-7B | 1 | 1 | 2190 | 2472 | 12.9 | 1440 | 1522 | 5.7

###  4.3 Inference validation for GPUs

Since inference involves skinny GEMM or GEMV, we first validate the model by
running them on A100 GPU as discussed in Section 4.1. We validate the
inference latencies reported by NVIDIA in [19] with the predicted values using
our performance model for Llama-2-7b, Llama2-13b, and Llama2-70b on A100-80GB
and H100-SXM GPU systems (Table 3). For each LLM model and a given GPU type,
we validate the strong scaling from 1 to 8 GPUs. Here, the batch size is set
to 1, and the prefill and generation involve 200 tokens. In all the cases, we
match the actual reported numbers within a relative error of 13%. The scaling
in performance from A100 to H100 is largely due to the change in DRAM
technology from HBM2e (bandwidth of 1.9 TBPs) to HBM3 (bandwidth of 3.35
TBPs). Interesting to note that, inference scales poorly with the number of
GPUs, unlike training due to the sequential nature of token generation in the
prediction phase that suffers from relatively less compute requirement and is
purely memory-bound. The only anomaly we see is in the case of the Llama2-7B
parameter model on 8 H100 GPUs: while scaling from 4 to 8 GPUs, the predicted
inference latency goes up while NVIDIA reports otherwise. Although the gain
reported by NVIDIA is very modest, we see an opposite trend. This is primarily
due to the absence of a rigorous network simulator in our performance model
which could give a more realistic utilization of the network communication
bandwidth based on the data volume.

##  5 Case studies: Training

Different aspects of LLM training workload analysis, such as memory profiling,
performance projection, and bottleneck analysis, can be conducted through the
Optimus framework. In this section, we performed case studies to profile the
LLM training memory footprint, project GPU performance scaling across multiple
GPU generations (i.e., from A100 to the latest B200), and assess the
performance scaling of semiconductor technology node from 12 nm (N12) to 1 nm
(N1). Here, the technology node refers to a specific generation of
manufacturing process technology.

Figure 4: Memory breakdown for training GPT models. The dash line indicates
the NVIDIA A100 memory capacity, 80 GB. For each GPT model, three activation
recomputation methods are compared: no recomputation, selective recomputation,
and full recomputation.

###  5.1 Memory dissection

One of the challenges in training billion to trillion-parameter LLMs is memory
overflow. Data parallelism and model parallelism reduce memory proportionally
but suffer from fundamental scaling limitation. It is important to know if an
LLM can fit in the device memory of a system before performing any other
analysis. The choice of parallelism mapping strategy largely depends on the
resulting model memory footprint. Memory profiling helps identify the
bottleneck, so that we can apply the corresponding memory optimization
techniques. We can also determine the best parallelism mapping or training
settings for an LLM model on a certain hardware system or the suitable
hardware system for desired settings.

As an example, we performed memory profiling on three different GPT models in
three activation recomputation scenarios: no recomputation, selective
recomputation, and full recomputation. The training configurations are
available in TABLE 1. We consider mixed-precision training with 2 bytes. Fig.
4 shows the memory breakdown. We can clearly observe the difference in memory
footprint made by the choice of the recomputation technique. With no
recomputation, an LLM can not generally fit in the device memory unless a very
small batch size or large degrees of parallelism are applied, which largely
degrades the training efficiency. Selective recomputation exhibits a small
difference from the full recomputation and causes very little computational
overhead. The extra memory space can be utilized for further optimizing
training efficiency.

###  5.2 Projected GPU performance scaling

With the dire demand for accelerated processing of AI workloads, NVIDIA GPUs
have experienced many generations of improvement, providing massive compute
and scalability. Using our performance model, we project the training time for
the GPT3-175B model run on various GPU systems, namely A100, H100, H200, and
B200. Table 3 lists the configurations used for the projection. The inter-node
communication in the A100 cluster is through HDR InfiniBand (IB) network (200
GB/s), while the nodes in the more advanced GPU clusters are connected through
the NDR IB network (400 GB/s) or NVLink switch system (NVS). Fig. 5 reveals a
clear trend of performance scaling across these generations. H200-NVS-L and
B200-NVS-L used a larger batch size of 4096 to demonstrate their large DRAM
capacity. Improved over A100, H100 triples the compute throughput and
introduces an FP8 transformer engine for mixed-precision training. H100-NDR
gives rise to around 4x speedup. An additional factor of 2 can be obtained by
using the NVLink Switch system. Due to a larger DRAM capacity, H200 can
accommodate an even larger batch size and hence, further accelerate the
training by 3x. The latest breakthrough from NVIDIA, B200, enables FP4
processing and further boosts the performance by 3x with NDR IB and by 14x
with NVS. B200 also features a larger DRAM capacity and faster memory
bandwidth. Using a larger batch size leads to 12x more acceleration. Our
projections from A100 to H100 and from H100 to B200 are aligned with the
reported data from NVIDIA [2, 20].

The scaling of compute capacity and communication speed from one generation to
another is not at the same rate. Through performance modeling, we found that
the LLM training configuration, such as the degrees of all the parallelisms,
can largely affect the scaling trend since it affects the ratio between
compute time and communication time. However, in general, training
configurations are not disclosed by NVIDIA. Through our performance model, we
can obtain all the details and gain insights into how the significant
performance gain is achieved and how to optimally utilize the computing power
and network throughput of a GPU.

Figure 5: Training performance scaling across multiple GPU generations for
GPT-3 175B. Training times are normalized against that of B200-NVS-L. The A100
cluster is connected through HDR InfiniBand network, while the others are
configured with NDR infiniBand network or NVLink switch system (NVS). L
indicates a larger batch size. Other includes weight update time + pipeline
bubble time.

###  5.3 Technology node scaling

The Optimus framework creates a full-stack platform linking semiconductor
technology specifics with the performance predictions of LLMs. This
facilitates the analysis of how advancements in low-level technology nodes
impact model performance, offering insights into performance bottlenecks and
future trends for better HW-SW codesign. This case study showcases training
performance evolution for an LLM with the advancement in technology
development and projects the performance of future nodes. Here, we consider
the GPT-7B parameter model distributed across 1024 GPUs. The training
configuration is displayed in TABLE 3. Seven logic technology nodes were
explored, ranging from N12 to N1. With technology logic scaling, transistors
effectively get smaller, allowing for more of them to fit onto a chip, thus,
enhancing the compute density. We followed the assumption of iso-performance
scaling between consecutive nodes for an optimistic prediction [29, 3], which
determined the scaling factors of 1.8 and 1.3 for area and power,
respectively. Four generations of HBM technology (HBM2 (1TB/s), HBM2E
(1.9TB/s), HBM3(2.6TB/s), and HBM4 (projected 3.3TB/s)) and three types of
inter-node IB network technology (NDR-x8 (100 GB/s), XDR-x8 (200GB/s) and
GDR-x8 (400GB/s)) were considered. We apply the DSE method to optimize the
architecture at each technology node based on the area and power budget. The
scaling results are shown in Fig. 6.

Figure 6: Scaling technology node for different memory and network types. Table 3: Training configurations of case studies for different GPT models. Here, total #GPUs = DP $\times$ TP $\times$ PP. Model | Batch size | Seq length | Vocab size | DP-TP-SP-PP
---|---|---|---|---
GPT-175B | 1024/4096 | 2048 | 51200 | 128-8-8-8
GPT-7B | 512 | 2048 | 51200 | 64-4-4-4

(a) HBM2-NDR

(b) HBM3-NDR

(c) HBM4-NDR

Figure 7: GEMM time breakdown for a single transformer layer in terms of bound
types for different HBM technologies: (a) HBM2, (b) HBM3 and (c) HBM4. The
GEMM times were extracted from the corresponding technology node scaling
experiments.

The node scaling improves the core throughput, cache capacity, and bandwidth.
The training time decreases drastically initially and then tends to saturate
at very advanced nodes beyond N5. During training phase, transformer layers
are generally compute bound. However, when the compute throughput of the
device improves significantly with the scaling, these layers start becoming
bounded by device memory, like DRAM. At this point, the node scaling hardly
contributes to the performance improvement. Fig. 7 shows the breakdown of
bound types in one attention layer. The impact of memory boundedness becomes
dominant gradually with the scaling. Using more advanced HBM technology, such
as moving from HBM2 to HBM2E, leads to a significant runtime reduction.
Further improvement with HBM3/4 is not observed because the model performance
becomes bounded by the network bandwidth. It is worth noting that DeepFlow
framework predicted that the model performance was bounded by L2 cache instead
of compute or device memory, which is not aligned with actual experiments
[11][7].

Logic node scaling and memory technology advancement reduce device kernel
time. However, in distributed training, the communication overhead between
compute nodes grows substantially at the same time due to parallelism. The
performance scaling is particularly constrained by the slow inter-node
network. Enhancing network bandwidth from 100 GB/s to 400 GB/s markedly
improves training times. Pushing for faster networks has been the focus of
industries. For example, NVIDIA’s NVLink Switch System has transformed inter-
node networking to match intra-node network performance, resulting in
significant performance gains.

Table 4: Identification of performance bottlenecks at different matrix
multiply functions per transformer layer in the summarization phase of
inference. The data are Llama2-13B model, single A100 and H100 systems with
half precision, B = 1, and summarization of 200 tokens.

GEMM-function | A100 | H100
---|---|---
t ($\mu$s) | bound-type | t ($\mu$s) | bound-type
merged-head X.WK/Q/V= K,Q,V | 82 | compute | 32 | memory
single head Q.KT = R | 3 | memory | 2 | memory
single head softmax(R).V = Z | 3 | memory | 2 | memory
Z.W=O | 42 | compute | 17 | memory
O.WMLP1 = O1 | 216 | compute | 81 | memory
O1.WMLP2 = O2 | 109 | compute | 42 | memory

Figure 8: GEMM time breakdown per layer in terms of bound types during the
summarization phase of inference for the batch size of 1 and 16. The
generation phase is completely memory bound. The inset shows the memory
capacities, required memory size for kv-cache and model weights.

##  6 Case studies: Inference

In this section, we primarily investigate the importance of compute versus
memory throughput on LLM inference.

###  6.1 Memory boundedness of inference

The autoregressive generation phase in inference is typically DRAM bandwidth
bound for both A100 and H100 GPUs even when batch size (B) $>1$ (for half-
precision). However, the prefill or summarization phase can be compute bound
depending on the accelerator type, precision, batch size, and the prompt
length. In Table 4, we analyze the performance bottleneck of all the GEMM
kernels in the summarization phase of inference using the Llama2-13B model
with half precision. We identify that on A100, the following GEMMs, single
head QKT = R and softmax(R)V = Z are particularly DRAM memory bandwidth bound
due to the limited reuse of the bytes transferred from the device memory. The
other GEMMs are compute bound due to their shape and dimensions. On H100, all
the GEMMs in both prefill and generation phases are DRAM-bound (see Table 4).
This implies that as the compute scales, performance for inference becomes
completely determined by the memory technology. However, it is worth noting
that at low precision, the memory transaction volume decreases and the compute
throughput goes up, thus impacting the arithmetic intensity of the GEMM– these
options are being exploited in NVIDIA GPUs. In Fig. 8, we show the GEMM time
breakdown for a single layer based on the compute versus memory boundedness
for two different batch sizes, B = 1 and 16. On A100, when B=1, roughly 67% of
the time is spent on compute-bound GEMMs – this percentage grows to 96% when B
= 16. Whereas, for H100, the compute dominated time fraction reduces to zero
when B = 1, but grows to 85% when B = 16. Larger batch sizes, thus, improve
inference throughput but at the cost of latency. However, the growth of
latency with B is rather modest. The inset shows the memory footprint of the
KV-cache and the model weights for the Llama2-13B model when B = 1 and 16.

Figure 9: Impact of memory technology scaling on inference latency as
predicted by the performance model. The data are for the Llama2-13B parameter
model, batch size = 1, prefill and generation of 200 tokens. The on-chip
specifications are same as A100. The horizontal dashed lines correspond to
2xH100-HBM3 and 8xH100-80GB-HBM3 connected via NVLink-Gen4.

###  6.2 Memory technology scaling

Since the performance of inference is primarily bound by off-chip memory, we
investigate the impact of DRAM memory technology scaling on the inference
latency for multi-GPU systems using our performance model. For this case
study, we again consider the Llama2-13B parameter model (kv-cache and weights
fit into GPU’s device memory), batch size = 1, prefill and generation of 200
tokens. We keep the compute fixed at the 7 nm technology node (A100) and vary
the DRAM technology from GDR6 (bandwidth = 600GB/s) to HBM3e (bandwidth =
4.8TB/s). Further, we consider a futuristic memory technology, HBMX with the
peak bandwidth of 6.8 TB/s. For all these cases, the network is considered as
NVLink-Gen3 (NV3). Lastly, we also consider the case of A100-HBMX GPUs
connected by NVLink-Gen4 (NV4) interconnect. The results for 2-GPU and 8-GPU
systems are presented in Fig. 9. We see that up to HBM3, the performance
almost scales linearly with the DRAM bandwidth and it slows down from HBM3 to
HBM3e. Beyond that, no such performance gain in memory time is observed since
the DRAM bandwidth surpasses the last level cache bandwidth and the problem
starts to become L2-bound. Further speed-up can be obtained by scaling the on-
chip memory bandwidth and capacity or the intra-node network. While changing
the network from NVLink-Gen3 (NV3) to NVLink-Gen4 (NV4) a modest gain in
communication $\sim 12\%$ is observed. The horizontal dashed lines correspond
to the latency of 2xH100-HBM3e and 8xH100-HBM3e systems as references. Note
that at HBM3e, H100 system is slightly faster than the projected A100-HBM3e
system– that’s primarily due to the faster on-chip memory (no difference in
DRAM technology) and faster network (NV4). It is also worth noting that
although the compute throughput of H100 (5 nm) for half-precision is 989.4
TFLOPS, more than 3x of A100’s compute throughput, it does not help improve
inference performance further, implying that in principle, one can use an
older technology node with the same memory technology without degrading the
performance at the same precision. Another important aspect while running
inference on a 4 or 8-GPU system, is the network communication overhead. We
identify that for 8 GPUs, communication time is roughly 1.6x of memory time
(for Llama2-13B). Solutions like optical networks can really alleviate this
communication bottleneck[23].

##  7 Conclusion

In this work, we analyze LLM training and inference workloads within an
analytical performance modeling framework. We describe the modeling approach,
extensively validate it with published data, and analyze different test cases.
In particular, we investigate the impact of compute and memory technology
scaling on the performance of LLM training and inference. At a single device
level, the performance model is based on a hierarchical roofline model, and
the communication across multiple devices (or nodes) is derived using the
Megatron-LM-like mapping. This combination leads to remarkably accurate
performance predictions at scale both for training and inference. Further
improvements can be made with a first principle approach of determining
realistic memory bandwidth utilization for memory-bound problems like,
inference and by complementing it with a network simulator to incorporate the
impact of network traffic and congestion. Our future investigation includes
evaluating future compute system architectures and technologies for LLM
training and inference, integrating a cost and an energy model into the
current performance modeling framework, and performing complete performance
per TCO analysis.

## References

  * [1] R. Y. Aminabadi, S. Rajbhandari, A. A. Awan, C. Li, D. Li, E. Zheng, O. Ruwase, S. Smith, M. Zhang, J. Rasley _et al._ , “Deepspeed-inference: enabling efficient inference of transformer models at unprecedented scale,” in _SC22: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 2022, pp. 1–15.
  * [2] M. Andersch, G. Palmer, R. Krashinsky, N. Stam, V. Mehta, G. Brito, and S. Ramaswamy, “Nvidia hopper architecture in-depth,” NVIDIA Technical Blog, March 2022. [Online]. Available: https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/
  * [3] N. Ardalani, S. Pal, and P. Gupta, “Deepflow: A cross-stack pathfinding framework for distributed ai systems,” _ACM Transactions on Design Automation of Electronic Systems_ , vol. 29, no. 2, pp. 1–20, 2024.
  * [4] T. Ben-Nun and T. Hoefler, “Demystifying parallel and distributed deep learning: An in-depth concurrency analysis,” _ACM Computing Surveys (CSUR)_ , vol. 52, no. 4, pp. 1–43, 2019.
  * [5] T. Chen, B. Xu, C. Zhang, and C. Guestrin, “Training deep nets with sublinear memory cost,” _arXiv preprint arXiv:1604.06174_ , 2016.
  * [6] T. Dao, “Flashattention-2: Faster attention with better parallelism and work partitioning,” _arXiv preprint arXiv:2307.08691_ , 2023.
  * [7] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré, “Flashattention: Fast and memory-efficient exact attention with io-awareness,” in _Proceedings of the 35th Neural Information Processing Systems Conference (NeurIPS)_ , 2022.
  * [8] K. Hegde, P.-A. Tsai, S. Huang, V. Chandra, A. Parashar, and C. W. Fletcher, “Mind mappings: enabling efficient algorithm-accelerator mapping space search,” in _Proceedings of the 26th ACM International Conference on Architectural Support for Programming Languages and Operating Systems_ , 2021, pp. 943–958.
  * [9] Q. Huang, M. Kang, G. Dinh, T. Norell, A. Kalaiah, J. Demmel, J. Wawrzynek, and Y. S. Shao, “Cosa: Scheduling by constrained optimization for spatial accelerators,” in _2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA)_. IEEE, 2021, pp. 554–566.
  * [10] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu _et al._ , “Gpipe: Efficient training of giant neural networks using pipeline parallelism,” _Advances in neural information processing systems_ , vol. 32, 2019.
  * [11] A. Ivanov, N. Dryden, T. Ben-Nun, S. Li, and T. Hoefler, “Data movement is all you need: A case study on optimizing transformers,” _Proceedings of Machine Learning and Systems_ , vol. 3, pp. 711–732, 2021.
  * [12] S. Jeaugey, “Massively scale your deep learning training with nccl 2.4,” 2019. [Online]. Available: https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/.
  * [13] S.-C. Kao and T. Krishna, “Gamma: Automating the hw mapping of dnn models on accelerators via genetic algorithm,” in _Proceedings of the 39th International Conference on Computer-Aided Design_ , 2020, pp. 1–9.
  * [14] V. A. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro, “Reducing activation recomputation in large transformer models,” _Proceedings of Machine Learning and Systems_ , vol. 5, 2023.
  * [15] Y. Liu, H. He, T. Han, X. Zhang, M. Liu, J. Tian, Y. Zhang, J. Wang, X. Gao, T. Zhong _et al._ , “Understanding llms: A comprehensive overview from training to inference,” _arXiv preprint arXiv:2401.02038_ , 2024.
  * [16] D. Moolchandani, J. Kundu, F. Ruelens, P. Vrancx, T. Evenblij, and M. Perumkunnil, “Amped: An analytical model for performance in distributed training of transformers,” in _2023 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS)_. IEEE, 2023, pp. 306–315.
  * [17] D. Narayanan, A. Phanishayee, K. Shi, X. Chen, and M. Zaharia, “Memory-efficient pipeline-parallel dnn training,” in _International Conference on Machine Learning_. PMLR, 2021, pp. 7937–7947.
  * [18] D. Narayanan, M. Shoeybi, J. Casper, P. LeGresley, M. Patwary, V. Korthikanti, D. Vainbrand, P. Kashinkunti, J. Bernauer, B. Catanzaro _et al._ , “Efficient large-scale language model training on gpu clusters using megatron-lm,” in _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ , 2021, pp. 1–15.
  * [19] NVIDIA, “Inference performance: Llama-2 results,” 2019. [Online]. Available: https://docs.nvidia.com/nemo-framework/user-guide/latest/llms/llama/performance.html.
  * [20] NVIDIA, “DGX B200,” 2024. [Online]. Available: https://www.nvidia.com/en-us/data-center/dgx-b200/.
  * [21] A. Parashar, P. Raina, Y. S. Shao, Y.-H. Chen, V. A. Ying, A. Mukkara, R. Venkatesan, B. Khailany, S. W. Keckler, and J. Emer, “Timeloop: A systematic approach to dnn accelerator evaluation,” in _2019 IEEE international symposium on performance analysis of systems and software (ISPASS)_. IEEE, 2019, pp. 304–315.
  * [22] P. Patarasuk and X. Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations,” _Journal of Parallel and Distributed Computing_ , vol. 69, no. 2, pp. 117–124, 2009.
  * [23] A. Patel, D. Biswas, J. Kundu, Y. Ban, N. Pantano, A. Mallik, J. Ryckaert, and J. Myers, “Accelerating large language model training with in-package optical links for scale-out systems,” _2024 IEEE Computer Society Annual Symposium on VLSI (ISVLSI)_ , 2024.
  * [24] S. Pati, S. Aga, M. Islam, N. Jayasena, and M. D. Sinclair, “Tale of two cs: Computation vs. communication scaling for future transformers on future hardware,” in _2023 IEEE International Symposium on Workload Characterization (IISWC)_. IEEE, 2023, pp. 140–153.
  * [25] P. Pitch and X. Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations,” _Journal of Parallel and Distributed Computing_ , vol. 69 (2), 2009.
  * [26] R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, J. Heek, K. Xiao, S. Agrawal, and J. Dean, “Efficiently scaling transformer inference,” _Proceedings of Machine Learning and Systems_ , vol. 5, 2023.
  * [27] P. Sanders, J. Speck, and J. L. Träff, “Two-tree algorithms for full bandwidth broadcast, reduction and scan,” _Parallel Computing_ , vol. 35, no. 12, pp. 581–594, 2009.
  * [28] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” _arXiv preprint arXiv:1909.08053_ , 2019.
  * [29] A. Stillmaker and B. Baas, “Scaling equations for the accurate prediction of cmos device performance from 180 nm to 7 nm,” _Integration_ , vol. 58, pp. 74–81, 2017.
  * [30] Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler, “Long range arena: A benchmark for efficient transformers,” _arXiv preprint arXiv:2011.04006_ , 2020.
  * [31] R. Thakur, R. Rabenseifner, and W. Gropp, “Optimization of collective communication operations in mpich,” _The International Journal of High Performance Computing Applications_ , vol. 19, no. 1, pp. 49–66, 2005.
  * [32] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [33] K. Wiggers, “Openai’s massive gpt-3 model is impressive, but size isn’t everything,” 2020. [Online]. Available: https://venturebeat.com/2020/06/01/ ai-machine-learning-openai-gpt-3-size-isnt-everything/.
  * [34] W. Won, T. Heo, S. Rashidi, S. Sridharan, S. Srinivasan, and T. Krishna, “Astra-sim2. 0: Modeling hierarchical networks and disaggregated systems for large-model training at scale,” in _2023 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS)_. IEEE, 2023, pp. 283–294.
  * [35] S. Xi, Y. Yao, K. Bhardwaj, P. Whatmough, G.-Y. Wei, and D. Brooks, “Smaug: End-to-end full-stack simulation infrastructure for deep learning workloads,” _ACM Transactions on Architecture and Code Optimization (TACO)_ , vol. 17, no. 4, pp. 1–26, 2020.
  * [36] X. Yang, M. Gao, Q. Liu, J. Setter, J. Pu, A. Nayak, S. Bell, K. Cao, H. Ha, P. Raina _et al._ , “Interstellar: Using halide’s scheduling language to analyze dnn accelerators,” in _Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems_ , 2020, pp. 369–383.
  * [37] Z. Yuan, Y. Shang, Y. Zhou, Z. Dong, C. Xue, B. Wu, Z. Li, Q. Gu, Y. J. Lee, Y. Yan _et al._ , “Llm inference unveiled: Survey and roofline model insights,” _arXiv preprint arXiv:2402.16363_ , 2024.
  * [38] H. Zhang, A. Ning, R. Prabhakar, and D. Wentzlaff, “A hardware evaluation framework for large language model inference,” _arXiv preprint arXiv:2312.03134_ , 2023.