开发者博客 | 在 PyTorch 中加速大规模混合专家训练

长期以来，训练大规模多专家（MoE）模型一直属于少数具备深厚基础设施和分布式系统专业知识的高级用户。对大多数开发者而言，真正的挑战不在于构建更智能的模型，而在于如何在数百乃至数千个 GPU 上高效地扩展模型，同时避免性能损失。

借助 NVIDIA NeMo Automodel（NVIDIA NeMo 框架中的开源库），开发者现在能够在 PyTorch 中直接使用熟悉的工具来训练大规模的 MoE 模型。NeMo Automodel 基于经过 NVIDIA 性能优化的加速 PyTorch distributed 构建，显著降低了大规模 MoE 训练的技术门槛，推动其实现普及化。

• 简单 – 无需管理复杂的并行架构或专用系统，即可直接在 PyTorch 中训练包含数十亿参数的模型。
• 可访问 – 让研究人员、初创企业及大型公司都能轻松尝试过去难以实现的 MoE 架构。
• 高效 – 支持从 8 个 GPU 扩展至 1000 多个 GPU，同时通过内置优化保持出色的性能与成本效益。

在本文中，您将了解 NeMo Automodel 如何结合 PyTorch 原生的分布式并行能力与 NVIDIA 的加速技术，让大规模 MoE 模型的训练变得更加高效、便捷且易于使用。文中还提供了详细的快速入门指南，帮助您复现基准测试结果、开展自定义实验，并探索各类配置选项，亲身体验这一技术带来的显著优势。

高效地进行大规模 MoE 训练需要应对多个相互关联的挑战。

1. 专家并行：在 GPU 上分布数百名专家，同时避免通信带宽过载。
2. Token 路由效率：快速且高效地将 token 分配给对应的专家。
3. 显存管理：在 GPU 显存限制内对大规模参数进行分片存储。
4. 通信与计算融合：显著降低多对多通信和 token 重排操作的延迟。

由于这些系统性挑战，在 H100 系统上以 BF16 精度实现每 GPU 超过 150 TFLOPS 的性能一直面临较大困难，导致 GPU 的性能未能被充分释放。

NVIDIA NeMo Automodel，是 NVIDIA NeMo 框架中的一个开源库，基于原生 PyTorch 并行架构构建，有效克服了传统实现中的诸多障碍。它将原本仅限于专业机器学习工程师使用的先进基础设施优化，无缝集成到 PyTorch 生态系统中。

开发者现在可以利用 PyTorch API，在配备 BF16 精度的 H100 GPU 上，为多种主流的百亿参数以上 MoE 架构实现单卡超过 200 TFLOPS 的性能。例如，DeepSeek V3 在 256 张 GPU 上实现了每张卡 250 TFLOPS/秒的计算效率。

这使得大规模 MoE 训练成为可能，让更广泛的社区能够基于包含数十亿参数的模型开展研究、实验与创新。

• 全分片数据并行（FSDP）: 对模型参数、梯度及优化器状态在数据并行维度进行分片，以显著降低内存占用。
• 专家并行（EP）: 将 MoE 专家高效地分布到多个 GPU 上，支持单个模型集成数百个专家。
• 流水线并行（PP）: 将模型层划分为多个阶段，实现跨节点的高效内存利用，适用于大规模模型训练。
• 上下文并行（CP）: 通过划分长序列来支持长上下文训练的扩展。

通过采用 NVIDIA Transformer 引擎 内核（包括 CUDNN RMSNorm、CUDNN Linear 和 DotProductAttention），NeMo Automodel 能够加速 Transformer 模块的运行，同时支持多种注意力机制，如多头潜在注意力（MLA）、分组查询注意力（GQA）和滑动窗口注意力（SWA）。

为实现大规模高效训练，NeMo Automodel 集成了 Megatron-Core 提供的先进 token 路由与专家计算组件，专为 MOE 模型训练而设计。

• DeepEP 令牌调度器 (实验性)：通过高效的多对多通信机制以及可选的 permute/unpermute 融合优化，支持将令牌路由扩展至 64 个及以上专家的并行规模。借助 DeepSeek 的 DeepEP 优化技术，NeMo Automodel 能有效降低通信开销，维持专家负载均衡，从而在数百个 GPU 上实现更顺畅的扩展。
• 面向 MoE 专家的 GroupedGEMM：将多个本地专家的计算整合为单一批量 GEMM 操作，有效减少内核启动开销，提升 GPU 利用率。该方法显著增强了吞吐能力与硬件效率，尤其适用于多个专家共享同一设备的场景。

下表展示了在主流 MoE 架构中，使用 BF16 精度的 DGX H100 系统进行预训练的性能基准。

NeMo AutoModel 在多种 MoE 架构和不同数量的 GPU 上均展现出卓越的效率与可扩展性。每个 GPU 的计算性能可达 190 至 280 TFLOPS/s，处理速度最高支持 13000 tokens/s，并可在 8 到 1024 个 GPU 的规模下实现接近线性的扩展效果。例如，DeepSeek V3 671B 模型在 256 个 GPU 上运行时，每个 GPU 可达到 250 TFLOPS/s 的性能表现。这些成果得益于原生 PyTorch 并行技术与 NVIDIA 优化方案的深度融合，为 PyTorch 社区用户提供了高效的硬件利用率和经济可行的大规模 MoE 模型训练解决方案。

更重要的是，它体现了 NVIDIA 对推动 PyTorch 及更广泛开源 AI 生态系统的坚定承诺，使大型模型的训练不仅更高效，而且更具开放性、互操作性，便于整个开发者社区使用。

