dLLM

简单扩散语言建模

📃 报告 | 🤗 模型

概述

dLLM 是一个统一 扩散语言模型 训练与评估的库，为整个开发流程带来透明性和可复现性：

• dLLM 提供基于 transformers Trainer 的可扩展训练流水线，支持 LoRA、DeepSpeed、FSDP 等。

• dLLM 基于 lm-evaluation-harness 提供统一的评估流水线，抽象了推理细节，使自定义变得简单。

• 在这些组件的基础上，dLLM 提供针对开放权重模型（如 LLaDA 和 Dream）的最小化 训练/推理/评估 流程，并实现了多种训练算法（如 MDLM（掩码扩散）、BD3LM（块扩散）、Edit Flows 等）。

新闻

[2026/02] ⚡Fast-dLLM：我们通过 Fast-dLLM（缓存、置信度阈值解码等）支持加速推理和评估 LLaDA 和 Dream。请参阅 examples/fastdllm 获取推理/评估说明。

[2025/12] 🤗Tiny-A2D：我们发布了一系列基于 AR 模型改造而来的 SOTA 小型（0.5B/0.6B）扩散模型，并提供了将 任意 AR 模型（如 Qwen、LLaMA、GPT-2 等）转换为扩散模型的完整开源流程。请参阅 examples/a2d 获取训练/推理/评估说明。

[2025/11] 🤗BERT-Chat：我们发布了一系列微调后可用于对话的 BERT 模型，提供了将 任意 BERT 编码器（如 BERT、RoBERTa、ModernBERT 等）转化为扩散模型的开源流程。请参阅 examples/bert 获取训练/推理/评估说明。

目录

• 功能
• 设置
• 文件
• 训练
• 推理
• 评估
• 引用

功能

• examples/llada：预训练、微调和评估 LLaDA / LLaDA-MoE。
• examples/dream：预训练、微调和评估 Dream。
• examples/a2d：将任意自回归模型微调为使用 掩码扩散 / 块扩散 生成文本。
• examples/bert：将任意 BERT 微调为轻量级聊天机器人。
• examples/editflow：用于训练 Edit Flows 模型的教育参考，展示了如何在现有 DLLM（如 LLaDA、Dream、BERT-Chat）基础上添加 编辑操作——插入、删除和替换——以及如何从头开始在公开数据上预训练或微调 Edit Flow 模型。
• examples/fastdllm：使用 Fast-dLLM（缓存、置信度阈值解码等）对 LLaDA 和 Dream 进行推理和评估。
• 更多内容即将推出。

设置

安装

# 创建并激活 conda 环境
conda create -n dllm python=3.10 -y
conda activate dllm

# 安装带有 CUDA 12.4 的 PyTorch（其他版本的 PyTorch/CUDA 也应适用）
conda install cuda=12.4 -c nvidia
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
    --index-url https://download.pytorch.org/whl/cu124

# 安装 dllm 包
pip install -e .

（可选）评估设置

# 初始化 `lm-evaluation-harness` 子模块
git submodule update --init --recursive

# 以可编辑模式安装子模块，并包含 IFEval 和 Math 依赖项
pip install -e "lm-evaluation-harness[ifeval,math]"

（可选）Slurm 设置

对于 Slurm 用户，请根据您的集群更新 scripts/train.slurm.sh：

- #SBATCH --partition=mllm_safety # 注意：请根据您的集群调整此项
- #SBATCH --quotatype=spot        # 注意：请根据您的集群调整此项
+ #SBATCH --partition=YOUR_PARTITION
+ #SBATCH --quotatype=YOUR_QUOTATYPE

接下来，为您的作业日志创建一个目录：

mkdir .logs

该文件夹将存储您通过 sbatch 提交的作业生成的日志文件。

文件

