Efficient Parallelization Layouts for Large-Scale Distributed Model Training

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. The findings can be summarized below - 1. FlashAttention-2 and the RMSNorm provide significant boost in training efficiency. 2. 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 3. Micro-batch size of 1 enables the better training by reducing the degree of model parallelization, avoiding activation checkpointing, and reducing pipeline bubbles. 4. Pipeline parallelism is better than tensor parallelism 5. Sequence parallelism improves training efficiency for models exceeding 30B parameters and 2k sequence length. 6. The authors achieve state-of-the-art training efficiency results, with a Model FLOPs utilization of 70.5% for a LLAMA 13B model.

Rating

7

Confidence

3

Ethics flag

1

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. 2. Incorporation of some recent advancements like FLASHATTENTION and sequence parallelism adds novelty 3. The paper distills the findings into several actionable recommendations for efficient large-scale language model training. 4. 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 2. The study focuses primarily on the LLAMA model architecture, and the recommendations may not directly apply to other language model architectures or domains. 3. 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? 2. Any reason why your findings on tensor vs pipeline parallelism will contradict to Narayanan et al's finding?

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.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

- The paper is well written and performs ablation on different hyperparameter of 3d parallelism training of LLMs. - 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. - 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. - The experiments should include 4k sequence length as well to better understand the impact of increasing sequence length. - 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. - The authors show the trade-off between between tensor and pipeline parallelism in Section 4.4 stating - " our results favor pipeline parallelism over tensor parallelism." However, there is discussion on why they observe this.

Questions to authors

- Please refer to "Reasons To Reject" section.

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.

Rating

5

Confidence

4

Ethics flag

1

Reasons to accept

Here are the following reasons why I would be excited to accept this paper + 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: - 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. - 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: 1. 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" 2. "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. 3. 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?

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.

Rating

7

Confidence

3

Ethics flag

1

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. 2. The authors provide concrete recommendations for improving MFU, backed up by their observations 3. The authors take care to also explicitly report "failed" training runs.

Reasons to reject

Some minor concerns and questions: 1. 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? 2. 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

Thank you for the feedback and answering the question

Thank you for the feedback and answering the question. I would encourage adding some of the answers to the questions to the paper for better clarity. I will still keep the score the same. Improving the score further would not be possible as I am not an expert in this domain. Best of luck!