Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

基于 SEDD 的 SFT 与 SEPO 提交包

本包是离散扩散语言模型 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/


实验结果

SFT 训练

项目 配置
初始化 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

训练曲线:

SFT training loss

训练 loss 的滑动趋势从约 0.90 下降至约 0.61,说明模型对条件 score 的拟合整体改善。

SFT all-mask validation

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 初始化。

SEPO 训练

Smoke 验证(100 update)

指标 数值
完成 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 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.jsonresults/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,在保持目标的同时降低方差。

参考文献

  1. Lou, A., Meng, C., Ermon, S. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution. arXiv:2310.16834v3, 2024.
  2. Zekri, O., Boullé, N. Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods. arXiv:2502.01384v3, NeurIPS 2025.
  3. Taori, R. et al. Stanford Alpaca: An Instruction-following LLaMA Model. 2023.
  4. Bai, Y. et al. Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback. 2022.

About

homework

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages