SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 to 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 is that the GLUE tasks are quite short by modern standards (circa 2023) and contain fewer than 512 tokens. Additionally, all GLUE finetuning 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 is possible, however, that these models are superior at tasks that involve long-context dependencies.
The original BERT paper pretrained with a 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 a 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].
We released the pretrained checkpoints for these models on the HuggingFace hub in April 2023:
- https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-256
- https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-512
- https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-1024
- https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-2048
Model	Bsz	Steps	SeqL	MNLI	QNLI	QQP	RTE	SST	MRPC	CoLA	STS	Av.
BERT-Large BookC + Wiki	256	1M	128	86.6	92.3	91.3	70.4	93.2	88	60.6	90	84.05
RoBERTa-Base all data+500k	8k	500k	512	87.6	92.8	91.9	78.7	94.8	90.2	63.6	91.2	86.35
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
2. all data	8k	100k	512	89.3	94	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.4	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	93	92	84.5	93.7	88.2	59.7	90.9	86.1
Table S3: Model performance comparison on GLUE (dev) tasks. Data from Table 8 of the RoBERTa paper [37]. Values for BERT and MosaicBERT in 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 D (data from Table 8 in the Appendix of the RoBERTa paper [37]). In particular, they trained on:
1. BookCorpus [68] and English Wikipedia for 100k steps with batch size 8192 and sequence length 512
2. BookCorpus [68] and English Wikipedia (16GB), plus "additional data" CC News, the English portion of the CommonCrawl News dataset (76GB) [41], OpenWebText 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 apples-to-apples comparison of RoBERTa and MosaicBERT experiments would keep all hyperparameters 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 of data constant. 100k steps with batch size 8192 is roughly equivalent to 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  $4 \times$  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  $4 \times$  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 a learning rate sweep for single-task GLUE finetuning  $\{1\mathrm{e} - 5,2\mathrm{e} - 5,3\mathrm{e} - 5\}$ , which we did not do. Finally, the differences between MosaicBERT and RoBERTa could also be due to differences in the underlying datasets.
F 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, we elaborate on the individual GLUE benchmarks for reference. We used the GLUE tasks from HuggingFace for finetuning and evaluation: https://huggingface.co/datasets/glue. All evaluation was done on the validation (a.k.a. dev) splits.
F.1 Large Finetuning Datasets
MNLI (Multi-Genre Natural Language Inference) [392,702 train | 19,643 validation | 19,643 test] is 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, or is neutral with respect to the first one. For example:
- Premise: "Buffet and a la carte available."
Hypothesis: "It has a buffet."
- Label: 0 (entailment)
MNLI has two subsets, "matched" (MNLI-m) and "mismatched" (MNLI-mm). All numbers reported in 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 question-paragraph pairs drawn from Wikipedia [47].
QQP (Quora Question Pairs 2) [363,846 train | 40,400 validation | 390,965 test]. The task is to determine whether two sentences are semantically equivalent [26].
F.2 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 sentence [10, 19, 4]. Note that in our work we use a checkpoint from the MNLI finetuning to finetune on RTE.
CoLA (Corpus of Linguistic Acceptability) [8,551 train | 1,040 validation | 1,063 test] [62] is a benchmark with sentences that are either linguistically acceptable or grammatically incorrect. For example:
- "The higher the stakes, the lower his expectations are." Label: 1 (acceptable)
"Mickey looked up it." Label: 0 (unacceptable)
SST-2 (Stanford Sentiment Treebank) [67,349 train | 872 validation | 1,821 test] consists of sentences from movie reviews. The task is to classify the sentiment as either positive or negative [54].
MRPC (Microsoft Research Paraphrase Corpus)[3,668 train | 408 validation | 1,725 test] [16] The dataset consists of sentence pairs extracted from online news sources. The task is to 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 to 5 [7].
Note that we excluded finetuning on the 9th GLUE task WNLI (Winograd NLI) [34], as in the original BERT study (it is a very small dataset [634 train, 146 test] with a high number of adversarial examples). Finetuning on RTE, MRPC and STSB starts from a checkpoint already finetuned on MNLI (following the example of [28] and other studies). This is done because all the above tasks deal with sentence pairs, and this staged finetuning leads to consistent empirical improvement.
G MosaicBERT-Large Multinode Throughput Scaling
The experiments in the main section of this paper were all performed on a single node with  $8 \times \mathrm{A}100$  GPUs. How well do our innovations to the BERT architecture maximize throughput at the multinode scale?
We measured the throughput of MosaicBERT-Large (430M) during training on 8, 16, 32, 64, 128 and 200 GPUs, and plotted the tokens per second for various global batch sizes. Global batch size is 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 to maintain high throughput.
If 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 S5; a global batch size of 4096 spread across 64 GPUs using Distributed Data Parallelism (DDP) means that each GPU will only apply matmul operations on matrices with a 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 on matrices with a dimension of 1024, leading to higher throughput. However, such a 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 to resource and time constraints.
G.1 How MosaicBERT modifications might scale to larger models
All of the architectural and configuration choices made in MosaicBERT 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 as well as MPT-7B and MPT-30B [39, 38]). For example, Gated Linear Units should have no issues with tensor parallelism as long as the matrix multiplications are a 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 as much.
Figure S5: MosaicBERT-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 are 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 forward/backwards pass of the model, and does not account for the additional compute required for
Model	Throughput (tokens /sec)	MFU	Hardware	Time to 79.6	Batch Size	Micro-batch Size
BERT Base	0.4e6	10.4%	8× A100 80	110.4 minutes (1.84 hours)	4096	512
MosaicBERT Base	1.1e6	39.97%	8× A100 80	67.8 minutes (1.13 hours)	4096	512
MosaicBERT Base	0.938e6	30.9%	8× A100 40	76.8 minutes	4096	128
MosaicBERT Base	1.88e6	31.0%	16× A100 40	38.5 minutes	4096	128
MosaicBERT Base	3.15e6	25.9%	32× A100 40	23.1 minutes	4096	128
MosaicBERT Base	4.77e6	19.6%	64× A100 40	15.7 minutes	4096	64
other implementation details such as activation checkpointing. Thus, MFU is independent of implementation and hardware. For more details, see [32]. All FLOP calculations exclude the operations required for normalization, activation, and residuals.
Following the notation in the PaLM paper [8], Model FLOPs Utilization (MFU) is approximated as:
$$
\mathrm {M F U} = \frac {\left(6 \cdot n _ {\text {p a r a m e t e r s}} \cdot T _ {\text {o b s e r v e d}}\right)}{n _ {\text {g p u s}} \cdot T _ {\text {t h e o r e t i c a l}}} \tag {3}
$$
where  $T_{\text{observed}}$  is the observed throughput and  $T_{\text{theoretical}}$  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 seen ( $2 \times$  for the forward pass and  $4 \times$  for the backward pass). This is then multiplied by the number of tokens seen per second. As a first-order approximation, we exclude the extra FLOPs per token due to dense self-attention.
In the denominator, the theoretical peak throughput is provided in the GPU hardware specs. For A100 GPUs using bf1oat16, this theoretical peak throughput is 312 teraFLOPs.
Table S4: Multinode Throughput scaling for MosaicBERT-Base
MosaicBERT-Base Ave. GLUE Score	8×A10080GB hours	8×A100 80GB cost ($2.50 GPU/hr)	8×A10040GB hours	8×A100 40GB cost ($2 GPU/hr)
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-Base 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 to $1.10 (40 GB) / $1.50 (80 GB) per GPU using Lambda labs. At an intermediate price of $2.50 an hour per A100 80 GB GPU, training to 79.6 GLUE average score takes 1.13 hours and costs roughly $22.60. $^{14}$  Some example costs are calculated in Table S5.
<sup>14</sup>See for example "Cloud GPU instances with the largest VRAM 2022" (medium.com/@aleixlopez/cloud-gpu-instances-to-solve-out-of-memory-error-2022-d5012883a272?)
C
Figure S6: Standard FeedForward Transformer Block and Gated Linear Unit Modifications. Each edge shows the tensor dimension with a batch size of 256, sequence length of 128 and a hidden dimension of 768. (A) A standard transformer feedforward block. (B) Naive implementation of a Gated Linear Unit. The number of parameters in this are the same as in (A). (C) Fused implementation of a Gated Linear Unit where the two matrix multiplications (Linear1_0 and Linear1_1) from (B) are fused into one (Linear1) with  $2 \times$  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 improvements over the standard Transformer block. There are multiple ways to implement GLUs and we experimented with two implementations. Figure S6 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 is expected to perform better in some domains.
Figure S7 shows the performance impact of the two GLU over standard feedforward transformer block (which would be  $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
Figure S7: Slowdown of different implementations of Gated Linear Unit. This slowdown is 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: 1) 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 GPU, this is in the regime of  $4096 / 8 = 512$ .
The main reason for slowness of GLUs 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-multiplications 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 we trained two different model sizes, we have not pretrained a MosaicBERT model in the  $>1\mathrm{B}$  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. It is possible that some of the Pareto properties change in the longer regime, although we suspect that this is unlikely.
K.2 Broader Impact
BERT models are highly used for NLP tasks. By open-sourcing 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.