Train ALBERT for natural language processing with TensorFlow on Amazon SageMaker Amazon Web Services

At re:Invent 2019, AWS shared the fastest training times on the cloud for two popular machine learning (ML) models: BERT (natural language processing) and Mask-RCNN (object detection). To train BERT in 1 hour, we efficiently scaled out to 2,048 NVIDIA V100 GPUs by improving the underlying infrastructure, network, and ML framework. Today, we're open-sourcing the optimized training code for ALBERT (A Lite BERT), a powerful BERT-based language model that achieves state-of-the-art performance on industry benchmarks while training 1.7 times faster and cheaper. This post demonstrates how to train a faster, smaller, higher-quality model called ALBERT on Amazon SageMaker, a fully managed service that makes it easy to build, train, tune, and deploy ML models. Although this isn't a new model, it's the first efficient distributed GPU implementation for TensorFlow 2. You can use AWS training scripts to train ALBERT in Amazon SageMaker on p3dn and g4dn instances for both single-node and distributed training.