突破 MLPerf 训练：优化 BERT 的案例研究
机器学习
2024-02-06 v1 计算与语言
摘要
加速大规模分布式训练具有挑战性，因为它需要改进训练的各个组件，包括负载均衡、通信、优化器等。我们提出了一种用于 BERT 模型快速大规模训练的新方法，该方法分别改进了每个组件，从而达到了新的 BERT 训练性能水平。在分布式 BERT 训练中，负载均衡至关重要，因为其训练数据集的特征是样本长度各异。与分布式训练规模成正比的通信成本需要被有用的计算所掩盖。此外，优化器（如 ADAM、LAMB 等）需要在大规模分布式训练的背景下进行仔细的重新评估。我们提出了两个新思路：（1）基于数据集分层的局部预排序以实现负载均衡；（2）在 allreduce 之前进行分桶梯度裁剪，这使我们能够受益于梯度计算与同步的重叠，以及 allreduce 之前进行梯度裁剪的快速训练。我们还通过超参数优化重新评估了现有的优化器，并采用了 ADAM，这通过比现有方法更大的批次量也促进了快速训练。我们提出的方法全部结合后，在 1,024 个 NVIDIA A100 GPU 上实现了最快的 MLPerf BERT 训练，耗时 25.1（22.3）秒，比 MLPerf v1.1（v2.0）的其他前两名（一名）提交结果分别快 1.33 倍（1.13 倍）和 1.57 倍。我们的实现和评估结果可在 MLPerf v1.1~v2.1 中获取。
引用
@article{arxiv.2402.02447,
title = {Breaking MLPerf Training: A Case Study on Optimizing BERT},
author = {Yongdeok Kim and Jaehyung Ahn and Myeongwoo Kim and Changin Choi and Heejae Kim and Narankhuu Tuvshinjargal and Seungwon Lee and Yanzi Zhang and Yuan Pei and Xiongzhan Linghu and Jingkun Ma and Lin Chen and Yuehua Dai and Sungjoo Yoo},
journal= {arXiv preprint arXiv:2402.02447},
year = {2024}
}
备注
Total 15 pages (Appendix 3 pages)