MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining

Summary

This work proposes RapidBERT to train BERT in a faster way. Different from the previous accelerated method, this work attempts to employ some recent popular transformer architecture efficient designs as the basic modification of RapidBERT architecture, such revisions including the introduction of FlashAttention, AliBi, Bf16, GLU, Low Precision LayerNorm, etc. The training dataset is C4. It cost 1.13 hours on 8 A100-80G GPU. Results on the GLUE NLU benchmark show that under fair comparison except for the training step, RapidBERT can get efficient training FLOPs compared to previous methods.

Strengths

- This work employs some recently released transformer-efficient architectures in BERT training, getting RapidBERT. - The RapidBERT got similar NLU ability compared to the vanilla-BERT.

Weaknesses

- For a fair comparison, The authors should also post the RapidBERT results pre-trained on the English Wikipedia and the Books Corpus. (The original BERT study trained on English Wikipedia and the Books Corpus). - Similar studies have been proposed in the same faster BERT pertaining track, it would be better to have a comprehensive comparison with them in detail to help readers learn more about this field. - In line 176, this work says that "we modify the vocab size from 30522 to 30528" inspired by MEGATRON so that the vocab size is a multiple of 64, and leads to a non-trivial throughput speedup. However, in MEGATRON's work, they pad the vocab because of the 8-way model parallelism. But this work does not employ any tensor parallelism or pipeline parallelism, could you explain why such an operation can lead to a speedup in the training of RapidBERT? - Most of the new techs (eg. FlashAttention) tried in this work have already been employed and integrated into the BERT pre-trained stage in some frameworks, e.g., Megatron, DeepSpeed [1][2]. In this manner, this work is not so innovative and may not be the first work to do this, it would be better to discuss differences with these similar works. (Actually, it is inspiring to see more works on efficient LM pertaining field, so this term is not very influential, I just brought it up.) [1] https://github.com/NVIDIA/Megatron-LM [2] https://github.com/microsoft/DeepSpeed

Questions

Except for the adding modules (eg. FlashAttention, AliBi, etc.), there are also some recently released efficient transformer architectures (eg, RoPE, DeepNorm, etc), have the authors try them in RapidBERT. Could such methods help accelerate the training period further? It would be interesting to see if these methods could help.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

Similar studies have been proposed in the same faster BERT pertaining track, this work does not have too much comparison with them in detail.

Summary

The paper benchmarks several architectural changes to BERT that allows for more efficient pretraining. More specifically, the paper adds flash attention, ALiBi position representations, and GLU activations to the original BERT architectures. For fair comparison, they re-implement the baseline BERT-base using the same pretraining hardwares. They show that with these architecture modifications they can achieve better performance on downstream tasks while using smaller amounts of pretraining time. Their base model’s performance is similar to BERT-base in (Devlin et al., 2018) using approximately 1 hour of pretraining on 8 A100 GPUs, costing $20 on a standard cloud provider. They also benchmark the training cost of each architecture modification, showing that reducing the sizes of GLU matrices help to improve throughput of training significantly

Strengths

The paper provides empirical evidence that combining flash attention, ALiBi position representations, and GLU activations reduces 50% of the training cost. Their ablation on GLU will benefit researchers working on the relevant areas

Weaknesses

The contribution of this paper is limited. The proposed architecture modifications are all adopted from prior work. Their primary findings mostly come from tuning the hyperparameters of GLU. And their current findings depend on multiple architecture modifications, which is complicated and makes their scientific findings unclear. I’d hope that the authors can provide more experiments showing why all these modifications are necessary and whether it’s possible to simplify it further for a clearer takeaway message

Questions

1. What is “Pareto Optimal”? This term is mentioned multiple times but never formally defined in the paper 2. There appears to be a throughput difference in Figure 5 between “lpLN+GLU+ALBi” and “GLU 3072”. What leads to the difference?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have discussed the limitations of their work.

Summary

The paper introduces RapidBERT, an architecture and training paradigm for pretraining BERT-style language models that is cost-effective. The proposed approach incorporates several modifications into the conventional transformer encoder block, including FlashAttention, Attention with Linear Biases, Gated Linear Units, Unpadding Module, and a low precision LayerNorm. The pretraining process avoids the Next Sentence Prediction task and follows the RoBERTa practices, using the C4 dataset with a 30% masking ratio in MLM. RapidBERT demonstrates faster convergence and achieves a better accuracy versus time Pareto curve during pretraining. Their base model consistently outperforms the original BERT model on average across the GLUE dev tasks. Remarkably, the training process only takes 1.13 hours using 8 A100 GPUs, resulting in a cost of approximately $20. The authors show that RapidBERT is Pareto-optimal when compared to BERT for both base and large models. They also highlight the necessity of extensive training for larger models, as RapidBERT-Base outperforms both BERT-Large and RapidBERT-Large for a significant portion of their training. In addition, the paper includes a comprehensive ablation analysis of the design choices made in RapidBERT and evaluates the throughput of each architecture. The results indicate that the GLU modification leads to a decrease in throughput, while ALiBi has minimal impact on throughput. On the other hand, incorporating low precision LayerNorm significantly improves throughput. Overall, the paper presents RapidBERT as an efficient and effective approach for pretraining BERT-style language models, showcasing its superior performance, cost-effectiveness, and the benefits of the proposed modifications.

