混合专家并行优化MoE训练通信

在大语言模型（LLM）训练中，针对超大规模混合专家（MoE）模型的专家并行（EP）通信具有挑战性。EP通信本质上是全互连（All-to-All）的，但由于其动态性和稀疏性（每个AI token仅路由到topk个专家，而非所有专家），实现和优化颇具难度。

本文详细介绍了一种高效的MoE EP通信解决方案——混合专家并行（Hybrid-EP），以及它在某机构的Megatron框架系列、某机构的Quantum InfiniBand和Spectrum-X以太网平台上的使用。本文还深入探讨了Hybrid-EP在实际模型训练中的有效性。

超大规模MoE模型训练的效率挑战

DeepSeek-V3是新一代大规模细粒度MoE模型的代表。此类模型通过“超参数规模稀疏激活”来平衡计算开销与模型性能，但也给现有大模型训练框架带来了严峻挑战：

• 通信效率瓶颈：MoE模型依赖并行专家，需要频繁的全互连通信。随着专家数量增加，EP通信负担加重。在DeepSeek-V3中，未经优化时通信时间可能占整体训练时间的50%以上。
• 负载不均衡：动态路由机制导致部分“热门专家”接收的token数多于平均水平，而“冷门专家”利用不足，造成设备间计算负载不均和算力浪费。在专家数量与激活专家数量持续增加的细粒度场景下，此问题更为突出。
• 框架适配挑战：当今的MoE模型对并行策略、低精度计算和动态资源调度提出了更高且更复杂的要求。同时需要优化以充分发挥新一代硬件架构（如某机构的Blackwell、Quantum InfiniBand和Spectrum-X以太网）的潜力。

MoE训练框架优化与通信方案

某机构的Megatron Core——一个开源的大规模模型训练库——是训练超大规模MoE模型的关键基础。其核心优势包括：

• 多维并行策略：支持张量并行（TP）、序列并行、流水线并行（PP）、MoE专家并行（EP）等多种策略，并可灵活组合以适应多样且复杂的训练任务。
• 资源与效率优化：集成FP8混合精度训练、激活值卸载、分布式优化器和细粒度重计算功能，以减少GPU显存消耗并全面支持模型训练。集成了多种高效算子（如MLA、Attention、MLP等），并提供多种融合优化和流水线调度策略以提升计算性能。
• MoE专用适配：完全支持主流MoE模型（如DeepSeek、Mixtral和Qwen），实现高效、可扩展的训练。

Hybrid-EP：一种高效的通信优化方案

Hybrid-EP是一个新设计的MoE EP通信库。它利用某机构平台的软硬件进步，在RDMA-NVLink混合网络架构中实现接近硬件极限的通信带宽，并最小化GPU硬件资源使用。

它实现了MoE EP通信中的两个核心算子：

• 分发（dispatch）：将注意力算子输出的token路由到对应的专家。
• 合并（combine）：将专家输出的token路由回注意力算子。

同时添加了路由和信息处理支持，以实现完整的EP通信。

设计目标与核心优化方向：

• 利用某机构平台最新通信技术，如用于NVLink横向扩展网络数据通信的TMA指令，以及用于RDMA网络的底层IBGDA网络技术。
• RDMA与NVLink混合网络通信：通过结合节点内NVLink和节点间RDMA，最大化跨域带宽以提升算法带宽。
• 数据流水线：将数据切割为细粒度块，并通过多级通信数据流水线流式处理，掩盖大部分通信和动态路由的延迟，使EP带宽接近高度优化的标准静态全互连。
• 最小化GPU流式多处理器（SM）使用率，以最大化通信与计算的重叠。Hybrid-EP用更少的SM达到峰值通信带宽，从而为计算留出更多SM。
• 低精度原生支持：分发算子支持FP8/BF16，合并算子支持BF16。

Hybrid-EP将每个CUDA块设计为一个独立的数据通道，占用一个SM运行完整的数据流水线。CUDA块内的不同线程束组处理不同的流水线阶段。各CUDA块并行运行，处理不同的数据块，块间无需同步和通信。

