冯瑶 居欣霞 张若涵 刘丽媛 郝世博 李某 董晨宇 王守航 沈叶龙 高建锋 尚景波
: 项目负责人； : 核心贡献者；（正在开发中）
UCSD, 微软
首次发布于 2025 年 6 月 30 日 | GitHub:
DenseMixer 
TL;DR
我们介绍了 DenseMixer —— 一种新颖且有效的 MoE 后训练方法，使 MoE 更易于训练并性能更佳。
通过将一个额外的前向传递在不活跃的专家上用于精确路由梯度，DenseMixer 在不同规模的 MoE（7B、14B、30B）、架构（有/没有共享专家）、预训练方法（从头开始/循环使用）和后训练数据类型（指令/长 CoT 数据）中始终优于传统方法。
我们提供了一个即插即用实现，让 MoE 后训练变得简单，只需 pip install densemixer。它与现有库（如 transformers, llama-factory, open-instruct, verl）完全兼容，并支持参数高效的方法（如 LoRA），无需更改推理过程。
图 1 展示了 Qwen3-30B MoE 在常规方法和 DenseMixer 后训练的性能提升。结果显示了解码参数：温度=0.6，top-p=0.95。其他配置下的结果见实证结果部分。
MoE 训练中的问题
MoE 的训练难度远高于密集模型。MoE 唯一的创新是稀疏路由机制——通常通过 Top-K 路由器实现，该机制在数学上不可微分。这一问题阻碍了直接反向传播，并使梯度计算复杂化。我们在技术细节部分详细说明了这一点。
介绍 DenseMixer
为了解决不可微问题，我们引入了用于 MoE 后训练的 DenseMixer，通过在前向过程中减少不活跃专家的计算，实现更精确的路由梯度估计。具体细节见技术细节部分。
DenseMixer 在不同规模的 MoE（7B、14B、30B）、架构（有/无共享专家）、预训练方法（从头开始/循环使用）以及后训练数据类型（指令/长 CoT）中，始终优于传统方法。
它适用于任何使用 Top-K 路由和反向传播的 MoE，可以即插即用，兼容现有训练库（transformers, llama-factory, open-instruct, verl）和参数高效方法（例如 LoRA）。
从传统方法切换到 DenseMixer，只需进行以下更改：
Shell
拷贝
# Your current MoE training
python your_moe_training_script.py
# Shift to DenseMixer (no code changes needed!)
pip install densemixer
densemixer setup
export DENSEMIXER_ENABLED=1
python your_moe_training_script.py
请参见我们的 GitHub 仓库获取更多详细信息。 →
DenseMixer 实证结果
我们进行了不同规模、架构和方法的实验，使用相对较短的指令和长推理数据集进行训练。
模型与数据集
我们选取以下三个基础模型进行后训练实验。
对于 OLMoE-1B-7B 和 Qwen1.5-MoE-A2.7B，我们主要采用了 DeepseekAI 最近 ESFT 论文 [Wang et al., 2024] 中的训练和测试数据集。
对于 Qwen3-30B-A3B-Base，我们发现用上述数据训练几乎无法提升，甚至可能降低其下游性能。因此，我们用长推理数据训练，并在需要推理的数学/编程挑战测试中测试它。
在数学领域，我们使用斯坦福 S1 数据集的一个过滤子集进行后训练，该子集包含约 1000 个训练样本，这些样本的推理路径是从 Deepseek-R1 中提炼出来的。
在编码领域，我们使用 Llama-Nemotron-Post-Training-Dataset 的筛选子集进行后训练，该数据集包含约 35K 个训练样本。
Baselines
常规 → 标准 MoE 训练
这是训练 MoE 的默认方法，大多数情况下采用这种方式。在这种方法中，非可微分的 Top-K 操作在自动微分中被视为零梯度。
ESFT → 专家精细调整
ESFT 由 DeepseekAI 提议，通过仅更新特定任务中预选的专家来利用每个专家的专业知识。它有两种变体：(1) ESFT-gate 根据平均评分选择专家；(2) ESFT-token 根据 token 选择率选择专家。专家选择非常微妙，对整体性能影响很大。
冻结路由器，只更新 MoE 中的其他组件。
Frozen Router 是由 Unsloth 在其关于如何微调 Qwen3-MoE 的教程中提出的，基于观察发现路由器训练时往往过于脆弱，因此被冻结。
此外，我们在参数高效微调（PEFT）环境中比较这些方法。其中，LoRA 仅应用于除路由器外的所有模块。由于 ESFT 本身已经是 PEFT 方法，我们没有为其添加 LoRA 版本。
实证结果
对于每种方法，我们对训练超参数（学习率和批量大小）进行网格搜索，并报告最佳性能。更多详细信息请见 Github 仓库。
OLMoE-1B-7B （参数总量 70 亿）
OLMoE-1B-7B 模型的结果如下。
Qwen1.5-MoE-A2.7B （参数总量 140 亿）
Qwen1.5-A2.7B 模型的结果如下。
Qwen3-30B-A3B-Base (参数总量 30B)
对于 Qwen3-30B-A3B 模型，我们报告了在三种解码超参数（温度和 top_p）下评估指标，因为长 CoT 评估不稳定。
使用 Nemotron-Code 数据集（35K 训练样本）后训练的检查点结果
使用斯坦福 S1 数据集（1000 个训练样本）后训练的检查点结果
Findings
广泛的实验结果表明，DenseMixer 在不同规模、架构和预训练方法的 MoE 模型中，以及在多样化的训练和评估数据集上，始终优于传统的 MoE 训练。这种性能优势在全微调和参数高效微调（如 LoRA）方法中都表现出来，并且在不同的生成超参数设置下依然保持稳定。
我们还发现，Frozen Router 有时能超越常规训练，因为常规训练使用不精确的梯度进行端到端的路由训练。然而，它始终不如 DenseMixer 表现好，后者得益于更准确的梯度信号。
技术细节
对于那些对背后的数学和实现方式感兴趣的人，我们提供了问题的详细步骤和解决方案。
核心问题：不可微分的路由
传统方法如何处理这个问题
DenseMixer: 用计算换梯度精度
处理更复杂的情况
效率分析
Memory
在标准的 MoE 训练中，模型权重始终加载到 GPU 上。因此，在这种情况下，DenseMixer 的内存使用几乎没有增加。这一结果通过所有实验的训练日志分析得到了验证。
FLOPs
虽然 DenseMixer 需要在不活跃的专家上额外计算，但时间消耗不会随着活跃专家的数量线性增长。它只需要额外一次前向传递，并且不需要对这些不活跃的专家进行梯度计算和参数更新。
例如，使用 Qwen3-30B-A3B 时，DenseMixer 的 FLOPs 约为传统训练的 1.46 倍（根据此脚本估算）。这代表了理论上的上限，假设没有其他额外开销。实际上，由于通信等额外时间成本，实际减速小于 1.46 倍。此外，在少量数据上进行后训练时，开销变得可以忽略不计。
代码运行结果的详细统计
Bash
拷贝
Model Training Cost Analysis -- Conventional Training for Qwen3-30B-A3B
Number of parameters: 30,431,444,992
Number of Forward TFLOPs per layer: 16.85
Number of Backward TFLOPs per layer: 33.70
Number of TFLOPs per layer: 50.54
Peak memory cost: 157.93 GBs
Model Training Cost Analysis -- DenseMixer Training for Qwen3-30B-A3B
Number of parameters: 30,431,444,992
Number of Forward TFLOPs per layer: 40.04 # densemixer increases this
Number of Backward TFLOPs per layer: 33.70 # densemixer barely changes this
Number of TFLOPs per layer: 73.74
Peak memory cost: 164.96 GBs
FLOPs: DenseMixer / Conventional = 1.46x
时间消耗
我们比较了传统方法和 DenseMixer 在使用 deepspeed ZeRO3 并行处理的完整微调设置下的时间消耗。
Note
DenseMixer 在 Qwen3-30B-A3B 中引入了更高的开销，比小规模模型更高。但这不代表大规模模型必然放大计算成本。如 Handling More Complicated Cases 部分所述，部分开销来自处理归一化 Top-K 的实现，涉及附加操作，如 detach，这会给 DenseMixer 带来额外开销。
我们相信更高效的实现——例如，使用自定义的 PyTorch autograd 函数 —— 可以进一步降低成本，使 DenseMixer 能够应用于现代 MoEs，同时保持一致的改进和较小的额外开销。
结论与未来研究方向
我们确认，Top-K 路由的不可微分性是实现更有效 MoE 训练的关键障碍。通过在不活跃的专家上进行额外的前向传递，DenseMixer 实现了更精确的路由梯度，显著提升了 MoE 后训练质量，具体表现为下游性能超越了传统 MoE 训练方法。
未来，我们计划将监督微调（SFT）扩展到强化学习（RL），并在参数超过 100B 的更大规模的 MoE 模型上评估 DenseMixer，如 Qwen3-235B-A22B、Llama-4 (109B)和 DeepSeekV3 (675B)。
Citation
LaTeX
拷贝
@misc{yao2025densemixer,
title = {DenseMixer: Solving MoE Post-Training with Precise Router Gradients},
url = {https://fengyao.notion.site/moe-posttraining},
author = {Yao, Feng and Cui, Junxia and Zhang, Ruohan and Liu, Liyuan and Hao, Shibo and Zhang, Li and Dong, Chengyu and Wang, Shuohang and Shen, Yelong and Gao, Jianfeng and Shang, Jingbo},
journal = {Feng Yao's Notion},
year = {2025},
month = jun
}
这项工作目前处于预览阶段。我们很高兴分享早期成果，并欢迎社区的反馈。我们将继续优化和扩展 DenseMixer 的功能。如果您有任何问题或反馈，请随时联系我们：fengyao@ucsd.edu
