File size: 3,753 Bytes
13c5606 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | # TAD: Temporal-Aware Trajectory Self-Distillation for Fast and Accurate Diffusion LLM
[](https://arxiv.org/abs/2605.09536)
[](https://huggingface.co/collections/Zhouhhy/tad)
<div align="center">
</div>
## 1. Environment Setup
```bash
conda create -n tad python=3.10 -y
conda activate tad
pip install -r requirements.txt
```
## 2. Data Preparation
### 2.1 Generate Math Trajectories
**For LLaDA:**
```bash
python prepare/llada_gen_math_traj.py \
--model_name <your_llada_model_path> \
--output_path data/llada_math_traj.jsonl \
--dataset_name gsm8k \
--dataset_split train \
--max_new_tokens 256 \
--steps 256 \
--block_length 32 \
--num_samples 1 \
--limit 0
```
**For Dream:**
```bash
python prepare/dream_gen_math_traj.py \
--model_name <your_dream_model_path> \
--output_path data/dream_math_traj.jsonl \
--dataset_name gsm8k \
--dataset_split train \
--max_new_tokens 256 \
--block_length 32 \
--top_p 0.95 \
--alg entropy \
--num_samples 1 \
--limit 0
```
### 2.2 Generate Code Trajectories
**For LLaDA:**
```bash
python prepare/llada_gen_code_traj.py \
--model_name <your_llada_model_path> \
--parquet_path <path_to_kodcode_parquet> \
--output_path data/llada_code_traj.jsonl \
--max_new_tokens 256 \
--steps 256 \
--block_length 32 \
--num_samples 1 \
```
**For Dream:**
```bash
python prepare/dream_gen_code_traj.py \
--model_name <your_dream_model_path> \
--parquet_path <path_to_kodcode_parquet> \
--output_path data/dream_code_traj.jsonl \
--max_new_tokens 256 \
--block_length 32 \
--top_p 0.95 \
--alg entropy \
--num_samples 1 \
```
## 3. Training
### 3.1 Configure
Edit the config file to set your model path and training hyperparameters:
- **LLaDA**: `LLaDA/configs/config_llada.yaml`
- **Dream**: `Dream/configs/config_dream.yaml`
### 3.2 Launch Training
**Train LLaDA:**
```bash
cd LLaDA
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 deepspeed train_llada.py
```
**Train Dream:**
```bash
cd Dream
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 deepspeed train_dream.py
```
## 4. Evaluation
### 4.1 Merge LoRA Weights
After training, merge the LoRA adapter into the base model. Edit the paths in `merge_lora.py` and run:
**For LLaDA:**
```bash
cd LLaDA
# Edit merge_lora.py:
# - name = "<your_base_model_path>"
# - PeftModel.from_pretrained(base_model, "<your_checkpoint_path>")
# - merged_model.save_pretrained("<save_path>")
python merge_lora.py
```
### 4.2 Run Evaluation
Edit the `model_path` variable in the evaluation script to point to your merged model, then run:
**Evaluate LLaDA:**
```bash
cd eval
# Edit eval_llada.sh: set model_path="<your_merged_model_path>"
bash eval_llada.sh
```
The evaluation scripts cover the following benchmarks:
- **Math**: GSM8K, MATH (Minerva)
- **Code**: HumanEval, MBPP
Results will be saved under `evals_results/`.
## Acknowledgments
This project builds upon excellent open-source work:
- [LLaDA](https://arxiv.org/abs/2502.09992) - Large Language Diffusion Models
- [Dream](https://arxiv.org/abs/2508.15487) - Diffusion Large Language Models
- [Fast-dLLM](https://arxiv.org/abs/2505.22618) - Training-free acceleration
- [D2F](https://arxiv.org/abs/2508.09192) - Discrete diffusion forcing
- [dParallel](https://arxiv.org/abs/2509.26488) - Distilled dLLMs
- [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) - Evaluation framework
- [d3LLM](https://github.com/hao-ai-lab/d3LLM) - Ultra-Fast Diffusion LLM using Pseudo-Trajectory Distillation
|