丽江市网站建设_网站建设公司_虚拟主机_seo优化

https://pytorch.org/blog/efficient-moe-pre-training-at-scale-with-torchtitan/

高效训练像 DeepSeek-V3 和 Llama 4-Scout 这样的大规模混合专家模型（MoE）是现代人工智能面临的挑战之一。这些模型将 GPU、网络和编译器的性能推向了极限。为了应对这一挑战，AMD 与 Meta 的 PyTorch 团队联手，为新款 Instinct MI325X GPU 优化了 AMD 的开源内核库 TorchTitan 和 Primus-Turbo。通过合作，他们在 1024 个 GPU 上实现了接近理想的扩展效果，证明了效率与规模并非不可兼得。

概览
通过利用 TorchTitan、Primus-Turbo 内核以及 MI325X 的硬件能力，我们实现了：

• 通过内核级优化，在 DeepSeek-V3-671B 模型上实现2.77 倍的速度提升。
• DeepSeek-V3-671B 从 128 个 GPU 扩展到 1024 个 GPU 时的扩展效率达到 96%。
• Llama 4-Scout 模型实现了完美的线性扩展（从 32 个 GPU 扩展到 512 个 GPU）。
• 基于 TorchTitan + Primus-Turbo 的完全开源的技术栈。

什么是 TorchTitan？

TorchTitan 是 Meta 推出的一款原生于 PyTorch 的蓝图，专为跨多 GPU 和多节点集群的大规模训练而设计。它将针对现代大语言模型（LLM）和混合专家（MoE）模型的成熟方案打包成一个单一、可配置的训练栈，让您可以将同一套代码路径从早期实验复用到全面规模的运行中。

配置优先的扩展方式

只需在一个 TOML 文件中设置流水线并行、张量并行、数据并行或专家并行的度数，TorchTitan 就能自动构建作业、连接 NCCL/RCCL 通信组，并在一个 GPU 或一千个 GPU 上运行相同的脚本。

广泛的架构覆盖

无论是训练密集型模型（Llama 3）、混合型模型（Qwen 3），还是稀疏型 MoE 模型（DeepSeek-V3, Llama 4），都可以通过统一的 TorchTitan 配置来完成。

模块化设计

像 AMD 的 FP8 Primus-Turbo 这样的优化内核可以无缝集成，无需更改代码。同样的钩子（hooks）也支持未来的训练后工作流。

简而言之，TorchTitan 为任何模型（无论是密集型还是稀疏型）提供了一条统一、可组合的扩展路径——从笔记本电脑上的原型开发到集群上的生产部署。

理解混合专家（Mixture-of-Experts）模型

混合专家（MoE）是一种稀疏激活的 Transformer 模型替代方案。在标准模型中，每个 token 都会流经相同的全连接前馈网络（feed-forward block），因此参数量翻倍意味着计算量也翻倍。而 MoE 层则用一组专业化的“专家”（experts）替换了单个的多层感知机（MLP）。一个可学习的“路由器”（router）会检查每个 token，并将其只发送给其中少数几个（例如 k 个）专家处理。

因为每个 token 只由少数几个专家运行，所以每个 token 的计算预算现在只与这个很小的 k 值成比例，而不是与专家总数成比例。内存占用也遵循相同的模式：如果应用了专家并行，每个 GPU 只存储它所负责的那些专家的参数，从而为提升隐藏层维度、序列长度或词汇表大小留出了空间，而不会超出设备显存限制。专家本身是普通的 MLP 模块，因此增加模型容量就像复制更多份专家一样简单，而无需加宽每一层。

这种经济性使得 MoE 模型的参数量可以攀升到数千亿级别，同时训练速度却能与密集型模型相当。DeepSeek-v3 和 Llama 4 就展示了这种优势，它们通过混合 MoE 层和密集层，在当今的硬件上平衡了模型容量、准确性和效率。

MoE 预训练的挑战

尽管 MoE 模型提供了引人注目的优势，但在大规模训练时也引入了独特的挑战，必须解决这些挑战才能实现高效的分布式训练。

内核效率（Kernel Efficiency）

• 微小的矩阵乘法（GEMM）→ GPU 利用率不足：经过路由后，每个专家只处理少数 token，导致后续的矩阵乘法（GEMM）运算规模变得极小。这些“微型 GEMM”无法让 GPU 的计算单元保持满负荷工作，并引入了显著的启动开销和数据移动开销。

