HEADLINE: MosaicBERT: Optimizing BERT for Faster Pretraining with FlashAttention and ALiBi

MosaicBERT, a custom BERT architecture developed by MosaicML in collaboration with Databricks, is designed to significantly speed up the pretraining process. This model introduces several key architectural modifications that can be applied to other transformer models, including MosaicML's own MPT-7B and MPT-30B. Here’s a deep dive into what makes MosaicBERT tick and why it matters for practitioners.

Key Architectural Modifications

1. FlashAttention:

   • FlashAttention is an efficient implementation of the attention mechanism that has quickly become the standard in transformer models.
   • It optimizes memory usage and parallelism, leading to faster training times without sacrificing accuracy.
   • This modification alone can lead to substantial speedups, especially on large datasets.

2. ALiBi (Attention with Linear Biases):

   • ALiBi allows you to train on shorter sequence lengths and then extrapolate to longer sequences.
   • This is particularly useful for reducing the computational load during pretraining while maintaining performance on longer sequences during inference.
   • By using linear biases, the model can learn positional information more effectively.

3. Gated Linear Units (GLUs):

   • GLUs are a type of activation function that can improve the efficiency and effectiveness of feedforward layers.
   • They help in controlling the flow of information through the network, leading to better performance and faster convergence.

4. Unpadding Inputs:

   • After the attention block, unpad your inputs to avoid wasting compute on PAD tokens.
   • This simple optimization can significantly reduce unnecessary computations, especially in datasets with varying sequence lengths.

5. Low Precision LayerNorm:

   • Set the LayerNorm operation to bfloat16 precision instead of the typical float32.
   • Using lower precision can speed up training without a significant loss in model performance.

Additional Efficiency Tips

• Increase Masked Language Modeling (MLM) Ratio:

  • Change the MLM ratio from the default 15% to 30%.
  • A higher masking rate can lead to more robust pretraining, although it may require more data to achieve the same level of performance.

• Remove Dropout in Attention Module:

  • Dropout is often used to prevent overfitting but can slow down training.
  • Removing dropout from the attention module can speed up the process without a significant impact on model quality.

• Use bfloat16 for Training:

  • Bfloat16 is a lower precision format that balances performance and accuracy.
  • It is particularly effective on modern GPUs and TPUs, leading to faster training times.

• Optimize Vocabulary Size:

  • Make the vocabulary size a multiple of 64.
  • This optimization aligns with hardware constraints and can lead to more efficient memory usage and faster training.

Why It Matters

Pretraining transformers from scratch is often prohibitively expensive and time-consuming. MosaicBERT's optimizations address these challenges by reducing computational requirements without compromising model performance. These modifications are not limited to BERT-style encoder models; many can be applied to decoder architectures like GPT and MPT as well.

For ML practitioners, adopting these techniques can lead to more efficient pretraining workflows, enabling faster experimentation and deployment of transformer models. Whether you're working on a small-scale project or scaling up to large datasets, the optimizations introduced by MosaicBERT are worth considering.

Source: mosaicbert.github.io