Inside Microsoft’s New Frameworks to Enable Large-Scale AI
By: Jesus Rodriguez
Tags: Artificial Intelligence, Data Science, Machine Learning, Deep Learning, Invector Labs

Inside Microsoft’s New Frameworks to Enable Large-Scale AI

ZeRO-2, a new ONNX Runtime and a supercomputer built in partnership with OpenAI are Microsoft’s latest efforts to enable large scale AI workloads.

Training models with massive datasets is becoming the norm in modern deep learning applications. Some of the latest innovations such as Transformers and other pre-training models massively large training infrastructures. This type of pre-training models are often seen as the latest breakthrough in artificial intelligence(AI) and it was triggered a race across some of the top technology companies in the world to dominate that space. At its recent Build conference, Microsoft unveiled several key efforts to enable the construction and use of pre-training models at scale.

Pre-training models and Transformers have been at the center of the latest breakthroughs in deep learning. Google open the doors to this movement with its famous BERT model just to be quickly followed by other massively large architectures such as OpenAI’s GPT-2 or Microsoft’s Turin-NLG. The creation of this type of models requires more than just large computational resources as it often requires architectures optimized for this type of large training exercises. Given its novelty, those architectures still haven’t been included in mainstream deep learning frameworks. Microsoft’s recent initiatives represent some of the most comprehensive efforts to streamline the creation and access of large scale models for the deep learning community. Let’s review a few of the efforts:

Zero-2 & DeepSpeed

Earlier this year, Microsoft’s open source the first version of DeepSpeed, a framework for training optimization at large scale. DeepSpeed was the fundamental architecture behind Microsoft’s record breaking Turing-NLG model. The release also included an optimization module called Zero Redundacy Optimizer(ZeRO). Conceptually, ZeRO tries to find a middleground between model and data parallelism in order to maximize the scalability of a model. The biggest benefit of ZeRO is that allows running large deep learning models in memory without the need of code refactoring. This is done by optimizing the memory used by the model. In that sense, ZeRO can be viewed as a memory optimization method for deep learning models.

ZeRO-2 is the second release of ZeRO that includes several innovations compared to its predecessor. Those innovation include all sorts of areas related to memory consumption in deep learning models such as activation memory, fragmented memory and model state memory.

· Model State Memory: ZeRO-2 optimizes the memory used to kept the state of the model saving up to 8x compared to traditional data parallelism methods. Any deep learning model can model its state in three fundamental stages: partitioning optimizer states, gradients, and parameters. The first implementation of ZeRO was able to optimized memory consumption for the Pos stage saving up to 4x compared to traditional methods. ZeRO-2 expands that approach enabling memory consumption optimization for the gradients stage adding an extra 2x of savings.

· Activation Memory: Even after optimizing the model state memory, there might be bottlenecks caused by activations stored in a forward pass in order to enable a backward pass. ZeRO-2 introduces an activation partitioning technique that removes these memory bottlenecks by offloading to the host CPU.

· Fragmented Memory: Some of the memory inefficiencies in deep learning models are caused by fragmentations due to the lifetime of the different tensors of the model. ZeRO-2 manages memory throughout the different lifetime of tensors, eliminating this inefficiency.

One of the most remarkable innovations of ZeRO-2 is that it enables memory optimizations on each individual GPU. This contrasts with alternative approaches that mostly optimize by distributing training across different GPUs. The value of this optimization was fairly visible when the ZeRO-2 team was able to train a version of Google’s BERT in 44 minutes using a single GPU. The chart below shows the improvements of the ZeRO-2 approach compared to previous attempts.

ZeRO-2 provides several advantages in areas such as the ability to support models with billions of parameters, achieving higher throughput and speed, scalability per GPU and across a GPU cluster and the ability of training super large models in a single GPU.

ZeRO-2 + ONNX Runtime

The ONNX Runtime is an open source framework that supports interoperability between different machine learning frameworks. One of the interesting questions about machine learning interoperability is whether is works for super large models. It’s fairly simple to use ONNX for a decision tree but would it work for a model with several billion parameters?

Microsoft decided to include DeepSpeed and ZeRO-2 as part of the ONNX Runtime in order to optimize the training and memory management of ONNX models. The results were tested on some of the largest pre-trained models in the world such as GPT-2 in Visual Studio, Turing-NLG in Office 365 and a large transformer with 500 million parameters in Bing. The following chart illustrates the benefits compared to the original implementations.

An AI-First Super Computer

Infrastructure is key to enable large scale pre-training models. Scientists from Microsoft and OpenAI collaborated in the implementation of a super computer specialized in large deep learning workloads. The super computer includes an unfathomable infrastructure of a single system with more than 285,000 CPU cores, 10,000 GPUs and 400 gigabits per second of network connectivity for each GPU server.

The infrastructure hosted by the Microsoft-OpenAI supercomputer places it within the top five super computers in the world. However, even more remarkable is the fact that this super computer is optimized for AI workloads. In that sense, the supercomputer is designed for executing large pretrained models such as Turing-NLG or BERT. Microsoft will be making the new supercomputer available to developers as part of the Azure cloud.

The race for large scale AI dominance promises to be one of the most fascinating trends in the next few years of the AI market. For now, Microsoft’s initiatives such as DeepSpeed-ZeRO-2, ONNX Runtime and the new OpenAI supercomputer represent an impressive first step in that race.