dLLM：復(fù)用自回歸模型權(quán)重快速訓(xùn)練擴散語言模型

大語言模型的文本生成方式一直都是以自回歸為主：一個token接一個token,從左往右,生成完就定了。

但現(xiàn)在有個不太一樣的思路開始在研究圈里流行起來，那就是擴散語言模型(Diffusion LMs)。擴散模型在圖像生成領(lǐng)域已經(jīng)證明了自己的可行性,但是問題是把這套東西用到文本上一直很麻煩——訓(xùn)練難、評估難、更別提怎么集成到現(xiàn)有的LLM工作流里了。

dLLM是一個開源的Python庫,它把擴散語言模型的訓(xùn)練、微調(diào)、推理、評估這一整套流程都統(tǒng)一了起來，而且號稱任何的自回歸LLM都能通過dLLM轉(zhuǎn)成擴散模型,還不需要什么特殊的算力。

擴散模型用在語言上有什么不同

做過圖像擴散模型的應(yīng)該能理解這個思路。

傳統(tǒng)自回歸是順序生成,擴散模型的玩法不一樣:先從噪聲或者masked tokens開始,然后一步步把整個序列細(xì)化出來。它不是一個token一個token往后走,而是對整個輸出做全局優(yōu)化。

擴散模型在幾個場景下表現(xiàn)特別好:需要復(fù)雜推理的任務(wù)、文本編輯重寫、結(jié)構(gòu)化生成,還有需要多輪迭代優(yōu)化的場景。

dLLM提供了什么

dLLM不是某個具體模型它是個框架，包括了下面的功能：

統(tǒng)一的訓(xùn)練流程

底層用的是Hugging Face的Trainer,所以常見的那些東西都支持:LoRA微調(diào)、DeepSpeed、FSDP、多節(jié)點Slurm集群、4-bit量化。

訓(xùn)練擴散模型和訓(xùn)練transformer沒什么區(qū)別用的都是同一套工具鏈。

統(tǒng)一的評估體系

評估部分基于lm-evaluation-harness搭建,好處是不同benchmark用同一套接口,不需要針對每個模型寫推理代碼,結(jié)果也能復(fù)現(xiàn)。

把AR模型轉(zhuǎn)成擴散模型

這是dLLM最核心的功能，LLaMA系列模型、instruction-tuned的LLM,甚至BERT這種encoder,都能拿來微調(diào)成擴散模型。而且支持的方法包括:Masked Diffusion(MDLM)、Block Diffusion(BD3LM)和Edit Flows。

支持的模型和訓(xùn)練方式

dLLM自帶了幾個參考實現(xiàn):LLaDA/LLaDA-MoE、Dream、BERT-Chat、Edit Flow模型。訓(xùn)練示例覆蓋預(yù)訓(xùn)練、監(jiān)督微調(diào)(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]"

訓(xùn)練代碼實際長什么樣

最簡單的訓(xùn)練腳本:

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，不用手動搞擴散循環(huán)，也不是那種只能在論文里跑的代碼。

還可以使用LoRA + 4-bit量化微調(diào)

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用統(tǒng)一的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這些細(xì)節(jié)。

Edit Flows:拿擴散做文本編輯

Edit Flows算是dLLM里比較有意思的一個方向。模型不是從零生成文本,而是學(xué)會對現(xiàn)有文本做操作:插入token、刪除token、替換token。這種方式特別適合代碼重構(gòu)、文檔編輯、可控的文本改寫這類任務(wù)，而dLLM提供了從頭訓(xùn)練Edit Flow模型的完整教程。

評估

評估擴散模型確實有點麻煩,dLLM用標(biāo)準(zhǔn)化的腳本解決這個問題。

在MMLU-Pro上跑個評估的示例如下:

accelerate launch --num_processes 4 \
dllm/pipelines/llada/eval.py \
--tasks "mmlu_pro" \
--model "llada" \
--apply_chat_template \
--num_fewshot 0

總結(jié)

擴散語言模型之前一直停留在研究階段,dLLM把它變成了能實際用起來的工程工具。現(xiàn)有的LLM可以直接復(fù)用,微調(diào)需要的算力也不夸張,模型之間的對比有了統(tǒng)一標(biāo)準(zhǔn),想做實驗也不用把整套東西重新搞一遍。

自回歸LLM能占主導(dǎo)地位,很大原因是它足夠?qū)嵱谩U散模型要是想在語言領(lǐng)域站穩(wěn)腳,就要做到訓(xùn)練簡單、評估方便、容易集成，dLLM在這個方向上走了不小一步。

對于在做next-gen語言模型的人來說,這個框架確實值得研究一下。

https://avoid.overfit.cn/post/5dc5d844044d404d868bf9512bca2f9b

作者：Sonu Yadav