MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining
Authors: Jacob Portes
Alexander R Trott
Sam Havens
DANIEL KING
Abhinav Venigalla
Moin Nadeem
Nikhil Sardana
Daya Khudia
Jonathan Frankle
Venue: NeurIPS 2023 poster (2023)
Keywords: BERT
Pretraining
Efficiency
FlashAttention
ALiBi
TL;DR: MosaicBERT is a BERT-style encoder architecture and training recipe optimized for fast pretraining

Abstract: Although BERT-style encoder models are heavily used in NLP research, many researchers do not pretrain their own BERTs from scratch due to the high cost of training. In the past half-decade since BERT first rose to prominence, many advances have been made with other transformer architectures and training configurations that have yet to be systematically incorporated into BERT. Here, we introduce MosaicBERT, a BERT-style encoder architecture and training recipe that is empirically optimized for fast pretraining. This efficient architecture incorporates FlashAttention, Attention with Linear Biases (ALiBi), Gated Linear Units (GLU), a module to dynamically remove padded tokens, and low precision LayerNorm into the classic transformer encoder block. The training recipe includes a 30% masking ratio for the Masked Language Modeling (MLM) objective, bfloat16 precision, and vocabulary size optimized for GPU throughput, in addition to best-practices from RoBERTa and other encoder models. When pretrained from scratch on the C4 dataset, this base model achieves a downstream average GLUE (dev) score of 79.6 in 1.13 hours on 8 A100 80 GB GPUs at a cost of roughly $20. We plot extensive accuracy vs. pretraining speed Pareto curves and show that MosaicBERT base and large are consistently Pareto optimal when compared to a competitive BERT base and large. This empirical speed up in pretraining enables researchers and engineers to pretrain custom BERT-style models at low cost instead of finetune on existing generic models. We open source our model weights and code.

MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining
Jacob Portes jacob. portes@
Alex Trott alex. trott@
Sam Havens' sam hanehaeese@
Daniel King' daniel. daniel.kingo
Abhinav Venigalla' ! abhi@
Moin Nadeem mo moinnadeem@
Nikhil Sardana nikhil. sardana@
Daya Khudia! daya. khudia@
Jonathan Frankle
jfrankle@
MMosacMLL Databricks
@databri iks. com 2 @m wwwiiinnndemmmm com
com
Abstract
Although BERT-style encoder models are heavily used in NLP research, many researchers do not pretrain their own BERTs from scratch due to the high cost training. the past half-decade since BERT first rose to prominence, many advances have been made with other transformer architectures and training con- figurations that have yet to be systematically incorporated into BERT. Here, we introduce MosaicBERT, a BERT-style encoder architecture and training recipe that is empirically optimized for fast pretraining. This efficient architecture incorporates FlashAttention, Attention with Linear Biases (ALiBi), Gated Linear Units (GLU), module dynamically remove padded tokens, and low precision LayerNorm into the classic transformer encoder block. The training recipe includes a 30% masking ratio for the Masked Language Modeling (MLM) objective, bfloat 16 precision, and vocabulary size optimized for GPU throughput, in addition to best-practices from RoBERTa and other encoder models. When pretrained from scratch on the C4 dataset, this base model achieves downstream average GLUE (dev) score of 79.6 in 1.13 hours on A 100 80 GB GPUs at a cost roughly $20. We plot extensive accuracy pretraining speed Pareto curves and show that MosaicBERT base and large are consistently Pareto optimal when compared to a competitive BERT base and large. This empirical speed up pretraining enables researchers and engineers to pretrain custom BERT-style models at low cost instead of finetune on existing generic models. We open source model weights and code.
Introduction
BERT has been the workhorse of modern natural language processing (NLP) since its introduction in 2018 [14.. Even in the era of large language models (LLMss,, BERT-style encoder models are still quite relevant; for example. encoder models are used for vector database embeddings and retrieval augmented generation in tandem with L.LMs 30 351
52 In the past half-decade since BERT first rose to prominence, however, many advances have been made with other transformer architectures and training configurations that have yet to be systematically incorporated into BERT 12 44 I1. In this study we empirically show that these speed optimizations can successfully be incorporated into the classic BERT architecture and training recipe.
"qqual contribution. Corresponding author jacob. wwwtesttlttabriccs.comm com. Code can be found at mosaicbert. github.io
wook done while at MosaicML.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

A
B
GLUE Pareto Curves
Gated Linear Unit
(GeLU[[W.. x
MosaicBERT
0.84
BERF (Devin 2015)
2018)
Encoder
Unpadding
Low Frecktinn LoyerNorm
Average GLUE (dev) Score
0.82
BERT-Sane (Devlis 2018
D MALLMRY
Gated Linear Unit
0.80
FlashAttention (IO-aware)
Iow Precitoon Layeniorm
Dropest
Unpodding
0.76
FlashAttention
Goned
ALiBI
BERT- BERT
MosaicBERT Base
ALIBi (i.e. Linear Blas)
Tokonization Embedaing
10
noftmor(QK B]
Pretraining Wall Clock Time (hours)
Figure (A) Schematic Mosa architecture (B) Pareto curves of average GLUE (dev) scores for Mosaic and the standard BERT-Base. Error bars indicate 95% confidence interval over n=5 pretraining seeds. All training was on 8xA100-80GB GPUs. FlashAttention schematic adapted from [II), and unpadding schematic adapted from [66]).
BERT-style models are typically trained in two stages: an initial self-supervised pretraining phase that builds general representations of language, and subsequent supervised finetuning phase that uses those representations to address a specific task. The pretraining stage for BERT models has historically been computationally expensive; in the original BERT study, for example, the authors trained their models for full days on 16 Google TPUs. In recent years, however, the time and cost to train BERT models has dropped significantly. One widely cited paper from 2021 successfully reduced the training time of BERT-Large to 24 hours on 8 Titan V-12 GPUs (28) and another very recent paper trained competitive BERT-Bas Base model single consumer GPU in only 24 hours 18. Our work builds on these trends.
In this study, we introduce our optimized MosaicBERT architecture and show that certain architectural choices for BERT-style encoders lead to accuracy vs. time Pareto improvements in pretraining. We do this by empirically comparing MosaicBERT with optimal BERT baseline that does not incorporate our architectural changes but does have non-architectural optimizations such as fast data streaming and optimal floating point precision. We then evaluate the classic GLUE benchmark [60].
can often be difficult to understand the effects of model architecture modifications and training hyperparameter choices by simply reporting few numbers in table, as is traditionally done in ML papers. few numbers in table can obscure differences in the training data, objective function, batch size, training duration, learning rate schedule and many other details. Since our main goal in this study is show how combinations of architecture choices lead to improvements in both accuracy and training time, we make the choice to plot accuracy vs. training time Pareto curves for all models (for example, in Figure BB. Certain architecture changes might lead to an increase in throughput, but a decrease in accuracy (e.g. changes to the floating point precision); other changes might lead an increase in accuracy but take a much longer time to converge (e.g. increasing the model size). Accuracy training time Pareto curves enable to adequately asses all these changes (see Appendix for a more detailed discussion on Pareto curves).
The contributions of this work are as follows:
We implement a new BERT-style encoder architecture that optimizes pretraining speed and accuracy. This architecture combines FlashAttention [II,, ALiBi [44], Gated Linear Units [12] 50] dynamic unpadding module [66), and low precision LayerNorm.
We show that MosaicBERT base achieves the downstream average GLUE (dev) score of 79.6 in 1.13 hours on AxA100 GB GPUs at a cost of roughly $20 on standard cloud provider.

---

We characterize the accuracy vs. pretraining time Pareto frontier for MosaicBERT Base and Large, and empirically show that the performance of MosaicBERT-BBas and Large Pareto optimal relative to BERT-Base and Large.
We characterize the relative throughput properties of cach of MosaicBERT's architecture and design choices via ablations.
We characterize the tradeoff between model size and training duration, and show that BERT- Large performance only surpasses BERT-Bass performance after extensive training. We open-source our model weights and code at mosaicbert. mossaccbert.githuu (ithub.io. io
Finally, we note that the ideas explored in this study are directly applicable to LLM pretraining, and directly motivated many of the architecture choices around MosaicML's MPT-7B and MPT-30B large language models 33 38)
Model Architecture
GLUE
Training Hardware (hours)
Pretraining
Params
(dev)
Corpus
MosaicBEERT-Base (ours)
79.6
1.13
8XA100-80
C4
137M
MosaicBERTT-Base (ours)
82.2
2.81
8 100-80
C4
137M
MosaicBERT-Basee (ours)
83.2
4.6
8xA100-80
C4
137M
MosaicBERTT-Base (ours)
83.4
5.27
8XA100-80
C4
137M
BERT-Base (our benchmark)
83.2
11.5
8XAT00-80
C4
110M
BERT-Base [4
81.0
96
16XTPU
Wiki+BookC
110M
BERT-Base 18
73.6
24
1 RTX A6000
Wiki+BookC
110M
WammingBERT-- Base
18
80.4
24
1 RTX A6000
Wiki+BookC
110M
BERT-Large [14 37
84.1
96
64XTPU
Wiki+BookC 340M
Table Average GLUE (dev) score across various efficient BERT implementations. Average includes
all GLUE tasks. BERT-Base average GLUE scores the dev set as reported by [18), Table 3.
BERT-Large average GLUE scores on the dev set as reported by 37 Table 5. Note that the original training set for BERT was pretrained on 40 epochs of English Wikipedia and BookCorpus [68]. while MosaicBERT was pretrained on the Colossal Cleaned Common Crawl (C4) corpus [46].
Methods
In order to build MosT,cBERT, we incorporated architectural choices from the recent transformer literature. These include FlashAttention [1.. ALiBi [44], training with dynamic unpadding (66),
low precision LayerNorm, and Gated Linear Units [12 50 (Figure D Before describing these
modifications in detail, first review the classic BERT architecture and how chose strong baseline.
Since our goal here is to show relative improvements in training time and final accuracy, we do not attempt to beat state of the art models for finetuning benchmarks such as GLUE [60] These SOTA models are often trained for much longer (e.g. (37]) and are larger than the models we explore in this study 23
2.1 Choosing a Strong BERT Baseline
The basic transformer block used in BERT models consists of (1) the attention mechanism and (2) the feed forward layers. This block is then repeated depending on the model size; BERT-Base has 12 repeated transformer blocks, while BERT-Large has 24.
For our baseline BERT-Base, we used the exact architecture of BERT from 14 this includes hidden size of 768, an intermediate size of 3072. 12 attention heads and 12 hidden layers, as well as the GeLU activation function, and learned positional embeddings. For our baseline BERT-Large, we used the exact architecture of BERT-Large from [14) which has hidden size of 1024. intermediate size of 4096, 16 attention heads, and 24 hidden layers.
'The Hugging Face bert-base bas---aceeeeeddd model has been downloaded 62 million times, more than any other model the Hugging Face hub.

---

While MosaicBERT-Baae (i.e. 12 hidden layers) and Large (i.e. 24 hidden layers) stay true to this general structure, we introduce modifications that affect both the attention mechanism and the feedforward layers.
2.2 MosaicBERT Architecture Modifications for Fast Pretraining
2.2.1 MosaicBERT Modifications to the Atfention Mechanism
The recently proposed FlashAttention layer reduces the number of read/write operations between the GPU HBM (high bandwidth memory, i.e. long-term memory) and the GPU
FlashAttention:
SRAM (i.e. short-term memory) [ (see also 455). We modified the FlashAttention module built by Hazy Research with OpenAl's triton library in order to flexibly incorporate ALiBi [57
Attention with Linear Biases (ALiBi): most BERT models, the positions of tokens sequence are encoded using learned position embeddings, which are added to the learned token embeddings at the start of the forward pass. ALiBi eliminates position embeddings and instead encodes position information directly through the attention operation [44]. It adds negative bias to the attention score between each token pair, which grows linearly with the relative distance between the tokens.
Intuitively, this biases attention to nearby tokens and allows for extrapolation to context lengths longer than those used for training [44.
Following the notation in 44, the attention block computes the attention scores between the ith query q: F2 and keys E IR Axd where d the head dimension and L is the sequence length. ALiBi adds fixed bias with m as a head-specific slope controlling how the bias grows with absolute distance between tokens, yielding attention weights as
softmax m abs((i (1- - -,,1 1.i-2, 2,
LI)).
The slopes m follow a geometric sequence such that for n heads, each head has ratio of 2-8nn
44
During finetuning or inference, the static bias can be increased to accommodated longer sequence lengths. For example, model pretrained using ALiBi with maximum sequence length of 128 tokens can then extrapolate to a task with 256 tokens with little to no decrease in zero-shot performance. This was the original motivation for AliBi (i.e. "train short, test long") 44] Since pretraining model with a maximum sequence length of 128 has much higher throughput than pretraining a model with sequence length of 256, ALiBi can be considered an indirect speedup method.
2.2.2 Modifications to the Feedforward Layers
Gated Linear Units (GLU): We used Gated Linear Units for the feedforward sublayer of a trans- former. GLUs were first proposed in 2016 [12]. and incorporate an extra learnable matrix that "gates" the outputs of the feedforward layer (Figure TA). More recent work has shown that GLUs can improve performance quality in transformers [50 42). We used the GeLU (Gaussian-error Linear Unit f aactvaation function with GLU, which is sometimes referred as GeGLU. The module can be described by the following equation
(GeLU(xW.,) & W2)
(2)
where x is the input to the feedfoward layer and the matrix gates the output of the GeLU activation function. The extra gating matrix in GLU model potentially adds additional parameters model; we chose to augment our MosaicBERT-BBal model with additional parameters due to GLU modules, as it leads to Pareto improvement across all timescales. While BERT-Base has 110 million parameters, MosaicBERTT-HE -ase has 137 million parameters. Note that MosaicBERT-BBan reaches higher accuracy faster than BERT-Base despite having more parameters (Figure (.B). Similarly, BERT-Large has 340 million parameters, and MosaicBERT-Large has 430 million parameters.
github. wwt.com///aaae comm/Haazyeeeearrchh/cooooooacee http///fasaaaatttttiee Note that while this research was being completed, PyTorch 2.0 was released with support for FlashAttention. However, FlashAttention PyTorch 2.0 does not currently support ALiBi integration.
'GeLU is a fully differentiable approximation ReLU, and was used in the original BERT study 14.

---

2.2.3 Additional Modifications
Low Precision LayerNorm: LayerNorm is a bandwidth-boumi operation, which means that its speed depends on how quickly data can be loaded from memory the compute units for element- wise operations. Typically the LayerNorm operation is set to float32 precision, which requires 4-bytes per-element. In MosaicBERT, we modify LayerNorm modules to in bfloat16 precision instead of float32. This reduces the amount of data that needs to be loaded from memory, as only 2-bytes are required per-element. PyTorch's by default, run LayerNorm in lower precision because this can lead to numerical instabilities for certain models. However, our experimental results show that MosaicBERT does not experience any numerical instabilities with bfloat16 precision LayerNorm.
package does not,
Unpadding: Standard NLP practice is to combine text samples of different lengths into batch, and pad the sequences with special padding tokens so that all sample sequence lengths are the same (Figure [A). During training, however, this leads many wasted operations on the padding tokens. In MosaicBERT, we take a different approach and instead concatenate all the examples from a minibatch into single sequence of batch size 1. Results from NVIDIA and others have shown that this approach leads to speed improvements during training, since operations are not performed on padding tokens 666.
MLM Masking Ratio and Dropout:
We used the standard Masked Language Modeling (MLM) pretraining objective. While the original BERT paper also included Next Sentence Prediction (NSP) task in the pretraining objective, subsequent papers have shown this to be unnecessary [37] 28). used the standard 15% masking ratio. However,
For our BERT baselines, masking ratio led to slight accuracy improvements in both pretraining ML. and downstream GLUE performance. We therefore included this simple change as part of our MosaieBERT training recipe. Recent studies have also found that this straightforward change can lead to downstream improvements 63
found that 30%
For the baseline BERT, we applied the standard 0.1 dropout to both the attention and feedforward layers of the transformer block. For MosaicBERT, however, we applied 0.1 dropout to the feedforward layers but did not apply dropout to the FlashAttention module, as this was not possible with the OpenAI triton implementation [57 The removal of dropout in the attention layers also leads to small speed up.
2.3 Pretraining Optimizations for Both MosaicBERT and the BERT baseline
Data: Pretraining data is an important factor when comparing BERT models; while the original BERT study trained on English Wikipedia and BookCorpus (68] subsequent models such as RoBERTa trained on much larger datasets (e.g. (37) trained on 160 GBs of text while [1] only trained on 16GB of text). Here we chose to train all models on the more modern Colossal Cleaned Common Crawl (C4) corpus [46). For all experiments, we used a maximum sequence length of 128 tokens. Note that in our implementation, samples longer than 128 tokens were simply truncated.
Streaming Dataset: As part of our efficiency pipeline, we converted the C4 dataset to the StreamingDataset format and used this for both MosaicBERT-Bass and the bascline BERT- Base. This ensured that our wall clock time measurements were not hindered by data streaming.
Bí Precision: We use bfloat16 mixed precision training for all the models. bfloat16 is custom 16-bit floating point format for machine learning that has one sign bit, eight exponent bits, and seven mantissa bits, and has the dynamic range of float32. For mixed precision training, matrix
NVIDIA removes padding their MLPerff1..I benchmarking training rww.reeuuttta. v1. 1 in the padding- py file
http::///gttttub..om./. commonss/
In "Should You Mask 15% in Masked Language Modeling?" Wettig al. find that constant ML. masking ratios above 15% lead to improved average GLUE and SQuAD scores for bert-base and bart- Similarly, in the recently published paper "Dynamic Masking Rate Schedules for MLM Pretraining." Ankner et al. find that a constant MLM masking ratio of 30% consistently outperforms MLM masking ratio of 15% for bert-base.
github. cttchub.com/aee http:/mossiimm/eeeemmmmmm streaming It is drop in replacement for PyTorch's EterableDatasee that allows for fast streaming from cloud storage.
5

---

MNLLI
QNL
QDP
RTE
0.05
0825
0.00
0.000
Accurscy
Accurscy
2.80
Accrrcc
0.70
0.70
cane
co0n
MosalcDERT-Dase
0.50
MoseicBERT-Basee
Pretraining WCT (hours) COLA
Pretraining WCT (hours)
Pretraining WCT thours) MRPCC
Pretraining WCT (hours)
SST-2
STSB
1,06
2.02
CorCaaCCrCCaut
0.000
MelthensConCon
04-
Acturacy
Acturasy
0.005
03
0373
BERT MOSTTRRRT--866
BERT
BERT
BERT-Base MosairaFIT--Ba
MosairBEiT--aaG
0300
23
12
10
Pretraining WCT (hours)
Pretraining WCT (hours)
Pretraining WCT (hours)
Pretraining WCT (hours)
Figure 2: Performance individual GLUE (dev) finetuning tasks. Our MosaicBERT-Baan consis- tently outperforms BERT-Base on MNLI-m, QNLI, QQP and RTE, and has comparable performance on CoLA, SST-2, MRPC and STSB. Wall clock time is for 8xA100-80C1 GPUs, and does not include finetuning time. Error bars are plotted with 95% confidence interval across n = 5 pretraining seeds, and all models trained for 70,000 steps with batch size 4096.
multiplication layer uses bfloat16 for the multiplication and 32-bit IEER floating point (float32) for gradient accumulation. We found this be more stable than using float16 mixed precision
Vocab Size as a Multiple of 64: We increased the vocab size to be a multiple 30.522 30,528) This small constraint is something established in the MEGATRON work by Shoeybi et al. [53] and leads to a non-trivial throughput speedup, as CUDA more efficient at computing matrix multiplications with these dimensions 12 Across all experiments, use the standard BERT-Baae and Large tokenizers.
(i.e. from
Pretraining Hyperparameters: For all models, we use global batch size of 4096, and microbatch size 128. We set the maximum sequence length during pretraining to 128, and we used the standard embedding dimension of 768. These hyperparameters were the same for MosaicBERT-Baae baseline BERT-Base. More hyperparameter details are included the Appendix.
and the
3 Results
In our first set experiments, we pretrained BERT- Base and MosaicBERT-Bee for 70,000 steps of batch size 4096, which roughly corresponds 78% of English C4 (Figure [ and Figure We experiments with = pretraining seeds for each model class. We then finetuned these models on the GLUE benchmark suite using identical finetuning parameters for all models and experiments; the details can be found in the Appendix.
3.1 MosaicBERT-Band Achieves 79,6 Average GLUE (dev) Score in 1.13 Hours
MosaicBERT -aaSe achieves the downstream average GLUE (dev) score of 79.6 in 1.13 hours on 8x A100 80 GB GPUs at cost of roughly $20 standard cloud provider. More details cost estimates are included in the Appendix.
1he NVIDIA Apex library and Megatron 53 both use but is not documented in any papers that we could find.
precision LayerNorm their code,
" Note that in the original BERT study, the vocabulary size was 30,000 [1] However, the default vocab size for the bert-base- uncased tokenizer in HuggingFace is 30,522.
With vocab size that is divisible by 64, the calculations go down a different kernel path with much higher occupancy.

---

Average GLUE (dev) Score
0.86
BERTI Large (Devin 2018:
0.84
Average GLUE (dev)
BERT-Bate (Devinn 2018)
0.80
BERT-Base 178k steps
0.78
cBERT-Base steps
BERT-Large 178k steps
Mosa cBERI-Large 1/8k steps
0.76
10
20
30
40
50
60
70
Pretraining WCT (hours)
Figure 3: Average GLUE (dev) score Pareto curves for MosaicBERT Base and Large trained for roughly epochs C4 (i.e. 178,000 steps with batch size 4096 with maximum sequence length 128 tokens), MosaicBERT-Bass and Large are Pareto optimal relative to BERT-Base and Large. All pretraining done on
80GB devices (n=2-3 pretraining seeds). Note that BERT-Base and MosaicBERT Base took much less time to train than BERT-Large and MosaicBERT-Large.
The baseline BERT-Bast Base reached average GLUE (dev) score of 83.2% in 11.5 hours (A100-80GB). while Mosa reached the same accuracy roughly 4.6 hours on the same hardware, which roughly a 2.38x speedup (Table T and Figure BB).
MosaicBERT-Base Pareto Optimal
As can seen Figure BB MosaicBERT -ase consistently achieves higher average GLUE accuracy more quickly than the standard BERT-Base across all training durations. The performance of MosaicBERT on individual GLUE finetuning tasks can be seen Figure 2 MosaicBERT-BBal
outperforms BERT-Base in four out of eight GLUE tasks across pretraining durations.
The space NLP benchmarks and tasks has exploded in recent years; we include more information on the individual tasks in the classic GLUE benchmark in the Appendix. MNLI, QNLI and QQP are largest datasets in the GLUE benchmark, with 100k-400k training examples, and MosaicBERT- Base strictly Pareto- Pareto-optimal for these tasks relative to BERT Base (Figure 2). We interpret this to mean that the architectural changes and training recipe we chose resulted in an optimized, efficient BERT-Base model.
The quality of both models on smaller datasets (3k-67k training samples) much more variable, as shown by the large error bars (standard deviation across n=5 pretraining seeds) for SST-2, MRPC and STSB. Regardless of this variation, MosaicBERT Base performs equivalently to BERT-Bass on these tasks across training duration.
3.3 MosaicBERT-Large is Pareto Optimal
While BERT-Bass is one of the most popular BERT models, BERT-Large comes in a close second. therefore curious whether
All of our model development done MosaicBERT BaaacBBRTT-Bse;;; our architecture and pretraining choices also generalized to a larger model.
In second set of experiments, and Large for two epochs of the C4 dataset. The training duration is 178,000 steps with batch size 4096, and is more than twice long the duration of the models in Figures I [nd
pretrained MosaicBERT-Bass and Large as well as BERT-

---

B
0,850
Ablation
0.835
MMLJim (Accurocy)
0.94
ONU (Accuracy)
00P (Acouracy)
ATE (Accu asy)
0.92
0.045
0.830
0.80
0.845
0.01
0.78
0.68
0.84
0.76
Avorgge GUE Score
0.885
0.835
0.82
0.74
0.830
12
12
10
0,830
(MArthwwr's Goaf)
(Accuraty)
NOFC (Accuracy)
Coef)
0.60
0.90
0.91
0.925
0.50
0.00
BERT
0.36
0.92
BCRT+drpt=0
0.54
0.86
BERT+GLU
0.32
BERT FigLN
0.815
mim30
0.50
0.80
MEsaicBENT
0.85
0.810
0.86
0.80
0.64
10.0
10.5 11.0
11.5
12.0
12.5
13.0
Pretram WET (hours)
(hours)
(hours)
(hours)
WCT (hours)
Figure 4: Ablation Experiments (A) Average GLUE score and (B) Individual GLUE tasks. BERT-base: standard BERT-base (110M parameters) with attention dropout=0.1 and feedforward dropout=0.1, vocab size set to 30522, MLM=15% (all Hugging Face standard configurations). <dRT+d standard BERT-base, except that the attention in the dropout layer is set to 0 instead of the default 0. BERT+GLU: standard BERT-base, with GLU for the feedforward component of the encoder block. BERT+IpLN: standard BERT-base, except with low precision LayerNorm (6) BERT+mIm30: standard BERT-base, except with a masked language modeling masking ratio of 30%. MosaicBERT: the complete MosaicBERT-Bas including GLU (where the dimension of the intermediate layer is 3072 resulting in 137M total parameters)) ALiBi, low precision Lay- erNorm, unpadding, MLM 30%, vocab size 30528 (a multiple of 64) and the attention dropout=O. MosaicBERT--Faase (lnrFllaaahAt--drpt-o..... MosaicBERT-Boa. without Flash Attention and with the atten- tion dropout set 0.1.
BERT-Large has 24 repeated transformer layers (BERT-Basc has 12), and as result consumes much more memory and takes longer train. We found that MosaicBERT--arg. reached an average GLUE score of 83.2 2 in 15.85 hours, while BERT-Large took 23.35 hours (Figure 3). MosaicBERTT-Lage therefore had a 1.47 speedup over BERT-Large this training regime.
While Bos is optimal for a constrained budget, the MosaicBERT-Large average GLUE score eventually surpasses MosaicBERT-BEE Base after 25 hours of training on 8 A 100-80GB node, and reaches an average score of 85.5 in roughly 50 hours. In our experiments, the MosaieBERTT-Lare architecture and pretraining was the same as MosaicBERT- heads, number of hidden layers, and intermediate size of the feedforward units.
outside of the number of attention
A striking result from Figures and S2 is that MosaicBERT- Base is Pareto optimal relative to BERT-large this training regime, and is Pareto optimal to MosaicBERT-Large during the first half of training. PoosicBERT-BBE -ase takes only 25 hours complete 2 epochs, while MosaicBERT-LLc. takes close to 70 to complete 2 epochs. potential takeaway from this that MosaicBERT Large only surpasses Base performance in the large data regime. For certain tasks such QQP. SST-2, and MRPC MosaicBERT- -aae achieves a maximum accuracy on par with the maximum accuracy of MosaicBERT--L.. for far fewer pretraining hours. When building encoders for specific domains and tasks, bigger is not always better.
3.4 Ablation Experiments
When deciding on the final setup for MosaicBERT, we tried to make educated guesses about what combination modifications could work well together to both increase FLOP utilization and maintain or improve accuracy. Our approach was inspired by Amdahl's Law, i.e. the idea that optimizing portion of a system responsible for N% of the costs can lead to, at most, an N% speedup [88), and chose each of the modifications such that they did not target the same exact part of the of the stack.
8

---

For example, low precision LayerNorm and GLU affect different parts of the encoder block, and so potentially compose well.
In the following ablation experiments we benchmark the individual effects of various architecture and training choices on pretraining wall clock time. We plot downstream GLUE accuracy as a function of measured pretraining wall clock time in Figure
The patterns in Figure A shed light on the individual effects of various architectures (c.g. BERT+GLU, HEW precision LayerNorm) and training configurations (e.g. BERT + 30% masking ratio). On average, all methods seem to provide a slight accuracy boost to BERT-base. Increasing the masking ratio to 30% leads to a slight accuracy boost while not affecting the WCT, while turning off dropout in the attention layer (dERT+)rpt=O) leads to slight improvement in both accuracy and WCT. Low precision I LayerNorm (ERT)+)pLNN leads to a significant speedup (i.e. a shift to the left). Gated Linear Units (BERT+GLU) add more parameters the model and lead to significant slowdown while providing an accuracy boost. As point of reference, we also henchmark the full BERT as well MosaicBERT without FlashAttention and with attention dropout set to 0.1 (the standard BERT-base configuration).
In these experiments, all BERT-Base models here were pretrained on C4 for 70,000 steps with batch size 4096, and microbatch size 256 on AxA100 80GB GPUs. All models were initialized with the same seed and shared all other hyperparameters including the bert-base uncased tokenizer, the learning rate schedule, Adam W as the optimizer, etc. Final pretraining checkpoints were then finetuned on GLUE following the details in the appendix. The points represented in these GLUE plots are final finetuning checkpoints.
These plots highlight the importance of benchmarking with Pareto curves, as it is not possible to tell from these plots alone whether training BERT-base for 2 more hours leads to better performance than BERT+GLU, for example.
Related Work
Various studies rightly focus on a single method or modification that improves throughput, such as FlashAttention [1]] and unpadding (66). In this study, incorporate many of these techniques to investigate whether they combine advantagcously.
RoBERTa ("Robustly optimized BERT approach") is the most influential work in this regard 37 In this study, they kept the exact BERT architecture but changed the training recipe by removing the next sentence prediction objective, training for longer on much larger datasets, and changing the batch size, among other things. They showed that the original BERT was significantly undertrained while the original BERT trained on 16GB worth of text, the top accuracy RoBERTa (Large) was trained on 160GB of data for 500,000 steps with batch size 8192. Many of the training choices in RoBERTa have become standard practice; training recipe therefore more closely resembles RoBERTa than the original BERT. See Appendix for a more detailed comparison of RoBERTa and MosaicBERT.
Improvements in transformer architecture and GPU hardware have caused the cost of pretraining BERT models to decline precipitously. The recent paper "Cramming: Training a Language Model on Single GPU in One Day" [18] exemplifies this trend. The goal of this rigorous study was to train the best BERT in 24 hours on single GPU. Similar to us, they tweaked the BERT architecture to incorporate FlashAttention and Gated Linear Units (but without increasing the dimensionality of the hidden block). Unlike Mosa,cBERT, they used scaled sinusoidal positional embeddings (59) as well as Pre-LayerNome (applying LayerNorm before the attention and feedforward layers) and did not change the pretraining masking rate. With this setup. they were able to train their modified BERT-Base to average GLUE score of 80.4 in 24 hours hours, as per Table of (18).. Our study similar in spirit, but asks what is the fastest architecture for pretraining, and expands on this in greater detail by showing that MosaicBERT is Pareto optimal relative to BERT.
single A6000 GPU (i.c. 24 GPU
While we believe that our training optimization choices for MosaicBERT go long way to improve BERT training efficiency, there are likely further modifications that could lead to an increase in accuracy with no effect on throughput, such as replacing LayerNorm with RMSNorm [67] and
GeGLU with SwiGL.U (50) 42. Removing all biases from the linear layers can also potentially lead to slight speed ups in training time without significant hit to accuracy. "NarrowBERT" [36

---

suggested clever change to the way encoder blocks are stacked so that computation is not wasted on unmasked tokens. Another recent study showed that dynamically masking the masking ratio during pretraining leads to downstream accuracy gains in BERT []] Note, however, that accounts such as 18] and [29] document many examples where certain modifications the BERT architecture and training recipe fail to improve accuracy. We also note here that there is an ongoing debate about
the benefits of RoPE (rotary positional embeddings) [55] versus AliBi [44]. Finally, using a more
modern tokenizer will likely have an effect on downstream accuracy [8]
Approaches such as knowledge distillation [3] might additionally push the Pareto frontier of BERT models during pretraining and finetuning. Progressive layer stacking is one example technique; by initializing larger BERT model with the weights from smaller pretrained BERT models, pretraining can be reliably accelerated [2] 29]. As the field is learning how to optimize stable model training at the billion parameter scale 13), we expect some of these innovations to cycle back to smaller models such as BERT-Base. For example, it has been hypothesized that incorporating more Layer- Norm modules at the QK matrix output (i.e. QK-LayerNorm (24) can lead improved stability during training; combining this with more aggressive learning rate schedule could lead to faster convergence.
Conclusion
In this study we show how combinations of architecture choices improve BERT pretraining speed and accuracy. We built MosaicBERT to enable ML researchers and engineers to pretrain BERT models from scratch their own data and build better models for their specific domains without facing time and cost restrictions. We ultimately hope that by making BERT training faster and cheaper, our work contributes the trend in NLP research away from finetuning generic models and towards of pretraining custom encoders on domain specific data. A large body of research has highlighted the success of BERT-style models pretrained
specific domains such biomedicine 3 33, 22, [7]. math [51], chemistry [2] (25], financial communications [49], and code [56]. This
seems hold true even in the age of LLMs 6 63].
Code
A stable webpage for this work can be found at mosaicbert. beesaiibert.githu.... io Code for pretrain- ing and finetuning MosaicBERT can be found in the MosaicMI. JessaicM,/examppless examples repository [htpps] The exact code for this study was pinned to v0. 4 of the MosaicML mws.acmlmmmmmmcm..m...m......mmm / repository https: thub cwm/mosai cmt/examplee/eeeeeeeememmceeeeee
//github. httpb///c//mmaacc/eeeeeeaaae
tree/v0. 0. 4/ //t////aaplesseeeeeeeeeeeeeeeeee berttetss////etee/eeeeee bert All pretraining and finetuning was done in PoTorc 1.13 using the
MosaicML Composer library https:://tttt wttp:tthhw...ooooiiiiiieoeomaaoe compose Model weights for co comoosaicml/
MosaicBERT-Baas Base can be found on the HuggingFace hub
Dosaic-bert-base
Code for pretraining and finetuning transformers such as MPT and derivative versions of Mos be found in the MosaicMI. 1lm-foundry repository https://giis t/g.thub...bbt.... com con/mosaiccm
Ilm-foundry
Acknowledgements
The authors would like thank Erica Yuen, Vitaliy Chiley, Connor Jennings, and Aaron Gokaslan for their feedback on the code and manuscript. We would also like to thank Ishana Shastri and Vlad Ivanchuk for some of their early work interns that led to MosaicBERT. The majority of this work appeared as a blogpost with the title ""Mos:icEERT: Pretraining BERT from Scratch for $20" in March 2023 [3 This work also builds the MosaicMI. MI Perf v2.1 results for BERT-Laage 31
10

---

References
[1] Z. Ankner, N. Saphra, D. Blalock, Frankle, and M. L. Leavitt. Dynamic masking rate schedules for mlm pretraining. arXiv preprint arXiv:2305. 15096, 2023.
[2] F. Bai, A. Ritter, and W. Xu. Pre-train annotate? domain adaptation with constrained budget. arXiv preprint 20X0:
[3] I. Beltagy, K. Lo, and A. Cohan. Scibert: pretrained language model for scientific text. arXiv preprint arXiv: 10676, 2019.
[4] L. Bentivogli, P. Clark, I. Dagan, and D. Giampiccolo. The fifth pascal recognizing textual entailment challenge. TAC. Citeseer, 2009.
[5] C. Blakeney, J. Z. Forde, Frankle, Z. Zong, and M. L. Leavitt. Reduce, reuse, recycle: Improving training efficiency with distillation. arXiv preprint arXiv:2211.000 2022.
[6] E. Bolton, D. Hall, M. Yasunaga, T. Lee, C. Manning, and P. Liang. BiomedIm. https 7/crfm. stanforarrrr edu/edd2/ 12t11///.m////mm..m//m/ biomedlm. html December 2022. Accessed on 15th May 2023.
[7] D. Cer, M. Diab, E. Agirr, I. Lopez-Gaziio, and L. Specia. Semeval-2017 task 1: Se- mantic textual conilariy-muultiiignn and cross-lingual focused evaluation. arXiv preprins 1XXv:/708.00053 2017.
[8] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint 2rXiv:2204.03111 2022.
[9] K. Clark, M.-T. Luong, Q. V. Le, and C. D. Manning. Electra: Pre-training text encoders as discriminators rather than generators. arXiv preprint arXiv:2003. 10555, 2020.
[10] I. Dagan, O. Glickman, and B. Magnini. The pascal recognising textual entailment challenge. In Machine Learning Challenges. Evaluating Predictive Uncertainty, Visual Object Classification, and Recognising Tectual Entailment: First PASCAL Machine Learning Challenges Workshop, MLCW 2005, Southampton, UK. April 11-13, 2005, Revised Selected Papers, pages 177-190. Springer, 2006.
[11] T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35: 16344-16359, 2022.
[12] Y. N. Dauphin, A. Fan, M. Auli, and D. Grangier. Language modeling with gated convolutional networks. International conference on machine learning, pages 933 941. PMLR, 2017.
[13] M. Dehghani, J. Djolonga,
B. Mustafa, Padlewski, J. Heek. Gilmer, A. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. arXiv preprint 23Xi:22002.05442. 2023.
[14] Devlin, M.-W. Chang, K. Lee, and Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0488)) 2018.
[15] N. Dey, D. Soboleva, F. Alattee,,,,, B. Yang, R. Pathria, H. Khachane, S. Muhammad, R. Myers, R. Steeves, N. Vassilieva, et Btlm-3b-8k: parameter performance 3b parameter model. arXiv preprint arXiv::309. 1566,, 2023.
[16] B. Dolan and C. Brockett. Automatically constructing a corpus of sentential paraphrases. In Third International Workshop on Paraphrasing (IWP2005), 2005.
[17] H. El Boukkouri, O. Ferret, Lavergne, and P. Zweigenbaum. Re-train or train from scratch? comparing pre- pre-training strategies of bert in the medical domain. In LREC 2022, pages 2626 2633, 2022.
[18] Geiping and T. Goldstein. Cramming: Training a language model on single gpu in one day. In International Conference Machine Learning, pages 11117-11143, PMLR, 2023.
11

---

[19] D. Giampiccolo, B. Magnini, I. Dagan, and W. B. Dolan. The third pascal recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing, pages 1-9, 2007.
[20] A. Gokaslan and V. Cohen. Openwebtext corpus, 2019. URL/httpp//Sk///ioo00...g
io/ oo/DpeniWebTextCorpuss
[21] L. Gong, D. He, Li, Qin, L. Wang, and Liu. Efficient training of bert by progressively stacking. In International conference machine learning, pages 2337-2346. PMLR, 2019.
[22] Y. Gu, R. Tinn, H. Cheng, M. Lucas, N. Usuyama, X. Liu, T. Naumann, J. Gao, and H. Poon. Domain-specifie language model pretraining for biomedical natural language processing. ACM Transactions on Computing for Healthcare (HEALTH), 3(1):1-23, 2021.
[23] P. He, J. Gao, and W. Chen. Debertav3: Improving deberta using clectra-style pre-training with gradient-disentangler embedding sharing. arXiv preprint arXiv:2/ 2021.
[24] A. Henry, P. R. Dachapally, S. Pawar, and Y. Chen. Query-key normalization for transformers. arXiv preprint arXiv:2010.04245, 2020.
[25] S. Horawalavithana, E. Ayton, S. Sharma, S. Howland, M. Subramanian, S. Vasquez, R. Cosbey, M. Glenski, and S. Volkova. Foundation models of scientific knowledge for chemistry: Oppor- tunities, challenges and lessons learned. In Proceedings of BigScience Episode ( 5-Workshop on Challenges & Perspectives in Creating Large Language Models, pages 160-172, 2022.
[26] S. Iyer, N. Dandekar, and K. Csernai First quora dataset relcase: Question pairs, 2017. URL oataaaattttttte--laaaaeeeeeeeiit
[27] G. Izacard, M. Caron, L. Hosseini, Riedel, Bojanowski, A. Joulin, and Grave. Unsuper- vised dense information retrieval with contrastive learning. 2022.
[28] Izsak, M. Berchansky, and O. Levy. How to train bert with an academic budget. arXiv preprint arXiv:2104.07705, 2021.
[29] J. Kaddour, O. Key, P. Nawrot, P. Minervini, and M. J. Kusner. No train no gain: Revis- iting efficient training algorithms for transformer-based language models. arXiv preprint 7XXv:2307.066440 2023.
[30] V. Karpukhin, B. Oguz, S. Min, P. Lewis, L. Wu, S. Edunov, D. Chen, and W.-t. Yih. Dense passage retrieval for open-domain question answering. arXiv preprint K1v:2004.04900 2020.
[31] D. Khudia, N. Sardana, S. Havens, A. Trott, and E. J. Yuen. Mosaicml delivers lead- ing nlp performance in miperf v2.1, 2022. URL https: // www..wosaccooommmmmmm.....w..o.m cml, cwwccom///bll /blm///g///// Aiperf-nlp-nov22
[32] V. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro. Re- ducing activation recomputation in large transformer models. arXiv preprint 0rXiv:2205.05198. 2022.
[33] Lee, W. Yoon, S. Kim, D. Kim, S. Kim, C. H. So, and J. Kang. Biobert: a pre-trained biomedical language representation model for biomedical text mining. Bioinformatics, 36(4): 1234-1240, 2020.
[34] H. Levesque, E. Davis, and L. Morgenstem. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
[35] P. Lewis. E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W.-t. Yih, T. Rocktäschel, al. Retrieval-augmenttd generation for nnowledge-nntensive nlp tasks. Advances in Neural Information Processing Systems, 33:9459- 9474, 2020.
[36] H. Li, P. Keung, D. Cheng, Kasai, and N. Smith. Narrowbert: Accelerating masked language model pretraining and inference. arXiv preprint arXiv:2301.04760 2023.
12

---

[37] Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint arXiv: /907. 1692, 2019.
[38] MosaicML NLP Team. Introducing mpt-30b: Raising the bar for open-source foundation
models. WWw www.iosaiiicoooo cm1. wtt.ccom/booooo www.co///oooopttaeeee 30b 2023. Accessed: 2023-06-22.
[39] MosaicML NLP Team. Introducing A new standard for openes commercially
usable llms. www..o/a..mm...mmm/ooo//////////
2023. Accessed: 2023-05-05.
[40] MosaicML Research Team. Methodology, 2022. URL https: www.wmosaicccmm http://wwwoaaicml.ccooooo
methodology
[41] S. Nagel. Common crawl blog news dataset available, 2023. URL WRLLhttts://commnccart org/blog/news htwp://wwwwaassseeeeee
available -aaasaa-aaaaet--aaaaalleeleee Accessed: 2016- 10-4.
[42] S. Narang, H. W. Chung, Y. Tay, W. Fedus, T. Fevry, M. Matena, K. Malkan, N. Fiedel, N. Shazeer, Z. Lan, et al. Do transformer modifications transfer across implementations and applications? arXiv preprint 1tXiv:2/02.11992. 2021.
[43] J. Portes, D. Blalock, C. Stephenson, and J. Frankle. Fast benchmarking of accuracy vs. training time with cyclic learning rates. arXiv preprint arXiv:2206.000 00832, 2022.
[44] O. Press, N. A. Smith, and M. Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. /2400,, 2021.
[45] M. N. Rabe and C. Staats. Self-attention does not need o(n²) memory. arXiv preprint arXiv:2112.05682, 2021.
[46] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 5551, 2020.
[47] Rajpurkar. Zhang, K. Lopyrev, and Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: 160.00250, 2016.
[48] D. P. Rodgers. Improvements in multiprocessor system design. ACM SIGARCH Computer Architecture News, 13(3):22--31,, 1985.
[49] R. S. Shah, K. Chawla, D. Eidnani, A. Shah, W. Du, S. Chava, N. Raman, C. Smiley, J. Chen, and D. Yang. When flue meets flang: Benchmarks and large pre-trained language model for financial domain. arXiv preprint ar 1xXiv:22/1.0000 1.00083, 2022.
[50] N. Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.052002, 2020.
[51] J. T. Shen, M. Yamashita, E. Prihar, N. Hefferman, X. Wu, B. Graff, and D. Lee. Mathbert: pre-trained language model for general nlp tasks in mathematics education. arXiv preprint arXiv: /06.:2/006.73440, 2021.
[52] W. Shi, S. Min, M. Yasunaga, M. Seo, R. James, M. Lewis, L. Zettlemoyer, and W.-t. Yih. Replug: Retrieval- -uggeette black-box language models, arXiv preprint 230.::2301.125222 2023.
[53] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint 1909 2019.
[54] R. Socher, A. Perelygin, J. Wu. J. Chuang, C. D. Manning, A. Y. Y. and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference empirical methods in natural language processing, pages 1631-1642, 2013.
[55] Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063 2024.
13

---

[56] J. Tabassum, M. Maddela, W. Xu, and A. Ritter. Code and named entity recognition in stackoverflow. arXiv preprint arXiv:2005. 01634, 2020.
[57] Tillet, H.-T. Kung, and D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10-19, 2019.
[58] T. H. Trinh and Q. V. Le. A simple method for commonsense reasoning. arXiv preprint 1806 02847, 2018.
[59] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[60] A. Wang, Singh, Michael, F. Hill, O. Levy, and S. R. Bowman. Glue: multi- task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 1804. 07461, 2018.
[61] L. Wang, N. Yang, X. Huang, B. Jiao, L. Yang, D. Jiang, R. Majumder, and F. Wei. Text embeddings by seakly-supervised contrastive pre-training. arXiv preprint arXiv:2212.03533, 2022.
[62] A. Warstadt, A. Singh, and S. R. Bowman. Neural network acceptability judgments. Transac- lions of the Association for Computational Linguistics, 7:625 2019.
[63] A. Wettig. T. Gao, Z. Zhong, and Chen. Should you mask 15% in masked language modeling? arXiv preprint arXiv:2202.08000 2022.
[64] A. Williams, N. Nangia, and S. R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint 17Xiv:1704.05426 2017.
[65] S. Wu, O. Irsoy, S. Lu, V. Dabravolski, M. Dredze, S. Gehrmann, Kambadur, D. Rosen- berg, and G. Mann. Bloomberggpt: A large language model for finance. arXiv preprint ar3iv:2303. 17564, 2023.
[66] J. Zeng, M. Li, Z. Wu, J. Liu, Y. Liu, D. Yu, and Y. Ma. Boosting distributed training performance of the unpadded bert model. arXiv preprint 2rXiv:2208.08124 2022.
[67] B. Zhang and R. Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[68] Y. Zhu, R. Kiros, R. Zemel, R. Salakhutdinov, R. Urtasun, A. Torralba, and S. Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.
14

---

A Pretraining Hyperparameters
We detail our pretraining hyperparameter values in Table ST below.
Model
Optimizer
LR
beta
eps
WD
MB
Warmup Final LR
MLM
BERT-Base
Decoupled
5e-4
[0.9,0998]
le-6
le-5
512
0.02LR
0.15
Adam W
MosaicBERT
Decoupled
5e-4
[0.9,0998]
le-6
le-5
512
0.02LR
0.3
Base
AdamW
BERT-Large
Decoupled
2c-4
[0.9,0.98]
1e-6
le-5
256
6%
0.02LR
0.15
AdamW
MosaicBERT-
Decoupled 2e-4
[0.9,0.98]
le-6
le-5
256
6%
0.02LR
0.3
Large
Adam
Table SI: Pretraining hyperparameters for BERT-Base MosaicBERT-Base. BERT-Large and MosaicBERT- LR = learning rate, WD weight decay, MB Microbbach and MLM
Masked Language Modeling ratio. All pretraining runs were done on 8x A100-80 GB GPUs.
All pretraining runs followed warmup linear decay learning rate schedule. Evaluation was done every 2000 batches, and checkpoints were saved every 3500 batches both of which slow down training slightly.
B Finetuning Hyperparameters
We used the hyperparameters in Table S2 for finetuning all BERT and MosaicBERT models. All finetuning datasets used a maximum sequence length of 256 tokens. We found that these values worked well across all tasks for BERT-Base, MosaicBERT-Base and MosaicBERT-Large BERT- Large however was somewhat less performant on QQP for some pretraining seeds.
Task
learning rate
beta
epsilon
weight decay
epochs
seeds
MNLI
5e-5
[0.9, 0.98]
Ie-6
5e-6
3
QNLI
le-5
[0.9, 0.98]
Ie-6
le-6
10
QQP
3e-5,1.5e-5
[0.9.0.988]
Ic-6
3e-6
5
RTE
le-5
[0.9, 0.98]
le-6
le-5
3
CoLA
5e-5
[0.9, 0.98]
le-6
5e-6
10
4
SST-2
3e-5
[0.9,0.988]
le-6
3e-6
3
3
MRPC
8e-5
[0.9, 0.98]
le-6
8e-6
10
5
STS
3e-5
[0.9, 0.98]
le-6
3e-6
5
Table S2: Finetuning hyperparameters for BERT and MosaicBERT across Base and Large. All
finetuning was done with sequence length 256. Note that finetuning BERT-Large on QQP with learning rate 3e-5 was highly unstable; we therefore lowered the learning rate in this instance to 1.5e-5 for BERT-Large (but not MosaicBERT-LLrree).
As can be seen from Figure S2 we found the finetuning of BERT-Large be highly variable. Finetuning runs that did not converge, or that led to outlier final values were dropped. For this reason, all checkpoints
represented in the BERT-Large average GLUE plot in Figure 3
We did not do a careful hyperparameter sweep for BERT-Large finetuning, and instead tried to keep values consistent across BERT-Base and Large well MosaicBERT Base and Large.
C Pareto Curves for Benchmarking Neural Network Training
In order to make meaningful improvements in neural network training efficiency, ML practitioners must be able to compare between different choices of network architectures, hyperparameters, and training algorithms. One straightforward way to do this is characterize the tradcoff between
15

---

Average GLUE (dev) Score
0.86
LeRe- (Deviin 2018)
0.84
Average GLUE (dev) Score
0.82
2018)
0.80
0.78
0.76
0.74
BERT-Base 178k steps
MosaicBERT-Base 178k steps
0.72
BERT-Base steps
MosaicBERT-Base 70k steps
0.70
5
10
15
20
25
Pretraining WCT (hours)
Figure S1: Pareto curves for BERT-Bare and MosaicBERT-BBE Base for runs trained for 70,000 and 178,000 steps with batch size 4096 and sequence length 128. Same data as Figure B and
accuracy and training time with "tradeoff curve" or a "Pareto curve." Pareto curves can be generated by varying the length of training for each model configuration; longer training runs take more time but tend to reach higher quality. For a fixed model and task configuration, this method of generating tradeoff curves is an estimate of the theoretical Pareto frontier, i.c. the set of all of the best possible tradeoffss between training time and accuracy, where any further attempt to improve one these for more details.
metrics worsens the other. See 40 43
We therefore consider one model Pareto optimal relative another if has superior accuracy across different training budgets while keeping everything fixed. Many studies advertise novel architecture approaches without measuring wall clock time, or simply show an improvement for single training is Pareto optimal relative to our BERT-Base
duration. In this study we show that MosaicBERT-Bee baseline for both short training durations and long training durations (Figures I 2 We also show that BERT-Large and MosaicBERTT-Lang are not Pareto optimal relative to BERT-Base and MosaicBERT- Base for training durations under 30 hours (Figure 3).
The Pareto plots in Figures D 3 and S2 are constructed by taking points from the same runs (with the same learning schedule), instead of having one run per pretraining budget on the plot. We did this for cost reasons; it would have been too expensive to do every run separately.
For ResNets, it has been established that the most accurate way generate Pareto curves is to completely separate training runs for each point on the Pareto curve, where each point also has their own learning rate schedule 4 However, the best of our knowledge, this has not been clearly established for BERT models. Can Pareto curves for BERT simply be constructed by evaluating checkpoints of single training run?
Since we ran BERT-Bass and MosaicBERT Base experiments for two separate learning rate schedules (70,000 steps and 178,000 steps with batch size 4096), we can plot the Pareto curves for all of these runs on top of each other (Figure(ST) Somewhat surprisingly, the 70k and 178k curves for BERT- Base (and MosaicBERT-B.... BERT models, a reasonable Pareto curve can be constructed using checkpoints from a single training run.
mostly lie top of each other. This is strong evidence that least for
We note that the Pareto curve approach discussed here also related to "tokens per parameter" (TPP) considerations that often discussed in LLM pretraining (c.g. (15)..
16

---

MNU-m (Accuracy)
QNLI (Accuracy)
(Accuracy)
RTE (Accuracy)
0.94
0.85
0.86
0.84
0.02
0.68
0.00
0.75
0.70
0.86
DERT Oase 170k
DERS (Datin
0.76
MosaicBERT-Base-
0.76
DERT-Large-18k
0.52
LoraaceERT-Large--1 178k
0.72
COLA (Matthew's Corr Coef)
5ST2 (Accuracy)
MRPC (Accuracy)
STSB (Spearman's Corr Coef)
0.60
0.94
0:90
0.91
0.93
0.55
0.86
0.90
0.97
0.50
0.90
0.82
0.00
0.87
0.70
0.40
0.87
0.80
0.74
0.85
60
20
40
40
60
Pretraining WCT (hours)
Pretreiming WCT (hours)
Pretraining WCT (hours)
Pretraining WCT (hours)
Figure S2: MosaicBERT-Base and Large accuracy (dev) vs. pretraining speed Pareto curves for individual GLUE benchmarks. All models were trained roughly 2 epochs of C4 (178,000 steps with batch size 4096 with maximum sequence length of 128). Error bars represent 95% confidence interval for n=2-3 pretraining seeds. Dashed black line represents BERT-Large accuracy the GLUE (dev) data as reported in 37 The average GLUE score can be found in Figure 3
D MosaicBERT-Bass Experiments with Increasing Sequence Length
Batch size, maximum sequence length, and number of training steps are important hyperparameters that affect both downstream GLUE accuracy as well as wall clock time. While increasing the maximum number of tokens in sequence leads to direct increase in wall clock time, it also exposes the model more data. For example, increasing sequence length from 128 to 256 while keeping all other hyperparameters constant leads to 2x the amount of total training tokens (assuming no padding). However, this will lead a commensurate increase in pretraining wall clock time. Given our Pareto-based approach, we can therefore ask: is it better to train with fewer steps and longer sequence length, or more sleps and shorter sequence length?
Here we find that training for more steps (i.e. 178,000) on shorter sequence length (i.e. 128 tokens) is Pareto optimal for GLUE scores.
As part of our experimental benchmarking, we trained MosaicBERT Base models with varying maximum sequence lengths ranging from 256 tokens to 2048 tokens for 70,000 steps with batch size 4096. We then plotted the downstream GLUE accuracy in Figures [S3] and SS4 Besides maximum sequence length, all other pretraining and finetuning hypeparameters are the same MosaicBERT. Base in the main text (and Tables ST and S2). Each point represents 12 pretraining checkpoints finetuned with multiple seeds according Table S2 The curve representing MosaicBERT--Ban. Base trained with sequence length 128 and 178,000 steps is the same as Figure
When looking at average GLUE accuracy sequence lengths 256, 512, 1024 and 2048 nominally do better than maximum sequence length of 128 tokens when trained for 70,000 steps. However, plotting steps the x-axis can be misleading, as increasing sequence length leads to increased wall clock time (ceteris paribus).
function of pretraining steps in Figure $3 longer
Instead, it is much more informative to plot wall clock time on the x-axis, in Figure S3B. Here, the order the curves is flipped and the curve representing maximum sequence length of 128 tokens is mostly Pareto optimal (red curve). Interestingly, training with batch size and sequence length of (4096,128) for 178,000 steps leads to the same final average GLUE accuracy and wall clock time as training with (4096, 512) for 70,000 steps (purple curve in Figure S3B). Note that the x-axis here is
17

---

A
MosaicBERT-Baal Across Sequence Lengths
0.86
0.85
Average (dev) Score
0.84
0.82
0.81
178k steps, seqien
0.80
70k steps, seqlen 256
70k steps, seqlen 512
0.79
70k steps, seqlen 1024
70k steps, seqlen 2048
0.78
25000
50000
75000
100000 Pretraining Steps
125000
150000
175000
B
MosaicBERT-Banl Across Sequence Lengths
0.86
0.85
Average GLUE (dev) Score
0.84
0.83
0.82
0.81
178k seqlen 128
0.80
steps, seqlen
70k steps, seqlen 512
0.79
70k steps, seqlen 1024
70k steps, seqlen 2048
0.78
10a
10a
Pretraining WCT (hours)
Figure S3: BRR--BBE Pareto curves maximum sequence lengths ranging from 128 to 2048 tokens, with all other hyperparameters fixed. (A) Average GLUE score function of pretraining steps. (B) Same data as in A, but with wall clock time on the x-axis. The curve representing maximum sequence length of 128 tokens mostly Pareto optimal (red curve). Interestingly, training with batch size and sequence length of !4096, 128) ) for 178,000 steps (red curve) leads to the same final average GLUE accuracy and wall clock time training with (4096, 512) for 70,000 steps (purple curve). Top dashed line represents BERT-Larg average (dev) GLUE score from Devlin et al. 2018, while bottom dashed line represents the BERT-Base average (dev) GLUE score from Devlin et al. 2018 [14].
plotted on logarithmic scale, as training with sequence length 2048 for 70,000 steps takes more than 100 hours on 8x A100 80GB GPUs.
18

---

MNUI-m (Accuracy)
QNLI (Accuracy)
(Accuracy)
RTE (Accuracy)
0.80
0.925
ECTT Larpe Doste 2018)
0.225
0.920
0.800
0.82
0.90
0.915
0.80
0.910
0.78
0.725
0.905
73k stops, 220ton
---
70% steps, seglen 512
POK staps, suqlen 2024
0.57
0.900
0.72
70K steps. soglen 2048
0.96
0.095
0.050
COLA (Matthew's Cor Coef)
5ST2 (Accuracy)
MRPC (Accuracy)
STSB (Spearman's Corr Coef)
0.94
0.93
0.60
0.93
0.00
0.55
0.89
0.90
0.85
0.82
0.45
0.57
0.80
0.85
101
107
101
107
10F
107
207
Protraining WCT (hours)
Pretraining (hours)
(hours)
Protroining WCT
Figure S4: MosaicBERT-Bass Pareto curves for maximum sequence lengths ranging from 128 to 2048 tokens, with all other hyperparameters fixed. Individual GLUB scores are plotted function of pretraining wall clock time. All curves used batch size 4096 and were pretrained for 70,000 steps, except for the curve representing sequence length 128 tokens, which was trained for 178,000 steps (and is the same as MosaieBERT Base in Figure
S4 plots the GLUE scores for individual tasks function of pretraining wall clock time. The curve representing maximum sequence length of 128 tokens is mostly Pareto optimal (red curve). However, QNLI and STSB seem to benefit from longer maximum sequence lengths.
Figure
An important consideration when increasing the maximum sequence length is the issue of padding. Many text samples in our C4 pretraining dataset contained fewer than 2048 tokens and needed to be padded. The "unpadding" approach detailed in the main text applies the feedforward layers of the transformer but not the attention layers. As sequence length is increased, there is potentially wasted computation on padding tokens in the attention layers. This means that wall clock time increases without significant boosts in accuracy. This is likely the case for the curves representing sequence lengths 1024 and 2048.
Another important consideration that the GLUE tasks are quite short by modern standards (circa 2023) and contain fewer than 512 tokens. Additionally, all GLUE finetuming in this study was done with maximum sequence length of 256 tokens. Training with sequence length 1024 or 2048 tokens is likely unnecessary for this task. It possible, however, that these models superior at tasks that involve long-context dependencies.
The original BERT paper pretrained with maximum sequence length of 128 tokens for 90% of training, and then completed pretraining with a sequence length of 512 [14]. The RoBERTa paper pretrained with maximum sequence length of 512 for the full duration of pretraining and also "packed" each batch sample with full sentences up to the limit of 512 tokens to decrease the number of padding tokens [37].
released the pretrained checkpoints for these models on the HuggingFace hub in April 2023:
19

---

Model
Bsz
Steps
SeqL
MNLI
QNLI
QQP
RTE
SST
MRPC CoLA
STS
Av.
BERT-Lave
256
1M
128
86.6
92.3
91.3
70.4
93.2
88
60.6
90
84.05
BookC + Wiki
RoBERTa-Base
8k
500k
512
87.6
92.8
91.9
78.7
94.8
90.2
63.6
91.2
86.35
all 5otk+
RoBERTa-L:
1. BookC+Wiki
8k
100k
512
89
93.9
91.9
84.5
95.3
90.2
66.3
91.6
87.8
all data
8k
100k
512
89.3
92
82.7
95.6
91.4
66.1
92.2
87.9
3. all data+300k
8k
300k
512
90
94.5
92.2
83.3
96.1
91.1
67.4
92.3
88.4
4. all data+500k
8k
500k
512
90.2
94.7
92.2
86.6
96.4
90.9
68
92.4
88.9
BERT-Base
4k
178k
128
85.44
91.6
91.4
78.2
93.10
89.5
59
89.4
84.7
MosaicBERT-B
4k
178k
128
85.6
91.4
92
83
93.5
89
58.2
90.3
85.4
BERT-Large
4k
178k
128
86.3
92.8
90.9
83.8
93.3
87.8
56.2
90.6
85.2
MosaicBERT-L
4k
178k
128
86.9
92
84.5
93.7
88.2
59.7
90.9
86.1
Table S3:
Model performance comparison GLUE (dev) tasks. Data from Table 8 of the RoBERTa
paper 33
Values for BERT and MosT bottom rows are best average values per task from
Figure S2
E
Comparison to RoBERTa
As mentioned in the main text. RoBERTa ("Robustly optimized BERT approach") was an important follow-up to the original BERT (37. In this study, they kept the exact BERT architecture the same but changed the training recipe by removing the next sentence prediction (NSP) objective, training for longer on much larger datasets, and changing the batch size, among other things. Many of the training choices in RoBERTa have become standard practice; our training recipe therefore more closely resembles RoBERTa than the original BERT. The RoBERTa paper also did careful experiments and ablations that were helpful for the research community.
The main pretraining results from RoBERTa-Base and Large are shown in Table (data from Table 8 in the Appendix of the RoBERTa paper (37). In particular, they trained on:
BookCorpus 688 and English Wikipedia for 100k steps with batch size 8192 and sequence length 512
2. BookCorpus 168 and English Wikipedia (16GB), plus "additional data" CC News, the
English portion of the CommonCrawi News dataset (76GB) [41), Open WebText Corpus (38GB) (20). and STORIES (31GB) (58 for 100k steps with batch size 8192 and sequence length 512
3. The full dataset mixture (160GB) for 300k steps with batch size 8192 and sequence length 512
4. The full dataset mixture (160GB) for 500k steps with batch size 8192 and sequence length 512
The RoBERTa- Base and Large values are better than the MosaicBERT- Base and Large values. An aoples-to-appll comparison of RoBERTa and MosaicBERT experiments would keep all hyperparam- eters consistent, including training data and batch size. Since we cannot do this post hoc, a somewhat reasonable comparison would be to keep the total amount data constant. 100k steps with batch size 8192 is roughly equivalent to 1 178k steps with batch size 4096 (i.e. it is almost equivalent to 89k steps with batch size 8192). However, RoBERTa also pretrained with a maximum sequence length of 512 tokens, which is 4x longer than the maximum sequence length of 128 we used for pretraining in this study. This means that the "worst" RoBERTa models trained with batch size 8192 for 100k steps still saw 4x more tokens than the MosaicBERT models trained with batch size 4096 for 178k steps. This is one likely one reason that the final RoBERTa values on GLUE are higher.
Additionally, the RoBERTa study formatted the data to minimize padding tokens by packing each sample with full sentences sampled contiguously from one or more documents (with an extra separator token) 37 They also did learning rate sweep for single-task GLUE finetuning (le-5, 2e-5, 3e-5), which we did not do. Finally, the differences between MosaicBERT and RoBERTa could also due to differences in the underlying datascts.
20

---

GLUE Benchmark Details
The GLUE benchmark consists of 8 (originally 9) tasks [60]. Since there has been a Cambrian explosion of benchmarks since the halcyon days of GLUE, elaborate the individual GLUE benchmarks for reference. We used the GLUE tasks from HuggingFace for finetuning and evaluation: https:/ haaps://huugingee ace. co/ datasatss/gsleeeee All evaluation was done the validation (a.k.a. dev) splits.
F11 Large Finetuning Datasets
MNLI (Multi-Genre Natural Language Inference) [392,702 train I 19,643 validation 19,643 test] a large crowd-sourced entailment classification task [64]. The model is given two sentences and has to predict whether the second sentence is entailed by, contradicts, is neutral with respect to the first one. For example:
Premise: "Buffet and la carte available."
Hypothesis: "It has a buffet."
Label: 0 (entailment)
MNII has two subsets, "matched" (MNLI-m) and "mismatched" (MNLI-mm). All numbers reported this study are for MNLI-m, unless otherwise stated.
QNLI [104,743 train 5,463 validation 5,463 test] this Stanford Question Answering dataset consists of areestio--paragrappe pairs drawn from Wikipedia 447
QQP (Quora Question Pairs [363,846 train 40,400 validation 390,965 test]. The task is to determine whether two sentences are semantically equivalent [26].
F22 Small Finetuning Datasets
RTE (Recognizing Textual Entailment) [2.490 train 277 validation 3,000 test] Given two sentences, the model has to predict whether the second sentence is or is not entailed by the first Note that in our work we use a checkpoint from the MNLI finetuning to finetune
sentence 10 19
on RTE.
CoLA (Corpus of Linguistic Acceptability) [8,551 train 1,040 validation 1,063 test] 62 is benchmark with sentences that are either linguistically acceptable or grammatically incorrect. For example:
"The higher the stakes, the lower his expectations are." Label: (acceptable)
"Mickey looked up it." Label: (unacceptable)
SST-2 (Stanford Sentiment Treebank) [67,349 train 872 validation I 1,821 test] consists of sentences from movie reviews. The task is to classify sentiment as either positive or negative (54)
(3.66 MRPC (Microsoft Research Paraphrase [3rpus) 668 train I 408 validation 1,725 test] [6] The
dataset consists of sentence pairs extracted from online news sources. The task classify whether the sentences in the pair are semantically equivalent.
STSB (Semantic Textual Similarity Benchmark) [5,749 train 1,500 validation 1,379 test] This dataset contains sentence pairs that are given similarity scores from 0 7
Note that we excluded finetuning on the 9th GLUE task WNLI (Winograd NLJ) 34], as in the original BERT study (it is a very small dataset [634 train, 146 test] with a high number adversarial examples). Finetuning on RTE, MRPC and STSB starts from checkpoint already finetuned on MNLI (following the example of (8 and other studies). This is done because all the above tasks deal with sentence pairs, and this staged finetuning leads to consistent empirical improvement.
21

---

G MosaicBERT-Large Multinode Throughput Scaling
The experiments in the main section of this paper were all performed on a single node with 8x AxkAUu A00 GPUs. How well do our innovations the BERT architecture maximize throughput at the multinode scale?
We measured the throughput of MosaicBERT--arg (430M) during training 8. 16, 32, 64, 128 and 200 GPUs, and plotted the tokens per second for various global batch sizes. Global batch size an important factor in the throughput measurements; in general, cranking up the batch size increases the GPU utilization and raw throughput. As the number of nodes increases, the global batch size needs to be increased as well in order maintain high throughput.
the global batch size is kept constant while increasing the number of nodes, the throughput does not increase linearly. This can be seen in Figure S3 global batch size of 4096 spread across 64 GPUs using Distributed Data Parallelism (DDP) means that each GPU will only apply matmul operations matrices with dimension of 64, which leads to suboptimal throughput. If the global batch size is increased to 65,536 across 64 GPUs, this roughly means that each GPU will apply matmul operations matrices with dimension of 1024, leading to higher throughput. However, such large global batch size might not lead to the best downstream accuracy; this is a question that we were not able to address in this study due resource and time constraints.
G.1 How MosaicBERT modifications might scale to larger models
All the architectural and configuration choices made in MosaicBERT can in principle be applied to larger models, although they might have different effects throughput. Additionally, they should not pose any issues with different forms of parallelism (some of the modifications we explored were used successfully at large scale in frameworks like MEGATRON as well as MPT-7B and MPT-30B 39 38). ). For example, Gated Linear Units should have no issues with tensor parallelism as long the matrix multiplications are multiple of the Tensor Parallel world size. As models are scaled, LayerNorm becomes a smaller portion of the total compute, so low precision LayerNorm might not matter much.
166
throughput (tokens/second
global
2-
global bs 8192 (A10080GB)
16394 (A10080G8)
global 65536 (A10080GH))
theoretical 2x throughput
128 GPU (A10040GB)
75
150
175
200
#
GPUs (A100 40GB and 80GB)
Figure S5: MosaicBEERT-Large (430M) multinode throughput scaling
H MosaicBERT-Base Model FLOPS Utilization (MFU)
Model FLOPs Utilization (MFU) is an estimate of what percentage of the hardware's FLOPs being used during training. The estimate is based on the measured throughput and the known FLOPs of the computation.
MFU calculates the utilization from the floating point operations required for a single for- ward/bacckwrrss pass of the model, and does not account for the additional compute required for
22

---

Model
Throughput (tokens /sec)
MFU
Hardware
Time to 79.6
Batch
Micro-
Size
batch
Size
BERT Base
0.4e6
10.4%
8x
1 110.4 minutes
4096
512
A100 80
(1.84 hours)
MosaicBERT
1.1e6
39.97%
8x
67.8 minutes
4096
512
Base
A100 80
(1.13 hours)
MosaicBERT
0.938c6
30.9%
8x
76.8 minutes
4096
128
Base
A100 40
MosaicBERT
1.88e6
31.0%
16x
38.5 minutes
4096
128
Base
A100 40
MosaicBERT
3.15e6
25.9%
32x
23.1 minutes
4096
128
Base
A100 40
MosaicBERT
4.77e6
19.6%
64x
15.7 minutes
4096
64
Base
A100 40
Table S4: Multinode Throughput scaling for MosaicBERI-Bas
other implementation details such as activation checkpointing. Thus, MFU is independent of imple- mentation and hardware. For more details, see (32). All FLOP calculations exclude the operations required for normalization, activation, and residuals.
Following the notation the PaL.M paper
Model FLOPPs Utilization (MFU) is approximated as:
(6 mparometers Tobserved)
MFU
(3)
Rgpus Thhereticcal
where Tobserved is the observed throughput and Tiheoretioal is the theoretical peak throughput.
In the numerator, the number of learnable parameters in the model is multiplied by a factor of 6 to estimate the matmul FLOPs per token This is then multiplied by the number of tokens seen per second. As first-order approximation, we exclude the extra FLOPs per token due dense self-attention.
(2x for the forward pass and 4x for the backward pass).
In the denominator, the thcoretical peak throughput is provided in the GPU hardware specs. For A100 GPUs using bfloat16, this theoretical peak throughput is 312 teraFLOPs.
MosaicBERT-
8XA1100
8xA100 80GB cost
8xA100
8xA100 40GB
Base Ave. GLUE
80GB hours
($2.50 GPU/hr)
40GB hours
cost ($2 GPU/hr)
Score
79.6
1.13
$22.60
1.28
$20.00
82.2
2.81
$56.20
3.19
$51.00
83.4
5.27
$105.40
5.99
$95.78
Table S5: MosaicBERT-Bass GLUE (dev) scores, time and cost comparison
I GPU Pricing
As of mid-2023, A100 GPU pricing ranges from $4.10 (40 GB) for on demand cloud compute on AWS, to $2.46 (40 GB) / $5.00 (80 GB) per GPU on GCP $1.10 (40 GB) / $1.50 (80 GB) per GPU using Lambda labs. At an intermediate price of $2.50 an hour per A 100 80 GB GPU, training to 79.6 GLUE average takes 1.13 hours and costs roughly $22.60 Some example costs calculated in Table 55
"See for example "Cloud GPU instances with the largest VRAM 2022" (medium Foom//aleiin ettp///eexxlpp/// cloud- gpu- conu-instances to-solve out of--or-memory mem.ry-moooo--ooo eerorrr--eerrr--eee error- 25122-24551288327277
23

---

(350, 138, 2881
[266,
1896 128, 739]
Liward 786x3072
(256, 120, 7988
Linew]] 768-2248
Linor 766
(256, 5072)
Uisard 77842043
PSA, 120, 2048]
(256.
Actvaten Fare
1208: 125, 2040)
AStNEDoI
[256, 30722
skereeniita mal
1290 12A, 20481
Addi-s001 FARYS
120 20H4)
[256. 3072)
Liwwar 2002
(2A 125L 204E
(208, 120, 2011
2000 150, 269)
(280, 120, 708)
Einadt
(256 2040
pS6 120, 7881
(256, 120, 7041
(256 120, 793)
Lina 2048x7E5
inpui
(250, 1201 700)
p2u rea]
[256,
(206 121, 2911
Layerhorm
domentwia
oup
(256
ouppt
Figure S6: Standard FeedForward Transformer Block and Gated Linear Unit Modifications. Each edge shows the tensor dimension with batch size of 256, sequence length of 128 and hidden dimension of 768. (A) A standard transformer feedforward block. (B) Naive implementation of Gated Linear Unit. The number of parameters in this are the same in (A). (C) Fused implementation of Gated Linear Unit where the two matrix multiplications (Linear1 O and Linear1_ 1) from (B) fused into (Linear1) with 2x the parameters. Here the output is sliced, and is functionally equivalent to (B).
J Gated Linear Units (GLU) Optimizations
GLU adds elementwise multiplication of two linear projections and it leads to qualitative improve- ments over the standard Transformer block. There multiple ways to implement GLUs and we experimented with two implementations. Figure shows standard feedforward transformer block (A) and two implementations of GLUs (B-C). "Fused GLU" in (C) fuses the two matrix multiplications into one and expected to perform better in some domains.
Figure S7 shows the performance impact of the two GLU over standard feedforward transformer block (which would 0% slowdown) for a single GPU. This figure only shows the performance of the forward pass, and the backward pass is expected to behave similarly. We can draw two conclusions

---

GLU Slowdown
Fused GLU Slowdown
Slowdown Wit Standard (%)
20%
0%
166
64
128
256
512
1024
Batch
Figure S7: Slowdown of different implementations of Gated Linear Unit. This slowdown with respect to standard feedforward transformer block. The number of parameters between standard feedforward transformer block and the two GLU implementations are the same.
from this chart: For smaller batch sizes, both GLU implementations add significant overhead over the standard block 2) For batch sizes 128, Fused GLU implementation is better than regular GLU implementation and beyond 128 it's slightly worse. The implementation used in the main text is the "Fused GLU" implementation (C) with global batch size 4096. Since the profiling in Figure S7 is per 512.
GPU, this is in the regime of 4096/8
The main reason for slowness of GL.Us over standard block is extra elementwise multiplication in GLU layers. As for why fused implementation is slower, profiling analysis shows that the Linear layer ends up calling different CUDA kernels for matrix-multiphee multipliciiiis and their relative performance varies for different sizes. While the MosaicBERT architecture in this work uses the fused GLU implementation, the analysis here indicates that it would be slightly more efficient to use the standard GLU implementation instead.
K Limitations and Broader Impact
K.1
Limitations
While trained two different model sizes, have not pretrained MosaicBERT model in the 1B parameter range. In this regime, it is possible there will be training stability issues; this is an area of future work.
We also only trained models for 70,000 steps and 178,000 steps of batch size 4096. possible that some of the Pareto properties change in the longer regime, although suspect that this is unlikely.
K.2 Broader Impact
BERT models are highly used for NLP tasks. By open-soureing this work, we hope that our code and models will be used by the wider research community. We recognize however that models like BERT and MosaicBERT are tools that can be used for nefarious purposes, and that biases inherent in the training data can be reflected in the final model artefacts.
25


Reviews:
[{"review_id": "0CSYpaJEZV", "reviewer": "Reviewer_Ss7L", "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.", "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?", "limitations": "I didn’t see a clearly potential negative societal impact of this paper.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n* The paper presents ablation experiments that quantify the impact of each proposed change.\n* The paper is well written and easy to follow. Replicating the results should not be hard.\n", "weaknesses": "* The experiments are only conducted on GLUE tasks; it is not clear how well the trained model transfers to other datasets.\n* There is no fair comparison (i.e., under the same hardware setup) of the proposed recipe and other efficient training methods such as CrammingBERT.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and questions.\n\n> The experiments are only conducted on GLUE tasks; it is not clear how well the trained model transfers to other datasets.\n\nWe chose to focus exclusively on the GLUE benchmark for multiple reasons. Since GLUE is a classic benchmark, fine-tuning on GLUE makes it easy to compare results across both recent and more “classic” papers (e.g. the original BERT). On a more practical level, the data presented in our paper represents hundreds of GPU hours, and we had to make the pragmatic decision to extensively benchmark on one dataset instead of benchmarking on multiple datasets in a limited way.\n\n> There is no fair comparison (i.e., under the same hardware setup) of the proposed recipe and other efficient training methods such as CrammingBERT.\n\nWhile some of the architecture modifications in our paper are shared with Cramming BERT, including FlashAttention and Gated Linear Units (GLU), we consider our work complimentary in the following ways:\n* We try to build and benchmark the fastest pretraining recipe with the highest performance on top-end hardware, while Cramming BERT focuses on the low-resource scenario (one consumer GPU for one day). In that sense, it focuses on one lower-performance point on the pareto frontier, whereas we focus on the entire pareto frontier.\n* We evaluate multiple model sizes (i.e. both RapidBERT-Base and RapidBERT-Large), while Cramming BERT only analyzes BERT-Base\n* We share full Pareto curves for BERT Base and Large as well as throughput and accuracy ablations (see Author Rebuttal PDF for further ablations). We believe this extensive data will be exceedingly useful to practitioners. We have already open-sourced our code and hope that the results from hundreds of GPU hours will be useful to the ML community\n\n> 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?\n\nMultiple studies have shown that the simple change to the MLM masking ratio from 15% (in the original BERT) to 30% leads to a small but significant improvement in downstream GLUE performance. In “Should You Mask 15% in Masked Language Modeling?” Wettig et al. find that constant MLM masking ratios above 15% lead to improved average GLUE and SQuAD scores for bert-base and bert-large. Similarly, in the recently published paper “Dynamic Masking Rate Schedules for MLM Pretraining,” Ankner et al. find that a constant MLM masking ratio of 30% consistently outperforms a MLM masking ratio of 15% for bert-base. We have elaborated on this in the updated manuscript.\n\nWe have also included an updated figure in the Author Rebuttal PDF showing that an MLM ratio of 30% leads to a small improvement in the downstream GLUE score over MLM 15% without affecting the wall clock time.\n"}, {"review_id": "FLKq3cRBCv", "reviewer": "Reviewer_xi9s", "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. ", "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.\n\n**EDIT: following rebuttal, I have updated my score to a Weak Accept (6).**\n\n* **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)?\n\n* **Q2.** Could the authors explain why the baseline used is so much above the original BERT-Base?\n\n* **Q3.** Could the authors clarify whether they use for RapidBERT a shorter training seqlen to accelerate training as proposed?\n\n* **Q4.** Could the authors provide further performance ablations for FlashAttention/unpadding?\n\n* **Q5.** (more of a suggestion) Could the authors better describe the limitation of their approach regarding Pareto-optimality? \n\nSmall suggestions:\n* 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\".", "limitations": "While there is no explicit section, some limitations in terms of scope are discussed in the conclusion.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n\n* **S2.** The authors detail their contributions and open-source their code, making these results reproducible and enabling the community to build upon them.\n\n* **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.**\n    * **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. \n   * **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?\n   * **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. \n   * **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.\n\n* **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).\n\n* **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. \n\n* **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments and suggestions, and we try to address them here.\n\n**W1** One of the primary goals in this paper is to develop a high-performance BERT architecture and recipe for pretraining on commercially available high-end hardware (A100-80GB GPUs). We agree with the reviewer that ablations should focus on downstream performance in addition to throughput, and have included updated plots with the individual effects of GLU, low precision LayerNorm, etc. on both wall clock time and GLUE scores.\n\n**W1.1 (How ablations affect GLUE score)** We ran further ablations in response to the reviewer’s comments and have plotted downstream GLUE accuracy as a function of measured pretraining wall clock time. We provide figures in the Author Rebuttal PDF. The patterns here shed light on the individual effects of various architectures (e.g. BERT+GLU, BERT+low precision LayerNorm) and training configurations (e.g. BERT + 30% masking ratio). On average, all methods seem to provide a slight accuracy boost to BERT-base. Increasing the masking ratio to 30% leads to a slight accuracy boost while not affecting the WCT, while turning off dropout in the attention layer leads to a slight improvement in both accuracy and WCT. Low precision LayerNorm leads to a significant speedup. Gated Linear Units (GLU) add more parameters to the model and lead to a significant slowdown while providing an accuracy boost. As a point of reference, we also benchmark the full RapidBERT as well as RapidBERT without FlashAttention and with dropout set to 0.1 (the standard BERT-base  configuration).\n\n**W1.2 (Strong baseline)** Thank you for this suggestion. We will include the time it takes for our strong baseline to reach a 79.6 GLUE score in Table 1. We are confident that the C4 dataset is the main reason for the strong BERT baseline relative to the original BERT (Devlin et al. 2018). The original BERT was trained for 40 epochs on Wikipedia + Books Corpus (16GB of text for a single epoch). C4 contains Wikipedia and Books Corpus, as well as much more data (roughly 350 GB). Besides the use of C4, we note some other differences from the original BERT paper: 8xA100 80 GB GPUs, MLM pretraining objective instead of MLM + NSP, a sequence length of 128 throughout the entirety of pretraining (instead of starting with 128 and then expanding to 512 towards the end of pretraining), larger batch size (the original paper used a batch size of 256), an increased vocab size of 30,528 (the original used a vocab size of 30,000, while the default Hugging Face vocab size for the bert-base-uncased tokenizer is 30,522), and the use of PyTorch 1.13. We used the original bert-base-uncased tokenizer for all our models; it is likely that a more modern tokenizer would improve GLUE results as well.\n\n**W1.3 (ALiBi)** Our main motivation for including ALiBi was to allow for long sequence extrapolation and-or finetuning in downstream use cases. There is increasing demand in the ML community for long context windows and alternatives to positional embedding schemes. We agree with the reviewer that our submission did not decisively show that ALiBi extrapolates well at test time. Both BERT and RapidBERT pretraining was done on a sequence length of 128, while finetuning was done for a sequence length of 512 (and have made this clear in the updated manuscript). Unfortunately the GLUE benchmark does not have many examples with long sequence lengths above 512 tokens, and is therefore not a great benchmark for testing sequence length extrapolation.\n\n**W1.4 (Ablations)** We agree with the reviewer that ablations here should not only focus on throughput but also downstream GLUE accuracy. We have provided updated experiments on GLUE score vs. wall clock time for BERT, BERT+GLU, BERT+low precision LayerNorm, BERT+ mlm 30% ratio, BERT+no attention dropout, RapidBERT, and RapidBERT-FlashAttention in the Author Rebuttal PDF. In the final version of the manuscript, we will also include results for BERT+ALiBi, BERT+ALiBi+GLU, etc.\n\n**W2 (Pareto Curve)** The Pareto plots in Figures 1-4 are constructed by taking points from the same runs (with the same learning schedule), instead of having one run per pretraining budget on the plot. We did this for cost reasons; it would have been too expensive to do every run separately. We emphasize that this approximation penalizes intermediate budgets, as they are evaluated with an incomplete LR schedule, so we expect that our results would be even stronger had we used the learning rate schedule corresponding to the training time for each point. \n\n**W3 (Repetitive)** We appreciate this feedback, and we will cut some of the unnecessary language from sections 3.2 and 3.3\n\n**W4 (References)** Thank you for these suggestions. We have updated the manuscript to properly cite the Triton paper and the correct NVIDIA reference for unpadding (which was NVIDIA’s MLPerf v1.1 submission).\n\n\n\n\n\n**Q1:** We hope that the additional ablation experiments shed light on the large gap between RapidBERT and BERT in Figure 3. Each of the individual architecture modifications lead to either a downstream accuracy boost or a speedup; combined (RapidBERT) they lead to the full Pareto improvement over BERT.\n\n**Q2:** We have addressed this question in comment W1.2\n\n**Q3:** We pretrain with sequence length 128 for both the baseline BERT and for RapidBERT. We therefore do not count this as a speedup method. We then finetune both BERT and RapidBERT with sequence length 512.\n\n**Q4:** We have provided some FlashAttention ablations in the Author Rebuttal PDF.\n\n**Q5:** We address the question of Pareto optimality in W2\n"}, {"review_id": "yrUuWzsrQF", "reviewer": "Reviewer_NjCS", "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.\n\nRapidBERT 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.\n\nThe 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.\n\nIn 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.\n\nOverall, 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.", "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?\n~- 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?~", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "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.\n- 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.\n- 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.\n- 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.~\n~- 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.~\n- 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.\n- 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%.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments and suggestions.\n\n> approach primarily combines existing techniques…the overall novelty of the architectural choices is limited\n\nWe strongly but respectfully disagree with the reviewer’s concerns about novelty. The literature is full of papers that propose a single, novel improvement to an architecture like BERT. (E.g., the many papers we cite and build on.) However, there is very little work assessing what this collection of papers actually amounts to: whether these improvements fit together and combine to produce something even stronger, or whether they are different, mutually exclusive paths to the same gains. We believe that it is crucial to take stock of scientific progress in the manner that we have in our paper, and we emphasize the novelty of (1) the nontrivial work we needed to perform to get these many improvements to work well together and (2) the fact that we sought to combine many existing methods rather than invent yet another new technique of uncertain value. We hope the reviewer recognizes the importance of this kind of contribution and the novelty inherent in performing this work.\n\n> It would have been beneficial to include a discussion comparing the proposed approach with the original RoBERTa model as a baseline\n\nWe appreciate the reviewer’s comments regarding RoBERTa, and have elaborated on this in the updated version of the manuscript. RoBERTa is a “training recipe” for BERT that preserves the network architecture but changes hyperparameters and training data. The RoBERTa paper showed that the training dataset is incredibly important; for example, RoBERTa trained on 350GB of data while BERT only trained on 16GB of data. The RoBERTa paper also showed that the next sentence prediction (NSP) objective was superfluous, and showed that increasing the batch size led to accuracy gains. These RoBERTa training choices have become the standard in the field, and our general training recipe choices resemble RoBERTa more than the original BERT (as we mention in the paper).\n\nSince our goal was to show that both our architectural and training recipe choices led to strict accuracy vs. wall clock time Pareto improvements, we wanted to make sure that we were comparing our models with a very strong baseline. Therefore both BERT and RapidBERT used the same dataset (C4), tokenizer, learning rate schedule, batch size and device microbatch size, hardware, etc. The initial untrained BERT architecture and the RoBERTa architecture are the same, so the setup would be the same for our strong baseline. We have included more reported RoBERTa values in the Appendix of our updated manuscript as a helpful reference point.\n\n> how the architectural choices … could facilitate the development of larger language models\n\nWe thank the reviewer for making this point, and have expanded on this in our updated manuscript. All of the architectural and configuration choices made in RapidBERT can in principle be applied to larger models, although they might have different effects on throughput. Additionally, they should not pose any issues with different forms of parallelism (some of the modifications we explored were used successfully at large scale in frameworks like MEGATRON). For example, Gated Linear Units should have no issues with tensor parallelism as long as the matrix multiplications are a multiple of the TP world size. As models are scaled, LayerNorm becomes a smaller portion of the total compute, so low precision LayerNorm might not matter as much. \n\n> The ablation studies do not investigate the effects of changing the masking ratio\n\nMultiple studies have shown that the change to the MLM masking ratio from 15% (in the original BERT) to 30% leads to a small but significant improvement in downstream GLUE performance. In “Should You Mask 15% in Masked Language Modeling?” Wettig et al. find that constant MLM masking ratios above 15% lead to improved average GLUE and SQuAD scores for bert-base and bert-large. Similarly, in the recently published paper “Dynamic Masking Rate Schedules for MLM Pretraining,” Ankner et al. find that a constant MLM masking ratio of 30% consistently outperforms a MLM masking ratio of 15% for BERT-base. We have elaborated on this in the updated manuscript.\n\nWe have also included an updated figure in the Author Rebuttal PDF showing that an MLM ratio of 30% leads to a small improvement in the downstream GLUE score over MLM 15% without affecting the wall clock time.\n\n> the absence of numerical instabilities in bfloat16. \n\nAs we were prototyping our RapidBERT architecture and training recipe, we were surprised to find that low precision (bfloat16) LayerNorm afforded significant speedups without compromising stability. While instabilities such as loss spikes are strongly coupled to the learning rate schedule and other hyperparameters, we found low precision LayerNorm to be a strict Pareto improvement when benchmarking wall clock time and downstream accuracy in our hyperparameter regime. We don’t have theoretical intuitions for why numerical instabilities didn’t occur in our hyperparameter regime. We should note however that our learning rate schedule was not overly aggressive (warmup followed by linear decay with peak values well within the range of published studies), and that this could have mitigated loss spikes. Finally, we should also note that the NVIDIA Apex library and Megatron both use a form of low precision LayerNorm in their code, but it is not documented in any papers that we could find. We have elaborated on this in our updated manuscript.\n"}, {"review_id": "NUY0IDmi0q", "reviewer": "Reviewer_WcGd", "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", "questions": "1. What is “Pareto Optimal”? This term is mentioned multiple times but never formally defined in the paper\n2. There appears to be a throughput difference in Figure 5 between “lpLN+GLU+ALBi” and “GLU 3072”. What leads to the difference?", "limitations": "The authors have discussed the limitations of their work.", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The contribution of this paper is limited. \n> I’d hope that the authors can provide more experiments showing why all these modifications are necessary.\n\nWe have added many additional ablations to demonstrate the individual contributions of each of the changes to the architecture in the Author Rebuttal PDF. This is our best attempt to address the reviewer’s feedback. However, this feedback is vague and unactionable; we hope the reviewer will give us concrete and constructive feedback after the author response period so that we can improve our paper.\n\n> The contribution of this paper is limited. The proposed architecture modifications are all adopted from prior work.\n\nWe strongly but respectfully disagree with this comment. The literature is full of papers that propose a single, novel improvement to an architecture like BERT (e.g., the many papers we cite and build on). However, there is very little work assessing what this collection of papers actually amounts to: whether these improvements fit together and combine to produce something even stronger, or whether they are different, mutually exclusive paths to the same gains. We believe that it is crucial to take stock of scientific progress in the manner that we have in our paper. We hope the reviewer recognizes the importance of this kind of contribution.\n> their current findings depend on multiple architecture modifications, which is complicated and makes their scientific findings unclear.\n> simplify it further for a clearer takeaway message\n\nWe strongly but respectfully disagree with the reviewer’s assumption that there should be a simple, straightforward result. None of the neural network architectures we train today are simple; each comprises a multitude of small design choices and tweaks that combine to lead to a strong result. ResNets, for example, rely on a combination of residual connections, batchnorm, global average pooling, specific initialization schemes, and particular sets of data augmentation. The goal of our paper is establishing a new baseline for future work on improving BERT, and that baseline must combine all of the best techniques available. Our proposed baseline is no more complicated than any other state-of-the-art architecture in use today.\n\n> Their primary findings mostly come from tuning the hyperparameters of GLU.\n\n\n\nOur primary contribution is building and extensively benchmarking one of the fastest architectures for BERT pretraining. We believe this is an important contribution to the academic and engineering communities, as fast, cheap pretraining enables ML practitioners.\n\nWe believe that the reviewer has misunderstood our primary contribution to be the change of the dimension of the intermediate layer for the Gated Linear Unit.\n\n> What is “Pareto Optimal”? \n\nWe thank the reviewer for bringing to our attention that we do not explicitly define the term “Pareto optimal,” and have fixed this in the updated manuscript.\n\nIn order to make meaningful improvements in neural network training efficiency, ML practitioners\nmust be able to compare between different choices of network architectures, hyperparameters, and training algorithms. One straightforward way to do this is to characterize the tradeoff between accuracy and training time with a “tradeoff curve” or a “Pareto curve.” Pareto curves can be generated by varying the length of training for each model configuration; longer training runs take more time but tend to reach higher quality. For a fixed model and task configuration, this method of generating tradeoff curves is an estimate of the theoretical Pareto frontier, i.e. the set of all of the best possible tradeoffs between training time and accuracy, where any further attempt to improve one of these metrics worsens the other.\n\nWe therefore consider one model Pareto optimal relative to another if it has superior accuracy across different training budgets while keeping everything else fixed. Many studies advertise novel architecture approaches without measuring wall clock time, or show an improvement for a single training duration. In our paper we show that RapidBERT-Base is Pareto optimal relative to our BERT-Base baseline for both short training durations and long training durations (Figure 2). We also show that BERT-Large and RapidBERT-Large are not Pareto optimal relative to BERT-Base and RapidBERT-Base for training durations under ~30 hours.\n\n> throughput difference in Figure 5\n\nIn Figure 5A, “lpLN+GLU+ALiBi” is bert-base with the additional pieces of low precision LayerNorm, GLU and ALiBi on top of the classic BERT architecture. In Figure 5B, “GLU 3072” indicates the complete RapidBERT architecture, which includes full-sized GLU, unpadding, FlashAttention, optimal vocab size, masked language modeling ratio of 30%, etc. We indicated this in the caption for Figure 5B “Throughput of the “complete” RapidBERT-Base with different intermediate sizes for GLU,” but will elaborate on this to be more clear.\n"}, {"review_id": "5AsiSMoHuj", "reviewer": "Reviewer_VemN", "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.\nResults 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.", "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.", "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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- This work employs some recently released transformer-efficient architectures in BERT training, getting RapidBERT.\n- 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).\n- 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.\n- 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?\n- 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.)\n\n[1] https://github.com/NVIDIA/Megatron-LM\n\n[2] https://github.com/microsoft/DeepSpeed", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments.\n\n> The authors should also post the RapidBERT results pre-trained on the English Wikipedia and the Books Corpus\n\nThe primary focus of our work was to show that certain architectural modifications and training choices lead to both a speed up as well as an improvement in downstream accuracy. We were less concerned with the well-worn path of “beating” the GLUE accuracy of the original BERT, or with getting the highest GLUE scores (which requires larger models trained for much longer). Even if we had chosen to train on 40 epochs of English Wikipedia and Books Corpus, it would be difficult to benchmark directly against the original BERT, which had many differences. For example, our BERT baseline and many other BERT variants such as RoBERTa include a MLM pretraining objective instead of MLM + NSP, a larger batch size (the original paper used a batch size of 256), etc.\n\nWe chose to run all experiments on the same, high quality, contemporary dataset (C4) so that we could explore the combined effects of architectures and training configuration (as opposed to data quality).\n\n> Similar studies have been proposed…comprehensive comparison with them in detail to help readers learn more about this field.\n\nWe have included an extended comparison with similar studies in the updated manuscript.\n\n> 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...could you explain why such an operation can lead to a speedup in the training of RapidBERT?\n\nThis was simply an empirical observation that we made when benchmarking our models, even in the absence of model parallelism. GPUs compute matrix multiplication by breaking matrices into tiles, and padding the vocab size aligns the matrix dimensions to a multiple of a fast tile size. CUDA has a heuristic which it uses to pick what tile size it uses for matrix multiplication, and it often chooses a suboptimal one when there is a “weird” matrix side length. Because of this, we and others have found that manually padding matrix dimensions can lead to non-trivial performance gains.\n\n> it would be better to discuss differences with these similar works\n\nWe thank the reviewer for this comment, and have updated our manuscript to include a discussion of DeepSpeed and MEGATRON\n\n> some recently released efficient transformer architectures (eg, RoPE, DeepNorm, etc)\n\nWe have a growing list of proposed efficiency improvements to transformer architectures, and are excited to explore this in future work.\n"}]

Author rebuttal:
We ran further ablations in response to reviewer comments and have plotted downstream GLUE accuracy as a function of measured pretraining wall clock time. The comments here describe the additional experiments in detail.

The patterns in Figures R1 and R2 shed light on the individual effects of various architectures (e.g. BERT+GLU, BERT+low precision LayerNorm) and training configurations (e.g. BERT + 30% masking ratio). On average, all methods seem to provide a slight accuracy boost to BERT-base. Increasing the masking ratio to 30% leads to a slight accuracy boost while not affecting the WCT, while turning off dropout in the attention layer (BERT+drpt=0) leads to a slight improvement in both accuracy and WCT. Low precision LayerNorm (BERT+lpLN) leads to a significant speedup (i.e. a shift to the left). Gated Linear Units (BERT+GLU) add more parameters to the model and lead to a significant slowdown while providing an accuracy boost. As a point of reference, we also benchmark the full RapidBERT as well as RapidBERT without FlashAttention and with attention dropout set to 0.1 (the standard BERT-base  configuration).

All BERT-Base models here were pretrained on C4 for 70,000 steps with batch size 4096, and microbatch size 256 on 8xA100 80GB GPUs. All models were initialized with the same seed and shared all other hyperparameters including the bert-base uncased tokenizer, the learning rate schedule, AdamW as the optimizer, etc. Final pretraining checkpoints were then finetuned on GLUE following the details in the appendix of our paper. The points represented in these GLUE plots are final finetuning checkpoints.

These plots highlight the importance of benchmarking with Pareto curves, as it is not possible to tell from these plots alone whether training BERT-base for 2 more hours leads to better performance than BERT+GLU (for example).

* BERT-base: standard BERT-base (110M parameters) with attention dropout=0.1 and feedforward dropout=0.1, vocab size set to 30522, MLM=15% (all Hugging Face standard configurations)

* BERT+drpt=0: standard BERT-base, except that the attention in the dropout layer is set to 0 instead of the default 0.1

* BERT+GLU: standard BERT-base, with GLU for the feedforward component of the encoder block

* BERT+lpLN: standard BERT-base, except with low precision LayerNorm (bfloat16)

* BERT+mlm30: standard BERT-base, except with a masked language modeling masking ratio of 30%

* RapidBERT: the complete RapidBERT detailed in the paper, including GLU (where the dimension of the intermediate layer is 3072), ALiBi, low precision LayerNorm, unpadding, MLM 30%, vocab size 30528 (a multiple of 64) and the attention dropout=0.

* RapidBERT-FlashAttn+drpt=0.1: RapidBERT _without_ Flash Attention and _with_ the attention dropout=0.1


Meta-review: The paper introduces RapidBERT, a BERT-style encoder architecture and training recipe optimized for fast pretraining in NLP research. Incorporating various advanced techniques and best practices, RapidBERT achieves competitive performance with significantly reduced pretraining time and cost, e.g., the base model achieves the downstream average GLUE score of the original BERT study in 1.13 hours on 8 A100 80 GB GPUs at a cost of roughly $20 on a standard cloud provider.

Overall it is a solid work with solid and replicable results, and potential value for practitioners. Added experiments (e.g., beyond GLUE), ablations, benchmarking (w.r.t. RoBERTa and CrammingBERT), and simplification were advised to strengthen the methodology, analysis, and contributions.
Decision: Accept (poster)