[论文解读] Multi-node Bert-pretraining: Cost-efficient Approach

Jiahuang Lin, Xin Li+1 more|arXiv (Cornell University)|Aug 1, 2020

Topic Modeling参考文献 26被引用 7

一句话总结

本文提出了一种成本高效的多节点 BERT-large 预训练框架，在 32 个节点的学术集群上使用 256 张 NVIDIA T4 GPU，12 天内完成训练。通过算法与软件优化，提升了单设备吞吐量，高效分配训练工作负载，并缓解了 10 Gbps 网络下的通信瓶颈，实现了 70% 的弱扩展效率，且模型性能与工业级设置相当，但成本仅为后者的几分之一。

ABSTRACT

Recently, large scale Transformer-based language models such as BERT, GPT-2, and XLNet have brought about exciting leaps in state-of-the-art results for many Natural Language Processing (NLP) tasks. One of the common trends in these recent models is a significant increase in model complexity, which introduces both more weights and computation. Moreover, with the advent of large-scale unsupervised datasets, training time is further extended due to the increased amount of data samples within a single training epoch. As a result, to train these models within a reasonable time, machine learning (ML) programmers often require advanced hardware setups such as the premium GPU-enabled NVIDIA DGX workstations or specialized accelerators such as Google's TPU Pods. Our work addresses this limitation and demonstrates that the BERT pre-trained model can be trained within 2 weeks on an academic-size cluster of widely available GPUs through careful algorithmic and software optimizations. In this paper, we present these optimizations on how to improve single device training throughput, distribute the training workload over multiple nodes and GPUs, and overcome the communication bottleneck introduced by the large data exchanges over the network. We show that we are able to perform pre-training on BERT within a reasonable time budget (12 days) in an academic setting, but with a much less expensive and less aggressive hardware resource requirement than in previously demonstrated industrial settings based on NVIDIA DGX machines or Google's TPU Pods.

研究动机与目标

• 在硬件预算有限的学术环境中，实现高效且成本低廉的 BERT-large 预训练。
• 解决训练大规模 Transformer 模型（如 BERT-large）所带来的高计算与财务门槛。
• 优化商品化硬件上的单 GPU 吞吐量与多节点分布式训练性能。
• 通过梯度累积与系统级调优，降低分布式训练中的通信开销。
• 证明通过显著更低的硬件成本，也能实现与工业基准相当的最先进模型质量。

提出的方法

• 使用由 32 个节点组成的集群，每个节点配备 8 张 NVIDIA T4 GPU，总计 256 张 GPU，通过 10 Gbps 网络连接。
• 在 4 个步骤内应用梯度累积，以减少节点间通信频率与网络流量。
• 通过混合精度训练与高效的内存管理，优化单 GPU 训练吞吐量。
• 实施两阶段预训练：第一阶段使用序列长度 128（36 个周期），第二阶段使用序列长度 512（6 个周期），分别采用不同的批量大小与初始学习率。
• 在节点间采用模型并行与数据并行，使用同步随机梯度下降进行参数更新。
• 通过损失曲线分析与超参数调优，诊断第二阶段的收敛问题并提升训练稳定性。

实验结果

研究问题

• RQ1 仅使用学术级硬件与 10 Gbps 网络，能否在 12 天内完成 BERT-large 的预训练？
• RQ2 在商品化多 GPU 集群上实现高扩展效率，需要哪些系统级优化？
• RQ3 在分布式 BERT 预训练中，梯度累积如何影响通信开销与训练吞吐量？
• RQ4 为何 BERT 预训练的第二阶段表现出收敛不稳定性？是否可在不增加硬件成本的前提下解决？
• RQ5 在低功耗 T4 GPU 上训练时，与高端 DGX 或 TPU 系统相比，模型性能能保留到何种程度？

主要发现

• 作者在 32 个节点的学术集群上，使用 256 张 T4 GPU，12 天内完成 BERT-large 预训练，实现了 70% 的弱扩展效率。
• 总硬件成本约为 62.4 万美元，仅为同等规模 DGX 或 TPU Pod 设置的 1/24（1500 万至 1280 万美元）。
• 微调后，模型在 SQuAD v1.1 上的 F1 得分为 81% 至 83%，与 Google 的 90.9% 和 NVIDIA 的 90%–91% 相差仅 9%–10%。
• 第二阶段的收敛问题源于超参数配置不当，而非系统优化问题；经更优调优后，第二阶段可由 6 个周期缩短至 4 个周期。
• 优化与非优化运行的损失曲线显示，性能优化未损害收敛性，且优化运行在第二阶段表现更稳定。
• 在 Google Cloud 上租用 256 张 T4 GPU 12 天的成本为 25,739 美元，仅为硬件购置成本的 1/24，且硬件寿命达 3 年，可支持 90 次此类实验。