使用高级融合内核提高 MoE 训练吞吐量

NVIDIA 发布了针对 MoE(混合专家)模型的自定义融合内核,通过将多个计算步骤合并为单一内核,可带来 1.3 到 2 倍的算子级加速,在 DeepSeek-V3 预训练场景中实现 8% 的端到端性能提升,在 GPT-OSS 预训练场景中提升达 93%。

使用高级融合内核提高 MoE 训练吞吐量

一句话看懂:NVIDIA 发布了针对 MoE(混合专家)模型的自定义融合内核,通过将多个计算步骤合并为单一内核,可带来 1.3 到 2 倍的算子级加速,在 DeepSeek-V3 预训练场景中实现 8% 的端到端性能提升,在 GPT-OSS 预训练场景中提升达 93%。

事件核心:发生了什么

NVIDIA 在 Generative AI Blog 上宣布,利用其 CuTE(CUDA 模板表达式)领域特定语言,开发了针对稠密模型和 MoE 模型的高级融合 MLP(多层感知机)内核。这些内核专门用于攻克 MoE 块训练中的三个主要系统级瓶颈:激活函数导致的内存受限、CPU 同步开销(因专家路由计算在 CPU 上进行)以及张量量化成本。

新内核通过将 GroupGemm(分组通用矩阵乘法)与激活函数(如 SwiGLU、GeGLU、sReLU)以及量化/转置操作融合到单一内核的“epilogue”(后处理)中,消除了中间张量的全局内存读写。特别地,为了解决 GLU 类激活函数无法简单融合的问题,NVIDIA 通过重新打包权重矩阵(将输入权重和门控权重按列拼接),使得同一线程块可以同时访问这两部分数据。这些内核已通过 NVIDIA cuDNN Frontend 提供,并可经由 Transformer Engine 和 Megatron-Core 轻松调用。

为什么重要

MoE 模型已成为当前大语言模型扩展容量的关键技术路径,但因其动态专家路由特性,训练效率受限于数据搬运和同步开销。NVIDIA 此次优化直接命中了训练瓶颈。一方面,1.3x–2x 的算子级加速对于动辄数千 GPU 卡、运行数周的大型训练任务而言,累计节省的时间成本非常可观;另一方面,该方案允许实现“无同步”的 MoE 执行,可配合完整的 CUDA Graphs 技术,减少 CPU 干预,这是将 GPU 利用效率推向极限的重要一步。

相较于此前的“先写后读”模式,NVIDIA 通过硬件感知的软硬件协同设计,将内存操作与 GEMM 计算本身重叠,等于是从底层算力调度逻辑上提升了硬件吞吐。这一举措将进一步拉大拥有完整工具链(CuTE、cuDNN、Transformer Engine)的平台与开源替代方案之间的性能差距。

对用户/开发者/创作者的影响

对于使用 NVIDIA 生态进行大模型训练的开发者和企业来说,这是一个低门槛的性能提升机会。开发者无需手动编写复杂的内核代码,只需升级到 Transformer Engine 或 Megatron-Core 的最新版本,即可自动受益于这些融合内核的加速效果。尤其对于正在训练 DeepSeek-V3 类 MoE 模型或 GPT-OSS 架构的团队,端到端性能的提升意味着同等算力预算下训练时间缩短,或者相同训练时间内模型规模可以扩大。

GamsGo AI

AI 工具推荐

想把多个 AI 模型放在一个入口?

GamsGo AI 集成 ChatGPT、DeepSeek、Gemini、Claude、Midjourney、Veo 等常用模型,适合写作、绘图、视频和日常 AI 工作流。

了解 GamsGo AI

推广链接:通过此链接购买,我可能获得佣金,不影响你的价格。

对于非自建训练基础设施的内容创作者或初创团队,这一技术的影响间接体现为云 GPU 租用成本的潜在下降——如果云服务商跟进部署这些优化内核,用户将获得更快的训练速度或更低的账单。

值得关注的后续

第一,这些内核的实际效果是否能在更广泛的 MoE 模型族(如 Mixtral 8x7B 等)上复现?当前数据仅覆盖 DeepSeek-V3 和 GPT-OSS 两个典型设置,通用性有待社区验证。第二,融合内核是否会影响模型训练的可调试性或精度控制?尤其是在更低精度(FP8/INT8)场景下,Quantize 与 GEMM 的融合对数值稳定性有何影响,值得关注后续更新的技术文档。第三,AMD ROCm 或 PyTorch 原生 Inductor 等竞争方案是否会跟进类似的融合策略?如果开源社区无法获得同等性能的融合实现,NVIDIA 在训练硬件+软件栈上的壁垒将进一步加强。

来源:NVIDIA Generative AI Blog

celebrityanime
celebrityanime
文章: 7899

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注