混合专家（MoE）已成为大规模模型训练的主流架构之一。DeepSeek、Qwen 与 Mixtral 等模型用更少的训练算力，达到了与稠密模型相当甚至更好的效果。MoE 的计算效率来自条件计算：把原本所有 token 共享的单个稠密前馈网络，换成多个更小的专家网络加一个可学习路由器，由路由器决定每个 token 激活哪几个 Top-K 专家。
但要在规模化时保持高效并不容易。生产级 MoE 训练会引入稠密模型不存在的问题：token 路由、专家分发与聚合、all-to-all 通信，以及长度不齐的专家 GEMM。路由权重是学出来的，训练过程中分布会明显倾斜，没有两个 batch 产生相同的专家负载；同一个 batch 内，某个专家收到的 token 可能远多于另一个，于是张量长度不齐，而大多数库恰恰针对规整的矩形张量做过高度优化。专家并行下，若分发与合并路径没有优化，通信就会占据主导，GPU 只能空等数据。
NVIDIA 给出的实测数字说明了差距：在 GB300 上训练 DeepSeek-V3，未优化的基线只有 103 TFLOPS/GPU，GPU 间通信占累计 kernel 时间的 84%；在 JAX Python 库与 NVIDIA Transformer Engine 的定向内核优化之后，同一负载提升到 1068 TFLOPS/GPU，约为 10.4 倍。端到端来看，DeepSeek-V3 671B 的吞吐提升约 10 倍，测试在 GB300 NVL72 机架硬件上完成，覆盖 128 到 2048 个 GPU。
第一项关键优化是 grouped GEMM。此前的做法有两种：循环调用多个 GEMM 内核，需要把各专家的 token 计数从设备拷回主机，这一步落在关键路径上，既带来传输延迟又会打断 CUDA graph；或者用 batched GEMM 按最坏情况的 token 容量计算，即便实际用量更少也要补齐，做了额外的无用功。grouped GEMM 把全部专家的矩阵乘放进一次内核调用，每个专家按真实 token 数计算，只处理有有效 token 的区域。Transformer Engine 的 grouped_gemm 与 ragged_dot 由 cuBLAS、cuBLASLt 支撑，即使专家形状不规则也能跑满 Tensor Core；在 Blackwell GPU 上，这条路径还可对专家矩阵乘启用 MXFP8 块缩放，使用 Transformer Engine 的分组量化内核。
第二项关键优化是把专家并行的分发（Dispatch）与合并（Combine）融合进紧耦合的内核路径，底层由专门针对专家并行不规则、不均衡流量调优的 NCCL EP 支撑。NCCL EP 还带有 token 去重机制：当一个 token 被发往同一 rank 上的多个专家，或发往远端 IB 节点上的多个 rank 时，它在网络上只走一次，到达接收节点后再复制，从而节省带宽。此外还有两项配套优化：JAX host offloading 用重物化 API 把中间激活卸载到主机内存，例如在 DeepSeek-V3 训练中把 query 与 value 投影结果卸载出去以省显存；XLA 多流集合通信让编译器把互相独立的集合通信并发调度到不同 CUDA 流，重叠跨节点 InfiniBand 与节点内 NVLink 传输。Latency Hiding Scheduler 会分析复制组并检查死锁风险，判断哪些集合通信可以安全重叠，无需手工标注。
扩展性方面，多机架通常是系统最容易失速的地方，因为通信开销往往比计算增长更快。采用这套 JAX MoE 与 Transformer Engine 组合后，系统在 1024 个 GPU 时仍保持 97% 的扩展效率。NVIDIA 表示这些优化已随内置 Transformer Engine 的 NGC MaxText 容器发布，可以直接复现与扩展；后续计划加入 NVFP4、量化与 GEMM 融合，以及 A2A 重叠。实践建议是先跑参考配置，在小规模 MoE 模型上验证正确性，再逐步放大规模，同时跟踪 step time、TFLOPS/GPU、MFU，以及 grouped GEMM 与 MoE dispatch/combine 的延迟。
WeChat
Profile