图2中的虚线框表示RDMA网络通信使用的流水线阶段。RDMA线程束组负责使用IBGDA技术将网络流量传输到RDMA网卡（NIC），完成同一机架内GPU之间以及不同节点之间（例如所有节点中的0号GPU）的网络通信和token数据传输。

G2S线程束组负责将本GPU拥有的token数据以及同一机架上其他节点GPU传输来的token数据，读入SM内部的共享内存先进先出（FIFO）队列。S2G线程束组将来自SM内部共享内存FIFO队列的token数据，写入本节点所有GPU（包括本GPU自身）输出缓冲区的对应位置。

在此过程中，token根据路由映射中的信息进行路由和传输，避免发送不需要的token数据。每个CUDA块使用此数据流水线按分配顺序处理数据块中的token数据。不同CUDA块处理不同的数据块，但使用相同的数据流水线。

与分发算子类似，虚线框部分仅用于RDMA网络通信。由于合并算子需要对token执行高精度累加操作（目前仅在SM内的CUDA核心上执行），这些操作必须分层累加。

在多节点情况下，相关的节点内线程束组首先完成token在节点内部的部分累加工作，然后RDMA线程束组将部分累加后的token发送到节点间同一机架上的GPU。最后，节点间相关的线程束组完成全局累加工作以得到最终结果。

在单节点情况下，节点内的累加工作直接由节点间相关的线程束组完成。在此过程中，输入token由对应的G2S线程束组读入SM内部的G2S FIFO队列，然后对应的归约线程束组在CUDA核心中累加token。结果存储在SM内部的S2G FIFO队列中，并交给TMA单元保存到GPU显存。

性能测试

Hybrid-EP在多种硬件平台上进行了测试。测试条件包括：隐藏维度8192，数据类型BF16（仅传输token），每卡注意力token数4096，每卡专家数2，路由映射从均匀分布随机生成，TopK为8，使用某机构Quantum InfiniBand网络。

• 单节点（某机构DGX Hopper平台，8张H100 GPU）：Hybrid-EP仅用8个SM即可填满NVLink带宽。
• 四节点（共32 GPU）：同一机架每4张DGX H100 GPU配有一个400 Gbps的ConnectX-7 NIC，通过Quantum InfiniBand网络连接。Hybrid-EP仅需约4个SM即可接近NIC的最大带宽。
• 大规模NVLink网络（某机构Grace Blackwell，GB200NVL36，NVLink域包含36 GPU）：Hybrid-EP仅需16个SM即可填满NVLink带宽。

实际案例与集成

Hybrid-EP基于模板和CUDA C++实现，接受输入输出缓冲区地址。要在基于PyTorch的Megatron Core框架中使用，需要进行一些集成工作。它现已可在DeepEP/Hybrid-EP分支获取，并提供了可直接调用的PyTorch算子，方便用户快速完成集成和测试。

由于Hybrid-EP内核仅接受指针参数且不负责内存管理，需要设计合理的缓冲区管理和分配机制。根据使用场景，Hybrid-EP缓冲区大致分为两类：

• 注册缓冲区：经特殊注册、可被其他rank内核访问的GPU内存，是唯一的全局静态缓冲区。
• 普通缓冲区：通过cudaMalloc分配的GPU内存，可由PyTorch分配器管理，通常不全局唯一。

由于缓冲区申请和注册操作耗时，理想情况下仅在Hybrid-EP初始化阶段完成。但MoE模型是动态的，每轮迭代当前rank接收的token数量会变化，从而改变所需缓冲区大小。为此，可采用最坏情况预分配策略，申请一个上限的大缓冲区，以确保所有token都能汇聚到同一rank。由于此缓冲区全局唯一，整体GPU显存使用仍然可控。

在Grace Blackwell上的优化实践

某机构的Megatron Core已在Grace Blackwell平台上集成Hybrid-EP，并可针对不同类型的MoE模型进行优化。实际测试结果包括：

• DeepSeek-V3（256专家，topk-8场景）：使用Hybrid-EP相比DeepEP获得了约14%的性能提升（未增加MTP）。
• Megatron-FSDP：Hybrid-EP仍带来约8%的性能提升。
• Qwen 3 235B场景：BF16场景下提升约5.5%，MXFP8场景下提升约9.9%。