Architecting for Hardware Diversity: Meta's End-to-End Co-Design for High-Performance AI Training on AMD GPUs

AI Systems and Hardware OLPP TVM等 Other AI Sys & HW
Meta 公司部署了一套异构计算集群,其中包含英伟达(NVIDIA)GPU、超威半导体(AMD)GPU 以及自研的 MTIA 加速器,旨在降低硬件采购与运维成本,并缓解对单一供应商的技术锁定风险。然而,长期以来针对英伟达 GPU 的深度优化积累,严重制约了模型与系统软件在其他硬件平台上的性能可移植性。本文以推荐系统大模型训练为应用场景,以 AMD Instinct MI350X 与 NVIDIA B200 的性能对比为案例,系统阐述了我们在适配非英伟达加速器过程中所采用的优化方法。我们提炼并提出了三条具备高度复用价值的核心原则: 第一,硬件与模型协同设计(hardware-model co-design)至关重要——我们并非简单地将原本面向英伟达 GPU 设计的模型直接迁移至新平台,而是主动重构模型结构并调整软件实现方案,从而充分释放目标加速器的独特优势(例如,MI350X 的高带宽内存(HBM)容量较 B200 高出 1.6 倍); 第二,我们构建了一套分析型性能建模方法,可在无需预热(zero-warmup)的前提下,针对生产环境中各类数据类型与矩阵形状,自动选择最优内核(kernel),极大简化了系统调优流程; 第三,日益成熟的“智能体驱动”(agentic)工作流显著降低了从 CUDA 生态向其他平台迁移的技术门槛;我们进一步扩展了此类工作流,支持面向具体微架构的内核级深度优化,以及定制化接口层(shim)的自动生成。 相较于未经针对性优化、仅做简单移植的基线方案,上述方法在端到端训练吞吐量上最高可提升 1.5 倍。据我们所知,这是首家超大规模数据中心厂商(hyperscaler)公开披露其在大规模训练场景中成功落地 AMD GPU 的完整实践报告。本工作所贡献的关键技术已全部合入开源社区,包括 PyTorch、Triton 及 ROCm 等主流框架与生态。
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.
许愿