提升无填充BERT模型的分布式训练性能
分布式、并行与集群计算
2022-08-18 v1
摘要
预训练模型是自然语言处理（NLP）中的重要工具，而BERT模型是经典的预训练模型，其结构已被后续工作广泛采用。它甚至被选为MLPerf训练基准的参考模型。BERT模型的分布式训练性能优化对加速大多数NLP任务的求解具有重要作用。BERT模型常使用填充张量作为输入，导致过多冗余计算。因此，消除这些冗余计算对提升分布式训练性能至关重要。本文设计了一种高效训练变长输入BERT模型的新方法。首先，我们提出变长BERT模型的通用结构，并通过分组多流FMHA（融合多头注意力）方法加速编码器层。其次，通过数据交换解决变长输入导致的负载不均衡问题，该过程与训练过程高度重叠。最后，我们优化BERT模型的整体性能，如核融合与算子优化。实验结果表明，我们高度优化的BERT模型达到了最先进的吞吐量，并在相同GPU配置下位居MLPerf Training v2.0榜首。本文的优化可应用于我们未来工作中更多类BERT模型。
引用
@article{arxiv.2208.08124,
title = {Boosting Distributed Training Performance of the Unpadded BERT Model},
author = {Jinle Zeng and Min Li and Zhihua Wu and Jiaqi Liu and Yuang Liu and Dianhai Yu and Yanjun Ma},
journal= {arXiv preprint arXiv:2208.08124},
year = {2022}
}