用 NVIDIA Transformer Engine
来自 NVIDIA Developer · 期号 2026-09-15
摘要
MoE 模型训练时，每个 token 被动态分配给不同的专家，导致每个专家收到的 token 数量不固定，形成不规则的张量。传统做法要么丢弃溢出 token（损害模型质量），要么填充到固定容量（浪费算力）。 NVIDIA Transformer Engine 与 JAX 配合，用 grouped GEMM 内核一次性处理所有专家的矩阵乘法，每个专家只算自己实际拿到的 token，不填充也不丢弃。专家并行中的 dispatch 和 combine 阶段被融合成一条紧密路径，由 NCCL EP 驱动，重复 token 只传一次，减少网络流量。 在 DeepSeek-V3 671B 训练上，这套优化把吞吐从 103 TFLOPS/GPU 提升到 1,068 TFLOPS/GPU，提升 10.4 倍。在 GB300 NVL72 硬件上，1024 块 GPU 的扩展效率达到 97%。 原文是技术详解，附了配置指南和代码路径。
深读
导读 MoE 模型训练的核心难题在于 token 到专家的动态分配会产生不规则张量，传统方案要么丢弃 token 损害模型质量，要么填充到固定容量浪费算力。NVIDIA Transformer Engine 与 JAX 配合，通过 grouped GEMM 和优化后的专家并行通信，在 DeepSeek-V3 671B 训练上将吞吐从 103 TFLOPS/GPU 提升至 1,068 TFLOPS/GPU，提升 10.4 倍，并在 1,024 块 GPU 上维持 97% 的扩展效率。读完能理解这套无丢弃 MoE 训练栈的具体技术组件和性能收益。 关键信息 - 在 NVIDIA GB200 上，未优化的 DeepSeek-V3 训练基线仅达到 103 TFLOPS/GPU，其中 GPU 间通信占累计内核时间的 84%。 - 应用 Transformer Engine 优化后，吞吐提升至 1,068 TFLOPS/GPU，提升 10.4 倍。 - 在 NVIDIA GB300 NVL72 硬件上，1,024 块 GPU 的扩展效率达到 97%。 - 优化栈包含四个主要组件：cuBLAS GroupedGEMM、XLA 多流集合通信、MXFP8 分组量化、主机激活卸载，以及优化后的专家并行（EP）实现。 - NCCL EP 通信后端专为专家并行中不均衡的流量模式设计，采用 token 去重机制：同一 token 被路由到同一节点上的多个专家时，只在网络中传输一次。 逐节讲解 MoE 训练的核心挑战 MoE 模型用多个小型专家网络替代单个密集前馈网络，由学习到的路由器决定每个 token 激活哪些专家。这带来了密集模型没有的瓶颈：token 路由、专家分发与收集、全对全通信，以及不规则专家 GEMM。由于路由器是学习得到的，训练过程中专家负载分布可能严重倾斜，每个专家收到的 token 数量不固定，形成不规则张量。大多数库高度优化了期望均匀矩形数据结构的张量操作，无法高效处理这种不规则性。在专家并行中，token 必须被分发到对应专家的 GPU，输出再被合并并恢复原始顺序，如果分发与合并路径未优化，通信将占主导地位，GPU 利用率低下。 无丢弃 MoE 与容量式 MoE 的对比 容量式 MoE 为每个专家分配固定 token 预算，溢出部分要么丢弃（损害模型质量），要么填充（浪费算力和内存）。无丢弃 MoE 则处理每个 token，无论负载多么不均，这对系统提出了更高要求。MegaBlocks 论文通过将专家计算重新表述为块稀疏矩阵乘法来解决此问题，允许每个专家处理不同数量的 token 而无需丢弃或填充，这需要新的块稀疏 GPU 内核、优化的 grouped GEMM，以及专为可变 token 数量设计的分发与合并原语。 Transformer Engine 的关键…