在大规模混合专家模型（MoE，如 DeepSeek-V3 和 Llama 4-Scout）的训练中，保持高效率是当代人工智能面临的挑战之一。这些模型将 GPU、网络和编译器的性能推向了极限。为了解决这一问题，AMD 与 Meta 的 PyTorch 团队联手，针对新型 Instinct MI325X GPU 优化了 TorchTitan 和 AMD 的开源内核库 Primus-Turbo。双方共同实现了 1,024 个 GPU 近乎理想的扩展效率，证明了效率与规模并非不可兼得。
概览
通过利用 TorchTitan、Primus-Turbo 内核以及 MI325X 的硬件能力，我们实现了：
- 2.77 倍加速：通过内核级优化提升 DeepSeek-V3-671B 训练速度
- 96% 的扩展效率：DeepSeek-V3-671B 在 128 到 1,024 个 GPU 之间的扩展表现
- 完美的线性扩展：Llama 4-Scout（32 到 512 个 GPU）
- 完全开源的技术栈：基于 TorchTitan + Primus-Turbo 构建
什么是 TorchTitan？
TorchTitan 是 Meta 推出的原生 PyTorch 蓝图，旨在支持跨多 GPU 和多节点集群的大规模训练。它将适用于现代大语言模型 (LLM) 和 MoE 模型的成熟配方整合进一个可配置的训练栈中，使开发者能够从早期实验到全规模运行复用同一套代码路径。
配置优先的扩展性
只需在一个 TOML 文件中设置流水线并行、张量并行、数据并行或专家并行程度，TorchTitan 即可自动构建作业、连接 NCCL/RCCL 组，并在单个 GPU 或上千个 GPU 上运行相同的脚本。
广泛的架构覆盖
通过统一的 TorchTitan 配置，即可训练稠密模型（Llama 3）、混合模型（Qwen 3）和稀疏 MoE 模型（DeepSeek-V3, Llama 4）。
模块化设计
诸如 AMD FP8 Primus-Turbo 等优化内核可直接插入，无需更改代码，且同一接口可兼容未来的训练后工作流。
简而言之，TorchTitan 为扩展任何模型（无论是稠密还是稀疏）提供了一条统一、可组合的路径，从笔记本电脑原型到集群生产环境均适用。
理解混合专家 (MoE) 模型
混合专家模型是一种稀疏激活的 Transformer 替代方案。在标准模型中，每个 token 都会流经同一个前馈块，因此参数加倍意味着计算量加倍。MoE 层用一组专业专家替换了单个 MLP。一个学习型路由程序会检查每个 token，并仅将其发送给其中少数（例如 k 个）专家。
图 1：稠密层与 MoE 层架构对比
由于每个 token 仅由少数几个专家处理，因此每个 token 的计算预算现在随着较小的 k 值而不是专家总数来扩展。内存遵循同样的模式：如果应用了专家并行，每个 GPU 只存储其拥有的专家参数，从而留出空间来增加隐藏维度、序列长度或词汇量，而不会超出设备限制。专家本身是普通的 MLP 块，因此增加容量就像复制更多副本一样简单，而无需拓宽每一层。
这种经济性使得 MoE 模型能够在以稠密模型速度进行训练的同时，参数规模扩展到数千亿级别。DeepSeek-V3 和 Llama 4 展示了这种优势，通过融合 MoE 层和稠密层，在当前硬件上平衡了容量、准确性和效率。
MoE 预训练的挑战
虽然 MoE 模型具有极具吸引力的优势，但在大规模训练时，必须解决一些独特的挑战才能实现高效的分布式训练。
内核效率
微小矩阵乘法 (GEMM) → GPU 利用率低：路由后，每个专家仅处理极少数 token，导致矩阵乘法 (GEMM) 操作变得极小。这些微型 GEMM 无法使 GPU 的计算单元保持满负荷，并引入了显著的启动和数据移动开销。
通信
繁重的 All-to-All 通信 → 网络瓶颈：专家并行将 token 嵌入分散到拥有所选专家的 GPU 上。因此，每次前向传播都会触发繁重的 All-to-All 集体通信。随着集群规模扩大，通信时间可能超过计算时间，成为主要瓶颈。
并行策略
复杂的并行化 → 因管道不同步导致的空转：现代训练通常结合完全分片数据并行 (FSDP) 以节省内存，流水线并行以增加深度，以及专家并行以增加宽度。如果它们的微批次或调度不同步，就会出现“气泡”——即部分 GPU 空闲而其他 GPU 工作的时间段，导致利用率低下。
路由与稳定性
不稳定的路由 → 专家负载不均：学习型路由程序必须在遵守容量限制的同时，在专家之间均匀分配 token。偏斜的 token 计数、专家过载或丢弃的 token 可能会减缓收敛速度或降低最终模型质量。
我们从三个方面解决了这些预训练障碍：AMD Instinct MI325X GPU 提供 256GB HBM3E 和超过 6TB/s 的带宽，加上 FP8 张量核心，使 MoE 专家保持本地化和计算密集型；TorchTitan 并行策略（FSDP + PP/VPP + EP）平衡了计算和通信；Primus-Turbo 的 FP8 注意力和分组 GEMM 内核将 MoE 的微型 GEMM 转变为高占用、高吞吐量的操作。
AMD Instinct™ MI325X：硬件与软件栈
图 2：AMD Instinct MI325X 规格
Instinct MI325X 将 256 GB HBM3E 与超过 6 TB/s 的封装内带宽以及 Petaflop 级的 FP8/BF16 矩阵核心相结合。如此巨大的快速内存使得 MoE 可以在本地设备上运行其大部分专家，而不是分散在整个集群中，这同时带来了三个好处：
- 单 GPU 专家数量更多 —— 在不将模型碎片化成更小分片的情况下增加容量。
- 更高的容量因子，更少的 token 丢弃 —— 路由程序可以将额外的 token 分配给忙碌的专家而不会耗尽内存，从而稳定训练。
- 更少的 All-to-All 流量 —— 较小的专家并行度缩小了集合通信的占用空间，使得通信不再主导墙上时钟时间。
在实践中，MI325X 的内存裕量和带宽将通常的网络瓶颈转化为插槽内工作负载，即使在大规模环境下也能使 MoE 训练保持在计算密集型状态。
并行策略设计
我们支持在 AMD GPU 上进行多维并行策略：
- 专家并行 (EP) 将专家池划分到节点内的 GPU 上，并在模型无法适应单个节点时跨节点划分。这使我们能够增加专家数量，同时保持每个 GPU 的专家状态在有限范围内，并避免过于微小的专家分片。
- 完全分片数据并行 (FSDP) 在数据并行组中切分参数、梯度和优化器状态，将每个 GPU 的模型占用空间减少约数据并行度倍，从而使大型检查点能够舒适地放入 MI325X 的 256 GB HBM 中。
- 流水线 + 虚拟流水线并行 (PP + VPP) 将层划分为多个阶段并交错进行前向和后向步骤；VPP 进一步将每个阶段细分为虚拟块，以便微批次交错运行并减少流水线空闲时间。
通过 EP 处理宽度，FSDP 处理内存，PP + VPP 重叠前向和后向传递以减少流水线气泡，该技术栈协调了 GPU 算术、HBM 带宽和 400 Gb RoCE 流量——确保没有任何单一维度成为瓶颈。
图 3：PP 调度——蓝色/橙色热图显示了 VPP 如何重叠前向/后向传递
Primus-Turbo：AMD 加速库
Primus-Turbo 是一个致力于 AMD GPU 上大规模模型训练的高性能加速库。在此技术栈中，Primus-Turbo 是本工作所使用的 ROCm 优化内核库。它提供高性能算子（GEMM、注意力机制、分组 GEMM）、通信原语（包括 All-to-All/DeepEP）、优化器实用程序以及用于在 AMD Instinct™ GPU 上进行高效训练的低精度内核（如 FP8）。
图 4：Primus-Turbo 架构
关键组件
| 组件 | 功能 |
| GEMM | – 支持 BF16/FP16
– 支持 FP8 (E4M3/E5M2)，具有张量级、行级和块级量化 |
| 分组 GEMM | – 支持 BF16/FP16
– 支持低精度（如 FP8 (E4M3/E5M2)），具有张量级、行级和块级量化 |
| 注意力 | – 支持 BF16/FP16
– 支持低精度（如 FP8 (E4M3/E5M2)），具有块级量化 |
| DeepEP | – 支持节点内和节点间
– 支持 Mellanox；Broadcom NIC 和 AMD AINIC（均在开发中） |
| All2All | – 支持低精度（如 FP8 (E4M3/E5M2)），具有张量级和行级量化 |
| 逐元素算子 (Elementwise Ops) | – 支持归一化、激活函数、RoPE 等 |
集成示例
在 MI325X 集群上的大规模预训练
所有实验均在配备 1024 个 AMD Instinct MI325X GPU 的 TensorWave 集群上运行。每个 GPU 连接一个 400Gb RoCE v2 Broadcom Thor 2 NIC，节点采用三层胖树拓扑结构。该设计提供了全分段带宽，因此专家并行 All-to-All 流量永远不会遇到网络瓶颈。TorchTitan 和 Primus-Turbo 提供了软件栈。
图 5：具有 3 层胖树网络拓扑的 MI325X 集群
通过 Primus-Turbo 优化 DeepSeekV3-671B 性能
首先，我们在 64 个 MI325X 节点上对 DeepSeek-V3-671B 进行了完整运行分析，以确定内核热点，如注意力内核和图 6 中的分组 GEMM。
实验配置
| 层数 | 序列长度 | EP (专家并行度) | PP (流水线并行度) | VPP (虚拟流水线并行度) | 批次大小 |
|---|---|---|---|---|---|
| 64 | 4096 | 8 | 8 | 4 | 16 |
图 6：显示注意力内核和 FP8 分组 GEMM 内核优化目标的内核分析
使用 Primus-Turbo，我们解决了这些问题并测量了每一步后的吞吐量，如图 7 中的条形图所示。
图 7：64 个节点（512 个 GPU）上的逐步性能优化
- 启用 AITER 注意力。 吞吐量提升约 15%。（要了解有关 AITER 的更多信息： AITER 博客）
- 用于稠密层的 FP8 GEMM。 将线性层切换到张量级 FP8，减少了内存流量并使有效 FLOPS 加倍，在此前的运行基础上又增加了约 102%。
- 用于专家的 FP8 分组 GEMM。 最后，结合置换和计算的融合内核消除了微型 GEMM 瓶颈，并额外带来了约 60% 的增益。
叠加在一起，这些优化使得端到端训练速度比基准提升了 2.77 倍，将 MI325X 集群转变为一台计算密集型而非带宽密集型的 MoE 训练机器。
DeepSeekV3-671B 预训练扩展
图 8：DeepSeekV3-671B 在 MI325 集群上的扩展效率
从 256 个 GPU 开始，吞吐量保持在接近理想水平的范围内：在 512 个 GPU 上为 97%，在 1024 个 GPU 的满规模下为 96%，这证明了软硬件结合的技术栈能够协调八路流水线阶段、专家路由和 FP8 内核，且不会浪费时钟周期。
DeepSeekV3 预训练收敛
图 9：DeepSeek-V3 在 MI325X 集群上的损失
为验证速度并未以牺牲质量为代价，我们使用 allenai/c4 数据集在 32 个 MI325X 节点（256 个 GPU）上跟踪了 DeepSeek-V3 FP8 的训练损失。曲线显示 900 次迭代中呈现平滑、单调的下降，与其他平台上的参考表现相匹配，证实了 FP8 内核、专家路由和并行调度按预期收敛。
Llama4-Scout 预训练扩展
图 10：Llama4-Scout 在 MI325X 集群上的扩展效率
随着 FSDP 与 EP=8 的结合，当我们将集群规模从 32 个增加到 512 个 GPU 时，Llama-4 Scout 保持了 100–102% 的扩展效率，证明了相同的优化可以干净利落地迁移到其他 MoE 架构中。
DeepEP — 有效控制专家流量
在上述实验中，我们应用了 EP=8，因为增加专家并行度会使 All-to-All 通信步骤变得代价高昂。为了在增加 EP 度时缓解这种影响，Primus-Turbo 还提供了 DeepEP 来加速 All-to-All 性能。下图比较了在 16 节点 MI325X 集群上训练 DeepSeek-V3 671B（24 层，PP 4 + VPP 3，本地 batch 16）时，标准 All-to-All 实现（红色）与 DeepEP（蓝色）的表现。随着专家并行度 (EP) 从 8 增加到 32，普通 All-to-All 的吞吐量从约 2000 TPS 下降到约 750 TPS，其占总运行时间的比例从 10% 激增至近 50%。DeepEP 将吞吐量保持在 2000–2100 TPS 左右，即便在 EP 32 时，通信的总运行时间占比也控制在约 18%，从而使可扩展性曲线保持平坦。
图 11：不同 EP 度下的 DeepEP 性能分析（DeepSeek-V3 671B, EP=8/16/32, LBS=16, PP=4, VPP=3, 24 层, 16 节点)
Primus-Turbo 构建在 ROCm DeepEP (https://github.com/ROCm/DeepEP) 之上，并进行了额外的性能改进，例如在训练工作负载中无需 CPU/GPU 同步（可在 https://github.com/AMD-AGI/Primus-Turbo 获取）。
结论与后续步骤
我们与 Meta 的联合工作表明，AMD Instinct MI325X 硬件结合 TorchTitan 和开源 Primus-Turbo 库，为现代 MoE 预训练模型提供了生产级的吞吐量。DeepSeek-V3 在高达 1024 个 GPU 上达到了超过 96% 的扩展效率；Llama-4 Scout 显示出更高的线性扩展效率。内核级调整（FP8 GEMM、分组 GEMM 和 Aiter 注意力）将端到端性能提升了 2.77 倍，且我们所有的实现均为开源。
下一步工作
- 流水线演进。 DualPIPE 和其他高级调度将进一步减少深流水线中的空闲气泡。
- 更广泛的内核覆盖。 正在进行的 FP6/FP4 支持和新算子旨在瞄准下一波 MoE 和稀疏张量架构。
- 下一代硬件。 正在为 MI450 GPU 和 Helios AI 机架做准备，以将今天的成果扩展到更大的集群。
- 开源扩展。 对 TorchTitan 的持续支持，以及对 Monarch, TorchForge 等项目的贡献。
总之，这些努力旨在保持 AMD 平台在大规模、高成本效益 MoE 训练领域的前沿地位。
了解更多
开源仓库
- Primus-Turbo: https://github.com/AMD-AGI/Primus
- TorchTitan: https://github.com/pytorch/torchtitan
- AMD 优化的 torchtitan 分支: github.com/AMD-AGI/torchtitan-amd
文档
- Primus-Turbo 简介: https://rocm.blogs.amd.com/software-tools-optimization/primus-large-models/README.html
- 流水线并行可视化: https://github.com/AMD-AGI/Primus/tree/main/tools/visualization/pp_vis
- ROCm 文档: https://rocm.docs.amd.com/en/latest/ https://rocm.docs.amd.com/
- AMD Instinct MI325X: https://www.amd.com/en/products/accelerators/instinct/mi300/mi325x.html
免责声明
第三方内容由拥有该内容的第三方直接许可给您，而非由 AMD 许可给您。所有链接的第三方内容均按“原样”提供，不提供任何形式的担保。使用此类第三方内容完全由您自行决定，在任何情况下，AMD 均不对您因使用第三方内容而产生的任何损害负责。您承担所有风险，并对因使用第三方内容可能引起的任何损害负全部责任。