研究摘要
专家混合模型（Mixture of Experts, MoE）架构通过扩展模型参数提升了模型质量。然而，在分布式训练场景中，其发展受到显著通信开销和专家负载不平衡的限制。现有方法仅允许通信和计算的粗粒度重叠，虽然略微缓解了通信成本，但同时也显著降低了计算效率。此外，当前解决负载不平衡的方法通常会损害模型质量。我们提出了CCFuser，这是一个专为高效训练MoE模型设计的新框架。CCFuser通过高效的GPU间共享内存访问，取代了MoE架构中典型的昂贵的All2All操作。这使得本地和远程数据可以在融合内核中并发计算，从而显著提高了GEMM（通用矩阵乘法）操作的计算FLOPS。此外，CCFuser通过一种资源高效的专家重分配策略解决了负载不平衡问题，该策略通过等效图变换优化了专家重分配中的计算资源使用，而不会牺牲统计准确性。通过整合这些优化，CCFuser显著提升了GPU的利用效率。
研究背景
图1： MoE专家并行流程
MoE通过动态路由机制将输入分配给不同的子模型（专家），显著扩展模型参数量（如万亿级模型）的同时保持单次推理的计算量恒定，成为训练超大规模模型的高效范式如图1所示。而在在分布式训练中，专家需分布在多个GPU上，导致跨设备通信开销（如All2All操作）和专家负载不均衡，成为制约训练效率的主要瓶颈。
研究现状
在训练MoE模型时，传统的专家并行执行需要在置换和逆置换阶段进行All2All通信，这导致了显著的通信开销。All2All通信严重依赖网络资源，而MoE模型中置换、计算和逆置换阶段之间的固有数据依赖性，使得在这些通信期间GPU计算资源处于空闲状态。All2All通信的持续时间越长，GPU的空闲时间就越长，从而降低了整体效率。
现有的方法如FasterMoE、PipeMoE和ScheMoE等建议将输入数据划分为较小的块（小张量），并策略性地调度这些块的通信和计算子任务。这些策略旨在最大化计算和通信操作之间的重叠。然而，这种重叠仍然是粗粒度的，无法完全消除通信瓶颈，且可能引入额外的调度开销。并且小块计算可能导致GPU计算单元利用率不足，进一步影响效率。
研究方法
图2： CCFuser的通信-计算算子融合
为了解决现有 MoE训练框架的局限性，我们提出了 CCFuser，一个旨在显著提升训练效率和负载均衡的全新框架。CCFuser 主要包含两大创新点：
基于 GPU 共享内存的通用矩阵乘法算子融合技术，有效消除 MoE 训练中的 All2All 通信开销；高效的专家重分配策略，优化 GPU 负载均衡并提升数据局部性，从而提高 GPU 资源利用率。
1. 基于共享内存的 GEMM 算子融合
传统的 MoE 训练框架依赖粗粒度的算子间重叠，而 CCFuser 通过通信与计算一体化融合 的方式，提出了一种优化的GEMM算子，该方法充分利用GPU共享内存，从根本上避免了 GPU 计算资源因独立通信操作而空闲的情况。在此基础上，我们在 GEMM 计算内部的跨线程块（Threadblock）级别 同步执行通信与计算，实现通信开销的隐藏，最大化 GPU 计算资源的利用率。为了实现这一点，我们在 跨 GPU 共享内存中组织 All2All 操作的数据，并设计了一种基于数据块划分的新型数据结构。此外，我们将所有 GPU 内部专家计算全部集成至 GEMM 计算内核中，从而大幅提升 MoE 训练的专家计算效率，解决 MoE 训练的核心瓶颈-GPU 利用率低 的问题。
2. 资源高效的专家重分配策略
在 GPU 之间实现负载均衡的关键是 资源高效的专家重分配算法。我们 重新构造 MoE 层结构，使得 专家分配算法 与 注意力层计算 可并行执行，最大化通信与计算资源的利用率。根据 token 分布特性，我们针对 热门专家和 非热门专家设计了不同的数据分配策略，以提升单 GPU 内部的数据局部性：对于热门专家，采用 专家复制机制，增强数据局部性，减少通信开销。对于非热门专家，允许 跨 GPU 进行专家重分配，有效缓解内存压力，提高整体计算效率。
综合来看，CCFuser 通过高效的 GEMM 算子融合与智能专家调度算法，从根本上优化了 MoE 训练框架的计算与通信效率，为大规模专家模型的训练提供了强有力的支撑。
研究成果
图3： MoE层端到端训练时间对比
本研究提出的 CCFuser框架在分布式MoE模型训练领域取得了显著成果，成功解决了训练过程中通信开销大和负载不平衡问题，为大规模MoE模型的性能训练提供了全新方案。我们将 CCFuser 集成至 PyTorch 以支持 MoE模型训练。在 NVIDIA A100 服务器 上进行测试时，CCFuser 相较于 FastMoE 和 FasterMoE 取得了显著的性能提升，平均加速比分别达到 2.96 倍 和 2.48 倍，最高可达 4.34 倍。
相关研究成果已收录于 CCF-A 类顶级会议 PPOPP 25（ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming）上，论文标题《Harnessing Inter-GPU Shared Memory for Seamless MoE Communication-Computation Fusion》。