KO | EN

GitLite

Minimalist Open Source Repository Search Engine

gitlite — search

#typescript #ai-agents #developer-tools #cli #codex #claude-code #react #dsh-plugin #mcp #open-source #dsh #deepseek-harness

Router-Tuning-Mixture-of-Depths

★ 31 Open GitHub ↗

The open-source Mixture of Depths code and the official implementation of the paper "Router-Tuning: A Simple and Effective Approach for Enabling Dynamic Depth in Transformers. (EMNLP 2025)"

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1 align="center">[EMNLP 2025] Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers</h1>

<a href="https://shwai-he.github.io/">Shwai He</a>, <a href="https://getao.github.io/">Tao Ge</a>, <a href="https://s1gh.alphaxiv.io/">Guoheng Sun</a>, <a href="https://bowei.netlify.app/#about">Bowei Tian</a>, <a href="https://xyang0.github.io/">Xiaoyang Wang</a>, <a href="https://sites.google.com/view/dongyu888/">Dong Yu</a>

<a href="https://case-lab-umd.github.io/Router-Tuning-Mixture-of-Depths/">🏠 Homepage</a> •
 <a href="#-introduction">📖 Introduction</a> •
 <a href="#-news">📰 News</a> •
 <a href="#-why-this-repo">✨ Why</a> •
 <a href="#-results">📈 Results</a> •
 <a href="#-quick-start">🚀 Quick Start</a> •
 <a href="#-citation">📄 Citation</a>

## 📖 Introduction