# 用于训练/采样的模块
dllm
├── core                   # 核心可重用模块，供 `dllm/pipelines` 共享
│   ├── samplers
│   ├── schedulers
│   └── trainers
├── data
├── pipelines              # 针对特定应用的训练与推理流水线
|   ├── bert
│   ├── dream
│   ├── editflow
│   ├── fastdllm
│   └── llada
│       ├── models         # 模型架构与配置
│       ├── sampler.py     # 推理模块
│       ├── trainer.py     # 训练模块
│       └── eval.py        # 评估模块
├── tools
└── utils

# 训练/采样的入口点
examples
├── bert
├── dream
├── editflow
├── fastdllm
└── llada
    ├── chat.py            # 交互式推理示例
    ├── sample.py          # 推理示例
    ├── pt.py              # 预训练示例
    ├── README.md          # 文档（您当前所在）
    ├── sft.py             # 监督微调示例
    └── eval.sh            # 评估脚本

训练

典型的训练入口脚本（例如 examples/llada/sft.py）如下所示：

import transformers

import dllm

model_args, data_args, training_args = parser.parse_args_into_dataclasses()
# ----- 模型 ------------------------------------------------------------------
模型 = dllm.utils.get_model(model_args=model_args)
# ----- 分词器 --------------------------------------------------------------
分词器 = dllm.utils.get_tokenizer(model_args=model_args)
# ----- 数据集 ----------------------------------------------------------------
数据集 = "..."

# ----- 训练 --------------------------------------------------------------
trainer = dllm.core.trainers.MDLMTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    args=training_args,
    data_collator=transformers.DataCollatorForSeq2Seq(
        tokenizer,
        return_tensors="pt",
        padding=True,
        label_pad_token_id=tokenizer.pad_token_id, 
    ),
)
trainer.train()

您可以在本地使用 accelerate 启动训练任务，或通过 sbatch 将其提交到 Slurm 集群。

# 在本地运行（8 张 GPU 上启用 ZeRO-2、4bit 量化和 LoRA）
accelerate launch \
    --config_file scripts/accelerate_configs/zero2.yaml \
    examples/llada/sft.py \
    --num_train_epochs 4 \
    --load_in_4bit True --lora True

# 提交到 Slurm 集群（单节点 8 张 GPU 上启用 FSDP）
sbatch --gres=gpu:8 scripts/train.slurm.sh \
    --accelerate_config "fsdp" \
    --script_path "examples/llada/sft.py" \
    --num_train_epochs 4

# 提交到 Slurm 集群（双节点 16 张 GPU 上启用 FSDP）
sbatch --nodes=2 --gres=gpu:8 scripts/train.slurm.sh \
    --accelerate_config "fsdp" \
    --script_path "examples/llada/sft.py" \
    --num_train_epochs 4

具体训练方案请参阅 Features。

有用的训练技巧：

• 使用数据子集： --dataset_args "allenai/tulu-3-sft-mixture[train:10000,test:1000]"
• 拼接多个数据集： --dataset_args "allenai/tulu-3-sft-mixture+HuggingFaceTB/smoltalk"
• 结合 LoRA 和 4bit 量化进行训练： --load_in_4bit True --lora True
• 使用不同的分布式训练方法： --accelerate_config "ddp,zero-{1,2,3},fsdp"
• 以流式方式加载预训练数据集： --streaming True
• 在训练前预处理 SFT 数据集（例如 LLaDA）：

  # 预处理 SFT 数据
  + python dllm/tools/preprocess_sft_dataset.py \
  +     --model_name_or_path "GSAI-ML/LLaDA-8B-Base" \
  +     --sft_map_fn_path "dllm.utils.default_sft_map_fn" \
  +     --dataset_args "allenai/tulu-3-sft-mixture" \
  +     --output_dir ".data/sft/llada/tulu-3-sft-mixture" \
  +     --num_proc 64
  
  # 使用预处理后的数据进行 SFT
  accelerate launch \
      --config_file scripts/accelerate_configs/fsdp.yaml \
      examples/llada/sft.py \
      --model_name_or_path "GSAI-ML/LLaDA-8B-Base" \
  -   --dataset_args "allenai/tulu-3-sft-mixture" \
  +   --dataset_args ".data/sft/llada/tulu-3-sft-mixture" \
  +   --load_preprocessed_data True \
      ...
  

