{"id":90699,"date":"2022-03-21T09:00:00","date_gmt":"2022-03-21T16:00:00","guid":{"rendered":"https:\/\/cloudblogs.microsoft.com\/opensource\/?p=90699"},"modified":"2024-06-19T10:50:28","modified_gmt":"2024-06-19T17:50:28","slug":"supporting-efficient-large-model-training-on-amd-instinct-gpus-with-deepspeed","status":"publish","type":"post","link":"https:\/\/opensource.microsoft.com\/blog\/2022\/03\/21\/supporting-efficient-large-model-training-on-amd-instinct-gpus-with-deepspeed\/","title":{"rendered":"Supporting efficient large model training on AMD Instinct\u2122 GPUs with DeepSpeed"},"content":{"rendered":"\n

This post was co-authored by Jithun Nair and Aswin Mathews, members of technical staff at AMD<\/em>.<\/p>\n\n\n\n

In recent years, large-scale deep learning models have demonstrated impressive capabilities, excelling at tasks across natural language processing, computer vision, and speech domains. Companies now use these models to power novel AI-driven user experiences across a whole spectrum of applications and industries. However, efficiently training large models with 10\u2019s or 100\u2019s of billions of parameters is difficult\u2014the sheer size of these models requires them to be distributed across multiple nodes with careful orchestration of compute and communication.     <\/p>\n\n\n\n

DeepSpeed<\/a>, as part of Microsoft\u2019s AI at Scale initiative<\/a>, is a popular open-source library for PyTorch that addresses these difficulties and vastly improves the scale, speed, cost, and usability of large model training and inference. It addresses the scaling challenges by allowing users to easily apply a powerful suite of compute, memory, and communication optimization techniques with minimal code changes. With these techniques, DeepSpeed has enabled training the largest transformer model with 530 billion parameters for language generation<\/a> and helped speed-up training and inference time by a factor of two times to 20 times for real-life scenarios. It is also integrated into popular training libraries like HuggingFace Transformers and PyTorch Lightning.<\/p>\n\n\n\n

Since 2006, AMD has been developing and continuously improving their GPU hardware and software technology for high-performance computing (HPC) and machine learning. Their open software platform, ROCm, contains the libraries, compilers, runtimes, and tools necessary for accelerating compute-intensive applications on AMD GPUs. Today, the major machine learning frameworks (like PyTorch, TensorFlow) have ROCm supported binaries that are fully upstreamed so that users can directly run their code written using these frameworks on AMD Instinct GPU hardware<\/a> and other ROCm compatible GPU hardware<\/a>\u2014without any porting effort.<\/p>\n\n\n\n

AMD has worked closely with the Microsoft DeepSpeed team to bring the suite of parallelization and optimization techniques for the training of large models efficiently on AMD GPUs supporting ROCm. This unlocks the ability to efficiently train models with hundreds of billions of parameters on a wide choice of GPU hardware and system configurations ranging from a single desktop to a distributed cluster of high-performance AMD Instinct\u2122 MI100\/MI200 accelerators.<\/p>\n\n\n\n

Enabling state-of-the-art DL Stack on AMD GPU<\/h2>\n\n\n\n

As an effective enabler of large model training, DeepSpeed provides a suite of powerful parallelism and memory optimizations, such as ZeRO<\/a>, ZeRO-Offload<\/a>, ZeRO-Infinity<\/a>, and 3D parallelism<\/a>, which are crucial for efficient training at massive model scales. With DeepSpeed, model scientists can significantly scale up their model sizes on AMD GPUs well beyond the limits of pure data parallelism. As an example, figure 1 shows model sizes that can be trained using 128 MI100 GPUs (on eight nodes) using different DeepSpeed optimizations. In general, each DeepSpeed optimization enables model scaling of two orders of magnitude compared to the 1.5 billion parameters limits of data parallelism. At the extreme, ZeRO-Infinity powers models with nearly 2 trillion parameters.<\/p>\n\n\n

\"DeepSpeed
Figure 1: DeepSpeed enables over three orders of magnitude of model scaling on 128 MI100 GPUs.<\/figcaption><\/figure>\n\n\n\n

To achieve the optimizations in compute, memory, and communication, DeepSpeed makes use of HIP<\/a> (language\/runtime), rocBLAS<\/a> (for GEMMs), and RCCL<\/a> (for communication) libraries in the ROCm stack. The following figure shows how DeepSpeed interacts with AMD\u2019s ROCm software stack. It requires a version of PyTorch that is built for ROCm.<\/p>\n\n\n

\"Diagram
Figure 2: DeepSpeed and the AMD ROCm stack.<\/figcaption><\/figure><\/div>\n\n\n\n

Through careful work across both AMD and Microsoft engineers, we are proud to announce that DeepSpeed v0.6<\/a> works natively with ROCm-enabled GPUs. This new release of DeepSpeed uses the same APIs as prior releases and does not require any user code changes to leverage the full features of DeepSpeed on ROCm-enabled GPUs. DeepSpeed’s python-level code remains unchanged primarily due to the seamless ROCm experience on PyTorch. DeepSpeed’s CUDA-specific kernels are exposed to users through ROCm’s automatic hipification tools embedded in the PyTorch runtime. This automatic hipification allows DeepSpeed users to continue to enjoy a simple install through PyPI<\/a> and just-in-time (JIT) hipification and compilation at runtime if or when kernels are utilized by end-users. In addition, ROCm has become an integral part of DeepSpeed’s continuous integration (CI) testing, which will elevate ROCm support in DeepSpeed for all future pull requests and new features.<\/p>\n\n\n\n

Power efficient distributed training of large models<\/h2>\n\n\n\n

DeepSpeed enables high training efficiency while running distributed training for large models with billions of parameters across multiple MI100 GPUs and nodes. For example, figure 3 shows that on 8 MI100 nodes\/64 GPUs, DeepSpeed trains a wide range of model sizes, from 0.3 billion parameters (such as Bert-Large) to 50 billion parameters, at efficiencies that range from 38TFLOPs\/GPU to 44TFLOPs\/GPU.<\/p>\n\n\n

\"DeepSpeed
Figure 3: DeepSpeed enables efficient training for a wide range of real-world model sizes.<\/figcaption><\/figure>\n\n\n\n

DeepSpeed also empowers MI100 GPUs to obtain good training scalability for large models as the number of GPUs increases. For example, the plot below in figure 4 shows that for training a 10B parameter model on up to 64 MI100 GPUs, DeepSpeed achieves a throughput scaling (weak scaling) that is close to perfect linear speedup.<\/p>\n\n\n

\"chart,
Figure 4: DeepSpeed achieves good training throughput scalability on 64 MI100 GPUs.<\/figcaption><\/figure>\n\n\n\n

Democratizing large model training<\/h2>\n\n\n\n

An important democratizing feature of DeepSpeed is the ability to reduce the number of GPUs required to fit large models by offloading model states to the central processing unit (CPU) memory and NVMe memory. Offloading makes large models accessible to users with a limited GPU budget by enabling the training (or finetuning) of models with 10s or 100s of billions of parameters on a single node. Below, we briefly provide a flavor of the model scaling that DeepSpeed enables on a single MI100 GPU.<\/p>\n\n\n\n

Efficient model scaling on single GPU <\/h3>\n\n\n\n

The figure below shows that ZeRO-Offload (such as offloading to CPU memory) can train much larger models (such as 12B parameters), on a single MI100 GPU, compared to the baseline PyTorch which runs out of memory (OOM) for models larger than 1.2B parameters. Moreover, ZeRO-Offload sustains higher training throughput (41\u201451 TFLOPs) than PyTorch (30 TFLOPs) by enabling larger batch sizes. In summary, ZeRO-Offload supports model sizes ten times larger than otherwise possible and at higher performance.<\/p>\n\n\n

\"chart,
Figure 5: DeepSpeed enables 10 times model scaling on a single MI100 GPU with great efficiency. PyTorch runs out of memory for models larger than 1.2B.<\/figcaption><\/figure>\n\n\n\n

Extreme model scaling on single GPU <\/h3>\n\n\n\n

The figure below shows that ZeRO-Infinity (such as offloading to NVMe memory) enables even more dramatic scaling in model size on a single MI100 GPU. ZeRO-Infinity utilizes the available 3.5TB NVMe memory in the server to train models as large as 120B parameters on one GPU, thereby scaling the model size by two orders of magnitude compared to the baseline.<\/p>\n\n\n

\"\"
Figure 6: DeepSpeed enables two orders of magnitude of model scaling on a single MI100 GPU.<\/figcaption><\/figure>\n\n\n\n

We also noticed that the training throughput with a single NVMe device (6.2GB\/sec reads and 3.2GB\/sec writes) was 12 TFLOPs and we can increase the throughput by using more NVMe devices. For example, we observed that by using four NVMe devices we were able to double the throughput to 24 TFLOPs.  <\/p>\n\n\n\n

Using DeepSpeed on AMD GPUs\u2014Getting started<\/h2>\n\n\n\n

It is convenient to use DeepSpeed with the latest ROCm software on a range of supported AMD GPUs<\/a>.<\/p>\n\n\n\n

Installation<\/h3>\n\n\n\n

The simplest way to use DeepSpeed for ROCm is to use the pre-built docker image (rocm\/deepspeed:latest) available on Docker Hub<\/a>.<\/p>\n\n\n\n

You can also easily install DeepSpeed for ROCm by just using \u201cpip install deepspeed<\/code>\u201d. For details and advanced options please refer to the installation section of the DeepSpeed documentation page<\/a>.<\/p>\n\n\n\n

Using DeepSpeed on ROCm with HuggingFace models <\/h3>\n\n\n\n

The HuggingFace Transformers<\/a> is compatible with the latest DeepSpeed and ROCm stack.<\/p>\n\n\n\n