[论文解读] Accelerated Large Batch Optimization of BERT Pretraining in 54 minutes

Shuai Zheng, Haibin Lin+2 more|arXiv (Cornell University)|Jun 24, 2020

Topic Modeling参考文献 33被引用 10

一句话总结

该论文提出 LANS，一种用于 BERT 预训练的加速大规模批量优化方法，结合了逐块梯度归一化与一种新型学习率调度策略，该策略在预热阶段后引入恒定学习率阶段。通过实现 96K 和 33K 的小批量大小下的稳定训练，LANS 在 192 个 AWS EC2 P3dn.24xlarge 实例上将 BERT 预训练时间缩短至 54 分钟，创下云环境中无精度损失的全新速度纪录。

ABSTRACT

BERT has recently attracted a lot of attention in natural language understanding (NLU) and achieved state-of-the-art results in various NLU tasks. However, its success requires large deep neural networks and huge amount of data, which result in long training time and impede development progress. Using stochastic gradient methods with large mini-batch has been advocated as an efficient tool to reduce the training time. Along this line of research, LAMB is a prominent example that reduces the training time of BERT from 3 days to 76 minutes on a TPUv3 Pod. In this paper, we propose an accelerated gradient method called LANS to improve the efficiency of using large mini-batches for training. As the learning rate is theoretically upper bounded by the inverse of the Lipschitz constant of the function, one cannot always reduce the number of optimization iterations by selecting a larger learning rate. In order to use larger mini-batch size without accuracy loss, we develop a new learning rate scheduler that overcomes the difficulty of using large learning rate. Using the proposed LANS method and the learning rate scheme, we scaled up the mini-batch sizes to 96K and 33K in phases 1 and 2 of BERT pretraining, respectively. It takes 54 minutes on 192 AWS EC2 P3dn.24xlarge instances to achieve a target F1 score of 90.5 or higher on SQuAD v1.1, achieving the fastest BERT training time in the cloud.

研究动机与目标

• 通过实现高效的超大规模批量优化，将 BERT 预训练的训练时间从传统数天大幅缩短。
• 克服现有大规模批量优化器（如 LAMB）在小批量大小超过 64K 时出现的训练不稳定与精度下降问题。
• 设计一种学习率调度策略，在不发散的前提下延长高学习率的使用时间，从而实现更快收敛。
• 在云环境中实现 SOTA 的 BERT 训练速度，同时不牺牲模型性能。

提出的方法

• LANS 引入逐块梯度归一化，扩展 LAMB 的分块自适应学习率机制，确保不同层之间更新的稳定性。
• 通过利用经典动量与 Nesterov 动量之间的联系，改进动量项的计算方式，以提升收敛性能。
• 采用新型学习率调度策略，包含预热后的恒定学习率阶段，将高学习率的使用时间延长至初始预热阶段之后。
• 学习率采用平方根规则进行缩放，但经过调整以避免在大批次下发散，通过经验调参确保稳定性。
• 框架支持分布式训练，通过在各工作节点间分片数据，最小化梯度方差并确保采样非冗余。
• 实现针对云基础设施的优化，使用 192 个 AWS EC2 P3dn.24xlarge 实例并启用 EFA，以实现高效的通信性能。

实验结果

研究问题

• RQ1 经修改的大规模批量优化器能否在 BERT 预训练中将小批量大小扩展至 96K 和 33K 的同时保持模型精度？
• RQ2 在预热后引入恒定学习率阶段是否相比线性预热调度策略能提升训练稳定性与收敛速度？
• RQ3 LANS 是否能在云环境的 GPU 集群上实现比 LAMB 更快的训练速度，且不造成精度下降？
• RQ4 在大批次下实现快速收敛的最优超参数设置（学习率、预热长度、恒定阶段比例）为何？

主要发现

• LANS 在 1536 块 V100 GPU 上仅用 53.6 分钟即在 SQuAD v1.1 开发集上达到 90.60 的 F1 分数，创下云环境中 BERT 预训练的最快纪录。
• 该方法成功将小批量大小扩展至 BERT 预训练第一阶段的 96K 和第二阶段的 33K，且未造成性能下降。
• LAMB 在 96K/33K 批次下无法收敛，而 LANS 仅用 4301 次迭代即达到相同目标 F1 分数，远少于 LAMB 的 8599 次迭代。
• 所提出的调度策略将 0.007 与 0.01 学习率之间的 AUC 差距从 5.28 个百分点降低至 1.91 个百分点，证明了更强的鲁棒性。
• 在预热后引入恒定学习率阶段，使高学习率下的训练进度得以持续，从而实现更快收敛。
• 在 192 个 AWS EC2 P3dn.24xlarge 实例上，训练时间缩短至 54 分钟，优于以往方法在 TPU 上所需 76.2 分钟的训练时间。