[论文解读] Boosting Distributed Training Performance of the Unpadded BERT Model

Jinle Zeng, Min Li+5 more|arXiv (Cornell University)|Aug 17, 2022

Topic Modeling被引用 4

一句话总结

该论文提出了一种高度优化的无填充 BERT 模型，通过使用可变长度输入，利用分组多流 FMHA、重叠数据交换实现的工作负载均衡，以及内核/操作符优化，消除了填充标记带来的冗余计算。该方法在 8 块 A100 GPU 上实现了 2578 个样本/秒的最先进吞吐量，在 MLPerf Training v2.0 中排名第一，仅用 16.598 分钟即完成收敛。

ABSTRACT

Pre-training models are an important tool in Natural Language Processing (NLP), while the BERT model is a classic pre-training model whose structure has been widely adopted by followers. It was even chosen as the reference model for the MLPerf training benchmark. The distributed training performance optimization of BERT models plays an important role in accelerating the solutions of most NLP tasks. BERT model often uses padding tensors as its inputs, leading to excessive redundant computations. Thus, removing these redundant computations is essential to improve the distributed training performance. This paper designs a new approach to train BERT models with variable-length inputs efficiently. Firstly, we propose a general structure for the variable-length BERT models, and accelerate the encoder layer via our grouped multi-stream FMHA (Fused Multi-Head Attention) method. Secondly, through data exchange, we address the unbalanced workload problem caused by the variable-length inputs, which overlaps highly with the training process. Finally, we optimize the overall performance of the BERT model, such as kernel fusion, and operator optimization. Our experimental results show that our highly optimized BERT model achieves state-of-the-art throughput and ranks first in MLPerf Training v2.0 within the same GPU configuration. The optimizations in this paper can be applied to more BERT-like models in our future works.

研究动机与目标

• 消除 BERT 训练中因填充标记导致的冗余计算，这些计算会降低分布式训练性能。
• 解决在分布式训练环境中，由于输入序列长度可变导致的工作负载不平衡问题。
• 在相同硬件配置下，实现 BERT-Large 的最先进训练吞吐量。
• 优化完整训练流水线，包括注意力机制、内核融合和操作符级别的性能。

提出的方法

• 设计一种支持动态序列长度、无需填充的可变长度 BERT 模型通用架构。
• 引入分组多流 FMHA（融合多头注意力）以优化注意力计算，加速编码器层。
• 通过将数据交换与 GPU 计算重叠，实现工作负载均衡，最小化空闲时间和通信开销。
• 应用内核融合和操作符级别优化，减少内存带宽并提高 GPU 利用率。
• 采用混合精度训练，配合 LAMB 优化器，并确保设备间批量处理的一致性。
• 使用 PaddlePaddle 作为深度学习框架，并在 8 块 NVIDIA A100 400W GPU 上进行基准测试。

实验结果

研究问题

• RQ1 如何在 BERT 训练中有效消除填充标记带来的冗余计算？
• RQ2 在分布式训练中，当序列长度可变时，有哪些技术可以实现设备间的工作负载均衡？
• RQ3 将数据交换与计算重叠如何提升分布式训练效率？
• RQ4 哪些内核和操作符优化组合能为 BERT-Large 实现最高吞吐量？
• RQ5 完全无填充的 BERT 模型是否能在 MLPerf 等标准化基准测试中超越现有优化实现？

主要发现

• 优化后的无填充 BERT 模型实现了 2578 个样本/秒的吞吐量，在相同 GPU 配置下超越所有其他实现。
• 该模型在 MLPerf Training v2.0 中排名第一，仅用 16.598 分钟即收敛至 72% MLM 准确率，优于 NVIDIA、HazyResearch 及其他提交方案。
• 仅移除与填充相关的计算，相比基线填充实现即获得 2.3 倍的加速。
• 内核融合和操作符优化分别带来 8.9% 和 11.3% 的性能提升。
• 将数据交换与计算重叠可实现 2.8% 的性能提升，有效缓解通信瓶颈。
• 分组多流 FMHA 优化相比标准无填充 FMHA 实现了 3.6% 的性能增益。