|
--- |
|
license: apache-2.0 |
|
datasets: |
|
- c4 |
|
language: |
|
- en |
|
inference: false |
|
--- |
|
|
|
# MosaicBERT-Base model |
|
|
|
MosaicBERT-Base is a custom BERT architecture and training recipe optimized for fast pretraining. |
|
MosaicBERT trains faster and achieves higher pretraining and finetuning accuracy when benchmarked against |
|
Hugging Face's [bert-base-uncased](https://huggingface.co/bert-base-uncased). |
|
|
|
This study motivated many of the architecture choices around MosaicML's [MPT-7B](https://huggingface.co/mosaicml/mpt-7b) and [MPT-30B](https://huggingface.co/mosaicml/mpt-30b) models. |
|
|
|
## Model Date |
|
|
|
March 2023 |
|
|
|
## Documentation |
|
|
|
* [Project Page (mosaicbert.github.io)](mosaicbert.github.io) |
|
* [Github (mosaicml/examples/tree/main/examples/benchmarks/bert)](https://github.com/mosaicml/examples/tree/main/examples/benchmarks/bert) |
|
* [Paper (NeurIPS 2023)](https://openreview.net/forum?id=5zipcfLC2Z) |
|
* Colab Tutorials: |
|
* [MosaicBERT Tutorial Part 1: Load Pretrained Weights and Experiment with Sequence Length Extrapolation Using ALiBi](https://colab.research.google.com/drive/1r0A3QEbu4Nzs2Jl6LaiNoW5EumIVqrGc?usp=sharing) |
|
* [Blog Post (March 2023)](https://www.mosaicml.com/blog/mosaicbert) |
|
|
|
## Community Adoption |
|
|
|
* [DNABERT-2-117M](https://huggingface.co/zhihan1996/DNABERT-2-117M) for genome classification |
|
|
|
## How to use |
|
|
|
```python |
|
import torch |
|
import transformers |
|
from transformers import AutoModelForMaskedLM, BertTokenizer, pipeline |
|
from transformers import BertTokenizer, BertConfig |
|
|
|
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') # MosaicBERT uses the standard BERT tokenizer |
|
|
|
config = transformers.BertConfig.from_pretrained('mosaicml/mosaic-bert-base') # the config needs to be passed in |
|
mosaicbert = AutoModelForMaskedLM.from_pretrained('mosaicml/mosaic-bert-base',config=config,trust_remote_code=True) |
|
|
|
# To use this model directly for masked language modeling |
|
mosaicbert_classifier = pipeline('fill-mask', model=mosaicbert, tokenizer=tokenizer,device="cpu") |
|
mosaicbert_classifier("I [MASK] to the store yesterday.") |
|
``` |
|
|
|
Note that the tokenizer for this model is simply the Hugging Face `bert-base-uncased` tokenizer. |
|
|
|
In order to take advantage of ALiBi by extrapolating to longer sequence lengths, simply change the `alibi_starting_size` flag in the |
|
config file and reload the model. |
|
|
|
```python |
|
config = transformers.BertConfig.from_pretrained('mosaicml/mosaic-bert-base') |
|
config.alibi_starting_size = 1024 # maximum sequence length updated to 1024 from config default of 512 |
|
|
|
mosaicbert = AutoModelForMaskedLM.from_pretrained('mosaicml/mosaic-bert-base',config=config,trust_remote_code=True) |
|
``` |
|
|
|
This simply presets the non-learned linear bias matrix in every attention block to 1024 tokens (note that this particular model was trained with a sequence length of 128 tokens). |
|
|
|
**To continue MLM pretraining**, follow the [MLM pre-training section of the mosaicml/examples/benchmarks/bert repo](https://github.com/mosaicml/examples/tree/main/examples/benchmarks/bert#pre-training). |
|
|
|
**To fine-tune this model for classification**, follow the [Single-task fine-tuning section of the mosaicml/examples/benchmarks/bert repo](https://github.com/mosaicml/examples/tree/main/examples/benchmarks/bert#fine-tuning). |
|
|
|
### [Update 1/2/2024] Triton Flash Attention with ALiBi |
|
|
|
Note that by default, triton Flash Attention is **not** enabled or required. In order to enable our custom implementation of triton Flash Attention with ALiBi from March 2023, |
|
set `attention_probs_dropout_prob: 0.0`. We are currently working on supporting Flash Attention 2 (see [PR here](https://github.com/mosaicml/examples/pull/440)) and replacing the custom triton implementation. |
|
|
|
|
|
### Remote Code |
|
|
|
This model requires that `trust_remote_code=True` be passed to the `from_pretrained` method. This is because we train using [FlashAttention (Dao et al. 2022)](https://arxiv.org/pdf/2205.14135.pdf), which is not part of the `transformers` library and depends on [Triton](https://github.com/openai/triton) and some custom PyTorch code. Since this involves executing arbitrary code, you should consider passing a git `revision` argument that specifies the exact commit of the code, for example: |
|
|
|
```python |
|
mosaicbert = AutoModelForMaskedLM.from_pretrained( |
|
'mosaicml/mosaic-bert-base', |
|
trust_remote_code=True, |
|
revision='24512df', |
|
) |
|
``` |
|
|
|
However, if there are updates to this model or code and you specify a revision, you will need to manually check for them and update the commit hash accordingly. |
|
|
|
## Model description |
|
|
|
In order to build MosaicBERT, we adopted architectural choices from the recent transformer literature. |
|
These include [FlashAttention (Dao et al. 2022)](https://arxiv.org/pdf/2205.14135.pdf), [ALiBi (Press et al. 2021)](https://arxiv.org/abs/2108.12409), |
|
and [Gated Linear Units (Shazeer 2020)](https://arxiv.org/abs/2002.05202). In addition, we remove padding inside the transformer block, |
|
and apply LayerNorm with low precision. |
|
|
|
### Modifications to the Attention Mechanism |
|
1. **FlashAttention**: Attention layers are core components of the transformer architecture. 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) [[Dao et al. 2022]](https://arxiv.org/pdf/2205.14135.pdf). We used the FlashAttention module built by |
|
[hazy research](https://github.com/HazyResearch/flash-attention) with [OpenAI’s triton library](https://github.com/openai/triton). |
|
|
|
2. **Attention with Linear Biases (ALiBi)**: In most BERT models, the positions of tokens in a sequence are encoded with a position embedding layer; |
|
this embedding allows subsequent layers to keep track of the order of tokens in a sequence. ALiBi eliminates position embeddings and |
|
instead conveys this information using a bias matrix in the attention operation. It modifies the attention mechanism such that nearby |
|
tokens strongly attend to one another [[Press et al. 2021]](https://arxiv.org/abs/2108.12409). In addition to improving the performance of the final model, ALiBi helps the |
|
model to handle sequences longer than it saw during training. Details on our ALiBi implementation can be found [in the mosaicml/examples repo here](https://github.com/mosaicml/examples/tree/main/examples/benchmarks/bert/src/bert_layers.py#L425). |
|
|
|
3. **Unpadding**: Standard NLP practice is to combine text sequences of different lengths into a batch, and pad the sequences with empty |
|
tokens so that all sequence lengths are the same. During training, however, this can lead to many superfluous operations on those |
|
padding tokens. In MosaicBERT, we take a different approach: we concatenate all the examples in a minibatch into a 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 (see for example [Zeng et al. 2022](https://arxiv.org/pdf/2208.08124.pdf)). |
|
Details on our “unpadding” implementation can be found [in the mosaicml/examples repo here](https://github.com/mosaicml/examples/tree/main/examples/benchmarks/bert/src/bert_padding.py). |
|
|
|
4. **Low Precision LayerNorm**: this small tweak forces LayerNorm modules to run in float16 or bfloat16 precision instead of float32, improving utilization. |
|
Our implementation can be found [in the mosaicml/examples repo here](https://docs.mosaicml.com/projects/composer/en/stable/method_cards/low_precision_layernorm.html). |
|
|
|
### Modifications to the Feedforward Layers |
|
|
|
5. **Gated Linear Units (GLU)**: We used Gated Linear Units for the feedforward sublayer of a transformer. GLUs were first proposed in 2016 [[Dauphin et al. 2016]](https://arxiv.org/abs/1612.08083), |
|
and incorporate an extra learnable matrix that “gates” the outputs of the feedforward layer. More recent work has shown that |
|
GLUs can improve performance quality in transformers [[Shazeer, 2020](https://arxiv.org/abs/2002.05202), [Narang et al. 2021](https://arxiv.org/pdf/2102.11972.pdf)]. We used the GeLU (Gaussian-error Linear Unit) |
|
activation function with GLU, which is sometimes referred to as GeGLU. The GeLU activation function is a smooth, fully differentiable |
|
approximation to ReLU; we found that this led to a nominal improvement over ReLU. More details on our implementation of GLU can be found here. |
|
The extra gating matrix in a GLU model potentially adds additional parameters to a model; we chose to augment our BERT-Base model with |
|
additional parameters due to GLU modules as it leads to a Pareto improvement across all timescales (which is not true of all larger |
|
models such as BERT-Large). While BERT-Base has 110 million parameters, MosaicBERT-Base has 137 million parameters. Note that |
|
MosaicBERT-Base trains faster than BERT-Base despite having more parameters. |
|
|
|
|
|
|
|
|
|
## Training data |
|
|
|
MosaicBERT is pretrained using a standard Masked Language Modeling (MLM) objective: the model is given a sequence of |
|
text with some tokens hidden, and it has to predict these masked tokens. MosaicBERT is trained on |
|
the English [“Colossal, Cleaned, Common Crawl” C4 dataset](https://github.com/allenai/allennlp/discussions/5056), which contains roughly 365 million curated text documents scraped |
|
from the internet (equivalent to 156 billion tokens). We used this more modern dataset in place of traditional BERT pretraining |
|
corpora like English Wikipedia and BooksCorpus. |
|
|
|
## Pretraining Optimizations |
|
|
|
Many of these pretraining optimizations below were informed by our [BERT results for the MLPerf v2.1 speed benchmark](https://www.mosaicml.com/blog/mlperf-nlp-nov2022). |
|
|
|
1. **MosaicML Streaming Dataset**: As part of our efficiency pipeline, we converted the C4 dataset to [MosaicML’s StreamingDataset format](https://www.mosaicml.com/blog/mosaicml-streamingdataset) and used this |
|
for both MosaicBERT-Base and the baseline BERT-Base. For all BERT-Base models, we chose the training duration to be 286,720,000 samples of sequence length 128; this covers 78.6% of C4. |
|
|
|
|
|
2. **Higher Masking Ratio for the Masked Language Modeling Objective**: 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 [Liu et al. 2019](https://arxiv.org/abs/1907.11692). |
|
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 simple |
|
change can lead to downstream improvements [Wettig et al. 2022](https://arxiv.org/abs/2202.08005). |
|
|
|
3. **Bfloat16 Precision**: We use [bf16 (bfloat16) mixed precision training](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus) for all the models, where a matrix multiplication layer uses bf16 |
|
for the multiplication and 32-bit IEEE floating point for gradient accumulation. We found this to be more stable than using float16 mixed precision. |
|
|
|
4. **Vocab Size as a Multiple of 64**: We increased the vocab size to be a multiple of 8 as well as 64 (i.e. from 30,522 to 30,528). |
|
This small constraint is something of [a magic trick among ML practitioners](https://twitter.com/karpathy/status/1621578354024677377), and leads to a throughput speedup. |
|
|
|
5. **Hyperparameters**: For all models, we use Decoupled AdamW with Beta_1=0.9 and Beta_2=0.98, and a weight decay value of 1.0e-5. |
|
The learning rate schedule begins with a warmup to a maximum learning rate of 5.0e-4 followed by a linear decay to zero. |
|
Warmup lasted for 6% of the full training duration. Global batch size was set to 4096, and microbatch size was 128; since global batch size was 4096, full pretraining consisted of 70,000 batches. |
|
We set the maximum sequence length during pretraining to 128, and we used the standard embedding dimension of 768. |
|
For MosaicBERT, we applied 0.1 dropout to the feedforward layers but no dropout to the FlashAttention module, as this was not possible with the OpenAI triton implementation. |
|
Full configuration details for pretraining MosaicBERT-Base can be found in the configuration yamls [in the mosaicml/examples repo here](https://github.com/mosaicml/examples/blob/main/examples/benchmarks/bert/yamls/main/mosaic-bert-base-uncased.yaml). |
|
|
|
|
|
## Evaluation results |
|
|
|
When fine-tuned on downstream tasks (following the [finetuning details here](https://github.com/mosaicml/examples/tree/main/examples/benchmarks/bert/yamls/finetuning/glue/mosaic-bert-base-uncased.yaml)), the MosaicBERT model achieves the following GLUE results: |
|
|
|
| Task | MNLI-(m/mm) | QQP | QNLI | SST-2 | CoLA | STS-B | MRPC | RTE | Average | |
|
|:----:|:-----------:|:----:|:----:|:-----:|:----:|:-----:|:----:|:----:|:-------:| |
|
| | 0.8495 | 0.9029 | 0.9074| 0.9246 | 0.5511 | 0.8927 | 0.9003 | 0.8136 | 0.8428 | |
|
|
|
Note that this is averaged over n=5 pretraining seeds. |
|
|
|
## Collection of MosaicBERT-Base models trained using ALiBi on different sequence lengths |
|
|
|
ALiBi allows a model trained with a sequence length n to easily extrapolate to sequence lengths >2n during finetuning. For more details, see [Train Short, Test Long: Attention with Linear |
|
Biases Enables Input Length Extrapolation (Press et al. 2022)](https://arxiv.org/abs/2108.12409) |
|
|
|
This model is part of the **family of MosaicBERT-Base models** trained using ALiBi on different sequence lengths: |
|
|
|
* mosaic-bert-base (trained on a sequence length of 128 tokens) |
|
* [mosaic-bert-base-seqlen-256](https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-256) |
|
* [mosaic-bert-base-seqlen-512](https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-512) |
|
* [mosaic-bert-base-seqlen-1024](https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-1024) |
|
* [mosaic-bert-base-seqlen-2048](https://huggingface.co/mosaicml/mosaic-bert-base-seqlen-2048) |
|
|
|
The primary use case of these models is for research on efficient pretraining and finetuning for long context embeddings. |
|
|
|
## Intended uses & limitations |
|
|
|
This model is intended to be finetuned on downstream tasks. |
|
|
|
## Citation |
|
|
|
Please cite this model using the following format: |
|
|
|
``` |
|
@article{portes2023MosaicBERT, |
|
title={MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining}, |
|
author={Jacob Portes, Alexander R Trott, Sam Havens, Daniel King, Abhinav Venigalla, |
|
Moin Nadeem, Nikhil Sardana, Daya Khudia, Jonathan Frankle}, |
|
journal={NeuRIPS https://openreview.net/pdf?id=5zipcfLC2Z}, |
|
year={2023}, |
|
} |
|
``` |
