Axolotl Diffusion LM 训练插件基于 LLaDA 随机掩码范式的大语言模型预训练与指令微调实战【免费下载链接】axolotlGo ahead and axolotl questions项目地址: https://gitcode.com/GitHub_Trending/ax/axolotl导读本文系统讲解 Axolotl 内置的扩散语言模型Diffusion LM训练插件axolotl.integrations.diffusion.DiffusionPlugin它以 LLaDALarge Language Diffusion Models为思路用随机掩码 双向注意力取代自回归的下一个 Token 预测覆盖从配置启用、预训练/指令微调、训练中样本生成到 CLI 交互式推理的完整链路。读完本文你将掌握 diffusion 训练块的全部参数语义、其底层前向加噪与去噪损失的计算原理以及如何用现成命令跑通 Llama-3.2-1B 的扩散预训练与 SFT。什么是 Diffusion LM从自回归到随机掩码传统的自回归语言模型在每一步只能看到左侧上下文而 LLaDA 风格的扩散语言模型把生成文本重新定义为从纯噪声全掩码序列逐步去噪到完整文本的扩散过程。Axolotl 的 Diffusion 插件在训练阶段采用了 LLaDA 的三大核心设计随机 Token 掩码Random token masking训练时不再做 next-token prediction而是按随机的掩码概率把序列中的部分 Token 替换为专用的[MASK]token要求模型根据上下文补全被遮蔽的位置双向注意力Bidirectional attention掩码后的序列用双向注意力编码让每个位置都能同时参考左右两侧的上下文这正是模型能猜出被掩 Token 的基础重要性加权Importance weighting按掩码概率p对每个掩码位置的损失进行loss / p的加权得到无偏的扩散损失估计使训练在不同掩码率下保持稳定。从源码结构看该功能是一个标准的 Axolotl 插件代码全部位于 src/axolotl/integrations/diffusion 目录核心模块包括模块职责plugin.pyDiffusionPlugin注册配置参数并挂载自定义 Trainerargs.pyPydantic 模型定义嵌套的diffusion:配置块trainer.pyDiffusionTrainer覆写compute_loss实现扩散损失generation.py反向扩散采样生成训练中样本生成与推理共用callbacks.py训练过程中定时生成样本并记录日志utils.py掩码 token 解析、双向注意力掩码、logits 对齐等工具安装与快速开始插件随 Axolotl 一同分发无需额外安装包直接使用项目自带的两个开箱即用示例Llama-3.2-1B即可体验# 扩散预训练Pretrain axolotl train examples/llama-3/diffusion/pretrain-1b.yaml # 扩散指令微调SFT axolotl train examples/llama-3/diffusion/sft-1b.yaml对应的真实配置文件位于 examples/llama-3/diffusion/pretrain-1b.yaml预训练使用pretraining_datasetwikitext-103-raw-v1和 examples/llama-3/diffusion/sft-1b.yamlSFT使用teknium/GPT4-LLM-Cleanedalpaca 数据集。在已有配置上启用扩散训练你不需要从零写配置只需在任意 Axolotl 配置中做两处修改通过plugins字段挂载插件# 启用 diffusion LM 训练插件 plugins: - axolotl.integrations.diffusion.DiffusionPlugin配置嵌套的diffusion块下面展示的是默认值diffusion: noise_schedule: linear # 或 cosine min_mask_ratio: 0.1 max_mask_ratio: 0.9 num_diffusion_steps: 128 eps: 1e-3 importance_weighting: true # Mask token训练时若缺失会自动添加请避免使用 pad/eos mask_token_str: |diffusion_mask| # 或直接使用已有特殊 token 的 id例如 Llama-3.x 的 128002 # mask_token_id: 128002 # 训练过程中采样生成可选 generate_samples: true generation_interval: 100 num_generation_samples: 3 generation_steps: 128 generation_temperature: 0.0 generation_max_length: 100配置项的完整定义、默认值、取值范围与约束都集中在 args.py 中下一节逐一展开。配置参数全解diffusion 嵌套块以下参数由 Pydantic 模型DiffusionConfig见 args.py解析校验所有字段均有默认值因此即使只写diffusion: {}也能用默认配置跑起来。噪声调度Noise Schedule参数默认值类型/范围说明noise_schedulelinearlinear/cosine噪声调度类型决定掩码概率随时间步的变化曲线min_mask_ratio0.1float[0.0, 1.0]最小掩码比例max_mask_ratio0.9float[0.0, 1.0]最大掩码比例num_diffusion_steps128int 1扩散时间步数量也是反向去噪时的迭代步数eps1e-3float[0.0, 1.0]前向过程最小掩码概率的下界防止掩码率降为 0校验器会在min_mask_ratio max_mask_ratio时抛出ValueErrorargs.py所以两个比例必须保持min ≤ max。训练配置Training参数默认值说明importance_weightingtrue是否对损失按掩码概率做重要性加权loss / p关闭后改为对掩码 token 取平均mask_token_idNone掩码 token 的 id未设置或无效时走字符串解析/自动添加逻辑mask_token_strNone掩码 token 字符串若mask_token_id无效或未设置会确保该 token 以 additional special token 形式存在并使用若两者都缺省自动添加默认的|diffusion_mask|mask_token_id的解析逻辑在 utils.py 的 resolve_mask_token_id 中解析优先级为显式配置的mask_token_id且小于词表大小直接采用尝试把mask_token_str或默认字符串映射为已注册的 special token id训练阶段allow_addTrue若 tokenizer 支持自动把该字符串加入additional_special_tokens并调用model.resize_token_embeddings(len(tokenizer))扩充嵌入层兜底回退到unk_token_id或 0。注意推理阶段allow_addFalse见 cli/utils/diffusion.py不会修改 tokenizer 与模型词表因此推理配置中最好显式给出mask_token_id。训练中采样生成Sample Generation参数默认值范围说明generate_samplestruebool是否在训练过程中生成样本用于观测generation_interval100int 1每隔多少训练步生成一次样本num_generation_samples3int 1每次生成的样本条数generation_steps128int 1反向扩散去噪的迭代步数generation_temperature0.0float 0.0采样温度0.0为确定性贪心generation_max_length100int 1生成序列的最大长度底层原理前向加噪、扩散损失与双向注意力本节的实现依据是 trainer.py 与 utils.py它们精确对应 README 中随机掩码与扩散损失两节描述的公式。随机掩码Forward Process每个训练 batch 内对 batch 中的每条样本独立采样一个随机时间步t ∈ [0, 1]掩码概率按p (1 - eps) * t eps计算eps保证最小掩码率即 README 中的公式。源码实现见_forward_processtrainer.py其中还有几处工程细节不掩 padding若提供了attention_maskpadding 位置被排除不掩特殊 tokenbos/eos/pad等特殊 token 通过缓存的_special_token_ids集合排除trainer.pySFT 只掩答案当传入labels即 SFT 数据时掩码位置还必须满足labels ! -100即只对答案部分做掩码提示词部分保持原样trainer.py最后用mask_token_id填充被掩位置得到加噪序列noisy_batch。扩散损失Diffusion LossDiffusionTrainer覆写了compute_losstrainer.py核心计算在_compute_diffusion_losstrainer.py对noisy_batch做一次前向得到 logits并调用shift_logits_to_input_positions将 next-token logits 对齐到输入位置因为双向注意力下位置i的 logits 应预测位置i本身而非i1见 utils.py只取掩码位置上的 logits 与目标计算F.cross_entropy(..., reductionnone)若开启重要性加权按token_loss / p_mask加权归一化方式随训练模式不同而不同非 SFT预训练开启加权时使用无偏估计sum(loss / p) / total_tokens关闭加权时对掩码 token 取平均以在变化的掩码率下保持尺度稳定trainer.pySFT按每条样本的答案 token 数归一化后求平均trainer.py。这正对应 README 中给出的损失公式loss sum(cross_entropy(pred, target) / p_mask) / total_tokens。双向注意力掩码扩散训练必须打破因果掩码。create_bidirectional_attention_maskutils.py构造[batch, 1, seq_len, seq_len]的 4D 掩码没有attention_mask时是全局全连接有attention_mask时位置i与j可互相注意当且仅当它们共享相同的非零段 id这样既能排除 padding也能在sample_packing: true打包多条样本时天然阻断跨样本注意力——这也解释了 README支持 4D attention mask 的模型可开箱即用的说法。训练过程中的样本生成当diffusion.generate_samples: true时post_set_axolotl_cfgtrainer.py会把DiffusionGenerationCallback挂到 Trainer 上。该回调在on_step_end中按generation_interval触发callbacks.py优先从 eval dataloader、否则从 train dataloader 随机抽取序列交给generate_samplesgeneration.py执行反向扩散Sample 1: Original (45 tokens): The quick brown fox jumps over the lazy dog... Masked (18/45 tokens, 40.0%): The [MASK] [MASK] fox [MASK] over [MASK] lazy [MASK]... Generated: The quick brown fox jumps over the lazy dog...生成结果会同时输出到控制台带颜色标记绿色正确补全、红色错误补全、暗色未改动并写入 wandb若开启use_wandb以generated_samples表格记录见 callbacks.py。反向去噪是怎么一步步揭晓的generate/_diffusion_stepgeneration.py实现了逐步揭晓progressive unmasking每个去噪步先对当前序列前向一次仅在被掩位置取 logitstemperature 0时缩放 logits 并叠加 Gumbel 噪声采样否则直接 argmax 贪心同时把输出中预测为mask_token_id的位置抑制为-inf防止模型把掩码 token 当答案输出除最后一步外每一步只揭晓remaining_masked / (num_diffusion_steps - step)比例的 token并且只揭晓置信度最高的那些按 softmax 概率 topk 选择最后一步一次性揭晓全部剩余位置。推理CLI 交互式扩散生成扩散推理已集成进标准 Axolotl CLIinference.py会自动检测配置中的 Diffusion 插件并切换到扩散推理模式见 src/axolotl/cli/inference.py。使用与训练相同的配置文件即可axolotl inference path/to/your-config.yaml可选加--gradio参数启动一个简单的 Web 界面Gradio UI 实现在 src/axolotl/cli/utils/diffusion.py支持 random / completion 两种模式切换、掩码率滑块与结果 HTML 着色展示。命令行交互支持在提示词前加前缀指令可链式使用解析逻辑见 cli/utils/diffusion.py 的 _parse_commands:complete N→ 补全模式在末尾追加 N 个新掩码 token默认 64:mask R→ 随机掩码模式目标掩码率 R ∈ [0.0, 1.0]示例会话 Commands: :complete N - completion mode with N tokens (default 64) :mask R - random masking with ratio R (0.0–1.0) Give me an instruction (Ctrl D to submit): :mask 0.4 The quick brown fox jumps over the lazy dog Masked (40.0%): The [MASK] brown [MASK] jumps over the [MASK] dog Generated: The quick brown fox jumps over the loud dog推理时同样走generate()generation.py先按模式构造掩码序列SFT 场景还会识别 labels 中-100标记只掩答案再迭代num_diffusion_steps步去噪最后输出带掩码率统计与逐 token 着色对比的生成结果。注意推理用allow_addFalse解析掩码 token因此推理配置中应确保mask_token_id或mask_token_str能直接命中现有词表。训练监控指标_compute_diffusion_loss会通过store_metrics记录一组指标trainer.pyREADME 列出的指标与源码一一对应指标含义train/loss加权后的扩散损失train/accuracy掩码 token 上的预测准确率argmax 对比train/mask_ratio平均掩码 token 占比train/num_masked_tokens被掩码的 token 数量train/avg_p_mask平均掩码概率train/ce_loss未加权的交叉熵损失train/importance_weight_avg平均重要性权重开启加权时记录此外SFT 训练还会额外记录answer_ratio答案 token 占比与avg_answer_length平均答案长度便于观察指令数据的答案分布。模型兼容性与已知限制兼容性任何支持 4D 注意力掩码的模型理论上都可以直接使用掩码构造见上文双向注意力掩码一节若遇到不兼容的模型README 建议到仓库提交 issue 或 PR 协助适配。已知限制README 明确列出不支持 Flash Attention示例配置中统一使用attn_implementation: sdpa扩散训练需以支持双向 4D 掩码的实现运行不支持 RL 训练当前插件仅覆盖 SFT/预训练形式的扩散训练无法与 RL 训练管线组合。实战配置示例预训练与 SFT 对照仓库自带的两个示例配置展示了 Pretrain 与 SFT 两种典型用法可直接复制改造。扩散预训练examples/llama-3/diffusion/pretrain-1b.yaml要点使用pretraining_dataset加载 wikitextnoise_schedule: cosine、min/max_mask_ratio: 0.15/0.85、eps: 5e-4显式mask_token_id: 128002Llama-3.x 已有的特殊 tokensample_packing: true打包attn_implementation: sdpabf16: autotf32: true混合精度训练。扩散指令微调examples/llama-3/diffusion/sft-1b.yaml要点使用datasets alpaca 格式数据val_set_size: 0.05划分验证集eval_sample_packing: true保证评估时打包一致性save_strategy: besteval_strategy: epoch按验证集选择最优 checkpoint。由于 SFT 只掩码答案 tokenlabels ! -100扩散模型实际被训练为在提示词条件下补全答案。两套配置都在special_tokens.pad_token中指定了|end_of_text|避免与扩散掩码 token 混淆。质量保障端到端冒烟测试仓库在 tests/e2e/test_diffusion.py 中提供了两个端到端测试可作为自定义配置的参考基准test_diffusion_smoke_test用 SmolLM2-135M 小型 alpaca 数据集完整跑通插件加载、配置校验、数据集加载、训练与模型产物检查generate_samples/generation_interval/generation_steps等参数全部置为小值以加速test_diffusion_sft_labels专门验证 SFT 场景下数据集的labels正确包含-100提示词标记且扩散训练能正确处理这些标签。这两个测试也证明了插件接入 Axolotl 主流程的路径prepare_plugins → validate_config → normalize_config → load_datasets → train与普通训练完全一致只是compute_loss被替换为扩散损失。总结Axolotl 的 Diffusion LM 插件以 LLaDA 的随机掩码 双向注意力 重要性加权范式为核心把扩散训练无缝嵌入既有 Axolotl 工作流一份plugins声明、一个diffusion嵌套块即可在预训练或 SFT 场景下开启扩散训练并复用 CLI/Gradio 完成交互式扩散推理。阅读 README 与 args.py 可以快速掌握全部参数深入 trainer.py、generation.py 与 utils.py 则能完全理解其掩码调度、损失归一化与逐步揭晓的采样机制为二次定制打下基础。【免费下载链接】axolotlGo ahead and axolotl questions项目地址: https://gitcode.com/GitHub_Trending/ax/axolotl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考