Introduction

4 months in the past, we shared how AMD had emerged as a succesful platform for generative AI and demonstrated the way to simply and effectively practice LLMs utilizing AMD Intuition GPUs. Immediately, we’re excited to share that the hits maintain coming!

Group adoption of AMD GPUs has exploded: AI startups reminiscent of Lamini are utilizing AMD MI210 and MI250 methods to finetune and deploy customized LLMs, and Moreh was in a position to practice a 221B parameter language mannequin on their platform utilizing 1200 AMD MI250 GPUs. Moreover, open-source LLMs reminiscent of AI2’s OLMo are additionally being skilled on giant clusters of AMD GPUs.

In the meantime, at AMD, the ROCm software program platform has been upgraded from model 5.4 to five.7, and the ROCm kernel for FlashAttention has been upgraded to FlashAttention-2, delivering important efficiency positive aspects.

In parallel, AMD has been making important contributions to the Triton compiler, enabling ML engineers to write down customized kernels as soon as that run performantly on a number of {hardware} platforms, together with NVIDIA and AMD.

Lastly, at MosaicML/Databricks, we acquired early entry to a brand new multi-node MI250 cluster constructed as a part of the AMD Accelerator Cloud (AAC). This new cluster has 32 nodes, every containing 4 x AMD Intuition™ MI250 GPUs, for a complete of 128 x MI250s. With a high-bandwidth 800Gbps interconnect between nodes, the cluster is ideal for testing LLM coaching at scale on AMD {hardware}.

We’re excited to share our first multi-node coaching outcomes on MI250 GPUs!

• When coaching LLMs on MI250 utilizing ROCm 5.7 + FlashAttention-2, we noticed 1.13x greater coaching efficiency vs. our ends in June utilizing ROCm 5.4 + FlashAttention.
• On AAC, we noticed robust scaling from 166 TFLOP/s/GPU at one node (4xMI250) to 159 TFLOP/s/GPU at 32 nodes (128xMI250), after we maintain the worldwide practice batch dimension fixed.
• To check convergence, we skilled two MPT fashions with 1B and 3B params from scratch with Chinchilla-optimal token budgets, every on 64 x MI250 GPUs. We discovered that coaching was steady, and the ultimate fashions had comparable eval metrics to compute-matched open-source fashions (Cerebras-GPT-1.3B and Cerebras-GPT-2.7B).

As earlier than, all outcomes are measured utilizing our open-source coaching library, LLM Foundry, constructed upon Composer, StreamingDataset, and PyTorch FSDP. Due to PyTorch’s assist for each CUDA and ROCm, the identical coaching stack can run on both NVIDIA or AMD GPUs with no code adjustments.

Waiting for the next-gen AMD Intuition MI300X GPUs, we anticipate our PyTorch-based software program stack to work seamlessly and proceed to scale effectively. We’re additionally excited concerning the potential of AMD + Triton, which can make it even simpler to port customized mannequin code and kernels. For instance, quickly we will use the identical FlashAttention-2 Triton kernel on each NVIDIA and AMD methods, eliminating the necessity for a ROCm-specific kernel.

Learn on for extra particulars about AMD multi-node coaching, and ask your CSP if/when they are going to provide the AMD MI300X!

AMD Platform

When you aren’t acquainted with the AMD platform, under is a fast overview. For a deeper dive, take a look at our AMD Half 1 weblog.

The AMD MI250 and (future) MI300X are datacenter accelerators just like the NVIDIA A100 and H100, with Excessive Bandwidth Reminiscence (HBM) and Matrix Cores which are analogous to NVIDIA’s Tensor Cores for quick matrix multiplication. When evaluating GPUs inside “generations”, it is most acceptable to check the MI250 to the A100, whereas the MI300X is just like the H100. See Desk 1 for particulars.

Desk 1: {Hardware} specs for NVIDIA and AMD GPUs. Please word that solely a subset of the MI300X specs have been launched publicly as of Oct 2023.

On high of those accelerators, AMD has developed a set of software program and networking infrastructure:

• ROCm: A library of drivers, instruments, and high-performance GPU kernels.
• RCCL: A communications library for high-performance cross-GPU operations like collect, scatter, and scale back which are used for distributed coaching.
• Infinity Cloth: excessive bandwidth networking inside a node.
• Infiniband or RoCE: excessive bandwidth networking throughout nodes.

At every layer of the stack, AMD has constructed software program libraries (ROCm, RCCL) or networking infrastructure (Infinity Cloth) or adopted present networking infrastructure (Infiniband or RoCE) to assist workloads like LLM coaching. When you’re acquainted with NVIDIA’s platform, you will see that many parts of AMD’s platform immediately map to these of NVIDIA’s. See Determine 2 for a side-by-side comparability.

One of the crucial widespread libraries for ML programmers is PyTorch, and as we reported in our final weblog, PyTorch runs seamlessly in keen mode on AMD GPUs like MI250. Even superior distributed coaching methods like Absolutely-Sharded Information Parallelism (FSDP) work out of the field, so when you’ve got a PyTorch program (like LLM Foundry) that works on NVIDIA GPUs, there’s a excessive likelihood that it’ll additionally work on AMD GPUs. That is made doable by PyTorch’s inside mapping of operators like `torch.matmul(...)` to both CUDA or ROCm kernels.

