在大规模人工智能应用中，高效训练诸如 DeepSeek-V3 和 Llama 4-Scout 等海量混合专家模型 (MoE) 是一项重大挑战。这些模型将 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 个 GPU 扩展至 1,024 个 GPU
- 完美的线性扩展：Llama 4-Scout (32 → 512 个 GPU)
- 完全开源的栈：基于 TorchTitan + Primus-Turbo 构建
什么是 TorchTitan？
TorchTitan 是 Meta 为多 GPU 和多节点集群的大规模训练而设计的 PyTorch 原生蓝图。它将现代大语言模型 (LLM) 和 MoE 模型的成熟配方封装在一个可配置的统一训练栈中，使开发者能够将代码路径从早期实验无缝复用到全规模运行中。
配置优先的扩展
只需在单个 TOML 文件中设置流水线并行、张量并行、数据并行或专家并行度，TorchTitan 即可自动构建作业、连接 NCCL/RCCL 组，并在单个或一千个 GPU 上运行相同的脚本。
广泛的架构覆盖
通过统一的 TorchTitan 配置，即可训练稠密模型 (Llama 3)、混合模型 (Qwen 3) 和稀疏 MoE 模型 (DeepSeek-V3, Llama 4)。
模块化设计
诸如 AMD 的 FP8 Primus-Turbo 等优化内核可直接插入，无需更改代码；相同的钩子（hooks）还支持未来的训练后工作流。
简而言之，TorchTitan 为扩展任何模型（无论是稠密还是稀疏）提供了一条统一、可组合的路径，从笔记本原型到集群生产环境均可适用。
理解混合专家模型 (MoE)
混合专家模型 (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）、优化器实用程序以及低精度内核（如 FP8），以在 AMD Instinct™ GPU 上进行高效训练。
图 4：Primus-Turbo 架构
关键组件
| 组件 | 特性 |
| GEMM | – 支持 BF16/FP16
– 支持 FP8 (E4M3/E5M2)，包括张量级、行级和块级量化 |
| 分组 GEMM (Grouped 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 网卡，节点以三层胖树 (Fat-Tree) 拓扑结构连接。该设计提供了全分段带宽，因此专家并行 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：MI325 集群上 DeepSeekV3-671B 的扩展效率
从 256 个 GPU 开始，吞吐量保持在理想值的范围内：在 512 个 GPU 时为 97%，在 1024 个 GPU 全规模下为 96%，证明该软硬件组合能够协调八路流水线阶段、专家路由和 FP8 内核，而不会浪费计算周期。
DeepSeekV3 预训练收敛
图 9：MI325X 集群上的 DeepSeek-V3 loss（损失函数值）
为了验证速度没有牺牲质量，我们在 32 个 MI325X 节点（256 个 GPU）上使用 allenai/c4 数据集跟踪了 DeepSeek-V3 FP8 的训练损失。曲线显示在 900 次迭代中呈现平滑、单调的下降，与其他平台上的参考表现相匹配，证实了 FP8 内核、专家路由和并行调度均按预期收敛。
Llama4-Scout 预训练扩展
图 10：MI325X 集群上 Llama4-Scout 的扩展效率
随着 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，本地批次 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 fork: 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 均不对因任何第三方内容而给您造成的损害负责。您承担所有风险，并对因使用第三方内容而可能产生的任何损害负全部责任。