Efficient Parallelization Layouts for Large-Scale Distributed Model Training
Authors: Johannes Hagemann
Samuel Weinbach
Konstantin Dobler
Maximilian Schall
Gerard de Melo
Venue: COLM (2024)
Keywords: model training efficiency, large models, parallelization, distributed
TL;DR: We conduct a large sweep over LLM training setups and distill the results into recommendations for the most efficient training.

Abstract: Efficiently training large language models requires parallelizing across hundreds of hardware accelerators and invoking various compute and memory optimizations.  
When combined, many of these strategies have complex interactions regarding the final training efficiency. Prior work tackling this problem did not have access to the latest set of optimizations, such as FlashAttention or sequence parallelism. In this work, we conduct a comprehensive ablation study of possible training configurations for large language models. We distill this large study into several key recommendations for the most efficient training.
For instance, we find that using a micro-batch size of 1 usually enables the most efficient training layouts. Larger micro-batch sizes necessitate activation checkpointing or higher degrees of model parallelism and also lead to larger pipeline bubbles.
Our most efficient configurations enable us to achieve state-of-the-art training efficiency results over a range of model sizes, most notably a Model FLOPs utilization of 70.5% when training a LLaMA 13B model.

Published as a conference paper at COLM 2024
Efficient Parallelization Layouts for Large-Scale Distributed Model Training
Johannes Hagemann
Samuel Weinbach
Aleph Alpha / Hasso Plattner Institute SohannesSprimmaintellee.....
Aleph Alpha samuel. www.webacchheelppppppph....om
Konstantin Dobler
Maximilian Schall
Gerard de Melo Hasso Plattner Institute gerard. demelo@hpi de
Hasso Plattner Institute konstantin. dobl coblerrhpi..o fohpp..dd de
Hasso Plattner Institute
maximili lian. schallhhhhhhi
Abstract
Efficiently training large language models requires parallelizing across hun- dreds of hardware accelerators and invoking various compute and memory optimizations. When combined, many of these strategies have complex interactions regarding the final training efficiency. Prior work tackling this problem did not have access to the latest set of optimizations, such as ATTENTTTTNTION or sequence parallelism. In this work, conduct comprehensive ablation study of possible training configurations for large language models. We distill this large study into several key recommen- dations for the most efficient training. For instance, we find that using a micro-batch size of usually enables the most efficient training layouts. Larger micro-batch sizes necessitate activation checkpointing or higher degrees of model parallelism and also lead to larger pipeline bubbles. Our most efficient configurations enable us to achieve tate-of-theeart training efficiency results over a range of model sizes, most notably a Model FLOPs utilization of 70.5% when training a LLAMA 13B model.
1 Introduction
The number of parameters and computational resources spent on training deep neural networks is growing rapidly (Brown et al., 2020; Chowdhery et al., 2022; OpenAI, 2023). The largest models consisting of hundreds of billions of parameters do not even fit onto a single hardware accelerator. Thus, training these models requires various ways of re- ducing the memory requirements, such as ZeRO (Rajbhandari et al., 2020), activation checkpointing (Chen et al., 2016), and 3D-parallel (data, tensor, and pipeline parallel) train- ing (Narayanan al., 2021b). 3D parallelism, in particular, has been demonstrated to be effective for the training of Transformer-based large language models (LLMs) with hundreds of billions of parameters (Narayanan et al., 2021b).
However, training these models efficiently with 3D parallelism requires significant domain expertise and extensive manual effort to determine the ideal configurations. These config- urations not only need to combine data, model, and pipeline parallelism most efficiently, but also consider complex interactions with other memory and compute optimizations. FLASHATTENTION (Dao et al., 2022) in particular has had a notable impact since its release, enabling us to train models at previously impossible degrees of training efficiency. In light of these developments, we conduct systematic study via a large-scale training efficiency sweep of these interactions. We consider up to 256 GPUs and LLAMA models with up to 65 billion parameters.
We expand on previous work in this direction (Narayanan et al., 2021b), but include more complex interactions, such as varying the micro-batch size alongside the 3D-parallel config- uration. We also investigate the impact newer methods, such as FLASHATTENTION (Dao et al., 2022) and sequence parallelism (Korthikanti et al., 2022), finding that these affect

---

Published as a conference paper at COLM 2024
the optimal training configuration considerably. Our paper provides several actionable insights for efficiently training LLMs. In summary, the contributions of our work are as follows:
We conduct large sweep over possible configurations for efficiently training LLMs.
Our work considers more degrees of freedom in the training configurations than pre- vious work (Narayanan et al., 2021b) and incorporates important recent techniques such as FLASHATTENTION and sequence parallelism.
We distill our findings into several, actionable insights that enable a more efficient large-scale training of LLMs.
Background
Training very large models requires the combination of various techniques for paralleliza- tion across devices and other memory and compute optimizations. In the following, we provide an overview of the techniques implemented in our in-house training framework AA-Scaling, which we use to conduct the experiments in this paper. These techniques are also implemented in various other frameworks (Zheng et al., 2022; Li et al., 2022; Shoeybi et al., 2019; Rasley et al., 2020; Xu et 2021).
Data Parallelism. Data parallelism (Valiant, 1990) splits the dataset across GPUs during training. Each GPU holds a full model copy, computing loss and gradients for its data shard in parallel. Gradients then synchronized across devices before weight updates. However, this requires that the model fits entirely within a single GPU's memory.
Tensor Parallelism.
Tensor parallelism splits individual weight matrices across multiple GPUs and computes the matrix multiplication in parallel across them. As each GPU only holds a shard of the full matrix, we can fit larger models into memory. For Transformer mod- els, the self-attention and MLP blocks can be parallelized this way with little communication overhead (Shoeybi et al., 2019).
Pipeline Parallelism. Pipeline parallelism splits the model's 's layers into subsequent stages across GPUs. Activations are transferred between these stages. As each GPU only holds some of the layers of the model, we can again fit larger models into memory. However, it can introduce pipeline bubbles of GPU inactivity due to processing delays. PipeDream (Narayanan et al., 2021a) is scheduling algorithm to reduce these by using micro-batches and scheduling their forward and backward computations across pipeline stages. By interleaving forward and backward passes for each micro-batch, PipeDream reduces memory usage, discarding activations after the specific micro-batch's backward pass.
3D Parallelism. As shown by Megatron-LM (Shoeybi et al., 2019), data, tensor, and pipeline parallelism can be combined, which is also referred to as 3D parallelism. In this paper, we use model parallelism as an umbrella term for both tensor and pipeline parallelism. With an efficient combination of these techniques, we can scale the training of models up to 1 trillion parameters (Narayanan et al., 2021b).
Sequence Parallelism. Sequence parallelism (Korthikanti et al., 2022) builds on tensor par- allelism (Shoeybi et al., 2019) by further parallelizing normalization and dropout operations along the sequence dimension. This reduces activation memory usage, especially for longer sequences. Efficiently implemented, sequence parallelism does not introduce additional communication overhead when used together with tensor parallelism.
Activation Checkpointing. Activation checkpointing (Chen et al., 2016) enables a tradeoff between memory and compute. Rather than storing all activations for gradients, they are recalculated during the backward pass. This enables fitting larger models into memory and can improve training throughput by enabling larger batch sizes (Narayanan et al., 2021b).
2

---

Published as a conference paper at COLM 2024
Model
Seq. Len.
GPUs
TP sizes
PP sizes
MB Sizes
Act. Checkpointing
RMSNorm Kemel
(1, 2)
(1, 2)
(1, 2, 4, 8)
(yes, no)
(yes, no)
138
128
(1, 2, 43
(1, 2,
{1, 4)
(yes, no)
(yes, no]
30B
2k
256
(1, 2, 4)
(1, 2,
(1, 4)
(yes, no)
(yes, no}
30B
8k
128
(2, 4)
(2, 4, 8, 16)
(I, 2, 4)
(yes, no}
(yes, no}
65B
128
(2, 4, B)
(2, 4, B)
(1, 2, 4) (yes, no]
(yes, no}
Table Search space of our training efficiency sweep. We sweep over the Cartesian product of all options given in set notation. In particular, we sweep over different tensor parallelization (TP), pipeline parallelization (PP), and micro-batch (MB) sizes, and also whether activation checkpointing was used. Models with sequence length of 2k use global batch size of 2,048, whereas models with a sequence length of 8k use a global batch size of 512. All runs use FLASHATTENTION-2 For runs using activation checkpointing, the RMSNorm kernel caused an error. Therefore, this combination is omitted.
Fused Kernels. Fusing sequential operations into a single kernel enhances the efficiency of memory-bound computations. By executing multiple operations concurrently within a single kernel, data is loaded only once, minimizing memory accesses and optimizing computational overhead.
Flash Attention. Dao et al. (Dao et al., 2022; Dao, 2023) introduce an IO-aware attention algorithm that builds on kernel fusion. Their method provides speedups compared to conventional implementation by minimizing read/ write operations between the slower high-bandwidth memory and the quicker on-chip SRAM in GPUs. Additionally, selective activation recomputation during the backward pass alleviates the ((n²) memory cost in the sequence length.
Experimental Setup
Our experiments are conducted on up to 32 NVIDIA DGX A100 nodes, each equipped
with eight NVIDIA A100 80GB GPUs, resulting in total of 256 GPUs. The GPUs within each node are interconnected via third-generation NVLink! which provides 600GB/s of bandwidth. Cross-node communication is facilitated by NVIDIA Mellanox 200Gb/s
HDR I connections. We utilize torch. ditsdibutebd package with NCCL for the communication backend.
We chose the LLAMA (Touvron et al., 2023a) model architecture for our experiments, due to its recent popularity. The LLAMA architecture introduces minor improvements over the standard Transformer architecture (Vaswani et al., 2017), which have been incorporated into other models over the past few years. The primary architecture modifications include pre-normalization and RMSNorm (Zhang & Sennrich, 2019), the SwiGLU activation func- tion (Shazeer, 2020), and rotary positional embeddings (Su et al., 2022). Our LLAMA models use 128k token vocabulary. The LLAMA models have sequence length of 2k tokens. How- ever, the growing trend of training LLMs with longer sequences (OpenAI, 2023; Touvron et al., 2023b) led us to assess the training efficiency of our LLAMA models on sequences of up to 8k in length. We use AdamW optimization (Loshchilov & Hutter, 2019) following the training setup of LLAMA (Touvron et al., 2023a). All training runs are conducted with our in-house large-scale training framework A--Scclinng using mixed-precision with bfloat 6. We use ZeRO-1 (Rajbhandari al., 2020) to shard the optimizer states across all data parallel ranks based on the results of previous scaling experiments (Narayanan et al., 2021b).
We aim provide a systematic analysis of different combinations of parallelization strategies and other memory and compute optimizations. To this end, we conducted large-scale training efficiency sweep. We ran this analysis for the following model types: LLAMA 13B (2k
NVLin nvidia.com) wwwcoooomnnnaeee wwiiia..om//eruuss/taaccctttee// http://daaacceeeee//iiee nvlink Iinfiniband: nvidia.com httt:///wm/m././/tttwtot////nnna
3

---

Published as a conference paper at COLM 2024
Inffrrnce Opcimized Implemeriatioon Model FLOP' Unilzottoo
70-
015.11
Thalltrettadd Xenel
FlabAmaakee
CISAA Keael
Model FLOPs Velleorte (9)
00-
0..34
01.2 42
(.444)
30-
10-
13B
153 Ak
808
300 Sk
GGB
Model Type
Figure Comparison of the MFU with different attention layer optimizations. The optimal 3D layout was selected for each respective setting. Each optimal layout is annotated with its (micro-batch size, tensor parallelism size, pipeline parallelism size). The kernel from Megatron-LM failed to operate with an 8k sequence length.
8k sequence length), length),,,)))), LLAMA 30B (2k & 8k sequence length), and LLAMA 65B (2k sequence sequeeeeeeee
length). Depending on the model size and availability of compute, we used 64 to 256 GPUs. Table lists the different configuration options for each of the model types. For our training efficiency sweep, we build the Cartesian product of possible options and benchmark each individual configuration. For each configuration, we train for 10 global steps and measure the Model FLOPS Utilization (MFU) (Chowdhery et al., 2022). We exclude the first step, as its performance is significantly impacted by warm-up phase, and report the mean of the last 9 steps. 3 We chose MFU (Chowdhery et al., 2022) over other metrics such as measured hardware TFLOPS, since the latter are system- and implementattion-eeeedeen A higher MFU results in higher throughput and shorter training duration, displaying its efficiency in evaluating computational performance.
Specifically, we compare different tensor parallelization, pipeline parallelization, and micro- batch sizes, as well as the use of activation checkpointing (yes/no no). Since we operate with fixed number of GPUs and global batch size for each model, the data parallelization size and the number of necessary accumulation steps directly follow from the other specified options and are automatically calculated. For example, using 128 GPUs with tensor parallelization size of 4 and pipeline parallelization size of results
rank 16 data parallelization (with 4 x 2 x 16 128), each with 2 pipeline stages and each pipeline stage sharded across 4 tensor parallel splits. We provide the full results of our training efficiency sweep in Table B.1.
Additionally, we conducted preliminary sweep over different attention kernels (native Torch implementation, Megatron-LM kernel FLASHATTENTION-1.00 and PLASHATTENTION-2). Based on the results, we concluded that FLASHATTENTION-2 is superior and thus always used it for our main sweep.
Efficient LLM Training Analysis
4.1 Fused Kernels and Flash Attention
Our evaluation of FLASHATTENTION expands on the original papers (Dao et al., 2022; Dao, 2023) by testing larger model sizes and up to 256 GPUs, compared to the prior focus models up to 2.7B parameters single node. We benchmark against an optimized baseline, the Megatron-LM softmax attention kernel, and assess an optimized RMSNorm kernel from the FLASHATTENTION repository.
In Figure 1, we present results from our main and preliminary sweeps over attention implementations, detailed in Section 3. We compare FLASHATTENTION FLASHATTEN-
Our detailed MFU calculation is reported in Appendix A.1.
Megatron--.N M softmax attention kernel: from here (Date: 25 July 2024)
4

