Efficient Parallelization Layouts for Large-Scale Distributed Model Training

Thank you for the review for recognizing our comprehensive analysis, the incorporation of recent advancements like FlashAttention, and the practical recommendations for efficient training.

Hardware Generalization

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 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.

Model Generalization

Firstly, we chose the Llama-style model architecture to maximize the impact of our findings for a large target audience. We 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. We 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.

ZeRO / FSDP

We 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. We did not include FSDP, as it is a different style of sharded training and not as widely used for very large models.

Contradiction to Narayanan

Our 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.

[1] https://arxiv.org/abs/2402.17834 [2] https://arxiv.org/abs/2401.02954

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.

Comparison with other frameworks

We 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.

Running for more steps

We observed consistent times after the first step and therefore did not find it necessary to run more steps.

More sequence lengths

We 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.

Selective Activation Checkpointing

We 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.

Tensor vs. pipeline parallelism

Our 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.

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.

Deeper Insights

We 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.

Evolving hardware

Firstly, 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. We will add more discussions of the expected generalization of our findings to different hardware accelerators.

Results of Table 2

Please 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.).

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."

What are the models trained on?

The 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.

Generalization to other devices

We 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).