中文

基于完整 AMD 平台的基础模型训练:计算、网络与系统设计

计算与语言 2025-12-05 v2 人工智能 分布式、并行与集群计算

摘要

我们报告了首个在纯 AMD 硬件上进行的大规模混合专家(MoE)预训练研究,采用 MI300X GPU 和 Pollara 网络。我们提炼出针对系统和模型设计的实用指导。在系统层面,我们提供了完整的集群和网络特征轮廓:针对所有核心聚合操作(all-reduce、reduce-scatter、all-gather、broadcast)在不同消息大小和 GPU 数量下的微基准测试。至我们知识,这是首个规模如此之大。我们进一步提供 MI300X 的微基准测试,涉及 kernel 规模和内存带宽,以为模型设计提供参考。在建模方面,我们引入并应用 MI300X 感知的 transformer 规模规则,用于注意力和 MLP 模块,并论证 MoE 宽度如何在训练吞吐量和推理延迟之间进行联合优化。我们详细描述了我们的训练堆栈,包括鲁棒性和 checkpoint-reshaping 等常被忽视的实用工具,以及我们的训练配方的详细信息。我们还提供了我们模型架构和基础模型 - ZAYA1 (760M 活跃参数,83B 总参数 MoE,可在 https://huggingface.co/Zyphra/ZAYA1-base 获取) 的预览,该模型将在后续论文中进一步改进。ZAYA1 基础模型在规模及更大规模下实现了与 Qwen3-4B 和 Gemma3-12B 等主流基础模型相当的性能,并在推理、数学和编码基准测试中超越 Llama-3-8B 和 OLMoE 等模型。总体而言,这些结果表明 AMD 硬件、网络和软件堆栈已成熟且足够优化,可实现竞争性的大规模预训练。

关键词

引用

@article{arxiv.2511.17127,
  title  = {Training Foundation Models on a Full-Stack AMD Platform: Compute, Networking, and System Design},
  author = {Quentin Anthony and Yury Tokpanov and Skyler Szot and Srivatsan Rajagopal and Praneeth Medepalli and Anna Golubeva and Vasu Shyam and Robert Washbourne and Rishi Iyer and Ansh Chaurasia and Tomas Figliolia and Xiao Yang and Abhinav Sarje and Drew Thorstensen and Amartey Pearson and Zack Grossbart and Jason van Patten and Emad Barsoum and Zhenyu Gu and Yao Fu and Beren Millidge},
  journal= {arXiv preprint arXiv:2511.17127},
  year   = {2025}
}