[Paper Review] Boosting Distributed Training Performance of the Unpadded BERT Model

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

Topic Modeling4 citations

TL;DR

This paper proposes a highly optimized unpadded BERT model that eliminates redundant computations from padding tokens using variable-length inputs, leveraging grouped multi-stream FMHA, workload balancing via overlapping data exchange, and kernel/operator optimizations. The approach achieves state-of-the-art throughput of 2578 samples/second and ranks first in MLPerf Training v2.0 on 8x A100 GPUs, converging in just 16.598 minutes.

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.

Motivation & Objective

• To eliminate redundant computations from padding tokens in BERT training, which degrade distributed training performance.
• To address workload imbalance caused by variable-length inputs in distributed training environments.
• To achieve state-of-the-art training throughput for BERT-Large under the same hardware configuration.
• To optimize the full training pipeline, including attention mechanisms, kernel fusion, and operator-level performance.

Proposed method

• Designs a general architecture for variable-length BERT models to support dynamic sequence lengths without padding.
• Introduces grouped multi-stream FMHA (Fused Multi-Head Attention) to accelerate the encoder layer by optimizing attention computation.
• Overlaps data exchange for workload balancing with GPU computations to minimize idle time and communication overhead.
• Applies kernel fusion and operator-level optimizations to reduce memory bandwidth and improve GPU utilization.
• Employs mixed-precision training with LAMB optimizer and ensures consistent batch processing across devices.
• Uses PaddlePaddle as the deep learning framework and benchmarks on 8x NVIDIA A100 400W GPUs.

Experimental results

Research questions

• RQ1 How can redundant computations from padding tokens be effectively eliminated in BERT training?
• RQ2 What techniques can balance workloads across devices when sequences have variable lengths in distributed training?
• RQ3 How does overlapping data exchange with computation improve distributed training efficiency?
• RQ4 What combination of kernel and operator optimizations yields the highest throughput for BERT-Large?
• RQ5 Can a fully unpadded BERT model outperform existing optimized implementations in standardized benchmarks like MLPerf?

Key findings

• The optimized unpadded BERT model achieves a throughput of 2578 samples/second, surpassing all other implementations in the same GPU configuration.
• The model ranks first in MLPerf Training v2.0, converging to 72% MLM accuracy in 16.598 minutes, outperforming NVIDIA, HazyResearch, and other submissions.
• The removal of padding-related computations alone yields a 2.3x speedup over the baseline padded implementation.
• Kernel fusion and operator optimizations contribute an 8.9% and 11.3% performance improvement, respectively.
• Overlapping data exchange with computation provides a 2.8% improvement, reducing communication bottlenecks.
• The grouped multi-stream FMHA optimization delivers a 3.6% performance gain over standard unpad FMHA.