推理

我们提供了统一的 samplers，用于抽象化推理细节。 典型的推理入口脚本（例如 examples/llada/sample.py）如下所示：

import dllm

模型 = dllm.utils.get_model(model_args=script_args).eval()
分词器 = dllm.utils.get_tokenizer(model_args=script_args)
采样器 = dllm.core.samplers.MDLMSampler(model=model, tokenizer=tokenizer)

messages = [
    [{"role": "user", "content": "莉莉以每小时 12 公里的速度跑了 4 小时。那么 8 小时能跑多远？"}],
    [{"role": "user", "content": "请编写一个教育性的 Python 函数。"}],
]

inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
)

outputs = sampler.sample(inputs, return_dict=True)
sequences = dllm.utils.sample_trim(tokenizer, outputs.sequences.tolist(), inputs)

您还可以尝试交互式聊天脚本（例如 examples/llada/chat.py），以可视化多轮对话：

python -u examples/llada/chat.py --model_name_or_path "GSAI-ML/LLaDA-8B-Instruct"

借助 Fast-dLLM，您可以加速 LLaDA 和 Dream 的推理。

python examples/fastdllm/llada/sample.py --model_name_or_path "GSAI-ML/LLaDA-8B-Instruct" --use_cache prefix --threshold 0.9

评估

在运行评估之前，请阅读(可选) 评估设置。

例如，要在 MMLU_Pro 数据集上使用 4 张 GPU 评估 LLaDA-8B-Instruct，请运行：

accelerate launch --num_processes 4 \
    dllm/pipelines/llada/eval.py \
    --tasks "mmlu_pro" \
    --model "llada" \
    --apply_chat_template \
    --num_fewshot 0 \
    --model_args "pretrained=GSAI-ML/LLaDA-8B-Instruct,is_check_greedy=False,mc_num=1,max_new_tokens=256,steps=256,block_size=256,cfg_scale=0.0"

我们还提供了脚本，用于在所有基准测试上自动评估 LLaDA、Dream 和 BERT-Chat。例如，您可以直接使用以下命令运行 examples/llada/eval.sh：

bash examples/llada/eval.sh --model_name_or_path "GSAI-ML/LLaDA-8B-Instruct" --instruct True
bash examples/llada/eval.sh --model_name_or_path "GSAI-ML/LLaDA-8B-Base" --instruct False

我们还提供了使用 Fast-dLLM 评估 LLaDA 和 Dream 的脚本：

bash examples/fastdllm/llada/eval.sh --model_name_or_path "GSAI-ML/LLaDA-8B-Instruct" --instruct True --num_gpu 1
bash examples/fastdllm/dream/eval.sh --model_name_or_path "Dream-org/Dream-v0-Base-7B" --instruct False --num_gpu 1

引用

@misc{zhou2026dllm,
      title={dLLM: 简单的扩散语言建模}, 
      author={Zhanhui Zhou 和 Lingjie Chen 和 Hanghang Tong 和 Dawn Song},
      year={2026},
      eprint={2602.22661},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2602.22661}, 
}

dLLM 快速上手指南

dLLM 是一个用于统一训练和评估扩散语言模型（Diffusion Language Models）的开源库。它基于 Hugging Face transformers 和 lm-evaluation-harness，提供了可扩展的训练流水线（支持 LoRA、DeepSpeed、FSDP）和统一的评估方案，旨在让扩散模型的开发更加透明和可复现。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐)
• Python 版本: 3.10
• GPU: 支持 CUDA 的 NVIDIA 显卡（示例基于 CUDA 12.4，其他版本通常也兼容）
• 依赖管理: 推荐使用 conda 进行环境隔离

安装步骤

1. 创建并激活 Conda 环境

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

