Why Databricks Chooses to Use AMD GPU

Despite the worldwide frenzy for NVIDIA GPUs like the H100, Databricks has taken a strategic move by choosing to leverage AMD GPUs to enhance their LLM training capabilities, and it has proven to be very effective for the company.

Last year, Databricks partnered with AMD and used their third-generation EPYC Instance processors. In June of this year, Databricks acquired MosaicML, which is using AMD MI250 GPUs to train AI models. This enterprise software company sees the potential in AMD and is looking forward to the release of the next-generation MI300X from the chip manufacturing company to make a mark in the generative AI field.

AMD GPUs seem to have gained widespread adoption in the community, proving their strength in the field of artificial intelligence. Well-known AI startups like Lamini and Moreh have adopted AMD MI210 and MI250 systems to fine-tune and deploy custom LLM models. Just a week ago, Lamini revealed that their LLM model runs on AMD's Instinct GPU.

Moreh successfully trained a language model with an astonishing 221 billion parameters using 1200 AMD MI250 GPUs. AMD recently invested $22 million in Series B funding for Moreh. The open-source LLM model OLMo from AI2 is also trained using a large cluster of AMD GPUs.

Databricks announced that they have gained early access to a significant multi-node MI250 cluster from AMD Accelerated Cloud (AAC). The cluster consists of 32 nodes, each equipped with 4 AMD Instinct MI250 GPUs and has 800Gbps interconnectivity. This setup is ideal for scalable LLM training on AMD hardware.

The company has now expanded to using 128 MI250 GPUs.

The reason is simple - software

We all know that NVIDIA's true monopoly lies in its software CUDA, which powers all of its AI capabilities. AMD has also recognized this and has been at the forefront of software innovation, particularly with the AMD Radeon Open eCosystem (ROCm) software platform, which is an alternative to CUDA.

Recently, Vamsi Boppana, the Senior Vice President of AI at AMD, stated that ROCm is currently a top priority for the company. "We have more manpower in software, and AMD CEO Lisa Su has made it clear that she wants to make significant and sustained investments in software," he said.

ROCm has undergone a significant upgrade from version 5.4 to 5.7. Additionally, Databricks emphasizes that the ROCm kernel for FlashAttention has been upgraded to FlashAttention-2, resulting in a substantial performance improvement.

Lamini also states that AMD's ROCm is production-ready and claims that "it has tremendous potential to accelerate LLM fine-tuning and AI development in various domains, comparable to or even better than CUDA."

Databricks also praises AMD's active involvement in OpenAI's Triton compiler. This contribution enables machine learning engineers to develop custom kernels that run efficiently on various hardware platforms, including NVIDIA and AMD systems, simplifying the process of porting custom model code and kernels without the need for ROCm-specific kernels.

How does AMD GPU perform for Databricks?

Databricks has seen significant improvements in training performance when using ROCm 5.7 and FlashAttention-2, achieving a 1.13x improvement compared to previous results with ROCm 5.4 and FlashAttention. Furthermore, Databricks demonstrates strong scalability, with performance increasing from 166 TFLOP/s/GPU on a single node to 159 TFLOP/s/GPU on 32 nodes, while maintaining consistent global training batch sizes.

After successfully training MPT models with 1 billion and 3 billion parameters on 64 MI250 GPUs, the training process remains stable, and the evaluation metrics of the final models are comparable to renowned open-source models such as Cerebras-GPT-1.3B and Cerebras-GPT-2.7B.

To conduct the training, Databricks utilizes open-source training libraries based on PyTorch, such as LLM Foundry, built on top of Composer, StreamingDataset, and PyTorch FSDP. This is made possible by PyTorch's support for CUDA and ROCm, enabling seamless operations on both NVIDIA and AMD GPUs without code modifications.

Databricks states that they are excited about the upcoming release of the next-generation AMD Instinct MI300X GPU. They expect their PyTorch-based software stack to continue running seamlessly and scaling effectively.

In addition, the integration of AMD and Triton will simplify the process of porting custom model code and kernels, eliminating the need for ROCm-specific kernels.

Abhi Venigalla, a researcher at Databricks, said, "While the H100 still reigns supreme, we look forward to testing AMD's new MI300X as we believe it will be highly competitive!"

Lamini is also eagerly awaiting the launch of the MI300X, which comes with 192GB of high-bandwidth memory (HBM), expected to enhance their model's performance.

In summary, Databricks has taken an important step towards AMD GPUs in the LLM training field. This development also demonstrates AMD's growing position in the GPU market.