Search papers, labs, and topics across Lattice.
Meta addresses performance portability challenges in heterogeneous AI hardware by developing an end-to-end co-design approach for AMD GPUs, using recommendation model training as a case study. This involves hardware-model co-design to leverage accelerator-specific strengths, an analytical performance model for optimal kernel selection, and agentic workflows for microarchitecture-aware optimizations. The resulting methods achieve up to 1.5x end-to-end training throughput improvement over direct porting and have been upstreamed to PyTorch, Triton, and ROCm.
Hyperscalers can achieve 1.5x training throughput gains on AMD GPUs by moving beyond direct CUDA porting to hardware-model co-design and agentic optimization workflows.
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.