AMD GPUs Reach B200 Training Parity Through Co-Design
Model and kernel changes exploit MI350X’s larger HBM capacity, lifting production recommendation-training throughput to 101.2% of a B200 baseline.
Underlying Paper
Architecting for Hardware Diversity: Meta's End-to-End Co-Design for High-Performance AI Training on AMD GPUs
Meta deploys a heterogeneous fleet comprising NVIDIA GPUs, AMD GPUs, and proprietary MTIA accelerators to reduce costs and mitigate vendor lock-in. However, the extensive history of NVIDIA-specific optimizations hinders performance portability to other hardware. This paper presents our methods for optimizing alternative accelerators, using AMD Instinct MI350X versus NVIDIA B200 for recommendation model training as a case study. We propose three reusable principles. First, hardware-model co-design is essential: rather than directly porting models from NVIDIA GPUs, revising model structures and software recipes allows full exploitation of accelerator-specific strengths (e.g., MI350X's 1.6x HBM capacity relative to B200). Second, our analytical performance model enables zero-warmup optimal kernel selection across diverse datatypes and matrix shapes in production models, drastically simplifying tuning. Third, increasingly mature agentic workflows have lowered the barrier to porting from the CUDA ecosystem; we extend these workflows for microarchitecture-aware kernel optimizations and custom shim generation. Compared to the directly ported baseline, these methods deliver up to 1.5x improvement in end-to-end training throughput. To our knowledge, this is the first hyperscaler report on productionizing AMD GPUs for large-scale training. Our contributions have been upstreamed to the open-source PyTorch, Triton, and ROCm ecosystems.
Heterogeneous AI fleets reduce dependence on a single accelerator supplier, but training stacks tuned over years for CUDA do not transfer cleanly to AMD hardware. Meta evaluates that gap on production deep learning recommendation models, comparing AMD’s MI350X with NVIDIA’s B200. The paper’s central argument is that portability is not a matter of recompiling kernels: model shapes, embedding capacity, distributed communication, kernel selection, and operations all need to be designed around the target device.
Core Contribution
The authors present an end-to-end co-design program rather than a single compiler or kernel optimization. The production comparison covers ten recommendation-model configurations, spanning 8–128 GPUs, sequence lengths from 512 to 8192, and models from roughly 100M to 1T parameters. Across these workloads, MI350X reaches a geometric-mean 101.2% of B200 end-to-end throughput, despite a directly ported implementation initially leaving hardware-specific capacity and execution behavior underused.
The relevant hardware asymmetry is memory. MI350X has 1.6× the HBM capacity of B200, and the paper reports that larger HBM-backed batch sizes improve per-GPU throughput in Models C and D by 98%–102%, Models E and F by 94%–110%, and Model H by 92%–95%. This is a useful correction to a simplistic device-versus-device comparison: the reported parity depends partly on changing the workload configuration to use memory that the NVIDIA-oriented recipe did not require.
Technical Approach
The first component is hardware-model co-design. The authors revise model structures and training recipes to increase usable batch size and preserve stable, comparable kernel shapes. For recommendation workloads dominated by embedding tables, this also means treating placement and sharding as system-level choices: tables have uneven access frequencies, while memory capacity and bandwidth are distributed across GPU HBM, CPU DRAM, and SSDs.
Second, the team builds an analytical performance model for kernel selection across data types and matrix shapes. Its stated purpose is zero-warmup selection of the best kernel in production, avoiding empirical tuning runs for each shape. The paper couples this with customized Triton and CUDA-derived kernel work, plus microarchitecture-aware agentic workflows that assist porting and shim generation.
Kernel-level results explain why a uniform port is insufficient. Hand-written Triton kernels reduce execution latency by 6%–52% across categories, while the Triton Inductor path is reported to gain 2.1× in one Model-H case. Embedding operations are mixed: MI350X ranges from 0.61× to 1.04× of B200 latency, showing that some CUDA-derived optimizations still favor NVIDIA rather than producing an across-the-board AMD advantage.
The communication findings are similarly conditional. The paper reports broadly comparable collective performance, with MI350X ahead on multi-host, small-message AllToAll collectives and especially on Models C and J. Figure 7 measures AllToAll bandwidth across group and message sizes against theoretical upper bounds, making the communication result a bandwidth-and-topology question rather than a generic claim about collectives.
Results and Analysis
Table III gives the clearest end-to-end evidence. MI350X exceeds B200 on Models B (114%), C (121%), H (102%), and J (130%); it trails on A (82%), D (98%), E (94%), F (92%), G (92%), and I (96%). The strongest result is therefore workload parity in aggregate, not dominance on every production configuration. Model J’s 242,325 examples/s on MI350X versus 185,728 on B200 is substantial, but Model A remains at 754,936 versus 921,714 examples/s.
The category breakdown also identifies a persistent constraint. GEMM operations favor B200, with MI350X showing 1.3×–1.7× higher latency; the authors attribute much of Model A’s 18% deficit to three BF16 GEMMs. MI350X lacks native hardware support for TF32, so those operations carry a penalty. As recommendation models move toward BF16 or lower precision, that penalty may narrow, but the evidence here does not establish parity for TF32-heavy training.
The operational work matters because a second vendor changes more than model throughput. The paper describes reliability measures for high-bandwidth-memory RAS, memory-access faults, silent data corruption, device quarantine, and host burn-in. It reports training-loss divergence within 0.3% and numerical error within about 10⁻⁴ at multi-thousand-GPU scale. Those checks support deployment credibility, although they do not replace evidence from model families beyond DLRMs.
Limits in Practice
This is a focused hyperscaler case study, not a general AMD-versus-NVIDIA benchmark. Results center on Meta’s recommendation models, its software stack, and an MI350X/B200 comparison. The paper also acknowledges remaining GEMM disadvantages and relies on overlapping GPU streams, which complicates attributing individual kernel latency to end-to-end throughput. The reported 1.5× maximum gain is credible as a result of coordinated workload and systems changes; it should not be read as a portable expectation for arbitrary CUDA workloads.
Evidence Box
strongKey Claims
- •Hardware-model co-design can recover production training performance on AMD GPUs
- •Analytical kernel selection removes warmup tuning across production matrix shapes
- •Microarchitecture-aware porting workflows lower CUDA-to-ROCm migration effort
- •Multi-vendor operations can meet large-scale training reliability requirements
Key Results
- •101.2% geometric-mean MI350X/B200 throughput across 10 recommendation-model configurations
- •Up to 130% MI350X/B200 throughput for Model J, with 242,325 vs. 185,728 examples/s
- •Hand-written Triton kernels reduce category latency by 6%–52%
- •Training-loss divergence stays within 0.3% and numerical error within about 10⁻⁴ at multi-thousand-GPU scale
Limitations & Caveats
- •Evaluation is limited to Meta production DLRMs on MI350X versus B200
- •MI350X GEMMs show 1.3×–1.7× higher latency, especially for TF32 workloads
- •Several models remain below B200 throughput, including Model A at 82%
- •Overlapped GPU streams complicate attribution of kernel latency to end-to-end throughput