Strengths

- The paper introduces a time and compute-efficient approach for pretraining BERT-like models, offering significant advantages over previous works. This approach holds promise for pretraining task-specific language models that do not require high parameter requirements, making it highly practical and valuable in various applications. - The implementation details provided in the paper are clear and comprehensive, contributing to the reproducibility of the research. The authors effectively explain the concepts and methodologies, ensuring a thorough understanding of the proposed approach. - The research conducted on related work demonstrates a comprehensive exploration of the existing literature. The design choices made in this study are grounded in well-established prior works, and the ablation study further strengthens the validity and effectiveness of the proposed modifications. - The paper presents an intriguing finding regarding the impact of model size on performance in specific domains. The authors discover that larger models may not always yield superior results due to limited data availability and increased compute time. This insight, demonstrated by the Pareto optimality of RapidBERT-Base over BERT-Large and RapidBERT-Large for a significant portion of the training, adds a valuable contribution to the understanding of model performance and scalability.

Weaknesses

~- The approach primarily combines existing techniques, such as FlashAttention, GLUs, and low precision LayerNorm. While the authors introduce the novel aspect of maintaining a 30% masking ratio, the overall novelty of the architectural choices is limited.~ ~- It would have been beneficial to include a discussion comparing the proposed approach with the original RoBERTa model as a baseline. This comparison would provide a clearer understanding of the improvements achieved by RapidBERT and highlight its unique contributions.~ - The paper lacks a discussion on how the architectural choices made in RapidBERT could facilitate the development of larger language models, such as GPT 3.5. Exploring the potential scalability and benefits of these choices for larger models would enhance the practical relevance and implications of the research. - The ablation studies do not investigate the effects of changing the masking ratio. Including an analysis of different masking ratios would offer insights into the choice behind keeping the ratio to be 30%. - The paper would benefit from discussions on the comparison with RoBERTa, the scalability of the approach to larger models, and the effects of different masking ratios in the ablation studies.

Questions

- Can you provide a rationale for considering BERT as a "strong" baseline in your study? Why was RoBERTa not chosen as the baseline for comparison, and what factors led to the selection of BERT as the benchmark model? ~- It would be helpful to gain insights into the absence of numerical instabilities in bfloat16. Could you provide any intuition or theoretical reasoning behind why it doesn't happen in your case?~

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

The discussion on limitations is satisfactory. There are some limitations discussed in the appendix regarding potential training stability issues on the larger models. Potential negative societal impacts (ease of pre-training could lead to ease of more biased/inappropriate models) is also discussed in broader impacts section. General limitations with any pre-trained LMs will apply to RapidBERT.

Summary

This paper proposes a new efficient recipe for training BERT, matching the original performance of BERT on GLUE in ~1h on 8x A100. To do so, the authors leverage a number of architectural/implementation improvements: FlashAttention, ALiBi, GeGLU, `bf16` layer norm, unpadding, and tweaks to the masking ratio. The authors find that their recipe is optimal even when considering longer runs, and that it transfers well to BERT-Large.

Strengths

* **S1.** Reducing the costs associated with training language models can help practitioners iterate faster, improving downstream research outcomes. This makes this paper potentially valuable to the community, as it improves upon previously introduced similar recipe such as CrammingBERT. * **S2.** The authors detail their contributions and open-source their code, making these results reproducible and enabling the community to build upon them. * **S3.** The authors feature an updated baseline (BERT-Base) that helps for fair comparisons in their setup.

Weaknesses

