  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.