通信（Communication）

• 繁重的 All-to-All 通信 → 网络瓶颈：专家并行（Expert Parallelism）会将 token 的嵌入（embeddings）分散到拥有被选中专家的那些 GPU 上。因此，每次前向传播都会触发繁重的 All-to-All 集体通信操作。随着集群规模的增长，通信时间可能会超过计算时间，成为主要瓶颈。

并行策略（Parallelism）

• 复杂的并行策略 → 流水线不同步导致的空闲“气泡”：现代训练通常会堆叠使用多种并行策略：用于节省内存的完全分片数据并行（FSDP）、用于处理模型深度的流水线并行（PP），以及用于扩展模型宽度的专家并行（EP）。如果它们的微批次（micro-batches）或调度计划出现不同步，就会出现“气泡”——即某些 GPU 闲置而其他 GPU 在工作的时段——导致利用率低下。

路由与稳定性（Routing and Stability）

• 不稳定的路由 → 专家负载不均：可学习的路由器必须在遵守容量限制的同时，将 token 均匀地分配给各个专家。倾斜的 token 数量、专家过载或 token 丢失都可能减慢收敛速度或降低最终模型的质量。

我们从三个方面应对这些预训练的障碍：

1. AMD Instinct MI325X GPU：提供 256GB 的 HBM3E 显存和超过 6TB/s 的带宽，外加 FP8 张量核心，用于将 MoE 专家尽可能地保留在本地，并使训练保持计算密集型。
2. TorchTitan 并行策略(FSDP + PP/VPP + EP)：平衡计算和通信。
3. Primus-Turbo 的 FP8 注意力机制和分组 GEMM (grouped-GEMM) 内核：将 MoE 的微型 GEMM 转化为高占用率、高吞吐量的操作。

AMD Instinct™ MI325X：硬件与软件栈

Instinct MI325X 将 256 GB 的 HBM3E 显存、超过 6 TB/s 的片上带宽以及 Petaflop 级的 FP8/BF16 矩阵核心相结合。如此大的高速显存使得 MoE 模型可以将其大部分专家运行在本地设备上，而不是分散到整个集群中，这同时带来了三个好处：

1. 每个 GPU 可容纳更多专家— 模型容量得以增长，而无需将模型分割成越来越小的碎片。
2. 更高的容量因子，更少的 token 丢失— 路由器可以为繁忙的专家分配额外的 token 而不会耗尽显存，从而稳定训练过程。
3. 更少的 All-to-All 流量— 较小的专家并行度缩小了集体通信的规模，因此通信不再主导总耗时。

在实践中，MI325X 的显存空间和带宽将通常是网络瓶颈的问题转化为了一个在插槽内部（in-socket）即可处理的工作负载，即使在非常大的规模下，也能保持 MoE 训练是计算密集型（compute-bound）的。

并行策略设计

我们在 AMD GPU 上支持多维度并行策略：

• 专家并行 (EP)：将专家池在节点内的 GPU 之间进行分区，当模型大到无法装入单个节点时，则跨节点进行分区。这使我们能够增加专家数量，同时保持每个 GPU 上的专家状态有界，并避免专家分片过小。
• 完全分片数据并行 (FSDP)：将参数、梯度和优化器状态在数据并行组内进行切分，从而将每个 GPU 的模型占用空间减少大约数据并行的度数，使得大型模型检查点可以轻松地放入 MI325X 的 256 GB HBM 显存中。
• 流水线并行 + 虚拟流水线并行 (PP + VPP)：将模型的层拆分成多个阶段，并交错执行前向和后向传播步骤；VPP 进一步将每个阶段细分为虚拟块（virtual chunks），使微批次错开执行，从而减少流水线的空闲时间。

通过 EP 来吸收模型宽度，FSDP 来处理内存占用，以及 PP + VPP 来重叠前向和后向传播以减少流水线气泡，该技术栈协调了 GPU 的算术运算、HBM 带宽和 400 Gb RoCE 网络流量——确保没有任何一个维度成为瓶颈。

图 3：PP 调度——蓝色/橙色热图展示了 VPP 如何重叠前向/后向传播过程。

Primus-Turbo：AMD 的加速库