为开发者带来的主要优势：

• 更快的迭代周期：通过提升吞吐量，加速实验进程与模型开发。
• 更低的训练成本：更高的 GPU 利用率显著减少每次训练所需的 GPU 小时数。
• 可扩展的性能表现：支持从 8 个 GPU 到 1000 多个 GPU 的一致且接近线性的扩展，便于灵活规划基础设施。
• 原生 PyTorch 集成：依托分布式 PyTorch 技术，无需依赖外部模型并行框架，完整保留于 PyTorch 工作流之中。
• 对生态系统的持续投入：体现 NVIDIA 在推动 PyTorch 发展方面的长期承诺，确保未来创新能够直接融入核心框架。
• 面向生产的成熟方案：提供经过实际验证的稳定配置，全面支持主流开源 MoE 架构。

快速开始：训练大型 MoE 模型并对其进行基准测试

对于任何 PyTorch 开发者来说，NeMo Automodel 的入门过程既快速又亲切。

您可以使用提供的基准测试脚本和配置文件复现结果，或借助 NVIDIA 优化的技术来训练您自己的大规模 MoE 模型。

最低要求

建议使用至少 8 个 GPU（每个 GPU 具备 80 GB 显存）以复现基准测试结果，并支持高效的微调实验运行。

按照以下简单步骤，即可运行基准测试或微调实验：

# 1. Pull the NeMo docker image and start a containerdocker pull nvcr.io/nvidia/nemo:25.09docker run -it -d --ulimit memlock=-1 --ulimit stack=67108864 --gpus all nvcr.io/nvidian/nemo:25.09 bash
# 2. Once  inside the container, clone the repo and navigate to Automodelgit clone https://github.com/NVIDIA-NeMo/Automodel.gitcd Automodel

示例： 在 8 个 GPU 上对 Qwen3 MoE 30B 进行基准测试

torchrun --nproc-per-node 8 nemo_automodel/recipes/llm/benchmark.py \  --config examples/benchmark/configs/qwen3_moe_30b_te_deepep.yaml

运行微调

示例： 微调 Qwen3 MoE 30B

注意：

• 请先从 Hugging Face 下载模型检查点： hf download Qwen/Qwen3-30B-A3B
• 若遇到数据集实例化错误，建议升级数据集库至最新版本： pip install --upgrade datasets

torchrun --nproc-per-node 8 examples/llm_finetune/finetune.py --config examples/llm_finetune/qwen/qwen3_moe_30b_te_deepep.yaml

可用配置文件：

• deepseek_v3_te_deepep.yaml – DeepSeek V3 (671B parameter)
• kimi_k2_te_deepep.yaml – Optimized configuration for Kimi K2
• qwen3_moe_30b_te_deepep.yaml – Qwen3 MoE 30B with full NVIDIA optimizations
• gptoss_20b_te_deepep.yaml – GPT-OSS 20B with FlexAttention
• gptoss_120b_te_deepep.yaml – GPT-OSS 120B production configuration

查看文档，以获取完整的性能说明和实施细节。

这一版本标志着通过加速的 PyTorch 实现大规模混合专家（MoE）训练迈出了重要一步，但这仅仅是个开端。

我们正在积极开展以下工作：

扩展模型支持：引入新的 MoE 架构与混合架构。

• 更深层次的优化：持续改进内核层级设计与通信机制，以进一步提升运行效率。
• 技术深度解析：深入阐述 NeMo AutoModel 的 MoE 架构设计及其性能优化技术。
• 更广泛的基准测试：在多种硬件平台与集群配置中扩展性能验证，全面评估系统表现。

我们期待您开始使用 NeMo Automodel 并加入这一探索之旅——尝试不同配置，分享使用成果，并通过 GitHub Issues 提交反馈。您的见解将有助于推动新一代可扩展的开源 AI 训练工具的发展。

Hemil Desai

是 NVIDIA 的高级软件工程师，专注于为生成式 AI 和大规模模型构建可扩展的高性能基础设施。他的技术兴趣涵盖深度学习系统、分布式训练框架和 GPU 级 PyTorch 工作负载优化。Hemil 拥有加州大学洛杉矶分校计算机科学硕士学位和普渡大学计算机科学学士学位。

Wenwen Gao

是 NVIDIA NeMo 的高级产品经理，专注于 LLM 训练框架和微服务。她过去的经验包括 LLM 推理 (NIM) 和推荐系统 (Merlin) 。她拥有学士学位。多伦多大学计算机科学专业博士学位和麻省理工学院斯隆管理学院工商管理硕士学位。

Gordana Neskovic 

是AI/DL产品营销团队的成员，负责 NVIDIA Maxine。在加入 NVIDIA 之前，Gordana曾在VMware、Wells Fargo、Pinterest、SFO-ITT和KLA Tencor担任过各种产品营销、数据科学家、AI架构师和工程职位。她拥有博士学位。圣塔 Clara 大学，塞尔维亚贝尔格莱德大学电气工程硕士和学士学位。

Bernard Nguyen

是 NVIDIA 的深度学习工程总监。他负责开发 NVIDIA NeMo 框架，专注于从单个 GPU 集群到 1000 个规模 GPU 集群的高效生成式 AI 模型开发，涵盖数据管护、分布式并行训练、微调、对齐、部署和评估。此前，他曾在 Meta 领导过 PyTorch 分布式和大规模 AI 系统的开发。