* **W1. It is difficult to untangle individual contributions to the final result.** * **W1.1.** The proposed recipe visibly has some positive impact on GLUE score, as Figure 3 shows that it achieves significantly better performance than BERT-Base after the same amount of training. However, that impact is never quantified in the paper. Here, the ablations should not only focus on throughput, but also on how the proposed interventions might impact the GLUE score. * **W1.2.** The proposed baseline is very strong (which is a positive point), and it would be good in Table 1 to also showcase the time it take for it to reach a 79.6 GLUE score. Furthermore it would be interesting to identify what makes the baseline such a strong one (with a final score much higher than BERT-Base). Is it the change in data to C4 (which authors acknowledge as an important factor l161)? Something else? * **W1.3.** l112 the authors discuss using ALiBi to pretrain with a shorter sequence length and extrapolate at test time. It's unclear if this end up being used, and if it is included in the ALiBi ablations. * **W1.4.** l251 it is disappointing to not ablate adequately every component, especially since the value of this paper lies in having a potentially systematic approach to performance improvements of BERT models. * **W2. The Pareto frontiers described may be slightly misleading, as they do not account for LR schedule.** The so-called Pareto optimality is obtained by taking points from the same run, instead of having one run per pretraining budget on the plot. This approximation penalises intermediate budgets, as they are evaluated with an incomplete LR schedule. While I don't think this has a significant influence in this work, since the authors discuss Pareto optimality so much, this should at the very least be clearly discussed as a limitation to avoid misleading other authors. This issue is particularly relevant, as it has lead to significant misunderstandings around scaling laws for instance (see Hoffmann et al., 2022). * **W3.** The paper feels a bit repetitive, as if it had been stretched to fit the 9 pages of content. Section 3.2 and 3.3 are particularly egregious in this regard, and more time could instead be spent on ablations. * **W4.** (minor nits) l98 reference to Triton should cite "Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations" (Tillet et al., 2019); l145 the sentence "results from NVIDIA and others" is confusing, as the final work cited is not from NVIDIA -- there should be a citation somewhere for the NVIDIA results.

Questions

This has the potential to be a valuable paper to the community, by systematically identifying and detailing interventions that can accelerate the training of large language models like BERT. Reducing iteration costs is in particular a great enabler for researchers. Unfortunately, the lack of systematic ablations on GLUE score and the limited throughput ablations make this paper fall somewhat short of its promise. Accordingly, I am rating it as a **Borderline Reject (4)** but would be willing to increase my score to an accept should some of my concerns be addressed. **EDIT: following rebuttal, I have updated my score to a Weak Accept (6).** * **Q1.** Could the authors explain why Figure 3 shows a significant improvement in GLUE score for RapidBERT over BERT? Would it be possible to ablate for this improvement, to ultimately better understand which intervention improves throughput and which improves "modelling performance" (i.e., GLUE score)? * **Q2.** Could the authors explain why the baseline used is so much above the original BERT-Base? * **Q3.** Could the authors clarify whether they use for RapidBERT a shorter training seqlen to accelerate training as proposed? * **Q4.** Could the authors provide further performance ablations for FlashAttention/unpadding? * **Q5.** (more of a suggestion) Could the authors better describe the limitation of their approach regarding Pareto-optimality? Small suggestions: * l13/14 in the abstract "When pretrained from scratch on the C4 dataset, this base model achieves the downstream average GLUE score of 79.6 in 1.13 hours on 8 A100 80 GB GPUs at a cost of roughly USD 20" it would be good to mention this is the GLUE score achieved by the original BERT; for instance "RapidBERT achieves the same downstream average GLUE score as the original BERT (79.6) in 1.13 hours on 8 A100 80 GB GPUs at a cost of roughly USD 20".

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

While there is no explicit section, some limitations in terms of scope are discussed in the conclusion.

Summary

This paper presents a training recipe that can train a BERT-style encoder model efficiently (1.13 hours on 8 GPUs). The recipe combines several techniques including a higher masking rate for MLM, bf16, optimized vocabulary size. The model is trained on the C4 dataset for 1.13 hours and achieves 79.6 on GLUE. The paper conducts ablation studies to characterize properties of the architecture design choices.

Strengths

* The paper presents very impressive results of pre-training a BERT-like model very efficiently without loss in accuracy. In general, I believe that this is a solid outcome and it is definitely a good addition to the community, as it can enable more researchers to pretrain custom BERT models from scratch. * The paper presents ablation experiments that quantify the impact of each proposed change. * The paper is well written and easy to follow. Replicating the results should not be hard.

Weaknesses

* The experiments are only conducted on GLUE tasks; it is not clear how well the trained model transfers to other datasets. * There is no fair comparison (i.e., under the same hardware setup) of the proposed recipe and other efficient training methods such as CrammingBERT.

Questions

* The paper uses a 30% MLM masking ratio for RapidBERT instead of 15% for BERT. How do you compare the improvement from this change to the gains from your architectural modifications?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

I didn’t see a clearly potential negative societal impact of this paper.

Answer to the rebuttal

First, I would like to thank the authors for providing an extensive rebuttal to each reviewer, as well as additional results. I believe that this additional information+discussions significantly improve the quality and value of the paper. I will trust that the authors update their paper based on the rebuttals they have provided. Accordingly, I have updated my score to a **Weak Accept (6)**. Re: as a small clarification on W2, I agree that this does not impact the results presented -- I just pointed to the need for mentioning this explicitly (that intermediary results are penalised), as I believe this is an important point of confusion in the community.