This is the official implementation of the paper [**Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers**](https://arxiv.org/abs/2410.13184), accepted at **EMNLP 2025**. We provide a practical framework for efficient dynamic-depth training and inference in Transformers.

Router-Tuning enables dynamic-depth inference by fine-tuning only router-related parameters. Compared with standard full-model dynamic-depth tuning, it significantly reduces training cost while keeping model quality competitive.

## 📰 News
- Aug 2025: Router-Tuning accepted to **EMNLP 2025** main conference.
- Oct 2024: arXiv preprint and code release.

## ✨ Why This Repo
Traditional transformers execute a fixed number of layers for every token, which wastes computation on easy tokens.

[Mixture of Depths (MoD)](https://arxiv.org/abs/2404.02258) addresses this by dynamically skipping less important computations, but two practical issues remain:

1. Existing methods usually tune the whole model, causing high training cost.
2. Aggressive skipping can hurt quality if routing is not well calibrated.

Router-Tuning tackles both by focusing optimization on routing components and introducing routing strategies that better preserve performance-efficiency tradeoffs.

## 📈 Results

### 🏁 Main Results

 <img src="docs/figures/main_results.png" alt="Main benchmark results of Router-Tuning" width="92%">

Router-Tuning consistently improves the efficiency-quality tradeoff over full-parameter dynamic-depth tuning baselines.
The reported best setting reaches notable speedup while keeping quality degradation small.

### 🔬 Expert Routing Analysis

 <img src="docs/figures/expert_rt.png" alt="Expert-level routing behavior analysis" width="82%">

Router specialization becomes clearer after tuning: the model learns more stable token-to-layer routing patterns.
This supports dynamic-depth execution with lower unnecessary computation.

### 🔗 LoRA Compatibility

 <img src="docs/figures/lora_rt.png" alt="LoRA and Router-Tuning combined results" width="92%">

Router-Tuning is compatible with LoRA-based adaptation and can be composed for a better efficiency-performance balance.
In practice, this enables lightweight deployment recipes without full-model retraining.

## 🔍 Core Methods
1. **Router-Only Fine-Tuning**
- Tune router-related parameters instead of full-model updates.
- Strongly reduces optimization cost for dynamic-depth adaptation.

2. **Attention-Based Dynamic Depth**
- Uses attention-based routing granularity to improve compute and memory efficiency.
- Preserves output quality under dynamic-depth execution.

## 📦 Repository Layout
- `entrypoints/finetune/finetune_router_tuning.py`: main training entrypoint.
- `scripts/finetune_router_tuning.sh`: reproducible launcher with `accelerate` + DeepSpeed.
- `entrypoints/data/reformat_datasets.py`: convert raw datasets to unified `messages` format.
- `entrypoints/data/mix_datasets.py`: build mixed instruction-tuning data.
- `utils/pipeline/customized_trainer.py`: router-focused trainer logic.
- `configs/accelerate/`: distributed training launcher configs.
- `configs/deepspeed/`: DeepSpeed runtime configs.

## ⚙️ Installation
```bash
conda create -n router-tuning python=3.10 -y
conda activate router-tuning

git clone https://github.com/CASE-Lab-UMD/Router-Tuning-Mixture-of-Depths.git
cd Router-Tuning-Mixture-of-Depths

pip install -r requirements.txt
```

`flash-attn` is optional and is intentionally installed in a second step. Its
build process imports PyTorch, so installing it in the same isolated pip
transaction as PyTorch can fail while pip is "getting requirements to build
wheel".

```bash
# Optional: install only after the base environment, including torch, is ready.
pip install --no-build-isolation flash-attn==2.6.3
```

The dependency snapshot used by the original repository records:

- Python 3.10
- PyTorch 2.1.1
- Transformers 4.40.1
- FlashAttention 2.6.3

The original CUDA toolkit/driver minor version was not recorded in this
repository. Choose the PyTorch build matching the CUDA driver on the target
machine, verify `python -c "import torch; print(torch.cuda.is_available(),
torch.version.cuda)"`, and only then install FlashAttention.

For this dependency set, keep `numpy<2` and `setuptools<70`. The optional
FlashAttention installation also needs a discoverable CUDA toolkit and `nvcc`
during package metadata/build setup.

FlashAttention is not required for correctness. The launcher uses the model's
default Transformers attention implementation unless explicitly enabled:

```bash
USE_FLASH_ATTN=True bash scripts/finetune_router_tuning.sh
```

If FlashAttention is requested but unavailable, the training entrypoint logs a
warning and falls back to the model's default attention implementation.

## 🚀 Quick Start
### 1) 🧹 Prepare Data
Put raw datasets under `data/raw/` using the expected subdirectory names:
- `vicuna_sharegpt`
- `evol_instruct`
- `slim_orca`
- `meta_math_qa`
- `evol_code_alpaca`
- `alpaca`

Then run:

```bash
python entrypoints/data/reformat_datasets.py \
  --raw_data_root ./data/raw \
  --save_path ./data/reformatted

python entrypoints/data/mix_datasets.py \
  --reformatted_dir ./data/reformatted \
  --save_path ./data/mixed
```

### 2) 🏃 Run Router-Tuning
```bash
bash scripts/finetune_router_tuning.sh
```

The training entrypoint is now Hugging Face first:
- Preferred: pass a Hugging Face model ID such as `Qwen/Qwen2.5-7B`.
- Also supported: a normal local model directory.

### 3) 🖥️ Minimal Single-Node Override Example
```bash
NUM_PROCESSES=4 PORT=29501 bash scripts/finetune_router_tuning.sh
```

### 4) ☁️ Hugging Face Example
```bash
MODEL_NAME_OR_PATH="Qwen/Qwen2.5-7B" \
RUN_NAME="qwen2.5-7b" \
NUM_PROCESSES=4 PORT=29501 \
bash scripts/finetune_router_tuning.sh
```

## 🎛️ Training Knobs
`finetune_router_tuning.sh` is the recommended launcher. Commonly adjusted fields:

- `model_name_or_path`: preferred Hugging Face model ID, or a normal local model path.
- `run_name`: optional output folder alias; defaults to `basename(model_name_or_path)`.
- `trust_remote_code`: defaults to `False`; enable only for architectures that require it.
- `data_type`: one dataset under `data/reformatted/` or `mixed`.
- `max_train_samples`: training subset size for quick experiments.
- `router_layers`: number of deep layers enabled for router tuning in the current implementation.
- `target_capacity`: optional target activation ratio for router regularization.
- `granularity`: routing granularity (`attn_sequence`, `attn_token`, `mlp_sequence`, `mlp_token`, `block_sequence`, `block_token`).
- `router_only`: enable router-only training (default `True`).
- `learning_rate`, `weight_decay`, `num_epochs`.

Distributed launch overrides:
- `NUM_PROCESSES`: number of GPU processes.
- `PORT`: distributed master port.
- `USE_FLASH_ATTN`: set to `True` to request optional FlashAttention 2; defaults
  to `False` and safely falls back when unavailable.

## 🧪 Evaluation
Evaluation is compatible with [EleutherAI/lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).
For strict reproduction used in earlier experiments, see [s1ghhh/lm-evaluation-harness](https://github.com/s1ghhh/lm-evaluation-harness).

## ✅ Repro Checklist
- Python 3.10 environment with `pip install -r requirements.txt`.
- Valid `model_name_or_path`: preferably a Hugging Face model ID or a normal local model path.
- Reformatted/mixed data exists at `data/reformatted/*/data.jsonl` or `data/mixed/data.jsonl`.
- `accelerate` config selected in `scripts/finetune_router_tuning.sh` matches your hardware.
- `NUM_PROCESSES` and GPU memory are consistent with `max_seq_length` and batch setup.

## 📄 Citation
```bibtex
@misc{he2024routertuningsimpleeffectiveapproach,
  title={Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers},
  author={Shwai He and Tao Ge and Guoheng Sun and Bowei Tian and Xiaoyang Wang and Dong Yu},
  year={2024},
  eprint={2410.13184},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2410.13184}
}
```

## 📬 Contact
- Shwai He: `shwaihe@umd.edu`

Router-Tuning-Mixture-of-Depths

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README