【Task 41】Relax-SDPO / Reinforcement Learning via Self-Distillation 论文复现进展

[ZiyiTsang]

任务信息

项目	内容
任务	Task 41 - SDPO（高级）
论文	Reinforcement Learning via Self-Distillation
参考实现	lasgroup/SDPO
开发分支	sdpo
已推送基线 commit	6f794e5 test(megatron): guard unavailable backend import
对比基线	Relax ordinary OPD、Relax GRPO，以及参考 SDPO 路径
当前状态	PARTIAL：最小工程路径和定向测试已完成，正式训练矩阵和论文级结果尚未完成

当前工作区另有测试整理、AGENTS.md 和 PR_DESCRIPTION.md 的本地未提交修改；这些修改不属于上面的已推送 commit。

当前进展

目前已完成 Relax-SDPO 的最小可复用训练路径接入：

• 在 reward 之后，由 FeedbackProvider、FeedbackRecord 和 SdpoPromptBuilder 在同一 rollout group 内路由反馈或成功 response，构造 sample.teacher_prompt。
• 使用两个独立的 Ray actor 管理 student 和 teacher。teacher 使用静态 checkpoint，不做权重同步和 EMA 更新；两个 actor 可以先加载同一个 checkpoint，但实例和生命周期独立。
• 使用 TeacherPrefillBuilder 和 OpdManager 统一 plain OPD、ordinary OPSD 与 SDPO 的 teacher-prefill 生命周期，并保留原 response suffix 与 response offset 对齐。
• Relax-SDPO 只支持 student_topk：student rollout 产生 global token ids，teacher 在同一组 ids 上重打分；不在 SDPO 路径中使用 teacher_topk 或 union support。
• 新增独立 compute_sdpo_loss 与 compute_sdpo_topk_divergence，实现 student-selected Top-K 上的 KL/JSD、tail bucket、teacher detach、payload 校验和 CP-aware reduction。
• 普通 OPD 保留原有 loss、KL/JSD endpoint、sampled/Top-K/union 语义；SDPO 的 endpoint 约定单独处理，避免改变现有 OPD 数值行为。
• OPDDistributionTopKMaskEstimator、topk_layout 和 teacher-prefill 输入对象作为共享底层组件复用；普通 OPD 和 SDPO 的高层 criterion 保持分离。
• 处理 TP global token id、CP response-row layout、padding、ragged Top-K row、empty local shard 和 micro-batch 输入契约；当前不考虑 PP。
• 对缺少 feedback、successful response、student Top-K 或 teacher payload 的 SDPO 样本 fail-fast，不构造伪造的零梯度样本，也不使用 sample-level mask。
• SDPO 只支持文本输入；遇到 image、video、audio 或结构化多模态消息在 teacher 边界直接报错退出。
• 示例目录已包含数据准备、GRPO baseline、Relax-SDPO 启动脚本和 README；数据归一化入口支持参考 SDPO 数据、SciKnowEval 和 ToolAlpaca 的 train/test schema。

已完成验证

当前已完成的定向验证如下：

验证层级	结果
ordinary OPD 独立 oracle、SDPO loss、Megatron loss aggregation	40 passed
CPU Gloo 的 OPD/SDPO CP layout、reducer、TP/CP contract	6 passed
SDPO payload、prompt、mask、teacher offset、Top-K layout 等 focused tests	已通过既有 focused suite；当前工作树正在整理重复的 ordinary OPD 测试
CUDA/NCCL contract	已有双卡 GPU contract 记录；尚未替代真实模型 forward/backward 验收
GitHub CPU CI	已修复测试在无 Megatron 环境中的导入失败；完整 CI 尚未对最新本地测试整理重新运行

测试覆盖的关键不变量包括：

• ordinary OPD 重构前后的 masked JSD、KL endpoint、sampled-token gradient 和 low-variance KL 一致；
• SDPO 的 student gradient 保留，teacher target 不产生梯度；
• loss_masks 仍是 response/token-level mask，Top-K support mask 只作用于 vocabulary-level distribution；
• CP 只切 response rows，K 列和 global token ids 不切分；TP/CP 重组后与单卡 oracle 一致；
• padding、empty local shard、ragged row 和 micro-batch 切分不改变有效 token 上的统计；
• 普通 OPD 不会因为新增 SDPO teacher prompt 或 payload 字段而改变 loss；
• 缺失 SDPO teacher payload 会在 loss 前失败，而不是静默参与更新。

与论文原始方案的差异

当前实现遵循已经确定的最小方案，不等同于论文的完整 EMA self-teacher 复现：

论文方案	当前 Relax-SDPO
EMA Student 作为 self-teacher	独立静态 teacher actor
EMA teacher 参数更新	不实现 EMA，不同步 teacher 权重
根据有效反馈动态训练	已实现 reward 后的 feedback/successful-response teacher prompt 路由
teacher 对增强上下文下的 response 进行重打分	已实现 teacher 对原 response 的 student Top-K ids 重打分
student/teacher Top-K 近似蒸馏	当前只支持 student-selected Top-K
论文完整训练和 benchmark	尚未完成正式训练矩阵和效果对比

因此当前可以称为 Relax-SDPO 的静态 teacher 最小实现，不能把结果描述为论文原始 EMA teacher 的完整复现结果。

正式复现方案

模型

• Student：Qwen3-4B-Instruct-2507。
• Teacher：同一 checkpoint 的独立静态 teacher 实例，或用户后续指定的独立 teacher checkpoint。
• 初始训练设备：双卡训练环境，TP/CP 路径；不考虑 PP。

