加速训练定制化BERT，ExtremeBERT工具包助力轻松快速预训练

2022-12-06 05:12

前言

预训练语言模型为自然语言处理带来了革命性的进展，预训练-微调的范式提高了许多下游任务的性能。在实际中，我们有时并不满足于在自己的数据上微调预训练模型，而是想拿一些内部数据和领域特定数据去重新训练一个自己的模型，以此获取更好的表征。

然而预训练动辄需要几十上百张 GPU，训练速度也非常缓慢。因此，即使是在 BERT 已经提出四年之后的今天，个人研究者和小型科研机构想要从头训练一个 BERT 依然十分困难。虽然我们可以很容易地拿到开源训练代码，但是数据量大、计算资源短缺给语言模型预训练的研究和应用带来了困难和挑战。

在本文中，来自港科大统计与机器学习实验室的研究者们提出了 ExtremeBERT 工具包，极大加快了预训练的速度，只需要 8 张 3090 GPU，就可以在一天左右从头训练一个模型达到 BERT-base 的效果，训练加速接近一个数量级。该工具包提供了灵活易用的接口、具有一键轻松训练、优化加速、支持海量数据集的特点，适合预训练研究和训练 domain specific 的模型。

论文标题：

ExtremeBERT: A Toolkit for Accelerating Pretraining of Customized BERT

论文链接：

https://arxiv.org/abs/2211.17201

代码链接：

https://github.com/extreme-bert/extreme-bert

文档链接：

https://extreme-bert.github.io/extreme-bert-page

工具包特点

2.1 易用：一键预训练

预训练流程步骤多且复杂，一般分为环境搭建、数据预处理、预训练、微调和结果后处理。大多数预训练框架需要用户手动输入每个阶段的命令去运行，对于新人用户尤其不友好。单单一个环境搭建可能就会劝退很多小白（安装各种环境依赖，nvidia/apex 需要 C++ 编译）。

ExtremeBERT 提供了一键式预训练，只需要一行命令，运行所有阶段，包括上述提到的全部五个步骤。预训练配置也支持灵活修改，用户可以根据自己需要修改 .yaml 配置。ExtremeBERT 支持自动安装，让用户远离底层 CUDA 配置、编译依赖，轻松安装所有依赖库。

▲ 图1：一键完成 ExtremeBERT 预训练的五个步骤

▲ 图2：灵活的参数配置。左图：训练 BERT Large 模型的 YAML 配置文件 bert-large.yaml；右图：自定义 BERT 模型的 YAML 配置文件示例 bert-customized.yaml

2.2 高效：优化加速

为了实现训练加速，ExtremeBERT 引入了多种加速技巧。包括模型过参数化、混合精度训练、大 batch 训练等。在优化方面，ExtremeBERT 使用了 Pan 等人提出的 elastic step decay [1] 来调整学习率, 进一步提升了训练收敛速度，同时还带来了效果上的提升。

2.3 定制化：支持海量数据集

ExtremeBERT 同时支持标准数据和用户自定义数据集，具有良好的扩展性。首先，它无缝接入 huggingface datasets 接口，可以支持 huggingface 官方提供的上千个数据集来进行预训练。

其次，用户自定义数据是重中之重，为了简化数据处理，ExtremeBERT 只需要用户将数据整理成每行一篇文章的纯文本格式，后续可以自动接入预先配置好的自动化预处理流程，一键完成数据掩码（pre-masking）、打乱（shuffling）、切分（sharding）和 token化（tokenization）。

实验

ExtremeBERT 比 BERT 的原始实现提供了显著的加速。作者将训练出来的模型与BERT Base 和 BERT Large 进行了实验对比和分析。在 base 和 large 两种模型中，ExtremeBERT 均达到或超过 BERT 的 GLUE 得分，如表 1 所示。实验结果表明，ExtremeBERT 提供了大致一个数量级的加速，使研究人员和工程师可以使用有限的资源从头开始预训练 BERT 模型。

▲ 表1：GLUE 基准评测结果

总结和展望

在本文中，作者提出了 ExtremeBERT，一个加速模型模型预训练的工具包。作者承诺未来会持续不断地维护，加入更多好的特性。包括但不限于：更多模型的支持（例如 RoBERTa、DeBERTa）、更多更大的预训练数据集（例如 Pile）和更多更好的优化算法。目前 ExtremeBERT 正在 Beta Test，大家在使用过程中有任何问题，欢迎随时提 issue 告诉作者。作者也欢迎提PR、以及更多 Contributor 帮助 ExtremeBERT 变的更强。

全部代码均已开源，可以点击链接体验:

https://github.com/extreme-bert/extreme-bert

参考文献

[1] Pan, R., Ye, H., & Zhang, T. (2021, September). Eigencurve: Optimal Learning Rate Schedule for SGD on Quadratic Objectives with Skewed Hessian Spectrums. In International Conference on Learning Representations.

更多阅读