Multi-node Bert-pretraining: Cost-efficient Approach

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.

References in corpus (13)

• Neural Machine Translation by Jointly Learning to Align and Translate
• Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
• Language Models are Few-Shot Learners
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• XLNet: Generalized Autoregressive Pretraining for Language Understanding
• Mixed Precision Training
• One weird trick for parallelizing convolutional neural networks
• SQuAD: 100,000+ Questions for Machine Comprehension of Text
• Gradient Sparsification for Communication-Efficient Distributed Optimization
• MLPerf Training Benchmark
• Measuring the Effects of Data Parallelism on Neural Network Training
• Massive Exploration of Neural Machine Translation Architectures
• Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Training

Cited by in corpus (3)

• Distributed Deep Learning in Open Collaborations
• Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices
• LV-BERT: Exploiting Layer Variety for BERT