数据

训练阶段优先使用参考实现 lasgroup/SDPO 已整理的数据格式，并通过
examples/on_policy_distillation/sdpo/prepare_data.py 转换为 Relax JSONL：

• SciKnowEval L3 推理子集：Chemistry、Physics、Biology、Materials；
• ToolAlpaca：train/test split，用于工具规范与用户请求到工具调用的映射；
• 参考 SDPO 的 tooluse 数据作为可选兼容输入；
• 当前只计划训练和验证上述单轮文本任务，不纳入 LiveCodeBench 训练或评测，也不引入 code executor。

数据来源记录：

• feedback 构造和参考 split：https://github.com/lasgroup/SDPO；
• SciKnowEval：https://github.com/HICAI-ZJU/SciKnowEval；
• ToolAlpaca：https://huggingface.co/datasets/Ahren09/ToolAlpaca。

示例 reward 已覆盖 SciKnowEval 的答案判定和 ToolAlpaca 的工具名/JSON 参数判定；错误样本生成反馈文本，正确样本可作为同 group 的 successful response context。

训练设置

• 普通 baseline：Relax GRPO。
• SDPO：--opd-loss-mode sdpo、--opd-token-selection student_topk、独立静态 teacher、tail Top-K 分布。
• SDPO loss：支持 forward_kl、reverse_kl 和 JSD；JSD 使用参考路径的 endpoint 约定，alpha=0.5 为对称 JSD。
• 并行：TP/CP；Top-K payload 使用 CP response-row layout，不切 K 维。
• 数据契约：有效反馈或 successful response、完整 student Top-K ids 和 teacher log-prob payload 必须存在。
• 计划 seeds：42、1234、2026。
• 计划统计：最终准确率、KL/JSD、Top-K coverage、有效反馈比例、训练吞吐、峰值显存和 teacher 额外计算耗时。

正式实验当前状态

正式训练矩阵尚未完成。当前没有可用于填写论文结果的稳定训练曲线、三 seed 均值/方差或 GRPO 对比结果。

数据集	Split	Student	Teacher	Seeds	GRPO	Relax-SDPO	有效反馈比例	Top-K coverage	吞吐/显存	状态
SciKnowEval Chemistry	train/test	Qwen3-4B-Instruct-2507	static teacher	42/1234/2026						待运行
SciKnowEval Physics	train/test	Qwen3-4B-Instruct-2507	static teacher	42/1234/2026						待运行
SciKnowEval Biology	train/test	Qwen3-4B-Instruct-2507	static teacher	42/1234/2026						待运行
SciKnowEval Materials	train/test	Qwen3-4B-Instruct-2507	static teacher	42/1234/2026						待运行
ToolAlpaca	train/test	Qwen3-4B-Instruct-2507	static teacher	42/1234/2026						待运行
Mean ± std.	—	—	—	3 seeds						待运行

由于当前 teacher 不更新，EMA update count/time 在本实现中为 N/A，不应填成 0 或声称已经完成。

待完成

• 用实际本地 SciKnowEval 四个 domain 和 ToolAlpaca train split 生成非空 Relax JSONL，逐项核对 metadata、label、reward 和 feedback；
• 在双卡环境完成一次包含有效 feedback 和 successful response 的 Relax-SDPO 真实训练闭环，确认两个 actor、teacher prefill、student Top-K transfer、Megatron backward 全部连通；
• 验证 TP/CP、padding、ragged response、micro-batch 和 empty local shard 在真实模型 forward/backward 上的统计一致性；
• 记录有效反馈比例、KL/JSD、Top-K coverage、teacher extra forward 次数和耗时；EMA 指标按当前静态 teacher 方案标记为 N/A；
• 在相同模型、数据、rollout 数和预算下运行 GRPO 与 Relax-SDPO，完成 3 个 seed 的均值/方差、准确率、吞吐和峰值显存比较；
• 运行完整 Linux 依赖环境中的全仓 pytest、pre-commit 和 GitHub CI；当前 gitleaks hook 仍受本地网络/工具安装条件影响，不能把 focused hooks 结果当作全量 pre-commit 通过；
• 在正式结果完成后更新 PR 材料，并将所有“已完成”描述与真实日志、commit 和实验表对齐。

已知限制与风险

• 当前是静态 teacher 最小实现，不包含论文 EMA Student self-teacher；这是实现范围的明确选择，不是隐藏的未完成参数。
• 缺少反馈或 teacher payload 会 fail-fast；这会暴露数据/reward/teacher prefill 问题，但不会自动过滤样本或生成零梯度目标。
• SDPO 当前只支持文本、student_topk、TP/CP 和无 PP；普通 OPD 的既有能力不因该限制改变。
• 真实 Ray/SGLang teacher prefill、Megatron TP/CP forward/backward 和正式训练效果仍未完成验收。
• 当前 smoke 中若样本没有 feedback 或 successful response，会按数据契约在 prompt routing 阶段退出；这不是成功的 SDPO 训练结果。

交付判断

当前状态：PARTIAL。

Relax-SDPO 的最小工程路径、普通 OPD 隔离、反馈 teacher prompt、student Top-K 重打分、独立 KL/JSD loss、TP/CP layout 和定向测试已经具备 Draft 审查条件；但由于 EMA teacher、真实训练闭环、三 seed 正式结果和完整 CI/硬件验收尚未完成，当前不能标记为 Task 41 完成交付，也不能宣称已经复现论文的稳定性主结果。