本包是离散扩散语言模型 SEDD(Score Entropy Discrete Diffusion)的指令微调与强化学习后训练完整链路:Alpaca 指令数据处理 → SEDD SFT → SFT 评估 → HH-RLHF 数据处理 → SEPO → SEPO 评估。
SEDD 不学习 next-token probability,而是学习相邻离散状态之间的概率比值(concrete score)。本项目从 Score Entropy 的逐点最优性出发,完成了条件 SFT 的数学分析、mask 审计、训练诊断、生成评估和基于 concrete score 的 PPO 训练。
包内不包含模型权重、数据集、处理后的张量、checkpoint、测试集、调试脚本或 W&B 缓存;results/ 只保存可审阅的 JSON 结果和训练曲线。
请先下载以下资源,并在 .env 中填写绝对路径:SEDD Medium、Alpaca JSON、HH-RLHF 原始数据、GPT-2 Medium tokenizer/model、Helpful GPT-2 Large reward model,以及 MMLU 子集 JSON。当前 SEDD Medium 运行时强依赖 FlashAttention。
先安装与机器 CUDA 匹配的 PyTorch,再安装其余依赖:
pip install -r requirements.txt
cp .env.example .env
# 只修改 .env 中的路径变量。所有公开入口统一为:
bash scripts/run_pipeline.sh --stage check
bash scripts/run_pipeline.sh --stage prepare-sft
bash scripts/run_pipeline.sh --stage train-sft
bash scripts/run_pipeline.sh --stage eval-sft
bash scripts/run_pipeline.sh --stage prepare-rl
bash scripts/run_pipeline.sh --stage train-rl
bash scripts/run_pipeline.sh --stage eval-rl--stage all 会按相同顺序执行。SEPO 固定使用 SFT 生成的 $OUTPUT_DIR/sft/checkpoints/best_step_9000.pt;最终 SFT 配置每 250 step 验证一次,因此该 checkpoint 会在 12,400 step 训练中生成。每一阶段都会先检查资源,缺失时会指出具体路径,且不会先占用显卡。
third_party/sedd/:最小 SEDD 运行时及许可证。sedd_sft/:SFT 数据读取、loss、训练和生成评估。sepo_rl/:SEPO actor-critic、冻结奖励模型适配器和分片数据读取。configs/:最终 SFT 与 SEPO 超参数。scripts/:数据处理、资源检查、训练、评估和总入口。results/:SFT/RL 的三组输出、MMLU 与训练曲线。
具体数学目标、资源约定和结果来源见 docs/。
| 项目 | 配置 |
|---|---|
| 初始化 | SEDD Medium Base(约 4.24 亿参数) |
| 数据 | Alpaca legacy ~50k / 1k / 1k |
| 最大长度 | 384(prompt 上限 128) |
| 目标 | token-mean unweighted conditional DSE |
| all-MASK 比例 | 20% |
| 有效 batch | 32(micro 8 × 累计 4) |
| 学习率 | 2×10⁻⁵ cosine → 2×10⁻⁶ |
| 总更新数 | 12,400(约 8 epoch) |
| 保存策略 | valid_all_mask top-3 |
训练曲线:
训练 loss 的滑动趋势从约 0.90 下降至约 0.61,说明模型对条件 score 的拟合整体改善。
all-MASK 验证曲线由约 0.645 下降到约 0.568,后期略有反弹。
Top-3 checkpoint(valid_all_mask,越低越好):
| 排名 | Step | valid_all_mask |
|---|---|---|
| 1 | 9000 | 0.5406 |
| 2 | 10500 | 0.5551 |
| 3 | 8500 | 0.5592 |
Step 9000 被选为 demo 模型和 SEPO actor 初始化。
| 指标 | 数值 |
|---|---|
| 完成 update | 100 |
| 平均 step time | 4.72 s |
| 峰值 CUDA 显存 | 17.75 GiB |
| Advantage std | 0.9403 |
| 参数变化量 | 0.000104 |
| Tail KL mean | 0.00305 |
| Tail clipped fraction | 0.0 |
所有 smoke gate 通过:actor 参数变化、advantage 有信号、数值有限、KL/ratio 稳定、clipping 未饱和、梯度受控。
| 项目 | 配置 |
|---|---|
| Actor 初始化 | SFT step 9000 |
| 数据 | HH-RLHF prompts |
| Completion tokens | 128 |
| Denoising steps | 128 |
| PPO epochs | 2 |
| SNIS samples | 4 |
| PPO epsilon | 0.2 |
| KL beta | 0.02 |
| Actor LR | 1×10⁻⁶ |
| Critic LR | 9×10⁻⁶ |
| Gradient clip | 1.0 |
训练曲线:
| Reward | KL | Actor Loss |
|---|---|---|
![]() |
![]() |
![]() |
- Reward:早期从约 0.2 提高到 0.48,随后回落,峰值出现在约 800-1400 step。
- KL:从接近 0 持续增长,后期超过 0.2,表明 actor 逐渐偏离 SFT 参考策略。
- Actor Loss:从约 -1 快速下降到约 -20,结合 reward 回落和 KL 增长,更符合策略漂移而非正常收敛。
Top-3 checkpoint(held-out reward,越高越好):
| 排名 | Step | Held-out Reward |
|---|---|---|
| 1 | 800 | 1.0679 |
| 2 | 1400 | 1.0530 |
| 3 | 1600 | 0.9868 |
约 2.4k update 后出现 reward 回落、KL 增长和 actor loss 绝对值急剧增大,触发 early stopping。保留 step 800、1400、1600 三个早期 checkpoint。
| 指标 | Base | SFT-9000 | PPO-1400 | 说明 |
|---|---|---|---|---|
| MMLU accuracy | 8.4% | 18.4% | 17.0% | SFT 显著改善知识+格式;PPO 未进一步 |
| Invalid output rate | 58.2% | 16.0% | 14.4% | 格式遵循大幅提升 |
| Mean output tokens | 128.0 | 88.7 | 80.8 | 输出更短,拖尾减少 |
| Reference word F1 | 0.1126 | 0.2084 | 0.3404 | PPO 更贴近参考措辞 |
McNemar 检验 SFT vs PPO 的 MMLU correctness:exact p = 0.4638,无显著差异。
关键结论:
- SFT 相比 Base,有效输出大幅增加,回答开始遵循指令格式和主题。
- PPO 相比 SFT,输出更短、F1 更高,部分开放题更直接,但未可靠提升知识准确率。
- Helpful RM 优化的是 "helpful 风格",不能替代通用事实正确性评测。
基于 results/sft/step_9000/chat_10q.json 和 results/rl/step_1400/chat_10q.json:
- Base 模型:对指令基本无响应,输出跑题、重复或无关文本。
- SFT step 9000:恢复任务格式和主题相关性,能给出定义、故事、列表等结构化回答,但存在杜撰词(如
pinjabbing)和语义错误。 - PPO step 1400:表达更直接、格式更规范(如编号列表),词汇更接近参考回答,但长文本规划能力仍有限,代码生成等任务停留在表面关键词层面。
| 现象 | 根因 | 修复 |
|---|---|---|
| loss 近似横线,grad 剧烈跳变 | dsigma 在 t→1 时权重爆炸(Θ(ε⁻¹)) |
使用 uniform-t 无权重 DSE |
| loss 下降但生成混乱 | padding 被纳入监督,mask 边界错误 | 审计四个 tensor mask,100 条过拟合验证 |
| CUDA OOM | 多个 loss 辅助项产生额外大计算图 | 合并目标,控制 auxiliary forward |
| categories > 2²⁴ | 位置×词表邻居展平后 multinomial 超限 | 按位置分层分块处理 |
| PPO 后期 reward collapse | KL 漂移 + RM 偏置 | top-3 早停 + 独立生成评测 |
- 模型仅约 0.4B 参数,基础知识和复杂推理能力有限。
- 无权重 DSE 是工程折中,不再具有 DWDSE likelihood-bound 的严格解释。
- Helpful RM 与通用指令质量不完全对齐,建议引入独立 judge 或多奖励模型。
- 长训练容易 KL 漂移和 reward hacking,可引入自适应 β 和更短 rollout 视野。
- 建议使用 uniform-σ 重新验证严格 DWDSE,在保持目标的同时降低方差。
- Lou, A., Meng, C., Ermon, S. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution. arXiv:2310.16834v3, 2024.
- Zekri, O., Boullé, N. Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods. arXiv:2502.01384v3, NeurIPS 2025.
- Taori, R. et al. Stanford Alpaca: An Instruction-following LLaMA Model. 2023.
- Bai, Y. et al. Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback. 2022.