2. 安装 PyTorch (CUDA 12.4)

根据您的硬件环境安装对应的 PyTorch 版本。以下是官方推荐的 CUDA 12.4 安装命令：

conda install cuda=12.4 -c nvidia
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
    --index-url https://download.pytorch.org/whl/cu124

提示：如果您在中国大陆，建议使用国内镜像源加速下载。例如使用清华源：

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
    --index-url https://pypi.tuna.tsinghua.edu.cn/simple

(注：PyTorch 的 CUDA 版本包通常仍需从官方源或特定渠道获取，上述命令主要加速纯 Python 依赖)

3. 安装 dLLM 库

克隆仓库后，以可编辑模式安装：

pip install -e .

4. (可选) 配置评估工具

如果需要运行基准测试评估，需初始化子模块并安装 lm-evaluation-harness：

git submodule update --init --recursive
pip install -e "lm-evaluation-harness[ifeval,math]"

基本使用

dLLM 提供了针对不同模型架构（如 LLaDA, Dream, BERT-Chat 等）的示例脚本。以下以 LLaDA 模型的监督微调（SFT）为例，展示最基础的训练流程。

1. 准备训练脚本

典型的训练入口脚本位于 examples/llada/sft.py。其核心逻辑如下（无需手动修改，了解即可）：

import transformers
import dllm

# 解析参数
model_args, data_args, training_args = parser.parse_args_into_dataclasses()

# 加载模型和分词器
model = dllm.utils.get_model(model_args=model_args)
tokenizer = dllm.utils.get_tokenizer(model_args=model_args)

# 构建数据集 (此处省略具体加载代码)
dataset = "..."

# 初始化 Trainer (以 MDLMTrainer 为例)
trainer = dllm.core.trainers.MDLMTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    args=training_args,
    data_collator=transformers.DataCollatorForSeq2Seq(
        tokenizer,
        return_tensors="pt",
        padding=True,
        label_pad_token_id=tokenizer.pad_token_id, 
    ),
)

# 开始训练
trainer.train()

2. 启动训练

您可以使用 accelerate 在本地单机多卡启动，或通过 sbatch 提交到 Slurm 集群。

方式 A：本地运行 (推荐初学者)

以下命令演示如何在 8 张 GPU 上使用 ZeRO-2 策略、4bit 量化和 LoRA 进行微调：

accelerate launch \
    --config_file scripts/accelerate_configs/zero2.yaml \
    examples/llada/sft.py \
    --num_train_epochs 4 \
    --load_in_4bit True --lora True

常用参数技巧：

• 使用数据子集测试: --dataset_args "allenai/tulu-3-sft-mixture[train:10000,test:1000]"
• 合并数据集: --dataset_args "dataset1+dataset2"
• 流式加载大数据集: --streaming True

方式 B：提交到 Slurm 集群

如果您拥有集群权限，可以修改 scripts/train.slurm.sh 中的分区配置后提交任务：

# 单节点 8 卡 (FSDP)
sbatch --gres=gpu:8 scripts/train.slurm.sh \
    --accelerate_config "fsdp" \
    --script_path "examples/llada/sft.py" \
    --num_train_epochs 4

# 多节点 16 卡 (FSDP)
sbatch --nodes=2 --gres=gpu:8 scripts/train.slurm.sh \
    --accelerate_config "fsdp" \
    --script_path "examples/llada/sft.py" \
    --num_train_epochs 4

3. 推理与采样

训练完成后，可以使用提供的采样器进行文本生成。以 LLaDA 为例：

python examples/llada/sample.py \
    --model_name_or_path <your_model_path> \
    --prompt "Once upon a time" \
    --max_new_tokens 100

对于交互式对话体验（如 BERT-Chat），可使用：

python examples/bert/chat.py --model_name_or_path <your_chat_model_path>

更多针对特定模型（如 Dream, EditFlow, Fast-dLLM）的详细用法，请参阅 examples/ 目录下对应文件夹的 README.md。