---

Published as a conference paper at COLM 2024
Influence of Activation Checkpointing on Model FLOPs Utilization
(.
Disabled
Enabled
Model FLOPs Utilization (%)
(4.1.
(1, 1, 9)
(L. 2,
[1. 1.
1..2,
(1.
[2. 2,
40
CUDA Out of Memory
20
10
13B
13B 8k
30B Model Type
30B 8k
65B
Figure 2: Comparing MFU of the optimal 3D layout with and without activation checkpoint- ing. LLAMA 30B with 8k sequence length did not fit into memory without checkpointing. The reported results do not use the RMSNorm kernel. Each optimal layout annotated with its (micro-batch size, tensor parallelism size, pipeline parallelism size).
1rON-1.0.8, the Megatron-LMM kernel, and the standard PyTorch implementation, noting the fused kernel's limit 2,048 tokens sequence length. Given the poor performance of pure PyTorch attention in our LLAMA 13B evaluation, we excluded it for larger models. For LLAMA 65B and 30B with 8k sequence length, we focused solely on FLASHATTENTION.
Unsurprisingly, we find that FLASHÂTTENTION vastly outperforms native PyTorch. How- ever, we also find that FLASHÂTTENTION significantly outperforms the kemnel from Megatron-LM. Between the two different FLASH outperforms FI. ALASHATTENION FLASHATTENTION-11.. already contains many of the optimizations introduced in the FLASHATTENTION-2 paper, which measures 2 improvement (Dao, 2023).
versions, FLASHATTENTION: by 4 to 13 percentage points across model sizes.
It is important to note that FLASHATTENTION' improvements are two-fold: FLASHAT- TENTION's 's improved tiling method for an efficient 10-aware SRAM cache utilization and reduced memory requirements through its activation recomputation approach in the atten- tion block. Notably, all best-performing FLASHATTENTION layouts reported in Figure 1 do not make use of activation checkpointing, thereby also benefiting from FTTSNTTTNNSSONNS own activation recomputation.
We further evaluate the effect of FLASHATTENTION'S optimized RMSNorm kernel in Fig- ure 1. The RMSNorm kernel provides a significant boost in training efficiency, up to 14 percentage points compared to ALASHATTENNION-22 without the RMSNorm kernel. No- tably, with the use of the kernel, can fit the entire LLAMA 13B model into single GPU without model parallelization during training (although we still employ ZeRO-1 and shard the optimizer states). In general, the RMSNorm kernel allows to choose more efficient parallelization layouts due to its memory savings. We do not have results combining ac- tivation checkpointing with the RMSNorm kernel, the combination caused an error in our experiments. We control for this and only consider runs without the RMSNorm kernel whenever necessary for fair comparison.
4.2 Activation Checkpointing
In Figure 2, we report the MFU of the best configurations across model sizes, with activation checkpointing of every layer and without. Overall, not using activation checkpointing while compensating with smaller batch sizes higher model parallelism achieves the best training throughputs. We exclude runs with the RMSNorm kernel for a fair comparison due to compatibility issues with checkpointing.

---

Published as a conference paper at COLM 2024
Model FL.PPs Utiliatiin (35)
idiaable),
Model FLOPs Utiliattio (%)
50
(disahled, 2,
disabled, 2,
10-
(enabled, 1. 2)
62
(enabled.
30
50
20
48
10
45
(mabled,,
Micro-bateh Size
Miero-hatch Size
(a) LLAMA 13B
(b) LLAMA 13B 8k seq len
Model Mode FLOPs Udlization (96)
45
(dieabled, 2,4)
40
(mableel, 2,
(enabled,
50
(diiahhled 2,
13-
Model FLOPs Utllzzation
30
Model Model FLOPs Utitization
48
46
38-
(disabled, 4,
14
35
42-
(enabled, 2.
(mabled, 2, 2)
0-
OOM
40
(suabled, 2,
2 Micro-batch
2
2 Micro-batch Size
Micro-batch Size
(c) LLAMA 30B
(d) LLAMA 30B 8k seq len
(e) LLAMA 65B
Figure 3: MFU of the best best-performing run configurations at different fixed micro-batch sizes, visualized by the (activation checkpointing, tensor parallelism size, pipeline parallelism size) triple. The reported results do not use the RMSNorm kernel.
For the LLAMA 30B with 8k sequence length, activation checkpointing was necessary to fit the model into memory, even with tensor parallelism up and pipeline parallelism up to 16. The tensor parallelism could not be increased because the model has 52 attention heads, not divisible by 8. In Section 4.1, we show that adding the RMSNorm kernel further reduces memory, making activation checkpointing unnecessary for the best throughput.5
It is crucial to underline that efficient performance without activation checkpointing is only feasible due to ATTENTIOO..ON. Without it, models larger than LLAMA 13B required activation checkpointing due to memory constraints despite extensive parallelization.
FLASHATTENTION already uses selective activation checkpointing in the attention block, suggesting a need for more targeted activation checkpointing in large-scale Transformer training. Previous work (Korthikanti et al., 2022) also suggests focusing on selective re- computation. This was addressed by PLASHATENTION-aaaa checkpointing strategy in recent research (Li et al., 2023) and the incorporation into NeMO-Megatron Harper et al.. Future work could benchmark the performance and memory efficiency of large-scale model training against frameworks in combination with other parallelization strategies.
4.3 Micro-batch size
In this section, we evaluate the tradeoff between the micro-batch size, required de- gree of model (tensor or pipeline) parallelism, and activation checkpointing. Previous work (Narayanan et al., 2021b) benchmarked different micro-batch sizes with fixed de- grees of tensor and pipeline parallelism, showing that larger micro-batch sizes yield higher throughput. However, a smaller micro-batch size may enable a different, more efficient
Detailed results reported in Appendix B.5.

---

Published as a conference paper at COLM 2024
Model FLOPs Utllization (96)
65-
LLMA 136 8k seq
LLaMA 30B
60
LLaMA 65B
45
40
36
(2.8)
(2: [4, 2)
(4, (Tensor Parallel Size, Pipeline Paralle) Size)
(4, 4)
71
08,
Figure 4: MFU for various model and pipeline parallel configurations for the LLAMA 13B with 8k sequence length, LLAMA 30B, and LLAMA 65B models. Only runs with micro- batch size of 1, activation checkpointing disabled, PLASHATTENTION-2 and the RMS norm kernel are included; runs that out of memory excluded. The LLAMA 13B and the LLAMA 30B with 8k sequence length models are excluded due to limited model parallel configuration options in our sweep.
parallelization configuration. Also, the degree of tensor and pipeline parallelism are often not fixed in practice.
In Figure 3, present the optimal (activation checkpointing, tensor paralleliss iss size, pipeline parallelism size) configuration for each assessed model type. To ensure fair comparison, runs with the RMSNorm kemnel are excluded due error when coupled with checkpointing. For all model types, micro-batch size of achieves the highest MFU. Generally, smaller micro-batch sizes correlate with better MFU performance. Models with an 8k sequence length did not fit into memory with a micro-batch size over 2.
Thus, a micro-batch size of 1 is beneficial in most scenarios, attributable to three factors:
1. Minimal Degree of Model Parallelization: The most efficient training typically requires the least amount of model (tensor or pipeline) parallelization, which is achieved when the micro-batch size is smallest.
2. Avoiding Activation Checkpointing:
For some models (e.g,, LLAMA 65B), a micro- batch size of was the only configuration allowing training without activation checkpointing. As discussed in the previous section, not using activation check- pointing often allows for the highest throughput configurations. The LLAMA 30B 8k model did not fit into memory without using the RMSNorm kernel.
3. Reduced Pipeline Bubble Time: A smaller micro-batch size reduces the time spent in the pipeline bubble at the beginning and end of batch. We already use the teeter--hann-nive Pipedream 1FIB scheduling method (Narayanan et al., 2021a)
discussed in Section 2.
4.4 Tensor vs. Pipeline Parallelism
Narayanan al. (2021b) ablation studies show that neither tensor nor pipeline parallelism, when used in isolation, can achieve the performance of utilizing both at the same time. Our empirical data, especially from the LLAMA 65B model where higher degrees of parallelism are necessary validate this to some extent even in combination with the newly introduced optimizations, as depicted in Figure 4. Their results suggest that an even distribution between the tensor and pipeline parallelism size is optimal, up until the tensor parallelism size reaches the GPU limit in a single node. contrast, our results favor pipeline parallelism over tensor parallelism. The LLAMA 65B model performed best with (tensor, pipeline) parallelism size of (2,8) compared to an evenly distributed (4, 4). Also, the (8,2) configuration was considerably less efficient. This trend was also observed in the LLAMA 13B with Sk sequence length and LLAMA 30B model, where the configurations with higher pipeline parallel size outperform the configurations with
7

---

Published as a conference paper at COLM 2024
larger tensor parallelism size. Our findings differ from Narayanan et al. (2021b) due to the inclusion of recent optimizations like FlashAttention-22 and the RMS norm kernel, which significantly impact the efficiency of parallelism configurations. The training efficiency measured by Megatron-LM (Shoeybi et al., 2019) was comparable when the tensor and pipeline parallelism sizes were interchanged.
Sequence Parallelism
Model FLOPs Utilization with Sequence Parallelism
(1
Enabled
Duaaledd
Model FLOPs Utilization (%)
60
t.. B.
1,1,4)
(1,4,
[L. 2,
(1, 4,
40
20
13B
13B 8k
30B
30B 8k
65B
Model Type
Figure 5: Comparing MFU of the optimal 3D layout with and without sequence parallelism. The reported results use the RMSNorm kernel. Each optimal layout (micro-batch size, tensor parallelism size, pipeline parallelism size).
layout is annotated with its
Based on the findings from previous sections, we performed an additional efficiency sweep to assess the impact of sequence parallelism. We limited the search space to consistently use PLASHATTENTION-2 and the RMSNorm kernel, while omitting the use of activation checkpointing. Furthermore, we restricted the number of GPUs for each model type due to computational constraints. The full configuration sweep is documented in Table 9.
In Figure 5, we report the MFU of the best configurations across model sizes, both with sequence parallelism enabled or disabled. For the LLAMA 13B and 30B models with 2k sequence length, top configurations do use tensor parallelism, hence showing no effect from sequence parallelism. For the 13B model with an 8k sequence length, employing tensor parallelism size of shows no training efficiency improvement. However, for the 30B with 8k sequence length and the 65B models, when using sequence parallelism. both cases, sequence parallelism enables a lower degree of model parallelization due to reduced memory needs, which leads to higher training efficiency.
see 2-6 percentage point improvements
Therefore, conclude that the of sequence parallelism, paired with other optimiza- tions, only shows a notable difference training efficiency for model sizes exceeding 30B parameters 2k sequence length.
4.6 End-to-End Performance
We benchmark the recommendations distilled from extensive sweeps against other publicly reported results in Table 2. For our runs using the in-house AA-Scaling frame- work, report the MFU of the best configuration for each model size, following our recommendations.
We evaluate against publicly available benchmarks from MosaicMI.5, Megatron- DeepSpeed (Korthikanti et al., 2022), Meta's LLAMA (Touvron et al., 2023a), and
MosaicMM. MPT Training Benchmarks: from here (Date: 25 July 2024)
8

---

Published as a conference paper at COLM 2024
Model
GPUs
Seq. Len.
Batch Size
MFU (t)
AA-Scaling LLAMA 13B (ours)
64
2048
2048
70.5%
MPT 13B
64
2048
2048
52.5%
Megatron-LM 18B*
256
2048
1024
34.2%
AA-Scaling LLAMA 13B (ours)
64
8192
512
62.7%
MPT 13B
8
8192
120
52.8%
AA-Scaling LLAMA 30B (ours)
64
2048
2048
61.9%
MPT 30B
64
2048
3072
52.9%
Megaaron--eeppSeed 22B
8
2048
4
41.5%
Megatron-LM 39B*
512
2048
1536
34.5%
AA-Scaling LLAMA 30B (ours)
64
8192
512
60.2%
MPT 30B
8
8192
168
42.6%
AA-Scaling LLAMA 65B (ours)
64
2048
2048
59.6%
MPT 70B
64
2048
2048
53.3%
LLAMA 65B by Meta
2048
2048
2048
49.4%
Megatron-LM 76B*
1024
2048
1792
34.7%
Table 2:
Best achieved end-to-end training efficiency numbers using our recommendations compared to other public training efficiency numbers. We group across comparable model sizes and sequence length. Batch size refers to the Global Batch size. t: MFU numbers were calculated by us based on published training times, as detailed in Appendix A. We provide the exact configurations of our runs included this table in Appendix B.1.
Megatron-LM (Narayanan et al., 2021b). Other frameworks were excluded from our compar- isons, because they either lacked publicly available training efficiency scores, used entirely different hardware, or trained models with vastly different parameter sizes. To the best of our knowledge, we have gathered the best performing, publicly available training efficiency benchmarks for LLMs with comparable architectures and parameter counts.
In general, our configurations set new benchmarks for all assessed model sizes, with MFU improvements of 6-18 MFU percentage points over the previous state-of-the-art Noticeably our LLAMA 13B model achieves an MFU of 70.5% outperforming the MPT and Megratron- LM models. For the 13B and 30B models with an 8k sequence length, our only point of comparison are the models by MosaicML's framework. Here, we outperform the MPT models by 9-17 percentage points. Within the 65B parameter model range, we outperform MPT-70B, the original LLAMA 65B model by Meta, and the Megatron-LM 76B model with an MFU of 59.6% compared to 53.3%, 49.4%, and 34.7%, respectively.
Note on comparability. Most of the comparisons in Table are not one-to-one, due to further differences such as model architecture, employed global batch size, number of used GPUs, and hardware interconnect. For example, the MPT model family employs additional efficiency optimizations, such as the of ALiBi (Press et al., 2022) instead of RoPE (Su et al., 2022). On the other hand, our models use 128k token vocabulary, which can result in more optimistic MFU estimate compared to smaller vocabulary sizes. The comparison made in Table is not meant to be directly one-to-one but a wholesale evaluation of the achieved end-to-end training efficiency. We hope showcase that careful evaluation of the training layout and optimizations used can enable significant boost to training efficiency.
Conclusion
We conducted an exhaustive search over possible combinations of tensor, pipeline, and data parallelism layouts, fused kernels, FLASHATTENTION, activation checkpointing, micro- batch sizes, and sequence parallelism. Our experiments are designed to reflect real-world scenarios. We expect that the distilled findings of our experiments will help researchers and

---

Published as a conference paper at COLM 2024
practitioners to perform more efficient training 3D parallel training configurations, without having to repeat our extensive analysis:
similar frameworks by using similar
Use a micro-batch size of 1 to enable the least degree of model parallelization, to avoid activation checkpointing, and to reduce pipeline bubbles.
Prefer increasing the degree of tensor and pipeline parallelization over the use of activation checkpointing.
Only scale the micro-batch size when you cannot further reduce the degree of model parallelization.
Use sequence parallelization for models exceeding 30B parameters and 2k sequence length.
We experimentally verify that the efficacy of the ATTENTION--- kernel remains as we scale the model size to tens of billions of parameters and to training on multiple nodes. Additionally, we compared the end-to-end training efficiency of our most efficient configu- rations against several other frameworks. We achieved state-of- state-of-the-art of-the-art efficiency in five out of the five model configurations we evaluated, reaching up to 70.5% MFU.
broader range of computational frameworks and hardware setups. This involves examining the applicability of our findings with different parallelization strategies such as various ZeRO stages and FSDP, as well as their performance on recently introduced hardware such as NVIDIA's H100 GPUs and GH200 superchips. Since those offer improved support for fp8 precision and better data transfer efficiency, it is necessary to re-evaluate and possibly refine training strategies to take full advantage of these advancements.
Future work could expand recommendations
We publish the full data of our sweeps at ggtto..cmm/ohameesh////////sssss officeetppaalllizzatoon--yyo
Limitations
Hardware generalization. We deliberately choose the Llama-style model architecture and Nvidia A100 GPU nodes as our setting: These are widely used, giving our findings large target audience. Further, the currently deployed clusters will remain in use for some time. We also believe that our findings can be extrapolated to clusters with similar accelerators and VRAM sizes (e.g., 80GB devices), such as clusters based on H100s with Infiniband interconnect. However, different hardware characteristics, such the Nvidia GH200 and its coherent memory model, may lead to different conclusions. For example, these characteristics make offloading more beneficial. As we only benchmark on Nvidia A100 80GB GPUs, we do not derive "hardware scaling laws". We recommend running small-scale performance benchmark based on the recommendations presented in this work on different systems.
Applicability to other Model Architectures and Domains. We perform our training analysis using a Transformer language model with the LLAMA architecture. Some of the benchmarked optimizations FLASTATTENTION, while others are tailored to specific architectural design choices, such as RMSNorm. In our experiments, only considered the language modeling task. Applica- tions of the Transformer architecture to other domains, such as vision (Dosovitskiy et al., 2021), might also benefit from our recommendations. However, we do not experimentally verify this.
specific to the general Transformer architecture, such as
Additional methods not considered in this study. We evaluate 3D-parallel training with ZeRO-1. Using different ZeRO stages or FSDP (Zhao et al., 2023) might enable even more efficient configurations due to the saved memory. In similar vein, employing selective activation checkpointing (Korthikanti et al., 2023) to improve the interplay the FLASHAI- TENTIoN's activation checkpointing might enable more efficient configurations.
10

---

Published as a conference paper at COLM 2024
Acknowledgements
Konstantin Dobler thanks the German Federal Ministry for Education and Research (BMBF) through the project iKl-Servicezentrum Berlin Brandenburga (011S22092) and the European Laboratory for Learning and Intelligent Systems (ELLIS) PhD program for support. We further thank the COLM 2024 and WANT@Neurlls 2023 reviewers, who reviewed an earlier version of this paper, for their helpful comments as well as the research team within the lab at Aleph Alpha for insightful discussions and feedback.
References
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vosse Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tiee Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL http:: http:///ooooeeen.../.o/ wwwnnngs..nuriiie neurips cc/ ec/paper/202e) hhsh:114776eddbrrc444444444 (attaaccc2fff444 42f Abstract.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost, 2016.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Isvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Hen- ryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanu- malayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022. URL https: / /ttpxvaaxiivvaa////i//ee///aaava arg/abs/22e/ 2204. 02311.
Tri Dao. F Faster attention with better parallelism and work partitioning, 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention Fast and memory-efficient exact attention with TO-awareness. In Advances in Neural Information Processing Systems, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: wwwpeeewwwweeeeooonnneiiiiiieeie
Eric Harper, Somshubra Majumdar, Oleksii Kuchaiev, Li Jason, Yang Zhang, Evelina Bakh- turina, Vahid Noroozi, Sandeep Subramanian, Koluguri Nithin, Huang Jocelyn, Fei Jia, Jagadeesh Balam, Xuesong Yang, Micha Livne, Yi Dong, Sean Naren, and Boris Ginsburg. NeMo: a toolkit for Conversational AI and Large Language Models. URL https://gi wwt.///IIIA/NN///....
11

---

Published as a conference paper at COLM 2024
Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models, 2022.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5:341-353, 2023,
Dacheng Li, Rulin Shao, Anze Xie, Eric P. Xing, Joseph E. Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. Lightseq: Sequence level parallelism for distributed training long context transformers, 2023.
Shenggui Li, Jiarui Fang, Zhengda Bian, Hongxin Liu, Yuliang Liu, Haichen Huang, Boxiang Wang, and Yang You. Colossal-ai: A unified deep learning system for large-scale parallel training, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th Interna- fional Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https: wwtp:e/eeew.ruee//e/ net/forrum?dd-BegRiCCYY
Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory- efficient pipeline-parallel DNN training. In Marina Meila and Tong Zhang (eds.), Pro- ceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 Proceedings of Machine Learning Research, pp. 7937-7947. PMLR, 2021a.
http:tt:://ccceeiiiiiie mlr. press/v1 wwtpsssss...pl//v/ 39/ 39/narayanan21 a html.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei A. Zaharia. Efficient large-scale language model training on gpu clusters using megatron-Im. SC21: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14, 2021b. URL https://api. semantimiiiiieeeee wes.mnttiicchllrr........... www.cccccccccoooooooo orggCoorpsID: 236635565.
OpenAL. Gpt-4 technical report, 2023.
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openrevieew wttp::/weeeeefooreeeeeee met/FoorumiddR&S@PPOCW
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, Nov 2020. doi: 10.1109/ sc41405.2020.0007 URL http: / //tx./xxx/// doi. org/ 10. 1109/SC44 405 2020. 00024.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Rajesh Gupta, Yan Liu, Jiliang Tang, and B. Aditya Prakash (eds.), KDD 20: The 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Viriual Event, CA, USA, August 23-27, 2020, pp. 3505-3506. ACM, 2020. URL https: //ll, acm. wrg//oi//0....... 1145/ 3394486. 3406703.
Noam Shazeer. Glu variants improve transformer, 2020.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019. URL htttps worg/or//aaao
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2022.
12

---

Published as a conference paper at COLM 2024
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Leslie G. Valiant. A bridging model for parallel computation. Commun, ACM, 3301)1003111,, 1990. ISSN 0001-0782. doi: 10.1145/79173.791811 URL 79181.
11455799779
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Aduances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6008, 2017. URL /ttttps:/prrccedinnng neur iprriss.oo cc/pccce/pp////////ppee 7/hash/ 3f5ee243547deee fbdoddc fod053c1c44845ar c4a44aaa Abstract html.
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. Gspmd: General and scalable parallelization for ml computation graphs, 2021.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 12360-12371, 2019. URL https: httpps:/ htwp:ooeeediinsnnnriiissss... cc/ pappp/ html.
Yanli Zhao, Andrew Gu, Rohan Varma, Liangchen Luo, Chien chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, and Shen Li. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proc. VLDB Endow. (838848-860, 2023. URL https://api. semanti cscholar
semanticscholarr wwwchlaa...ggg/ooousS...... ptgp/CorpuSDD:2282777111
Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P. Xing, Joseph E. Gonzalez, and Ion Stoica. Alpa: Automating inter- and intra-operator parallelism for distributed deep learning, 2022.
13

---

Published as a conference paper at COLM 2024
A Measuring Model Training Efficiency Calculations
A.1 Model FLOPs Utilization (MFU)
The calculation of Model FLOPs Utilization (MFU) follows that of PaLM (Chowdhery et al., 2022). We consider the theoretical peak matrix multiplication throughput of P FLOP's per second (e.g., A100 GPUs with 312 peak matmul TFLOPs). Then, the model FLOPs utilization the ratio of the achieved throughput in tokens per second the theoretical peak throughput R = P/ (6N + 12LHQT), where L is the number of layers, H the number of attention heads, Q the attention head size, and T the sequence length. Note that Q H equal to the hidden layer size of the model.
HardwareType
A100 "a108"
H108 "h108"
RTX3090 "rtx3090"
Sproperty
def (sel:))) *ae
Mappings for Maximun Throughout numbers of each GPU. only includes FP16 for now.
pax__tflo__neeee mapppngnnngggg ("alde": 312e12, "h1de": 989.4e12, "rix3098": 58e value]
return nax_
def wtw.lll.ooo.uiilzzatiooo
iter_tine_s: float.
paraneter_count: int,
topology: Any,
architecture: Any,
hardware: HardwareType HardwareType.A110
tokens_per_second
topology. config-global,ee sithsi__atcch_szee architecture. wttp:/uereeeeeeeeeeeeeee.g
/ itte_ttine. tine__
theoretical_ pat.ulaaaakll_eec_ccppaea_ throughput
hardware. nax_tflops topology.com config world_size
attention_flops
12
urchitecture layerayers
architecture hidden_size
pequence_lenngte
nodel_flops parameter_count attention_flon
theoretical_ppake throughput theoretical_peat peak_____a_attttiuu throughhut nodel_flopt
nodel_flops_ut flops_utilization tokens_per_secome theoretical peak, throughput
return modil__llopsuutiizzatie
A.2
LLAMA MFU
From LLAMA paper (Touvron et al., 2023a):
"When training 65B-parameter model, our code processes around 380 tokens/sed GPU on 2048 A100 GPU with 80GB RAM. This means that training over dataset containing 1.4T tokens takes approximately 21 days." "
/GP/
LLAMA 65B model FLOPs utilization (MFU):
tokens_per__ee ceres__per_second 388 2048
theoretical, www.ppaaa..aa fetaklnttaallaa..eeeaeaaalll cestuul_throughen 312e12 2848
attention_flops
12
30
8192
modal__llae _flops attention_
14

---

Published as a conference paper at COLM 2024
theoretical, peak, throughput theoretical peak_ppppaaanaltlleaeaaaa throughhtttt flops model, Soatll_flopsuutiizaaae = tokens_perr peccsdd theoretical_pen throughutt
LLAMA 65B MFU 49.46%
A.3
Megatron-LM MFU
Based on Megatron-LM's (Narayanan et al., 2021b) provided formula, the end-to-end training time is given by ~ STP Rx where 7 represents the number of tokens, P the number of parameters, л the number GPUs, and X the achieved TFLOPs per GPU. Since they provide the achieved TFLOPs per GPU, we can determine the step time through the formula. Subsequently, we can compute the MFU using the specified architecture configuration of the model.
Megatron- 18B:
Step time- 1124-2048-18.4
8.93s
256 135-1012
Lekens_peer_eeeee (1024 2848) 8.93
theoretical_peat pat___peeanatttuueee aaataallntttuaeupaaaaeeaaaaaeeee throughput
312e12 256
iiteentio_ffopp
48
6144
2048
nodel_flops 48. attention_fleps
Sheoretical_peak,teee throukhtt theoretical thetnul_throoggput.. / nodel_flops
moddl_flops_utiliaate tokens_per_secae cheoretical Be__eeakktttthheeeee
Megatron-LM 18B MFU 34.24%
Megatron-LM 39B:
Step time 8:1536-2048 39.1-10"
13.92s
512-138-1012
tokens_per Bertens_perreeee _er_ssccod____dd (1526 2848) 13.92
theoretical, Bacaallppaa___aaae
312e12 512
sttentiion_lone 12
8192
flops 1e9 AicentiionFFops
theoretical_peat teaas__eetttthuhhpuuee fttttiitttttiitiioon tokens_per per_second Sepoorttcal__eaa_thheee
httpsco__eeek_tttl_teeeeeeee nodel_flops
Megatron-LM 39B MFU = 34.56%
Megatron-LM 76B:
Step time 8-17 81192-2008-76.11100
15.59s
1024-140-10/
tokens_per_secomt (1792 2848) 15.59
theoretical_ pec__aa__pp___atnul_____ tattkknntttt_throuughorr 312e12 1824
attention_flops 12
18249
2048
model__loppe 6 76.1e8 attention_ flops
theoretical_peale thauouuuhuut theoretical.peak. www....nttul_tthouuugeeoo nodel_flops model_ fopeuuloooouuttiiatioon tokens_per. second theoretical_peak. httpttcaa__pea_thhhgggge
Megatron-LM 76B MFU 34.76%
15

---

Published as a conference paper at COLM 2024
B
Training Efficiency Sweeps
B.1 End-to-End Best-Performing Runs
Model
GPUs
Step Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
AA-Scaling LLAMA 13B
64
3.32
70.57
1
1
1
False
AA-Scaling LLAMA 13B Sk seq len
34.84
62.78
1
2
True
AA-Scaling LLAMA 30B
72.00
61.98
False
AA-Scaling LLAMA 30B Sk seq len
64
84.37
60.22
1
4
2
True
AA-Scaling
147.02
59.62
True
Table 3: Configurations for the end-to-end performance results of Table 2. All best- performing layouts use FLASHATTENTION-2, the RMS norm kernel, and do not make use of activation checkpointing.
B.2
LLAMA 13B
Step Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
26.54
20.57
disabled
flash_attn2 RMS kern.
1
29.70
63.05
disabled
flash_att22 RMS kern.
31.06
60.26
disabled
flash_attm)) RMS kern.
2
31.29
59.82
disabled
flash_att22 RMS kern.
2
1
33.11
56.55
disabled
flash_atin2 RMS
2
33.58
55.71
disabled
flush_attrl.0.5
2
33.70
55.53
disabled
flash_atth2
2
33.90
55.19
disabled
flash_att22
34.84
53.69
disabled
flash_attn2 RMS kern.
2
35.05
53.37
disabled
flash_attn2 RMS kem.
2
35.54
52,64
disabled
flash_attn2
36.66
51.4
every_laye
flash_atin2
36.69
51.02
every layer
flash_att12
49.89
every layer
flash
37.57
49.80
disabled
flash_tt122
RMS kern.
2
37.71
49.59
disabled
flash_attn1.0.8
2
38.27
48.86
disabled
flash_attn2
2
39.64
47.19
disabled
flash_attn2
2
46.97
disabled
flash_
40.25
46.46
disabled
flash_attn1.0.5
2
40.56
46.11
disabled
flash_attn..0.8
41.81
44.74
disabled
flash_attn1.0.8
42.30
44.21
every layer
flash_attn1.0.3
2
42.36
44.14
every layer
attn110
2
42.54
43.96
every layer
flash_attal.0..
2
every every_layee layer
flash, att.l..8
43.24
43.26
every layer
flash_att22 attn2
43.36
43.13
disabled
fused
43.89
42,61
every every_layer layer
flash_attn..0.8
44.10
42.40
every_ layer
flash_att22
44.13
42.39
every every_layer layer
flash atin2
2
44.35
42.18
every layer
attnl.0
every layer
flash
44.72
41.82
every layer
aash_attm1.0.8
2
45.96
40.69
disabled
fused
46.86
39.90
every every_layer layer
fused
47.12
39.72
every layer
flssh_ttn22
47.09
39.71
every layer
attn..0
47.17
39.64 every layer
fused
39.51
every layer
47.50
39.36
every Jayer
flash_attn1.00
47 82
39.11
every every_layer
fused
47.95
39.00
every every_layer
fused
2
48.93
38.23
every layer
fused
37.5
disabled
torch
2
49.39
37 87
every layer
fused
49.77
37.57
every layer
a...
2
50.05
37.37
every_ layer
flash_attn1.03
1
50.16
37.28
every layer
aats1
2
51.44
36.41
every every_layer
flash
2
2
51.79
36.11
every layer
fused
2
1
52.55
35.59
disabled
torch
16

---

Published as a conference paper at COLM 2024
Step Time
MFU
Activation
Kernel
MB Size
TP size
PP Size
53.50
34.95
every. leery_lyerr
flash 1asshatttn1.0.8
2
53.57
31.94
every layer
flash_att22
2
2
every. layer
54.35
34.40 every layer
fused
2
1
54.56
34.27
every_ layer
fused
2
55.13
33.95
every layer
fused
2
35.35
33.78
every layer
torch
2
35.61
33.64
every every_layer layer
_nsh_attn2
2
36.21
33.41 every Jayer
torch
2
every layer
36.61
33.06
every_ layer
flash_att22
2
57.07
32.78
every layer
torch
57.25
32,66
every_layer
fused
58.92
31.77
every layer
fused
2
2
60.24
31.04
every every_aayer
torch
2
60.75
30.78
every layer
torch
29.87
every layer
torch
62.69
29.83
every_ layer
torch
2
63.28
24.58
every layer
torch
2
65.25
28.66
every every_layer
torch
2
65.47
28.57
every layer
flash_att22
2
72.82
25.69
every_ layer
flash_atin2
2
OOM Error
disabled
flash_att22 RMS kern.
Error
disabled
flash aath2 RMS kern.
2
OOM Error
disabled
flash_att22 RMS kern.
2
OOM Error
disabled
flash_att22 RMS kern.
OOM Frror
disabled
flash_ RMS kern.
OOM Error
disabled
flash_att22 RMS kern.
OOM Error
disabled
flash_atin2 RMS
OOM Errur
disabled
flash_att22
Erroz
disabled
flash
OOM Error
disabled
flash_att22
OOM Error
disabled
flash_attri
OOM Error
disabled
flash_ aash_attm2
OOM Error
disabled
flash_atto
OOM Error
disabled
flash attn2
OOM Errur
disabled
flash_att22
Error
disabled
flash attn2
2
OOM Errur
layer
turch
2
OOM Error
disabled
torch
2
OOM Error
every_ layer
fused
2
OOM Error
disabled
fused
2
OOM Errur
disabled
flash attr1.0.8
OOM Error
disabled
torch
2
OOM Error
disabled
fused
2
2
OOM Error
disabled
8ash_attm1.0.88
2
2
OOM Error
disabled
torch
2
OOM Error
disabled
fused
2
OOM Error
every layer
torch
2
OOM
disabled
torch
OOM Frror
every layer
fused
2
Error
disabled
fused
2
OOM Error
disabled
flash 1.0.
2
OOM Error
disabled
torch
2
OOM Error
disabled
fused
2
OOM Error
disabled
flash_attnn attn1.0.8
2
OOM
disabled
torch
OOM Error
disabled
fused
2
Error
disabled
flash flash_ att
2
OOM Error
disabled
torch
2
OOM Error
disabled
fused
2
OOM Error
every layer
torch
OOM Errur
disabled
turch
OOM
every layer
OOM Error
disabled
fused
OOM Error
every layer
flash_attal attm1.0..
OOM Error
disabled
1.5
OOM Error
disabled
torch
OOM
disabled
fused
OOM Errur
disabled
attn1..
0OM
disabled
torch
OOM Error
disabled
fused
OOM Error
every layer_
torch
OOM Error
disabled
torch
OOM Error
every layer
fused
OOM
disabled
fused
OOM Frror
every layer
aush_attn1.0..
Error
disabled
flash 1.8
OOM Error
every_layer
torch
17

---

Published as a conference paper at COLM 2024
Slep Time
MFU
Activation
Kernel
MB Size
TP size
PP Size
OOM Frror
disabled
torch
OOM Errur
disabled
fused
OOM
disabled
attn1.0.8
OOM Error
disabled
torch
OOM Error
disabled
fused
OOM Error
disabled
OOM Error
disabled
torch
OOM Error
disabled
fused
OOM Error
disabled
attn118
Table 4: Performance analysis of LLAMA 13B model trained on 64 A100 GPUs with the AA-Scaling codebase: A comprehensive sweep of 3D parallel layout configurations. The table displays variations in step time and MFU across different activation checkpointing types, kernels, micro batch (MB) sizes, tensor parallel (TP) sizes, and pipeline parallel (PP) sizes. The analysis also includes Out of Memory (OOM) error occurrences.
B.3 LLAMA 13B with 8k sequence length
Step Time
MFU
Activation
Kernel
MB Size
TP Size
59.41
disabled
flash_atte
2
19.32
56.61
disabled
flash_atth2 RMS kern.
4
21.36
51.21
disabled
flash_att22 RMS kern.
21.94
49.93
disabled
flash_att22 RMS kern.
2
21.93
49.88
disabled
flash_att22
4
23.4
46.90
disabled
flash_atin2 RMS
23.78
46.27
disabled
flash_attn2 RMS kern.
4
44.42
every layer
flash_att12
2
24.84
44.03
disabled
flask_attii.0.88
4
42.08
disabled
flash_att22
2
26.29
41.81
every layer
flash flash_attn2
26.63
41.08
every layer
flash_att22
.4
40.59
disabled
flash
27 42
39.89
every layer
attn2
2
27.53
39.73
every_ every_layer
flash atth2
2
28.08 08
38.94
every layer
flash_att22
4
28.15
38.85
every layer
flash attn2
28.52
38.34
disabled
flash at0.
2
28.73
38.06
every layer
flash_attn1.0
2
37.6
every every_layer layer
flash attn2
29.19
37.49
every layer
Bash_attz22
2
29.21
37.43
disabled
flsshatts1.00.
30.03
36.42
every_layer layer
fltsh_atto...88
30.08
36.36
every layer
flash_attn2 attm2
30.14
36.28
every Jayer
1ash_attm1.0.8
31.65
34.55
every layer
fash_atta1.0.
2
every every_laye layer
flash a8t.
32.13
34.04 every layer
flash_att22
32.57
33.58
every layer
alts
32,91
33.23
every lveer_layer
flash_attn1.0.8
2
34.75
32.47
every layer
fash_attm1.00
4
34.33
31.85
every Jayer
Bash_attn1.0.8
34.61
31.60
every layer
fassh_tttl.0..
4
31.58
every every_layee layer
flash,
34.76
31.46
every layer
flash_attmi
2
35.86
30.49
every layer
flash_atte1.0.
1
36.07
30.32
every layer
flash attn2
4
36.33
30.10
every layer
flash lnsh_attm1.008
36.71
29.92
every every_layer layer
flash
37.63
29.06
every layer
flash
2
28.59
every layer
at0.
38.61
28.33
every layer
flash_attn..0$
4
39.76
27.50
every every_layer layer
altsh_atto..088
4
51.95
21.05
every every_layer layer
torch
34.97
19.89
every layer
torch
2
35.08
19.85
every layer
torch
2
58.37
18.73
every layer
turch
every layer
39.00
18.56
every
layer_layer layer
torch
4
60.64
18.03
every layer
torch
62.00
17.64
every every_layer
torch
2
64.25
17.02
every every_layer
torch
2
4
layer torch
18

---

Published as a conference paper at COLM 2024
Slep Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
66.48
16.45
every layer
torch
OOM Errur
disabled
fltsh_att22 RMS kern.
OOM
disabled
OOM Error
disabled
flash_attn2 RMS kern.
4
OOM Error
disabled
flash RMS kern.
2
OOM Error
disabled
flash_att22 RMS kern.
OOM Error
disabled
Bash_attn2 RMS kern.
4
OOM Error
disabled
flash_attn2 RMS kern.
OOM Error
disabled
flash_ RMS kern.
2
2
OOM
disabled
OOM Error
disabled
flash RMS kern.
2
OOM Error
disabled
flash RMS kern.
OOM Errur
disabled
flash_ RMS kern.
OOM Error
disabled
flash_attn2 RMS kern.
4
OOM Error
disabled
flash_attn2 RMS
OOM Error
disabled
flash RMS kern.
2
Error
disabled
attm2 RMS flash_attn2 RMS kern.
OOM Erroz
disabled
OOM Error
disabled
flash_att22 RMS kern.
OOM Error
disabled
flash_attn2 RMS kern.
OOM Error
disabled
flash_attmy RMS kern.
OOM Error
every every_laye layer
flash_atin2
OOM Error
every layer
flash_att22
Error
every layer
flash
2
OOM Error
every every_layer layer
flash_atin2
2
OOM Errur
every every_layer layer every layer every layer every Jayer_laee every layer
flash_att22
2
OOM Frror
flash_ _
OOM Error
flash_athi
OOM Error
flash
OOM Errur
flash_att22
Erroz
every every_laye layer
flash attn2
OOM Error
every layer
flash_att22
OOM Error
disabled
Bash_attn2
4
OOM Error
disabled
flash_ _
OOM Error
disabled
flash_ath22
OOM Error
disabled
flash ath22
OOM Errur
disabled
flash_attr
2
OOM Error
disabled
flash atta2
OOM Errur
disabled
flash_att22
OOM Error
disabled
Rash_att22
OOM Error
disabled
flash_attm2
OOM Error
disabled
flash_
4
OOM Errur
disabled
flash attn2
OOM Error
disabled
flash_attmi
OOM Error
disabled
aath
4
OOM Errur
disabled
flash_att22
OOM Error
disabled
flash_att22
OOM Error
disabled
flash_attn22
OOM Error
disabled
flash
OOM
disabled
flash, att22
OOM Error
disabled
flash_attni
Error
disabled
flash attn2
2
OOM Error
disabled
flash_attn2 attn2
OOM Error
disabled
flash_att22
OOM Error
disabled
flash
OOM Error
disabled
flash
OOM
every every_ layer
OOM Error
every layer
flash_attn....8
Error
every every_layer
attn0.
2
OOM Frror
every layer
Aash_attn1.0.8
OOM Error
every layer
fash__ttnn1.088
OOM Error
every every_layer layer
attn1.
OOM Errur
every layer
flash_attnl.0
OOM
every layer
a.0.
OOM Error
every layer
faash_attn.0.88
2
OOM Error
every_layer
attn..0
4
OOM Error
every every_layer layer
1ash_attn1.00.8
OOM Error
every layer
fassh_tttn1.0.8
OOM
disabled
attsl.0
OOM Errur
disabled
flash_attn1.0.8
2
0OM
disabled
1.00
OOM Error
disabled
flask_attn1.0.8
4
OOM Error
disabled
flash_attal.0..
2
OOM Error
disabled
attt_00
OOM Error
disabled
flash
4
OOM
disabled
flash att.1.08
OOM Frror
disabled
Mush_attml.0
Error
disabled
flash 1...
OOM Error
disabled
attsk_attn1...8
4
2
19

---

Published as a conference paper at COLM 2024
Slep Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
OOM Frror
disabled
Jash_attn1.0.8
OOM Errur
disabled
flash_ttto..0..
2
4
OOM
disabled
OOM Error
disabled
flash_tttn1...8
Error
disabled
faash_ttt1.0..
4
OOM Error
disabled
fltsh_atti1...8
2
OOM Error
disabled
flash, attt__
OOM Error
disabled
1ash_attn1.0088
OOM Error
disabled
faash_ttta1.088
2
OOM
disabled
a.0
OOM Error
disabled
flash_tttx..0.8
4
OOM Error
disabled
fash_atts1.0.8
OOM Errur
disabled
flash_attn...88
OOM Error
every layer
torch
4
OOM Error
every layer_ayerr
torch
OOM Error
every layer
torch
2
Error
every layer
torch
OOM Erroz
every_ layer
torch
2
OOM Error
every layer
torch
OOM Error
every layer
torch
OOM Error
every layer
torch
4
OOM Error
every_ layer
torch
OOM Error
every Jayer
torch
Error
every_ layer
torch
OOM Error
every_ Jayer
torch
OOM Errur
every every_layer
turch
OOM Frror
every layer
torch
OOM Error
every. layer
torch
OOM Error
every_laye
torch
OOM Errur
disabled
turch
4
Erroz
disabled
torch
OOM Error
disabled
torch
OOM Error
disabled torch
4
OOM Error
disabled torch
OOM Error
disabled torch
OOM Error
disabled torch
OOM Errur
disabled turch
2
Error
disabled torch
OOM Errur
disabled turch
OOM Error
disabled torch
OOM Error
disabled torch
OOM Error
disabled torch
4
OOM Errur
disabled torch
OOM Error
disabled torch
OOM Error
disabled torch
4
OOM Errur
disabled torch
2
OOM Error
disabled torch
OOM Error
disabled torch
OOM Error
disabled torch
2
OOM
disabled torch
OOM Frror
disabled torch
4
Error
disabled torch
0OM Error
disabled torch
OOM Error
disabled torch
4
OOM Error
disabled
torch
2
OOM Error
disabled
turch
Table 5: Performance analysis of LLAMA 13B model with sequence length of 8k trained on 128 A100 GPUs with the AA-Scaling codebase. All measurements use the FLASHATTENTION or Py Torch kernel. The analysis also includes Out of Memory (OOM) error occurrences.
B.4
LLAMA 30B
Step Time
Activation
Kernel
Size
22.67
49.22
disabled
flash_attn2 RMS kern.
24.00
46.76
disabled
aath2 RMS kem.
24.26
46.01
disabled
flash_att22 RMS kem.
24.70
45.16
disabled
alts2_
25.05
44.55
disabled
flash_att22 RMS kern.
4
25.32
disabled
attn2
26.06
42,80
disabled
flash_attn1.00
4
26.26
42.48
disabled
flash_attn2 RMS ker..
26.28
42.45
disabled
flash_attn2 RMS kem.
20

---

Published as a conference paper at COLM 2024
Step Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
26.94
41.42
disabled
aash_attn2... RMS kern.
2
28.34
39.39
disabled
alash_att22 RMS kern.
2
every layer
_ash_
29.15 13
38.32
disabled
Basshattn2
2
38.26
every every_layer
flashatttn1.0s
29.45
37.88
disabled
alash_att22
29.55
37.75
disabled
Blash_attnl..
29.6 64.64
37.67
disabled
Jash_attni.08
29.76
37.48
every every_layer
fltsh_ttt22
4
37.23
every layer
flash lash_attni.a
30.08
37.09
disabled
fass_attn2
31.13
35.84
every every_layer
altshattt2
31.38
35.54
disabled
flash_ttt11.08
31.47
35.46
disabled
flash_attn2
31,66
35.24
disabled
flash_attn2
31.77
35.15
disabled
flash_
RMS kem.
31.81
35.08
every rayer_layer
flash_att22
31.94
34.93
every_ every_layer
flash_att22
2
32.17
34.68
every layer
aats_at
32.27
34.57
every layer
flash_attn2
32.42
34.41
disabled
flash a.t
32.73
34.08
every Jayer
fused
32.84
33.96
every eayer_layer
aathhat
2
33.08
33.72
every Jayer
fused
33.29
33.51
every Jayer
fasshattni..08
33.47
33.33
every every_layer
flts_attt1.0.8
33.70
33.11
every layer
flash_attmy
33,75
33.05
every_layer
flash_attni 1.8
4
33.79
33.01
layer
flash_tttn..0s
33.97
32.85
disabled
flssh_ttt1...88
2
32.61
every Jayer
fltshatth2
34.45
32.39
every_layer Jayer
flash_att22
35.07
31.84
every every_layer
flash_att22
35.23
31.67
layer
flash_attn
35.32
31.61
disabled
flash RMS kem.
31.09
every Jayer
fused
36.42
30.66 every layer
flash_attn1.088
2
30.58 every every_layer
fused
4
37.12
30.05 every layer
flash_att22
2
37.76
29.58
every layer
flash_attn2
37.81
29.50 every layer
attsh_
37.99 94
29.40 every layer
fused
4
every layer
fused
38.07
29.32
every layer
flash_attn1.00
2
38.28
29.15
every every_layer
fltshattn2
4
38.51
28.97
every layer
flash_attn2
4
38 60
28.90
every layer
flash flash_attn1.00
39.37
28.19
every Jayer
flash_att22
39.60 61
28.17
every layer
fused
2
28.14 every every_layer layer
flash_attn1.08
39.70
28.10 every layer
flash_
39.74
28.07
every every_layer
fused
39.81
28.03
disabled
flash_attmy
40.07
27.8 27.84
every every_layer layer
flash_att22
2
40.22
27.73
every every_aaee Jayer
fused
40.29
27 68
every layer
ftash_attn1.0.8
2
27.20
every layer
flash_
41.74
26.73
every_ layer
fused
41.80
26.69
every every_layer
flashattt1.008
41.84
26.66
every layer
flash_attn 1.0
41.94
26.60
every_layer
fash_attn1.03
every every_layer layer
alash
42.22
26.42
every layer
fused
26.30
every layer
att_1.8
43.24
25.80
disabled
flash_attn1.0s
43.51
25.64
every layer
fatsh
43.92
25.40
every_ layer
flash_att22
44.09
25.30
every Jayer
flash_ttt22
45.56
24.49
every layer
fassh
0OM Error
disabled
flash_attn2 RMS kem.
disabled
flash_att22 RMS
OOM Errr
disabled
flash_att22 RMS kem.
0OM Error
disabled
flash RMS kem.
OOM Eror
disabled
flash_ RMS kem.
OOM Error
disabled
flash_ttt22 RMS kem.
2
OOM
disabled
flash RMS
0OM Eror
disabled
flash_attn2 RMS ker..
2
Eror
disabled
altn2 RMS
OOM Ervor
disabled
alash2 RMS kem.
21

---

Published as a conference paper at COLM 2024
Step Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
0OM Eror
disabled
aash_attm2 RMS ker.
OOM Error
disabled
flash_att22 RMS kern.
disabled
OOM Eror
disabled
flash_att22 RMS kem.
2
Error
disabled
flash_att22 RMS kem.
OOM Error
every_layer
flashattt22
0OM Error
every_ layer
flash_att22
0OM Error
every- every_layer
flash_attn2
OOM Error
every_layer
altn2_
every layer
attn2
OOM Error
every layer
flash_att22
OOM Error
disabled
altsh_ttt22
OOM Error
disabled
alash_att22
OOM Error
disabled
lash_attn2
OOM Erro
disabled
flash_att22
OOM Error
disabled
flash_
Error
disabled
OOM Error
disabled
flash_att22
OOM Error
disabled
alth
0OM Error
disabled
flash_attng
0OM Error
disabled
fash_attn2
OOM Error
disabled
flash_att22
OOM Error
disabled
flash_attn2
Error
disabled
alash_att22
COM Error
disabled
alash_attn2
OOM Error
disabled
alash_ttt2
0OM Error
disabled
flash attm2
COM Error
disabled
flash_atton
OOM Error
disabled
fltsh_
OOM Error
disabled
alash_att22
Error
disabled
althh_
COM Error
layer
fused
OOM Error
every every_layer
fused
0OM Emor
every_ layer
fused
OOM Error
every layer
fused
COM
every every_layer layer
fused
OOM Error
every layer
fused
OOM Error
every every_layer
flash_ttt1.008
OOM Error
every every_layer layer every layer
flash_ttt1.0.8
0OM Erro
flash laash_attn1.08
OOM Emor
every every_layer layer
flash_attn1.08
OOM Error
every layer
fash_attn1.00
COM
every layer
flssh_
OOM Error
disabled
fused
OOM Error
disabled
fused
0OM Error
disabled
fused
COM Erro
disabled
fused
OOM Error
disabled
fused
OOM Error
disabled
fused
COM
disabled
OOM Emor
disabled
fused
2
Error
disabled
fused
(OM Eror
disabled
fused
COM Error
disabled
fused
OOM Error
disabled
fused
OOM Error
disabled
fused
4
disabled
fused
OOM Emor
disabled
flash_attn1.0s
Error
disabled
flash_
(OM Emor
disabled
1ash_attn1.008
OOM Error
disabled
flash_attn1.03
4
OOM
disabled
dassh
OOM Error
disabled
flash_att1...88
disabled
flash_attn1.05
OOM Emor
disabled
flash_attn1.0s
2
0OM Errrr
disabled
fatsh
OOM Emor
disabled
flash_attn1.08
OOM Error
disabled
flash
4
COM
disabled
fash_
0OM Error
disabled
flash
flash_attn1.03
OOM Error
disabled
flash
2
0OM Error
disabled
fatth
OOM Eroor
disabled
attsh
OOM Error
disabled
aatn
4
OOM
disabled
0OM Eror
disabled
flash attn1.00
Kernel unavail.
every Jryer_layer
fused
Kernel unavail.
every Jayer
fused
2
22

---

Published as a conference paper at COLM 2024
Step Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
Kernel unavail.
every_ layer
fused
Kernel unavail.
every layer
fused
2
unavail.
every layer
Kernel unavail.
pvery_layee
fused
2
Kernel unavall.
every_layer
fused
Kernel unavail.
every_ every_layer
fused
Kernel unavail.
every. every_layer layer
fused
Kernel unavail.
every_ Jayer
fused
Kernel unavail.
disabled
fused
4
Kernel unavail.
disabled
Kernel unavail.
disabled
fused
Kernel unavail.
disabled
fused
Kernel unavail.
disabled
fused
2
Kernel unavoil.
disabled
fused
Kernel unavail.
disabled
fused
Kernel unavail.
disabled
fused
2
Kernel unavail.
disabled
fused
Table 6: Performance analysis of a LLAMA 30B model trained on 256 A100 GPUs with the AA-Scaling codebase. The analysis also includes Out of Memory (OOM) error occurrences and errors where the fused kernel does not support this specific tensor parallel configuration.
B.5
LLAMA 30B with 8k sequence length
Step Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
49.43
51.40
disabled
flash_attnis RMS kern.
4
50.23
50.57
disabled
flash_atth2 RMS kern.
8
54.78
46.37
disabled
flash_att22
RMS kern.
62.84
40.43
every every_laye layer
flash_attn2
4
62.99
40.33
every every_layer layer
flash_attr?
2
63.99
39.70
every layer
flash_att12
2
B
every layer
66.71
38.08 every_ layer
flash_attnt
2
4
68.99
36.82
every_layer
flash_att22
S
69.45
36.58
every every_layer
fltsh_atts1...8
2
70.01
36.29
every layer
flash Mash_attn1.0..
4
71.81
35.38
every Jayer
flash_attn2
71.87
35.34
every layer
flash_attnl.0.8
B
every layer
attn2
2 2 4
73.33
34.64 every_ layer
flash_atth?
75.02
33.86
every every_laeer
flash_ttte1.0
75.53
33.64
every layer
flash_att22
76.53
33.20
every layer
flash_attm
42
32.81 every_ Jayer
77.86
32,63
every layer
flash_att22
B 2 B
79.71
31.87
every layer
flash attnl..t
79.74
3186
every every_laerr layer
flash_attn2
80.25
31.65
every every_layer
fltsh_ttte....
2
82.16
30.92
every layer
flash atth_attn1.0.8
82.62
30.75
every layer
flash_attm
30.24
every every_layer layer
Bash_atin1.00
85.34
29.77
every layer
flash_attne
16
85.53
29.70
every_layer layer
attsh_
B
87.18
29.14
every layer
fassk_ttt1.088
B
OOM Errur
disabled
flash RMS kerm.
OOM Ernor
disabled
flash aath_attn2-40 RMS kem.
8
OOM Error
disabled
flash_att22 RMS kern.
4
8
OOM Error
disabled
flash RMS
OOM Errur
disabled
flash RMS kern.
2
16
OOM Error
disabled
flash attn2 RMS kern.
OOM Errur
disabled
alt2h_att22 RMS kern.
OOM Error
disabled
flash_attn2 RMS kern.
B
OOM Error
disabled
flash_attn2 RMS kern.
8
OOM Error
every layer
flash_attn2
4
16
OOM Error
every every_layer layer
flash
2
OOM Errur
every layer
flash_attri
4
OOM Error
every every_layer
alts2
OOM Errur
every every_layer layer
att22
2
OOM Frnor
every layer
flash_att22
B
OOM Error
every Jayer
alash_attn22
OOM Error
every layer
flash_ttt22
2
2
OOM
every layer
flash attn2
OOM Frror
disabled
flash_atth
4
16
23

---

Published as a conference paper at COLM 2024
Slep Time
MFU
Activation
Kernel
MB Size
TP Size
PP Size
OOM Frror
disabled
flash_att2
OOM Errur
disabled
fltsh_
OOM
disabled
OOM Error
disabled
flash_atthe
16
OOM Error
disabled
flash
B
OOM Error
disabled
alash_att22
OOM Error
disabled
flash_attn2
2
OOM Error
disabled
flash_attn2
OOM Error
disabled
flash_att22
B
OOM
disabled
OOM Error
disabled
flash
4
2
OOM Error
disabled
flash
OOM Errur
disabled
flash_att2
B
OOM Error
disabled
flash_attn2
OOM Error
disabled
flash_atin2
2 16
OOM Error
disabled
flash_att22
Error
disabled
flash attm2
OOM Erroz
disabled
flash
4
OOM Error
disabled
flash_att22
OOM Error
disabled
flash att2
OOM Error
disabled
flash_attmy
8
OOM Error
disabled
flash
OOM Error
disabled
flash_attal
2
Error
disabled
flash a.8
OOM Error
disabled
flask_atta1.0.8
OOM Error
every every_layer
flash_atto.0.88
4
OOM Frror
every layer
flash_attn1.0.5
2
OOM Error
every layer
fasshattml.088
2
OOM Error
every_ layer
attn10
OOM Errur
every layer
flash_
B S
Error
every_ layer
flash att
OOM Error
layer
fatsh_att1..88
2
OOM Error
disabled
flash_attn1.0.8
OOM Error
disabled
flash_attn1.0
2 4
OOM Error
disabled
fassb_tttn1.0.8
OOM Error
disabled
altsh_
OOM Errur
disabled
flash_attn1.0.8
B
OOM Error
disabled
flash attn
4
OOM Errur
disabled
fltsh_attn1.088
2 B
OOM Error
disabled
1ash_attn1.0.8
OOM Error
disabled
flash_attn1.0.8
OOM Error
disabled
faash_atta1.0.8
2
OOM Errur
disabled
flash attnl.0
OOM Error
disabled
flash_attn1.0.8
4
OOM Error
disabled
fltsh_att1....8
OOM Errur
disabled
Bash_attn11.0.8
OOM Error
disabled
flash_attn1.0.8
4
OOM Error
disabled
flash_attn1.0..
2
OOM Error
disabled
flash RMS kem..
4
OOM
disabled
flash, attn2 RMS
OOM Error
disabled
flash_att22 RMS ker..
4
Error
disabled
flash attn2 RMS
2
0OM Error
disabled
flash aath__ttn2240 RMS ker..
4
OOM Error
disabled
flash_attni RMS kern.
4
OOM Error
disabled
flash RMS
2
OOM Error
disabled
flash_ RMS kern.
2
2
OOM
disabled
flash,
OOM Error
disabled
flash_att22 RMS kern.
4
2
OOM Error
disabled
flash attn2 RMS kern.
2
Table 7: Performance analysis of a LLAMA 30B model with 8k sequence length trained on 128 A100 GPUs with the AÅ-Scaling codebase. All measurements use the FLASHATTENTION kernel. The analysis also includes Out of Memory (OOM) error occurrences.
B.6
LLAMA 65B
Step Time
MFU
Activation
Kernel
MB Size
TP size
PP Size
79.31
35.26
disabled
flash RMS
79.54
55.10
disabled
flash_attrz, - RMS kern.
disabled
86.55
50.63
disabled
flash RMS kem.
86.61
50.60
disabled
flash attn2 RMS kern.
87.14
50.30
disabled
flash_att22 RMS kern.
24

---

Published as a conference paper at COLM 2024
Step Time
MFU
Activation
Kemel
MB Size
TP size
PP Size
88
49.71
disabled
flash_attn2
8
92.61
47.32
disabled
flash_ RMS kern.
disabled
flash
RMS
101.97
42.98
disabled
flash_attne
4
103.17
42.48
disabled
flash_attn2 RMS kem.
106.59
41.11
disabled
flash_attn1...8
107.39
40.81
every layer
flash flash_attn2
107 83
40.64
disabled
flash_attn2
RMS kern
107 97
40.39
disabled
flash_attn2
disabled
attm1.0.8
109.04
40.19
Jayer
flash_att22
110.87
39.53
every evyre_layer
flash_
111.35
39.36
every layer
flash_attn2
111.83
39.19
disabled
flash ahtnatt22 RMS ker.
111.83
39.19 every layer
0
112.77
38.86 every every_ Jayer
flash_attn2
38.30 every layer
flash attm2
115.03
38.09
every_ Jayer
0ash_attn1.08
115.58
37.92 disabled
flash_attn2 RMS kem.
116.25
37.70 every layer
flash_attnl.088
116.27
37.69 every. layer
Mash_attn|.08
118.39
37.02 every_ layer
flash_att22
118.68
36.93 every every_layer
flash_attn2
121.70
36.01 every layer
flash_attn1.08
121.93
35.95 disabled
flash RMS kem.
125.11
35.03
every every_layer layer
flash_attn2
126.61
34.61
every layer
flash_att22
128.00
34.22
every layer
flash_attnl
129.68
33.80
every layer
flanh_att22
129.9 92
33.73
every layer
flash_attnl.03
132.34
33.11
every layer
flash_att11.08
133.07
32,93
every_layer
flash_att1
4
133.56
32.81
every layer
flash flash_attn2
136.43
32.12
every layer
flash att22
137.73
31.82
every layer
flash aash_attn1.0.8
8
31.38
disabled
flash
140.15
31.28
disabled
flash_
2
140.12
31.28
every every_layer
aast1attm1.03
140.95
31.09
every_ layer
altsh_att22
141.87
30.89
every layer
flash_attri attn2
142.77
30.70
disabled
lash_attm2
144.59
30.31
every Jayer
Dash_attn1.03
every layer
flash Bash_attn1.0.
151.95
28.84
disabled
flash_attn2
8
153.61
28.53
disabled
flashattt1..0
156.54
28.00
every_ layer
flash_attn2
159.57
27.46
every layer
aash_attnlle .08
164.2
26.69 every Jayer
flash_attn
168.66 6-
25.99 every layer
flash_attn1.0
2
every layer
flash att22
173.16
25.31
every layer
flash_attnl
4
174.43
25.13
every layer
flash attm2
178.66
24.53
every layer
flash Aash_attnl.05
2
180.93
24.22
every layer
flash_att22
23.69
every Jayer
alts2_
187 81
23.33
every layer
flash_ttn1.003
4
every layer
flash atch_attn1...
192.69
22.75
every layer
flash_attmi
2
200.12
21.90
every Jayer
flash_att11.08
203.08
21.58
every layer
flash_attn2
209.24
20.95
every layer
flash_att22
every layer
flash
215.12
20.37
every layer
flash_attnl
every layer
247 97
17.67
every layer
flash_atm1.0.8
8
OOM Error
disabled
flash RMS kem.
OOM Error
disabled
lash_attm2 RMS kem.
8
OOM Error
disabled
flash RMS kem.
OOM
disabled
flash RMS
0OM Error
disabled
flash_ RMS kern.
4
disabled
OOM Error
disabled
flash RMS kem.
2
0OM Error
disabled
flash RMS kern.
0OM Error
disabled
flash RMS ker.
0OM Error
disabled
flash - RMS kem.
0OM
disabled
flash RMS
0OM Error
disabled
flash_attn2- - RMS kern.
2
Error
disabled
flash attm2
OOM Error
every_layer
flash_
2
25

---

Published as a conference paper at COLM 2024
Step Time
MFU
Activation
Kemel
MB Size
TP size
PP Size
0OM Error
every. layer
flash_attn2
OOM Error
disabled
fltsh_att22
disabled
OOM Error
disabled
flash_att22
Error
disabled
flash_
OOM Erzor
disabled
flash_att22
0OM Error
disabled
flash attn2
OOM Error
disabled
flash attm2
OOM Error
disabled
flash_attn?
disabled
OOM Error
disabled
_lash_att22
OOM Erzor
disabled
alash_
0OM Erzor
disabled
altsh_attn2
0OM Error
disabled
flash attn2
OOM Error
disabled
flash attm2
OOM Error
disabled
flash
Error
disabled
flash attm2
OOM Error
disabled
_lash_att22
COM Error
disabled
flash_att22
0OM Erzor
disabled
flash_attr)
0OM Error
disabled
flash_attn2
OOM Error
disabled
flash_attn)
OOM Error
disabled
flash_attn1.0.
Error
disabled
0.
COM Error
disabled
flashatttn.0.8
COM Error
every laaer
flash_attn1.088
OOM Error
layer
flash_attnl.03
OOM Error
disabled
flash_attn.088
0OM
disabled
flash_attn1.03
OOM Error
disabled
flash_attn1...8
Error
disabled
flash_attn..08
COM Error
disabled
flash_attn1.0.8
OOM Error
disabled
flash_attnl.08
OOM Error
disabled
flash_attnl.03
OOM Error
disabled
flash_attn1.08
4
OOM Error
disabled
flash_attn1.02
OOM Error
disabled
flash_attn1.0.8
2
COM Error
disabled
flash_att1...8
0OM Erzor
disabled
flash_attn1.0.8
0OM Error
disabled
flash_attn1.00
0OM Error
disabled
flash_attl.088
OOM Error
disabled
flsshattnn.003
COM
disabled
attn1_0
OOM Error
disabled
flash_
OOM Error
disabled
flsshattn1.03
OOM Erzor
disabled
flash_atttl.0.5
Table 8: Performance analysis of LLAMA 65B model trained on 128 A100 GPUs with the A AA-Scaling codebase. All measurements use the FLASHATTENTION kernel. The analysis also includes Out of Memory (OOM) error occurrences.
26

---

Published as a conference paper at COLM 2024
C
Sequence Parallelism Training Efficiency Sweep
C.1 Sweep Configurations
Model
Seq. Len.
GPUs
sizes
sizes
MB Sizes
Seq. Parallelism
13B
2k
32
(1, 2}
4, 8)
(yes, no)
138
8k
64
(1, 2,
(1, 2,
4]
(yes,
30B
2k
64
(1, 2,
(1, 2,
(1, 4)
Cyes, no)
30B
8k
64
(2, 4)
(2, 4, 8, 16)
(1, 2, 4)
(yes, no)
65B
(2, 4, 83
(2, 4, 8)
(1, 2, 4)
(yes, no)
Table 9: Search space of additional sequence parallel training efficiency sweep. We sweep over the Cartesian product of all mentioned options similar to Table 1. All runs use FLASHATTENTION-2 the RMSNorm kernel, and do not use activation checkpointing.
C.2 LLAMA 13B
Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
53.81
69.66
False
53.99
69.45
True
58.01
64.54
True
58.62
63.88
False
59.35
62.86
True
61.12
62.28
2
True
60.23
False
60.96
61.43
True
62.09
60.34
False
65.04
57.55
True
66.13
56.59
False
66.26
36.47
True
67.47
55.44
True
55.08
False
68.48
54.62
False
Error
True
0OM Ervor
True
COM Error
True
OOM Error
True
OOM Error
True
COM Error
True
COM Error
False
OOM Error
False
COM Error
2
False
OOM
False
OOM Erzor
2
False
False
COM Error
False
OOM Error
False
Table 10: Performance analysis of LLAMA 13B model trained on 32 A100 GPUs, with and without sequence parallelism. All measurements use PLASHATTENTION-2 the RMS norm kernel, and do not make use of activation checkpointing. The analysis also includes Out of Memory (OOM) error occurrences.
C.3 LLAMA 13B 8k sequence length
Step Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
34.84
62.78
2
True
34.85
False
35,80
61.10
2
True
36.60
59.76
False
36.99
59.13
True
38.85
56.31
True
38.90
56.23
False
40.70
53.74
False
27

---

Published as a conference paper at COLM 2024
Step Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
40.82
53.58
True
41.06
53.27
2
True
False
44.03
49.68
False
46.37
47.18
True
46.89
46.65
True
49.96
43.79
True
51.03
42.87
False
51.91
42.14
2
True
39,93
56.23
38.90
False
59.62
36.69
False
62.04
35.26
False
0OM Error
True
COM Eroo
True
0OM Error
True
True
COM Error
True
COM Error
True
0OM Error
True
0OM Error
True
COM Error
True
0OM Error
2
True
Error
True
COM Error
False
0OM Error
False
(OM Error
False
COM Error
False
0OM Error
False
OOM Error
False
Error
False
0OM Error
False
0OM Error
False
0OM Error
False
COM Error
False
0OM Error
False
0OM Error
8
True
COM Error
True
0OM Error
True
0OM Error
True
COM Eroo
True
OOM Error
2
True
COM
True
OOM Erzor
False
COM Error
False
OOM Error
False
0OM Error
False
0OM
False
OOM Error
2
2
False
OOM
False
OOM Erro
False
Table 11: Performance analysis of a LLAMA 13B model with 8k sequence length trained on 64 A100 GPUs, with and without sequence parallelism. All measurements use FLASHATTENTION-2, the RMS norm kernel, and do not make of activation check- pointing. The analysis also includes Out of Memory (OOM) error occurrences.
C.4 LLAMA 30B
Step Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
False
72.60
61.47
True
72.66
61.42
True
74.86
59.61
False
76.38
58 38.27
2
True
76.84
38.07
True
77.59
57.52
True
False
80.17
55.66
False
81.06
55.05
True
62.22
54.28
True
83.11
53.71
True
83.76
53.30
True
28

---

Published as a conference paper at COLM 2024
Step Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
84.88
52.57
False
85.21
52.37
2
True
86. 44
51.63
False
86.67
51.49
False
67.68
50.89
True
89.33
49.95
True
89.76
49.71
False
89.86
49.66
False
False
95.32
46.72
False
OOM Error
True
0OM Error
True
0OM Error
True
COM Erro
True
OOM Error
True
True
COM Error
True
0OM Error
True
0OM Error
False
0OM Error
False
0OM Error
False
0OM Error
False
Error
False
COM Error
False
COM Error
False
(OM Error
False
COM Error
False
OOM Error
False
OOM Error
2
False
Error
True
OOM Error
True
0OM Error
2
True
0OM Error
True
OOM Error
2
False
OOM Error
False
OOM Error
2
False
OOM Error
False
Table 12: Performance analysis of a LLAMA 30B model trained on 64 A100 GPUs, with and without sequence parallelism. All measurements use FI.ASHATTENTION-2 the RMS norm kernel, and do not make use of activation checkpointing. The analysis also includes Out of Memory (OOM) error occurrences.
C.5 LLAMA 30B 8k sequence length
Step Time
MFU
Size
TP size
PP Size
Sequence Parallel
84.37
60.22
True
39.35
True
86.
58.97
16
True
58.97
True
93.84
54.15
False
96.70
52.54
False
98.37
51.65
16
False
0OM Error
2
True
OOM Error
True
OOM Error
True
OOM Error
True
0OM Error
True
OOM Emor
True
0OM Error
True
OOM
True
OOM Error
False
16
False
OOM Error
False
0OM Error
16
False
0OM Emor
False
0OM Error
False
OOM
False
OOM Error
False
False
OOM Error
2
False
29

---

Published as a conference paper at COLM 2024
Step Time
MFU
MB Size
TP size
PP Size
Sequence Parallel
0OM Error
False
OOM Error
2
False
False
OOM Error
2
False
OOM Error
False
OOM Error
16
False
0OM Error
16
False
OOM Error
False
OOM Error
False
Table 13: Performance analysis of a LLAMA 30B model with 8k sequence length trained on A100 GPUs, with and without sequence parallelism. All measurements use FLASHATTENTION-2, the RMS norm kernel, and do not make use of activation check- pointing. The analysis also includes Out of Memory (OOM) error occurrences.
C.6 LLAMA 65B
Step Time
MFU
Size
TP size
PP Size
Sequence Parallel
147.02
59.62
True
149.92
58.47
True
149.97
38.44
2
True
57.42
False
156.40
56.04
True
158.74
55.22
False
159.57
54.93
True
162.32
54.00
True
166.36
52.69
True
166.49
52.65
8
False
52.27
False
168.70
51.96
False
169.39
51.75
False
172.11
50.93
False
178.
49.07
True
179.78
48.76
False
179.8 87
48.74
8
True
True
192.45
45.55
True
193.69
45.26
False
198.11
44.25
True
201.44
43.52
True
False
207 26
42.29
False
223.09
39.29
False
233.00
37.62
False
COM Emoo
True
OOM Error
2
True
COM
True
COM Eroo
True
True
COM Error
True
0OM Error
True
COM Error
True
0OM Error
2
False
0OM
False
OOM Eror
False
Errrr
False
COM Error
False
COM Error
False
OOM Error
False
OOM Error
False
0OM
False
OOM Error
False
Error
False
OOM Ervor
False
OOM Error
2
False
Table 14: Performance analysis of a LLAMA 65B model trained on 64 A100 GPUs, with and without sequence parallelism. All measurements use FLASHAITENTION-2. the RMS norm kernel, and do not make use of activation checkpointing. The analysis also includes Out of Memory (OOM) error occurrences.
30


Reviews:
[{"review_id": "n3d5QF0Pwl", "reviewer": "Reviewer_XpzU", "summary": "The authors conduct a large-scale empirical study on the effectiveness of various combinations of efficient parallelization strategies for LLM distributed training. They experiment primarily on Llama models of sizes 13B, 30B, and 65B, with sequence lengths of 2k or 8k tokens, and use up between 64 and 256 GPUs for their experiments. The authors report Model FLOPs Utilization (MFU) as their metric for efficiency of a training configuration. The authors include a set of concrete recommendations to researchers and practitioners based on their findings, including a micro-batch size of 1 whenever possible, prioritizing tensor and pipeline parallization, and reserving sequence parallelization for models larger than 30B parameters with sequences longer than 2k tokens.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for taking the time to review our paper and provide thoughtful feedback. We are happy to hear that you found our careful experimental design extremely valuable in verifying today's \"common sense.\"\n\n\n**What are the models trained on?**\n\nThe models are trained on a mix of real text sampled from common data sources (e.g. Common Crawl) with tensor packing. Overall, it does not matter since the input to the transformer block after the embeddings are just vectors, and the input semantics do not change the computational efficiency.\n\n**Generalization to other devices**\n\nWe want to highlight that we do not include just a single model or system setup: We benchmark 13B, 30B, and 70B models with varying numbers of GPUs; our results hold across these variations. Therefore, we have already shown generalization to a certain degree. We also believe that our findings can be extrapolated to clusters with similar accelerators and VRAM sizes (e.g. 80GB devices), such as clusters based on H100s with Infiniband interconnect, which is a common setup. We believe our findings have a large target audience but acknowledge that different hardware characteristics, such as the Nvidia GH200 and its coherent memory model, can lead to different conclusions (e.g. making offloading more beneficial).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The careful experimental design in combination with the sheer scale of the experimentation. Most researchers and practitioners do not have access to the resources necessary to experiment extensively with different LLM training recipes, if they are even able to train LLMs at all. For those do have access to and firsthand experience with the large-scale systems used to train LLMs, it may be disincentivized or discouraged by their organizations to share insights gained and intuitions developed about the strategies that work for training efficiently. It is therefore extremely valuable, in my opinion, to have a published paper that explicitly states the empirical observations that form today's \"common sense\" of LLM training.\n2. The authors provide concrete recommendations for improving MFU, backed up by their observations\n3. The authors take care to also explicitly report \"failed\" training runs.", "reasons_to_reject": "Some minor concerns and questions:\n\n1. What exactly are the sequences that the model is trained on? Are they \"real\" text? Are they uniformly sampled tokens? If not \"real\" text, could this have affected the results at all? \n2. All experiments reported use a single type of GPU. While this GPU is commonly used in practice today for distributed LLM training, it leaves open the question of how hardware-dependent the results of this study are. Is there any reason to believe that the conclusions drawn would generalize or differ given a different kind of hardware?", "questions_to_authors": "See above", "extra_scores": {}}, {"review_id": "w7nhrkL9tS", "reviewer": "Reviewer_WJx2", "summary": "This paper presents a comprehensive review of the different distributed training optimizations for LLMs. The authors consider 32 NVIDIA DGX A100 nodes and measure techniques including data/model/tensor/3d/pipeline/sequence parallelism, activation check-pointing, fusing kernels and flash attention. The authors ablate along different axis for LLaMa, MPT, and Megatron (Deepspeed/LM) models.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper and providing feedback on verb tenses and anonymization. We will carefully proofread the paper to ensure that the usage of tenses is consistent. We are glad you found our analysis of optimization techniques exciting.\n\n**Deeper Insights**\n\nWe disagree that our work lacks deeper insights. We do not propose a new model or training method, however, we design a careful experimental evaluation of current LLM training techniques. We believe our results are far from trivial: e.g. the finding that the best possible configurations do not use activation checkpointing and that increasing model parallelization is preferable over activation checkpointing. Our results using the FlashAttention RMSNorm kernel are also non-trivial, i.e. a 14% point improvement in MFU by using the kernel for Llama 13B is not immediately obvious. We will rework individual discussions of experiments to emphasize the key findings further. \n\n**Evolving hardware**\n\nFirstly, we want to highlight that we carefully chose the Llama-style model architecture and A100 GPU nodes as our setting: These are very widely used, giving our findings a large target audience (and the currently deployed clusters will remain in use for quite some more time). Also, we believe that our findings can be extrapolated to clusters with similar accelerators and VRAM sizes (e.g. 80GB devices), such as clusters based on the H100 with Infiniband interconnect. We stress our findings’ very large target audience. Substantially different hardware characteristics, such as the Nvidia GH200 and its coherent memory model, may lead to different conclusions (e.g. making offloading more beneficial). Nevertheless, our findings could be a starting point for efficiency research on such hardware and can serve as a point of reference for comparative studies.\nWe will add more discussions of the expected generalization of our findings to different hardware accelerators.\n\n**Results of Table 2**\n\nPlease note that Table 2 is not meant to be a full grid-search (see also the “Note on Comparability” accompanying the table). We compare efficiency numbers obtained using our recommendations with our framework to other publicly available efficiency numbers from prior work (which possibly use different model architecture variations of the Transformer, global batch sizes, numbers of GPUs, etc.).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Here are the following reasons why I would be excited to accept this paper\n\n+ Early work and comprehensive analysis of tying all the optimizations together", "reasons_to_reject": "Here are the following factors, I would like to see more clarity on:\n\n- I often find, many an experiments to be of form, \"we tried X, and observed Y\". While a great first step, I'd have liked to see some deeper insights. There is an attempt towards this in the conclusion, however, the rest of the paper reads as a log of experiments. I'd be happy to revise my score, if the authors can present technical insights from their observations. \n- I expect many of the observations to be obsolete sooner rather than later as hardware evolve. One way to sidestep this is to consider what are the trends, and how will that impact what we observe. So, as a community we want to accept papers that present generalizations beyond today's hardware. I urge the authors to present their work factoring in evolving hardware.", "questions_to_authors": "Refer above, but here are some minor nits: \n\n1. The comprehension of the paper can be vastly improved if the authors stick to a given tense. In some parts of the experiment, I saw wavering between tenses. For e.g., \"We experimentally verify .. Additionally, we compared\" \n2. \"We publish the full data of our sweeps on GitHub at <anonymized>.\" For future reference, https://anonymous.4open.science/ is one way to share anonymous data. \n3. Table 2 is not a full grid-search (I tried looking in the Appendix but its quite possible I got lost), is there any reason why not? Like what happens if we train Megatron-LM 18B with Batch size 512?", "extra_scores": {}}, {"review_id": "uvCX8pLsrN", "reviewer": "Reviewer_dncv", "summary": "The paper presents an ablation study of possible configurations of 3d parallelism used for large language model training (LLM). They also study the effect of activation checkpointing and RMSNorm kernel on training of llama models and finally presents some recommendations for efficient training.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. We appreciate your recognition of our work on ablation studies of 3D parallelism configurations in large language model training. We are glad that you found our exploration of activation checkpointing, RMSNorm kernel effects, and our recommendations for efficient training valuable. \n\n**Comparison with other frameworks**\n\nWe stress that the point of Table 2 is not a one-to-one comparison (which we carefully conduct in all previous sections) but rather to contextualize the MFU numbers achieved in our experiments with regard to other publicly reported numbers.\n\n**Running for more steps**\n\nWe observed consistent times after the first step and therefore did not find it necessary to run more steps. \n\n**More sequence lengths**\n\nWe already include the 8k sequence length dimension for most experiments to show the impact of longer sequence lengths. Ideally, we would have enough compute to evaluate many more interactions. Please understand that we already spent considerable computational resources on our controlled experimental design evaluating many different setups, and we have to compromise at certain points. We chose our evaluation dimensions to maximize impact for a large target audience and we do believe that our experiments allow for insight into the impact of sequence length.\n\n\n**Selective Activation Checkpointing**\n\nWe agree that selective activation checkpointing is interesting to evaluate in future iterations of this work. At the time of starting our experimentations, most common frameworks did not yet implement this.\n\n**Tensor vs. pipeline parallelism**\n\nOur findings differ from Narayanan et al.'s due to the inclusion of recent optimizations like FlashAttention-2 and the RMS norm kernel, which significantly impact the efficiency of parallelism configurations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper is well written and performs ablation on different hyperparameter of 3d parallelism training of LLMs.\n- The authors present their recommendation that can be helpful for anyone using 3d parallelism for efficient training of LLMs.", "reasons_to_reject": "- The comparison with other frameworks in Section 4.6 should include experiments by the authors, instead of relying on the reported numbers to control for different factors that can impact the training speed e.g. hyperparameters, GPU connectivity, health of the machines etc as discussed later in the same section.\n- The experiments should run for more steps (e.g. 100) to collect more reliable statistics instead of current 10 steps, because these values can vary a lot.\n- The experiments should include 4k sequence length as well to better understand the impact of increasing sequence length.\n- Other training framework e.g. NeMO-Megatron from NVIDIA already supports selective activation checkpointing to better utilize GPU and memory speed up training. This is missing in the current experiment.\n- The authors show the trade-off between between tensor and pipeline parallelism in Section 4.4 stating - \" our results favor\npipeline parallelism over tensor parallelism.\" However, there is discussion on why they observe this.", "questions_to_authors": "- Please refer to \"Reasons To Reject\" section.", "extra_scores": {}}, {"review_id": "jH6MLBmMnC", "reviewer": "Reviewer_UsrP", "summary": "This is a detailed study on efficient parallelization layouts for training large-scale language models. The authors perform an exhaustive sweep over various training configurations, including tensor parallelism, pipeline parallelism, data parallelism, micro-batch sizes, activation checkpointing, sequence parallelism and various fast attention kernels. They do so on models of three sizes 13B, 30B and 65B. And also different sequence lengths. \n\nThe findings can be summarized below -\n\n1. FlashAttention-2 and the RMSNorm provide significant boost in training efficiency.\n2. Interestingly, they find that not using activation checkpointing and thus using smaller batch sizes or higher model parallelism achieves the best training throughputs. However for larger sequence length, activation checkpoint is necessary\n3. Micro-batch size of 1 enables the better training by reducing the degree of model parallelization, avoiding activation checkpointing, and reducing pipeline bubbles.\n4.  Pipeline parallelism is better than tensor parallelism \n5. Sequence parallelism improves training efficiency for models exceeding 30B parameters and 2k sequence length.\n6. The authors achieve state-of-the-art training efficiency results, with a Model FLOPs utilization of 70.5% for a LLAMA 13B model.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review for recognizing our comprehensive analysis, the incorporation of recent advancements like FlashAttention, and the practical recommendations for efficient training.\n\n\n**Hardware Generalization**\n\nWe believe that our findings can be extrapolated to clusters with similar accelerators and VRAM sizes (e.g. 80GB devices), such as clusters based on H100s with Infiniband interconnect, which is a common setup. We believe our findings have a large target audience, but acknowledge that different hardware characteristics, such as the Nvidia GH200 and its coherent memory model, may lead to different conclusions (e.g. making offloading more beneficial). Nevertheless, our findings could be a starting point for efficiency research on such hardware.\n\n**Model Generalization**\n\nFirstly, we chose the Llama-style model architecture to maximize the impact of our findings for a large target audience.\nWe also want to highlight that we do not include just a single model or system setup: We benchmark 13B, 30B, and 70B models with varying numbers of GPUs; our results hold across these variations. Therefore, we have already shown generalization to a certain degree. \nWe do expect our findings to generalize across many of the common slight variations of the Transformer architecture that do not change the fundamental nature of the model’s compute/memory usage profile. Whether our findings generalize to entirely different model classes with different memory/compute usage profiles (such as e.g. Mamba) is not clear. We will add this point to the paper.\n\n**ZeRO / FSDP**\n\nWe use ZeRO-1, which follows a setup employed by many other researchers using 3D-Parallelism [1,2]. However, indeed ZeRO-2/3 might allow for more efficient layouts due to the saved memory. We will add a discussion of this point to the paper.\nWe did not include FSDP, as it is a different style of sharded training and not as widely used for very large models.\n\n**Contradiction to Narayanan**\n\nOur findings differ from Narayanan et al.'s due to the inclusion of recent optimizations like FlashAttention-2 and the RMS norm kernel, which significantly impact the efficiency of parallelism configurations.\n\n[1] https://arxiv.org/abs/2402.17834\n[2] https://arxiv.org/abs/2401.02954", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Comprehensive study: The authors conduct a thorough analysis of various training configurations, considering complex interactions between different parallelization techniques and optimizations.\n\n2. Incorporation of some recent advancements like FLASHATTENTION and sequence parallelism adds novelty \n\n3. The paper distills the findings into several actionable recommendations for efficient large-scale language model training.\n\n4. 70.5% MFU for 13B during training is an impressive feat and shows the value of this study", "reasons_to_reject": "1. Experiments are only conducted on A100 and results are not associated with any numerical artifacts. This likely makes the study less generalizable to other HW and need to redo it on newer generations of GPU\n\n2. The study focuses primarily on the LLAMA model architecture, and the recommendations may not directly apply to other language model architectures or domains.\n\n3. Optimization techniques does not include more latest optimizations like FSDP and Zero-Optimizer", "questions_to_authors": "1. Do you expect your findings to generalize to other language model architectures beyond the LLAMA model?\n2. Any reason why your findings on tensor vs pipeline parallelism will contradict to  Narayanan et al's finding?", "extra_scores": {}}]

Author rebuttal:

Meta-review: Quality
- Pro
  - general thorough experimental results
  - comprehensive analysis
- Con
  - experiments were run with too few steps for robust experimental validation

Clarity
- Pro
  - well-written easy to understand

Originality
- Pro
  - incorporates some combinations of common training procedures, the originality mostly stems from the broad analysis which is novel
- Con
  - some common ways of training LLMs are not analyzed

Significance
- Pro
  - very important results for anyone involved in LLM pertaining
  - large-scale empirical study that yields valuable insights for the community and further study
- Con
  - some limitations with respect to hardware and model choice can make this work obsolete with time

**Decision**

All reviewers agree that the breadth and depth of experimentation and analysis yields results that are highly useful for the pertaining sub-community. In particular, the rigor and details are highlighted. The paper has some faults though. Some newer pertaining techniques are not explored and some ablations and the length of experiments could improve the robustness and usefulness of the results. Overall, reviewers agree that this work is worthy of acceptance and will prove useful for the community.
Decision: Accept