🎨 DenseMixer 🎨

Improving MoE Post-Training with Precise Router Gradients (Blog)

What is DenseMixer? • Key Features • Experiments • Quick Start • Efficiency • Citation

DenseMixer is a novel MoE post-training technique that empowers MoE training with more precise router gradient estimation, consistently outperforming conventional MoE training in downstream tasks.

What is DenseMixer?

DenseMixer addresses the non-differentiable Top-K routing problem in MoE training via straight-through estimator (STE). This enables more precise router gradients by computing all experts's output during forward pass for better gradient estimation during backward pass. For more technical details, please refer to our blog.

🚀 Key Features

• Plug-and-play: Zero code changes required
• Universal compatibility: Works with any MoE using Top-K routing
• Performance gains: Consistently outperforms conventional MoE training
• Parameter-efficient: Compatible with LoRA and other PEFT methods
• No inference overhead: Zero impact on model inference speed

📈 Experiments

DenseMixer consistently outperforms conventional MoE training across:

• Model scales: 7B, 14B, 30B parameters
• Architectures: With/without shared experts
• Training methods: From scratch and up-cycling
• Data types: Instruction tuning and long reasoning data

Reproducible Experiments: For detailed training scripts, configurations, and evaluation code, please refer to the experiments folder.

⚡ Quick Start

1. Installation

pip install densemixer

2. Setup (One-time)

densemixer setup

3. Enable DenseMixer

export DENSEMIXER_ENABLED=1

4. Use Your MoE Models

from transformers import Qwen3MoeForCausalLM

# DenseMixer automatically patches the model
model = Qwen3MoeForCausalLM.from_pretrained("Qwen/Qwen3-MoE-30B-A3B")

# Train as usual - no code changes needed!

🔧 Configuration

DenseMixer currently supports the following models.

• Qwen3-MoE (30B parameters)
• Qwen1.5-MoE (14B parameters)
• OLMoE (7B parameters)

DenseMixer uses environment variables for configuration:

Usage Examples

Enable for all models:

export DENSEMIXER_ENABLED=1
python your_training_script.py

Enable only for specific models:

export DENSEMIXER_ENABLED=1
export DENSEMIXER_QWEN3=1
export DENSEMIXER_QWEN2=0
export DENSEMIXER_OLMOE=0
python your_training_script.py

Disable (default behavior):

# No environment variables needed
python your_training_script.py

📊 Logging

DenseMixer provides intelligent logging to track when custom forward methods are used:

INFO - densemixer - DenseMixer: Using custom forward method for Qwen3-MoE
INFO - densemixer - DenseMixer: Using custom forward method for OLMoE

You can also customize the logging as below.

import logging

# Set logging level
logging.getLogger("densemixer").setLevel(logging.INFO)

# Or disable logging entirely
logging.getLogger("densemixer").setLevel(logging.WARNING)

⚡ Efficiency Analysis

FLOPs: 1.46x overhead vs conventional training (theoretical analysis on Qwen3-30B-A3B)

Model Training Cost Analysis Results --- 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 Results --- DenseMixer Training for Qwen3-30B-A3B ---
Number of parameters: 30,431,444,992
Number of Forward TFLOPs per layer: 40.04
Number of Backward TFLOPs per layer: 33.70 # we assume DenseMixer doesn't change backward significantly
Number of TFLOPs per layer: 73.74
Peak memory cost: 164.96 GBs

FLOPs: DenseMixer / Conventional = 1.46x

Memory: Negligible overhead - model weights are already loaded on GPU

Time: Negligible when training with small scale of data

Detailed FLOPs analysis available in efficiency_analysis/flops_compute.py

🚧 Roadmap & Future Improvements

While DenseMixer already delivers significant performance gains, we're working on several improvements to make it more efficient and production-ready:

• Optimize backward FLOPs: Currently backward FLOPs increase alongside forward FLOPs, but this is unnecessary. We'll optimize the backward pass to reduce training overhead while maintaining performance gains
• Integrate modern MoE kernels: Replace transformers' native for-loop MoE implementation with efficient kernels (e.g., grouped_gemm) for better large-scale training performance
• Industry integration: Extend support beyond open-instruct/llama-factory to MoE-optimized frameworks like Megatron for seamless production adoption

📚 Citation

If you find our work useful, please cite us:

@misc{yao2025densemixer,
  title = {DenseMixer: Improving MoE Post-Training with Precise Router Gradient},
  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
}

Questions?

If you have any questions related to the code or the blog, feel free to reach out to us at fengyao@ucsd.edu.