54 分钟加速大规模批量优化 BERT 预训练
机器学习
2020-09-21 v2 计算与语言
分布式、并行与集群计算
机器学习
摘要
BERT 近期在自然语言理解（NLU）中备受关注，并在多项 NLU 任务上取得了最先进（state-of-the-art）结果。然而，其成功依赖于大型深度神经网络与海量数据，导致训练时间长并阻碍开发进展。使用具有大 mini-batch 的随机梯度方法被提倡为减少训练时间的高效工具。在此研究路线中，LAMB 是一个突出例子，它将 BERT 在 TPUv3 Pod 上的训练时间从 3 天缩短至 76 分钟。本文中，我们提出一种称为 LANS 的加速梯度方法，以改进大 mini-batch 训练的使用效率。由于学习率理论上以函数 Lipschitz 常数倒数为上界，不能通过选取更大学习率来总是减少优化迭代次数。为在不损失精度下使用更大 mini-batch 规模，我们开发了一种新的学习率调度器，克服了使用大学习率的困难。利用所提 LANS 方法与学习率方案，我们在 BERT 预训练的阶段 1 与阶段 2 分别将 mini-batch 规模扩展至 96K 与 33K。在 192 个 AWS EC2 P3dn.24xlarge 实例上达到 SQuAD v1.1 上 90.5 或更高的目标 F1 分数耗时 54 分钟，实现了云上最快的 BERT 训练时间。
引用
@article{arxiv.2006.13484,
title = {Accelerated Large Batch Optimization of BERT Pretraining in 54 minutes},
author = {Shuai Zheng and Haibin Lin and Sheng Zha and Mu Li},
journal= {arXiv preprint arXiv:2006.13484},
year = {2020}
}
备注
Technical Report (not under reviewed in any venue)