One other thrilling software program enchancment to the AMD platform is its integration with the Triton compiler. Triton is a Python-like language that permits customers to write down performant GPU kernels whereas abstracting away a lot of the underlying platform structure. AMD has already added assist for its GPUs as a third-party backend in Triton. And kernels for numerous efficiency vital operations, like FlashAttention-2, can be found for AMD platforms, for each inference and coaching, with aggressive efficiency (See Determine 3).

AMD’s assist for Triton additional will increase interoperability between their platform and NVIDIA’s platform, and we stay up for upgrading our LLM Foundry to make use of Triton-based FlashAttention-2 for each AMD and NVIDIA GPUs. For extra data on Triton assist on ROCm, please see AMD’s presentation on the 2023 Triton Developer Convention.

LLM Coaching Efficiency

Beginning with single-node LLM coaching efficiency, ROCm 5.7 + FlashAttention-2 is considerably quicker than the sooner ROCm 5.4 + FlashAttention package deal. When coaching MPT fashions of assorted sizes, we noticed a median speedup of 1.13x versus our coaching ends in June. See Determine 4.

On every accelerator, we ran the identical coaching scripts from LLM Foundry utilizing MPT fashions with a sequence size of 2048, BF16 blended precision, FlashAttention-2 (Triton-based on NVIDIA methods and ROCm-based on AMD methods at the moment) and PyTorch FSDP with sharding_strategy: FULL_SHARD. We additionally tuned the microbatch dimension for every mannequin on every system to attain most efficiency. All coaching efficiency measurements are reporting mannequin FLOPs somewhat than {hardware} FLOPs. See our benchmarking README for extra particulars.

When evaluating MI250 towards the same-generation A100, we discovered that the 2 accelerators have comparable efficiency. On common, MI250 coaching efficiency is inside 85% of A100-40GB and 77% of A100-80GB.

When evaluating the AMD MI250 to the brand new NVIDIA H100-80GB, there’s a important hill to climb. However based mostly on public details about the upcoming MI300X, which has a reminiscence bandwidth of 5.2TB/s vs. 3.35TB/s on H100-80GB, we anticipate that the MI300X can be very aggressive with the H100.

Shifting on to multi-node efficiency, the 128 x MI250 cluster reveals glorious scaling efficiency for LLM coaching. See Determine 5. We skilled an MPT-7B mannequin with fastened international practice batch dimension samples on [1, 2, 4, 8, 16, 32] nodes and located near-perfect scaling from 166 TFLOP/s/GPU at one node (4xMI250) to 159 TFLOP/s/GPU at 32 nodes (128xMI250). That is partly as a result of FSDP parameter sharding, which frees up GPU reminiscence because it scales, permitting barely bigger microbatch sizes at bigger machine counts.

For extra particulars about our coaching configs and the way we measure efficiency, see our public LLM Foundry coaching benchmarking web page.

Given these outcomes, we’re very optimistic about MI250 efficiency at greater machine counts, and we won’t wait to profile and share outcomes on bigger MI250/MI300X clusters sooner or later!

LLM Convergence

To check coaching stability on AMD MI250, we determined to coach MPT-1B and MPT-3B fashions from scratch on the C4 dataset utilizing Chinchilla-optimal token budgets to verify that we might efficiently practice high-quality fashions on the AMD platform.

We skilled every mannequin on 64 x MI250 with BF16 blended precision and FSDP. Total, we discovered that coaching was steady and FSDP + distributed checkpointing labored flawlessly (See Determine 6).

Once we evaluated the ultimate fashions on commonplace in-context-learning (ICL) benchmarks, we discovered that our MPT-1B and MPT-3B fashions skilled on AMD MI250 achieved comparable outcomes to Cerebras-GPT-1.3B and Cerebras-GPT-2.7B. These are a pair of open-source fashions skilled with the identical parameter counts and token budgets, permitting us to check mannequin high quality immediately. See Desk 2 for outcomes. All fashions have been evaluated utilizing the LLM Foundry eval harness and utilizing the identical set of prompts.

These convergence outcomes give us confidence that prospects can practice high-quality LLMs on AMD with no important points as a result of floating level numerics vs. different {hardware} platforms.

Desk 2: Coaching particulars and analysis outcomes for MPT-[1B, 3B] vs. Cerebras-GPT-[1.3B, 2.7B]. Each pairs of fashions are skilled with comparable configurations and attain comparable zero-shot accuracies on commonplace in-context-learning (ICL) duties.

What’s Subsequent?

On this weblog, we have demonstrated that AMD MI250 is a compelling possibility for multi-node LLM coaching. Our preliminary outcomes present robust linear scaling as much as 128 x MI250, steady convergence, and because of latest software program enhancements, the MI250 is closing the efficiency hole with the A100-40GB.

We consider these outcomes strengthen the AI coaching story for AMD GPUs, and because of the interoperability of PyTorch, Triton, and open-source libraries (e.g., Composer, StreamingDataset, LLM Foundry), customers can run the identical LLM workloads on both NVIDIA or AMD and even swap between the platforms, as we exhibit in Determine 7.

Keep tuned for future blogs on AMD MI300X and coaching at a fair bigger scale!