Primus-Turbo 是一个专为 AMD GPU 上的大规模模型训练而设计的高性能加速库。在此技术栈中，Primus-Turbo 是本次工作中使用的经过 ROCm 优化的内核库。它为在 AMD Instinct™ GPU 上进行高效训练提供了高性能算子（GEMM、注意力、分组 GEMM）、通信原语（包括 All-to-All/DeepEP）、优化器工具以及如 FP8 等低精度内核。

图 4：Primus-Turbo 架构

关键组件

集成示例：
（此处原文有代码示例，但未在文本中提供）

在 MI325X 集群上进行大规模预训练

所有实验均在一个配备了 1024 个 AMD Instinct MI325X GPU 的 TensorWave 集群上运行。每个 GPU 连接到一个 400Gb RoCE v2 Broadcom Thor 2 网卡，节点之间采用三层胖树（fat-tree）拓扑结构连接。这种设计提供了全对分带宽（full-bisection bandwidth），因此专家并行的 All-to-All 流量永远不会遇到网络瓶颈。软件栈由 TorchTitan 和 Primus-Turbo 提供。

图 5：采用 3 层胖树网络拓扑的 MI325X 集群

通过 Primus-Turbo 优化 DeepSeekV3-671B 性能

首先，我们在 64 个 MI325X 节点上对 DeepSeek-V3-671B 的一次完整运行进行了性能分析，以识别内核热点，如图 6 中的注意力内核和分组 GEMM。

实验配置：

图 6：内核分析显示注意力内核和 FP8 分组 GEMM 内核是优化目标

我们使用 Primus-Turbo 解决了这些瓶颈，并在每一步后测量吞吐量，结果如图 7 的条形图所示。

图 7：在 64 个节点（512 个 GPU）上逐级优化的性能表现

1. 启用 AITER attention：将吞吐量提升了约15%。（要了解更多关于 AITER 的信息，请参见 AITER 博客）
2. 为密集层启用 FP8 GEMM：将线性层切换到张量级（tensorwise）FP8，减少了内存流量，并将有效浮点运算能力翻倍，比上一次运行额外增加了约102%的性能。
3. 为专家层启用 FP8 分组 GEMM：最后，一个融合了置换（permutation）和计算的内核消除了微型 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 集群上的损失曲线

为了验证速度没有牺牲质量，我们在 32 个 MI325X 节点（256 个 GPU）上使用allenai/c4数据集追踪了 DeepSeek-V3 FP8 训练的损失（loss）。曲线显示，在 900 次迭代中，损失平稳、单调下降，与其他平台上的参考行为相匹配，证实了 FP8 内核、专家路由和并行调度能够按预期收敛。

Llama4-Scout 预训练扩展性

图 10：Llama4-Scout 在 MI325X 集群上的扩展效率

通过将 FSDP 与 EP = 8 相结合，当我们将集群规模从 32 个 GPU 增加到 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（每秒处理的 token 数）下降到约 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 attention——将端到端性能提升了 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 均不对任何第三方内容对您负责。您承担所有风险，并对因使用第三方内容而可能产生的任何损害负全部责任。

“Scaling efficiency”（扩展效率）是一个衡量系统性能如何随着资源（在这里是 GPU 数量）增加而提升的指标。

简单来说，它回答了这样一个问题：“如果我把 GPU 的数量翻倍，我的训练速度能接近翻倍吗？”

一个完美的、理想的系统，其扩展效率是 100%。这意味着：

• GPU 数量从 256 个增加到 512 个（翻倍），训练速度（吞吐量）也正好翻倍。
• GPU 数量从 256 个增加到 1024 个（变为 4 倍），训练速度也正好变为 4 倍。

结合原文例子来解读

原文中说：

“Starting at 256 GPUs, throughput stayed within striking distance of ideal: 97% at 512 GPUs and 96% at the full 1024-GPU scale”

这句话的意思是，他们以 256 个 GPU 时的性能作为基准（100% 效率的基础）：

1. 扩展到 512 个 GPU 时，效率为 97%：

   • 理想情况：GPU 数量翻倍（从 256 到 512），训练速度也应该翻倍。
   • 实际情况：训练速度达到了理想速度的97%。也就是说，速度几乎翻倍了，只损失了 3% 的效率。这是一个非常出色的结果。

