dMoE: dLLMs with Learnable Block Experts

![](https://img.shields.io/github/stars/fscdc/dmoe?color=yellow) ![](https://img.shields.io/github/forks/fscdc/dmoe?color=lightblue) ![](https://img.shields.io/github/last-commit/fscdc/dmoe?color=green) ![](https://img.shields.io/badge/PRs-Welcome-blue) arXiv HuggingFace
https://github.com/user-attachments/assets/fb2fb91d-5d25-4cb1-8d46-2f2fe4248f05 > **[dMoE: dLLMs with Learnable Block Experts](https://arxiv.org/abs/2605.30876)** \ > [Sicheng Feng](https://fscdc.github.io/)1, [Zigeng Chen](https://czg1225.github.io/chenzigeng99/)1, [Gongfan Fang](https://fangggf.github.io/)1, [Xinyin Ma](https://horseee.github.io/)1, [Xinchao Wang](https://sites.google.com/site/sitexinchaowang/)1,* \ > 1National University of Singapore, Singapore \ > ∗Corresponding author: xinchao@nus.edu.sg --- ## ⭐ Updates - **[4.26.2026]**: Paper, code and model are released. --- ## 💪 Highlights - **Learnable Block Experts**: Introduces block-level MoE routing into dLLMs, drastically compressing the number of activated unique experts across diffusion steps — directly targeting memory footprint reduction. - **Reduced MoE Bandwidth**: By constraining expert activation at the block level, dMoE significantly reduces memory bandwidth consumed by expert weight loading during the block diffusion process. - **Improved Efficiency-Accuracy Trade-off**: dMoE achieves competitive performance on reasoning and general benchmarks while reducing unnecessary computation through adaptive expert activation. - **Plug-and-play on LLaDA-2.0**: Built directly on top of LLaDA-2.0-mini without architectural changes, enabling straightforward extension to other masked dLLMs. --- ## 📚 Table of Contents - [💡 Introduction](#-introduction) - [💻 Model and Datasets](#-model-and-datasets) - [🚀 Quick Start](#-quick-start) - [🔧 Installation](#-installation) - [🔥 Training](#-training) - [⚡ Evaluation](#-evaluation) - [☀️ Acknowledgement](#-acknowledgement) - [📖 Citation](#-citation) --- ## 💡 Introduction We present **dMoE**, a framework that introduces Learnable Block Experts into diffusion large language models (dLLMs). ![Overview](assets/overview.png) --- ## 💻 Model and Datasets | Model | Description | Source Model | Link | |-------|-------------|--------------|------| | 🤖 dMoE-16B | General-purpose dLLM with learnable block experts | LLaDA-2.0-mini | [Hugging Face](https://huggingface.co/FSCCS/dMoE-16B) | --- ## 🚀 Quick Start ```bash git clone https://github.com/fscdc/dMoE.git cd dMoE # not include all but main libraries conda create -n dmoe python==3.12 conda activate dmoe pip install -r ./evaluations/requirements.txt ``` ```python import torch from transformers import AutoTokenizer from evaluations.models.modeling_llada2_moe_be_adaptive import LLaDA2MoeModelLM MODEL_NAME = "FSCCS/dMoE-16B" device = "cuda:0" model = LLaDA2MoeModelLM.from_pretrained( MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16 ).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) prompt = "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?" + "\nLet's think step by step\n" messages = [[{"role": "user", "content": prompt}]] input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) inputs = tokenizer(input_text, return_tensors="pt", padding_side="left") input_ids = inputs["input_ids"].to(device) with torch.no_grad(): out, unique_experts_count = model.generate( input_ids, steps=32, gen_length=2048, block_length=32, temperature=0.0, eos_early_stop=True, ) generated = out[:, input_ids.shape[1]:] result = tokenizer.batch_decode(generated, skip_special_tokens=True) print("Output:", result[0]) print("Unique experts count:", unique_experts_count) ``` --- ## 🔧 Installation Clone the dMoE repository: ```bash git clone https://github.com/fscdc/dmoe.git --recursive cd dmoe ``` Install the **training** environment (based on dFactory): ```bash cd training conda create -n dmoe-training python==3.12 conda activate dmoe-training pip install -e VeOmni/ ``` Install the **evaluation** environment: ```bash cd evaluations conda create -n dmoe python==3.12 conda activate dmoe pip install -r requirements.txt ``` --- ## 🔥 Training Our training pipeline is based on [dFactory](https://github.com/inclusionAI/dFactory). ```bash cd training ``` ### 1. Download and Merge Model Weights Download the base model and convert it to the merged-expert format required for training: ```bash # Download the original LLaDA-2.0-mini weights python scripts/download_hf_model.py \ --repo_id inclusionAI/LLaDA2.0-mini \ --local_dir /path/to/separate_expert_model # Convert to merged format for training python scripts/moe_convertor.py \ --input-path /path/to/separate_expert_model \ --output-path /path/to/merged4moe/LLaDA2.0-mini \ --mode merge ``` ### 2. Prepare Training Data ```bash # gsm8k as an example python scripts/build_gsm8k_dataset.py ``` ### 3. Modify Training Config Edit `configs/moe/llada2_mini_adaptive.yaml`: ```yaml model: model_path: "/path/to/merged4moe/LLaDA2.0-mini" tokenizer_path: "/path/to/merged4moe/LLaDA2.0-mini" data: train_path: "/your/data/path" train: output_dir: "/your/output/path" ``` ### 4. Run Training ```bash PYTHONPATH=$(pwd)/VeOmni:$PYTHONPATH sh train.sh tasks/train_llada2_bd.py configs/moe/llada2_mini_adaptive.yaml # Or bash scripts_moe/train_adaptive.sh ``` ### 5. Convert the Checkpoint After training, convert the checkpoint back to the standard MoE format: ```bash python scripts/moe_convertor.py \ --input-path ./logs/llada2_mini_dmoe/checkpoints/global_step_XXX/hf_ckpt/ \ --output-path /path/to/output/llada2_mini_dmoe \ --mode split ``` --- ## ⚡ Evaluation ```bash cd evaluations ``` ### Evaluate dMoE (Adaptive Block Experts) Set `--model-name` to your local model path in `scripts_moe/run_be_adaptive.sh`, then run: ```bash bash scripts_moe/run_be_adaptive.sh ``` This evaluates on four benchmarks: - ✅ GSM8K - ✅ MATH500 - ✅ MMLU - ✅ ARC-C ### Evaluate Baseline (Original LLaDA-2.0-mini) ```bash bash scripts_moe/run_origin.sh ``` --- ## ☀️ Acknowledgement We sincerely thank Huawei for their support and contribution to the research and development of this model and algorithm. Furthermore, our code builds on [dFactory](https://github.com/inclusionAI/dFactory). We acknowledge these great works for laying the groundwork that made our approach possible. --- ## 📖 Citation If our research assists your work, please give us a star ⭐ or cite us using: ```bibtex @article{feng2026dmoe, title={dMoE: dLLMs with Learnable Block Experts}, author={Feng, Sicheng and Chen, Zigeng and Fang, Gongfan and Ma, Xinyin and Wang, Xinchao}, journal={arXiv preprint arXiv:2605.30876}, year={2026} } ```