EDIT: Early Diffusion Inference Termination for dLLMs Based on Dynamics of Training Gradients

tags: diffusion language models, early termination, training metadata, reasoning benchmarks

This repo is the official implementation of "EDIT: Early Diffusion Inference Termination for dLLMs Based on Dynamics of Training Gradients."

TL;DR: EDIT uses training-time metadata to enable early termination during inference in diffusion language models, reducing cost while maintaining or improving accuracy.

---

🔥 Updates

• Accepted to NeurIPS OPT 2025 Workshop

🛠️ 1. Setup

The project is tested on Python 3.10 and supports two hardware platforms:

• NVIDIA CUDA GPUs
• Intel XPU

All installation commands should be run from the project’s root directory: EDIT/

1.1 Create Conda Environment

conda create -n EDIT python=3.10 -y
conda activate EDIT

1.2 Installation (Choose One Platform)

1.2.A | NVIDIA GPU Installation

• Command Line for Installing PyTorch 2.6 and Dependencies

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt

(Optional) Verify Installation

python setup_check.py


# Example output:

🔍 Checking Python environment...

✅ torch: 2.6.0+cu124
...
🖥️ Device Detection:
✅ CUDA GPU detected: NVIDIA A100-SXM4-80GB (CUDA 12.4)

1.2.B | Intel XPU Installation

• Command Line for Installing PyTorch 2.6 and Dependencies

python -m pip install \
    torch==2.6.0 \
    torchvision==0.21.0 \
    torchaudio==2.6.0 \
    --index-url https://download.pytorch.org/whl/xpu

python -m pip install \
    --proxy http://proxy-dmz.intel.com:912 \
    intel-extension-for-pytorch==2.6.10+xpu \
    oneccl_bind_pt==2.6.0+xpu \
    --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/

pip install -r requirements.txt

(Optional) Verify Installation

python setup_check.py


# Example output:

🔍 Checking Python environment...

✅ torch: 2.6.0+xpu

🖥️ Device Detection:
✅ Intel XPU detected
   XPU devices available: 12

🗺️ 1.3 Environment Summary

Library versions

• Framework
  • pytorch: 2.6.0 (CUDA 12.4 or XPU backend)
  • torchvision: 0.21.1
  • python: 3.10
• Hardware
  • GPU: NVIDIA A100-SXM4-80GB x 2
  • or Intel XPU x 12

---

The examples below use 2 GPUs

For SFT, please specify your number of GPUs with --nproc_per_node
For inference, please specify your available GPUs using --device_ids

🔥 2. Supervised Fine-tuning (SFT)

To run SFT on LLaDA, use the following command.

Reminder: Please ensure you are in the project root directory (EDIT/) and that your environment is active (conda activate EDIT) before running the commands below.

cd SFT
torchrun \
  --nproc_per_node 2 \
  --master_port 29410 \
  train_llada_sft.py \
  --model_name "GSAI-ML/LLaDA-8B-Instruct" \
  --output_dir logs/logs_llada_sft \
  --num_epochs 2 \
  --debugging

Argument notes

• --nproc_per_node : number of processes per node (e.g., 2 for two GPUs)
• --debugging: optional flag to disable wandb logging

🚀 3. Run Early Termination on Fine-tuned Model

Before running early termination to shorten inference steps, switch to the eval directory:

cd ../eval

(Optional) For NVIDIA GPU users, please run:

export CUBLAS_WORKSPACE_CONFIG=":4096:8"

Inference on Countdown Dataset

Countdown (Sequence 128)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset countdown \
  --output_path EDIT_results/countdown/seq128 \
  --config_file configs/EDIT_countdown_seq128.json \
  --gen_length 128 \
  --device_ids 0 1

Countdown (Sequence 256)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset countdown \
  --output_path EDIT_results/countdown/seq256 \
  --config_file configs/EDIT_countdown_seq256.json \
  --gen_length 256 \
  --device_ids 0 1

Countdown (Sequence 512)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset countdown \
  --output_path EDIT_results/countdown/seq512 \
  --config_file configs/EDIT_countdown_seq512.json \
  --gen_length 512 \
  --device_ids 0 1

Inference on Sudoku Dataset

Sudoku (Sequence 128)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset sudoku \
  --output_path EDIT_results/sudoku/seq128 \
  --config_file configs/EDIT_sudoku_seq128.json \
  --gen_length 128 \
  --device_ids 0 1

Sudoku (Sequence 256)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset sudoku \
  --output_path EDIT_results/sudoku/seq256 \
  --config_file configs/EDIT_sudoku_seq256.json \
  --gen_length 256 \
  --device_ids 0 1

Sudoku (Sequence 512)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset sudoku \
  --output_path EDIT_results/sudoku/seq512 \
  --config_file configs/EDIT_sudoku_seq512.json \
  --gen_length 512 \
  --device_ids 0 1

Inference on MATH500 Dataset

MATH500 (Sequence 128)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset math \
  --output_path EDIT_results/math/seq128 \
  --config_file configs/EDIT_math_seq128.json \
  --gen_length 128 \
  --device_ids 0 1

