大語言模型的文本生成方式一直都是以自迴歸為主：一個token接一個token,從左往右,生成完就定了。
但現在有個不太一樣的思路開始在研究圈裏流行起來，那就是擴散語言模型(Diffusion LMs)。擴散模型在圖像生成領域已經證明了自己的可行性,但是問題是把這套東西用到文本上一直很麻煩——訓練難、評估難、更別提怎麼集成到現有的LLM工作流裏了。
dLLM是一個開源的Python庫,它把擴散語言模型的訓練、微調、推理、評估這一整套流程都統一了起來，而且號稱任何的自迴歸LLM都能通過dLLM轉成擴散模型
擴散模型用在語言上有什麼不同
做過圖像擴散模型的應該能理解這個思路。
傳統自迴歸是順序生成,擴散模型的玩法不一樣:先從噪聲或者masked tokens開始,然後一步步把整個序列細化出來。它不是一個token一個token往後走,而是對整個輸出做全局優化。
擴散模型在幾個場景下表現特別好:需要複雜推理的任務、文本編輯重寫、結構化生成,還有需要多輪迭代優化的場景。
dLLM提供了什麼
dLLM不是某個具體模型它是個框架，包括了下面的功能：
統一的訓練流程
底層用的是Hugging Face的
Trainer,所以常見的那些東西都支持:LoRA微調、DeepSpeed、FSDP、多節點Slurm集羣、4-bit量化。
訓練擴散模型和訓練transformer沒什麼區別用的都是同一套工具鏈。
統一的評估體系
評估部分基於
lm-evaluation-harness搭建,好處是不同benchmark用同一套接口,不需要針對每個模型寫推理代碼,結果也能復現。
把AR模型轉成擴散模型
這是dLLM最核心的功能，LLaMA系列模型、instruction-tuned的LLM,甚至BERT這種encoder,都能拿來微調成擴散模型。而且支持的方法包括:Masked Diffusion(MDLM)、Block Diffusion(BD3LM)和Edit Flows。
支持的模型和訓練方式
dLLM自帶了幾個參考實現:LLaDA/LLaDA-MoE、Dream、BERT-Chat、Edit Flow模型。訓練示例覆蓋預訓練、監督微調(SFT)、評估這幾個階段。
# Create environment
conda create -n dllm python=3.10 -y
conda activate dllm
# Install PyTorch (CUDA 12.4 example)
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
# Install dLLM
pip install -e .如果要跑評估:
git submodule update --init --recursive
pip install -e "lm-evaluation-harness[ifeval,math]"訓練代碼實際長什麼樣
最簡單的訓練腳本:
import transformers
import dllm
model = dllm.utils.get_model(model_args)
tokenizer = dllm.utils.get_tokenizer(model_args)
trainer = dllm.core.trainers.MDLMTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=train_data,
eval_dataset=eval_data,
args=training_args,
data_collator=transformers.DataCollatorForSeq2Seq(
tokenizer,
padding=True,
return_tensors="pt",
),
)
trainer.train()就這些，不用寫自定義loss，不用手動搞擴散循環，也不是那種只能在論文裏跑的代碼。
還可以使用LoRA + 4-bit量化微調
accelerate launch \
--config_file scripts/accelerate_configs/zero2.yaml \
examples/llada/sft.py \
--num_train_epochs 4 \
--load_in_4bit True \
--lora True推理怎麼做
擴散推理是分步驟迭代的和自迴歸的greedy decoding完全是不同的概念，dLLM用統一的sampler把這層抽象掉了:
import dllm
model = dllm.utils.get_model(model_args).eval()
tokenizer = dllm.utils.get_tokenizer(model_args)
sampler = dllm.core.samplers.MDLMSampler(
model=model,
tokenizer=tokenizer
)
inputs = tokenizer.apply_chat_template(
[{"role": "user", "content": "Explain diffusion models simply."}],
add_generation_prompt=True,
tokenize=True,
)
outputs = sampler.sample(inputs)sampler會處理mask schedule、refinement steps、decoding、output cleanup這些細節。
Edit Flows:拿擴散做文本編輯
Edit Flows算是dLLM裏比較有意思的一個方向。模型不是從零生成文本,而是學會對現有文本做操作:插入token、刪除token、替換token。這種方式特別適合代碼重構、文檔編輯、可控的文本改寫這類任務，而dLLM提供了從頭訓練Edit Flow模型的完整教程。
評估
評估擴散模型確實有點麻煩,dLLM用標準化的腳本解決這個問題。
在MMLU-Pro上跑個評估的示例如下:
accelerate launch --num_processes 4 \
dllm/pipelines/llada/eval.py \
--tasks "mmlu_pro" \
--model "llada" \
--apply_chat_template \
--num_fewshot 0總結
擴散語言模型之前一直停留在研究階段,dLLM把它變成了能實際用起來的工程工具。現有的LLM可以直接複用,微調需要的算力也不誇張,模型之間的對比有了統一標準,想做實驗也不用把整套東西重新搞一遍。
自迴歸LLM能占主導地位,很大原因是它足夠實用。擴散模型要是想在語言領域站穩腳,就要做到訓練簡單、評估方便、容易集成，dLLM在這個方向上走了不小一步。
對於在做next-gen語言模型的人來説,這個框架確實值得研究一下。
[https://avoid.overfit.cn/post/5dc5d844044d404d868bf9512bca2f9b
](https://avoid.overfit.cn/post/5dc5d844044d404d868bf9512bca2f9b)
作者：Sonu Yadav
