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.
Figures & tables
Fig. 1 : Total training compute for deep learning models.
Specification
MI350X
MI355X
B200
Architecture
CDNA 4
CDNA 4
Blackwell
Compute Units / Streaming Multiprocessors
256 (8 XCD × 32)
256 (8 XCD × 32)
148 (2 chiplets)
Peak FP16/BF16 Dense Performance
2300 TFLOPS
2500 TFLOPS
2250 TFLOPS
HBM Type / Capacity
HBM3e / 288 GB
HBM3e / 288 GB
HBM3e / 180 GB
HBM Bandwidth
Up to 8 TB/s
Up to 8 TB/s
Up to 8 TB/s
LDS per CU or SM
160 KB
160 KB
256 KB
Table I : Comparison of AMD Instinct MI350X, MI355X, and NVIDIA Blackwell B200 GPUs. MI355X is a 1400 W design using direct liquid cooling, while Meta uses the 1000 W air-cooled MI350X. Meta’s deployment of B200 operates at 750 W, which is different from the general B200 HGX solution that operates at 1000 W. Meta operates both GPUs below their vendor-specified TDPs on the same air-cooled platform. B200 is capped at a lower power limit to enable air cooling because its power density is higher than that of MI350X.
Fig. 2 : Hipification process to enable ML jobs on AMD GPUs.
Fig. 3 : Pipelined GEMM execution. Loads/stores use vector units; compute uses matrix units. The K-loop overlaps load and compute.
Fig. 4 : Wave quantization tail effect. A 4 × 5 output (20 tiles) on 8 CUs requires 3 timesteps with 20mod8=4 underutilized CUs in the final timestep. Origami addresses this via two strategies: tile selection to change the number of output tiles, and Stream-K to partition work along the K dimension.
Fig. 5 : GEMM performance on MI350X for production shapes, shown as percentage of peak achievable FLOPS.
Fig. 6 : GEMM performance ( M,N,K dimensions =8192 ) with GPU power excursion across kernel durations and duty ratios on MI350X.
Fig. 7 : Measured AllToAll bandwidth with different group and message sizes relative to the theoretical upper bounds.
Fig. 8 : Eigen decomposition performance.
Fig. 9 : GDPA kernel performance with incremental optimization techniques applied on AMD MI350X, compared to optimal configuration on B200. Lower latency and higher throughput are better.
Fig. 10 : HSTU kernel performance with incremental optimization techniques applied on AMD MI350X, compared to optimal configuration on B200. Lower latency and higher throughput are better.
Fig. 11 : Bounds check forward kernel latency vs. violation rate. Tree reduction eliminates atomic contention, enabling MI350X to outperform B200 at high violation rates.
Fig. 12 : GIS kernel performance. For the bandwidth-bound forward kernel, MI350X exceeds B200. For backward, multi-row packing helps on mixed-dimension configs but cannot fully overcome atomic contention on uniform, high-collision workloads.
B200
MI350X
B
T
L
D
Prec
α
FWD
BWD
FWD
BWD
(ms)
(ms)
(ms)
Before
After
65536
5
55
192
fp16
1.0
1.75
14.40
1.48 (0.85 × )
23.6 (1.64 × )
7.72 (0.54 × )
65536
5
55
192
fp16
1.15
0.84
4.61
0.78 (0.93 × )
9.92 (2.15 × )
3.97 (0.86 × )
65536
10
26
256
fp16
1.0
1.86
15.79
1.84 (0.99 × )
24.4 (1.55 × )
8.06 (0.51 × )
65536
10
26
256
fp16
1.15
0.87
4.85
0.91 (1.05 × )
11.4 (2.35 × )
4.36 (0.90 × )
Table II : TBE training performance: MI350X forward and backward before/after optimization vs. B200 with representative production configs: B (batch size), T (number of tables), L (lookups per table), D (embedding dimension), Prec (precision), α (Zipf skew factor). Parenthetical values show latency relative to B200 (<1.0 × means MI350X is faster).
Fig. 13 : Triton Inductor kernel latencies by category for three representative models, showing MI350X before/after AMD optimizations compared to B200. Lower latency is better.
Fig. 14 : Training throughput vs. HBM tradeoff on a GPT-2 1.3B model with activation checkpointing.
Fig. 15 : (a) Batch size scaling yields up to 51% throughput improvement. (b) Combined co-design strategies (batch size increase, AC removal, embedding placement optimization) yield 26% speedup.
Model
GPUs
Seq. len
Params
B200 ex/s
MI350X ex/s
MI350X/B200
A
8
512
∼100 M
921,714
754,936
82%
B
8
256
∼100 B
1,592,926
1,810,036
114%
C
64
200
∼200 B
728,001
884,216
121%
D
64
1000
∼300 B
774,324
755,221
98%
E
64
200
∼300 B
283,641
267,635
94%
F
128
1000
∼3.6 T
330,481
304,781
92%
Table III : End-to-end training throughput (examples/second) on MI350X vs B200 across ten production models. Higher is better.
Fig. 16 : Kernel latency ratio (MI350X / B200) by category. For each model, the sum of kernel latencies on B200 is normalized to 100%. This analysis does not consider multi-stream overlap.