[论文解读] MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining

Jacob Portes, Alex Trott+7 more|arXiv (Cornell University)|Dec 29, 2023

Topic Modeling被引用 4

一句话总结

MosaicBERT 是一种针对快速预训练进行优化的 BERT 风格双向编码器，通过集成 FlashAttention、ALiBi、门控线性单元（GLU）、动态去填充（dynamic unpadding）以及低精度 LayerNorm 实现加速。其在 8 张 A100 80GB GPU 上仅用 1.13 小时即达到 79.6 的平均 GLUE 分数，成本约为 20 美元，相较于标准 BERT 模型，在准确率与训练时间之间实现了帕累托最优。

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.

研究动机与目标

• 为解决从零开始预训练 BERT 风格模型所面临的高计算成本与长训练时间问题。
• 系统性地将近期 Transformer 技术进展（如 FlashAttention、ALiBi 和 GLU）整合进经典 BERT 架构，以提升效率。
• 通过实证研究证明在多种模型规模下，准确率与预训练速度之间存在帕累托改进。
• 使研究人员和工程师能够以较低成本微调自定义 BERT 风格模型，而非依赖通用预训练模型。
• 开源模型权重与代码，以促进可复现性与更广泛的应用。

提出的方法

• 集成 FlashAttention，通过优化内核操作加速自注意力计算。
• 应用带线性偏置的注意力（ALiBi），实现相对位置编码而无需位置嵌入。
• 在前馈层中使用门控线性单元（GLU），并采用融合实现方式以提升表征能力。
• 引入动态去填充模块，跳过对填充 token 的计算，提升可变长度序列的吞吐量。
• 采用低精度 LayerNorm（bfloat16）以减少内存带宽与计算开销。
• 为掩码语言建模（MLM）目标采用 30% 的掩码比例，并对词表大小进行优化以提升 GPU 吞吐量。

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$ A100-80GB GPUs. FlashAttention schematic adapted fro

实验结果

研究问题

• RQ1 能否有效将现代 Transformer 变体中的架构创新整合进经典 BERT 架构中，以在不损失准确率的前提下加速预训练？
• RQ2 在 BERT 风格模型中结合多种效率技术时，模型准确率与预训练速度之间的权衡关系如何？
• RQ3 与标准 BERT-Base 和 BERT-Large 相比，MosaicBERT 在不同预训练时长下是否实现了帕累托最优？
• RQ4 各个独立架构组件（如 FlashAttention、GLU、去填充）对吞吐量与收敛速度的贡献如何？
• RQ5 将 BERT-Base 模型训练至 BERT-Large 水平的性能是否需要显著更长的训练时间？这一问题能否通过架构改进缓解？

主要发现

• MosaicBERT-Base 在 8 张 A100 80GB GPU 上仅用 1.13 小时即在下游任务中达到 79.6 的平均 GLUE（开发集）分数，总成本约为 20 美元。
• 该性能在使用 8× A100 80GB GPU、30% MLM 掩码比例与 bfloat16 精度时，训练成本为 22.60 美元。
• MosaicBERT-Base 与 MosaicBERT-Large 在多个预训练时长下，相较于标准 BERT-Base 与 BERT-Large，均表现出实证上的帕累托最优性。
• 吞吐量消融实验表明，FlashAttention、ALiBi 与动态去填充均对训练加速有显著贡献。
• 融合 GLU 实现方式虽在某些场景下略慢于标准 GLU，但因其更强的表征能力且兼容高批量训练，仍被采用。
• BERT-Large 仅在经过长时间训练后才超越 BERT-Base 的性能，凸显了架构优化带来的效率提升。

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$ A100-80GB GPUs, and does not include finetuning time. E