MATH500 (Sequence 256)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset math \
  --output_path EDIT_results/math/seq256 \
  --config_file configs/EDIT_math_seq256.json \
  --gen_length 256 \
  --device_ids 0 1

MATH500 (Sequence 512)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset math \
  --output_path EDIT_results/math/seq512 \
  --config_file configs/EDIT_math_seq512.json \
  --gen_length 512 \
  --device_ids 0 1

Inference on GSM8K Dataset

GSM8K (Sequence 128)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset gsm8k \
  --output_path EDIT_results/gsm8k/seq128 \
  --config_file configs/EDIT_gsm8k_seq128.json \
  --gen_length 128 \
  --device_ids 0 1

GSM8K (Sequence 256)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset gsm8k \
  --output_path EDIT_results/gsm8k/seq256 \
  --config_file configs/EDIT_gsm8k_seq256.json \
  --gen_length 256 \
  --device_ids 0 1

GSM8K (Sequence 512)

python dllm_eval_mp.py \
  --model_path "GSAI-ML/LLaDA-8B-Instruct" \
  --checkpoint_path ../SFT/logs/logs_llada_sft/llada-s1/checkpoint-200 \
  --dataset gsm8k \
  --output_path EDIT_results/gsm8k/seq512 \
  --config_file configs/EDIT_gsm8k_seq512.json \
  --gen_length 512 \
  --device_ids 0 1

📊 4. Perform Evaluation

Evaluate the Countdown Dataset

Countdown (Sequence 128)

python parse_and_get_acc.py -d EDIT_results/countdown/seq128/

Countdown (Sequence 256)

python parse_and_get_acc.py -d EDIT_results/countdown/seq256/

Countdown (Sequence 512)

python parse_and_get_acc.py -d EDIT_results/countdown/seq512/

Evaluate the Sudoku Dataset

Sudoku (Sequence 128)

python parse_and_get_acc.py -d EDIT_results/sudoku/seq128/

Sudoku (Sequence 256)

python parse_and_get_acc.py -d EDIT_results/sudoku/seq256/

Sudoku (Sequence 512)

python parse_and_get_acc.py -d EDIT_results/sudoku/seq512/

Evaluate the MATH500 Dataset

MATH500 (Sequence 128)

python parse_and_get_acc.py -d EDIT_results/math/seq128/

MATH500 (Sequence 256)

python parse_and_get_acc.py -d EDIT_results/math/seq256/

MATH500 (Sequence 512)

python parse_and_get_acc.py -d EDIT_results/math/seq512/

Evaluate the GSM8K Dataset

GSM8K (Sequence 128)

python parse_and_get_acc.py -d EDIT_results/gsm8k/seq128/

GSM8K (Sequence 256)

python parse_and_get_acc.py -d EDIT_results/gsm8k/seq256/

GSM8K (Sequence 512)

python parse_and_get_acc.py -d EDIT_results/gsm8k/seq512/

🔢 5. Compute Reduced Denoising Steps

Calculate on Countdown Dataset

Countdown (Sequence 128)

python compute_early_exit_diffusion_steps.py -d EDIT_results/countdown/seq128/

Countdown (Sequence 256)

python compute_early_exit_diffusion_steps.py -d EDIT_results/countdown/seq256/

Countdown (Sequence 512)

python compute_early_exit_diffusion_steps.py -d EDIT_results/countdown/seq512/

Calculate on Sudoku Dataset

Sudoku (Sequence 128)

python compute_early_exit_diffusion_steps.py -d EDIT_results/sudoku/seq128/

Sudoku (Sequence 256)

python compute_early_exit_diffusion_steps.py -d EDIT_results/sudoku/seq256/

Sudoku (Sequence 512)

python compute_early_exit_diffusion_steps.py -d EDIT_results/sudoku/seq512/

Calculate on MATH500 Dataset

MATH500 (Sequence 128)

python compute_early_exit_diffusion_steps.py -d EDIT_results/math/seq128/

MATH500 (Sequence 256)

python compute_early_exit_diffusion_steps.py -d EDIT_results/math/seq256/

MATH500 (Sequence 512)

python compute_early_exit_diffusion_steps.py -d EDIT_results/math/seq512/

Calculate on GSM8K Dataset

GSM8K (Sequence 128)

python compute_early_exit_diffusion_steps.py -d EDIT_results/gsm8k/seq128/

GSM8K (Sequence 256)

python compute_early_exit_diffusion_steps.py -d EDIT_results/gsm8k/seq256/

GSM8K (Sequence 512)

python compute_early_exit_diffusion_steps.py -d EDIT_results/gsm8k/seq512/

🔖 6. Citation

If you find this work helpful for your research, please cite the following paper:

@article{HsiehWK25,
  author       = {He{-}Yen Hsieh and
                  Hong Wang and
                  H. T. Kung},
  title        = {EDIT: Early Diffusion Inference Termination for dLLMs Based on Dynamics of Training Gradients},
  booktitle    = {Proceedings of the NeurIPS Workshop on Optimization for Machine Learning},
  year         = {2025}
}

🙌 7. Acknowledgements

Big thanks to the authors of LLaDA and d1 for their great work. Really appreciate it!