AI 词典：扩散语言模型">扩散语言模型（dLLM）近年常被当作“并行生成、双向看上下文”的另一条路线。但从零预训练一个 dLLM 代价不小，更实际的做法是：拿一个已经训好的自回归（AR）模型，把它的权重原样搬过来，只改注意力方向、训练目标和采样器。这类“借用自己已有解码器”的配方，下文按惯例称 BYOD（Bring Your Own Decoder）。不同实现的细节会有差异，具体超参与脚本以对应仓库 README 和论文为准。
这篇能做出什么
照着下面走完，你会得到一个可跑通的完整链路：
1. 权重映射：一个 AR 检查点的全部参数（词嵌入、Transformer 块、输出头）无损迁移到 dLLM 上，不新增任何结构参数，因此不需要“重新学一遍语言”。
2. 注意力改造：把因果掩码换成双向掩码，并用一个可复现的扰动测试证明“因果性真的关掉了”。
3. 训练改造：吸收态（absorbing-state）离散扩散目标——随机挖掉一部分 token，让模型在 [MASK] 位置还原原文，损失按 1/t 加权。
4. 并行解码：一次前向同时解出多个位置，通过调节步数在“速度”和“质量”之间取舍。
5. 实测对照：同一模型、同一提示、同一生成长度下，对比 AR 与 dLLM 的吞吐（tokens/s、前向次数）、显存和生成质量。
最终你会得到一张自己的对照表，而不是别人论文里的数字。
前置条件清单
- 硬件：一张能装下目标模型的 GPU。流程验证阶段建议 1B 参数以下、序列长度 256–512，消费级显卡也能跑通。
- 软件：PyTorch、Transformers、Datasets、Accelerate。版本以官方文档当前版本为准，注意
attention_mask传 4 维张量时的处理逻辑在版本间有过变化。 - 模型：任意一个 HF 格式的 decoder-only 检查点（本地目录即可）。选小模型，先把流程跑通再换大模型。
- 数据：一份纯文本语料，几百 MB 足够做继续训练；验证阶段 100 条样本就够。
- 知识：看得懂 Transformer 的 attention 计算，知道 KV cache 是什么。
一个心理预期：BYOD 的“权重迁移”几乎没有损失，因为参数一一对应；“能力恢复”才是主要成本——原来的 AR 模型只见过从左往右的上下文，现在要它适应“中间有洞、两侧都能看”，还是需要一定量的继续训练。
分步骤
步骤 0：装好环境并加载基线模型
```bash
pip install torch transformers datasets accelerate
```
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
BASE = "your-local-model-dir"
tok = AutoTokenizer.from_pretrained(BASE)
ar = AutoModelForCausalLM.from_pretrained(
BASE,
torch_dtype=torch.bfloat16,
attn_implementation="eager", # 先用 eager，方便手工传 4 维掩码
device_map="cuda",
)
ar.eval()
```
attn_implementation="eager" 是为了让后面的显式掩码行为可预测；跑通之后再换 sdpa 或 flash_attention_2 做性能优化。
步骤 1：权重映射——先把“搬家”这件事做干净
BYOD 的核心优势是参数形状完全不变，所以映射逻辑就是一次 deepcopy 加一次词表扩容。
```python
import copy
dlm = copy.deepcopy(ar) # 参数逐位复制，这就是"权重映射"
MASK_TOKEN = "<|mask|>"
tok.add_special_tokens({"additional_special_tokens": [MASK_TOKEN]})
mask_id = tok.convert_tokens_to_ids(MASK_TOKEN)
dlm.resize_token_embeddings(len(tok))
with torch.no_grad():
in_emb = dlm.get_input_embeddings().weight
in_emb[mask_id] = in_emb[:-1].mean(0) # 用已有词向量均值初始化
out_emb = dlm.get_output_embeddings()
if out_emb is not None and out_emb.weight.shape[0] == len(tok):
out_emb.weight[mask_id] = out_emb.weight[:-1].mean(0)
```
为什么用均值初始化：随机初始化一个“全新”的 [MASK] 向量，会让训练前几百步的 loss 剧烈震荡，因为这个向量参与的注意力打分和输出 logits 都是噪声。用已有词向量的均值初始化，等于告诉模型“这是一个中性的占位符”，收敛会平稳很多。
验证映射是否无损：保留原模型，用带因果掩码的 4 维 mask 跑一遍 dLLM，两者 logits 应该逐位一致（词表扩容只是追加了一行，不影响已有行）。
```python
def build_4d_mask(attention_mask, dtype, causal=False):
"""attention_mask: (B, L)，1=有效 token，0=padding。返回加性 mask (B,1,L,L)。"""
B, L = attention_mask.shape
dev = attention_mask.device
valid = attention_mask[:, None, None, :].bool() # (B,1,1,L)
if causal:
allow = torch.ones(L, L, dtype=torch.bool, device=dev).tril()[None, None]
allow = allow & valid
else:
allow = valid.expand(B, 1, L, L)
m = torch.zeros(B, 1, L, L, dtype=dtype, device=dev)
return m.masked_fill(~allow, torch.finfo(dtype).min)
ids = tok("The quick brown fox jumps", return_tensors="pt").input_ids.cuda()
am = torch.ones_like(ids)
with torch.no_grad():
a = ar(input_ids=ids, attention_mask=am).logits.float()
b = dlm(input_ids=ids,
attention_mask=build_4d_mask(am, dlm.dtype, causal=True)).logits.float()
print("max diff:", (a - b).abs().max().item()) # 期望在 1e-3 量级以内（bf16 精度）
```
步骤 2：注意力改造——关掉因果性
把 causal=True 换成 causal=False，就是 dLLM 的注意力：每个位置都能看到整条序列。此时 token 之间的信息是双向流动的，这正是扩散式生成能“先填中间、再补两边”的前提。
注意：masked_fill 里 pad 位置仍然被屏蔽，因为 padding 不该参与注意力。这两件事要分开处理——很多人一起改，结果 padding 污染了 logits。
用一个扰动测试确认因果性真的没了：
```python
with torch.no_grad():
ids2 = ids.clone()
ids2[0, -1] = tok("cat").input_ids[0] # 只改最后一个 token
a2 = dlm(input_ids=ids2, attention_mask=build_4d_mask(am, dlm.dtype, causal=True)).logits.float()
b2 = dlm(input_ids=ids2, attention_mask=build_4d_mask(am, dlm.dtype, causal=False)).logits.float()
print("causal :", (a[:, :-1] - a2[:, :-1]).abs().max().item()) # 应≈0
print("bidirect :", (b[:, :-1] - b2[:, :-1]).abs().max().item()) # 应明显>0
```
第一行约等于 0，说明因果模型里“后面的 token 影响不到前面”；第二行明显大于 0，说明双向注意力生效了。这个测试比看 loss 曲线可靠得多——掩码写错时训练 loss 照样会降，只是采样出来全是重复文本。
步骤 3：训练目标——吸收态扩散
训练过程只有三步：
1. 每个样本独立采样一个掩码比例 t ~ U(0,1)（把 t 下限截到 1e-2，避免除零放大）。
2. 把这个样本中比例为 t 的 token 换成 [MASK]。
3. 只在这些被掩码的位置上算交叉熵，并按 1/t 加权。
```python
import torch.nn.functional as F
def diffusion_loss(model, input_ids, attention_mask, mask_id, t_min=1e-2):
B, L = input_ids.shape
dev = input_ids.device
t = torch.rand(B, device=dev).clamp(min=t_min) # (B,)
mask = torch.rand((B, L), device=dev) < t[:, None]
mask &= attention_mask.bool() # 不动 padding
empty = ~mask.any(dim=1) # 保证每行至少挖一个
if empty.any():
rnd = torch.randint(0, L, (B,), device=dev)
mask[empty, rnd[empty]] = attention_mask[empty, rnd[empty]].bool()
noisy = torch.where(mask, torch.full_like(input_ids, mask_id), input_ids)
logits = model(
input_ids=noisy,
attention_mask=build_4d_mask(attention_mask, model.dtype, causal=False),
).logits
ce = F.cross_entropy(logits[mask].float(), input_ids[mask], reduction="none")
w = (1.0 / t).repeat_interleave(mask.sum(dim=1)) # 与展平顺序对齐
return (ce * w).mean()
```
1/t 的来历是连续时间扩散的变分下界：掩码比例越小，剩下的上下文越多，这一步的预测越“容易”，所以要用更大的权重平衡；反之 t 大时模型几乎只能靠语言先验猜，权重小。如果 t 不截断，t→0 时权重会发散，这是新手最常见的一个 NaN 来源。
步骤 4：跑通训练循环
先不要上大规模语料，用几百条文本把流程对一个过拟合：
```python
from torch.utils.data import DataLoader
from datasets import load_dataset
ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
def encode(ex):
out = tok(ex["text"], truncation=True, max_length=512, padding="max_length")
return out
ds = ds.map(encode, batched=True, remove_columns=ds.column_names)
ds = ds.filter(lambda x: sum(x["attention_mask"]) > 32)
loader = DataLoader(ds, batch_size=4, shuffle=True)
opt = torch.optim.AdamW(dlm.parameters(), lr=2e-5, weight_decay=0.01)
dlm.train()
step = 0
for batch in loader:
ids = torch.tensor(batch["input_ids"]).cuda()
am = torch.tensor(batch["attention_mask"]).cuda()
loss = diffusion_loss(dlm, ids, am, mask_id)
loss.backward()
torch.nn.utils.clip_grad_norm_(dlm.parameters(), 1.0)
opt.step()
opt.zero_grad(set_to_none=True)
if step % 50 == 0:
print(step, loss.item())
step += 1
if step > 2000:
break
```
学习率要比预训练小一到两个数量级——这是在“改造”而不是“重训”。理想情况下 2000 步内 loss 应该稳步下降；如果前 100 步就掉到很低，反而要怀疑掩码没生效（见排错部分）。
步骤 5：并行解码采样器
采样是一个从“全 [MASK]”到“全部填满”的迭代过程。每一步做一次前向，得到所有掩码位置上的预测分布，选出置信度最高的一批位置一次性写回。
```python
@torch.no_grad()
def dlm_generate(model, prompt_ids, gen_len, steps=32, mask_id=mask_id, temperature=1.0):
dev = prompt_ids.device
P = prompt_ids.shape[1]
L = P + gen_len
x = torch.full((1, L), mask_id, dtype=torch.long, device=dev)
x[:, :P] = prompt_ids
am = torch.ones((1, L), dtype=torch.long, device=dev)
n_forward = 0
for s in range(steps):
masked = x.eq(mask_id)
if not masked.any():
break
logits = model(input_ids=x,
attention_mask=build_4d_mask(am, model.dtype, causal=False)).logits
n_forward += 1
remaining = int(masked.sum())
k = max(1, remaining // (steps - s)) # 剩余位置平均分摊到剩余步数
probs = (logits.float() / temperature).softmax(-1)
conf, pred = probs.max(dim=-1) # (1,L)
conf = conf.masked_fill(~masked, -1.0) # 只在掩码位置上选
idx = conf.topk(k, dim=1).indices
x.scatter_(1, idx, pred.gather(1, idx))
return x, n_forward
```
几个可调项：
- 步数：
steps = gen_len时接近“逐 token 生成”（但仍是双向上下文）；steps = 1时一次前向填完所有位置，速度最快、质量最差。实践里常用 8–64 步。 - 每步解多少：上面的线性分摊最直观；也可以换成余弦调度（前期少解、后期多解），或者用“置信度阈值”——只要置信度超过阈值就解，步数变成自适应。
- 温度：并行解码时同一批位置的预测来自同一次前向，温度偏高容易出现同一个词被填进多个位置。
步骤 6：KV cache 与工程细节
双向注意力的 cache 逻辑和 AR 不同，值得单独说清楚：
- 已经确定（非
[MASK]）的位置，其输入 token 不再变化，因此它们的 K/V 在后续步骤中可以复用。 - 仍是
[MASK]的位置，输入向量每步都可能变，K/V 必须每步重算。 - 但每个 query 都要对全部 key 做注意力，所以只能省“重算 K/V 的矩阵乘”，省不掉“完整注意力”。
结论：cache 有收益，但收益随掩码比例变化——最后几步几乎没有掩码时，收益接近 AR 的 cache；第一步全是掩码时，几乎没有收益。先把不带 cache 的版本跑通，再考虑这层优化。
步骤 7：实测吞吐与质量
用同一段 prompt、同样生成长度、同样 batch 做对照。AR 走 model.generate（带 KV cache），dLLM 走上面的采样器。
```python
import time
def bench(fn, warmup=3, iters=10):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(iters):
fn()
torch.cuda.synchronize()
return (time.perf_counter() - t0) / iters
prompt = tok("The history of language models", return_tensors="pt").input_ids.cuda()
GEN = 128
ar.eval()
t_ar = bench(lambda: ar.generate(prompt, max_new_tokens=GEN, do_sample=False,
pad_token_id=tok.eos_token_id))
print(f"AR : {t_ar*1000:.1f} ms, {GEN/t_ar:.1f} tok/s")
for st in (8, 16, 32, 64):
t_d, nf = None, None
def run():
global t_d, nf
_, nf = dlm_generate(dlm, prompt, GEN, steps=st)
t_d = bench(run)
print(f"dLLM s={st:3d}: {t_d*1000:.1f} ms, {GEN/t_d:.1f} tok/s, 前向次数={nf}")
```
怎么读这组数：
- 前向次数之比容易给人过高期待（32 步 vs 128 步 = 少 4 倍），但每一步都要对全长序列做前向，而且掩码位置的计算一点没省。实际墙钟加速通常明显小于步数比，小模型 + 短序列时甚至可能没有优势，因为 GPU 没被打满。
- batch 越大，并行解码的优势越明显：AR 是串行的，加 batch 只是摊薄单步开销；dLLM 的每一步本身就包含大量可并行的位置。
- 长序列（512 以上）时，注意力开销占比上升，dLLM 每步的成本更接近一次完整前向，加速比会进一步被压缩。
质量对照建议两条路：
1. 生成困惑度：用另一个 AR 参考模型给 dLLM 生成的文本打分。分数越低越像自然语言。注意这只衡量“像不像人话”，不衡量“对不对”。
2. 完形填空：从验证集里随机挖掉句中的若干 token，让模型补全，看 top-1 命中率。这个场景天然对 dLLM 有利（双侧上下文可用），所以同时也要给 AR 一个公平设置——比如只挖句尾附近的词，或者明确标注这是“双向模型的主场”。不要拿不对等的设置下结论。
一个典型的观察结果：steps 从 64 降到 16，困惑度上升不多、速度提升明显；降到 4 以下，质量会明显变差，出现词语重复和语法断裂。
常见坑与排错
1. loss 掉得飞快但采样全是重复词
十有八九是因果掩码没关掉：模型仍在做“看前文猜下一个”，而你的输入里到处是 [MASK]，它学会了“复制右侧最近的 token”。用步骤 2 的扰动测试确认。
2. loss 出现 NaN
1/t 在 t 接近 0 时爆炸。检查 t.clamp(min=1e-2) 是否生效；另外确认 ce 用的是 float() 而非 bf16。
3. padding 参与了 loss 或注意力
mask &= attention_mask.bool() 和 4 维掩码里的 valid 两处都要检查。padding 参与训练的表现是：短样本上 loss 异常低，采样时开场就是一堆空白。
4. 4 维掩码传进去没生效
不同 Transformers 版本对 4 维 attention_mask 的处理方式有差异，有些实现仍会叠加自己的因果掩码。排查办法就是步骤 2 的测试；仍不生效时，改用 attn_implementation="eager"，或者直接在 modeling 文件里把 is_causal=True 改成 False（改完记得记录你改的是哪个文件）。
5. 新增 [MASK] token 后保存模型，加载时报维度不匹配
resize_token_embeddings 改了模型，也要 tok.save_pretrained() 一起保存 tokenizer，否则下次加载时词表大小对不上。
6. 采样时多个位置填了同一个词
并行解码的固有现象。缓解手段：降低温度、对每个位置做 top-k/top-p 过滤、或者每步少解一些位置（增大 steps）。
7. 训练很久但效果不如原 AR 模型
正常。BYOD 迁移的是“语言知识”，不是“双向建模能力”。数据量太小（几千条）、学习率太大、或者序列太短，都会让它退化成“只会补高频词”。建议先用几百条数据确认能过拟合（即模型能复现训练句子），再放量。
8. 显存比 AR 高
训练时输入是全长序列且无 cache，显存占用比同长度 AR 推理高是正常的。开梯度检查点、减小 batch、缩短序列都能缓解。
下一步建议
- 块状扩散（半自回归）：把序列切成块，块内并行解码、块间从左往右。这是速度和质量之间的一个实用折中，工程改动不大。
- 采样器升级：加入 remask 机制（把低置信度的已填位置重新打回
[MASK]），或者用置信度阈值替代固定步数。 - 蒸馏减步：用一个多步采样的 dLLM 当教师，蒸馏出一个少步数的学生，这是把并行解码真正做成“低延迟”的关键一步。
- KV cache 落地：按步骤 6 的分析实现增量缓存，重点测“高掩码比例”和“低掩码比例”两种状态下的收益差异。
- 指令微调：在 dLLM 上做 SFT，观察它在补全、纠错、填空类任务上的表现是否优于同源的 AR 模型——这是 dLLM 相对有优势的战场。
- 混合推理：让 AR 负责开头（保证连贯），dLLM 负责中段批量填充，取两家之长。
走完这一遍，你手上有的不只是一个能跑的 dLLM，而是一套可复用的改造流程：换模型、换语料、换步数，都能用同一套脚本量出自己的吞吐—质量曲线。