2. 扩展到 1024 个 GPU 时，效率为 96%：

   • 理想情况：GPU 数量变为 4 倍（从 256 到 1024），训练速度也应该变为 4 倍。
   • 实际情况：训练速度达到了理想速度的96%。也就是说，在增加了这么多 GPU 后，系统仍然保持了极高的效率，只损失了 4% 的效率。

为什么扩展效率通常不是 100%？

在现实世界中，扩展效率几乎不可能达到 100%。当你增加更多的 GPU 时，会引入额外的开销，导致性能损失。主要原因包括：

• 通信开销 (Communication Overhead)：这是最大的瓶颈。GPU 之间需要不断地交换数据（如梯度、模型参数等）。GPU 越多，通信网络就越拥挤，协调工作的耗时也越长。就像一个会议室里人越多，大家互相沟通达成一致所需的时间就越长。
• 同步开销 (Synchronization Overhead)：所有 GPU 必须在特定步骤（例如，更新模型权重之前）等待彼此完成工作。总会有一个最慢的 GPU，导致其他所有 GPU 都在“空等”，造成资源浪费。
• 负载不均衡 (Load Imbalance)：任务可能没有被完美地平均分配到每个 GPU 上。有些 GPU 可能比其他 GPU 更忙，导致整体效率下降。

总结

“Scaling efficiency”（扩展效率）是衡量分布式系统可扩展性的关键指标。

一个高的扩展效率（如文中的 96%、97%）意味着该系统（包括 AMD 的 GPU 硬件、Primus-Turbo 软件库、TorchTitan 框架和网络架构）设计得非常出色，成功地将增加 GPU 带来的通信等额外开销降到了最低。这证明了他们的解决方案在进行超大规模模型训练时，能够有效地利用硬件资源，实现“人多力量大”的效果，而不是“人多添乱”。

这里的TPS指的是**训练（Training）**时的吞吐量。

更具体地说，TPS 在这个上下文中代表Tokens Per Second，即每秒处理的 Token 数量。

为什么是训练而不是推理？

1. 全文语境：整篇文章的核心是关于**预训练（Pre-training）**大规模 MoE 模型，讨论的是如何优化训练过程中的挑战，如内核效率、通信瓶颈和并行策略。推理（Inference）是模型训练完成后的应用阶段，文章并未涉及。

2. 指标的含义：在训练 LLM 时，“Tokens Per Second” 是衡量训练速度和效率的核心指标。它表示整个训练系统（包括所有 GPU、网络和软件）一秒钟内能够完成前向传播、计算损失、反向传播和参数更新等一系列操作所处理的 token 总量。TPS 越高，意味着训练同样的数据集所需的时间越短，训练效率越高。

3. 对引文的详细解读:

   “As the expert-parallel degree (EP) rises from 8 to 32, plain All-to-All throughput sinks from roughly 2000 TPS to about 750 TPS and its share of total runtime balloons from 10% to nearly 50%. DeepEP holds throughput steady around 2000–2100 TPS”

   这句话描述了一个训练中的性能实验：

   • 实验变量：他们增加了专家并行（EP）的度数，从 8 增加到 32。这意味着 MoE 模型中的“专家”被分散到了更多的 GPU 上。
   • 负面影响：增加 EP 会导致 GPU 之间的通信量（All-to-All 操作）急剧增加。使用标准的通信方法（plain All-to-All），这个通信开销变成了巨大的瓶颈。
   • 性能下降：由于系统大部分时间都在等待数据通信而不是进行计算，整个训练的吞吐量从大约 2000 Tokens/Second 急剧下降到 750 Tokens/Second。
   • 解决方案：当换用 AMD 优化的DeepEP通信库后，即使 EP 增加到 32，通信瓶颈也被有效解决了，训练吞吐量得以保持在 2000-2100 Tokens/Second 的高水平。

总结

所以，这里的 TPS 是一个衡量训练效率的指标。它清晰地展示了在分布式训练中，通信开销是如何直接影响整体训练速度的，以及 Primus-Turbo 中的 DeepEP 技术是如何成功解决这个问题的。

在推理场景下，TPS 同样表示 “Tokens Per Second”，但它衡量的是模型生成文本的速度，即一个已经训练好的模型每秒能产生多少个 token。这是衡量模型响应速度的指标，与衡量学习速度的训练 TPS 是两个不同阶段的概念。