SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jacob Portes $^{1*}$   
jacob.portes@
Alex Trot<sup>1*</sup> alex.trott@
Sam Havens<sup>1</sup>  
sam.havens@
Daniel King<sup>1</sup>  
daniel.king@
Abhinav Venigalla<sup>1</sup> abhi@
Moin Nadeem $^{2\dagger}$   
moinnadeem@
Nikhil Sardana<sup>1</sup> nikhil.sardana@
Daya Khudia<sup>1</sup>  
daya.khudia@
Jonathan Frankle<sup>1</sup> jfrankle@
$^{1}$ MosaicML × Databricks  $^{1}$ @databricks.com  $^{2}$ @moinnadeem.com
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.
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 (LLMs), 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 LLMs [30, 35, 27, 61, 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, 11]. In this study we empirically show that these speed optimizations can successfully be incorporated into the classic BERT architecture and training recipe.
*equal contribution. Corresponding author jacob.portes@databricks.com. Code can be found at mosaicbert.github.io
$\dagger$  work done while at MosaicML.
B
Figure 1: (A) Schematic of MosaicBERT architecture (B) Pareto curves of average GLUE (dev) scores for MosaicBERT-Base and the standard BERT-Base. Error bars indicate  $95\%$  confidence interval over  $n = 5$  pretraining seeds. All training was on  $8\times \mathrm{A}100 - 80\mathrm{GB}$  GPUs. FlashAttention schematic adapted from [11], 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 a 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 4 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 a competitive BERT-Base model on a 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 an 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 on the classic GLUE benchmark [60].
It can often be difficult to understand the effects of model architecture modifications and training hyperparameter choices by simply reporting a few numbers in a table, as is traditionally done in ML papers. A few numbers in a 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 to 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 1B). 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 to an increase in accuracy but take a much longer time to converge (e.g. increasing the model size). Accuracy vs. training time Pareto curves enable us to adequately asses all these changes (see Appendix for a more detailed discussion on Pareto curves).
- We implement a new BERT-style encoder architecture that optimizes pretraining speed and accuracy. This architecture combines FlashAttention [11], ALiBi [44], Gated Linear Units [12, 50], a 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  $8 \times$  A100 80 GB GPUs at a cost of roughly $20 on a 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-Base and Large is Pareto optimal relative to BERT-Base and Large.
- We characterize the relative throughput properties of each 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-Base performance after extensive training. We open-source our model weights and code at mosaicbert.github.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 [39, 38].
Model Architecture	GLUE (dev)	Training (hours)	Hardware	Pretraining Corpus	Params
MosaicBERT-Base (ours)	79.6	1.13	8×A100-80	C4	137M
MosaicBERT-Base (ours)	82.2	2.81	8×A100-80	C4	137M
MosaicBERT-Base (ours)	83.2	4.6	8×A100-80	C4	137M
MosaicBERT-Base (ours)	83.4	5.27	8×A100-80	C4	137M
BERT-Base (our benchmark)	83.2	11.5	8×A100-80	C4	110M
BERT-Base [14, 18]	81.0	96	16×TPU	Wiki+BookC	110M
BERT-Base [18]	73.6	24	1 RTX A6000	Wiki+BookC	110M
CrammingBERT-Base [18]	80.4	24	1 RTX A6000	Wiki+BookC	110M
BERT-Large [14, 37]	84.1	96	64×TPU	Wiki+BookC	340M
Table 1: Average GLUE (dev) score across various efficient BERT implementations. Average includes all 8 GLUE tasks. BERT-Base average GLUE scores on 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].
In order to build MosaicBERT, we incorporated architectural choices from the recent transformer literature. These include FlashAttention [11], ALiBi [44], training with dynamic unpadding [66], low precision LayerNorm, and Gated Linear Units [12, 50] (Figure 1). Before describing these modifications in detail, we first review the classic BERT architecture and how we chose a 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 [9, 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 a 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 a hidden size of 1024, an intermediate size of 4096, 16 attention heads, and 24 hidden layers.
3The Hugging Face bert-base-uncased model has been downloaded 62 million times, more than any other model on the Hugging Face hub.
While MosaicBERT-Base (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 Attention Mechanism
**FlashAttention:** 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 SRAM (i.e. short-term memory) [11] (see also [45]). We modified the FlashAttention module built by Hazy Research with OpenAI's triton library in order to flexibly incorporate ALiBi [57].<sup>4</sup>
Attention with Linear Biases (ALiBi): In most BERT models, the positions of tokens in a 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 a 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  $i$ th query  $q_{i} \in \mathbb{R}^{d}$  and keys  $\mathbf{K} \in \mathbb{R}^{L \times d}$  where  $d$  is the head dimension and  $L$  is the sequence length. ALiBi adds a fixed bias with  $m$  as a head-specific slope controlling how the bias grows with absolute distance between tokens, yielding attention weights as
$$
\operatorname {s o f t m a x} \left(q _ {i} \mathbf {K} ^ {\top} - m \cdot \operatorname {a b s} ([ i - 1, i - 2, \dots , i - L ])\right). \tag {1}
$$
The slopes  $m$  follow a geometric sequence such that for  $n$  heads, each head has a ratio of  $2^{-8/n}$  [44].
During finetuning or inference, the static bias can be increased to accommodate longer sequence lengths. For example, a model pretrained using ALiBi with a 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 a model with a maximum sequence length of 128 has much higher throughput than pretraining a model with a 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 transformer. GLUs were first proposed in 2016 [12], and incorporate an extra learnable matrix that "gates" the outputs of the feedforward layer (Figure 1A). More recent work has shown that GLUs can improve performance quality in transformers [50, 42]. We used the GeLU (Gaussian-error Linear Unit)<sup>5</sup> activation function with GLU, which is sometimes referred to as GeGLU. The module can be described by the following equation
$$
\left(\operatorname {G e L U} \left(\mathbf {x W} _ {1}\right) \otimes \mathbf {x V}\right) \mathbf {W} _ {2}, \tag {2}
$$
where  $\mathbf{x}$  is the input to the feedforward layer and the matrix  $\mathbf{V}$  gates the output of the GeLU activation function. The extra gating matrix in a GLU model potentially adds additional parameters to a model; we chose to augment our MosaicBERT-Base model with additional parameters due to GLU modules, as it leads to a Pareto improvement across all timescales. While BERT-Base has 110 million parameters, MosaicBERT-Base has 137 million parameters. Note that MosaicBERT-Base reaches higher accuracy faster than BERT-Base despite having more parameters (Figure 1B). Similarly, BERT-Large has 340 million parameters, and MosaicBERT-Large has 430 million parameters.
$^{4}$ github.com/HazyResearch/flash-attention. Note that while this research was being completed, PyTorch 2.0 was released with support for FlashAttention. However, FlashAttention in PyTorch 2.0 does not currently support ALiBi integration.
${}^{5}\mathrm{{GeLU}}$  is a fully differentiable approximation to ReLU, and was used in the original BERT study [14].
2.2.3 Additional Modifications
Low Precision LayerNorm: LayerNorm is a bandwidth-bound operation, which means that its speed depends on how quickly data can be loaded from memory to the compute units for elementwise operations. Typically the LayerNorm operation is set to float32 precision, which requires 4-bytes per-element. In MosaicBERT, we modify LayerNorm modules to run in bffloat16 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 automatic-mixed-precision package does not, 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 bffloat16 precision LayerNorm.
Unpadding: Standard NLP practice is to combine text samples of different lengths into a batch, and pad the sequences with special padding tokens so that all sample sequence lengths are the same (Figure 1A). During training, however, this leads to many wasted operations on the padding tokens. In MosaicBERT, we take a different approach and instead concatenate all the examples from a minibatch into a single sequence of batch size 1. Results from NVIDIA $^6$  and others have shown that this approach leads to speed improvements during training, since operations are not performed on padding tokens [66].
MLM Masking Ratio and Dropout: We used the standard Masked Language Modeling (MLM) pretraining objective. While the original BERT paper also included a Next Sentence Prediction (NSP) task in the pretraining objective, subsequent papers have shown this to be unnecessary [37, 28]. For our BERT baselines, we used the standard  $15\%$  masking ratio. However, we found that a  $30\%$  masking ratio led to slight accuracy improvements in both pretraining MLM and downstream GLUE performance. We therefore included this simple change as part of our MosaicBERT training recipe. Recent studies have also found that this straightforward change can lead to downstream improvements [63, 1].<sup>7</sup>
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 a 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 [14] 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-Base and the baseline BERT-Base. This ensured that our wall clock time measurements were not hindered by data streaming.
Bfloat16 Precision: We use bfloat16 mixed precision training for all the models. bfloat16 is a 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, a matrix
$^{6}$ NVIDIA removes padding in their MLPerfv1.1 benchmarking results https://github.com/mlcommons/training_results_v1.1 in the padding.py file
<sup>7</sup>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.
github.com/openai/triton
github.com/mosaicml/streaming. It is a drop in replacement for PyTorch's IteratorDataset that allows for fast streaming from cloud storage.
Figure 2: Performance on individual GLUE (dev) finetuning tasks. Our MosaicBERT-Base consistently 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  $8 \times \mathrm{A}100 - 80\mathrm{GB}$  GPUs, and does not include finetuning time. Error bars are plotted with  $95\%$  confidence interval across  $n = 5$  pretraining seeds, and all models are trained for 70,000 steps with batch size 4096.
multiplication layer uses bfloat16 for the multiplication and 32-bit IEEE floating point (float32) for gradient accumulation. We found this to be more stable than using float16 mixed precision. $^{10}$
Vocab Size as a Multiple of 64: We increased the vocab size to be a multiple of 64 (i.e. from 30,522 to 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 is more efficient at computing matrix multiplications with these dimensions. Across all experiments, we use the standard BERT-Base and Large tokenizers.
Pretraining Hyperparameters: For all models, we use a global batch size of 4096, and microbatch size of 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-Base and the baseline BERT-Base. More hyperparameter details are included in the Appendix.
In our first set of experiments, we pretrained BERT-Base and MosaicBERT-Base for 70,000 steps of batch size 4096, which roughly corresponds to  $78\%$  of English C4 (Figure 1B and Figure 2). We ran experiments with  $n = 5$  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-Base Achieves 79.6 Average GLUE (dev) Score in 1.13 Hours
MosaicBERT-Base achieves the downstream average GLUE (dev) score of 79.6 in 1.13 hours on  $8 \times$  A100 80 GB GPUs at a cost of roughly $20 on a standard cloud provider. More details on cost estimates are included in the Appendix.
<sup>10</sup>The NVIDIA Apex library and Megatron [53] both use a form of low precision LayerNorm in their code, but it is not documented in any papers that we could find.
<sup>11</sup>Note that in the original BERT study, the vocabulary size was 30,000 [14]. However, the default vocab size for the bert-base-uncased tokenizer in HuggingFace is 30,522.
<sup>12</sup>With vocab size that is divisible by 64, the calculations go down a different kernel path with much higher occupancy.
Figure 3: Average GLUE (dev) score Pareto curves for MosaicBERT-Base and Large trained for roughly 2 epochs of C4 (i.e. 178,000 steps with batch size 4096 with maximum sequence length 128 tokens). MosaicBERT-Base and Large are Pareto optimal relative to BERT-Base and Large. All pretraining is done on  $8 \times \mathrm{A}100$  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-Base reached an average GLUE (dev) score of  $83.2\%$  in 11.5 hours (A100-80GB), while MosaicBERT reached the same accuracy in roughly 4.6 hours on the same hardware, which is roughly a  $2.38 \times$  speedup (Table 1 and Figure 1B).
3.2 MosaicBERT-Base is Pareto Optimal
As can be seen Figure 1B, MosaicBERT-Base 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 in Figure 2. MosaicBERT-Base outperforms BERT-Base in four out of eight GLUE tasks across pretraining durations.
The space of 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 the largest datasets in the GLUE benchmark, with 100k-400k training examples, and MosaicBERT-Base is strictly 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) is 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-Base on these tasks across training duration.
3.3 MosaicBERT-Large is Pareto Optimal
While BERT-Base is one of the most popular BERT models, BERT-Large comes in a close second. All of our model development was done on MosaicBERT-Base; we were therefore curious whether our architecture and pretraining choices also generalized to a larger model.
In a second set of experiments, we pretrained MosaicBERT-Base and Large as well as BERT-Base 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 as long as the duration of the models in Figures 1 and 2.